[Mlir-commits] [mlir] 0480298 - [mlir][SCFToGPU] Fix crash when converting affine.for with iter_args to GPU (#185073)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-11T12:13:09+01:00
New Revision: 0480298e85180783f8e9923be0d3c0178b6c0b5b

URL: https://github.com/llvm/llvm-project/commit/0480298e85180783f8e9923be0d3c0178b6c0b5b
DIFF: https://github.com/llvm/llvm-project/commit/0480298e85180783f8e9923be0d3c0178b6c0b5b.diff

LOG: [mlir][SCFToGPU] Fix crash when converting affine.for with iter_args to GPU (#185073)

The convert-affine-for-to-gpu pass moved operations from the affine.for
loop body to the GPU launch kernel, then erased the original loop.
However, if the loop had iter_args (reduction loops), the moved
operations could still reference the loop body's block arguments (the
iter_args). When the loop was erased, those block arguments were
destroyed while still having live uses, triggering a use_empty()
assertion.

Fix this by detecting loops with iter_args in collectBounds and
returning an error. Reduction loops cannot be trivially converted to GPU
kernels without dedicated handling of the accumulator semantics.

Fixes #116044

Assisted-by: Claude Code

Added: 
    mlir/test/Conversion/SCFToGPU/no_iter_args.mlir

Modified: 
    mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
index 93e1e3f86279a..370457c85e797 100644
--- a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
+++ b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
@@ -184,6 +184,12 @@ AffineLoopToGpuConverter::collectBounds(AffineForOp forOp, unsigned numLoops) {
   steps.reserve(numLoops);
   AffineForOp currentLoop = forOp;
   for (unsigned i = 0; i < numLoops; ++i) {
+    if (currentLoop.getNumIterOperands() > 0) {
+      currentLoop.emitError(
+          "affine loop with iter_args cannot be converted to GPU kernel");
+      return std::nullopt;
+    }
+
     Value lowerBound = getOrEmitLowerBound(currentLoop, builder);
     Value upperBound = getOrEmitUpperBound(currentLoop, builder);
     if (!lowerBound || !upperBound) {

diff  --git a/mlir/test/Conversion/SCFToGPU/no_iter_args.mlir b/mlir/test/Conversion/SCFToGPU/no_iter_args.mlir
new file mode 100644
index 0000000000000..ab6cb7ffedf49
--- /dev/null
+++ b/mlir/test/Conversion/SCFToGPU/no_iter_args.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-opt %s --verify-diagnostics -pass-pipeline="builtin.module(func.func(convert-affine-for-to-gpu{gpu-block-dims=1 gpu-thread-dims=0}))"
+
+// Regression test: affine.for loops with iter_args (reduction loops) must be
+// rejected with a proper diagnostic rather than crashing. The pass cannot
+// convert reduction loops because the loop body is moved to the GPU kernel
+// while block arguments referencing iter_args would become dangling uses.
+
+func.func @reduction_loop(%arg0: memref<1024xf32>) -> f32 {
+  %c10 = arith.constant 10 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  // expected-error @+1 {{affine loop with iter_args cannot be converted to GPU kernel}}
+  %result = affine.for %i = 0 to %c10 iter_args(%acc = %cst) -> (f32) {
+    %val = affine.load %arg0[%i] : memref<1024xf32>
+    %sum = arith.addf %acc, %val : f32
+    affine.yield %sum : f32
+  }
+  return %result : f32
+}