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

[rriddle]

Summary

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.

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

private:
   //===--------------------------------------------------------------------===//
   // Operations
 
   LogicalResult writeBlock(EncodingEmitter &emitter, Block *block);
   LogicalResult writeOp(EncodingEmitter &emitter, Operation *op);
   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
 
   void writeResourceSection(Operation *op, EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
   // Strings
 
LogicalResult BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
   // With the mask emitted, we can now emit the regions of the operation. We do
   // this after mask emission to avoid offset complications that may arise by
   // emitting the regions first (e.g. if the regions are huge, backpatching the
   // op encoding mask is more annoying).
   if (numRegions) {
     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 (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();
 }
 
 void BytecodeWriter::writeUseListOrders(EncodingEmitter &emitter,
                                         uint8_t &opEncodingMask,
                                         ValueRange range) {

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() {
+  %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
 }
 
 // Before version 2, we can't support lazy loading.
 // OLD-BYTECODE-NOT: Has 1 ops to materialize
 // OLD-BYTECODE-NOT: Materializing
 // OLD-BYTECODE: Has 0 ops to materialize
 

 // 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: 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
 // CHECK-NOT:  test.consumer
 // CHECK: Materializing...
 // 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

 def IsolatedOneRegionOp : TEST_Op<"isolated_one_region_op", [IsolatedFromAbove]> {
   let arguments = (ins Variadic<AnyType>:$operands);
   let regions = (region AnyRegion:$my_region);
   let assemblyFormat = [{
     attr-dict-with-keyword $operands $my_region `:` type($operands)
   }];
 }
 
+def IsolatedRegionsOp : TEST_Op<"isolated_regions", [IsolatedFromAbove]> {
+  let regions = (region VariadicRegion<AnyRegion>:$regions);
+  let assemblyFormat = "attr-dict-with-keyword $regions";
+}
+
 //===----------------------------------------------------------------------===//
 // NoTerminator Operation
 //===----------------------------------------------------------------------===//
 
 def SingleNoTerminatorOp : TEST_Op<"single_no_terminator_op",
                                    GraphRegionNoTerminator.traits> {
   let regions = (region SizedRegion<1>:$my_region);