Add support for Lazyloading to the MLIR bytecode

IsolatedRegions are emitted in sections in order for the reader to be
able to skip over them. A new class is exposed to manage the state and
allow the readers to load these IsolatedRegions on-demand.

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

1 files changed, 12 insertions, 2 deletions

diff --git a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp b/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
index 801f3022d0e4..158dbe6d161d 100644
--- a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
+++ b/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
@@ -734,8 +734,18 @@ void BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
     bool isIsolatedFromAbove = op->hasTrait<OpTrait::IsIsolatedFromAbove>();
     emitter.emitVarIntWithFlag(numRegions, isIsolatedFromAbove);
 
-    for (Region &region : op->getRegions())
-      writeRegion(emitter, &region);
+    for (Region &region : op->getRegions()) {
+      // If the region is not isolated from above, or we are emitting bytecode
+      // targetting version <2, we don't use a section.
+      if (!isIsolatedFromAbove || config.bytecodeVersion < 2) {
+        writeRegion(emitter, &region);
+        continue;
+      }
+
+      EncodingEmitter regionEmitter;
+      writeRegion(regionEmitter, &region);
+      emitter.emitSection(bytecode::Section::kIR, std::move(regionEmitter));
+    }
   }
 }