tvm
op.h
Go to the documentation of this file.
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 
25 #ifndef TVM_IR_OP_H_
26 #define TVM_IR_OP_H_
27 
28 #include <tvm/ffi/function.h>
29 #include <tvm/ffi/reflection/registry.h>
30 #include <tvm/ir/attrs.h>
31 #include <tvm/ir/env_func.h>
32 #include <tvm/ir/expr.h>
33 #include <tvm/ir/type.h>
35 #include <tvm/runtime/logging.h>
36 
37 #include <string>
38 #include <utility>
39 #include <vector>
40 
41 namespace tvm {
42 
43 // forward declare name.
44 template <typename>
45 class OpAttrMap;
46 
47 // TODO(tvm-team): migrate low-level intrinsics to use Op
59 class OpNode : public RelaxExprNode {
60  public:
62  ffi::String name;
64  mutable FuncType op_type;
69  ffi::String description;
70  /* \brief Information of input arguments to the operator */
71  ffi::Array<AttrFieldInfo> arguments;
76  ffi::String attrs_type_key;
81  uint32_t attrs_type_index{0};
86  int32_t num_inputs = -1;
92  int32_t support_level = 10;
93 
94  static void RegisterReflection() {
95  namespace refl = tvm::ffi::reflection;
96  refl::ObjectDef<OpNode>()
97  .def_ro("name", &OpNode::name)
98  .def_ro("op_type", &OpNode::op_type, refl::AttachFieldFlag::SEqHashIgnore())
99  .def_ro("description", &OpNode::description, refl::AttachFieldFlag::SEqHashIgnore())
100  .def_ro("arguments", &OpNode::arguments, refl::AttachFieldFlag::SEqHashIgnore())
101  .def_ro("attrs_type_key", &OpNode::attrs_type_key, refl::AttachFieldFlag::SEqHashIgnore())
102  .def_ro("num_inputs", &OpNode::num_inputs, refl::AttachFieldFlag::SEqHashIgnore())
103  .def_ro("support_level", &OpNode::support_level, refl::AttachFieldFlag::SEqHashIgnore());
104  }
105 
106  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance;
108 
109  private:
111  uint32_t AttrRegistryIndex() const { return index_; }
113  std::string AttrRegistryName() const { return name; }
114 
115  // friend class
116  template <typename>
118  template <typename, typename>
119  friend class AttrRegistry;
120  friend class OpRegEntry;
121 
122  // Program internal unique index of operator.
123  // Used to help index the program.
124  uint32_t index_{0};
125 };
126 
131 class Op : public RelaxExpr {
132  public:
140  template <typename ValueType>
141  inline static OpAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
147  TVM_DLL static bool HasAttrMap(const ffi::String& attr_name);
154  TVM_DLL static const Op& Get(const ffi::String& op_name);
155 
157 
158  private:
164  TVM_DLL static const AttrRegistryMapContainerMap<Op>& GetAttrMapContainer(const ffi::String& key);
165 };
166 
171 class OpRegEntry {
172  public:
174  const Op& op() const { return op_; }
181  inline OpRegEntry& describe(const std::string& descr); // NOLINT(*)
189  inline OpRegEntry& add_argument(const std::string& name, const std::string& type,
190  const std::string& description);
196  template <typename AttrsType>
197  inline OpRegEntry& set_attrs_type();
203  inline OpRegEntry& set_attrs_type_key(const ffi::String& key);
209  inline OpRegEntry& set_num_inputs(int32_t n); // NOLINT(*)
215  inline OpRegEntry& set_support_level(int32_t level); // NOLINT(*)
229  template <typename ValueType>
230  inline OpRegEntry& set_attr(const std::string& attr_name, // NOLINT(*)
231  const ValueType& value, int plevel = 10);
232 
237  inline void reset_attr(const std::string& attr_name);
238 
239  // set the name of the op to be the same as registry
240  inline OpRegEntry& set_name() { // NOLINT(*)
241  if (get()->name.length() == 0) {
242  get()->name = name;
243  }
244  return *this;
245  }
251  TVM_DLL static OpRegEntry& RegisterOrGet(const ffi::String& name);
252 
253  private:
254  template <typename, typename>
255  friend class AttrRegistry;
256  // the name
257  std::string name;
259  Op op_;
260  // private constructor
261  TVM_DLL OpRegEntry(uint32_t reg_index);
262  // return internal pointer to op.
263  inline OpNode* get();
264  // update the attribute OpAttrMap
265  TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
266 };
267 
272 template <typename ValueType>
273 class OpAttrMap : public AttrRegistryMap<Op, ValueType> {
274  public:
282  inline ValueType get(const RelaxExpr& expr, ValueType def_value) const;
283 
285  using TParent::count;
286  using TParent::get;
287  using TParent::operator[];
288 
289  private:
290  friend class Op;
291  // constructor
292  explicit OpAttrMap(const AttrRegistryMapContainerMap<Op>& map) : TParent(map) {}
293 };
294 
295 // internal macros to make
296 #define TVM_OP_REGISTER_VAR_DEF static DMLC_ATTRIBUTE_UNUSED ::tvm::OpRegEntry& __make_##Op
297 
313 #define TVM_REGISTER_OP(OpName) \
314  TVM_STR_CONCAT(TVM_OP_REGISTER_VAR_DEF, __COUNTER__) = \
315  ::tvm::OpRegEntry::RegisterOrGet(OpName).set_name()
316 
317 // implementations
318 
319 template <typename ValueType>
320 inline OpAttrMap<ValueType> Op::GetAttrMap(const ffi::String& key) {
321  return OpAttrMap<ValueType>(Op::GetAttrMapContainer(key));
322 }
323 
324 inline OpNode* OpRegEntry::get() { return const_cast<OpNode*>(op_.operator->()); }
325 
326 inline OpRegEntry& OpRegEntry::describe(const std::string& descr) { // NOLINT(*)
327  get()->description = descr;
328  return *this;
329 }
330 
331 inline OpRegEntry& OpRegEntry::add_argument(const std::string& name, const std::string& type,
332  const std::string& description) {
333  auto n = ffi::make_object<AttrFieldInfoNode>();
334  n->name = name;
335  n->type_info = type;
336  n->description = description;
337  get()->arguments.push_back(AttrFieldInfo(n));
338  return *this;
339 }
340 
341 inline OpRegEntry& OpRegEntry::set_num_inputs(int32_t n) { // NOLINT(*)
342  get()->num_inputs = n;
343  return *this;
344 }
345 
346 template <typename AttrsType>
347 inline OpRegEntry& OpRegEntry::set_attrs_type() { // NOLINT(*)
348  get()->attrs_type_key = AttrsType::_type_key;
349  get()->attrs_type_index = AttrsType::RuntimeTypeIndex();
350  return *this;
351 }
352 
353 inline OpRegEntry& OpRegEntry::set_attrs_type_key(const ffi::String& key) { // NOLINT(*)
354  get()->attrs_type_key = key;
355  get()->attrs_type_index = tvm::ffi::TypeKeyToIndex(key.c_str());
356  return *this;
357 }
358 
359 inline OpRegEntry& OpRegEntry::set_support_level(int32_t n) { // NOLINT(*)
360  get()->support_level = n;
361  return *this;
362 }
363 
364 template <typename ValueType>
365 inline OpRegEntry& OpRegEntry::set_attr( // NOLINT(*)
366  const std::string& attr_name, const ValueType& value, int plevel) {
367  ICHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
368  UpdateAttr(attr_name, Any(value), plevel);
369  return *this;
370 }
371 
372 // member functions of OpAttrMap
373 
374 template <typename ValueType>
375 inline ValueType OpAttrMap<ValueType>::get(const RelaxExpr& expr, ValueType def_value) const {
376  ICHECK(expr.defined());
377  if (const OpNode* op = expr.as<OpNode>()) {
378  return this->map_.get(ffi::GetRef<Op>(op), def_value);
379  } else {
380  return def_value;
381  }
382 }
383 
384 } // namespace tvm
385 #endif // TVM_IR_OP_H_
Attribute map used in registry.
Helpers for attribute objects.
AttrFieldInfo.
Definition: attrs.h:91
Generic attribute map.
Definition: attr_registry_map.h:38
ffi::Map<Key, ValueType> used to store meta-data.
Definition: attr_registry_map.h:105
ValueType get(const Op &key, ValueType def_value) const
get the corresponding value element at key with default value.
Definition: attr_registry_map.h:136
int count(const Op &key) const
Check if the map has op as key.
Definition: attr_registry_map.h:117
Definition: instruction.h:30
Managed reference to FuncTypeNode.
Definition: type.h:274
ffi::Map<Op,ValueType> used to store meta-information about Op.
Definition: op.h:273
ValueType get(const RelaxExpr &expr, ValueType def_value) const
get the corresponding value element at op with default value.
Definition: op.h:375
Primitive Op(builtin intrinsics)
Definition: op.h:59
ffi::String attrs_type_key
The type key of the attribute field This can be empty, in which case it defaults to anything.
Definition: op.h:76
static void RegisterReflection()
Definition: op.h:94
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Op", OpNode, RelaxExprNode)
ffi::String description
detailed description of the operator This can be used to generate docstring automatically for the ope...
Definition: op.h:69
uint32_t attrs_type_index
attribute type index, this field varies in each run and is not exposed to frontend.
Definition: op.h:81
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: op.h:106
int32_t support_level
support level of the operator, The lower the more priority it contains. This is in analogies to BLAS ...
Definition: op.h:92
ffi::String name
name of the operator
Definition: op.h:62
int32_t num_inputs
number of input arguments to the operator, -1 means it is variable length
Definition: op.h:86
ffi::Array< AttrFieldInfo > arguments
Definition: op.h:71
FuncType op_type
the type of the operator
Definition: op.h:64
Helper structure to register operators.
Definition: op.h:171
OpRegEntry & describe(const std::string &descr)
setter function during registration Set the description of operator
Definition: op.h:326
static OpRegEntry & RegisterOrGet(const ffi::String &name)
Register or get a new entry.
OpRegEntry & set_name()
Definition: op.h:240
void reset_attr(const std::string &attr_name)
Resets an attr of the registry.
OpRegEntry & set_attrs_type_key(const ffi::String &key)
Set the attrs type key and index to be AttrsType.
Definition: op.h:353
OpRegEntry & add_argument(const std::string &name, const std::string &type, const std::string &description)
Add argument information to the function.
Definition: op.h:331
OpRegEntry & set_attrs_type()
Set the attrs type key and index to be AttrsType.
Definition: op.h:347
OpRegEntry & set_support_level(int32_t level)
Set the support level of op.
Definition: op.h:359
const Op & op() const
Definition: op.h:174
OpRegEntry & set_attr(const std::string &attr_name, const ValueType &value, int plevel=10)
Register additional attributes to operator.
Definition: op.h:365
OpRegEntry & set_num_inputs(int32_t n)
Set the num_inputs.
Definition: op.h:341
Managed reference class to OpNode.
Definition: op.h:131
static const Op & Get(const ffi::String &op_name)
Get an Op for a given operator name. Will raise an error if the op has not been registered.
static OpAttrMap< ValueType > GetAttrMap(const ffi::String &attr_name)
Get additional registered attribute about operators. If nothing has been registered,...
Definition: op.h:320
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Op, RelaxExpr, OpNode)
static bool HasAttrMap(const ffi::String &attr_name)
Checks if an attr map is present in the registry.
Base node of all non-primitive expressions.
Definition: expr.h:416
Managed reference to RelaxExprNode.
Definition: expr.h:439
Serializable global function used in IR.
Base expr nodes in TVM.
IR/AST nodes for the unified type system in TVM.
Definition: repr_printer.h:91
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37