-
Notifications
You must be signed in to change notification settings - Fork 13.4k
[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
Conversation
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.
@llvm/pr-subscribers-mlir-spirv @llvm/pr-subscribers-mlir Author: Igor Wodiany (IgWod-IMG) ChangesThe 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:
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
|
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.
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.
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.