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
67namespace tvm {
68namespace relay {
69namespace 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 */
76class 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 */
86class 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 */
108tvm::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 */
122tvm::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 */
136tvm::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 */
146tvm::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