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