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
36namespace 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 */
44enum 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 */
77class 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 */
143class 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 */
153namespace attr {
154/*!
155 * \brief Indicates the special calling convention.
156 *
157 * Type: Integer
158 *
159 * \sa tvm::CallingConv
160 */
161constexpr const char* kCallingConv = "calling_conv";
162
163/*!
164 * \brief Compilation target of the function.
165 *
166 * Type: Target
167 *
168 * \sa tvm::Target
169 */
170constexpr 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 */
191constexpr const char* kGlobalSymbol = "global_symbol";
192
193} // namespace attr
194} // namespace tvm
195#endif // TVM_IR_FUNCTION_H_
196