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 src/relay/transforms/compiler_function_utils.h |
22 | * \brief Helper passes for working with functions with the "Compiler" attribute. |
23 | * |
24 | * Those wishing to use the "RelayToTIR" custom pass machinery to do IRModule-at-a-time external |
25 | * codegen may find the following helpers useful: |
26 | * |
27 | * - The \p OutlineCompilerFunctionsWithExistingGlobalSymbols pass will lift inline functions with |
28 | * a matching "Compiler" attribute to be global functions, using the "global_symbol" attribute |
29 | * already assigned. Can be used before custom lowering. |
30 | * |
31 | * Note that ideally "Compiler" attributed functions would be made global functions as early as |
32 | * possible and would stay that way. However, the GraphExecutorCodegen and AOTExecutorCodegen |
33 | * assume the entire model can be represented by a single 'main' function, and the Inline pass |
34 | * is run to respect that assumption. So this pass is mostly just to undo that Pass after modules |
35 | * have passed through the 'codegen' keyhole. |
36 | * |
37 | * See also OutlineCompilerFunctionsMutator in src/relay/backend/contrib/ethosu/codegen.cc. |
38 | * |
39 | * - (The \p OutlineCompilerFunctions pass is a more general version of the above which can use |
40 | * a custom cache to both allocate "global_symbol" names and ensure two structurally equal |
41 | * functions are assigned the same name, and thus lowered only once. This is used by Collage |
42 | * when preparing the optimally partitioned IRModule). |
43 | * |
44 | * - The \p MarkCompilerFunctionsAsExtern pass will update the attributes of global functions |
45 | * with a matching "Compiler" attribute to have just the "Extern" attribute. That will signal |
46 | * the function has been dealt with. However calls to such functions will be left unchanged. |
47 | * Can be used after lowering to cleanup the IRModule. |
48 | * |
49 | * - The \p InlineCompilerFunctions pass can selectively inline global functions with a matching |
50 | * "Compiler" attribute who's name appears in the given set. Obviously it's more sensible to |
51 | * not create that function in the first place, however some external codegen have rules to |
52 | * accept or reject partitionings based on the overall partitioned function body. This pass |
53 | * can be used do the legwork, and will take care to not only inline the outer "Compiler" |
54 | * annotated funcition, but also any "Composite" annotated functions in its body. |
55 | */ |
56 | |
57 | #ifndef TVM_RELAY_TRANSFORMS_COMPILER_FUNCTION_UTILS_H_ |
58 | #define TVM_RELAY_TRANSFORMS_COMPILER_FUNCTION_UTILS_H_ |
59 | |
60 | #include <memory> |
61 | #include <string> |
62 | #include <unordered_map> |
63 | |
64 | #include "tvm/ir/transform.h" |
65 | #include "tvm/relay/function.h" |
66 | |
67 | namespace tvm { |
68 | namespace relay { |
69 | namespace transform { |
70 | |
71 | /*! |
72 | * \brief Abstract class representing a cache of unique global vars keyed by functions. This can |
73 | * be used to ensure structurally equal functions are assigned the same global var object, and |
74 | * thus lowered at most once. |
75 | */ |
76 | class GlobalSymbolCache { |
77 | public: |
78 | virtual ~GlobalSymbolCache(); |
79 | virtual GlobalVar GetGlobalSymbol(const Function& function) = 0; |
80 | }; |
81 | |
82 | /*! |
83 | * \brief A \p GlobalSymbolCache that requires every "Compiler" attributed function to already |
84 | * have a "global_symbol" attribute. |
85 | */ |
86 | class ExistingGlobalSymbolCache : public GlobalSymbolCache { |
87 | public: |
88 | ExistingGlobalSymbolCache() = default; |
89 | |
90 | GlobalVar GetGlobalSymbol(const Function& function) final; |
91 | |
92 | private: |
93 | /*! \brief Maps already seen global symbol names to their corresponding GlobalVar objects. */ |
94 | std::unordered_map<std::string, GlobalVar> global_vars_; |
95 | }; |
96 | |
97 | /*! |
98 | * \brief A pass to outline all let-bound and literal functions in direct call positions which have |
99 | * a "Compiler" attribute. The given \p GlobalSymbolCache is used to determine a unique global |
100 | * symbol for each function, which is also assigned to the "global_symbol" attribute of the new |
101 | * global function. |
102 | * |
103 | * At most one function with the same global symbol is outlined. |
104 | * |
105 | * If \p compiler_filter is non-empty only functions with that as their attribute value are |
106 | * outlined. |
107 | */ |
108 | tvm::transform::Pass OutlineCompilerFunctions(std::shared_ptr<GlobalSymbolCache> cache, |
109 | std::string compiler_filter = "" ); |
110 | |
111 | /*! |
112 | * \brief A pass to outline all let-bound and literal functions in direct call positions which have |
113 | * a "Compiler" attribute. The functions are bound to unique global vars according to their |
114 | * existing "global_symbol" attribute. At most one function with the same global symbol is outlined. |
115 | * |
116 | * If \p compiler_filter is non-empty only functions with that as their attribute value are |
117 | * outlined. |
118 | * |
119 | * This pass may be useful for external codegen using the "RelayToTIR" custom pass mechanism |
120 | * to prepare the IRModule before custom lowering. |
121 | */ |
122 | tvm::transform::Pass OutlineCompilerFunctionsWithExistingGlobalSymbols( |
123 | std::string compiler_filter = "" ); |
124 | |
125 | /*! |
126 | * \brief A pass to mark all global functions which have a "Compiler" attribute matching |
127 | * compiler_filter as 'extern' by replacing all attributes with a single "Extern" attribute. |
128 | * Calls to such functions are not changed. |
129 | * |
130 | * If \p compiler_filter is non-empty only functions with that as their attribute value are |
131 | * outlined. |
132 | * |
133 | * This pass may be useful for external codegen using the "RelayToTIR" custom pass mechanism to |
134 | * cleanup the IRModule after custom lowering. |
135 | */ |
136 | tvm::transform::Pass MarkCompilerFunctionsAsExtern(std::string compiler_filter = "" ); |
137 | |
138 | /*! |
139 | * \brief A pass to inline all global "Compiler" functions which are bound to a global var |
140 | * in \p global_vars. Both the global function and any calls to "Composite" functions it its body |
141 | * are inlined. |
142 | * |
143 | * This pass may be useful for external codegen which needs to undo partitioning based on |
144 | * properties of the entire partition. |
145 | */ |
146 | tvm::transform::Pass InlineCompilerFunctionsBoundTo(Array<GlobalVar> global_vars); |
147 | |
148 | } // namespace transform |
149 | } // namespace relay |
150 | } // namespace tvm |
151 | |
152 | #endif // TVM_RELAY_TRANSFORMS_COMPILER_FUNCTION_UTILS_H_ |
153 | |