[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

Author: River707

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

@@ -1317,11 +1317,11 @@ class mlir::BytecodeReader::Impl {
     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;

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 @@ func.func @op_with_passthrough_region_args() {
 // 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