[Mlir-commits] [mlir] [MLIR][Vector] Fix WarpOpScfForOp and WarpOpScfIfOp leaving invalid ops after region moves (PR #188951)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-vector

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

WarpOpScfForOp::matchAndRewrite called mergeBlocks() to move forOp's body block into the inner WarpOp. mergeBlocks() erases the source block, leaving forOp with an empty body region (0 blocks). Since scf.for requires exactly 1 body block, IR verification fails with "region with 1 blocks" after the pattern succeeds. Additionally, when forOp had no init args, the pattern was missing the scf.yield terminator in the new ForOp.

WarpOpScfIfOp::matchAndRewrite had the same issue: takeBody() emptied the ifOp's then/else regions, leaving scf.if with 0 blocks.

Fix:
- Restore the conditional scf.yield creation (only when newForOp has results).
- After merging/taking the regions, replace the remaining op's results with ub.poison and erase the now-invalid op from the new WarpOp's body.

Assisted-by: Claude Code
Fix a failure present with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.

---
Full diff: https://github.com/llvm/llvm-project/pull/188951.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp (+31) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index b4d500212c770..31d875e3a67de 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -12,6 +12,7 @@
 #include "mlir/Dialect/GPU/Utils/DistributionUtils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"
 #include "mlir/IR/AffineExpr.h"
@@ -2000,6 +2001,21 @@ struct WarpOpScfIfOp : public WarpDistributionPattern {
     for (auto [origIdx, newIdx] : ifResultMapping)
       rewriter.replaceAllUsesExcept(newWarpOp.getResult(origIdx),
                                     newIfOp.getResult(newIdx), newIfOp);
+
+    // The original `ifOp` was left inside `newWarpOp` with empty then/else
+    // regions (their blocks were moved into the inner WarpOps by takeBody).
+    // Replace its results with poison and erase it to restore IR validity.
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPoint(ifOp);
+      for (OpResult result : ifOp.getResults()) {
+        Value poison =
+            ub::PoisonOp::create(rewriter, ifOp.getLoc(), result.getType());
+        rewriter.replaceAllUsesWith(result, poison);
+      }
+      rewriter.eraseOp(ifOp);
+    }
+
     return success();
   }
 
@@ -2215,6 +2231,21 @@ struct WarpOpScfForOp : public WarpDistributionPattern {
     for (auto [origIdx, newIdx] : forResultMapping)
       rewriter.replaceAllUsesExcept(newWarpOp.getResult(origIdx),
                                     newForOp.getResult(newIdx), newForOp);
+
+    // The original `ForOp` was left inside `newWarpOp` with an empty body
+    // region (its body block was moved into `innerWarp` by `mergeBlocks`).
+    // Replace its results with poison and erase it to restore IR validity.
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPoint(forOp);
+      for (OpResult result : forOp.getResults()) {
+        Value poison =
+            ub::PoisonOp::create(rewriter, forOp.getLoc(), result.getType());
+        rewriter.replaceAllUsesWith(result, poison);
+      }
+      rewriter.eraseOp(forOp);
+    }
+
     // Update any users of escaping values that were forwarded to the
     // inner `WarpOp`. These values are now arguments of the inner `WarpOp`.
     newForOp.walk([&](Operation *op) {

``````````

</details>


https://github.com/llvm/llvm-project/pull/188951