Thanks to visit codestin.com
Credit goes to github.com

Skip to content

[mlir][spirv] Update mergeInfo of blocks nested in regions #137789

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 30, 2025

Conversation

IgWod-IMG
Copy link
Contributor

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.

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.
@llvmbot
Copy link
Member

llvmbot commented Apr 29, 2025

@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir

Author: Igor Wodiany (IgWod-IMG)

Changes

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.


Full diff: https://github.com/llvm/llvm-project/pull/137789.diff

1 Files Affected:

  • (modified) mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp (+62-24)
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

@IgWod-IMG IgWod-IMG merged commit 6ae7177 into llvm:main Apr 30, 2025
14 checks passed
@IgWod-IMG IgWod-IMG deleted the img_update-merge-info branch April 30, 2025 12:45
IanWood1 pushed a commit to IanWood1/llvm-project that referenced this pull request May 6, 2025
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.
IanWood1 pushed a commit to IanWood1/llvm-project that referenced this pull request May 6, 2025
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.
IanWood1 pushed a commit to IanWood1/llvm-project that referenced this pull request May 6, 2025
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.
GeorgeARM pushed a commit to GeorgeARM/llvm-project that referenced this pull request May 7, 2025
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants