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
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 62 additions & 24 deletions mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down