diff --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp --- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp +++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp @@ -1317,11 +1317,11 @@ 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 @@ childState.owningReader = std::make_unique(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; diff --git a/mlir/test/Bytecode/bytecode-lazy-loading.mlir b/mlir/test/Bytecode/bytecode-lazy-loading.mlir --- a/mlir/test/Bytecode/bytecode-lazy-loading.mlir +++ b/mlir/test/Bytecode/bytecode-lazy-loading.mlir @@ -3,6 +3,9 @@ func.func @op_with_passthrough_region_args() { + // Ensure we can handle nested non-isolated/non-lazy regions. + "test.one_region_op"() ({}) : () -> () + %0 = arith.constant 10 : index test.isolated_region %0 { "test.consumer"(%0) : (index) -> () @@ -36,6 +39,7 @@ // CHECK-NOT: arith // CHECK: Materializing... // CHECK: "func.func"() <{function_type = () -> (), sym_name = "op_with_passthrough_region_args"}> ({ +// CHECK: one_region_op // CHECK: arith // CHECK: isolated_region // CHECK-NOT: test.consumer