mlir/test/lib/IR/TestLazyLoading.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-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<>> {
  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 } // namespace
  90
  91 namespace mlir {
  92 void registerLazyLoadingTestPasses() { PassRegistration<LazyLoadingPass>(); }
  93 } // namespace mlir