[mlir][bytecode] Fix lazy loading of non-isolated regions

The bytecode reader currently assumes all regions can be lazy
loaded, which breaks reading any non-isolated region. This patch
fixes that by properly handling nested non-lazy regions, and only
considers isolated regions as lazy.

Differential Revision: https://reviews.llvm.org/D153795
This commit is contained in:
River Riddle
2023-06-26 11:36:06 -07:00
parent a18266473b
commit 6ee1aba8ac
2 changed files with 13 additions and 12 deletions

View File

@@ -1317,11 +1317,11 @@ private:
regionStack.push_back(std::move(it->getSecond()->second));
lazyLoadableOps.erase(it->getSecond());
lazyLoadableOpsMap.erase(it);
auto result = parseRegions(regionStack, regionStack.back());
assert((regionStack.empty() || failed(result)) &&
"broken invariant: regionStack should be empty when parseRegions "
"succeeds");
return result;
while (!regionStack.empty())
if (failed(parseRegions(regionStack, regionStack.back())))
return failure();
return success();
}
/// Return the context for this config.
@@ -2094,14 +2094,11 @@ BytecodeReader::Impl::parseRegions(std::vector<RegionReadState> &regionStack,
childState.owningReader =
std::make_unique<EncodingReader>(sectionData, fileLoc);
childState.reader = childState.owningReader.get();
}
if (lazyLoading) {
// If the user has a callback set, they have the opportunity
// to control lazyloading as we go.
if (!lazyOpsCallback || !lazyOpsCallback(*op)) {
lazyLoadableOps.push_back(
std::make_pair(*op, std::move(childState)));
// If the user has a callback set, they have the opportunity to
// control lazyloading as we go.
if (lazyLoading && (!lazyOpsCallback || !lazyOpsCallback(*op))) {
lazyLoadableOps.emplace_back(*op, std::move(childState));
lazyLoadableOpsMap.try_emplace(*op,
std::prev(lazyLoadableOps.end()));
continue;