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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/188951

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.

>From 5f0ca700afb9e9aa1fc014c9563bcc10f26abb46 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Mar 2026 15:57:49 -0700
Subject: [PATCH] [MLIR][Vector] Fix WarpOpScfForOp and WarpOpScfIfOp leaving
 invalid ops after region moves

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.
---
 .../Vector/Transforms/VectorDistribute.cpp    | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)

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) {