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
35namespace tvm {
36
37class IRModule;
38
39// Forward class for PassInstrumentNode methods
40namespace transform {
41class PassInfo;
42} // namespace transform
43
44namespace 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 */
102class 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 */
149class 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