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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp (+6) 
- (added) mlir/test/Conversion/SCFToGPU/no_iter_args.mlir (+18) 


``````````diff
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
+}

``````````

</details>


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