[Mlir-commits] [mlir] [mlir][spirv] Update mergeInfo of blocks nested in regions (PR #137789)

[Igor Wodiany]

https://github.com/IgWod-IMG created https://github.com/llvm/llvm-project/pull/137789

The current code that updates mergeInfo iterates only over constructBlocks, potentially leaving blocks nested in already structured regions with invalid mergeInfo. This patch adds walk for each block to ensure all nested blocks are considered.

It is not possible to add a unit test exercising this change as whether the problem occurs depends on the structuring order that is currently non-deterministic.

>From 3e62d0e7e451630808a842750958b39c34513b79 Mon Sep 17 00:00:00 2001
From: Igor Wodiany <igor.wodiany at imgtec.com>
Date: Tue, 8 Apr 2025 16:48:31 +0100
Subject: [PATCH] [mlir][spirv] Update mergeInfo of blocks nested in regions

The current code that updates mergeInfo iterates only over
constructBlocks, potentially leaving blocks nested in already
structured regions with invalid mergeInfo. This patch adds walk
for each block to ensure all nested blocks are considered.

It is not possible to add a unit test exercising this change
as whether the problem occurs depends on the structuring order
that is currently non-deterministic.
---
 .../SPIRV/Deserialization/Deserializer.cpp    | 86 +++++++++++++------
 1 file changed, 62 insertions(+), 24 deletions(-)

diff --git a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
index d471d9a8e3d6c..a1e757fcfb7f9 100644
--- a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
+++ b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
@@ -2105,37 +2105,75 @@ LogicalResult ControlFlowStructurizer::structurize() {
     // selection/loop. If so, they will be recorded within blockMergeInfo.
     // We need to update the pointers there to the newly remapped ones so we can
     // continue structurizing them later.
+    //
+    // We need to walk each block as constructBlocks do not include blocks
+    // internal to ops already structured within those blocks. It is not
+    // fully clear to me why the mergeInfo of blocks (yet to be structured)
+    // inside already structured selections/loops get invalidated and needs
+    // updating, however the following example code can cause a crash (depending
+    // on the structuring order), when the most inner selection is being
+    // structured after the outer selection and loop have been already
+    // structured:
+    //
+    //  spirv.mlir.for {
+    //    // ...
+    //    spirv.mlir.selection {
+    //      // ..
+    //      // A selection region that hasn't been yet structured!
+    //      // ..
+    //    }
+    //    // ...
+    //  }
+    //
+    // If the loop gets structured after the outer selection, but before the
+    // inner selection. Moving the already structured selection inside the loop
+    // will invalidate the mergeInfo of the region that is not yet structured.
+    // Just going over constructBlocks will not check and updated header blocks
+    // inside the already structured selection region. Walking block fixes that.
+    //
+    // TODO: If structuring was done in a fixed order starting with inner
+    // most constructs this most likely not be an issue and the whole code
+    // section could be removed. However, with the current non-deterministic
+    // order this is not possible.
+    //
     // TODO: The asserts in the following assumes input SPIR-V blob forms
     // correctly nested selection/loop constructs. We should relax this and
     // support error cases better.
-    auto it = blockMergeInfo.find(block);
-    if (it != blockMergeInfo.end()) {
-      // Use the original location for nested selection/loop ops.
-      Location loc = it->second.loc;
-
-      Block *newHeader = mapper.lookupOrNull(block);
-      if (!newHeader)
-        return emitError(loc, "failed control flow structurization: nested "
-                              "loop header block should be remapped!");
-
-      Block *newContinue = it->second.continueBlock;
-      if (newContinue) {
-        newContinue = mapper.lookupOrNull(newContinue);
-        if (!newContinue)
+    auto updateMergeInfo = [&](Block *block) -> WalkResult {
+      auto it = blockMergeInfo.find(block);
+      if (it != blockMergeInfo.end()) {
+        // Use the original location for nested selection/loop ops.
+        Location loc = it->second.loc;
+
+        Block *newHeader = mapper.lookupOrNull(block);
+        if (!newHeader)
           return emitError(loc, "failed control flow structurization: nested "
-                                "loop continue block should be remapped!");
+                                "loop header block should be remapped!");
+
+        Block *newContinue = it->second.continueBlock;
+        if (newContinue) {
+          newContinue = mapper.lookupOrNull(newContinue);
+          if (!newContinue)
+            return emitError(loc, "failed control flow structurization: nested "
+                                  "loop continue block should be remapped!");
+        }
+
+        Block *newMerge = it->second.mergeBlock;
+        if (Block *mappedTo = mapper.lookupOrNull(newMerge))
+          newMerge = mappedTo;
+
+        // The iterator should be erased before adding a new entry into
+        // blockMergeInfo to avoid iterator invalidation.
+        blockMergeInfo.erase(it);
+        blockMergeInfo.try_emplace(newHeader, loc, it->second.control, newMerge,
+                                   newContinue);
       }
 
-      Block *newMerge = it->second.mergeBlock;
-      if (Block *mappedTo = mapper.lookupOrNull(newMerge))
-        newMerge = mappedTo;
+      return WalkResult::advance();
+    };
 
-      // The iterator should be erased before adding a new entry into
-      // blockMergeInfo to avoid iterator invalidation.
-      blockMergeInfo.erase(it);
-      blockMergeInfo.try_emplace(newHeader, loc, it->second.control, newMerge,
-                                 newContinue);
-    }
+    if (block->walk(updateMergeInfo).wasInterrupted())
+      return failure();
 
     // The structured selection/loop's entry block does not have arguments.
     // If the function's header block is also part of the structured control