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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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

>From 04e987dfe13529bbd2813b120217556a0cee6e66 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 15:31:22 -0800
Subject: [PATCH] [mlir][SCFToGPU] Fix crash when converting affine.for with
 iter_args to GPU

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
---
 mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp      |  6 ++++++
 .../test/Conversion/SCFToGPU/no_iter_args.mlir | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+)
 create mode 100644 mlir/test/Conversion/SCFToGPU/no_iter_args.mlir

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