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

[llvmlistbot at llvm.org]

Author: Arseniy Obolenskiy
Date: 2026-07-28T14:25:34+02:00
New Revision: c273220275a7cbb8384812e1a4c84e773e1dd820

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

LOG: [mlir][SCFToSPIRV] Fix iter_args returning undef on zero trip scf.for (#206280)

Added: 
    

Modified: 
    mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
    mlir/test/Conversion/SCFToSPIRV/for.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index a6c5352d568eb..fc6b37798e42d 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -214,6 +214,16 @@ 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.
+    // Skip the stores if the loop is known to always execute at least once.
+    std::optional<APInt> tripCount = forOp.getStaticTripCount();
+    if (!tripCount || tripCount->isZero()) {
+      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..660f5eeaf93e6 100644
--- a/mlir/test/Conversion/SCFToSPIRV/for.mlir
+++ b/mlir/test/Conversion/SCFToSPIRV/for.mlir
@@ -89,6 +89,61 @@ 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>>) {
+  // CHECK: %[[LB:.*]] = spirv.Constant 42 : i32
+  %lb = arith.constant 42 : index
+  // CHECK: %[[UB:.*]] = spirv.Constant 4 : i32
+  %ub = arith.constant 4 : 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:.*]](%[[LB]], %[[INITVAR1]], %[[INITVAR2]] : i32, f32, f32)
+  // CHECK: ^[[HEADER]](%[[INDVAR:.*]]: i32, %{{.*}}: f32, %{{.*}}: f32):
+  // CHECK:   spirv.SLessThan %[[INDVAR]], %[[UB]] : i32
+  // 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_yield_dynamic_trip
+func.func @loop_yield_dynamic_trip(%lb : index, %ub : index, %step : index, %arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg3 : memref<10xf32, #spirv.storage_class<StorageBuffer>>) {
+  // 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