[Mlir-commits] [mlir] [mlir][SCFToSPIRV] Fix iter_args returning undef on zero trip scf.for (PR #206280)
Arseniy Obolenskiy
llvmlistbot at llvm.org
Tue Jul 28 04:36:32 PDT 2026
https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/206280
>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 1/3] [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
>From e8bff608f3fa69f58f5b41f349bfa18ad4e5462b Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Mon, 27 Jul 2026 20:06:00 +0200
Subject: [PATCH 2/3] Address comments
---
mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp | 12 ++++++++----
mlir/test/Conversion/SCFToSPIRV/for.mlir | 10 +++++++---
2 files changed, 15 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index c5859ca811cc2..d5c7bac5db0fa 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -216,10 +216,14 @@ struct ForOpConversion final : SCFToSPIRVPattern<scf::ForOp> {
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);
+ // 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 d5c1801ebba57..fd7b3f2e93381 100644
--- a/mlir/test/Conversion/SCFToSPIRV/for.mlir
+++ b/mlir/test/Conversion/SCFToSPIRV/for.mlir
@@ -91,8 +91,10 @@ func.func @loop_yield(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>
// 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
+ // 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
@@ -103,7 +105,9 @@ func.func @loop_yield_zero_trip(%arg2 : memref<10xf32, #spirv.storage_class<Stor
// 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: 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
>From 344b4d7015adbc9ce68027a6ee0190d08d704aca Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Tue, 28 Jul 2026 13:36:20 +0200
Subject: [PATCH 3/3] Add dyn trip count test
---
mlir/test/Conversion/SCFToSPIRV/for.mlir | 24 ++++++++++++++++++++++++
1 file changed, 24 insertions(+)
diff --git a/mlir/test/Conversion/SCFToSPIRV/for.mlir b/mlir/test/Conversion/SCFToSPIRV/for.mlir
index fd7b3f2e93381..660f5eeaf93e6 100644
--- a/mlir/test/Conversion/SCFToSPIRV/for.mlir
+++ b/mlir/test/Conversion/SCFToSPIRV/for.mlir
@@ -120,6 +120,30 @@ func.func @loop_yield_zero_trip(%arg2 : memref<10xf32, #spirv.storage_class<Stor
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
More information about the Mlir-commits
mailing list