-
Notifications
You must be signed in to change notification settings - Fork 3.4k
/
registry.h
302 lines (287 loc) · 9.79 KB
/
registry.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*!
* \file tvm/runtime/registry.h
* \brief This file defines the TVM global function registry.
*
* The registered functions will be made available to front-end
* as well as backend users.
*
* The registry stores type-erased functions.
* Each registered function is automatically exposed
* to front-end language(e.g. python).
*
* Front-end can also pass callbacks as PackedFunc, or register
* then into the same global registry in C++.
* The goal is to mix the front-end language and the TVM back-end.
*
* \code
* // register the function as MyAPIFuncName
* TVM_REGISTER_GLOBAL(MyAPIFuncName)
* .set_body([](TVMArgs args, TVMRetValue* rv) {
* // my code.
* });
* \endcode
*/
#ifndef TVM_RUNTIME_REGISTRY_H_
#define TVM_RUNTIME_REGISTRY_H_
#include <tvm/runtime/packed_func.h>
#include <string>
#include <utility>
#include <vector>
namespace tvm {
namespace runtime {
/*! \brief Registry for global function */
class Registry {
public:
/*!
* \brief set the body of the function to be f
* \param f The body of the function.
*/
TVM_DLL Registry& set_body(PackedFunc f); // NOLINT(*)
/*!
* \brief set the body of the function to be f
* \param f The body of the function.
*/
Registry& set_body(PackedFunc::FType f) { // NOLINT(*)
return set_body(PackedFunc(f));
}
/*!
* \brief set the body of the function to the given function.
* Note that this will ignore default arg values and always require all arguments to be
* provided.
*
* \code
*
* int multiply(int x, int y) {
* return x * y;
* }
*
* TVM_REGISTER_GLOBAL("multiply")
* .set_body_typed(multiply); // will have type int(int, int)
*
* // will have type int(int, int)
* TVM_REGISTER_GLOBAL("sub")
* .set_body_typed([](int a, int b) -> int { return a - b; });
*
* \endcode
*
* \param f The function to forward to.
* \tparam FLambda The signature of the function.
*/
template <typename FLambda>
Registry& set_body_typed(FLambda f) {
using FType = typename detail::function_signature<FLambda>::FType;
return set_body(TypedPackedFunc<FType>(std::move(f)).packed());
}
/*!
* \brief set the body of the function to be the passed method pointer.
* Note that this will ignore default arg values and always require all arguments to be
* provided.
*
* \code
*
* // node subclass:
* struct Example {
* int doThing(int x);
* }
* TVM_REGISTER_GLOBAL("Example_doThing")
* .set_body_method(&Example::doThing); // will have type int(Example, int)
*
* \endcode
*
* \param f the method pointer to forward to.
* \tparam T the type containing the method (inferred).
* \tparam R the return type of the function (inferred).
* \tparam Args the argument types of the function (inferred).
*/
template <typename T, typename R, typename... Args>
Registry& set_body_method(R (T::*f)(Args...)) {
auto fwrap = [f](T target, Args... params) -> R {
// call method pointer
return (target.*f)(params...);
};
return set_body(TypedPackedFunc<R(T, Args...)>(fwrap));
}
/*!
* \brief set the body of the function to be the passed method pointer.
* Note that this will ignore default arg values and always require all arguments to be
* provided.
*
* \code
*
* // node subclass:
* struct Example {
* int doThing(int x);
* }
* TVM_REGISTER_GLOBAL("Example_doThing")
* .set_body_method(&Example::doThing); // will have type int(Example, int)
*
* \endcode
*
* \param f the method pointer to forward to.
* \tparam T the type containing the method (inferred).
* \tparam R the return type of the function (inferred).
* \tparam Args the argument types of the function (inferred).
*/
template <typename T, typename R, typename... Args>
Registry& set_body_method(R (T::*f)(Args...) const) {
auto fwrap = [f](const T target, Args... params) -> R {
// call method pointer
return (target.*f)(params...);
};
return set_body(TypedPackedFunc<R(const T, Args...)>(fwrap));
}
/*!
* \brief set the body of the function to be the passed method pointer.
* Used when calling a method on a Node subclass through a ObjectRef subclass.
* Note that this will ignore default arg values and always require all arguments to be
* provided.
*
* \code
*
* // node subclass:
* struct ExampleNode: BaseNode {
* int doThing(int x);
* }
*
* // noderef subclass
* struct Example;
*
* TVM_REGISTER_GLOBAL("Example_doThing")
* .set_body_method<Example>(&ExampleNode::doThing); // will have type int(Example, int)
*
* // note that just doing:
* // .set_body_method(&ExampleNode::doThing);
* // wouldn't work, because ExampleNode can't be taken from a TVMArgValue.
*
* \endcode
*
* \param f the method pointer to forward to.
* \tparam TObjectRef the node reference type to call the method on
* \tparam TNode the node type containing the method (inferred).
* \tparam R the return type of the function (inferred).
* \tparam Args the argument types of the function (inferred).
*/
template <typename TObjectRef, typename TNode, typename R, typename... Args,
typename = typename std::enable_if<std::is_base_of<ObjectRef, TObjectRef>::value>::type>
Registry& set_body_method(R (TNode::*f)(Args...)) {
auto fwrap = [f](TObjectRef ref, Args... params) {
TNode* target = ref.operator->();
// call method pointer
return (target->*f)(params...);
};
return set_body(TypedPackedFunc<R(TObjectRef, Args...)>(fwrap));
}
/*!
* \brief set the body of the function to be the passed method pointer.
* Used when calling a method on a Node subclass through a ObjectRef subclass.
* Note that this will ignore default arg values and always require all arguments to be
* provided.
*
* \code
*
* // node subclass:
* struct ExampleNode: BaseNode {
* int doThing(int x);
* }
*
* // noderef subclass
* struct Example;
*
* TVM_REGISTER_GLOBAL("Example_doThing")
* .set_body_method<Example>(&ExampleNode::doThing); // will have type int(Example, int)
*
* // note that just doing:
* // .set_body_method(&ExampleNode::doThing);
* // wouldn't work, because ExampleNode can't be taken from a TVMArgValue.
*
* \endcode
*
* \param f the method pointer to forward to.
* \tparam TObjectRef the node reference type to call the method on
* \tparam TNode the node type containing the method (inferred).
* \tparam R the return type of the function (inferred).
* \tparam Args the argument types of the function (inferred).
*/
template <typename TObjectRef, typename TNode, typename R, typename... Args,
typename = typename std::enable_if<std::is_base_of<ObjectRef, TObjectRef>::value>::type>
Registry& set_body_method(R (TNode::*f)(Args...) const) {
auto fwrap = [f](TObjectRef ref, Args... params) {
const TNode* target = ref.operator->();
// call method pointer
return (target->*f)(params...);
};
return set_body(TypedPackedFunc<R(TObjectRef, Args...)>(fwrap));
}
/*!
* \brief Register a function with given name
* \param name The name of the function.
* \param override Whether allow oveeride existing function.
* \return Reference to theregistry.
*/
TVM_DLL static Registry& Register(const std::string& name, bool override = false); // NOLINT(*)
/*!
* \brief Erase global function from registry, if exist.
* \param name The name of the function.
* \return Whether function exist.
*/
TVM_DLL static bool Remove(const std::string& name);
/*!
* \brief Get the global function by name.
* \param name The name of the function.
* \return pointer to the registered function,
* nullptr if it does not exist.
*/
TVM_DLL static const PackedFunc* Get(const std::string& name); // NOLINT(*)
/*!
* \brief Get the names of currently registered global function.
* \return The names
*/
TVM_DLL static std::vector<std::string> ListNames();
// Internal class.
struct Manager;
protected:
/*! \brief name of the function */
std::string name_;
/*! \brief internal packed function */
PackedFunc func_;
friend struct Manager;
};
#define TVM_FUNC_REG_VAR_DEF static TVM_ATTRIBUTE_UNUSED ::tvm::runtime::Registry& __mk_##TVM
/*!
* \brief Register a function globally.
* \code
* TVM_REGISTER_GLOBAL("MyPrint")
* .set_body([](TVMArgs args, TVMRetValue* rv) {
* });
* \endcode
*/
#define TVM_REGISTER_GLOBAL(OpName) \
TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) = ::tvm::runtime::Registry::Register(OpName)
#define TVM_STRINGIZE_DETAIL(x) #x
#define TVM_STRINGIZE(x) TVM_STRINGIZE_DETAIL(x)
#define TVM_DESCRIBE(...) describe(__VA_ARGS__ "\n\nFrom:" __FILE__ ":" TVM_STRINGIZE(__LINE__))
/*!
* \brief Macro to include current line as string
*/
#define TVM_ADD_FILELINE "\n\nDefined in " __FILE__ ":L" TVM_STRINGIZE(__LINE__)
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_REGISTRY_H_