Program Listing for File attrs.h

↰ Return to documentation for file (tvm/ir/attrs.h)

/*
 * 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 <string>
#include <type_traits>
#include <unordered_map>
#include <utility>

namespace tvm {

class AttrsObj : public ffi::Object {
 public:
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO("ir.Attrs", AttrsObj, ffi::Object);
};

class Attrs : public ffi::ObjectRef {
 public:
  explicit Attrs(ffi::ObjectPtr<AttrsObj> n) : ffi::ObjectRef(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(n != nullptr);
    data_ = std::move(n);
  }

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Attrs, ffi::ObjectRef, AttrsObj);
};

class DictAttrsObj : public AttrsObj {
 public:
  ffi::Map<ffi::String, ffi::Any> dict;

  static void RegisterReflection() {
    namespace rfl = ffi::reflection;
    rfl::ObjectDef<DictAttrsObj>().def_ro("dict", &DictAttrsObj::dict);
  }

  // type info
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DictAttrs", DictAttrsObj, AttrsObj);
};

class DictAttrs : public Attrs {
 public:
  explicit DictAttrs(ffi::Map<ffi::String, Any> dict = {}) : Attrs(ffi::UnsafeInit{}) {
    ffi::ObjectPtr<DictAttrsObj> n = ffi::make_object<DictAttrsObj>();
    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<DictAttrsObj>();
  }

  DictAttrs& operator=(DictAttrs&& other) noexcept {
    if (this != &other) {
      data_ = std::move(other.data_);
      other.data_ = ffi::make_object<DictAttrsObj>();
    }
    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 DictAttrsObj* node = get();
    auto it = node->dict.find(attr_key);
    if (it != node->dict.end()) {
      return (*it).second.cast<TObjectRef>();
    } else {
      return default_value;
    }
  }
  template <typename T>
  ffi::Optional<T> GetAttr(ffi::reflection::AttrKey<T> key,
                           ffi::Optional<typename ffi::reflection::AttrKey<T>::ValueType>
                               default_value = std::nullopt) const {
    return GetAttr<T>(key.name, std::move(default_value));
  }
  template <typename T>
  ffi::Optional<T> GetAttr(ffi::reflection::AttrKey<T> key,
                           typename ffi::reflection::AttrKey<T>::ValueType default_value) const {
    return GetAttr(key, ffi::Optional<T>(std::move(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<DictAttrsObj::_type_mutable, DictAttrsObj*, const DictAttrsObj*>;
  __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 = DictAttrsObj;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DictAttrsObj);
};

template <typename TFunc>
inline TFunc WithAttr(TFunc input, const std::string& attr_key, Any attr_value) {
  using TNode = typename TFunc::ContainerType;
  if (!input.unique()) {
    static ffi::reflection::TypeAttrColumn shallow_copy(ffi::type_attr::kShallowCopy);
    ffi::AnyView copy_func = shallow_copy[input->type_index()];
    TVM_FFI_CHECK(copy_func.type_index() == ffi::TypeIndex::kTVMFFIFunction, TypeError)
        << "Type " << input->GetTypeKey()
        << " must register an ffi.Function for __ffi_shallow_copy__";
    Any copy = copy_func.cast<ffi::Function>()(input);
    TVM_FFI_CHECK(copy.type_index() == input->type_index() && copy.as<ffi::Object>() != input.get(),
                  TypeError)
        << "Shallow copy must return a distinct object of type " << input->GetTypeKey();
    input = std::move(copy).cast<TFunc>();
  }
  TNode* node = const_cast<TNode*>(input.operator->());
  // 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, typename T, typename Value,
          std::enable_if_t<(std::is_convertible_v<Value&&, T> &&
                            ((!std::is_same_v<std::decay_t<Value>, ffi::Any> &&
                              !std::is_same_v<std::decay_t<Value>, ffi::AnyView>) ||
                             std::is_same_v<T, ffi::Any> || std::is_same_v<T, ffi::AnyView>)),
                           int> = 0>
inline TFunc WithAttr(TFunc input, ffi::reflection::AttrKey<T> key, Value&& value) {
  T typed_value(std::forward<Value>(value));
  return WithAttr(std::move(input), key.name, ffi::Any(typed_value));
}

template <typename TFunc>
inline TFunc WithAttrs(TFunc input, ffi::Map<ffi::String, Any> attrs) {
  using TNode = typename TFunc::ContainerType;
  if (attrs.empty()) return input;
  if (!input.unique()) {
    static ffi::reflection::TypeAttrColumn shallow_copy(ffi::type_attr::kShallowCopy);
    ffi::AnyView copy_func = shallow_copy[input->type_index()];
    TVM_FFI_CHECK(copy_func.type_index() == ffi::TypeIndex::kTVMFFIFunction, TypeError)
        << "Type " << input->GetTypeKey()
        << " must register an ffi.Function for __ffi_shallow_copy__";
    Any copy = copy_func.cast<ffi::Function>()(input);
    TVM_FFI_CHECK(copy.type_index() == input->type_index() && copy.as<ffi::Object>() != input.get(),
                  TypeError)
        << "Shallow copy must return a distinct object of type " << input->GetTypeKey();
    input = std::move(copy).cast<TFunc>();
  }
  TNode* node = const_cast<TNode*>(input.operator->());
  // 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;
  if (!input.unique()) {
    static ffi::reflection::TypeAttrColumn shallow_copy(ffi::type_attr::kShallowCopy);
    ffi::AnyView copy_func = shallow_copy[input->type_index()];
    TVM_FFI_CHECK(copy_func.type_index() == ffi::TypeIndex::kTVMFFIFunction, TypeError)
        << "Type " << input->GetTypeKey()
        << " must register an ffi.Function for __ffi_shallow_copy__";
    Any copy = copy_func.cast<ffi::Function>()(input);
    TVM_FFI_CHECK(copy.type_index() == input->type_index() && copy.as<ffi::Object>() != input.get(),
                  TypeError)
        << "Shallow copy must return a distinct object of type " << input->GetTypeKey();
    input = std::move(copy).cast<TFunc>();
  }
  TNode* node = const_cast<TNode*>(input.operator->());
  // 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;
}

template <typename TFunc, typename T>
inline TFunc WithoutAttr(TFunc input, ffi::reflection::AttrKey<T> key) {
  return WithoutAttr(std::move(input), key.name);
}

}  // namespace tvm
#endif  // TVM_IR_ATTRS_H_