tvm
attrs.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  */
28 #ifndef TVM_IR_ATTRS_H_
29 #define TVM_IR_ATTRS_H_
30 
31 #include <dmlc/common.h>
32 #include <tvm/ffi/container/map.h>
33 #include <tvm/ffi/function.h>
34 #include <tvm/ffi/reflection/accessor.h>
35 #include <tvm/ffi/reflection/registry.h>
36 #include <tvm/ir/expr.h>
39 
40 #include <functional>
41 #include <string>
42 #include <type_traits>
43 #include <unordered_map>
44 #include <utility>
45 #include <vector>
46 
47 namespace tvm {
48 
54 template <typename TObjectRef>
55 inline TObjectRef NullValue() {
56  static_assert(TObjectRef::_type_is_nullable, "Can only get NullValue for nullable types");
57  return TObjectRef(ObjectPtr<Object>(nullptr));
58 }
59 
60 template <>
62  return DataType(DataType::kHandle, 0, 0);
63 }
64 
68 class AttrFieldInfoNode : public Object {
69  public:
71  String name;
73  String type_info;
75  String description;
76 
77  static void RegisterReflection() {
78  namespace rfl = ffi::reflection;
79  rfl::ObjectDef<AttrFieldInfoNode>()
80  .def_ro("name", &AttrFieldInfoNode::name)
81  .def_ro("type_info", &AttrFieldInfoNode::type_info)
82  .def_ro("description", &AttrFieldInfoNode::description);
83  }
84 
85  static constexpr const char* _type_key = "ir.AttrFieldInfo";
86  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
87 
89 };
90 
92 class AttrFieldInfo : public ObjectRef {
93  public:
95 };
96 
103 class BaseAttrsNode : public Object {
104  public:
106  virtual ~BaseAttrsNode() {}
112  template <typename... Args>
113  inline void InitBySeq(Args&&... args);
121  TVM_DLL virtual void InitByPackedArgs(const ffi::PackedArgs& kwargs,
122  bool allow_unknown = false) = 0;
123 
124  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
125 
126  static constexpr const char* _type_key = "ir.Attrs";
128 };
129 
134 class Attrs : public ObjectRef {
135  public:
137 };
138 
145 class DictAttrsNode : public BaseAttrsNode {
146  public:
148  Map<String, ffi::Any> dict;
149 
150  static void RegisterReflection() {
151  namespace rfl = ffi::reflection;
152  rfl::ObjectDef<DictAttrsNode>().def_ro("__dict__", &DictAttrsNode::dict);
153  }
154 
155  void InitByPackedArgs(const ffi::PackedArgs& args, bool allow_unknown) final;
156 
157  // type info
158  static constexpr const char* _type_key = "ir.DictAttrs";
160 };
161 
166 class DictAttrs : public Attrs {
167  public:
172  TVM_DLL explicit DictAttrs(Map<String, Any> dict = {});
173 
174  // Utils for accessing attributes
175  // This needs to be on DictAttrs, not DictAttrsNode because we return the default
176  // value if DictAttrsNode is not defined.
196  template <typename TObjectRef>
197  Optional<TObjectRef> GetAttr(
198  const std::string& attr_key,
199  Optional<TObjectRef> default_value = Optional<TObjectRef>(std::nullopt)) const {
200  if (!defined()) return default_value;
201  const DictAttrsNode* node = this->as<DictAttrsNode>();
202  auto it = node->dict.find(attr_key);
203  if (it != node->dict.end()) {
204  return (*it).second.cast<TObjectRef>();
205  } else {
206  return default_value;
207  }
208  }
209  // variant that uses TObjectRef to enable implicit conversion to default value.
210  template <typename TObjectRef>
211  Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
212  return GetAttr<TObjectRef>(attr_key, Optional<TObjectRef>(default_value));
213  }
233  bool HasNonzeroAttr(const std::string& attr_key) const {
234  return GetAttr<Integer>(attr_key, 0).value_or(0).IntValue() != 0;
235  }
236 
239 };
240 
251 DictAttrs WithAttrs(DictAttrs attrs, Map<String, Any> new_attrs);
252 
264 DictAttrs WithAttr(DictAttrs attrs, String key, Any value);
265 
266 inline DictAttrs WithAttr(DictAttrs attrs, const std::string& key, Any value) {
267  return WithAttr(std::move(attrs), String(key), std::move(value));
268 }
269 
279 DictAttrs WithoutAttr(DictAttrs attrs, const std::string& key);
280 
308 template <typename TFunc>
309 inline TFunc WithAttr(TFunc input, const std::string& attr_key, Any attr_value) {
310  using TNode = typename TFunc::ContainerType;
311  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
312  TNode* node = input.CopyOnWrite();
313  node->attrs = WithAttr(std::move(node->attrs), attr_key, attr_value);
314  return input;
315 }
316 
327 template <typename TFunc>
328 inline TFunc WithAttrs(TFunc input, Map<String, Any> attrs) {
329  using TNode = typename TFunc::ContainerType;
330  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
331  TNode* node = input.CopyOnWrite();
332 
333  node->attrs = WithAttrs(std::move(node->attrs), attrs);
334 
335  return input;
336 }
337 
364 template <typename TFunc>
365 inline TFunc WithoutAttr(TFunc input, const std::string& attr_key) {
366  using TNode = typename TFunc::ContainerType;
367  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
368 
369  TNode* node = input.CopyOnWrite();
370  node->attrs = WithoutAttr(std::move(node->attrs), attr_key);
371 
372  return input;
373 }
374 
383 template <typename DerivedType>
385  public:
386  void InitByPackedArgs(const ffi::PackedArgs& args, bool allow_unknown) final {
387  LOG(FATAL) << "`" << DerivedType::_type_key << "` uses new reflection mechanism for init";
388  }
389 
390  private:
391  DerivedType* self() const {
392  return const_cast<DerivedType*>(static_cast<const DerivedType*>(this));
393  }
394 };
395 
401 template <typename TAttrs>
402 inline TAttrs AttrsWithDefaultValues() {
403  static_assert(std::is_base_of_v<Attrs, TAttrs>, "Can only take attr nodes");
404  using ContainerType = typename TAttrs::ContainerType;
405  if constexpr (std::is_base_of_v<AttrsNodeReflAdapter<ContainerType>, ContainerType>) {
406  static auto finit_object = ffi::Function::GetGlobalRequired("ffi.MakeObjectFromPackedArgs");
407  AnyView packed_args[1];
408  packed_args[0] = ContainerType::RuntimeTypeIndex();
409  ffi::Any rv;
410  finit_object.CallPacked(ffi::PackedArgs(packed_args, 1), &rv);
411  return rv.cast<TAttrs>();
412  } else {
413  auto n = make_object<ContainerType>();
414  n->InitByPackedArgs(ffi::PackedArgs(nullptr, 0), false);
415  return TAttrs(n);
416  }
417 }
418 
419 } // namespace tvm
420 #endif // TVM_IR_ATTRS_H_
Information about attribute fields in string representations.
Definition: attrs.h:68
static void RegisterReflection()
Definition: attrs.h:77
String description
detailed description of the type
Definition: attrs.h:75
TVM_DECLARE_FINAL_OBJECT_INFO(AttrFieldInfoNode, Object)
String type_info
type docstring information in str.
Definition: attrs.h:73
String name
name of the field
Definition: attrs.h:71
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: attrs.h:86
static constexpr const char * _type_key
Definition: attrs.h:85
AttrFieldInfo.
Definition: attrs.h:92
TVM_DEFINE_OBJECT_REF_METHODS(AttrFieldInfo, ObjectRef, AttrFieldInfoNode)
Adapter for AttrsNode with the new reflection API.
Definition: attrs.h:384
void InitByPackedArgs(const ffi::PackedArgs &args, bool allow_unknown) final
Initialize the attributes by arguments.
Definition: attrs.h:386
Managed reference to BaseAttrsNode.
Definition: attrs.h:134
TVM_DEFINE_OBJECT_REF_METHODS(Attrs, ObjectRef, BaseAttrsNode)
Base class of all attribute class.
Definition: attrs.h:103
virtual ~BaseAttrsNode()
virtual destructor
Definition: attrs.h:106
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: attrs.h:124
static constexpr const char * _type_key
Definition: attrs.h:126
TVM_DECLARE_BASE_OBJECT_INFO(BaseAttrsNode, Object)
void InitBySeq(Args &&... args)
Initialize the attributes by sequence of arguments.
virtual void InitByPackedArgs(const ffi::PackedArgs &kwargs, bool allow_unknown=false)=0
Initialize the attributes by arguments.
Specialized attribute type that is backed by a map. The DictAttrsNode implements the Attrs behavior,...
Definition: attrs.h:145
TVM_DECLARE_FINAL_OBJECT_INFO(DictAttrsNode, BaseAttrsNode)
static void RegisterReflection()
Definition: attrs.h:150
Map< String, ffi::Any > dict
internal attrs map
Definition: attrs.h:148
void InitByPackedArgs(const ffi::PackedArgs &args, bool allow_unknown) final
Initialize the attributes by arguments.
static constexpr const char * _type_key
Definition: attrs.h:158
Managed reference to DictAttrsNode.
Definition: attrs.h:166
DictAttrs(Map< String, Any > dict={})
Consruct a Attrs backed by DictAttrsNode.
bool HasNonzeroAttr(const std::string &attr_key) const
Check whether the function has an non-zero integer attr.
Definition: attrs.h:233
Optional< TObjectRef > GetAttr(const std::string &attr_key, Optional< TObjectRef > default_value=Optional< TObjectRef >(std::nullopt)) const
Get a function attribute.
Definition: attrs.h:197
TVM_DEFINE_OBJECT_REF_COW_METHOD(DictAttrsNode)
Optional< TObjectRef > GetAttr(const std::string &attr_key, TObjectRef default_value) const
Definition: attrs.h:211
TVM_DEFINE_OBJECT_REF_METHODS_WITHOUT_DEFAULT_CONSTRUCTOR(DictAttrs, Attrs, DictAttrsNode)
Runtime primitive data type.
Definition: data_type.h:47
@ kHandle
Definition: data_type.h:61
Base expr nodes in TVM.
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
DictAttrs WithoutAttr(DictAttrs attrs, const std::string &key)
Copy the DictAttrs, but without a specific attribute.
DataType NullValue< DataType >()
Definition: attrs.h:61
TAttrs AttrsWithDefaultValues()
Create an Attr object with all default values.
Definition: attrs.h:402
DictAttrs WithAttr(DictAttrs attrs, String key, Any value)
Copy the DictAttrs, but overrides a single attribute.
runtime::DataType DataType
Definition: data_type.h:458
DictAttrs WithAttrs(DictAttrs attrs, Map< String, Any > new_attrs)
Copy the DictAttrs, but overrides attributes with the entries from attrs.
TObjectRef NullValue()
Create a NodeRef type that represents null.
Definition: attrs.h:55
Structural equality comparison.