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/function.h |
22 | * \brief Function nodes. |
23 | */ |
24 | #ifndef TVM_IR_FUNCTION_H_ |
25 | #define TVM_IR_FUNCTION_H_ |
26 | |
27 | #include <tvm/ir/attrs.h> |
28 | #include <tvm/ir/expr.h> |
29 | #include <tvm/runtime/container/array.h> |
30 | #include <tvm/runtime/container/map.h> |
31 | #include <tvm/runtime/container/string.h> |
32 | |
33 | #include <string> |
34 | #include <type_traits> |
35 | |
36 | namespace tvm { |
37 | |
38 | /*! |
39 | * \brief Possible Calling conventions. |
40 | * |
41 | * NOTE: The calling convention also implies |
42 | * the way we implement the function during lowering. |
43 | */ |
44 | enum class CallingConv : int { |
45 | /*! |
46 | * \brief Default calling convention. |
47 | * |
48 | * - Uses the native calling convention of the target. |
49 | * - Implementation: specified by the native target. |
50 | */ |
51 | kDefault = 0, |
52 | /*! |
53 | * \brief PackedFunc that exposes a CPackedFunc signature. |
54 | * |
55 | * - Calling by PackedFunc calling convention. |
56 | * - Implementation: Expose a function with the CPackedFunc signature. |
57 | */ |
58 | kCPackedFunc = 1, |
59 | /*! |
60 | * \brief Device kernel launch |
61 | * |
62 | * - Call by PackedFunc calling convention. |
63 | * - Implementation: defined by device runtime(e.g. runtime/cuda) |
64 | */ |
65 | kDeviceKernelLaunch = 2, |
66 | }; |
67 | |
68 | /*! |
69 | * \brief Base node of all functions. |
70 | * |
71 | * We support several variants of functions throughout the stack. |
72 | * All of the functions share the same type system(via checked_type) |
73 | * to support cross variant calls. |
74 | * |
75 | * \sa BaseFunc |
76 | */ |
77 | class BaseFuncNode : public RelayExprNode { |
78 | public: |
79 | /*! \brief Additional attributes storing the meta-data */ |
80 | DictAttrs attrs; |
81 | |
82 | /*! |
83 | * \brief Get a function attribute. |
84 | * |
85 | * \param attr_key The attribute key. |
86 | * \param default_value The default value if the key does not exist, defaults to nullptr. |
87 | * |
88 | * \return The result |
89 | * |
90 | * \tparam TOBjectRef the expected object type. |
91 | * \throw Error if the key exists but the value does not match TObjectRef |
92 | * |
93 | * \code |
94 | * |
95 | * void GetAttrExample(const BaseFunc& f) { |
96 | * auto value = f->GetAttr<Integer>("AttrKey", 0); |
97 | * } |
98 | * |
99 | * \endcode |
100 | */ |
101 | template <typename TObjectRef> |
102 | Optional<TObjectRef> GetAttr( |
103 | const std::string& attr_key, |
104 | Optional<TObjectRef> default_value = Optional<TObjectRef>(nullptr)) const { |
105 | return attrs.GetAttr(attr_key, default_value); |
106 | } |
107 | // variant that uses TObjectRef to enable implicit conversion to default value. |
108 | template <typename TObjectRef> |
109 | Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const { |
110 | return GetAttr<TObjectRef>(attr_key, Optional<TObjectRef>(default_value)); |
111 | } |
112 | |
113 | /*! |
114 | * \brief Check whether the function has an non-zero integer attr. |
115 | * |
116 | * This function can be used to check whether an optional |
117 | * attribute mark(e.g. inline) exists. |
118 | * |
119 | * \param attr_key The key to the attribute. |
120 | * \return The check result. |
121 | * |
122 | * \code |
123 | * |
124 | * void HasNonzeroAttrExample(const BaseFunc& f) { |
125 | * if (f->HasNonzeroAttr(attr::kInline)) { |
126 | * // inline the function. |
127 | * } |
128 | * } |
129 | * |
130 | * \endcode |
131 | */ |
132 | bool HasNonzeroAttr(const std::string& attr_key) const { return attrs.HasNonzeroAttr(attr_key); } |
133 | |
134 | static constexpr const char* _type_key = "BaseFunc" ; |
135 | static constexpr const uint32_t _type_child_slots = 2; |
136 | TVM_DECLARE_BASE_OBJECT_INFO(BaseFuncNode, RelayExprNode); |
137 | }; |
138 | |
139 | /*! |
140 | * \brief Managed reference to BaseFuncNode. |
141 | * \sa BaseFuncNode |
142 | */ |
143 | class BaseFunc : public RelayExpr { |
144 | public: |
145 | TVM_DEFINE_OBJECT_REF_METHODS(BaseFunc, RelayExpr, BaseFuncNode); |
146 | }; |
147 | |
148 | /*! |
149 | * \brief Generic attribute names that can be attached to any function. |
150 | * |
151 | * \sa tvm::tir::attr, tvm::relay::attr |
152 | */ |
153 | namespace attr { |
154 | /*! |
155 | * \brief Indicates the special calling convention. |
156 | * |
157 | * Type: Integer |
158 | * |
159 | * \sa tvm::CallingConv |
160 | */ |
161 | constexpr const char* kCallingConv = "calling_conv" ; |
162 | |
163 | /*! |
164 | * \brief Compilation target of the function. |
165 | * |
166 | * Type: Target |
167 | * |
168 | * \sa tvm::Target |
169 | */ |
170 | constexpr const char* kTarget = "target" ; |
171 | |
172 | /*! |
173 | * \brief Global linker symbol of the function in generated code. |
174 | * |
175 | * This option forces the code generator to name the |
176 | * function with the given. |
177 | * |
178 | * For example, we could set a global_symbol of a function |
179 | * early to make sure that we can always refer to it by |
180 | * the symbol name in the generated DLL. |
181 | * |
182 | * We should not set the attribute for local functions, |
183 | * so that the compiler can freely rename them. |
184 | * |
185 | * A unique global symbol will be automatically assigned |
186 | * to each function in the module before the target code |
187 | * generation phase. |
188 | * |
189 | * Type: String |
190 | */ |
191 | constexpr const char* kGlobalSymbol = "global_symbol" ; |
192 | |
193 | } // namespace attr |
194 | } // namespace tvm |
195 | #endif // TVM_IR_FUNCTION_H_ |
196 | |