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/registry.h>
#include <tvm/ir/cow.h>
 
#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>
 
namespace tvm {
 
class AttrsNode : public ffi::Object {
 public:
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO("ir.Attrs", AttrsNode, ffi::Object);
};
 
class Attrs : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Attrs, ffi::ObjectRef, AttrsNode);
};
 
class DictAttrsNode : public AttrsNode {
 public:
  ffi::Map<ffi::String, ffi::Any> dict;
 
  static void RegisterReflection() {
    namespace rfl = ffi::reflection;
    rfl::ObjectDef<DictAttrsNode>().def_ro("__dict__", &DictAttrsNode::dict);
  }
 
  // type info
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DictAttrs", DictAttrsNode, AttrsNode);
};
 
class DictAttrs : public Attrs {
 public:
  explicit DictAttrs(ffi::Map<ffi::String, Any> dict = {}) {
    ffi::ObjectPtr<DictAttrsNode> n = ffi::make_object<DictAttrsNode>();
    n->dict = std::move(dict);
    data_ = std::move(n);
  }
 
  DictAttrs(DictAttrs&& other) noexcept : Attrs(ffi::UnsafeInit{}) {
    data_ = std::move(other.data_);
    other.data_ = ffi::make_object<DictAttrsNode>();
  }
 
  DictAttrs& operator=(DictAttrs&& other) noexcept {
    if (this != &other) {
      data_ = std::move(other.data_);
      other.data_ = ffi::make_object<DictAttrsNode>();
    }
    return *this;
  }
 
  // Explicit copy ctor/assign defaults. Declaring the move members above
  // would otherwise suppress the implicit copy members.
  DictAttrs(const DictAttrs& other) = default;
  DictAttrs& operator=(const DictAttrs& other) = default;
 
  // Utils for accessing attributes
  template <typename TObjectRef>
  ffi::Optional<TObjectRef> GetAttr(
      const std::string& attr_key,
      ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
    const DictAttrsNode* node = get();
    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<int64_t>(attr_key, 0).value_or(0) != 0;
  }
 
  // Inline-expand TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE here, minus
  // the default copy/move it normally injects (we define our own move members
  // above so the moved-from instance stays defined-but-empty).
  explicit DictAttrs(::tvm::ffi::UnsafeInit tag) : Attrs(tag) {}
  using __PtrType =
      std::conditional_t<DictAttrsNode::_type_mutable, DictAttrsNode*, const DictAttrsNode*>;
  __PtrType operator->() const { return static_cast<__PtrType>(data_.get()); }
  __PtrType get() const { return static_cast<__PtrType>(data_.get()); }
  static constexpr bool _type_is_nullable = false;
  using ContainerType = DictAttrsNode;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DictAttrsNode);
};
 
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 is NOTNULLABLE by contract, but defend against a caller
  // that left a moved-from DictAttrs in place by re-initializing here.
  if (!node->attrs.defined()) node->attrs = DictAttrs();
  node->attrs.CopyOnWrite()->dict.Set(attr_key, std::move(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");
  if (attrs.empty()) return input;
  TNode* node = input.CopyOnWrite();
  // node->attrs is NOTNULLABLE by contract, but defend against a caller
  // that left a moved-from DictAttrs in place by re-initializing here.
  if (!node->attrs.defined()) node->attrs = DictAttrs();
  auto* dict_node = node->attrs.CopyOnWrite();
  for (const auto& [k, v] : attrs) {
    dict_node->dict.Set(k, v);
  }
  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 is NOTNULLABLE by contract, but defend against a caller
  // that left a moved-from DictAttrs in place; nothing to erase from an
  // empty dict.
  if (!node->attrs.defined()) {
    node->attrs = DictAttrs();
    return input;
  }
  node->attrs.CopyOnWrite()->dict.erase(attr_key);
  return input;
}
 
}  // namespace tvm
#endif  // TVM_IR_ATTRS_H_