[Mlir-commits] [mlir] 03e570b - [mlir][spirv] Fix control of nested sibling constructs in the deserializer (#225920)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 19 02:28:28 PDT 2026
Author: Stanislav Bardyuk
Date: 2026-10-19T10:28:23+01:00
New Revision: 03e570bd8d3b3a7947ab2d89a0943d48a99abcef
URL: https://github.com/llvm/llvm-project/commit/03e570bd8d3b3a7947ab2d89a0943d48a99abcef
DIFF: https://github.com/llvm/llvm-project/commit/03e570bd8d3b3a7947ab2d89a0943d48a99abcef.diff
LOG: [mlir][spirv] Fix control of nested sibling constructs in the deserializer (#225920)
When the deserializer structurizes an outer loop or selection,
updateMergeInfo re-keys the blockMergeInfo entries of the nested
constructs from the old blocks to the cloned ones. It erases the old
entry first and then reads it->second.control through the erased
iterator. blockMergeInfo is a MapVector, so the erase shifts the next
entry into that slot and the nested construct is re-created with its
sibling's control.
With two sibling loops inside an outer loop, the first one marked
DontUnroll and the second Unroll, the round trip comes back with both
marked Unroll. Same thing for two sibling selections with Flatten and
DontFlatten. If the erased entry was the last one the read is past the
end of the vector.
Copy the control value out of the entry before erasing it, like loc
already is. Adds round-trip tests with two sibling loops and two sibling
selections.
Fixes #218853
Added:
Modified:
mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
mlir/test/Target/SPIRV/loop.mlir
mlir/test/Target/SPIRV/selection.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
index 5813a8961e5e6..029c2f24de24f 100644
--- a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
+++ b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
@@ -2871,6 +2871,8 @@ LogicalResult ControlFlowStructurizer::structurize() {
if (it != blockMergeInfo.end()) {
// Use the original location for nested selection/loop ops.
Location loc = it->second.loc;
+ // Copy the control before the erase below invalidates `it`.
+ uint32_t nestedControl = it->second.control;
Block *newHeader = mapper.lookupOrNull(block);
if (!newHeader)
@@ -2889,10 +2891,9 @@ LogicalResult ControlFlowStructurizer::structurize() {
if (Block *mappedTo = mapper.lookupOrNull(newMerge))
newMerge = mappedTo;
- // The iterator should be erased before adding a new entry into
- // blockMergeInfo to avoid iterator invalidation.
+ // Erase the stale entry before inserting the remapped one.
blockMergeInfo.erase(it);
- blockMergeInfo.try_emplace(newHeader, loc, it->second.control, newMerge,
+ blockMergeInfo.try_emplace(newHeader, loc, nestedControl, newMerge,
newContinue);
}
diff --git a/mlir/test/Target/SPIRV/loop.mlir b/mlir/test/Target/SPIRV/loop.mlir
index b9a4295a04e7c..4173533cf8f0c 100644
--- a/mlir/test/Target/SPIRV/loop.mlir
+++ b/mlir/test/Target/SPIRV/loop.mlir
@@ -340,3 +340,79 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
spirv.Return
}
}
+
+// -----
+
+// Two sibling loops nested in a loop, each with its own loop control
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
+// CHECK-LABEL: @sibling_loops
+ spirv.func @sibling_loops(%count : i32) -> () "None" {
+ %zero = spirv.Constant 0: i32
+ %one = spirv.Constant 1: i32
+ %ivar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
+ %jvar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
+ %kvar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
+
+// CHECK: spirv.mlir.loop {
+ spirv.mlir.loop {
+ spirv.Branch ^header
+ ^header:
+ %ival0 = spirv.Load "Function" %ivar : i32
+ %icmp = spirv.SLessThan %ival0, %count : i32
+ spirv.BranchConditional %icmp, ^body, ^merge
+ ^body:
+ spirv.Store "Function" %jvar, %zero : i32
+// CHECK: spirv.mlir.loop control(DontUnroll)
+ spirv.mlir.loop control(DontUnroll) {
+ spirv.Branch ^header
+ ^header:
+ %jval0 = spirv.Load "Function" %jvar : i32
+ %jcmp = spirv.SLessThan %jval0, %count : i32
+ spirv.BranchConditional %jcmp, ^body, ^merge
+ ^body:
+ spirv.Branch ^continue
+ ^continue:
+ %jval1 = spirv.Load "Function" %jvar : i32
+ %jadd = spirv.IAdd %jval1, %one : i32
+ spirv.Store "Function" %jvar, %jadd : i32
+ spirv.Branch ^header
+ ^merge:
+ spirv.mlir.merge
+ }
+ spirv.Store "Function" %kvar, %zero : i32
+// CHECK: spirv.mlir.loop control(Unroll)
+ spirv.mlir.loop control(Unroll) {
+ spirv.Branch ^header
+ ^header:
+ %kval0 = spirv.Load "Function" %kvar : i32
+ %kcmp = spirv.SLessThan %kval0, %count : i32
+ spirv.BranchConditional %kcmp, ^body, ^merge
+ ^body:
+ spirv.Branch ^continue
+ ^continue:
+ %kval1 = spirv.Load "Function" %kvar : i32
+ %kadd = spirv.IAdd %kval1, %one : i32
+ spirv.Store "Function" %kvar, %kadd : i32
+ spirv.Branch ^header
+ ^merge:
+ spirv.mlir.merge
+ }
+ spirv.Branch ^continue
+ ^continue:
+ %ival1 = spirv.Load "Function" %ivar : i32
+ %iadd = spirv.IAdd %ival1, %one : i32
+ spirv.Store "Function" %ivar, %iadd : i32
+ spirv.Branch ^header
+ ^merge:
+ spirv.mlir.merge
+ }
+ spirv.Return
+ }
+
+ spirv.func @main() -> () "None" {
+ spirv.Return
+ }
+ spirv.EntryPoint "GLCompute" @main
+ spirv.ExecutionMode @main "LocalSize", 1, 1, 1
+}
diff --git a/mlir/test/Target/SPIRV/selection.mlir b/mlir/test/Target/SPIRV/selection.mlir
index 317c951f2ba8d..7fd50dda3e17d 100644
--- a/mlir/test/Target/SPIRV/selection.mlir
+++ b/mlir/test/Target/SPIRV/selection.mlir
@@ -375,3 +375,59 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
spirv.Return
}
}
+
+// -----
+
+// Two sibling selections nested in a loop, each with its own selection control
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
+// CHECK-LABEL: @sibling_selections
+ spirv.func @sibling_selections(%count : i32, %cond : i1) -> () "None" {
+ %zero = spirv.Constant 0: i32
+ %one = spirv.Constant 1: i32
+ %ivar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
+
+// CHECK: spirv.mlir.loop {
+ spirv.mlir.loop {
+ spirv.Branch ^header
+ ^header:
+ %ival0 = spirv.Load "Function" %ivar : i32
+ %icmp = spirv.SLessThan %ival0, %count : i32
+ spirv.BranchConditional %icmp, ^body, ^merge
+ ^body:
+// CHECK: spirv.mlir.selection control(Flatten)
+ spirv.mlir.selection control(Flatten) {
+ spirv.BranchConditional %cond, ^then, ^merge
+ ^then:
+ spirv.Store "Function" %ivar, %one : i32
+ spirv.Branch ^merge
+ ^merge:
+ spirv.mlir.merge
+ }
+// CHECK: spirv.mlir.selection control(DontFlatten)
+ spirv.mlir.selection control(DontFlatten) {
+ spirv.BranchConditional %cond, ^then, ^merge
+ ^then:
+ spirv.Store "Function" %ivar, %zero : i32
+ spirv.Branch ^merge
+ ^merge:
+ spirv.mlir.merge
+ }
+ spirv.Branch ^continue
+ ^continue:
+ %ival1 = spirv.Load "Function" %ivar : i32
+ %iadd = spirv.IAdd %ival1, %one : i32
+ spirv.Store "Function" %ivar, %iadd : i32
+ spirv.Branch ^header
+ ^merge:
+ spirv.mlir.merge
+ }
+ spirv.Return
+ }
+
+ spirv.func @main() -> () "None" {
+ spirv.Return
+ }
+ spirv.EntryPoint "GLCompute" @main
+ spirv.ExecutionMode @main "LocalSize", 1, 1, 1
+}
More information about the Mlir-commits
mailing list