1 //===- BytecodeReader.h - MLIR Bytecode Reader ------------------*- C++ -*-===// 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 // This header defines interfaces to read MLIR bytecode files/streams. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #ifndef MLIR_BYTECODE_BYTECODEREADER_H 14 #define MLIR_BYTECODE_BYTECODEREADER_H 15 16 #include "mlir/IR/AsmState.h" 17 #include "mlir/Support/LLVM.h" 18 #include <functional> 19 #include <memory> 20 21 namespace llvm { 22 class MemoryBufferRef; 23 class SourceMgr; 24 } // namespace llvm 25 26 namespace mlir { 27 /// The BytecodeReader allows to load MLIR bytecode files, while keeping the 28 /// state explicitly available in order to support lazy loading. 29 /// The `finalize` method must be called before destruction. 30 class BytecodeReader { 31 public: 32 /// Create a bytecode reader for the given buffer. If `lazyLoad` is true, 33 /// isolated regions aren't loaded eagerly. 34 explicit BytecodeReader( 35 llvm::MemoryBufferRef buffer, const ParserConfig &config, bool lazyLoad, 36 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef = {}); 37 ~BytecodeReader(); 38 39 /// Read the operations defined within the given memory buffer, containing 40 /// MLIR bytecode, into the provided block. If the reader was created with 41 /// `lazyLoad` enabled, isolated regions aren't loaded eagerly. 42 /// The lazyOps call back is invoked for every ops that can be lazy-loaded. 43 /// This let the client decide if the op should be materialized 44 /// immediately or delayed. 45 LogicalResult readTopLevel( 46 Block *block, llvm::function_ref<bool(Operation *)> lazyOps = 47 [](Operation *) { return false; }); 48 49 /// Return the number of ops that haven't been materialized yet. 50 int64_t getNumOpsToMaterialize() const; 51 52 /// Return true if the provided op is materializable. 53 bool isMaterializable(Operation *op); 54 55 /// Materialize the provide operation. The provided operation must be 56 /// materializable. 57 /// The lazyOps call back is invoked for every ops that can be lazy-loaded. 58 /// This let the client decide if the op should be materialized immediately or 59 /// delayed. 60 /// !! Using this materialize withing an IR walk() can be confusing: make sure 61 /// to use a PreOrder traversal !! 62 LogicalResult materialize( 63 Operation *op, llvm::function_ref<bool(Operation *)> lazyOpsCallback = 64 [](Operation *) { return false; }); 65 66 /// Finalize the lazy-loading by calling back with every op that hasn't been 67 /// materialized to let the client decide if the op should be deleted or 68 /// materialized. The op is materialized if the callback returns true, deleted 69 /// otherwise. The implementation of the callback must be thread-safe. 70 LogicalResult finalize(function_ref<bool(Operation *)> shouldMaterialize = 71 [](Operation *) { return true; }); 72 73 class Impl; 74 75 private: 76 std::unique_ptr<Impl> impl; 77 }; 78 79 /// Returns true if the given buffer starts with the magic bytes that signal 80 /// MLIR bytecode. 81 bool isBytecode(llvm::MemoryBufferRef buffer); 82 83 /// Read the operations defined within the given memory buffer, containing MLIR 84 /// bytecode, into the provided block. 85 LogicalResult readBytecodeFile(llvm::MemoryBufferRef buffer, Block *block, 86 const ParserConfig &config); 87 /// An overload with a source manager whose main file buffer is used for 88 /// parsing. The lifetime of the source manager may be freely extended during 89 /// parsing such that the source manager is not destroyed before the parsed IR. 90 LogicalResult 91 readBytecodeFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr, 92 Block *block, const ParserConfig &config); 93 94 } // namespace mlir 95 96 #endif // MLIR_BYTECODE_BYTECODEREADER_H 97