1//===- TestLazyLoading.cpp - Pass to test operation lazy loading ---------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "TestDialect.h"
10#include "mlir/Bytecode/BytecodeReader.h"
11#include "mlir/Bytecode/BytecodeWriter.h"
12#include "mlir/IR/BuiltinOps.h"
13#include "mlir/IR/OperationSupport.h"
14#include "mlir/Pass/Pass.h"
15#include "llvm/Support/MemoryBufferRef.h"
16#include "llvm/Support/raw_ostream.h"
17#include <list>
18
19using namespace mlir;
20
21namespace {
22
23/// This is a test pass which LazyLoads the current operation recursively.
24struct LazyLoadingPass : public PassWrapper<LazyLoadingPass, OperationPass<>> {
25 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LazyLoadingPass)
26
27 StringRef getArgument() const final { return "test-lazy-loading"; }
28 StringRef getDescription() const final { return "Test LazyLoading of op"; }
29 LazyLoadingPass() = default;
30 LazyLoadingPass(const LazyLoadingPass &) {}
31
32 void runOnOperation() override {
33 Operation *op = getOperation();
34 std::string bytecode;
35 {
36 BytecodeWriterConfig config;
37 if (version >= 0)
38 config.setDesiredBytecodeVersion(version);
39 llvm::raw_string_ostream os(bytecode);
40 if (failed(Result: writeBytecodeToFile(op, os, config))) {
41 op->emitError() << "failed to write bytecode at version "
42 << (int)version;
43 signalPassFailure();
44 return;
45 }
46 }
47 llvm::MemoryBufferRef buffer(bytecode, "test-lazy-loading");
48 Block block;
49 ParserConfig config(op->getContext(), /*verifyAfterParse=*/false);
50 BytecodeReader reader(buffer, config,
51 /*lazyLoad=*/true);
52 std::list<Operation *> toLoadOps;
53 if (failed(Result: reader.readTopLevel(block: &block, lazyOps: [&](Operation *op) {
54 toLoadOps.push_back(x: op);
55 return false;
56 }))) {
57 op->emitError() << "failed to read bytecode";
58 return;
59 }
60
61 llvm::outs() << "Has " << reader.getNumOpsToMaterialize()
62 << " ops to materialize\n";
63
64 // Recursively print the operations, before and after lazy loading.
65 while (!toLoadOps.empty()) {
66 Operation *toLoad = toLoadOps.front();
67 toLoadOps.pop_front();
68 llvm::outs() << "\n\nBefore Materializing...\n\n";
69 toLoad->print(os&: llvm::outs());
70 llvm::outs() << "\n\nMaterializing...\n\n";
71 if (failed(Result: reader.materialize(op: toLoad, lazyOpsCallback: [&](Operation *op) {
72 toLoadOps.push_back(x: op);
73 return false;
74 }))) {
75 toLoad->emitError() << "failed to materialize";
76 signalPassFailure();
77 return;
78 }
79 toLoad->print(os&: llvm::outs());
80 llvm::outs() << "\n";
81 llvm::outs() << "Has " << reader.getNumOpsToMaterialize()
82 << " ops to materialize\n";
83 }
84 }
85 Option<int> version{*this, "bytecode-version",
86 llvm::cl::desc("Specifies the bytecode version to use."),
87 llvm::cl::init(Val: -1)};
88};
89} // namespace
90
91namespace mlir {
92void registerLazyLoadingTestPasses() { PassRegistration<LazyLoadingPass>(); }
93} // namespace mlir
94

source code of mlir/test/lib/IR/TestLazyLoading.cpp