base_expr.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_BASE_EXPR_H_
#define TVM_IR_BASE_EXPR_H_
 
#include <tvm/ffi/cast.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/location.h>
 
#include <cstddef>
#include <cstdint>
#include <optional>
#include <type_traits>
 
namespace tvm {
 
class TypeNode : public ffi::Object {
 public:
  mutable Location loc = UnknownLoc();
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    // loc do not participate in structural equal and hash.
    refl::ObjectDef<TypeNode>().def_ro("loc", &TypeNode::loc, refl::DefaultValue(UnknownLoc()),
                                       refl::AttachFieldFlag::SEqHashIgnore());
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  static constexpr bool _type_s_eq_hash_subclass_kind_fixed = true;
 
  static constexpr const uint32_t _type_child_slots = 14;
  TVM_FFI_DECLARE_OBJECT_INFO("ir.Type", TypeNode, ffi::Object);
};
 
class Type : public ffi::ObjectRef {
 public:
  TVM_TIRX_DLL static Type Missing();
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef, TypeNode);
};
 
class MissingTypeNode final : public TypeNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<MissingTypeNode>();
  }
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.MissingType", MissingTypeNode, TypeNode);
};
 
class MissingType final : public Type {
 public:
  TVM_TIRX_DLL MissingType();
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MissingType, Type, MissingTypeNode);
};
 
class OpaqueTypeNode final : public TypeNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<OpaqueTypeNode>();
  }
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.OpaqueType", OpaqueTypeNode, TypeNode);
};
 
class OpaqueType final : public Type {
 public:
  TVM_TIRX_DLL OpaqueType();
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpaqueType, Type, OpaqueTypeNode);
};
 
class PrimTypeNode final : public TypeNode {
 public:
  DLDataType dtype;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PrimTypeNode>().def_ro("dtype", &PrimTypeNode::dtype);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.PrimType", PrimTypeNode, TypeNode);
};
 
/*
 * \brief Managed reference to PrimTypeNode.
 * \sa PrimTypeNode
 */
class PrimType final : public Type {
 public:
  TVM_TIRX_DLL explicit PrimType(DLDataType dtype);
 
  TVM_TIRX_DLL PrimType(DLDataTypeCode code, int bits, int lanes = 1);
 
  TVM_TIRX_DLL static PrimType Int(int bits, int lanes = 1);
  TVM_TIRX_DLL static PrimType UInt(int bits, int lanes = 1);
  TVM_TIRX_DLL static PrimType Float(int bits, int lanes = 1);
  TVM_TIRX_DLL static PrimType BFloat(int bits, int lanes = 1);
  TVM_TIRX_DLL static PrimType Bool(int lanes = 1);
  TVM_TIRX_DLL static PrimType Void();
  TVM_TIRX_DLL static PrimType ScalableVector(DLDataTypeCode code, int bits, int lanes);
 
  TVM_FFI_INLINE DLDataTypeCode code() const {
    return static_cast<DLDataTypeCode>(static_cast<int>(get()->dtype.code));
  }
 
  TVM_FFI_INLINE int32_t bits() const { return get()->dtype.bits; }
 
  TVM_FFI_INLINE int32_t lanes() const {
    int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
    if (TVM_FFI_PREDICT_FALSE(encoded_lanes < 0)) {
      TVM_FFI_THROW(InternalError)
          << "Can't fetch the lanes of a scalable vector at a compile time.";
    }
    return encoded_lanes;
  }
 
  TVM_FFI_INLINE bool MatchesElementType(DLDataTypeCode code, int bits) const {
    DLDataType dtype = get()->dtype;
    return dtype.code == static_cast<uint8_t>(code) && dtype.bits == bits;
  }
 
  template <typename... Codes>
  TVM_FFI_INLINE bool MatchesCode(Codes... codes) const {
    uint8_t dtype_code = get()->dtype.code;
    return ((dtype_code == static_cast<uint8_t>(codes)) || ...);
  }
 
  TVM_FFI_INLINE bool IsScalar() const {
    int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
    return encoded_lanes == 1;
  }
 
  TVM_FFI_INLINE bool IsVoid() const {
    DLDataType dtype = get()->dtype;
    return dtype.code == static_cast<uint8_t>(DLDataTypeCode::kDLOpaqueHandle) && dtype.bits == 0 &&
           static_cast<int16_t>(dtype.lanes) == 0;
  }
 
  TVM_FFI_INLINE bool IsScalableVector() const {
    return static_cast<int16_t>(get()->dtype.lanes) < -1;
  }
 
  TVM_FFI_INLINE bool IsFixedLengthVector() const {
    return static_cast<int16_t>(get()->dtype.lanes) > 1;
  }
 
  TVM_FFI_INLINE size_t StorageBytes() const {
    DLDataType dtype = get()->dtype;
    int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
    if (TVM_FFI_PREDICT_FALSE(encoded_lanes < 0)) {
      TVM_FFI_THROW(InternalError)
          << "Cannot compute compile-time storage bytes for non-fixed vector type " << dtype;
    }
    return static_cast<size_t>(
        (static_cast<uint64_t>(dtype.bits) * static_cast<uint64_t>(dtype.lanes) + 7) / 8);
  }
 
  TVM_FFI_INLINE PrimType WithCode(DLDataTypeCode code) const {
    DLDataType dtype = get()->dtype;
    int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
    if (encoded_lanes < -1) {
      return ScalableVector(code, dtype.bits, -encoded_lanes);
    }
    return PrimType(code, dtype.bits, encoded_lanes);
  }
 
  TVM_FFI_INLINE PrimType WithBits(int bits) const {
    DLDataType dtype = get()->dtype;
    int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
    if (encoded_lanes < -1) {
      return ScalableVector(this->code(), bits, -encoded_lanes);
    }
    return PrimType(this->code(), bits, encoded_lanes);
  }
 
  TVM_FFI_INLINE PrimType WithLanes(int lanes) const {
    return PrimType(this->code(), this->bits(), lanes);
  }
 
  TVM_FFI_INLINE int32_t VScaleFactor() const {
    int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
    if (encoded_lanes >= -1) {
      TVM_FFI_THROW(InternalError) << "A fixed length vector doesn't have a vscale factor.";
    }
    return -encoded_lanes;
  }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimType, Type, PrimTypeNode);
};
 
inline bool operator==(const PrimType& lhs, const PrimType& rhs) {
  return lhs->dtype == rhs->dtype;
}
 
inline bool operator!=(const PrimType& lhs, const PrimType& rhs) { return !(lhs == rhs); }
 
class ExprNode : public ffi::Object {
 public:
  mutable Location loc = UnknownLoc();
 
  mutable Type ty = MissingType();
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    // loc does not participate in structural equal and hash.
    refl::ObjectDef<ExprNode>()
        .def_ro("loc", &ExprNode::loc, refl::DefaultValue(UnknownLoc()),
                refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("ty", &ExprNode::ty, refl::DefaultValue(MissingType()));
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
 
  static constexpr const uint32_t _type_child_slots = 64;
  TVM_FFI_DECLARE_OBJECT_INFO("ir.Expr", ExprNode, ffi::Object);
};
 
class Expr : public ffi::ObjectRef {
 public:
  // Expressions do not implicitly compare by object identity or address.  Callers must name
  // whether they intend object identity, structural equality, or primitive symbolic comparison.
  bool operator==(const Expr& other) const = delete;
  bool operator!=(const Expr& other) const = delete;
  bool operator<(const Expr& other) const = delete;
 
  explicit Expr(ffi::ObjectPtr<ExprNode> node) : ffi::ObjectRef(std::move(node)) {
    TVM_FFI_CHECK(defined(), ValueError) << "Expr expects a defined node";
  }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Expr, ffi::ObjectRef, ExprNode);
};
 
class StagingExprNode : public ExprNode {
 public:
  static void RegisterReflection() { ffi::reflection::ObjectDef<StagingExprNode>(); }
  static constexpr const uint32_t _type_child_slots = 4;
  TVM_FFI_DECLARE_OBJECT_INFO("ir.StagingExpr", StagingExprNode, ExprNode);
};
 
class StagingExpr : public Expr {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StagingExpr, Expr, StagingExprNode);
};
 
class OpaqueExprNode : public ExprNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<OpaqueExprNode>();
  }
 
  static constexpr const uint32_t _type_child_slots = 2;
  TVM_FFI_DECLARE_OBJECT_INFO("ir.OpaqueExpr", OpaqueExprNode, ExprNode);
};
 
class OpaqueExpr : public Expr {
 public:
  explicit OpaqueExpr(ffi::ObjectPtr<OpaqueExprNode> node) : Expr(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpaqueExpr, Expr, OpaqueExprNode);
};
 
class Call;
 
template <typename ExpectedType>
class TypedExpr : public Expr {
 public:
  ExpectedType ty() const {
    const auto* node = get();
    TVM_FFI_DCHECK(node != nullptr);
    const auto* ty_node = node->ExprNode::ty.template as<typename ExpectedType::ContainerType>();
    TVM_FFI_DCHECK(ty_node != nullptr);
    return ffi::GetRef<ExpectedType>(ty_node);
  }
 
  explicit TypedExpr(ffi::ObjectPtr<ExprNode> node) : Expr(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypedExpr, Expr, ExprNode);
  static constexpr bool _type_container_is_exact = false;
};
 
class PrimExpr : public TypedExpr<PrimType> {
 public:
  using TypedExpr<PrimType>::ty;
 
  TVM_TIRX_DLL PrimExpr(Call call);  // NOLINT(*)
 
  TVM_TIRX_DLL PrimExpr(int32_t value);  // NOLINT(*)
  TVM_TIRX_DLL PrimExpr(float value);  // NOLINT(*)
 
  explicit PrimExpr(ffi::ObjectPtr<ExprNode> node) : TypedExpr<PrimType>(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimExpr, TypedExpr<PrimType>, ExprNode);
  static constexpr bool _type_container_is_exact = false;
};
 
class IntExpr : public PrimExpr {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntExpr, PrimExpr, ExprNode);
  static constexpr bool _type_container_is_exact = false;
};
 
namespace ffi {
template <>
inline constexpr bool use_default_type_traits_v<PrimType> = false;
 
template <>
struct TypeTraits<PrimType> : public ObjectRefWithFallbackTraitsBase<PrimType, DLDataType> {
  TVM_FFI_INLINE static PrimType ConvertFallbackValue(DLDataType dtype) { return PrimType(dtype); }
};
 
template <typename ExpectedType>
inline constexpr bool use_default_type_traits_v<TypedExpr<ExpectedType>> = false;
 
template <typename ExpectedType>
struct TypeTraits<TypedExpr<ExpectedType>>
    : public ObjectRefTypeTraitsBase<TypedExpr<ExpectedType>> {
  TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
    if (src->type_index == TypeIndex::kTVMFFINone) {
      return TypedExpr<ExpectedType>::_type_is_nullable;
    }
    if (src->type_index < TypeIndex::kTVMFFIStaticObjectBegin ||
        !details::IsObjectInstance<ExprNode>(src->type_index)) {
      return false;
    }
    // Non-owning: this only reads `ty`, and the owning form's incref/decref pair costs two
    // atomics per check on a path every typed field assignment takes.
    const auto* expr = details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
    return details::AnyUnsafe::CheckAnyViewStrict<ExpectedType>(expr->ty);
  }
 
  TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
    if (src->type_index >= TypeIndex::kTVMFFIStaticObjectBegin && src->v_obj != nullptr &&
        details::IsObjectInstance<ExprNode>(src->type_index)) {
      const auto* expr = details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
      if (expr->ty.defined()) {
        const auto* prim_ty = expr->ty.as<PrimTypeNode>();
        std::string ty =
            prim_ty ? std::string(ffi::DLDataTypeToString(prim_ty->dtype)) : expr->ty.GetTypeKey();
        return TypeIndexToTypeKey(src->type_index) + "[ty=" + ty + "]";
      }
    }
    return TypeTraitsBase::GetMismatchTypeInfo(src);
  }
};
 
}  // namespace ffi
 
}  // namespace tvm
 
#endif  // TVM_IR_BASE_EXPR_H_