[Mlir-commits] [mlir] 4575c05 - [mlir][scf] fix loop specialization crash when an scf.for result is used (#236046)

[llvmlistbot at llvm.org]

Author: Prashanth T
Date: 2026-11-01T15:57:55+08:00
New Revision: 4575c053e0519b968ad6117b5df4d248326725a7

URL: https://github.com/llvm/llvm-project/commit/4575c053e0519b968ad6117b5df4d248326725a7
DIFF: https://github.com/llvm/llvm-project/commit/4575c053e0519b968ad6117b5df4d248326725a7.diff

LOG: [mlir][scf] fix loop specialization crash when an scf.for result is used (#236046)

the replacement `scf.if` now has the original loop result types, if the
condition is satisfied (i.e bound equal to constant) the then branch
replaces the upper bound of the cloned loop with the constant, else the
original operation with bound is cloned as such.

the yield op is created inside both the branches if there is any
resullt, the original op's result users are redirected to the `scf.if`
results and the original loop is erased.

fixes #235552

assisted by: claude

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
    mlir/test/Dialect/SCF/for-loop-specialization.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index e232692ecc4e5a..6ca907415fbb57 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -103,9 +103,19 @@ static void specializeForLoopForUnrolling(ForOp op) {
   Value cond = arith::CmpIOp::create(b, op.getLoc(), arith::CmpIPredicate::eq,
                                      bound, constant);
   map.map(bound, constant);
-  auto ifOp = scf::IfOp::create(b, op.getLoc(), cond, /*withElseRegion=*/true);
-  ifOp.getThenBodyBuilder().clone(*op.getOperation(), map);
-  ifOp.getElseBodyBuilder().clone(*op.getOperation());
+  auto ifOp = scf::IfOp::create(b, op.getLoc(), op->getResultTypes(), cond,
+                                /*withElseRegion=*/true);
+
+  OpBuilder thenB = ifOp.getThenBodyBuilder();
+  Operation *thenLoop = thenB.clone(*op.getOperation(), map);
+  OpBuilder elseB = ifOp.getElseBodyBuilder();
+  Operation *elseLoop = elseB.clone(*op.getOperation());
+
+  if (op->getNumResults() > 0) {
+    scf::YieldOp::create(thenB, op.getLoc(), thenLoop->getResults());
+    scf::YieldOp::create(elseB, op.getLoc(), elseLoop->getResults());
+    op->replaceAllUsesWith(ifOp.getResults());
+  }
   op.erase();
 }
 

diff  --git a/mlir/test/Dialect/SCF/for-loop-specialization.mlir b/mlir/test/Dialect/SCF/for-loop-specialization.mlir
index 59e522f621963f..6485c9cf89cb2c 100644
--- a/mlir/test/Dialect/SCF/for-loop-specialization.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-specialization.mlir
@@ -63,3 +63,37 @@ func.func @for_constant_operand(%n: index) {
 // CHECK:           scf.for {{%.*}} = {{%.*}} to [[MIN]] step {{%.*}} {
 // CHECK:           }
 // CHECK:         }
+
+// -----
+
+// see https://github.com/llvm/llvm-project/issues/235552
+// CHECK-LABEL: func @specialize_with_iter_args
+// CHECK-SAME:    (%[[N:.*]]: index) -> index {
+func.func @specialize_with_iter_args(%n: index) -> index {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %ub = affine.min affine_map<(d0) -> (d0, 4)>(%n)
+  %result = scf.for %iv = %c0 to %ub step %c1
+      iter_args(%acc = %c0) -> (index) {
+    scf.yield %acc : index
+  }
+  return %result : index
+}
+
+// CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
+// CHECK:       %[[UB:.*]] = affine.min #{{.*}}(%[[N]])
+// CHECK:       %[[C4:.*]] = arith.constant 4 : index
+// CHECK:       %[[COND:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK:       %[[RES:.*]] = scf.if %[[COND]] -> (index) {
+// CHECK:         %[[THEN:.*]] = scf.for %{{.*}} = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[ACC1:.*]] = %[[C0]]) -> (index) {
+// CHECK:           scf.yield %[[ACC1]] : index
+// CHECK:         }
+// CHECK:         scf.yield %[[THEN]] : index
+// CHECK:       } else {
+// CHECK:         %[[ELSE:.*]] = scf.for %{{.*}} = %[[C0]] to %[[UB]] step %[[C1]] iter_args(%[[ACC2:.*]] = %[[C0]]) -> (index) {
+// CHECK:           scf.yield %[[ACC2]] : index
+// CHECK:         }
+// CHECK:         scf.yield %[[ELSE]] : index
+// CHECK:       }
+// CHECK:       return %[[RES]] : index