[mlir:bytecode] Fix bytecode lazy loading for ops with multiple regions

We currently encode each region as a separate section, but
the reader expects all of the regions to be in the same section.
This updates the writer to match the behavior that the reader
expects.

Differential Revision: https://reviews.llvm.org/D156198

Author: River707

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

@@ -587,6 +587,13 @@ class BytecodeWriter {
   LogicalResult writeRegion(EncodingEmitter &emitter, Region *region);
   LogicalResult writeIRSection(EncodingEmitter &emitter, Operation *op);
 
+  LogicalResult writeRegions(EncodingEmitter &emitter,
+                             MutableArrayRef<Region> regions) {
+    return success(llvm::all_of(regions, [&](Region &region) {
+      return succeeded(writeRegion(emitter, &region));
+    }));
+  }
+
   //===--------------------------------------------------------------------===//
   // Resources
 
@@ -938,20 +945,17 @@ LogicalResult BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
     bool isIsolatedFromAbove = op->hasTrait<OpTrait::IsIsolatedFromAbove>();
     emitter.emitVarIntWithFlag(numRegions, isIsolatedFromAbove);
 
-    for (Region &region : op->getRegions()) {
-      // If the region is not isolated from above, or we are emitting bytecode
-      // targeting version <kLazyLoading, we don't use a section.
-      if (!isIsolatedFromAbove ||
-          config.bytecodeVersion < bytecode::kLazyLoading) {
-        if (failed(writeRegion(emitter, &region)))
-          return failure();
-        continue;
-      }
-
+    // If the region is not isolated from above, or we are emitting bytecode
+    // targeting version <kLazyLoading, we don't use a section.
+    if (isIsolatedFromAbove &&
+        config.bytecodeVersion >= bytecode::kLazyLoading) {
       EncodingEmitter regionEmitter;
-      if (failed(writeRegion(regionEmitter, &region)))
+      if (failed(writeRegions(regionEmitter, op->getRegions())))
         return failure();
       emitter.emitSection(bytecode::Section::kIR, std::move(regionEmitter));
+
+    } else if (failed(writeRegions(emitter, op->getRegions()))) {
+      return failure();
     }
   }
   return success();

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

@@ -3,17 +3,26 @@
 
 
 func.func @op_with_passthrough_region_args() {
+  %0 = arith.constant 10 : index
+
   // Ensure we can handle nested non-isolated/non-lazy regions.
-  "test.one_region_op"() ({}) : () -> ()
+  "test.one_region_op"() ({
+    "test.consumer"(%0) : (index) -> ()
+  }) : () -> ()
 
-  %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) -> ()
   }
+  
+  test.isolated_regions {
+    "test.unknown_op"() : () -> ()
+  }, {
+    "test.unknown_op"() : () -> ()
+  }
   return
 }
 
@@ -39,11 +48,12 @@ 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: one_region_op
+// CHECK: test.consumer
 // CHECK: isolated_region
 // CHECK-NOT: test.consumer
-// CHECK: Has 2 ops to materialize
+// CHECK: Has 3 ops to materialize
 
 // CHECK: Before Materializing...
 // CHECK: test.isolated_region
@@ -52,12 +62,21 @@ func.func @op_with_passthrough_region_args() {
 // CHECK: test.isolated_region
 // CHECK: ^bb0(%arg0: index):
 // CHECK:  test.consumer
-// CHECK: Has 1 ops to materialize
+// CHECK: Has 2 ops to materialize
 
 // CHECK: Before Materializing...
 // CHECK: test.isolated_region
 // CHECK-NOT: test.consumer
 // CHECK: Materializing...
 // CHECK: test.isolated_region
 // CHECK: test.consumer
+// CHECK: Has 1 ops to materialize
+
+// CHECK: Before Materializing...
+// CHECK: test.isolated_regions
+// CHECK-NOT: test.unknown_op
+// CHECK: Materializing...
+// CHECK: test.isolated_regions
+// CHECK: test.unknown_op
+// CHECK: test.unknown_op
 // CHECK: Has 0 ops to materialize

mlir/test/lib/Dialect/Test/TestOps.td

@@ -451,6 +451,11 @@ def IsolatedOneRegionOp : TEST_Op<"isolated_one_region_op", [IsolatedFromAbove]>
   }];
 }
 
+def IsolatedRegionsOp : TEST_Op<"isolated_regions", [IsolatedFromAbove]> {
+  let regions = (region VariadicRegion<AnyRegion>:$regions);
+  let assemblyFormat = "attr-dict-with-keyword $regions";
+}
+
 //===----------------------------------------------------------------------===//
 // NoTerminator Operation
 //===----------------------------------------------------------------------===//