attrs.h

Go to the documentation of this file.

/*
 * 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.
 */
#ifndef TVM_IR_ATTRS_H_
#define TVM_IR_ATTRS_H_
 
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
 
#include <functional>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>
 
namespace tvm {
 
template <typename TObjectRef>
inline TObjectRef NullValue() {
  static_assert(TObjectRef::_type_is_nullable, "Can only get NullValue for nullable types");
  return TObjectRef(ffi::ObjectPtr<typename TObjectRef::ContainerType>(nullptr));
}
 
template <>
inline DataType NullValue<DataType>() {
  return DataType(DataType::kHandle, 0, 0);
}
 
class AttrFieldInfoNode : public ffi::Object {
 public:
  ffi::String name;
  ffi::String type_info;
  ffi::String description;
 
  static void RegisterReflection() {
    namespace rfl = ffi::reflection;
    rfl::ObjectDef<AttrFieldInfoNode>()
        .def_ro("name", &AttrFieldInfoNode::name)
        .def_ro("type_info", &AttrFieldInfoNode::type_info)
        .def_ro("description", &AttrFieldInfoNode::description);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.AttrFieldInfo", AttrFieldInfoNode, ffi::Object);
};
 
class AttrFieldInfo : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AttrFieldInfo, ffi::ObjectRef, AttrFieldInfoNode);
};
 
class BaseAttrsNode : public ffi::Object {
 public:
  virtual ~BaseAttrsNode() {}
  template <typename... Args>
  inline void InitBySeq(Args&&... args);
  TVM_DLL virtual void InitByPackedArgs(const ffi::PackedArgs& kwargs,
                                        bool allow_unknown = false) = 0;
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO("ir.Attrs", BaseAttrsNode, ffi::Object);
};
 
class Attrs : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Attrs, ffi::ObjectRef, BaseAttrsNode);
};
 
class DictAttrsNode : public BaseAttrsNode {
 public:
  ffi::Map<ffi::String, ffi::Any> dict;
 
  static void RegisterReflection() {
    namespace rfl = ffi::reflection;
    rfl::ObjectDef<DictAttrsNode>().def_ro("__dict__", &DictAttrsNode::dict);
  }
 
  void InitByPackedArgs(const ffi::PackedArgs& args, bool allow_unknown) final;
 
  // type info
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DictAttrs", DictAttrsNode, BaseAttrsNode);
};
 
class DictAttrs : public Attrs {
 public:
  explicit DictAttrs(ffi::UnsafeInit tag) : Attrs(tag) {}
  TVM_DLL explicit DictAttrs(ffi::Map<ffi::String, Any> dict = {});
 
  // Utils for accessing attributes
  // This needs to be on DictAttrs, not DictAttrsNode because we return the default
  // value if DictAttrsNode is not defined.
  template <typename TObjectRef>
  ffi::Optional<TObjectRef> GetAttr(
      const std::string& attr_key,
      ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
    if (!defined()) return default_value;
    const DictAttrsNode* node = this->as<DictAttrsNode>();
    auto it = node->dict.find(attr_key);
    if (it != node->dict.end()) {
      return (*it).second.cast<TObjectRef>();
    } else {
      return default_value;
    }
  }
  // variant that uses TObjectRef to enable implicit conversion to default value.
  template <typename TObjectRef>
  ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
    return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
  }
  bool HasNonzeroAttr(const std::string& attr_key) const {
    return GetAttr<Integer>(attr_key, 0).value_or(0).IntValue() != 0;
  }
 
  explicit DictAttrs(::tvm::ffi::ObjectPtr<DictAttrsNode> n) : Attrs(n) {}
  DictAttrs(const DictAttrs&) = default;
  DictAttrs(DictAttrs&&) = default;
  DictAttrs& operator=(const DictAttrs&) = default;
  DictAttrs& operator=(DictAttrs&&) = default;
  const DictAttrsNode* operator->() const { return static_cast<const DictAttrsNode*>(data_.get()); }
  const DictAttrsNode* get() const { return operator->(); }
  using ContainerType = DictAttrsNode;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DictAttrsNode);
};
 
DictAttrs WithAttrs(DictAttrs attrs, ffi::Map<ffi::String, Any> new_attrs);
 
DictAttrs WithAttr(DictAttrs attrs, ffi::String key, Any value);
 
inline DictAttrs WithAttr(DictAttrs attrs, const std::string& key, Any value) {
  return WithAttr(std::move(attrs), ffi::String(key), std::move(value));
}
 
DictAttrs WithoutAttr(DictAttrs attrs, const std::string& key);
 
template <typename TFunc>
inline TFunc WithAttr(TFunc input, const std::string& attr_key, Any attr_value) {
  using TNode = typename TFunc::ContainerType;
  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
  TNode* node = input.CopyOnWrite();
  node->attrs = WithAttr(std::move(node->attrs), attr_key, attr_value);
  return input;
}
 
template <typename TFunc>
inline TFunc WithAttrs(TFunc input, ffi::Map<ffi::String, Any> attrs) {
  using TNode = typename TFunc::ContainerType;
  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
  TNode* node = input.CopyOnWrite();
 
  node->attrs = WithAttrs(std::move(node->attrs), attrs);
 
  return input;
}
 
template <typename TFunc>
inline TFunc WithoutAttr(TFunc input, const std::string& attr_key) {
  using TNode = typename TFunc::ContainerType;
  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
 
  TNode* node = input.CopyOnWrite();
  node->attrs = WithoutAttr(std::move(node->attrs), attr_key);
 
  return input;
}
 
template <typename DerivedType>
class AttrsNodeReflAdapter : public BaseAttrsNode {
 public:
  void InitByPackedArgs(const ffi::PackedArgs& args, bool allow_unknown) final {
    TVM_FFI_THROW(InternalError) << "`" << DerivedType::_type_key
                                 << "` uses new reflection mechanism for init";
  }
 
 private:
  DerivedType* self() const {
    return const_cast<DerivedType*>(static_cast<const DerivedType*>(this));
  }
};
 
template <typename TAttrs>
inline TAttrs AttrsWithDefaultValues() {
  static_assert(std::is_base_of_v<Attrs, TAttrs>, "Can only take attr nodes");
  using ContainerType = typename TAttrs::ContainerType;
  if constexpr (std::is_base_of_v<AttrsNodeReflAdapter<ContainerType>, ContainerType>) {
    static auto finit_object = ffi::Function::GetGlobalRequired("ffi.MakeObjectFromPackedArgs");
    AnyView packed_args[1];
    packed_args[0] = ContainerType::RuntimeTypeIndex();
    ffi::Any rv;
    finit_object.CallPacked(ffi::PackedArgs(packed_args, 1), &rv);
    return rv.cast<TAttrs>();
  } else {
    auto n = ffi::make_object<ContainerType>();
    n->InitByPackedArgs(ffi::PackedArgs(nullptr, 0), false);
    return TAttrs(n);
  }
}
 
}  // namespace tvm
#endif  // TVM_IR_ATTRS_H_