diff --git a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp b/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp index f02389927019c..401629c739652 100644 --- a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp +++ b/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 regions) { + return success(llvm::all_of(regions, [&](Region ®ion) { + return succeeded(writeRegion(emitter, ®ion)); + })); + } + //===--------------------------------------------------------------------===// // Resources @@ -938,20 +945,17 @@ LogicalResult BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) { bool isIsolatedFromAbove = op->hasTrait(); emitter.emitVarIntWithFlag(numRegions, isIsolatedFromAbove); - for (Region ®ion : op->getRegions()) { - // If the region is not isolated from above, or we are emitting bytecode - // targeting version = bytecode::kLazyLoading) { EncodingEmitter regionEmitter; - if (failed(writeRegion(regionEmitter, ®ion))) + 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(); diff --git a/mlir/test/Bytecode/bytecode-lazy-loading.mlir b/mlir/test/Bytecode/bytecode-lazy-loading.mlir index 731d0e865ded3..d439f84db17d3 100644 --- a/mlir/test/Bytecode/bytecode-lazy-loading.mlir +++ b/mlir/test/Bytecode/bytecode-lazy-loading.mlir @@ -3,10 +3,13 @@ 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) -> () } @@ -14,6 +17,12 @@ func.func @op_with_passthrough_region_args() { 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,7 +62,7 @@ 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 @@ -60,4 +70,13 @@ func.func @op_with_passthrough_region_args() { // 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 diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td index 966896b27d1cb..4eb19e6dd6fe2 100644 --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/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:$regions); + let assemblyFormat = "attr-dict-with-keyword $regions"; +} + //===----------------------------------------------------------------------===// // NoTerminator Operation //===----------------------------------------------------------------------===//