94 lines · cpp
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-exception6//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(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} // namespace90 91namespace mlir {92void registerLazyLoadingTestPasses() { PassRegistration<LazyLoadingPass>(); }93} // namespace mlir94