Add support for Lazyloading to the MLIR bytecode

[mehdi_amini]

Summary

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.

[rriddle]

Nice! This is pretty much what I had in mind originally. Getting the bytecode format is the major part of this, we can tune the policy/API as we use it.

mlir/include/mlir/Bytecode/BytecodeReader.h

38–39

44–48

Why return a functor instead of driving this through BytecodeReader? The returned functor references the BytecodeReader, so what are the benefits to returning a functor? The only situation I can think of is that it makes things slightly simpler when reading multiple bytecode files (which I run into a lot for my needs), removing the need to remember which one holds each op.

47

Maybe getLazyOpMaterializer? Not blocking at all, but might be nice to keep consistent language for this.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1116

Can you document this? Can be a simple callout to the main method (just feels weird having methods undocumented).

1122

Can you add messages to these asserts?

1289

Please document this.

1653–1677

1660–1663

1700

Please add messages to the asserts.

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

742

Waiting for Jacques patch?

[jpienaar]

Nice! Thanks

mlir/docs/BytecodeFormat.md

345

Is region_section defined here? (Seems it should be ir_section from below)

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

742

Would be nice :-)

mlir/test/lib/IR/TestLazyLoading.cpp

52

So if I wanted to extract only a subgraph, one would materilze the op, materialize recursively everything in it while keeping track of referenced symbols and recursively materialize all of them?

60

Nit: error convention is sentence fragments (start lower case and no trailing punctuation)

[jpienaar]

mlir/docs/BytecodeFormat.md

345

Did you consider making it so that not all isolated from above are encoded same way? Perhaps stealing a bit from region encoding to indicate encoding. That way one could mark only what is needed and so writer could limit depth and also it's smaller in old encoding. That does complicate the writer though (writer config would need a set or somesuch to track what to encode).

[mfrancio]

Thanks a lot!

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1618–1621

This line seems unnecessary now?

[mehdi_amini]

mlir/docs/BytecodeFormat.md

345

Is region_section defined here? (Seems it should be ir_section from below)

Sections aren't defined, I'm not sure how to express it? I could make it:

region_section {
  region
}

But it is still implicit that it is a section, for example

op {
  name: varint,
  encodingMask: byte,
  ...

There is no section defined here.

As for the encoding, I like the uniformity of the format and the capabilities: it's not clear to me we need to push complexity on the emitter instead of just leaving the flexibility to the consumer. Do you have something in mind?

mlir/include/mlir/Bytecode/BytecodeReader.h

44–48

I was actually trying to have a single map lookup so that the API can do directly "check if this op can be materialized" and "materialize this op" without hitting the map twice! This may not be worth it though: the work to do is heavy enough that a map lookup is negligible.

I'm not super happy with the API right now though, as it's not easy to know which operation need materialization or not.

Another thing is how "unsafe" this all is from the client side, I'd like to think a bit more the API to provide some more guarantees maybe.
For example right now when we parse a file (bytecode or ASM) we verify it, but with lazy-loading we have to disable that (I am not doing it automatically, it has to be done in the config so that it is explicit for the client that verifier isn't performed!).
It is still easy for the client right now to move forward without any verification: what invariants do we want around lazy-loaded input? It is also possible that all use-cases will be so ad-hoc that we can't provide an "easy" API.

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

742

Yes, I was waiting for Jacques, I'll rebase and fix the TODO!

mlir/test/lib/IR/TestLazyLoading.cpp

52

Yes: you would materialize the module, build a symbol table, then worklist algorithm: add your entry point to the wordlist and while !empty, pop the current function, materialize and collect the symbols, using the symbol table to push the functions to the wordlist.

[jpienaar]

mlir/docs/BytecodeFormat.md

345

Good point, we need to tag those a bit differently. Perhaps section<region> and we could have op : Section { ... }.

(Can be follow up, as you mention it's not consistent and we can do a NOP change to make it).

Re encoding: I agree, uniform is better. I was just thinking it's probably rare than one would be lazily loading more than a handful of the regions but there is paying a cost on each. At least size wise, the uniformity does avoid additional cost during serialization/deserialization. I don't have anything quantitative though to back the tradeoff either way.

[mehdi_amini]

mlir/docs/BytecodeFormat.md

345

I've been wondering about the performance by the way, right now the cost is (if I don't miss anything):

• 1 varint for the section size per isolated region
• 1 entry in a map when lazyloading is enabled

Then in terms of runtime loading I'd be interested to benchmark "always lazy loading" vs "direct loading", it should be very close.
Maybe we need an API on the reader to access the ops to load instead of just getMaterializerOp(Operation *), that'll save traversing the IR.

[jpienaar]

mlir/docs/BytecodeFormat.md

345

1 byte for section id and 1 varint (that is if there is no alignment or padding). Not high, but means ~2.4 bytes per function additional. So if you crammed multiple models in ine file, it could add up (especially if you use encodings where everything is a function rather than region ... I've definitely been surprised by what python frontends end up producing). Probably still only at like 100 bytes larger kind of category.

[mehdi_amini]

• Rebase on Jacques' back-deployment support and test it
• Update the exposed APIs to avoid having to walk the IR to find ops to lazy-load.

[rriddle]

mlir/include/mlir/Bytecode/BytecodeReader.h

50–54

When will these be used outside of just testing? I can't think of a situation where this would actually get used.

[mehdi_amini]

mlir/include/mlir/Bytecode/BytecodeReader.h

50–54

You could have a list of functions to load and filter based on this?
That said the API isn't good enough yet for this, we can't "drop" things from the worklist without loading them..,.

[mehdi_amini]

Update the lazyloading C++ API

[jpienaar]

Nice - I recall we talked about reachability too and so materializing what is reachable from others materialized and doing that in finalize if needed. But that is probably a higher level function and a more expensive one that could be built on top of this instead.

mlir/include/mlir/Bytecode/BytecodeReader.h

44

This is only needs to be materialized given already materialized right? Ok that sentence is rough, i was mostly trying to think about recursive part and the user shouldn't expect that parser would decent into an unmaterialized op

48

Why couldn't this have same API as finalize and just return whether to materialize or not?

51

Why would a user care about this?

59

make sure to use

62

So upon callback of op to materialize, the callback could then call materialize?

[mehdi_amini]

Update the callbacks to return a bool to provide control to the client to materialize as we go

[mehdi_amini]

mlir/include/mlir/Bytecode/BytecodeReader.h

44

What about The lazyOps call back is invoked for every ops that contains regions that can be lazy loaded.

I would also change the callback to return a bool to immediately materialize as you suggested below.

51

I don't know, but it's cheap to expose :)

62

It technically can: we shouldn't because it is a recursion, so better use a worklist algorithm, like I did in mlir/test/lib/IR/TestLazyLoading.cpp (around line 70)

[uweigand]

The new test case Bytecode/bytecode-lazy-loading.mlir is failing on s390x due to what seems to be endian issues. (It looks likely that this has just exposed existing endian problems in the bytecode reader.)

Unfortunately the s390x builder was down for a couple of weeks due to maintenance, but now that it's back up, the failing test makes it red. Note that since this patch, three more patches have introduced additional failing test cases on s390x:
https://reviews.llvm.org/D149755 (Bytecode/uselist_orders.mlir)
https://reviews.llvm.org/D151386 (Bytecode/unregistered_dialect.mlir)
https://reviews.llvm.org/D151408 (Bytecode/./MLIRBytecodeTests/Bytecode/MultiModuleWithResource)
all for what appear to be similar (or even the same) endian problems.

See e.g. https://lab.llvm.org/buildbot/#/builders/199/builds/21088 for more details. Looking at the BytecodeReader.cpp file, I can see a number of places that obviously assume little-endian host byte order, e.g. in parseVarInt or parseMultiByteVarInt. Any suggestions on how to handle MLIR bytecode on big-endian hosts?

[rriddle]

In D149515#4396268, @uweigand wrote:

The new test case Bytecode/bytecode-lazy-loading.mlir is failing on s390x due to what seems to be endian issues. (It looks likely that this has just exposed existing endian problems in the bytecode reader.)

Unfortunately the s390x builder was down for a couple of weeks due to maintenance, but now that it's back up, the failing test makes it red. Note that since this patch, three more patches have introduced additional failing test cases on s390x:
https://reviews.llvm.org/D149755 (Bytecode/uselist_orders.mlir)
https://reviews.llvm.org/D151386 (Bytecode/unregistered_dialect.mlir)
https://reviews.llvm.org/D151408 (Bytecode/./MLIRBytecodeTests/Bytecode/MultiModuleWithResource)
all for what appear to be similar (or even the same) endian problems.

See e.g. https://lab.llvm.org/buildbot/#/builders/199/builds/21088 for more details. Looking at the BytecodeReader.cpp file, I can see a number of places that obviously assume little-endian host byte order, e.g. in parseVarInt or parseMultiByteVarInt. Any suggestions on how to handle MLIR bytecode on big-endian hosts?

MLIR bytecode doesn't support big endian hosts at all, so we'll need to just disable tests in those cases, e.g. https://github.com/llvm/llvm-project/blob/52ca6ad755b0cc2aa603cfb3124bf58c04a47005/mlir/test/Bytecode/resources.mlir#L4

We've discussed it before, but most of the engineers working on the bytecode (myself included) don't have easy big-endian platform access.

[mehdi_amini]

I assume Qemu can emulate big-endian systems?

A quick Google search yields https://aircrack-ng.blogspot.com/2018/10/to-be-or-not-to-be-using-qemu-to-run.html ; seems like Debian on MIPS is a possible option.

[uweigand]

In D149515#4396285, @rriddle wrote:

MLIR bytecode doesn't support big endian hosts at all, so we'll need to just disable tests in those cases, e.g. https://github.com/llvm/llvm-project/blob/52ca6ad755b0cc2aa603cfb3124bf58c04a47005/mlir/test/Bytecode/resources.mlir#L4

We've discussed it before, but most of the engineers working on the bytecode (myself included) don't have easy big-endian platform access.

Ah, right, I had forgotten about that. I've had a closer look now to identify the root causes of these failures, and I can see only two actual problems:

• Encoding/decoding of multi-byte integers in the bytecode reader/writer. This is trivial to fix, and actually in itself resolves nearly all bytecode related endian problems, including these new failures and existing testcases that were already marked unsupported.
• Handling of "dense resource element attributes". While "regular" attributes have a tryGetValues accessor that returns a custom iterator that performs endian conversion on access, the "resource" attributes use a tryGetAsArrayRef accessor that returns a standard ArrayRef, which does not do any conversion. It seems it should be possible to use a similar custom iterator approach here as well, which would fix the remaining issues.

I'd be happy to come up with patches for these issues, if you'd be willing to consider accepting such changes. Also, I'd be happy to investigate endian problems that might come up in the future. In the alternative (or in addition), we can also make access to an s390x machine available if you want to look into endian issues yourself.

[rriddle]

In D149515#4403158, @uweigand wrote:

In D149515#4396285, @rriddle wrote:

MLIR bytecode doesn't support big endian hosts at all, so we'll need to just disable tests in those cases, e.g. https://github.com/llvm/llvm-project/blob/52ca6ad755b0cc2aa603cfb3124bf58c04a47005/mlir/test/Bytecode/resources.mlir#L4

We've discussed it before, but most of the engineers working on the bytecode (myself included) don't have easy big-endian platform access.

Ah, right, I had forgotten about that. I've had a closer look now to identify the root causes of these failures, and I can see only two actual problems:

• Encoding/decoding of multi-byte integers in the bytecode reader/writer. This is trivial to fix, and actually in itself resolves nearly all bytecode related endian problems, including these new failures and existing testcases that were already marked unsupported.
• Handling of "dense resource element attributes". While "regular" attributes have a tryGetValues accessor that returns a custom iterator that performs endian conversion on access, the "resource" attributes use a tryGetAsArrayRef accessor that returns a standard ArrayRef, which does not do any conversion. It seems it should be possible to use a similar custom iterator approach here as well, which would fix the remaining issues.

I'd be happy to come up with patches for these issues, if you'd be willing to consider accepting such changes. Also, I'd be happy to investigate endian problems that might come up in the future. In the alternative (or in addition), we can also make access to an s390x machine available if you want to look into endian issues yourself.

I would love to review those patches, that would be amazing!!!

[jpienaar]

In D149515#4403405, @rriddle wrote:

In D149515#4403158, @uweigand wrote:

In D149515#4396285, @rriddle wrote:

MLIR bytecode doesn't support big endian hosts at all, so we'll need to just disable tests in those cases, e.g. https://github.com/llvm/llvm-project/blob/52ca6ad755b0cc2aa603cfb3124bf58c04a47005/mlir/test/Bytecode/resources.mlir#L4

We've discussed it before, but most of the engineers working on the bytecode (myself included) don't have easy big-endian platform access.

Ah, right, I had forgotten about that. I've had a closer look now to identify the root causes of these failures, and I can see only two actual problems:

• Encoding/decoding of multi-byte integers in the bytecode reader/writer. This is trivial to fix, and actually in itself resolves nearly all bytecode related endian problems, including these new failures and existing testcases that were already marked unsupported.
• Handling of "dense resource element attributes". While "regular" attributes have a tryGetValues accessor that returns a custom iterator that performs endian conversion on access, the "resource" attributes use a tryGetAsArrayRef accessor that returns a standard ArrayRef, which does not do any conversion. It seems it should be possible to use a similar custom iterator approach here as well, which would fix the remaining issues.

I'd be happy to come up with patches for these issues, if you'd be willing to consider accepting such changes. Also, I'd be happy to investigate endian problems that might come up in the future. In the alternative (or in addition), we can also make access to an s390x machine available if you want to look into endian issues yourself.

I would love to review those patches, that would be amazing!!!

+1, thanks!

[uweigand]

In D149515#4403405, @rriddle wrote:

In D149515#4403158, @uweigand wrote:

I'd be happy to come up with patches for these issues, if you'd be willing to consider accepting such changes. Also, I'd be happy to investigate endian problems that might come up in the future. In the alternative (or in addition), we can also make access to an s390x machine available if you want to look into endian issues yourself.

I would love to review those patches, that would be amazing!!!

I've now posted a patch for the bytecode reader/writer problem here: https://reviews.llvm.org/D153567
The resource decoding problem proved more complicated than I initially thought - I think I'll open an issue to discuss in more detail.

[uweigand]

In D149515#4441590, @uweigand wrote:

In D149515#4403405, @rriddle wrote:

In D149515#4403158, @uweigand wrote:

I'd be happy to come up with patches for these issues, if you'd be willing to consider accepting such changes. Also, I'd be happy to investigate endian problems that might come up in the future. In the alternative (or in addition), we can also make access to an s390x machine available if you want to look into endian issues yourself.

I would love to review those patches, that would be amazing!!!

I've now posted a patch for the bytecode reader/writer problem here: https://reviews.llvm.org/D153567
The resource decoding problem proved more complicated than I initially thought - I think I'll open an issue to discuss in more detail.

Issue now open as https://github.com/llvm/llvm-project/issues/63469 - any comments welcome!

mlir/docs/BytecodeFormat.md

 each element, we record the key, the value kind, and the encoded size. We avoid
 using the direct offset into the `resource_section`, as a smaller relative
 offsets provides more effective compression.
 
 ### IR Section
 
 The IR section contains the encoded form of operations within the bytecode.
 
+```
+ir_section {
+  block: block; // Single block without arguments.
+}
+```
+
 #### Operation Encoding
 
 ```
 op {
   name: varint,
   encodingMask: byte,
   location: varint,
 
   attrDict: varint?,
 
   numResults: varint?,
   resultTypes: varint[],
 
   numOperands: varint?,
   operands: varint[],
 
   numSuccessors: varint?,
   successors: varint[],
 
   regionEncoding: varint?, // (numRegions << 1) | (isIsolatedFromAbove)
-  regions: region[]
+
+  // regions are stored in a section if isIsolatedFromAbove
+  regions: (region | region_section)[] 
 }
 ```
 
 The encoding of an operation is important because this is generally the most
 commonly appearing structure in the bytecode. A single encoding is used for
 every type of operation. Given this prevelance, many of the fields of an
 operation are optional. The `encodingMask` field is a bitmask which indicates
 which of the components of the operation are present.

mlir/include/mlir/Bytecode/BytecodeReader.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BYTECODE_BYTECODEREADER_H
 #define MLIR_BYTECODE_BYTECODEREADER_H
 
 #include "mlir/IR/AsmState.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
+#include <functional>
+#include <memory>
 
 namespace llvm {
 class MemoryBufferRef;
 class SourceMgr;
 } // namespace llvm
 
 namespace mlir {
+
+/// The BytecodeReader allows to load MLIR bytecode files, while keeping the
+/// state explicitly available in order to support lazy loading.
+class BytecodeReader {
+public:
+  /// Create a bytecode reader for the given buffer. If `lazyLoad` is true,
+  /// isolated regions aren't loaded eagerly.
+  explicit BytecodeReader(
+      llvm::MemoryBufferRef buffer, const ParserConfig &config, bool lazyLoad,
+      const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef = {});
+  ~BytecodeReader();
+
+  /// Read the operations defined within the given memory buffer, containing
+  /// MLIR bytecode, into the provided block. If the reader was created with
+  /// `lazyLoad` enabled, isolated regions aren't loaded eagerly.
+  LogicalResult readTopLevel(Block *block);
+
+  /// If the reader was created with `lazyLoad` enabled, this function allows to
+  /// load the isolated region for the given operation. A nullptr is returned if
+  /// the operation doesn't have an isolated region to load.
+  std::function<LogicalResult()> getLazyOpMaterializer(Operation *op);
+
+  /// Return the number of ops that haven't been materialized yet.
+  int64_t getNumOpsToMaterialize() const;
+
+  /// Return the next operation to materialize, or nullptr if none.
+  Operation *getNextMaterializableOp() const;
+
+  /// Materialize the given operation.
+  LogicalResult materialize(Operation *op);
+
+  /// Materialize all operations.
+  LogicalResult materializeAll();
+
+  class Impl;
+
+private:
+  std::unique_ptr<Impl> impl;
+};
+
 /// Returns true if the given buffer starts with the magic bytes that signal
 /// MLIR bytecode.
 bool isBytecode(llvm::MemoryBufferRef buffer);
 
 /// Read the operations defined within the given memory buffer, containing MLIR
 /// bytecode, into the provided block.
 LogicalResult readBytecodeFile(llvm::MemoryBufferRef buffer, Block *block,
                                const ParserConfig &config);
 /// An overload with a source manager whose main file buffer is used for
 /// parsing. The lifetime of the source manager may be freely extended during
 /// parsing such that the source manager is not destroyed before the parsed IR.
 LogicalResult
 readBytecodeFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
                  Block *block, const ParserConfig &config);
+
 } // namespace mlir
 
 #endif // MLIR_BYTECODE_BYTECODEREADER_H

mlir/lib/Bytecode/Encoding.h

 // General constants
 //===----------------------------------------------------------------------===//
 
 enum {
   /// The minimum supported version of the bytecode.
   kMinSupportedVersion = 0,
 
   /// The current bytecode version.
-  kVersion = 1,
+  kVersion = 2,
 
   /// An arbitrary value used to fill alignment padding.
   kAlignmentByte = 0xCB,
 };
 
 //===----------------------------------------------------------------------===//
 // Sections
 //===----------------------------------------------------------------------===//

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

 #include "mlir/Bytecode/BytecodeReader.h"
 #include "../Encoding.h"
 #include "mlir/AsmParser/AsmParser.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Verifier.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/SourceMgr.h"
+#include <list>
+#include <memory>
 #include <optional>
 
 #define DEBUG_TYPE "mlir-bytecode-reader"
 
 using namespace mlir;
 
 /// Stringify the given section ID.
 static std::string toString(bytecode::Section::ID sectionID) {
LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
   }
   return success(!!entry.entry);
 }
 
 //===----------------------------------------------------------------------===//
 // Bytecode Reader
 //===----------------------------------------------------------------------===//
 
-namespace {
 /// This class is used to read a bytecode buffer and translate it into MLIR.
-class BytecodeReader {
+class mlir::BytecodeReader::Impl {
+  struct RegionReadState;
+  using LazyLoadableOpsInfo =
+      std::list<std::pair<Operation *, RegionReadState>>;
+  using LazyLoadableOpsMap =
+      DenseMap<Operation *, LazyLoadableOpsInfo::iterator>;
+
 public:
-  BytecodeReader(Location fileLoc, const ParserConfig &config,
-                 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef)
-      : config(config), fileLoc(fileLoc),
+  Impl(Location fileLoc, const ParserConfig &config, bool lazyLoading,
+       llvm::MemoryBufferRef buffer,
+      : config(config), fileLoc(fileLoc), lazyLoading(lazyLoading),
         attrTypeReader(stringReader, resourceReader, fileLoc),
         // Use the builtin unrealized conversion cast operation to represent
         // forward references to values that aren't yet defined.
         forwardRefOpState(UnknownLoc::get(config.getContext()),
                           "builtin.unrealized_conversion_cast", ValueRange(),
                           NoneType::get(config.getContext())),
-        bufferOwnerRef(bufferOwnerRef) {}
+        buffer(buffer), bufferOwnerRef(bufferOwnerRef) {}
 
   /// Read the bytecode defined within `buffer` into the given block.
-  LogicalResult read(llvm::MemoryBufferRef buffer, Block *block);
+  LogicalResult read(Block *block);
+
+  std::function<LogicalResult()> getLazyOpMaterializer(Operation *op) {
+    auto it = lazyLoadableOpsMap.find(op);
+    if (it == lazyLoadableOpsMap.end())
+      return nullptr;
+    return [=]() { return materialize(it); };
+  }
+
+  LogicalResult materialize(LazyLoadableOpsMap::iterator it) {
+    assert(it != lazyLoadableOpsMap.end());
+    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());
+    return result;
+  }
+
+  /// Return the number of ops that haven't been materialized yet.
+  int64_t getNumOpsToMaterialize() const { return lazyLoadableOpsMap.size(); }
+
+  /// Return the next operation to materialize, or nullptr if none.
+  Operation *getNextMaterializableOp() const {
+    if (lazyLoadableOps.empty())
+      return nullptr;
+    return lazyLoadableOps.front().first;
+  }
+
+  /// Materialize all operations.
+  LogicalResult materializeAll() {
+    while (!lazyLoadableOpsMap.empty()) {
+      if (failed(materialize(lazyLoadableOpsMap.begin())))
+        return failure();
+    }
+    return success();
+  }
 
 private:
   /// Return the context for this config.
   MLIRContext *getContext() const { return config.getContext(); }
 
   /// Parse the bytecode version.
   LogicalResult parseVersion(EncodingReader &reader);
 
parseResourceSection(EncodingReader &reader,
                        std::optional<ArrayRef<uint8_t>> resourceOffsetData);
 
   //===--------------------------------------------------------------------===//
   // IR Section
 
   /// This struct represents the current read state of a range of regions. This
   /// struct is used to enable iterative parsing of regions.
   struct RegionReadState {
-    RegionReadState(Operation *op, bool isIsolatedFromAbove)
-        : RegionReadState(op->getRegions(), isIsolatedFromAbove) {}
-    RegionReadState(MutableArrayRef<Region> regions, bool isIsolatedFromAbove)
-        : curRegion(regions.begin()), endRegion(regions.end()),
+    RegionReadState(Operation *op, EncodingReader *reader,
+                    bool isIsolatedFromAbove)
+        : RegionReadState(op->getRegions(), reader, isIsolatedFromAbove) {}
+    RegionReadState(MutableArrayRef<Region> regions, EncodingReader *reader,
+                    bool isIsolatedFromAbove)
+        : curRegion(regions.begin()), endRegion(regions.end()), reader(reader),
           isIsolatedFromAbove(isIsolatedFromAbove) {}
 
     /// The current regions being read.
     MutableArrayRef<Region>::iterator curRegion, endRegion;
+    /// This is the reader to use for this region, this pointer is pointing to
+    /// the parent region reader unless the current region is IsolatedFromAbove,
+    /// in which case the pointer is pointing to the `owningReader` which is a
+    /// section dedicated to the current region.
+    EncodingReader *reader;
+    std::unique_ptr<EncodingReader> owningReader;
 
     /// The number of values defined immediately within this region.
     unsigned numValues = 0;
 
     /// The current blocks of the region being read.
     SmallVector<Block *> curBlocks;
     Region::iterator curBlock = {};
 
     /// The number of operations remaining to be read from the current block
     /// being read.
     uint64_t numOpsRemaining = 0;
 
     /// A flag indicating if the regions being read are isolated from above.
     bool isIsolatedFromAbove = false;
   };
 
   LogicalResult parseIRSection(ArrayRef<uint8_t> sectionData, Block *block);
-  LogicalResult parseRegions(EncodingReader &reader,
-                             std::vector<RegionReadState> &regionStack,
+  LogicalResult parseRegions(std::vector<RegionReadState> &regionStack,
                              RegionReadState &readState);
   FailureOr<Operation *> parseOpWithoutRegions(EncodingReader &reader,
                                                RegionReadState &readState,
                                                bool &isIsolatedFromAbove);
 
-  LogicalResult parseRegion(EncodingReader &reader, RegionReadState &readState);
-  LogicalResult parseBlock(EncodingReader &reader, RegionReadState &readState);
+  LogicalResult parseRegion(RegionReadState &readState);
+  LogicalResult parseBlockHeader(EncodingReader &reader,
+                                 RegionReadState &readState);
   LogicalResult parseBlockArguments(EncodingReader &reader, Block *block);
 
   //===--------------------------------------------------------------------===//
   // Value Processing
 
   /// Parse an operand reference using the given reader. Returns nullptr in the
   /// case of failure.
   Value parseOperand(EncodingReader &reader);
private:
   };
 
   /// The configuration of the parser.
   const ParserConfig &config;
 
   /// A location to use when emitting errors.
   Location fileLoc;
 
+  /// Flag that indicates if lazyloading is enabled.
+  bool lazyLoading;
+
+  /// Keep track of operations that have been lazy loaded (their regions haven't
+  /// been materialized), along with the `RegionReadState` that allows to
+  /// lazy-load the regions nested under the operation.
+  LazyLoadableOpsInfo lazyLoadableOps;
+  LazyLoadableOpsMap lazyLoadableOpsMap;
+
   /// The reader used to process attribute and types within the bytecode.
   AttrTypeReader attrTypeReader;
 
   /// The version of the bytecode being read.
   uint64_t version = 0;
 
   /// The producer of the bytecode being read.
   StringRef producer;
private:
   /// references.
   Block forwardRefOps;
   /// A block containing previously created, and no longer used, forward
   /// reference operations.
   Block openForwardRefOps;
   /// An operation state used when instantiating forward references.
   OperationState forwardRefOpState;
 
+  /// Reference to the input buffer.
+  llvm::MemoryBufferRef buffer;
+
   /// The optional owning source manager, which when present may be used to
   /// extend the lifetime of the input buffer.
   const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef;
 };
-} // namespace
 
-LogicalResult BytecodeReader::read(llvm::MemoryBufferRef buffer, Block *block) {
+LogicalResult BytecodeReader::Impl::read(Block *block) {
   EncodingReader reader(buffer.getBuffer(), fileLoc);
 
   // Skip over the bytecode header, this should have already been checked.
   if (failed(reader.skipBytes(StringRef("ML\xefR").size())))
     return failure();
   // Parse the bytecode version and producer.
   if (failed(parseVersion(reader)) ||
       failed(reader.parseNullTerminatedString(producer)))
while (!reader.empty()) {
     bytecode::Section::ID sectionID;
     ArrayRef<uint8_t> sectionData;
     if (failed(reader.parseSection(sectionID, sectionData)))
       return failure();
 
     // Check for duplicate sections, we only expect one instance of each.
     if (sectionDatas[sectionID]) {
       return reader.emitError("duplicate top-level section: ",
-                              toString(sectionID));
+                              ::toString(sectionID));
     }
     sectionDatas[sectionID] = sectionData;
   }
   // Check that all of the required sections were found.
   for (int i = 0; i < bytecode::Section::kNumSections; ++i) {
     bytecode::Section::ID sectionID = static_cast<bytecode::Section::ID>(i);
     if (!sectionDatas[i] && !isSectionOptional(sectionID)) {
       return reader.emitError("missing data for top-level section: ",
-                              toString(sectionID));
+                              ::toString(sectionID));
     }
   }
 
   // Process the string section first.
   if (failed(stringReader.initialize(
           fileLoc, *sectionDatas[bytecode::Section::kString])))
     return failure();
 
if (failed(attrTypeReader.initialize(
           dialects, *sectionDatas[bytecode::Section::kAttrType],
           *sectionDatas[bytecode::Section::kAttrTypeOffset])))
     return failure();
 
   // Finally, process the IR section.
   return parseIRSection(*sectionDatas[bytecode::Section::kIR], block);
 }
 
-LogicalResult BytecodeReader::parseVersion(EncodingReader &reader) {
+LogicalResult BytecodeReader::Impl::parseVersion(EncodingReader &reader) {
   if (failed(reader.parseVarInt(version)))
     return failure();
 
   // Validate the bytecode version.
   uint64_t currentVersion = bytecode::kVersion;
   uint64_t minSupportedVersion = bytecode::kMinSupportedVersion;
   if (version < minSupportedVersion) {
     return reader.emitError("bytecode version ", version,
                             " is older than the current version of ",
                             currentVersion, ", and upgrade is not supported");
   }
   if (version > currentVersion) {
     return reader.emitError("bytecode version ", version,
                             " is newer than the current version ",
                             currentVersion);
   }
+  // Override any request to lazy-load if the bytecode version is too old.
+  if (version < 2)
+    lazyLoading = false;
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect Section
 
 LogicalResult BytecodeDialect::load(DialectReader &reader, MLIRContext *ctx) {
   if (dialect)
if (!versionBuffer.empty()) {
     loadedVersion = interface->readVersion(reader);
     if (!loadedVersion)
       return failure();
   }
   return success();
 }
 
 LogicalResult
-BytecodeReader::parseDialectSection(ArrayRef<uint8_t> sectionData) {
+BytecodeReader::Impl::parseDialectSection(ArrayRef<uint8_t> sectionData) {
   EncodingReader sectionReader(sectionData, fileLoc);
 
   // Parse the number of dialects in the section.
   uint64_t numDialects;
   if (failed(sectionReader.parseVarInt(numDialects)))
     return failure();
   dialects.resize(numDialects);
 
auto parseOpName = [&](BytecodeDialect *dialect) {
     return success();
   };
   while (!sectionReader.empty())
     if (failed(parseDialectGrouping(sectionReader, dialects, parseOpName)))
       return failure();
   return success();
 }
 
-FailureOr<OperationName> BytecodeReader::parseOpName(EncodingReader &reader) {
+FailureOr<OperationName>
+BytecodeReader::Impl::parseOpName(EncodingReader &reader) {
   BytecodeOperationName *opName = nullptr;
   if (failed(parseEntry(reader, opNames, opName, "operation name")))
     return failure();
 
   // Check to see if this operation name has already been resolved. If we
   // haven't, load the dialect and build the operation name.
   if (!opName->opName) {
     // Load the dialect and its version.
     EncodingReader versionReader(opName->dialect->versionBuffer, fileLoc);
     DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
                                 versionReader);
     if (failed(opName->dialect->load(dialectReader, getContext())))
       return failure();
     opName->opName.emplace((opName->dialect->name + "." + opName->name).str(),
                            getContext());
   }
   return *opName->opName;
 }
 
 //===----------------------------------------------------------------------===//
 // Resource Section
 
-LogicalResult BytecodeReader::parseResourceSection(
+LogicalResult BytecodeReader::Impl::parseResourceSection(
     EncodingReader &reader, std::optional<ArrayRef<uint8_t>> resourceData,
     std::optional<ArrayRef<uint8_t>> resourceOffsetData) {
   // Ensure both sections are either present or not.
   if (resourceData.has_value() != resourceOffsetData.has_value()) {
     if (resourceOffsetData)
       return emitError(fileLoc, "unexpected resource offset section when "
                                 "resource section is not present");
     return emitError(
LogicalResult BytecodeReader::Impl::parseResourceSection(
   return resourceReader.initialize(fileLoc, config, dialects, stringReader,
                                    *resourceData, *resourceOffsetData,
                                    dialectReader, bufferOwnerRef);
 }
 
 //===----------------------------------------------------------------------===//
 // IR Section
 
-LogicalResult BytecodeReader::parseIRSection(ArrayRef<uint8_t> sectionData,
-                                             Block *block) {
+LogicalResult
+BytecodeReader::Impl::parseIRSection(ArrayRef<uint8_t> sectionData,
   EncodingReader reader(sectionData, fileLoc);
 
   // A stack of operation regions currently being read from the bytecode.
   std::vector<RegionReadState> regionStack;
 
   // Parse the top-level block using a temporary module operation.
   OwningOpRef<ModuleOp> moduleOp = ModuleOp::create(fileLoc);
-  regionStack.emplace_back(*moduleOp, /*isIsolatedFromAbove=*/true);
+  regionStack.emplace_back(*moduleOp, &reader, /*isIsolatedFromAbove=*/true);
   regionStack.back().curBlocks.push_back(moduleOp->getBody());
   regionStack.back().curBlock = regionStack.back().curRegion->begin();
-  if (failed(parseBlock(reader, regionStack.back())))
+  if (failed(parseBlockHeader(reader, regionStack.back())))
     return failure();
   valueScopes.emplace_back();
   valueScopes.back().push(regionStack.back());
 
   // Iteratively parse regions until everything has been resolved.
   while (!regionStack.empty())
-    if (failed(parseRegions(reader, regionStack, regionStack.back())))
+    if (failed(parseRegions(regionStack, regionStack.back())))
       return failure();
   if (!forwardRefOps.empty()) {
     return reader.emitError(
         "not all forward unresolved forward operand references");
   }
 
   // Resolve dialect version.
   for (const BytecodeDialect &byteCodeDialect : dialects) {
BytecodeReader::Impl::parseIRSection(ArrayRef<uint8_t> sectionData,
   // Splice the parsed operations over to the provided top-level block.
   auto &parsedOps = moduleOp->getBody()->getOperations();
   auto &destOps = block->getOperations();
   destOps.splice(destOps.end(), parsedOps, parsedOps.begin(), parsedOps.end());
   return success();
 }
 
 LogicalResult
-BytecodeReader::parseRegions(EncodingReader &reader,
-                             std::vector<RegionReadState> &regionStack,
+BytecodeReader::Impl::parseRegions(std::vector<RegionReadState> &regionStack,
                                    RegionReadState &readState) {
   // Read the regions of this operation.
+  // Process regions, blocks, and operations until the end or if a nested
+  // region is encountered. In this case we push a new state in regionStack and
+  // return, the processing of the current region will resume afterward.
   for (; readState.curRegion != readState.endRegion; ++readState.curRegion) {
     // If the current block hasn't been setup yet, parse the header for this
-    // region.
+    // region. The current block is already setup when this function was
+    // interrupted to recurse down in a nested region and we resume the current
+    // block after processing the nested region.
     if (readState.curBlock == Region::iterator()) {
-      if (failed(parseRegion(reader, readState)))
+      if (failed(parseRegion(readState)))
         return failure();
 
       // If the region is empty, there is nothing to more to do.
       if (readState.curRegion->empty())
         continue;
     }
 
     // Parse the blocks within the region.
+    EncodingReader &reader = *readState.reader;
     do {
       while (readState.numOpsRemaining--) {
         // Read in the next operation. We don't read its regions directly, we
         // handle those afterwards as necessary.
         bool isIsolatedFromAbove = false;
         FailureOr<Operation *> op =
             parseOpWithoutRegions(reader, readState, isIsolatedFromAbove);
         if (failed(op))
           return failure();
 
-        // If the op has regions, add it to the stack for processing.
+        // If the op has regions, add it to the stack for processing and return:
+        // we stop the processing of the current region and resume it after the
+        // inner one is completed. Unless LazyLoading is activated in which case
+        // nested region parsing is delayed.
         if ((*op)->getNumRegions()) {
-          regionStack.emplace_back(*op, isIsolatedFromAbove);
+          RegionReadState childState(*op, &reader, isIsolatedFromAbove);
+          // Isolated regions are encoded as a section in version 2 and above.
+          if (version >= 2 && isIsolatedFromAbove) {
+            bytecode::Section::ID sectionID;
+            ArrayRef<uint8_t> sectionData;
+            if (failed(reader.parseSection(sectionID, sectionData)))
+              return failure();
+            if (sectionID != bytecode::Section::kIR) {
+              emitError(fileLoc, "expected IR section for region");
+              return failure();
+            }
+            childState.owningReader =
+                std::make_unique<EncodingReader>(sectionData, fileLoc);
+            childState.reader = childState.owningReader.get();
+          }
+
+          if (lazyLoading) {
+            lazyLoadableOps.push_back(
+                std::make_pair(*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.
           if (isIsolatedFromAbove)
             valueScopes.emplace_back();
           return success();
         }
       }
 
       // Move to the next block of the region.
       if (++readState.curBlock == readState.curRegion->end())
         break;
-      if (failed(parseBlock(reader, readState)))
+      if (failed(parseBlockHeader(reader, readState)))
         return failure();
     } while (true);
 
     // Reset the current block and any values reserved for this region.
     readState.curBlock = {};
     valueScopes.back().pop(readState);
   }
 
   // When the regions have been fully parsed, pop them off of the read stack. If
   // the regions were isolated from above, we also pop the last value scope.
-  if (readState.isIsolatedFromAbove)
+  if (readState.isIsolatedFromAbove) {
+    assert(!valueScopes.empty());
     valueScopes.pop_back();
+  }
+  assert(!regionStack.empty());
   regionStack.pop_back();
   return success();
 }
 
 FailureOr<Operation *>
-BytecodeReader::parseOpWithoutRegions(EncodingReader &reader,
+BytecodeReader::Impl::parseOpWithoutRegions(EncodingReader &reader,
                                             RegionReadState &readState,
                                             bool &isIsolatedFromAbove) {
   // Parse the name of the operation.
   FailureOr<OperationName> opName = parseOpName(reader);
   if (failed(opName))
     return failure();
 
   // Parse the operation mask, which indicates which components of the operation
   // are present.
   uint8_t opMask;
BytecodeReader::Impl::parseOpWithoutRegions(EncodingReader &reader,
 
   // If the operation had results, update the value references.
   if (op->getNumResults() && failed(defineValues(reader, op->getResults())))
     return failure();
 
   return op;
 }
 
-LogicalResult BytecodeReader::parseRegion(EncodingReader &reader,
-                                          RegionReadState &readState) {
+LogicalResult BytecodeReader::Impl::parseRegion(RegionReadState &readState) {
+  EncodingReader &reader = *readState.reader;
+
   // Parse the number of blocks in the region.
   uint64_t numBlocks;
   if (failed(reader.parseVarInt(numBlocks)))
     return failure();
 
   // If the region is empty, there is nothing else to do.
   if (numBlocks == 0)
     return success();
for (uint64_t i = 0; i < numBlocks; ++i) {
     readState.curRegion->push_back(readState.curBlocks.back());
   }
 
   // Prepare the current value scope for this region.
   valueScopes.back().push(readState);
 
   // Parse the entry block of the region.
   readState.curBlock = readState.curRegion->begin();
-  return parseBlock(reader, readState);
+  return parseBlockHeader(reader, readState);
 }
 
-LogicalResult BytecodeReader::parseBlock(EncodingReader &reader,
-                                         RegionReadState &readState) {
+LogicalResult
+BytecodeReader::Impl::parseBlockHeader(EncodingReader &reader,
   bool hasArgs;
   if (failed(reader.parseVarIntWithFlag(readState.numOpsRemaining, hasArgs)))
     return failure();
 
   // Parse the arguments of the block.
   if (hasArgs && failed(parseBlockArguments(reader, &*readState.curBlock)))
     return failure();
 
   // We don't parse the operations of the block here, that's done elsewhere.
   return success();
 }
 
-LogicalResult BytecodeReader::parseBlockArguments(EncodingReader &reader,
+LogicalResult BytecodeReader::Impl::parseBlockArguments(EncodingReader &reader,
                                                         Block *block) {
   // Parse the value ID for the first argument, and the number of arguments.
   uint64_t numArgs;
   if (failed(reader.parseVarInt(numArgs)))
     return failure();
 
   SmallVector<Type> argTypes;
   SmallVector<Location> argLocs;
   argTypes.reserve(numArgs);
LogicalResult BytecodeReader::Impl::parseBlockArguments(EncodingReader &reader,
   }
   block->addArguments(argTypes, argLocs);
   return defineValues(reader, block->getArguments());
 }
 
 //===----------------------------------------------------------------------===//
 // Value Processing
 
-Value BytecodeReader::parseOperand(EncodingReader &reader) {
+Value BytecodeReader::Impl::parseOperand(EncodingReader &reader) {
   std::vector<Value> &values = valueScopes.back().values;
   Value *value = nullptr;
   if (failed(parseEntry(reader, values, value, "value")))
     return Value();
 
   // Create a new forward reference if necessary.
   if (!*value)
     *value = createForwardRef();
   return *value;
 }
 
-LogicalResult BytecodeReader::defineValues(EncodingReader &reader,
+LogicalResult BytecodeReader::Impl::defineValues(EncodingReader &reader,
                                                  ValueRange newValues) {
   ValueScope &valueScope = valueScopes.back();
   std::vector<Value> &values = valueScope.values;
 
   unsigned &valueID = valueScope.nextValueIDs.back();
   unsigned valueIDEnd = valueID + newValues.size();
   if (valueIDEnd > values.size()) {
     return reader.emitError(
         "value index range was outside of the expected range for "
if (Value oldValue = std::exchange(values[valueID], newValue)) {
 
       oldValue.replaceAllUsesWith(newValue);
       forwardRefOp->moveBefore(&openForwardRefOps, openForwardRefOps.end());
     }
   }
   return success();
 }
 
-Value BytecodeReader::createForwardRef() {
+Value BytecodeReader::Impl::createForwardRef() {
   // Check for an avaliable existing operation to use. Otherwise, create a new
   // fake operation to use for the reference.
   if (!openForwardRefOps.empty()) {
     Operation *op = &openForwardRefOps.back();
     op->moveBefore(&forwardRefOps, forwardRefOps.end());
   } else {
     forwardRefOps.push_back(Operation::create(forwardRefOpState));
   }
   return forwardRefOps.back().getResult(0);
 }
 
 //===----------------------------------------------------------------------===//
 // Entry Points
 //===----------------------------------------------------------------------===//
 
+BytecodeReader::~BytecodeReader() = default;
+
+LogicalResult BytecodeReader::readTopLevel(Block *block) {
+  return impl->read(block);
+}
+BytecodeReader::BytecodeReader(
+    llvm::MemoryBufferRef buffer, const ParserConfig &config, bool lazyLoading,
+    const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
+  Location sourceFileLoc =
+      FileLineColLoc::get(config.getContext(), buffer.getBufferIdentifier(),
+                          /*line=*/0, /*column=*/0);
+  impl = std::make_unique<Impl>(sourceFileLoc, config, lazyLoading, buffer,
+                                bufferOwnerRef);
+}
+std::function<LogicalResult()>
+BytecodeReader::getLazyOpMaterializer(Operation *op) {
+  return impl->getLazyOpMaterializer(op);
+}
+
+int64_t BytecodeReader::getNumOpsToMaterialize() const {
+  return impl->getNumOpsToMaterialize();
+}
+
+Operation *BytecodeReader::getNextMaterializableOp() const {
+  return impl->getNextMaterializableOp();
+}
+
+LogicalResult BytecodeReader::materialize(Operation *op) {
+  auto materializer = getLazyOpMaterializer(op);
+  if (!materializer)
+    return success();
+  return materializer();
+}
+
+LogicalResult BytecodeReader::materializeAll() {
+  return impl->materializeAll();
+}
+
 bool mlir::isBytecode(llvm::MemoryBufferRef buffer) {
   return buffer.getBuffer().startswith("ML\xefR");
 }
 
 /// Read the bytecode from the provided memory buffer reference.
 /// `bufferOwnerRef` if provided is the owning source manager for the buffer,
 /// and may be used to extend the lifetime of the buffer.
 static LogicalResult
 readBytecodeFileImpl(llvm::MemoryBufferRef buffer, Block *block,
                      const ParserConfig &config,
                      const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
   Location sourceFileLoc =
       FileLineColLoc::get(config.getContext(), buffer.getBufferIdentifier(),
                           /*line=*/0, /*column=*/0);
   if (!isBytecode(buffer)) {
     return emitError(sourceFileLoc,
                      "input buffer is not an MLIR bytecode file");
   }
 
-  BytecodeReader reader(sourceFileLoc, config, bufferOwnerRef);
-  return reader.read(buffer, block);
+  BytecodeReader::Impl reader(sourceFileLoc, config, /*lazyLoading=*/false,
+                              buffer, bufferOwnerRef);
+  return reader.read(block);
 }
 
 LogicalResult mlir::readBytecodeFile(llvm::MemoryBufferRef buffer, Block *block,
                                      const ParserConfig &config) {
   return readBytecodeFileImpl(buffer, block, config, /*bufferOwnerRef=*/{});
 }
 LogicalResult
 mlir::readBytecodeFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
                        Block *block, const ParserConfig &config) {
   return readBytecodeFileImpl(
       *sourceMgr->getMemoryBuffer(sourceMgr->getMainFileID()), block, config,
       sourceMgr);
 }

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

void 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())
+    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));
+    }
   }
 }
 
 void BytecodeWriter::writeRegion(EncodingEmitter &emitter, Region *region) {
   // If the region is empty, we only need to emit the number of blocks (which is
   // zero).
   if (region->empty())
     return emitter.emitVarInt(/*numBlocks*/ 0);

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
+  test.isolated_region %0 {
+    "test.consumer"(%0) : (index) -> ()
+  }
+  %result:2 = "test.op"() : () -> (index, index)
+  test.isolated_region %result#1 {
+    "test.consumer"(%result#1) : (index) -> ()
+  }
+  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: Has 1 ops to materialize
+
+// CHECK: Before Materializing...
+// CHECK: "builtin.module"() ({
+// CHECK-NOT: func
+// CHECK: Materializing...
+// CHECK: "builtin.module"() ({
+// CHECK: "func.func"() <{function_type = () -> (), sym_name = "op_with_passthrough_region_args"}> ({
+// 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: 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
+// CHECK: Materializing...
+// CHECK: test.isolated_region
+// CHECK: ^bb0(%arg0: index):
+// CHECK:  test.consumer
+// CHECK: Has 1 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 0 ops to materialize

mlir/test/Bytecode/invalid/invalid-structure.mlir

 // This file contains various failure test cases related to the structure of
 // a bytecode file.
 
 // Bytecode currently does not support big-endian platforms
 // UNSUPPORTED: target=s390x-{{.*}}
 
 //===--------------------------------------------------------------------===//
 // Version
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-structure-version.mlirbc 2>&1 | FileCheck %s --check-prefix=VERSION
-// VERSION: bytecode version 127 is newer than the current version 1
+// VERSION: bytecode version 127 is newer than the current version 2
 
 //===--------------------------------------------------------------------===//
 // Producer
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-structure-producer.mlirbc 2>&1 | FileCheck %s --check-prefix=PRODUCER
 // PRODUCER: malformed null-terminated string, no null character found
 

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestBuiltinAttributeInterfaces.cpp
   TestClone.cpp
   TestDiagnostics.cpp
   TestDominance.cpp
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
+  TestLazyLoading.cpp
   TestOpaqueLoc.cpp
   TestOperationEquals.cpp
   TestPrintDefUse.cpp
   TestPrintInvalid.cpp
   TestPrintNesting.cpp
   TestSideEffects.cpp
   TestSlicing.cpp
   TestSymbolUses.cpp

mlir/test/lib/IR/TestLazyLoading.cpp

+//===- TestLazyLoading.cpp - Pass to test operation lazy loading  ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestDialect.h"
+#include "mlir/Bytecode/BytecodeReader.h"
+#include "mlir/Bytecode/BytecodeWriter.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/MemoryBufferRef.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+namespace {
+
+/// This is a test pass which LazyLoads the current operation recursively.
+struct LazyLoadingPass : public PassWrapper<LazyLoadingPass, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LazyLoadingPass)
+
+  StringRef getArgument() const final { return "test-lazy-loading"; }
+  StringRef getDescription() const final { return "Test LazyLoading of op"; }
+  LazyLoadingPass() = default;
+  LazyLoadingPass(const LazyLoadingPass &) {}
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    std::string bytecode;
+    {
+      BytecodeWriterConfig config;
+      if (version >= 0)
+        config.setDesiredBytecodeVersion(version);
+      llvm::raw_string_ostream os(bytecode);
+      BytecodeWriterResult result = writeBytecodeToFile(op, os, config);
+      if (version >= 0 && result.minVersion != version) {
+        op->emitError() << "Failed to write bytecode at version "
+                        << (int)version << ", emitted minimum version "
+                        << result.minVersion << " instead.";
+        signalPassFailure();
+        return;
+      }
+    }
+    llvm::MemoryBufferRef buffer(bytecode, "test-lazy-loading");
+    Block block;
+    ParserConfig config(op->getContext(), /*verifyAfterParse=*/false);
+    BytecodeReader reader(buffer, config,
+                          /*lazyLoad=*/true);
+    if (failed(reader.readTopLevel(&block))) {
+      op->emitError() << "Failed to read bytecode";
+      return;
+    }
+
+    llvm::outs() << "Has " << reader.getNumOpsToMaterialize()
+                 << " ops to materialize\n";
+
+    // Recursively print the operations, before and after lazy loading.
+    while (Operation *op = reader.getNextMaterializableOp()) {
+      llvm::outs() << "\n\nBefore Materializing...\n\n";
+      op->print(llvm::outs());
+      llvm::outs() << "\n\nMaterializing...\n\n";
+      if (failed(reader.materialize(op))) {
+        op->emitError() << "Failed to materialize";
+        signalPassFailure();
+        return;
+      }
+      op->print(llvm::outs());
+      llvm::outs() << "\n";
+      llvm::outs() << "Has " << reader.getNumOpsToMaterialize()
+                   << " ops to materialize\n";
+    }
+  }
+  Option<int> version{*this, "bytecode-version",
+                      llvm::cl::desc("Specifies the bytecode version to use."),
+                      llvm::cl::init(-1)};
+};
+} // namespace
+
+namespace mlir {
+void registerLazyLoadingTestPasses() { PassRegistration<LazyLoadingPass>(); }
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 
 using namespace llvm;
 using namespace mlir;
 
 // Defined in the test directory, no public header.
 namespace mlir {
 void registerConvertToTargetEnvPass();
 void registerCloneTestPasses();
+void registerLazyLoadingTestPasses();
 void registerPassManagerTestPass();
 void registerPrintSpirvAvailabilityPass();
 void registerLoopLikeInterfaceTestPasses();
 void registerShapeFunctionTestPasses();
 void registerSideEffectTestPasses();
 void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
 void registerRegionTestPasses();

 } // namespace test
 
 #ifdef MLIR_INCLUDE_TESTS
 void registerTestPasses() {
   registerCloneTestPasses();
   registerConvertToTargetEnvPass();
   registerPassManagerTestPass();
   registerPrintSpirvAvailabilityPass();
+  registerLazyLoadingTestPasses();
   registerLoopLikeInterfaceTestPasses();
   registerShapeFunctionTestPasses();
   registerSideEffectTestPasses();
   registerSliceAnalysisTestPass();
   registerSymbolTestPasses();
   registerRegionTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAffineReifyValueBoundsPass();