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
19 using namespace mlir;
20
21 namespace {
22
23 /// This is a test pass which LazyLoads the current operation recursively.
24 struct LazyLoadingPass : public PassWrapper<LazyLoadingPass, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anond11562330111::LazyLoadingPass25 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LazyLoadingPass)
26
27 StringRef getArgument() const final { return "test-lazy-loading"; }
getDescription__anond11562330111::LazyLoadingPass28 StringRef getDescription() const final { return "Test LazyLoading of op"; }
29 LazyLoadingPass() = default;
LazyLoadingPass__anond11562330111::LazyLoadingPass30 LazyLoadingPass(const LazyLoadingPass &) {}
31
runOnOperation__anond11562330111::LazyLoadingPass32 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(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(reader.readTopLevel(&block, [&](Operation *op) {
54 toLoadOps.push_back(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(llvm::outs());
70 llvm::outs() << "\n\nMaterializing...\n\n";
71 if (failed(reader.materialize(toLoad, [&](Operation *op) {
72 toLoadOps.push_back(op);
73 return false;
74 }))) {
75 toLoad->emitError() << "failed to materialize";
76 signalPassFailure();
77 return;
78 }
79 toLoad->print(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(-1)};
88 };
89 } // namespace
90
91 namespace mlir {
registerLazyLoadingTestPasses()92 void registerLazyLoadingTestPasses() { PassRegistration<LazyLoadingPass>(); }
93 } // namespace mlir
94