1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file tvm/ir/instrument.h |
22 | * |
23 | * This file introduces a pass instrument infrastructure, inspired by LLVM and MLIR. |
24 | * It inserts instrumentation points around passes. |
25 | */ |
26 | #ifndef TVM_IR_INSTRUMENT_H_ |
27 | #define TVM_IR_INSTRUMENT_H_ |
28 | |
29 | #include <tvm/node/reflection.h> |
30 | #include <tvm/runtime/container/string.h> |
31 | |
32 | #include <utility> |
33 | #include <vector> |
34 | |
35 | namespace tvm { |
36 | |
37 | class IRModule; |
38 | |
39 | // Forward class for PassInstrumentNode methods |
40 | namespace transform { |
41 | class PassInfo; |
42 | } // namespace transform |
43 | |
44 | namespace instrument { |
45 | |
46 | /*! |
47 | * \brief PassInstrumentNode forms an instrument implementation. |
48 | * It provides API for users to register callbacks at different instrumentation points. |
49 | * |
50 | * Within a PassContext, call sequence of a PassInstrument implementation is like: |
51 | * |
52 | * with PassContext(instruments=[pi]): # pi = a PassInstrument implementation |
53 | * pi.EnterPassContext() |
54 | * |
55 | * if pi.ShouldRun(Pass1): |
56 | * pi.RunBeforePass() |
57 | * Pass1() |
58 | * pi.RunAfterPass() |
59 | * |
60 | * if pi.ShouldRun(Pass2): |
61 | * pi.RunBeforePass() |
62 | * Pass2() |
63 | * pi.RunAfterPass() |
64 | * |
65 | * pi.ExitPassContext() |
66 | * |
67 | * `EnterPassContext` and `ExitPassContext` are only called once when entering/exiting a |
68 | * PassContext. `ShouldRun`, `RunBeforePass` and `RunAfterPass` are called multiple times depending |
69 | * on how many passes. |
70 | * |
71 | * If there are multiple pass instrumentations provided, the instrument points are the same. |
72 | * PassInstrument implementations' callbacks are called in order: |
73 | * |
74 | * with PassContext(instruments=[pi1, pi2]): # pi1, pi2 = two distinct PassInstrument impls |
75 | * pi.EnterPassContext() for pi in instruments |
76 | * |
77 | * should_run = all([pi.ShoudRun(Pass1) for pi in instruments)]) |
78 | * if (should_run) |
79 | * pi.RunBeforePass() for pi in instruments |
80 | * Pass1() |
81 | * pi.RunAfterPass() for pi in instruments |
82 | * |
83 | * should_run = all([pi.ShouldRun(Pass2) for pi in instruments)]) |
84 | * if (should_run) |
85 | * pi.RunBeforePass() for pi in instruments |
86 | * Pass2() |
87 | * pi.RunAfterPass() for pi in instruments |
88 | * |
89 | * pi.ExitPassContext() for pi in instruments |
90 | * |
91 | * Note: |
92 | * 1. Assume there is no dependency between PassInstrument implementations in `instruments` . |
93 | * 2. `EnterPassContext` and `ExitPassContext` have `with` behavior (see PassContext and its FFI): |
94 | * If there is any exception raised in `ShouldRun()`, `RunBeforePass()`, `RunAfterPass()` and |
95 | * `Pass()`, `ExitPassContext()` is still called. |
96 | * 3. In mutiple PassInstrument instances scenario, callbacks are called in order: |
97 | * If one throws exceptions, remainings will not be called. |
98 | * |
99 | * \sa PassInstrument |
100 | * \sa src/ir/transform.cc |
101 | */ |
102 | class PassInstrumentNode : public Object { |
103 | public: |
104 | /*! \brief Name of this pass instrument object. */ |
105 | String name; |
106 | |
107 | virtual ~PassInstrumentNode() {} |
108 | |
109 | /*! \brief Instrument when entering PassContext. Called once within a PassContext. */ |
110 | virtual void EnterPassContext() const = 0; |
111 | |
112 | /*! \brief Instrument when exiting PassContext. Called once within a PassContext. */ |
113 | virtual void ExitPassContext() const = 0; |
114 | |
115 | /*! |
116 | * \brief Determine whether to run the pass or not. Called multiple times depend on number of |
117 | * passes. |
118 | * \param mod The module that an optimization pass runs on. |
119 | * \param info The pass information. |
120 | * |
121 | * \return true to run the pass; false to skip the pass. |
122 | */ |
123 | virtual bool ShouldRun(const IRModule& mod, const transform::PassInfo& info) const = 0; |
124 | |
125 | /*! |
126 | * \brief Instrument before pass run. Called multiple times depend on number of passes. |
127 | * \param mod The module that an optimization pass runs on. |
128 | * \param info The pass information. |
129 | */ |
130 | virtual void RunBeforePass(const IRModule& mod, const transform::PassInfo& info) const = 0; |
131 | |
132 | /*! |
133 | * \brief Instrument after pass run. Called multiple time depend on number of passes. |
134 | * \param mod The module that an optimization pass runs on. |
135 | * \param info The pass information. |
136 | */ |
137 | virtual void RunAfterPass(const IRModule& mod, const transform::PassInfo& info) const = 0; |
138 | |
139 | void VisitAttrs(AttrVisitor* v) { v->Visit("name" , &name); } |
140 | |
141 | static constexpr const char* _type_key = "instrument.PassInstrument" ; |
142 | TVM_DECLARE_BASE_OBJECT_INFO(PassInstrumentNode, Object); |
143 | }; |
144 | |
145 | /*! |
146 | * \brief Managed reference class for PassInstrumentNode |
147 | * \sa PassInstrumentNode |
148 | */ |
149 | class PassInstrument : public ObjectRef { |
150 | public: |
151 | TVM_DEFINE_OBJECT_REF_METHODS(PassInstrument, ObjectRef, PassInstrumentNode); |
152 | }; |
153 | |
154 | } // namespace instrument |
155 | } // namespace tvm |
156 | |
157 | #endif // TVM_IR_INSTRUMENT_H_ |
158 | |