This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new b7555970 Consolidate expression headers and inline FFI conversion
traits (#10)
b7555970 is described below
commit b7555970a9aa35121c3b18354aebf0ef3bb7e04a
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 21 12:38:04 2026 -0400
Consolidate expression headers and inline FFI conversion traits (#10)
Place core type definitions in `type.h` and all expression definitions
in `expr.h`, removing the intermediate `base_expr.h` split. Define
expression FFI conversion traits alongside their complete node types and
remove their unnecessary exported functions and dead nullable branches.
Require explicit checked conversion from `Call` to `PrimExpr`,
preserving literal, tuple, string, and tensor-region conversion
behavior.
---
include/tvm/ir/base_expr.h | 596 -------------------------------
include/tvm/ir/expr.h | 340 ++++++++++++++++--
include/tvm/ir/type.h | 290 ++++++++++++++-
src/backend/cuda/transform/lower_iket.cc | 10 +-
src/ir/expr.cc | 39 --
src/tirx/ir/ir_mutator_with_analyzer.cc | 5 +-
src/tirx/ir/tensor.cc | 4 -
src/tirx/transform/vectorize_loop.cc | 3 +-
tests/cpp/ir/pattern_match_test.cc | 3 +-
9 files changed, 607 insertions(+), 683 deletions(-)
diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
deleted file mode 100644
index e7a15bc4..00000000
--- a/include/tvm/ir/base_expr.h
+++ /dev/null
@@ -1,596 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file tvm/ir/base_expr.h
- * \brief Base expression and primitive type nodes.
- */
-#ifndef TVM_IR_BASE_EXPR_H_
-#define TVM_IR_BASE_EXPR_H_
-
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/container/array.h>
-#include <tvm/ffi/container/tuple.h>
-#include <tvm/ffi/container/variant.h>
-#include <tvm/ffi/dtype.h>
-#include <tvm/ffi/extra/structural_mutate.h>
-#include <tvm/ffi/extra/structural_visit.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/optional.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 {
-
-/*!
- * \brief Type is the base type of all types.
- *
- * TVM's type system contains following subclasses:
- *
- * - PrimType: type of primitive type values used in the low-level IR.
- * - FuncType: type of a function.
- * - TensorType: type of certain Tensor values in the expression.
- *
- * There are also advanced types to support generic(polymorphic types).
- * \sa Type
- */
-class TypeNode : public ffi::Object {
- public:
- /*!
- * \brief Location that points to the original source code.
- * Reserved debug information.
- */
- 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::init(false),
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);
-};
-
-/*!
- * \brief Managed reference to TypeNode.
- * \sa TypeNode
- */
-class Type : public ffi::ObjectRef {
- public:
- /*! \brief Construct a MissingType for type information not yet populated. */
- TVM_TIRX_DLL static Type Missing();
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef,
TypeNode);
-};
-
-/*!
- * \brief Type information that has not been supplied or computed.
- *
- * MissingType is not a concrete type, wildcard, or inference variable. Unlike
- * AnyType and Void, it must be resolved before a fully typed IR boundary.
- */
-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);
-};
-
-/*! \brief Managed reference to the MissingTypeNode singleton. */
-class MissingType final : public Type {
- public:
- TVM_TIRX_DLL MissingType();
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MissingType, Type,
MissingTypeNode);
-};
-
-/*!
- * \brief Type marker for opaque construction-time expressions.
- *
- * Opaque values may be used while constructing IR, but must be lowered away
- * before the IR is considered complete.
- */
-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);
-};
-
-/*! \brief Managed reference to OpaqueTypeNode. */
-class OpaqueType final : public Type {
- public:
- TVM_TIRX_DLL OpaqueType();
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpaqueType, Type,
OpaqueTypeNode);
-};
-
-/*!
- * \brief Primitive data types used in the low-level IR.
- *
- * PrimType represents primitive POD values and the void sentinel.
- *
- * \sa PrimType
- */
-class PrimTypeNode final : public TypeNode {
- public:
- /*!
- * \brief The raw DLPack dtype represented by this primitive type.
- */
- 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:
- /*!
- * \brief Construct from a raw DLPack dtype.
- * \param dtype The corresponding DLPack dtype.
- */
- TVM_TIRX_DLL explicit PrimType(DLDataType dtype);
-
- /*!
- * \brief Construct from DLPack dtype fields.
- * \param code The DLPack dtype code.
- * \param bits The scalar bit width.
- * \param lanes The fixed lane count.
- */
- TVM_TIRX_DLL PrimType(DLDataTypeCode code, int bits, int lanes = 1);
-
- /*! \brief Construct a signed integer type with fixed lanes. */
- TVM_TIRX_DLL static PrimType Int(int bits, int lanes = 1);
- /*! \brief Construct an unsigned integer type with fixed lanes. */
- TVM_TIRX_DLL static PrimType UInt(int bits, int lanes = 1);
- /*! \brief Construct a floating-point type with fixed lanes. */
- TVM_TIRX_DLL static PrimType Float(int bits, int lanes = 1);
- /*! \brief Construct a bfloat type with fixed lanes. */
- TVM_TIRX_DLL static PrimType BFloat(int bits, int lanes = 1);
- /*! \brief Construct a boolean type with fixed lanes. */
- TVM_TIRX_DLL static PrimType Bool(int lanes = 1);
- /*! \brief Construct the void sentinel type, encoded as handle(0, 0). */
- TVM_TIRX_DLL static PrimType Void();
- /*!
- * \brief Construct a scalable vector type.
- * \param code The DLPack dtype code.
- * \param bits The scalar bit width.
- * \param lanes The positive vscale factor to encode in the DLPack lane
field.
- */
- TVM_TIRX_DLL static PrimType ScalableVector(DLDataTypeCode code, int bits,
int lanes);
-
- /*! \return The DLPack dtype code. */
- TVM_FFI_INLINE DLDataTypeCode code() const {
- return static_cast<DLDataTypeCode>(static_cast<int>(get()->dtype.code));
- }
-
- /*! \return The scalar bit width. */
- TVM_FFI_INLINE int32_t bits() const { return get()->dtype.bits; }
-
- /*!
- * \return The fixed lane count.
- * \note Throws on scalable vector types, where the encoded lane field
stores a vscale factor.
- */
- 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;
- }
-
- /*!
- * \brief Check the scalar element code and bit width.
- * \note Lane count and scalable-vector encoding are intentionally ignored.
- */
- 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;
- }
-
- /*!
- * \brief Check whether the dtype code matches any of the provided DLPack
codes.
- * \note Bit width and lanes are intentionally ignored.
- */
- 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)) || ...);
- }
-
- /*! \brief Whether this type is a scalar, excluding fixed and scalable
vectors. */
- TVM_FFI_INLINE bool IsScalar() const {
- int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
- return encoded_lanes == 1;
- }
-
- /*! \brief Whether this type is the void sentinel `handle(0, 0)`. */
- 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;
- }
-
- /*! \brief Whether this type is a scalable vector. */
- TVM_FFI_INLINE bool IsScalableVector() const {
- return static_cast<int16_t>(get()->dtype.lanes) < -1;
- }
-
- /*! \brief Whether this type is a fixed-length vector. */
- TVM_FFI_INLINE bool IsFixedLengthVector() const {
- return static_cast<int16_t>(get()->dtype.lanes) > 1;
- }
-
- /*!
- * \brief Return the number of bytes needed to store one value of this type.
- *
- * This uses the same packed sub-byte dtype sizing rule as runtime tensors.
- * Scalable vector types have no compile-time storage size and are rejected.
- */
- 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);
- }
-
- /*! \brief Return the same type with a different dtype code, preserving bits
and lanes. */
- 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);
- }
-
- /*! \brief Return the same type with a different scalar bit width,
preserving code and 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);
- }
-
- /*! \brief Return the same scalar element type with a fixed lane count. */
- TVM_FFI_INLINE PrimType WithLanes(int lanes) const {
- return PrimType(this->code(), this->bits(), lanes);
- }
-
- /*! \return The vscale factor encoded in a scalable vector type. */
- 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); }
-
-/*!
- * \brief Base type of all the expressions.
- * \sa Expr
- */
-class ExprNode : public ffi::Object {
- public:
- /*!
- * \brief Location that points to the original source code.
- * Reserved debug information.
- */
- mutable Location loc = UnknownLoc();
-
- /*!
- * \brief The deduced or annotated type of the expression.
- *
- * MissingType() denotes type information that will be populated by
- * later analysis passes instead of expression constructors.
- */
- 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);
-};
-
-/*!
- * \brief Managed reference to ExprNode.
- * \sa ExprNode
- */
-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);
-};
-
-/*! \brief Base node for traversable expressions eliminated during
compilation. */
-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);
-};
-
-/*! \brief Managed reference to a staging expression. */
-class StagingExpr : public Expr {
- public:
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StagingExpr, Expr,
StagingExprNode);
-};
-
-/*!
- * \brief Base node for opaque construction-time expressions.
- *
- * Subclasses are passed through by generic expression visitors and mutators.
- * They must not remain in finished IR.
- */
-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);
-};
-
-/*! \brief Managed reference to OpaqueExprNode. */
-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;
-
-/*!
- * \brief Typed reference/view over an expression whose result type is a
- * specific Type subtype.
- * \tparam ExpectedType The expected expression result type.
- */
-template <typename ExpectedType>
-class TypedExpr : public Expr {
- public:
- /*! \return the typed result of this expression. */
- 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;
-};
-
-/*!
- * \brief Typed reference/view over any Expr whose `ExprNode::ty` is PrimType.
- *
- * PrimExpr is a type category rather than a dedicated runtime node category.
- * It can contain intrinsic primitive nodes such as IntImmNode and
FloatImmNode,
- * or a general ExprNode such as CallNode, when that expression's `ty` field is
- * a PrimType. This keeps primitive-only APIs explicit while allowing shared
- * Expr nodes for cross-dialect values with richer result types when needed.
- */
-class PrimExpr : public TypedExpr<PrimType> {
- public:
- using TypedExpr<PrimType>::ty;
-
- /*!
- * \brief Construct from a call after checking that its result type is
- * PrimType.
- * \param call The call to view as a primitive expression.
- */
- TVM_TIRX_DLL PrimExpr(Call call); // NOLINT(*)
-
- /*!
- * \brief construct from integer.
- * \param value The value to be constructed.
- */
- TVM_TIRX_DLL PrimExpr(int32_t value); // NOLINT(*)
- /*!
- * \brief construct from float.
- * \param value The value to be constructed.
- */
- 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;
-};
-
-/*!
- * \brief View over a scalar signed or unsigned integer expression of any bit
width.
- *
- * IntExpr validates integer operands in operator signatures without changing
the
- * underlying expression node or its PrimExpr implementation API.
- */
-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);
- }
-};
-
-// Declare these traits before containers such as Array<Expr> and
SubscriptSlice instantiate them.
-// Conversion bodies that require literal or tuple nodes are defined in
expr.cc.
-// Generic expression arguments accept primitive/string literals and arrays of
expressions.
-template <>
-inline constexpr bool use_default_type_traits_v<Expr> = false;
-
-template <>
-struct TypeTraits<Expr> : public ObjectRefWithFallbackTraitsBase<Expr,
PrimExpr, ffi::String> {
- TVM_FFI_INLINE static Expr ConvertFallbackValue(PrimExpr value) { return
value; }
- TVM_TIRX_DLL static Expr ConvertFallbackValue(ffi::String value);
- TVM_TIRX_DLL static std::optional<Expr> TryCastFromAnyView(const TVMFFIAny*
src);
-};
-
-// Primitive expression arguments also accept bool, integer, and
floating-point literals.
-template <>
-inline constexpr bool use_default_type_traits_v<PrimExpr> = false;
-
-template <>
-struct TypeTraits<PrimExpr>
- : public ObjectRefWithFallbackTraitsBase<PrimExpr, StrictBool, int64_t,
double> {
- using Base = ObjectRefWithFallbackTraitsBase<PrimExpr, StrictBool, int64_t,
double>;
- TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
- if (src->type_index == TypeIndex::kTVMFFINone) return
PrimExpr::_type_is_nullable;
- return TypeTraits<TypedExpr<PrimType>>::CheckAnyStrict(src);
- }
-
- TVM_TIRX_DLL static std::optional<PrimExpr> TryCastFromAnyView(const
TVMFFIAny* src);
-
- TVM_FFI_INLINE static std::string TypeStr() { return "ir.PrimExpr"; }
-
- TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
- return TypeTraits<TypedExpr<PrimType>>::GetMismatchTypeInfo(src);
- }
-
- TVM_TIRX_DLL static PrimExpr ConvertFallbackValue(StrictBool value);
- TVM_TIRX_DLL static PrimExpr ConvertFallbackValue(int64_t value);
- TVM_TIRX_DLL static PrimExpr ConvertFallbackValue(double value);
-};
-
-} // namespace ffi
-
-/*! \brief Expression subscripts: point indices or (start, stop, step) slices.
*/
-using SubscriptSlice = ffi::Array<ffi::Variant<
- ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
- PrimExpr>>;
-
-namespace type_attr {
-/*!
- * \brief Realize a subscript of an expression, dispatched on the type of the
expression.
- *
- * The callback receives the expression, its point indices or slices, and the
source location.
- * It returns a defined object representing the realized subscription.
- */
-inline constexpr ffi::reflection::AttrKey<
- ffi::TypedFunction<ffi::ObjectRef(Expr, SubscriptSlice, Location)>>
- kExprSubscriptRealizeByTy{"__tvm_expr_subscript_realize__"};
-} // namespace type_attr
-
-} // namespace tvm
-
-#endif // TVM_IR_BASE_EXPR_H_
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index c19b5621..879f8168 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -25,30 +25,206 @@
#define TVM_IR_EXPR_H_
#include <tvm/ffi/big_int.h>
+#include <tvm/ffi/cast.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/container/variant.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/error.h>
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/native_function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/attrs.h>
-#include <tvm/ir/base_expr.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/location.h>
+#include <tvm/ir/type.h>
#include <algorithm>
+#include <cstddef>
#include <cstdint>
#include <functional>
#include <limits>
#include <optional>
#include <ostream>
#include <string>
+#include <type_traits>
#include <utility>
namespace tvm {
+/*!
+ * \brief Base type of all the expressions.
+ * \sa Expr
+ */
+class ExprNode : public ffi::Object {
+ public:
+ /*!
+ * \brief Location that points to the original source code.
+ * Reserved debug information.
+ */
+ mutable Location loc = UnknownLoc();
+
+ /*!
+ * \brief The deduced or annotated type of the expression.
+ *
+ * MissingType() denotes type information that will be populated by
+ * later analysis passes instead of expression constructors.
+ */
+ 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);
+};
+
+/*!
+ * \brief Managed reference to ExprNode.
+ * \sa ExprNode
+ */
+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);
+};
+
+/*! \brief Base node for traversable expressions eliminated during
compilation. */
+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);
+};
+
+/*! \brief Managed reference to a staging expression. */
+class StagingExpr : public Expr {
+ public:
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StagingExpr, Expr,
StagingExprNode);
+};
+
+/*!
+ * \brief Base node for opaque construction-time expressions.
+ *
+ * Subclasses are passed through by generic expression visitors and mutators.
+ * They must not remain in finished IR.
+ */
+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);
+};
+
+/*! \brief Managed reference to OpaqueExprNode. */
+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);
+};
+
+/*!
+ * \brief Typed reference/view over an expression whose result type is a
+ * specific Type subtype.
+ * \tparam ExpectedType The expected expression result type.
+ */
+template <typename ExpectedType>
+class TypedExpr : public Expr {
+ public:
+ /*! \return the typed result of this expression. */
+ 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;
+};
+
+/*!
+ * \brief Typed reference/view over any Expr whose `ExprNode::ty` is PrimType.
+ *
+ * PrimExpr is a type category rather than a dedicated runtime node category.
+ * It can contain intrinsic primitive nodes such as IntImmNode and
FloatImmNode,
+ * or a general ExprNode such as CallNode, when that expression's `ty` field is
+ * a PrimType. This keeps primitive-only APIs explicit while allowing shared
+ * Expr nodes for cross-dialect values with richer result types when needed.
+ */
+class PrimExpr : public TypedExpr<PrimType> {
+ public:
+ using TypedExpr<PrimType>::ty;
+
+ /*!
+ * \brief construct from integer.
+ * \param value The value to be constructed.
+ */
+ TVM_TIRX_DLL PrimExpr(int32_t value); // NOLINT(*)
+ /*!
+ * \brief construct from float.
+ * \param value The value to be constructed.
+ */
+ 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;
+};
+
+/*!
+ * \brief View over a scalar signed or unsigned integer expression of any bit
width.
+ *
+ * IntExpr validates integer operands in operator signatures without changing
the
+ * underlying expression node or its PrimExpr implementation API.
+ */
+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::details {
+// Array members require storage availability before the expression traits
below are complete.
+template <>
+inline constexpr bool storage_enabled_v<Expr> = true;
+template <>
+inline constexpr bool storage_enabled_v<PrimExpr> = true;
+} // namespace ffi::details
+
/*! \brief Base node for literal constants. */
class ConstantNode : public ExprNode {
public:
@@ -803,33 +979,6 @@ TVM_TIRX_DLL CallEffectKind SideEffect(const Expr& expr);
namespace ffi {
-template <>
-inline constexpr bool use_default_type_traits_v<IntExpr> = false;
-
-template <>
-struct TypeTraits<IntExpr> : public ObjectRefWithFallbackTraitsBase<IntExpr,
int64_t> {
- using Base = ObjectRefWithFallbackTraitsBase<IntExpr, int64_t>;
-
- TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
- if (!Base::CheckAnyStrict(src)) return false;
- if (src->type_index == TypeIndex::kTVMFFINone) return true;
- const auto* expr =
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
- const auto* ty = expr->ty.as<PrimTypeNode>();
- return ty != nullptr && ty->dtype.lanes == 1 &&
- (ty->dtype.code == kDLInt || ty->dtype.code == kDLUInt);
- }
-
- TVM_FFI_INLINE static IntExpr ConvertFallbackValue(int64_t value) {
- return
TypeTraits<PrimExpr>::ConvertFallbackValue(value).as_or_throw<IntExpr>();
- }
-
- TVM_FFI_INLINE static std::string TypeStr() { return "ir.IntExpr"; }
-
- TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
- return TypeTraits<TypedExpr<PrimType>>::GetMismatchTypeInfo(src);
- }
-};
-
template <>
inline constexpr bool use_default_type_traits_v<tvm::StringImm> = false;
@@ -874,6 +1023,114 @@ inline constexpr bool object_ref_contains_v<PrimExpr,
FloatImmNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, TensorLoadNode> = true;
+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::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);
+ }
+};
+
+// Primitive expression arguments also accept bool, integer, and
floating-point literals.
+template <>
+inline constexpr bool use_default_type_traits_v<PrimExpr> = false;
+
+template <>
+struct TypeTraits<PrimExpr>
+ : public ObjectRefWithFallbackTraitsBase<PrimExpr, StrictBool, int64_t,
double> {
+ TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
+ return TypeTraits<TypedExpr<PrimType>>::CheckAnyStrict(src);
+ }
+
+ TVM_FFI_INLINE static std::string TypeStr() { return "ir.PrimExpr"; }
+
+ TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
+ return TypeTraits<TypedExpr<PrimType>>::GetMismatchTypeInfo(src);
+ }
+
+ TVM_FFI_INLINE static PrimExpr ConvertFallbackValue(StrictBool value) {
+ return IntImm::Bool(value);
+ }
+ TVM_FFI_INLINE static PrimExpr ConvertFallbackValue(int64_t value) {
+ return TypeTraits<IntImm>::ConvertFallbackValue(value);
+ }
+ TVM_FFI_INLINE static PrimExpr ConvertFallbackValue(double value) {
+ return TypeTraits<FloatImm>::ConvertFallbackValue(value);
+ }
+};
+
+// Generic expression arguments accept primitive/string literals and arrays of
expressions.
+template <>
+inline constexpr bool use_default_type_traits_v<Expr> = false;
+
+template <>
+struct TypeTraits<Expr> : public ObjectRefWithFallbackTraitsBase<Expr,
PrimExpr, ffi::String> {
+ TVM_FFI_INLINE static Expr ConvertFallbackValue(PrimExpr value) { return
value; }
+ TVM_FFI_INLINE static Expr ConvertFallbackValue(ffi::String value) {
+ return StringImm(std::move(value));
+ }
+ static std::optional<Expr> TryCastFromAnyView(const TVMFFIAny* src) {
+ if (auto value =
+ ObjectRefWithFallbackTraitsBase<Expr, PrimExpr,
ffi::String>::TryCastFromAnyView(src)) {
+ return value;
+ }
+ if (auto fields = TypeTraits<ffi::Array<Expr>>::TryCastFromAnyView(src)) {
+ return tvm::Tuple(std::move(*fields));
+ }
+ return std::nullopt;
+ }
+};
+
+template <>
+inline constexpr bool use_default_type_traits_v<IntExpr> = false;
+
+template <>
+struct TypeTraits<IntExpr> : public ObjectRefWithFallbackTraitsBase<IntExpr,
int64_t> {
+ using Base = ObjectRefWithFallbackTraitsBase<IntExpr, int64_t>;
+
+ TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
+ if (!Base::CheckAnyStrict(src)) return false;
+ const auto* expr =
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
+ const auto* ty = expr->ty.as<PrimTypeNode>();
+ return ty != nullptr && ty->dtype.lanes == 1 &&
+ (ty->dtype.code == kDLInt || ty->dtype.code == kDLUInt);
+ }
+
+ TVM_FFI_INLINE static IntExpr ConvertFallbackValue(int64_t value) {
+ return
TypeTraits<PrimExpr>::ConvertFallbackValue(value).as_or_throw<IntExpr>();
+ }
+
+ TVM_FFI_INLINE static std::string TypeStr() { return "ir.IntExpr"; }
+
+ TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
+ return TypeTraits<TypedExpr<PrimType>>::GetMismatchTypeInfo(src);
+ }
+};
+
template <>
inline constexpr bool use_default_type_traits_v<PrimVar> = false;
@@ -889,9 +1146,6 @@ struct TypeTraits<PrimVar> : public
ObjectRefTypeTraitsBase<PrimVar> {
using Base::TypeStr;
TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
- if (src->type_index == TypeIndex::kTVMFFINone) {
- return PrimVar::_type_is_nullable;
- }
if (src->type_index != VarNode::RuntimeTypeIndex()) {
return false;
}
@@ -902,9 +1156,6 @@ struct TypeTraits<PrimVar> : public
ObjectRefTypeTraitsBase<PrimVar> {
TVM_FFI_INLINE static std::optional<PrimVar> TryCastFromAnyView(const
TVMFFIAny* src) {
if (CheckAnyStrict(src)) {
- if (src->type_index == TypeIndex::kTVMFFINone) {
- return details::ObjectUnsafe::ObjectRefFromObjectPtr<PrimVar>(nullptr);
- }
return details::ObjectUnsafe::ObjectRefFromObjectPtr<PrimVar>(
details::ObjectUnsafe::ObjectPtrFromUnowned<VarNode>(src->v_obj));
}
@@ -917,11 +1168,30 @@ inline constexpr bool
use_default_type_traits_v<TensorRegion> = false;
template <>
struct TypeTraits<TensorRegion> : public
ObjectRefWithFallbackTraitsBase<TensorRegion, TensorLoad> {
- TVM_TIRX_DLL static TensorRegion ConvertFallbackValue(TensorLoad value);
+ TVM_FFI_INLINE static TensorRegion ConvertFallbackValue(TensorLoad value) {
+ return TensorRegion::FromTensorLoad(std::move(value));
+ }
};
} // namespace ffi
+/*! \brief Expression subscripts: point indices or (start, stop, step) slices.
*/
+using SubscriptSlice = ffi::Array<ffi::Variant<
+ ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
+ PrimExpr>>;
+
+namespace type_attr {
+/*!
+ * \brief Realize a subscript of an expression, dispatched on the type of the
expression.
+ *
+ * The callback receives the expression, its point indices or slices, and the
source location.
+ * It returns a defined object representing the realized subscription.
+ */
+inline constexpr ffi::reflection::AttrKey<
+ ffi::TypedFunction<ffi::ObjectRef(Expr, SubscriptSlice, Location)>>
+ kExprSubscriptRealizeByTy{"__tvm_expr_subscript_realize__"};
+} // namespace type_attr
+
} // namespace tvm
/* \brief Allow tvm.Var and tvm.GlobalVar as keys in STL tables
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index b961b975..cada269a 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -27,7 +27,6 @@
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/base_expr.h>
#include <tvm/ir/location.h>
#include <tvm/runtime/memory_scope.h>
@@ -35,6 +34,295 @@
namespace tvm {
+/*!
+ * \brief Type is the base type of all types.
+ *
+ * TVM's type system contains following subclasses:
+ *
+ * - PrimType: type of primitive type values used in the low-level IR.
+ * - FuncType: type of a function.
+ * - TensorType: type of certain Tensor values in the expression.
+ *
+ * There are also advanced types to support generic(polymorphic types).
+ * \sa Type
+ */
+class TypeNode : public ffi::Object {
+ public:
+ /*!
+ * \brief Location that points to the original source code.
+ * Reserved debug information.
+ */
+ 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::init(false),
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);
+};
+
+/*!
+ * \brief Managed reference to TypeNode.
+ * \sa TypeNode
+ */
+class Type : public ffi::ObjectRef {
+ public:
+ /*! \brief Construct a MissingType for type information not yet populated. */
+ TVM_TIRX_DLL static Type Missing();
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef,
TypeNode);
+};
+
+/*!
+ * \brief Type information that has not been supplied or computed.
+ *
+ * MissingType is not a concrete type, wildcard, or inference variable. Unlike
+ * AnyType and Void, it must be resolved before a fully typed IR boundary.
+ */
+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);
+};
+
+/*! \brief Managed reference to the MissingTypeNode singleton. */
+class MissingType final : public Type {
+ public:
+ TVM_TIRX_DLL MissingType();
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MissingType, Type,
MissingTypeNode);
+};
+
+/*!
+ * \brief Type marker for opaque construction-time expressions.
+ *
+ * Opaque values may be used while constructing IR, but must be lowered away
+ * before the IR is considered complete.
+ */
+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);
+};
+
+/*! \brief Managed reference to OpaqueTypeNode. */
+class OpaqueType final : public Type {
+ public:
+ TVM_TIRX_DLL OpaqueType();
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpaqueType, Type,
OpaqueTypeNode);
+};
+
+/*!
+ * \brief Primitive data types used in the low-level IR.
+ *
+ * PrimType represents primitive POD values and the void sentinel.
+ *
+ * \sa PrimType
+ */
+class PrimTypeNode final : public TypeNode {
+ public:
+ /*!
+ * \brief The raw DLPack dtype represented by this primitive type.
+ */
+ 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:
+ /*!
+ * \brief Construct from a raw DLPack dtype.
+ * \param dtype The corresponding DLPack dtype.
+ */
+ TVM_TIRX_DLL explicit PrimType(DLDataType dtype);
+
+ /*!
+ * \brief Construct from DLPack dtype fields.
+ * \param code The DLPack dtype code.
+ * \param bits The scalar bit width.
+ * \param lanes The fixed lane count.
+ */
+ TVM_TIRX_DLL PrimType(DLDataTypeCode code, int bits, int lanes = 1);
+
+ /*! \brief Construct a signed integer type with fixed lanes. */
+ TVM_TIRX_DLL static PrimType Int(int bits, int lanes = 1);
+ /*! \brief Construct an unsigned integer type with fixed lanes. */
+ TVM_TIRX_DLL static PrimType UInt(int bits, int lanes = 1);
+ /*! \brief Construct a floating-point type with fixed lanes. */
+ TVM_TIRX_DLL static PrimType Float(int bits, int lanes = 1);
+ /*! \brief Construct a bfloat type with fixed lanes. */
+ TVM_TIRX_DLL static PrimType BFloat(int bits, int lanes = 1);
+ /*! \brief Construct a boolean type with fixed lanes. */
+ TVM_TIRX_DLL static PrimType Bool(int lanes = 1);
+ /*! \brief Construct the void sentinel type, encoded as handle(0, 0). */
+ TVM_TIRX_DLL static PrimType Void();
+ /*!
+ * \brief Construct a scalable vector type.
+ * \param code The DLPack dtype code.
+ * \param bits The scalar bit width.
+ * \param lanes The positive vscale factor to encode in the DLPack lane
field.
+ */
+ TVM_TIRX_DLL static PrimType ScalableVector(DLDataTypeCode code, int bits,
int lanes);
+
+ /*! \return The DLPack dtype code. */
+ TVM_FFI_INLINE DLDataTypeCode code() const {
+ return static_cast<DLDataTypeCode>(static_cast<int>(get()->dtype.code));
+ }
+
+ /*! \return The scalar bit width. */
+ TVM_FFI_INLINE int32_t bits() const { return get()->dtype.bits; }
+
+ /*!
+ * \return The fixed lane count.
+ * \note Throws on scalable vector types, where the encoded lane field
stores a vscale factor.
+ */
+ 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;
+ }
+
+ /*!
+ * \brief Check the scalar element code and bit width.
+ * \note Lane count and scalable-vector encoding are intentionally ignored.
+ */
+ 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;
+ }
+
+ /*!
+ * \brief Check whether the dtype code matches any of the provided DLPack
codes.
+ * \note Bit width and lanes are intentionally ignored.
+ */
+ 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)) || ...);
+ }
+
+ /*! \brief Whether this type is a scalar, excluding fixed and scalable
vectors. */
+ TVM_FFI_INLINE bool IsScalar() const {
+ int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
+ return encoded_lanes == 1;
+ }
+
+ /*! \brief Whether this type is the void sentinel `handle(0, 0)`. */
+ 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;
+ }
+
+ /*! \brief Whether this type is a scalable vector. */
+ TVM_FFI_INLINE bool IsScalableVector() const {
+ return static_cast<int16_t>(get()->dtype.lanes) < -1;
+ }
+
+ /*! \brief Whether this type is a fixed-length vector. */
+ TVM_FFI_INLINE bool IsFixedLengthVector() const {
+ return static_cast<int16_t>(get()->dtype.lanes) > 1;
+ }
+
+ /*!
+ * \brief Return the number of bytes needed to store one value of this type.
+ *
+ * This uses the same packed sub-byte dtype sizing rule as runtime tensors.
+ * Scalable vector types have no compile-time storage size and are rejected.
+ */
+ 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);
+ }
+
+ /*! \brief Return the same type with a different dtype code, preserving bits
and lanes. */
+ 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);
+ }
+
+ /*! \brief Return the same type with a different scalar bit width,
preserving code and 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);
+ }
+
+ /*! \brief Return the same scalar element type with a fixed lane count. */
+ TVM_FFI_INLINE PrimType WithLanes(int lanes) const {
+ return PrimType(this->code(), this->bits(), lanes);
+ }
+
+ /*! \return The vscale factor encoded in a scalable vector type. */
+ 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); }
+
+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); }
+};
+
+} // namespace ffi
+
/*! \brief The type whose only value is None. */
class NoneTypeNode : public TypeNode {
public:
diff --git a/src/backend/cuda/transform/lower_iket.cc
b/src/backend/cuda/transform/lower_iket.cc
index 70e395b4..aa056263 100644
--- a/src/backend/cuda/transform/lower_iket.cc
+++ b/src/backend/cuda/transform/lower_iket.cc
@@ -1097,14 +1097,16 @@ class InstrumentOfficialKernel : public StmtExprMutator
{
PrimExpr Event(PrimExpr event_id) const {
static const Op event_op = Op::Get("tirx.cuda.iket_official_event");
return Call(PrimType::UInt(32), event_op,
- {prim::cast(event_id, PrimType::UInt(32)),
StringImm(device_source_)});
+ {prim::cast(event_id, PrimType::UInt(32)),
StringImm(device_source_)})
+ .as_or_throw<PrimExpr>();
}
PrimExpr Event(PrimExpr event_id, PrimExpr payload) const {
static const Op event_op = Op::Get("tirx.cuda.iket_official_event");
- return Call(
- PrimType::UInt(32), event_op,
- {prim::cast(event_id, PrimType::UInt(32)), StringImm(device_source_),
std::move(payload)});
+ return Call(PrimType::UInt(32), event_op,
+ {prim::cast(event_id, PrimType::UInt(32)),
StringImm(device_source_),
+ std::move(payload)})
+ .as_or_throw<PrimExpr>();
}
PrimExpr NormalizePayload(PrimExpr payload, PayloadType type) const {
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index a23aa87f..3ac973f7 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -39,43 +39,6 @@
namespace tvm {
-namespace ffi {
-
-Expr TypeTraits<Expr>::ConvertFallbackValue(ffi::String value) {
- return StringImm(std::move(value));
-}
-
-std::optional<Expr> TypeTraits<Expr>::TryCastFromAnyView(const TVMFFIAny* src)
{
- if (auto value =
- ObjectRefWithFallbackTraitsBase<Expr, PrimExpr,
ffi::String>::TryCastFromAnyView(src)) {
- return value;
- }
- if (auto fields = TypeTraits<ffi::Array<Expr>>::TryCastFromAnyView(src)) {
- return tvm::Tuple(std::move(*fields));
- }
- return std::nullopt;
-}
-
-std::optional<PrimExpr> TypeTraits<PrimExpr>::TryCastFromAnyView(const
TVMFFIAny* src) {
- // A required expression cannot accept None.
- if (src->type_index == TypeIndex::kTVMFFINone) return std::nullopt;
- return Base::TryCastFromAnyView(src);
-}
-
-PrimExpr TypeTraits<PrimExpr>::ConvertFallbackValue(StrictBool value) {
- return IntImm::Bool(value);
-}
-
-PrimExpr TypeTraits<PrimExpr>::ConvertFallbackValue(int64_t value) {
- return TypeTraits<IntImm>::ConvertFallbackValue(value);
-}
-
-PrimExpr TypeTraits<PrimExpr>::ConvertFallbackValue(double value) {
- return TypeTraits<FloatImm>::ConvertFallbackValue(value);
-}
-
-} // namespace ffi
-
namespace {
template <typename TNode>
@@ -913,8 +876,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-PrimExpr::PrimExpr(Call call) :
PrimExpr(std::move(call).as_or_throw<PrimExpr>()) {}
-
PrimExpr::PrimExpr(int32_t value) : PrimExpr(IntImm::Int32(value)) {}
PrimExpr::PrimExpr(float value) : PrimExpr(FloatImm(PrimType::Float(32),
value)) {}
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc
b/src/tirx/ir/ir_mutator_with_analyzer.cc
index 4101094c..fd1e22bb 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.cc
+++ b/src/tirx/ir/ir_mutator_with_analyzer.cc
@@ -234,8 +234,9 @@ UnchangedOr<Stmt>
StmtExprMutatorWithAnalyzer::Mutate_(const RegionStmtNode* op,
if (field.rfind("grid.", 0) == 0) intrinsic = "tirx.cuda.block_idx";
if (field.rfind("block.", 0) == 0) intrinsic = "tirx.cuda.thread_idx";
if (intrinsic.empty()) continue;
- PrimExpr index = Call(PrimType::Int(32), Op::Get(intrinsic),
- {StringImm(field.substr(field.size() - 1))});
+ PrimExpr index =
+ Call(PrimType::Int(32), Op::Get(intrinsic),
{StringImm(field.substr(field.size() - 1))})
+ .as_or_throw<PrimExpr>();
constraint_scope_.Current().Emplace(analyzer_, index >= 0);
constraint_scope_.Current().Emplace(analyzer_, index <
op->args[i].as_or_throw<PrimExpr>());
}
diff --git a/src/tirx/ir/tensor.cc b/src/tirx/ir/tensor.cc
index 2537e867..e8e7b3bc 100644
--- a/src/tirx/ir/tensor.cc
+++ b/src/tirx/ir/tensor.cc
@@ -48,10 +48,6 @@ TensorRegion TensorRegion::FromTensorLoad(TensorLoad value) {
value->loc);
}
-TensorRegion ffi::TypeTraits<TensorRegion>::ConvertFallbackValue(TensorLoad
value) {
- return TensorRegion::FromTensorLoad(std::move(value));
-}
-
namespace tirx {
using namespace tvm::prim;
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index d21ffde0..b0428b2b 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -299,7 +299,8 @@ class TryPredicateTensorAccesses : public StmtExprMutator {
ffi::Array<Expr> args{load->source.as_or_throw<tvm::tirx::TensorVar>()};
for (const PrimExpr& index : load->indices) args.push_back(index);
args.push_back(mask.value());
- return Call(load->ty, tirx::masked_load_op(), args, {}, {}, load->loc);
+ return Call(load->ty, tirx::masked_load_op(), args, {}, {}, load->loc)
+ .as_or_throw<PrimExpr>();
}
return load;
}
diff --git a/tests/cpp/ir/pattern_match_test.cc
b/tests/cpp/ir/pattern_match_test.cc
index 24477477..d2c7073e 100644
--- a/tests/cpp/ir/pattern_match_test.cc
+++ b/tests/cpp/ir/pattern_match_test.cc
@@ -24,7 +24,8 @@
TEST(Pattern, Basic) {
using namespace tvm;
tvm::PrimVar x("x"), y("y"), z("z");
- PrimExpr scalable_lanes = prim::Mul(Call(PrimType::Int(32),
prim::vscale_op(), {}), 4);
+ PrimExpr scalable_lanes =
+ prim::Mul(Call(PrimType::Int(32), prim::vscale_op(),
{}).as_or_throw<PrimExpr>(), 4);
tvm::PVar<PrimExpr> px, py, pz;
tvm::PVar<DLDataType> pt;
tvm::PVar<PrimExpr> planes;
