[Mlir-commits] [mlir] [mlir][SCFToSPIRV] Fix iter_args returning undef on zero trip scf.for (PR #206280)

[Arseniy Obolenskiy]

https://github.com/aobolensk created https://github.com/llvm/llvm-project/pull/206280

None

>From 27b784f011f48b10a3432b0a98d151ff511c506b Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Sat, 27 Jun 2026 21:26:54 +0200
Subject: [PATCH] [mlir][SCFToSPIRV] Fix iter_args returning undef on zero-trip
 scf.for

---
 mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp |  6 +++++
 mlir/test/Conversion/SCFToSPIRV/for.mlir      | 27 +++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index d5140f3faa6ff..c5859ca811cc2 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -214,6 +214,12 @@ struct ForOpConversion final : SCFToSPIRVPattern<scf::ForOp> {
       initTypes.push_back(arg.getType());
     replaceSCFOutputValue(forOp, loopOp, rewriter, scfToSPIRVContext,
                           initTypes);
+
+    // Store init values so a zero-trip loop returns them instead of undef.
+    auto &allocas = scfToSPIRVContext->outputVars[loopOp];
+    rewriter.setInsertionPoint(loopOp);
+    for (auto [alloca, init] : llvm::zip(allocas, adaptor.getInitArgs()))
+      spirv::StoreOp::create(rewriter, loc, alloca, init);
     return success();
   }
 };
diff --git a/mlir/test/Conversion/SCFToSPIRV/for.mlir b/mlir/test/Conversion/SCFToSPIRV/for.mlir
index 702bee476668f..d5c1801ebba57 100644
--- a/mlir/test/Conversion/SCFToSPIRV/for.mlir
+++ b/mlir/test/Conversion/SCFToSPIRV/for.mlir
@@ -89,6 +89,33 @@ func.func @loop_yield(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>
   return
 }
 
+// CHECK-LABEL: @loop_yield_zero_trip
+func.func @loop_yield_zero_trip(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg3 : memref<10xf32, #spirv.storage_class<StorageBuffer>>) {
+  %lb = arith.constant 4 : index
+  %ub = arith.constant 42 : index
+  %step = arith.constant 2 : index
+  // CHECK: %[[INITVAR1:.*]] = spirv.Constant 0.000000e+00 : f32
+  %s0 = arith.constant 0.0 : f32
+  // CHECK: %[[INITVAR2:.*]] = spirv.Constant 1.000000e+00 : f32
+  %s1 = arith.constant 1.0 : f32
+  // CHECK: %[[VAR1:.*]] = spirv.Variable : !spirv.ptr<f32, Function>
+  // CHECK: %[[VAR2:.*]] = spirv.Variable : !spirv.ptr<f32, Function>
+  // CHECK-DAG: spirv.Store "Function" %[[VAR1]], %[[INITVAR1]] : f32
+  // CHECK-DAG: spirv.Store "Function" %[[VAR2]], %[[INITVAR2]] : f32
+  // CHECK: spirv.mlir.loop {
+  // CHECK:   spirv.Branch ^[[HEADER:.*]](%{{.*}}, %[[INITVAR1]], %[[INITVAR2]] : i32, f32, f32)
+  // CHECK: }
+  %result:2 = scf.for %i0 = %lb to %ub step %step iter_args(%si = %s0, %sj = %s1) -> (f32, f32) {
+    %sn = arith.addf %si, %si : f32
+    scf.yield %sn, %sn : f32, f32
+  }
+  // CHECK-DAG: %[[OUT1:.*]] = spirv.Load "Function" %[[VAR1]] : f32
+  // CHECK-DAG: %[[OUT2:.*]] = spirv.Load "Function" %[[VAR2]] : f32
+  memref.store %result#0, %arg3[%lb] : memref<10xf32, #spirv.storage_class<StorageBuffer>>
+  memref.store %result#1, %arg3[%ub] : memref<10xf32, #spirv.storage_class<StorageBuffer>>
+  return
+}
+
 // CHECK-LABEL: @loop_unroll
 func.func @loop_unroll(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg3 : memref<10xf32, #spirv.storage_class<StorageBuffer>>) {
   %lb = arith.constant 0 : index