op.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_OP_H_
#define TVM_IR_OP_H_
 
#include <tvm/ffi/error.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attr_registry_map.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/env_func.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/type.h>
 
#include <string>
#include <utility>
#include <vector>
 
namespace tvm {
 
// forward declare name.
template <typename>
class OpAttrMap;
 
class ArgumentInfoNode : public ffi::Object {
 public:
  ffi::String name;
  ffi::String type_info;
  ffi::String description;
 
  static void RegisterReflection() {
    namespace rfl = ffi::reflection;
    rfl::ObjectDef<ArgumentInfoNode>()
        .def_ro("name", &ArgumentInfoNode::name)
        .def_ro("type_info", &ArgumentInfoNode::type_info)
        .def_ro("description", &ArgumentInfoNode::description);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.ArgumentInfo", ArgumentInfoNode, ffi::Object);
};
 
class ArgumentInfo : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ArgumentInfo, ffi::ObjectRef, ArgumentInfoNode);
};
 
// TODO(tvm-team): migrate low-level intrinsics to use Op
class OpNode : public ExprNode {
 public:
  ffi::String name;
  ffi::String description;
  /* \brief Information of input arguments to the operator */
  ffi::Array<ArgumentInfo> arguments;
  ffi::String attrs_type_key;
  uint32_t attrs_type_index{0};
  int32_t num_inputs = -1;
  int32_t support_level = 10;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<OpNode>()
        .def_ro("name", &OpNode::name)
        .def_ro("description", &OpNode::description, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("arguments", &OpNode::arguments, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("attrs_type_key", &OpNode::attrs_type_key, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("num_inputs", &OpNode::num_inputs, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("support_level", &OpNode::support_level, refl::AttachFieldFlag::SEqHashIgnore());
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Op", OpNode, ExprNode);
 
 private:
  uint32_t AttrRegistryIndex() const { return index_; }
  std::string AttrRegistryName() const { return name; }
 
  // friend class
  template <typename>
  friend class AttrRegistryMapContainerMap;
  template <typename, typename>
  friend class AttrRegistry;
  friend class OpRegEntry;
 
  // Program internal unique index of operator.
  // Used to help index the program.
  uint32_t index_{0};
};
 
class Op : public Expr {
 public:
  explicit Op(ffi::ObjectPtr<OpNode> n) : Expr(std::move(n)) {
    TVM_FFI_CHECK(defined(), ValueError) << "Op expects a defined OpNode";
  }
 
  template <typename ValueType>
  inline static OpAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
  TVM_DLL static bool HasAttrMap(const ffi::String& attr_name);
  TVM_DLL static const Op& Get(const ffi::String& op_name);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Op, Expr, OpNode);
 
 private:
  TVM_DLL static const AttrRegistryMapContainerMap<Op>& GetAttrMapContainer(const ffi::String& key);
};
 
class OpRegEntry {
 public:
  const Op& op() const { return op_; }
  inline OpRegEntry& describe(const std::string& descr);  // NOLINT(*)
  inline OpRegEntry& add_argument(const std::string& name, const std::string& type,
                                  const std::string& description);
  template <typename AttrsType>
  inline OpRegEntry& set_attrs_type();
  inline OpRegEntry& set_attrs_type_key(const ffi::String& key);
  inline OpRegEntry& set_num_inputs(int32_t n);  // NOLINT(*)
  inline OpRegEntry& set_support_level(int32_t level);  // NOLINT(*)
  template <typename ValueType>
  inline OpRegEntry& set_attr(const std::string& attr_name,  // NOLINT(*)
                              const ValueType& value, int plevel = 10);
 
  inline void reset_attr(const std::string& attr_name);
 
  // set the name of the op to be the same as registry
  inline OpRegEntry& set_name() {  // NOLINT(*)
    if (get()->name.length() == 0) {
      get()->name = name;
    }
    return *this;
  }
  TVM_DLL static OpRegEntry& RegisterOrGet(const ffi::String& name);
 
 private:
  template <typename, typename>
  friend class AttrRegistry;
  // the name
  std::string name;
  Op op_;
  static Op MakeOp(uint32_t reg_index);
  // private constructor
  TVM_DLL OpRegEntry(uint32_t reg_index);
  // return internal pointer to op.
  inline OpNode* get();
  // update the attribute OpAttrMap
  TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
};
 
template <typename ValueType>
class OpAttrMap : public AttrRegistryMap<Op, ValueType> {
 public:
  inline ValueType get(const Expr& expr, ValueType def_value) const;
 
  using TParent = AttrRegistryMap<Op, ValueType>;
  using TParent::count;
  using TParent::get;
  using TParent::operator[];
 
 private:
  friend class Op;
  // constructor
  explicit OpAttrMap(const AttrRegistryMapContainerMap<Op>& map) : TParent(map) {}
};
 
// internal macros to make
#define TVM_OP_REGISTER_VAR_DEF [[maybe_unused]] static ::tvm::OpRegEntry& __make_##Op
 
#define TVM_REGISTER_OP(OpName)                              \
  TVM_FFI_STR_CONCAT(TVM_OP_REGISTER_VAR_DEF, __COUNTER__) = \
      ::tvm::OpRegEntry::RegisterOrGet(OpName).set_name()
 
// implementations
 
template <typename ValueType>
inline OpAttrMap<ValueType> Op::GetAttrMap(const ffi::String& key) {
  return OpAttrMap<ValueType>(Op::GetAttrMapContainer(key));
}
 
inline OpNode* OpRegEntry::get() { return const_cast<OpNode*>(op_.operator->()); }
 
inline OpRegEntry& OpRegEntry::describe(const std::string& descr) {  // NOLINT(*)
  get()->description = descr;
  return *this;
}
 
inline OpRegEntry& OpRegEntry::add_argument(const std::string& name, const std::string& type,
                                            const std::string& description) {
  auto n = ffi::make_object<ArgumentInfoNode>();
  n->name = name;
  n->type_info = type;
  n->description = description;
  get()->arguments.push_back(ArgumentInfo(n));
  return *this;
}
 
inline OpRegEntry& OpRegEntry::set_num_inputs(int32_t n) {  // NOLINT(*)
  get()->num_inputs = n;
  return *this;
}
 
template <typename AttrsType>
inline OpRegEntry& OpRegEntry::set_attrs_type() {  // NOLINT(*)
  get()->attrs_type_key = AttrsType::_type_key;
  get()->attrs_type_index = AttrsType::RuntimeTypeIndex();
  return *this;
}
 
inline OpRegEntry& OpRegEntry::set_attrs_type_key(const ffi::String& key) {  // NOLINT(*)
  get()->attrs_type_key = key;
  get()->attrs_type_index = tvm::ffi::TypeKeyToIndex(key.c_str());
  return *this;
}
 
inline OpRegEntry& OpRegEntry::set_support_level(int32_t n) {  // NOLINT(*)
  get()->support_level = n;
  return *this;
}
 
template <typename ValueType>
inline OpRegEntry& OpRegEntry::set_attr(  // NOLINT(*)
    const std::string& attr_name, const ValueType& value, int plevel) {
  TVM_FFI_ICHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
  UpdateAttr(attr_name, Any(value), plevel);
  return *this;
}
 
// member functions of OpAttrMap
 
template <typename ValueType>
inline ValueType OpAttrMap<ValueType>::get(const Expr& expr, ValueType def_value) const {
  TVM_FFI_ICHECK(expr.defined());
  if (const OpNode* op = expr.as<OpNode>()) {
    return this->map_.get(ffi::GetRef<Op>(op), def_value);
  } else {
    return def_value;
  }
}
 
}  // namespace tvm
#endif  // TVM_IR_OP_H_