[mlir][SPIR-V] Deserializer reads loop control through an invalidated MapVector iterator

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.

Description

The SPIR-V deserializer can change a loop's LoopControl while structurizing nested control flow.

BlockMergeInfoMap is an llvm::MapVector:

using BlockMergeInfoMap = llvm::MapVector<Block *, BlockMergeInfo>;

In Deserializer::structurizeControlFlow, the updateMergeInfo lambda erases an entry and then reads control through the erased iterator:

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

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:

spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
  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<i32, Function>
    %jvar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
    %kvar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
    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:

# 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:

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:

$ 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:

$ 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:

$ 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:

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:

None, DontUnroll, Unroll

Actual behavior

In the tested build, the first inner loop receives the following loop's control value:

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:

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.