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

[rriddle]

Summary

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.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

private:
   LogicalResult materialize(LazyLoadableOpsMap::iterator it) {
     assert(it != lazyLoadableOpsMap.end() &&
            "materialize called on non-materializable op");
     valueScopes.emplace_back();
     std::vector<RegionReadState> regionStack;
     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.
   MLIRContext *getContext() const { return config.getContext(); }
 
   /// Parse the bytecode version.
   LogicalResult parseVersion(EncodingReader &reader);
 
do {
             ArrayRef<uint8_t> sectionData;
             if (failed(reader.parseSection(sectionID, sectionData)))
               return failure();
             if (sectionID != bytecode::Section::kIR)
               return emitError(fileLoc, "expected IR section for region");
             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;
             }
           }
           regionStack.push_back(std::move(childState));
 
           // If the op is isolated from above, push a new value scope.

mlir/test/Bytecode/bytecode-lazy-loading.mlir

 // RUN: mlir-opt --pass-pipeline="builtin.module(test-lazy-loading)" %s -o %t | FileCheck %s
 // RUN: mlir-opt --pass-pipeline="builtin.module(test-lazy-loading{bytecode-version=1})" %s -o %t | FileCheck %s --check-prefix=OLD-BYTECODE
 
 
 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) -> ()
   }
   %result:2 = "test.op"() : () -> (index, index)
   test.isolated_region %result#1 {
     "test.consumer"(%result#1) : (index) -> ()
   }

 // CHECK-NOT: arith
 // CHECK: Has 1 ops to materialize
 
 // CHECK: Before Materializing...
 // CHECK: "func.func"() <{function_type = () -> (), sym_name = "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
 // CHECK: Has 2 ops to materialize
 
 // CHECK: Before Materializing...
 // CHECK: test.isolated_region
 // CHECK-NOT:  test.consumer