Program Listing for File op.h

↰ Return to documentation for file (tvm/ir/op.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_OP_H_
#define TVM_IR_OP_H_

#include <tvm/ffi/container/list.h>
#include <tvm/ffi/expected.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/native_function.h>
#include <tvm/ffi/reflection/registry.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 <type_traits>
#include <utility>

namespace tvm {

template <typename>
class OpAttrMap;

class ArgumentInfoObj : public ffi::Object {
 public:
  ffi::String name;
  ffi::String doc;

  static void RegisterReflection() {
    namespace refl = ffi::reflection;
    refl::ObjectDef<ArgumentInfoObj>()
        .def_ro("name", &ArgumentInfoObj::name)
        .def_ro("doc", &ArgumentInfoObj::doc);
  }

  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.ArgumentInfo", ArgumentInfoObj, ffi::Object);
};

class ArgumentInfo : public ffi::ObjectRef {
 public:
  explicit ArgumentInfo(ffi::UnsafeInit tag, ffi::ObjectPtr<ArgumentInfoObj> n)
      : ffi::ObjectRef(tag) {
    data_ = std::move(n);
  }

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ArgumentInfo, ffi::ObjectRef, ArgumentInfoObj);
};

class OpObj : public ExprObj {
 private:
  // Dense process-local index into attribute columns; never serialized.
  uint32_t index_{0};

 public:
  ffi::String name;
  ffi::String doc;
  ffi::Array<ArgumentInfo> args_info;
  ffi::Array<ArgumentInfo> ty_args_info;
  ffi::Optional<ArgumentInfo> var_args_info;
  ffi::Optional<ArgumentInfo> var_ty_args_info;
  ffi::String attrs_type_key;
  uint32_t attrs_type_index{0};
  ffi::Any validator;
  bool validator_is_custom{false};

  TVM_TIRX_DLL static void RegisterReflection();

  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Op", OpObj, ExprObj);

 private:
  friend class OpRegistry;
  friend class OpDef;
  friend class Op;
  template <typename>
  friend class OpAttrMap;
};

class Op : public Expr {
 public:
  explicit Op(ffi::ObjectPtr<OpObj> n) : Expr(std::move(n)) {
    TVM_FFI_CHECK(defined(), ValueError) << "Op expects a defined OpObj";
  }

  template <typename ValueType>
  static OpAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
  template <typename ValueType>
  static OpAttrMap<ValueType> GetAttrMap(ffi::reflection::AttrKey<ValueType> key);

  TVM_TIRX_DLL static bool HasAttrMap(const ffi::String& attr_name);
  template <typename ValueType>
  static bool HasAttrMap(ffi::reflection::AttrKey<ValueType> key) {
    return HasAttrMap(key.name);
  }

  TVM_TIRX_DLL static Op Get(const ffi::String& op_name);
  TVM_TIRX_DLL static ffi::Array<ffi::String> ListNames();

  TVM_FFI_INLINE void Validate(const CallObj* call) const {
    const auto& validator = get()->validator;
    if (validator != nullptr) {
      if (TVM_FFI_PREDICT_FALSE(!call)) {
        ThrowInvalidCall(get());
      }
      using View = ffi::reflection::NativeFunctionView<void(const CallObj*)>;
      // set_validator stores only an owning NativeFunction with this signature.
      ffi::detail::AnyUnsafe::CopyFromAnyViewAfterCheck<View>(validator).CallExpected(call).value();
    }
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Op, Expr, OpObj);

 private:
  friend class OpDef;
  TVM_TIRX_DLL static void Expose(const ffi::String& op_name);

  TVM_TIRX_DLL static void ThrowInvalidCall(const OpObj* op);
  TVM_TIRX_DLL static ffi::List<ffi::Any> GetAttrColumn(const ffi::String& attr_name);
};

namespace op_attr {
inline constexpr ffi::reflection::AttrKey<
    ffi::reflection::NativeFunctionView<Type(const CallObj* call)>>
    kInferType{"FInferType"};
inline constexpr ffi::reflection::AttrKey<Type> kFixedReturnType{"TFixedReturnType"};
}  // namespace op_attr

// Compatibility names; new callers can deduce the value type from the key.
using FInferType = decltype(op_attr::kInferType)::ValueType;
using TFixedReturnType = decltype(op_attr::kFixedReturnType)::ValueType;

namespace sig {

struct SignatureTrait {
  static constexpr size_t kArgsCount = 0;
  static constexpr size_t kTyArgsCount = 0;
  static constexpr size_t kVarArgsCount = 0;
  static constexpr size_t kVarTyArgsCount = 0;
  static constexpr size_t kCallAttrsCount = 0;
};

struct ValidateState {
  const CallObj* call;
  const OpObj* op;
  size_t args_index = 0;
  size_t ty_args_index = 0;
};

namespace detail {

template <typename T>
inline bool ReportTypeMismatch(ffi::Expected<void>* out, const ValidateState* state, bool type_arg,
                               size_t index, ffi::AnyView actual) {
  const auto& fixed = type_arg ? state->op->ty_args_info : state->op->args_info;
  const auto& tail = type_arg ? state->op->var_ty_args_info : state->op->var_args_info;
  std::string name;
  if (index < fixed.size()) {
    name = fixed[index]->name;
  } else if (tail.has_value()) {
    name = tail.value()->name;
  }
  std::string name_clause = name.empty() ? "" : " (`" + name + "`)";
  TVMFFIAny raw = actual.CopyToTVMFFIAny();
  *out = TVM_FFI_UNEXPECTED(TypeError)
         << "Op `" << state->op->name << "`: `" << (type_arg ? "Call.ty_args[" : "Call.args[")
         << index << "]`" << name_clause << " expected `" << ffi::TypeTraits<T>::TypeStr()
         << "`, got `" << ffi::TypeTraits<T>::GetMismatchTypeInfo(&raw) << "`.";
  return false;
}

}  // namespace detail

template <typename T = Expr>
struct arg : SignatureTrait {
  static_assert(std::is_base_of_v<Expr, T>, "arg<T> expects an expression view");
  static constexpr size_t kArgsCount = 1;
  ffi::String name;
  ffi::String doc;

  explicit arg(ffi::String name, ffi::String doc = "")
      : name(std::move(name)), doc(std::move(doc)) {}

  void CollectMetadata(OpObj* op) const {
    auto info = ffi::make_object<ArgumentInfoObj>();
    info->name = name;
    info->doc = doc;
    op->args_info.push_back(ArgumentInfo(ffi::UnsafeInit{}, std::move(info)));
  }

  static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
    if constexpr (!std::is_same_v<T, Expr>) {
      ffi::AnyView value(state->call->args.GetArrayObj()->begin()[state->args_index]);
      if (TVM_FFI_PREDICT_FALSE(value == nullptr ||
                                !ffi::detail::AnyUnsafe::CheckAnyViewStrict<T>(value))) {
        return detail::ReportTypeMismatch<T>(out, state, false, state->args_index, value);
      }
    }
    ++state->args_index;
    return true;
  }
};

template <typename T = Expr>
struct var_args : SignatureTrait {
  static_assert(std::is_base_of_v<Expr, T>, "var_args<T> expects an expression view");
  static constexpr size_t kVarArgsCount = 1;
  ffi::String name;
  ffi::String doc;

  explicit var_args(ffi::String name, ffi::String doc = "")
      : name(std::move(name)), doc(std::move(doc)) {}

  void CollectMetadata(OpObj* op) const {
    auto info = ffi::make_object<ArgumentInfoObj>();
    info->name = name;
    info->doc = doc;
    op->var_args_info = ArgumentInfo(ffi::UnsafeInit{}, std::move(info));
  }

  static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
    if constexpr (!std::is_same_v<T, Expr>) {
      const ffi::Any* values = state->call->args.GetArrayObj()->begin();
      for (; state->args_index != state->call->args.size(); ++state->args_index) {
        ffi::AnyView value(values[state->args_index]);
        if (TVM_FFI_PREDICT_FALSE(value == nullptr ||
                                  !ffi::detail::AnyUnsafe::CheckAnyViewStrict<T>(value))) {
          return detail::ReportTypeMismatch<T>(out, state, false, state->args_index, value);
        }
      }
    } else {
      state->args_index = state->call->args.size();
    }
    return true;
  }
};

template <typename T = Type>
struct ty_arg : SignatureTrait {
  static_assert(std::is_base_of_v<Type, T>, "ty_arg<T> expects a type view");
  static constexpr size_t kTyArgsCount = 1;
  ffi::String name;
  ffi::String doc;

  explicit ty_arg(ffi::String name, ffi::String doc = "")
      : name(std::move(name)), doc(std::move(doc)) {}

  void CollectMetadata(OpObj* op) const {
    auto info = ffi::make_object<ArgumentInfoObj>();
    info->name = name;
    info->doc = doc;
    op->ty_args_info.push_back(ArgumentInfo(ffi::UnsafeInit{}, std::move(info)));
  }

  static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
    if constexpr (!std::is_same_v<T, Type>) {
      ffi::AnyView value(state->call->ty_args.GetArrayObj()->begin()[state->ty_args_index]);
      if (TVM_FFI_PREDICT_FALSE(value == nullptr ||
                                !ffi::detail::AnyUnsafe::CheckAnyViewStrict<T>(value))) {
        return detail::ReportTypeMismatch<T>(out, state, true, state->ty_args_index, value);
      }
    }
    ++state->ty_args_index;
    return true;
  }
};

template <typename T = Type>
struct var_ty_args : SignatureTrait {
  static_assert(std::is_base_of_v<Type, T>, "var_ty_args<T> expects a type view");
  static constexpr size_t kVarTyArgsCount = 1;
  ffi::String name;
  ffi::String doc;

  explicit var_ty_args(ffi::String name, ffi::String doc = "")
      : name(std::move(name)), doc(std::move(doc)) {}

  void CollectMetadata(OpObj* op) const {
    auto info = ffi::make_object<ArgumentInfoObj>();
    info->name = name;
    info->doc = doc;
    op->var_ty_args_info = ArgumentInfo(ffi::UnsafeInit{}, std::move(info));
  }

  static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
    if constexpr (!std::is_same_v<T, Type>) {
      const ffi::Any* values = state->call->ty_args.GetArrayObj()->begin();
      for (; state->ty_args_index != state->call->ty_args.size(); ++state->ty_args_index) {
        ffi::AnyView value(values[state->ty_args_index]);
        if (TVM_FFI_PREDICT_FALSE(value == nullptr ||
                                  !ffi::detail::AnyUnsafe::CheckAnyViewStrict<T>(value))) {
          return detail::ReportTypeMismatch<T>(out, state, true, state->ty_args_index, value);
        }
      }
    } else {
      state->ty_args_index = state->call->ty_args.size();
    }
    return true;
  }
};

template <typename T>
struct call_attrs : SignatureTrait {
  static_assert(std::is_base_of_v<AttrsObj, T>, "call_attrs<T> expects an attribute node type");
  static constexpr size_t kCallAttrsCount = 1;

  void CollectMetadata(OpObj* op) const {
    uint32_t index = T::RuntimeTypeIndex();
    op->attrs_type_key = T::_type_key;
    op->attrs_type_index = index;
  }

  static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
    ffi::AnyView attrs(state->call->attrs);
    if (TVM_FFI_PREDICT_FALSE(!state->call->attrs.has_value() ||
                              !ffi::detail::AnyUnsafe::CheckAnyViewStrict<const T*>(attrs))) {
      TVMFFIAny raw = attrs.CopyToTVMFFIAny();
      std::string actual = state->call->attrs.has_value()
                               ? "`" + ffi::TypeTraits<const T*>::GetMismatchTypeInfo(&raw) + "`"
                               : "None";
      *out = TVM_FFI_UNEXPECTED(TypeError)
             << "Op `" << state->op->name << "`: Call.attrs expected `"
             << ffi::TypeTraits<const T*>::TypeStr() << "`, got " << actual;
      return false;
    }
    return true;
  }
};

}  // namespace sig

class OpDef {
 public:
  TVM_TIRX_DLL explicit OpDef(const ffi::String& name);
  TVM_TIRX_DLL OpDef(const ffi::String& name, const ffi::String& doc);
  OpDef(const OpDef&) = delete;
  OpDef& operator=(const OpDef&) = delete;
  OpDef(OpDef&&) = delete;
  OpDef& operator=(OpDef&&) = delete;
  Op op() const { return op_; }
  inline void expose() const { Op::Expose(op_->name); }
  TVM_TIRX_DLL OpDef& add_arg(const ffi::String& name, const ffi::String& doc);
  template <typename AttrsType>
  OpDef& attrs_type() {
    uint32_t index = AttrsType::RuntimeTypeIndex();
    get()->attrs_type_key = AttrsType::_type_key;
    get()->attrs_type_index = index;
    return *this;
  }
  TVM_TIRX_DLL OpDef& set_attrs_type_key(const ffi::String& key);
  TVM_TIRX_DLL OpDef& add_ty_arg(const ffi::String& name, const ffi::String& doc);
  template <typename... Specs>
  OpDef& signature(const Specs&... specs) {
    auto updated = ffi::make_object<OpObj>();
    (ApplySignatureTrait(updated.get(), specs), ...);
    if (op_->validator == nullptr) {
      using View = ffi::reflection::NativeFunctionView<void(const CallObj*)>;
      set_validator(View::FromNative<&ValidateSignature<Specs...>>());
      get()->validator_is_custom = false;
    }
    OpObj* target = get();
    target->args_info = std::move(updated->args_info);
    target->ty_args_info = std::move(updated->ty_args_info);
    target->var_args_info = std::move(updated->var_args_info);
    target->var_ty_args_info = std::move(updated->var_ty_args_info);
    target->attrs_type_key = std::move(updated->attrs_type_key);
    target->attrs_type_index = updated->attrs_type_index;
    return *this;
  }
  template <typename ValueType>
  OpDef& set_attr(const ffi::String& attr_name, const ValueType& value, bool override = false) {
    UpdateAttr(attr_name, ffi::Any(value), override);
    return *this;
  }
  template <typename ValueType, typename Value,
            std::enable_if_t<(std::is_convertible_v<Value&&, ValueType> &&
                              ((!std::is_same_v<std::decay_t<Value>, ffi::Any> &&
                                !std::is_same_v<std::decay_t<Value>, ffi::AnyView>) ||
                               std::is_same_v<ValueType, ffi::Any> ||
                               std::is_same_v<ValueType, ffi::AnyView>)),
                             int> = 0>
  OpDef& set_attr(ffi::reflection::AttrKey<ValueType> key, Value&& value, bool override = false) {
    ValueType typed_value(std::forward<Value>(value));
    return set_attr(key.name, typed_value, override);
  }
  template <auto Fn, typename ValueType>
  OpDef& set_native_fn_attr(ffi::reflection::AttrKey<ValueType> key, bool override = false) {
    return set_attr(key, ValueType::template FromNative<Fn>(), override);
  }
  OpDef& set_validator(ffi::reflection::NativeFunctionView<void(const CallObj*)> validator,
                       bool override = false) {
    TVM_FFI_CHECK(override || op_->validator == nullptr, ValueError)
        << "Validator of " << op_->name << " is already registered";
    get()->validator = ffi::reflection::NativeFunction<void(const CallObj*)>::From(validator);
    get()->validator_is_custom = true;
    return *this;
  }
  TVM_TIRX_DLL OpDef& reset_attr(const ffi::String& attr_name);
  template <typename ValueType>
  OpDef& reset_attr(ffi::reflection::AttrKey<ValueType> key) {
    return reset_attr(key.name);
  }

 private:
  template <typename Spec>
  static void ApplySignatureTrait(OpObj* op, const Spec& spec) {
    static_assert(std::is_base_of_v<sig::SignatureTrait, Spec>, "Unknown signature descriptor");
    spec.CollectMetadata(op);
  }

  template <typename... Specs>
  TVM_FFI_INLINE static ffi::Expected<void> ValidateSignature(const CallObj* call) noexcept {
    static_assert((std::is_base_of_v<sig::SignatureTrait, Specs> && ...),
                  "Unknown signature descriptor");
    constexpr size_t n_args = (size_t{0} + ... + Specs::kArgsCount);
    constexpr size_t n_ty_args = (size_t{0} + ... + Specs::kTyArgsCount);
    constexpr size_t n_var_args = (size_t{0} + ... + Specs::kVarArgsCount);
    constexpr size_t n_var_ty_args = (size_t{0} + ... + Specs::kVarTyArgsCount);
    constexpr size_t n_attrs = (size_t{0} + ... + Specs::kCallAttrsCount);
    static_assert(n_var_args <= 1 && n_var_ty_args <= 1 && n_attrs <= 1,
                  "Duplicate signature tail or call_attrs");
    const OpObj* op = call ? call->op.as<OpObj>() : nullptr;
    if (TVM_FFI_PREDICT_FALSE(!op)) {
      return TVM_FFI_UNEXPECTED(TypeError) << "Invalid Op Call";
    }
    if (TVM_FFI_PREDICT_FALSE(!call->args.defined() || !call->ty_args.defined())) {
      return TVM_FFI_UNEXPECTED(TypeError) << "Op `" << op->name << "`: invalid Call";
    }
    ffi::Expected<void> result;
    auto* out = &result;
    if (TVM_FFI_PREDICT_FALSE(n_var_args ? call->args.size() < n_args
                                         : call->args.size() != n_args)) {
      *out = TVM_FFI_UNEXPECTED(TypeError)
             << "Op `" << op->name << "`: Call.args expected " << (n_var_args ? "at least " : "")
             << n_args << (n_args == 1 ? " argument" : " arguments") << ", got "
             << call->args.size();
      return result;
    }
    if (TVM_FFI_PREDICT_FALSE(n_var_ty_args ? call->ty_args.size() < n_ty_args
                                            : call->ty_args.size() != n_ty_args)) {
      *out = TVM_FFI_UNEXPECTED(TypeError)
             << "Op `" << op->name << "`: Call.ty_args expected "
             << (n_var_ty_args ? "at least " : "") << n_ty_args
             << (n_ty_args == 1 ? " type argument" : " type arguments") << ", got "
             << call->ty_args.size();
      return result;
    }
    if constexpr (sizeof...(Specs) != 0) {
      sig::ValidateState state{call, op};
      if (TVM_FFI_PREDICT_FALSE(!(Specs::Validate(&state, out) && ...))) {
        return result;
      }
    }
    return result;
  }

  OpObj* get() { return const_cast<OpObj*>(op_.operator->()); }
  TVM_TIRX_DLL void UpdateAttr(const ffi::String& attr_name, ffi::Any value, bool override);
  Op op_;
};

template <typename ValueType>
class OpAttrMap {
 public:
  int count(const Op& op) const {
    return op->index_ < column_.size() && column_[op->index_] != nullptr;
  }
  ValueType operator[](const Op& op) const {
    TVM_FFI_ICHECK(count(op)) << "Attribute " << attr_name_ << " has not been registered for "
                              << op->name;
    if constexpr (std::is_same_v<ValueType, ffi::Any>) {
      return column_[op->index_];
    } else {
      return column_[op->index_].template cast<ValueType>();
    }
  }
  ValueType get(const Op& op, ValueType def_value) const {
    return count(op) ? (*this)[op] : def_value;
  }
  ValueType get(const Expr& expr, ValueType def_value) const {
    TVM_FFI_ICHECK(expr.defined());
    if (const auto* op = expr.as<OpObj>()) {
      return get(ffi::GetRef<Op>(op), def_value);
    }
    return def_value;
  }

 private:
  friend class Op;
  OpAttrMap(ffi::List<ffi::Any> column, ffi::String attr_name)
      : column_(std::move(column)), attr_name_(std::move(attr_name)) {}
  ffi::List<ffi::Any> column_;
  ffi::String attr_name_;
};

template <typename ValueType>
OpAttrMap<ValueType> Op::GetAttrMap(const ffi::String& attr_name) {
  return OpAttrMap<ValueType>(GetAttrColumn(attr_name), attr_name);
}

template <typename ValueType>
OpAttrMap<ValueType> Op::GetAttrMap(ffi::reflection::AttrKey<ValueType> key) {
  return GetAttrMap<ValueType>(key.name);
}

}  // namespace tvm
#endif  // TVM_IR_OP_H_