[mlir][SPIR-V] Deserializer reads loop control through an invalidated MapVector iterator
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Description
The SPIR-V deserializer can change a loop's `LoopControl` while structurizing nested control flow.
`BlockMergeInfoMap` is an `llvm::MapVector`:
```cpp
using BlockMergeInfoMap = llvm::MapVector;
```
In `Deserializer::structurizeControlFlow`, the `updateMergeInfo` lambda erases an entry and then reads `control` through the erased iterator:
```cpp
auto it = blockMergeInfo.find(block);
if (it != blockMergeInfo.end()) {
Location loc = it->second.loc;
Block *newHeader = mapper.lookupOrNull(block);
Block *newContinue = it->second.continueBlock;
// ... remap newContinue ...
Block *newMerge = it->second.mergeBlock;
// ... remap newMerge ...
blockMergeInfo.erase(it);
blockMergeInfo.try_emplace(newHeader, loc, it->second.control, newMerge,
newContinue);
}
```
`blockMergeInfo.erase(it)` invalidates `it`. `MapVector` stores its ordered entries in a vector, so erasing an entry can move the following entry into the erased slot. The subsequent `it->second.control` can then read the following loop's control value instead of the erased loop's value.
This is reproducible through the normal `mlir-translate` serialization/deserialization path: a first inner loop marked `DontUnroll` is deserialized as `Unroll` when it is followed by another inner loop marked `Unroll`.
The affected sequence is present at `llvm-project` commit `b18437588a1c66f95313ff8298b961f5b0fd79c9`: [`updateMergeInfo`](https://github.com/llvm/llvm-project/blob/b18437588a1c66f95313ff8298b961f5b0fd79c9/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp#L2819-L2846)
## End-to-end reproducer
The following MLIR serializes to a SPIR-V 1.0 compute module accepted by SPIRV-Tools' `spirv-val`. It has a parameterless `GLCompute` entry point and the required `LocalSize` execution mode.
Save it as `repro.mlir`:
```mlir
spirv.module Logical GLSL450 requires #spirv.vce {
spirv.func @three_loops() -> () "None" {
%count = spirv.Constant 4 : i32
%zero = spirv.Constant 0 : i32
%one = spirv.Constant 1 : i32
%ivar = spirv.Variable init(%zero) : !spirv.ptr
%jvar = spirv.Variable init(%zero) : !spirv.ptr
%kvar = spirv.Variable init(%zero) : !spirv.ptr
spirv.mlir.loop {
spirv.Branch ^outer_header
^outer_header:
%ival0 = spirv.Load "Function" %ivar : i32
%icmp = spirv.SLessThan %ival0, %count : i32
spirv.BranchConditional %icmp, ^outer_body, ^outer_merge
^outer_body:
spirv.Store "Function" %jvar, %zero : i32
spirv.mlir.loop control(DontUnroll) {
spirv.Branch ^inner1_header
^inner1_header:
%jval0 = spirv.Load "Function" %jvar : i32
%jcmp = spirv.SLessThan %jval0, %count : i32
spirv.BranchConditional %jcmp, ^inner1_body, ^inner1_merge
^inner1_body:
spirv.Branch ^inner1_continue
^inner1_continue:
%jval1 = spirv.Load "Function" %jvar : i32
%jadd = spirv.IAdd %jval1, %one : i32
spirv.Store "Function" %jvar, %jadd : i32
spirv.Branch ^inner1_header
^inner1_merge:
spirv.mlir.merge
}
spirv.Store "Function" %kvar, %zero : i32
spirv.mlir.loop control(Unroll) {
spirv.Branch ^inner2_header
^inner2_header:
%kval0 = spirv.Load "Function" %kvar : i32
%kcmp = spirv.SLessThan %kval0, %count : i32
spirv.BranchConditional %kcmp, ^inner2_body, ^inner2_merge
^inner2_body:
spirv.Branch ^inner2_continue
^inner2_continue:
%kval1 = spirv.Load "Function" %kvar : i32
%kadd = spirv.IAdd %kval1, %one : i32
spirv.Store "Function" %kvar, %kadd : i32
spirv.Branch ^inner2_header
^inner2_merge:
spirv.mlir.merge
}
spirv.Branch ^outer_continue
^outer_continue:
%ival1 = spirv.Load "Function" %ivar : i32
%iadd = spirv.IAdd %ival1, %one : i32
spirv.Store "Function" %ivar, %iadd : i32
spirv.Branch ^outer_header
^outer_merge:
spirv.mlir.merge
}
spirv.Return
}
spirv.EntryPoint "GLCompute" @three_loops
spirv.ExecutionMode @three_loops "LocalSize", 1, 1, 1
}
```
Use `mlir-translate` to serialize and deserialize the module, and use SPIRV-Tools to validate and inspect the generated binary:
```bash
# Serialize the top-level SPIR-V MLIR module into a SPIR-V binary.
mlir-translate --no-implicit-module --serialize-spirv \
repro.mlir -o repro.spv
# Verify that the generated binary is valid under the SPIR-V 1.0 rules.
spirv-val --target-env spv1.0 repro.spv
# Disassemble the validated binary into readable SPIR-V assembly.
spirv-dis repro.spv > repro.spvasm
# Show the entry point, execution mode, and loop controls encoded in the binary.
grep -E 'OpEntryPoint|OpExecutionMode|OpLoopMerge' repro.spvasm
# Deserialize the same validated binary back into SPIR-V dialect MLIR.
mlir-translate --deserialize-spirv \
repro.spv -o roundtrip.mlir
# Compare the input and round-tripped loop-control attributes.
grep -n 'spirv.mlir.loop' repro.mlir roundtrip.mlir
```
## Tested configuration and actual output
The reproducer was run with:
```text
llvm-project commit: d593279c0b2891f0b0c8af3f70a1a0383b4ad1b5
LLVM version 23.0.0git
Optimized build with assertions
SPIRV-Tools v2026.3 v2026.3-0-gb707790
SPIRV-Tools commit: b707790a898e44038547df54580022fc1cf89c3d
```
The generated binary passes the official SPIRV-Tools validator:
```text
$ spirv-val --target-env spv1.0 repro.spv
$ echo $?
0
```
`spirv-dis` confirms that the binary contains a compute entry point, its execution mode, and the intended three loop controls:
```text
$ grep -E 'OpEntryPoint|OpExecutionMode|OpLoopMerge' repro.spvasm
OpEntryPoint GLCompute %three_loops "three_loops"
OpExecutionMode %three_loops LocalSize 1 1 1
OpLoopMerge %16 %15 None
OpLoopMerge %23 %22 DontUnroll
OpLoopMerge %31 %30 Unroll
```
The serializer input and deserializer output differ at the first inner loop:
```text
$ grep -n 'spirv.mlir.loop' repro.mlir roundtrip.mlir
repro.mlir:9: spirv.mlir.loop {
repro.mlir:17: spirv.mlir.loop control(DontUnroll) {
repro.mlir:34: spirv.mlir.loop control(Unroll) {
roundtrip.mlir:11: spirv.mlir.loop {
roundtrip.mlir:23: spirv.mlir.loop control(Unroll) {
roundtrip.mlir:45: spirv.mlir.loop control(Unroll) {
```
Therefore:
```text
validated SPIR-V binary: None, DontUnroll, Unroll
deserialized MLIR: None, Unroll, Unroll
```
Both `mlir-translate` commands exit successfully.
## Expected behavior
Deserialization should preserve the three loop controls:
```text
None, DontUnroll, Unroll
```
## Actual behavior
In the tested build, the first inner loop receives the following loop's control value:
```text
None, Unroll, Unroll
```
The deserializer therefore changes optimization metadata represented by a validator-accepted SPIR-V module.
## Suggested fix
Save `control` before erasing the iterator, as is already done for the other fields:
```cpp
auto it = blockMergeInfo.find(block);
if (it != blockMergeInfo.end()) {
Location loc = it->second.loc;
uint32_t control = it->second.control; # ONE LINER FIX
Block *newHeader = mapper.lookupOrNull(block);
// ... save and remap continueBlock and mergeBlock ...
blockMergeInfo.erase(it);
blockMergeInfo.try_emplace(newHeader, loc, control, newMerge, newContinue);
}
```
## Disclosure
Assisted-by: OpenAI Codex
The bug is reported by static analyzers. The reproducer and this issue were prepared with Codex. I manually reviewed and carefully validated the reproducer and results.
Contributor guide
Research direction
Start in mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp at Deserializer::structurizeControlFlow and its updateMergeInfo lambda. Run the supplied mlir-translate reproducer, then compare loop controls in repro.mlir and roundtrip.mlir. Done means deserialization preserves None, DontUnroll, and Unroll without changing the validated SPIR-V metadata.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 84/100