This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit a44cbc5c545c454a0a1d66a198a631faae1cddd4 Author: Tianqi Chen <[email protected]> AuthorDate: Mon Oct 19 12:40:44 2026 -0400 [REFACTOR][IR] Colocate typed attribute contracts (#20795) Place expression, tensor access, region, and operator attributes beside their owning APIs. Share the typed subscript contract across dialects, use native tensor-store validation, and remove redundant aliases and Evaluate's type hook. Keep primitive operators in prim/op.h and organize expression traits by type, with foundational conversion contracts available from base_expr.h. --- tirx/include/tvm/ir/base_expr.h | 61 +++ tirx/include/tvm/ir/expr.h | 443 ++++----------------- tirx/include/tvm/ir/op.h | 32 +- tirx/include/tvm/ir/prim/op.h | 242 ++++++++++- tirx/include/tvm/ir/stmt.h | 60 +-- tirx/include/tvm/script/printer/doc_translator.h | 6 +- tirx/include/tvm/target/target.h | 12 +- tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc | 10 +- tirx/src/backend/cuda/codegen/target_kind.cc | 2 +- tirx/src/backend/cuda/op/iket.cc | 14 +- tirx/src/backend/cuda/op/legacy.cc | 8 +- tirx/src/backend/cuda/op/target_builtin.cc | 16 +- tirx/src/backend/cuda/op/tcgen05.cc | 6 +- tirx/src/backend/cuda/op/tensormap.cc | 2 +- tirx/src/backend/cuda/tile/exec_context.cc | 1 + .../src/backend/metal/codegen/intrin_rule_metal.cc | 6 +- tirx/src/backend/metal/op/cooperative_tensor.cc | 8 +- tirx/src/backend/metal/op/target_builtin.cc | 8 +- .../backend/opencl/codegen/intrin_rule_opencl.cc | 1 + tirx/src/backend/opencl/op/memory.cc | 6 +- tirx/src/backend/trn/op/target_builtin.cc | 2 +- .../backend/webgpu/codegen/intrin_rule_webgpu.cc | 6 +- tirx/src/ir/expr.cc | 43 +- tirx/src/ir/prim/builtin.cc | 13 +- tirx/src/ir/prim/math.cc | 76 ++-- tirx/src/ir/side_effect.cc | 3 +- tirx/src/ir/stmt.cc | 18 +- tirx/src/ir/subscript_expr.cc | 22 +- tirx/src/script/printer/ir/call.cc | 5 +- tirx/src/tirx/ir/stmt.cc | 13 +- tirx/src/tirx/ir/tensor.cc | 20 +- tirx/src/tirx/ir/tensor_load.cc | 5 +- tirx/src/tirx/op/abi.cc | 36 +- tirx/src/tirx/op/annotation.cc | 4 +- tirx/src/tirx/op/cpu.cc | 2 +- tirx/src/tirx/op/gpu.cc | 32 +- tirx/src/tirx/op/memory.cc | 31 +- tirx/src/tirx/op/vector.cc | 8 +- tirx/tests/cpp/ir/tensor_load_test.cc | 4 +- tirx/tests/cpp/target/target_test.cc | 4 +- .../python/tirx/script/test_error_handling.py | 9 +- 41 files changed, 659 insertions(+), 641 deletions(-) diff --git a/tirx/include/tvm/ir/base_expr.h b/tirx/include/tvm/ir/base_expr.h index 380a10c0ae..e7a15bc421 100644 --- a/tirx/include/tvm/ir/base_expr.h +++ b/tirx/include/tvm/ir/base_expr.h @@ -25,9 +25,14 @@ #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> @@ -528,8 +533,64 @@ struct TypeTraits<TypedExpr<ExpectedType>> } }; +// 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/tirx/include/tvm/ir/expr.h b/tirx/include/tvm/ir/expr.h index add9df21c8..c19b5621a3 100644 --- a/tirx/include/tvm/ir/expr.h +++ b/tirx/include/tvm/ir/expr.h @@ -249,119 +249,6 @@ class FloatImm : public PrimExpr { TVM_DEFINE_OBJECT_REF_COW_METHOD(FloatImmNode); }; -namespace ffi { -// Type traits to enable automatic conversion into IntImm, Integer, and Bool -// when called through the FFI -template <> -inline constexpr bool use_default_type_traits_v<IntImm> = false; - -// specialize to enable implicit conversion from const char* -template <> -struct TypeTraits<IntImm> : public ObjectRefWithFallbackTraitsBase<IntImm, int64_t> { - TVM_FFI_INLINE static IntImm ConvertFallbackValue(int64_t value) { - auto value_ty = - (value > std::numeric_limits<int>::max() || value < std::numeric_limits<int>::min()) - ? PrimType::Int(64) - : PrimType::Int(32); - return IntImm(value_ty, value); - } -}; - -template <> -inline constexpr bool use_default_type_traits_v<FloatImm> = false; - -template <> -struct TypeTraits<FloatImm> : public ObjectRefWithFallbackTraitsBase<FloatImm, double> { - TVM_FFI_INLINE static FloatImm ConvertFallbackValue(double value) { - return FloatImm(PrimType::Float(32), value); - } -}; -template <> -inline constexpr bool use_default_type_traits_v<tvm::StringImm> = false; - -template <> -struct TypeTraits<tvm::StringImm> - : public ObjectRefWithFallbackTraitsBase<tvm::StringImm, ffi::String> { - TVM_FFI_INLINE static tvm::StringImm ConvertFallbackValue(ffi::String value) { - return tvm::StringImm(value); - } -}; - -// define automatic conversion from bool, int64_t, double to PrimExpr -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_FFI_INLINE static std::optional<PrimExpr> TryCastFromAnyView(const TVMFFIAny* src) { - // A required expression cannot accept None. - if (src->type_index == TypeIndex::kTVMFFINone) return std::nullopt; - return Base::TryCastFromAnyView(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); - } -}; - -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); - } -}; - -// Generic Expr 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(value); } - static std::optional<Expr> TryCastFromAnyView(const TVMFFIAny* src); -}; -} // namespace ffi - /*! \brief Tuple container */ class TupleNode : public ExprNode { public: @@ -392,18 +279,6 @@ class Tuple : public Expr { TVM_DEFINE_OBJECT_REF_COW_METHOD(TupleNode); }; -// Use ordinary inline: array elements recursively invoke this conversion. -inline std::optional<Expr> ffi::TypeTraits<Expr>::TryCastFromAnyView(const TVMFFIAny* src) { - if (auto value = - ObjectRefWithFallbackTraitsBase<Expr, PrimExpr, ffi::String>::TryCastFromAnyView(src)) { - return value; - } - if (auto fields = ffi::TypeTraits<ffi::Array<Expr>>::TryCastFromAnyView(src)) { - return tvm::Tuple(std::move(*fields)); - } - return std::nullopt; -} - /*! \brief Get the index-th field out of a tuple. */ class TupleGetItemNode : public ExprNode { public: @@ -463,11 +338,15 @@ class TensorLoadNode : public ExprNode { TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TensorLoad", TensorLoadNode, ExprNode); }; +namespace type_attr { /*! * \brief Infer a load's primitive result type from its populated source, indices and location. * \note The callback must derive the result from those inputs, not the unpublished result type. */ -using FTensorLoadInferType = ffi::reflection::NativeFunctionView<PrimType(const TensorLoadNode*)>; +inline constexpr ffi::reflection::AttrKey< + ffi::reflection::NativeFunctionView<PrimType(const TensorLoadNode*)>> + kTensorLoadInferType{"__tvm_tensor_load_infer_type__"}; +} // namespace type_attr /*! \brief Managed reference to TensorLoadNode. */ class TensorLoad : public PrimExpr { @@ -487,220 +366,6 @@ class TensorLoad : public PrimExpr { TVM_DEFINE_OBJECT_REF_COW_METHOD(TensorLoadNode); }; -/*! - * \brief add operator - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator+(PrimExpr a, PrimExpr b); - -/*! - * \brief subtraction operator - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator-(PrimExpr a, PrimExpr b); - -/*! - * \brief negation. - * - * \param a input. - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator-(PrimExpr a); - -/*! - * \brief multiplication operator - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator*(PrimExpr a, PrimExpr b); - -/*! - * \brief division operator - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator/(PrimExpr a, PrimExpr b); - -/*! - * \brief left shift operator - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b); - -/*! - * \brief right shift operator - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b); - -/*! - * \brief greater - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator>(PrimExpr a, PrimExpr b); - -/*! - * \brief greater_equal - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b); - -/*! - * \brief less - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator<(PrimExpr a, PrimExpr b); - -/*! - * \brief less_equal - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b); - -/*! - * \brief equal - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator==(PrimExpr a, PrimExpr b); - -/*! - * \brief not_equal - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b); - -/*! - * \brief and - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note This operator does eager constant folding. - */ -TVM_TIRX_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b); - -/*! - * \brief or - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note This operator does eager constant folding. - */ -TVM_TIRX_DLL PrimExpr operator||(PrimExpr a, PrimExpr b); - -/*! - * \brief not - * - * \param a left operand - * \return The result expression. - * \note This operator does eager constant folding. - */ -TVM_TIRX_DLL PrimExpr operator!(PrimExpr a); - -/*! - * \brief take bitwise and of two values - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator&(PrimExpr a, PrimExpr b); - -/*! - * \brief take bitwise or of two values - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator|(PrimExpr a, PrimExpr b); - -/*! - * \brief take bitwise xor of two values - * - * \param a left operand - * \param b right operand - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator^(PrimExpr a, PrimExpr b); - -/*! - * \brief take bitwise negation of two values - * - * \param a the input expression. - * \return The result expression. - * \note this function does eager constant folding for - * index types(int32, int64) when possible. - */ -TVM_TIRX_DLL PrimExpr operator~(PrimExpr a); - /*! * \brief A local variable in the IR. * @@ -1008,6 +673,11 @@ class Call : public Expr { TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode); }; +namespace op_attr { +/*! \brief The runtime effect of an operator call. */ +inline constexpr ffi::reflection::AttrKey<CallEffectKind> kCallEffectKind{"TCallEffectKind"}; +} // namespace op_attr + /*! \brief range over one dimension */ class RangeNode : public ffi::Object { public: @@ -1134,13 +804,76 @@ TVM_TIRX_DLL CallEffectKind SideEffect(const Expr& expr); namespace ffi { template <> -inline constexpr bool use_default_type_traits_v<TensorRegion> = false; +inline constexpr bool use_default_type_traits_v<IntExpr> = false; template <> -struct TypeTraits<TensorRegion> : public ObjectRefWithFallbackTraitsBase<TensorRegion, TensorLoad> { - TVM_TIRX_DLL static TensorRegion ConvertFallbackValue(TensorLoad value); +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; + +template <> +struct TypeTraits<tvm::StringImm> + : public ObjectRefWithFallbackTraitsBase<tvm::StringImm, ffi::String> { + TVM_FFI_INLINE static tvm::StringImm ConvertFallbackValue(ffi::String value) { + return tvm::StringImm(value); + } +}; + +template <> +inline constexpr bool use_default_type_traits_v<IntImm> = false; + +template <> +struct TypeTraits<IntImm> : public ObjectRefWithFallbackTraitsBase<IntImm, int64_t> { + TVM_FFI_INLINE static IntImm ConvertFallbackValue(int64_t value) { + auto value_ty = + (value > std::numeric_limits<int>::max() || value < std::numeric_limits<int>::min()) + ? PrimType::Int(64) + : PrimType::Int(32); + return IntImm(value_ty, value); + } +}; + +template <> +inline constexpr bool object_ref_contains_v<PrimExpr, IntImmNode> = true; + +template <> +inline constexpr bool use_default_type_traits_v<FloatImm> = false; + +template <> +struct TypeTraits<FloatImm> : public ObjectRefWithFallbackTraitsBase<FloatImm, double> { + TVM_FFI_INLINE static FloatImm ConvertFallbackValue(double value) { + return FloatImm(PrimType::Float(32), value); + } +}; + +template <> +inline constexpr bool object_ref_contains_v<PrimExpr, FloatImmNode> = true; + +template <> +inline constexpr bool object_ref_contains_v<PrimExpr, TensorLoadNode> = true; + template <> inline constexpr bool use_default_type_traits_v<PrimVar> = false; @@ -1180,29 +913,15 @@ struct TypeTraits<PrimVar> : public ObjectRefTypeTraitsBase<PrimVar> { }; template <> -inline constexpr bool object_ref_contains_v<PrimExpr, IntImmNode> = true; -template <> -inline constexpr bool object_ref_contains_v<PrimExpr, FloatImmNode> = true; +inline constexpr bool use_default_type_traits_v<TensorRegion> = false; + template <> -inline constexpr bool object_ref_contains_v<PrimExpr, TensorLoadNode> = true; +struct TypeTraits<TensorRegion> : public ObjectRefWithFallbackTraitsBase<TensorRegion, TensorLoad> { + TVM_TIRX_DLL static TensorRegion ConvertFallbackValue(TensorLoad value); +}; } // namespace ffi -namespace op_attr { -/*! \brief Use integer to record the kind. */ -inline constexpr ffi::reflection::AttrKey<int64_t> kCallEffectKind{"TCallEffectKind"}; -} // namespace op_attr - -// Compatibility names; new callers can deduce the value type from the key. -using TCallEffectKind = decltype(op_attr::kCallEffectKind)::ValueType; - -namespace type_attr { -inline constexpr ffi::reflection::AttrKey<FTensorLoadInferType> kTensorLoadInferType{ - "__tvm_tensor_load_infer_type__"}; -inline constexpr ffi::reflection::AttrKey<ffi::Function> kSubscriptExprRealize{ - "__tvm_subscript_expr_realize__"}; -} // namespace type_attr - } // namespace tvm /* \brief Allow tvm.Var and tvm.GlobalVar as keys in STL tables diff --git a/tirx/include/tvm/ir/op.h b/tirx/include/tvm/ir/op.h index 9ae353a8c5..4b77eb316b 100644 --- a/tirx/include/tvm/ir/op.h +++ b/tirx/include/tvm/ir/op.h @@ -199,6 +199,22 @@ class Op : public Expr { TVM_TIRX_DLL static ffi::List<ffi::Any> GetAttrColumn(const ffi::String& attr_name); }; +namespace op_attr { +/*! \brief Infer a Call's result type from its explicit inputs without builder state. */ +inline constexpr ffi::reflection::AttrKey< + ffi::reflection::NativeFunctionView<Type(const CallNode* call)>> + kInferType{"FInferType"}; +/*! + * \brief An operator's result type independent of operands, attributes, and type arguments. + * When present, this concrete type takes precedence over 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 { /*! \brief Base counts for descriptors accepted by OpDef::signature. */ @@ -743,21 +759,5 @@ OpAttrMap<ValueType> Op::GetAttrMap(ffi::reflection::AttrKey<ValueType> key) { return GetAttrMap<ValueType>(key.name); } -namespace op_attr { -/*! \brief Infer a Call's result type from its explicit inputs without builder state. */ -inline constexpr ffi::reflection::AttrKey< - ffi::reflection::NativeFunctionView<Type(const CallNode* call)>> - kInferType{"FInferType"}; -/*! - * \brief An operator's result type independent of operands, attributes, and type arguments. - * When present, this concrete type takes precedence over 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 tvm #endif // TVM_IR_OP_H_ diff --git a/tirx/include/tvm/ir/prim/op.h b/tirx/include/tvm/ir/prim/op.h index 1d059dad9b..f1eca89c4e 100644 --- a/tirx/include/tvm/ir/prim/op.h +++ b/tirx/include/tvm/ir/prim/op.h @@ -30,6 +30,239 @@ #include <type_traits> #include <utility> namespace tvm { + +/*! + * \brief add operator + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator+(PrimExpr a, PrimExpr b); + +/*! + * \brief subtraction operator + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator-(PrimExpr a, PrimExpr b); + +/*! + * \brief negation. + * + * \param a input. + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator-(PrimExpr a); + +/*! + * \brief multiplication operator + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator*(PrimExpr a, PrimExpr b); + +/*! + * \brief division operator + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator/(PrimExpr a, PrimExpr b); + +/*! + * \brief left shift operator + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b); + +/*! + * \brief right shift operator + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b); + +/*! + * \brief greater + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator>(PrimExpr a, PrimExpr b); + +/*! + * \brief greater_equal + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b); + +/*! + * \brief less + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator<(PrimExpr a, PrimExpr b); + +/*! + * \brief less_equal + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b); + +/*! + * \brief equal + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator==(PrimExpr a, PrimExpr b); + +/*! + * \brief not_equal + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b); + +/*! + * \brief and + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note This operator does eager constant folding. + */ +TVM_TIRX_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b); + +/*! + * \brief or + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note This operator does eager constant folding. + */ +TVM_TIRX_DLL PrimExpr operator||(PrimExpr a, PrimExpr b); + +/*! + * \brief not + * + * \param a left operand + * \return The result expression. + * \note This operator does eager constant folding. + */ +TVM_TIRX_DLL PrimExpr operator!(PrimExpr a); + +/*! + * \brief take bitwise and of two values + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator&(PrimExpr a, PrimExpr b); + +/*! + * \brief take bitwise or of two values + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator|(PrimExpr a, PrimExpr b); + +/*! + * \brief take bitwise xor of two values + * + * \param a left operand + * \param b right operand + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator^(PrimExpr a, PrimExpr b); + +/*! + * \brief take bitwise negation of two values + * + * \param a the input expression. + * \return The result expression. + * \note this function does eager constant folding for + * index types(int32, int64) when possible. + */ +TVM_TIRX_DLL PrimExpr operator~(PrimExpr a); + +/*! \brief Add an expression in place. */ +inline PrimExpr operator+=(PrimExpr& a, PrimExpr b) { + a = a + b; + return a; +} + +/*! \brief Subtract an expression in place. */ +inline PrimExpr operator-=(PrimExpr& a, PrimExpr b) { + a = a - b; + return a; +} + +/*! \brief Multiply an expression in place. */ +inline PrimExpr operator*=(PrimExpr& a, PrimExpr b) { + a = a * b; + return a; +} + namespace prim { /*! @@ -565,12 +798,6 @@ inline PrimExpr MakeConst(PrimType dtype, ValueType value, Location loc) { } // namespace prim // additional const expression overloading -#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \ - inline PrimExpr Name(PrimExpr& a, PrimExpr b) { \ - a = OpFunc(a, b); \ - return a; \ - } - #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \ inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \ inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \ @@ -621,9 +848,6 @@ inline PrimExpr MakeConst(PrimType dtype, ValueType value, Location loc) { return Name(prim::MakeConst(b.ty(), a), b, loc); \ } -TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator+=, operator+); -TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator-=, operator-); -TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator*=, operator*); TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator+); TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator-); TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator*); diff --git a/tirx/include/tvm/ir/stmt.h b/tirx/include/tvm/ir/stmt.h index 16768fced7..53c832dcb8 100644 --- a/tirx/include/tvm/ir/stmt.h +++ b/tirx/include/tvm/ir/stmt.h @@ -233,6 +233,29 @@ class RegionStmt : public Stmt { TVM_DEFINE_OBJECT_REF_COW_METHOD(RegionStmtNode); }; +namespace op_attr { +/*! + * \brief Construct fresh typed variables for a region's lexical body parameters. + * + * The input carries only the operation, operands and attributes, without a body + * or builder state. Parameters are ordered, distinct definitions and may have + * name hints. Every region operation must register this hook, returning an empty + * array when it has no body parameters. Presence of the attribute identifies + * region support without invoking the hook or allocating variables. + */ +inline constexpr ffi::reflection::AttrKey< + ffi::reflection::NativeFunctionView<ffi::Array<Var>(const CallNode* call)>> + kRegionGetBodyParams{"FRegionGetBodyParams"}; +/*! \brief Optional dialect validation of a complete region statement. */ +inline constexpr ffi::reflection::AttrKey< + ffi::reflection::NativeFunctionView<void(const RegionStmtNode*)>> + kRegionValidate{"FRegionValidate"}; +} // namespace op_attr + +// Compatibility names; new callers can deduce the value type from the key. +using FRegionGetBodyParams = decltype(op_attr::kRegionGetBodyParams)::ValueType; +using FRegionValidate = decltype(op_attr::kRegionValidate)::ValueType; + /*! * \brief Assert condition, if an error occurs, return the error message. * @@ -328,6 +351,13 @@ class TensorStore : public Stmt { TVM_DEFINE_OBJECT_REF_COW_METHOD(TensorStoreNode); }; +namespace type_attr { +/*! \brief Validate a tensor store by its destination type before constructing the statement. */ +inline constexpr ffi::reflection::AttrKey<ffi::reflection::NativeFunctionView<void( + const Expr&, const ffi::Array<PrimExpr>&, const PrimExpr&)>> + kTensorStoreValidateByTy{"__tvm_tensor_store_validate__"}; +} // namespace type_attr + /*! * \brief Evaluates an expression. * This is mostly used for putting a Call node into Stmt. @@ -657,35 +687,5 @@ inline const char* ForKind2String(ForKind t) { TVM_FFI_UNREACHABLE(); } -namespace op_attr { -/*! - * \brief Construct fresh typed variables for a region's lexical body parameters. - * - * The input carries only the operation, operands and attributes, without a body - * or builder state. Parameters are ordered, distinct definitions and may have - * name hints. Every region operation must register this hook, returning an empty - * array when it has no body parameters. Presence of the attribute identifies - * region support without invoking the hook or allocating variables. - */ -inline constexpr ffi::reflection::AttrKey< - ffi::reflection::NativeFunctionView<ffi::Array<Var>(const CallNode* call)>> - kRegionGetBodyParams{"FRegionGetBodyParams"}; -/*! \brief Optional dialect validation of a complete region statement. */ -inline constexpr ffi::reflection::AttrKey< - ffi::reflection::NativeFunctionView<void(const RegionStmtNode*)>> - kRegionValidate{"FRegionValidate"}; -} // namespace op_attr - -// Compatibility names; new callers can deduce the value type from the key. -using FRegionGetBodyParams = decltype(op_attr::kRegionGetBodyParams)::ValueType; -using FRegionValidate = decltype(op_attr::kRegionValidate)::ValueType; - -namespace type_attr { -inline constexpr ffi::reflection::AttrKey<ffi::Function> kEvaluateValidate{ - "__tvm_evaluate_validate__"}; -inline constexpr ffi::reflection::AttrKey<ffi::Function> kTensorStoreValidate{ - "__tvm_tensor_store_validate__"}; -} // namespace type_attr - } // namespace tvm #endif // TVM_IR_STMT_H_ diff --git a/tirx/include/tvm/script/printer/doc_translator.h b/tirx/include/tvm/script/printer/doc_translator.h index 37e9926146..3287aac909 100644 --- a/tirx/include/tvm/script/printer/doc_translator.h +++ b/tirx/include/tvm/script/printer/doc_translator.h @@ -451,9 +451,9 @@ inline ffi::Optional<ExprDoc> DocTranslatorObj::DefaultTranslate(ffi::AnyView va tvm::script::printer::type_attr::kDocTranslate.name); ffi::AnyView attr = column[value.type_index()]; if (attr.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) { - return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<ExprDoc>>( - reinterpret_cast<decltype(DocTranslatorVTable::translate)>(attr.cast<void*>())( - this, value, destination)) + return type_attr::kDocTranslate.ValueFrom(attr) + .value() + .CallExpected(this, value, destination) .value(); } if (attr.type_index() == ffi::TypeIndex::kTVMFFIFunction) { diff --git a/tirx/include/tvm/target/target.h b/tirx/include/tvm/target/target.h index e6a9e21052..d1a2fe99c4 100644 --- a/tirx/include/tvm/target/target.h +++ b/tirx/include/tvm/target/target.h @@ -35,12 +35,6 @@ namespace tvm { -namespace backend::attr { -/*! \brief Optional backend configuration on a Target, distinct from serialized IR/runtime JSON. */ -inline constexpr ffi::reflection::AttrKey<ffi::Map<ffi::String, ffi::Any>> kBackendConfig{ - "backend_config"}; -} // namespace backend::attr - class TargetInternal; class Target; @@ -208,5 +202,11 @@ class Target : public ffi::ObjectRef { TVM_TIRX_DLL void ExitWithScope(); }; +namespace attr { +/*! \brief Optional backend configuration on a Target, distinct from serialized IR/runtime JSON. */ +inline constexpr ffi::reflection::AttrKey<ffi::Map<ffi::String, ffi::Any>> kBackendConfig{ + "backend_config"}; +} // namespace attr + } // namespace tvm #endif // TVM_TARGET_TARGET_H_ diff --git a/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc b/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc index c34bff568e..df263acd14 100644 --- a/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc +++ b/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc @@ -289,7 +289,7 @@ void RegisterCudaIntrinRules() { sig::arg<IntExpr>("width", "The warp thread width, must be a power of 2.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "__shfl_sync") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque) .set_attr(tvm::backend::cuda::op_attr::kNeedWarpShuffle, true); OpDef("tirx.cuda.__shfl_up_sync") @@ -299,7 +299,7 @@ void RegisterCudaIntrinRules() { sig::arg<IntExpr>("width", "The warp thread width, must be a power of 2.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "__shfl_up_sync") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque) .set_attr(tvm::backend::cuda::op_attr::kNeedWarpShuffle, true); OpDef("tirx.cuda.__shfl_down_sync") @@ -309,7 +309,7 @@ void RegisterCudaIntrinRules() { sig::arg<IntExpr>("width", "The warp thread width, must be a power of 2.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "__shfl_down_sync") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque) .set_attr(tvm::backend::cuda::op_attr::kNeedWarpShuffle, true); OpDef("tirx.cuda.__shfl_xor_sync") @@ -319,13 +319,13 @@ void RegisterCudaIntrinRules() { sig::arg<IntExpr>("width", "The warp thread width, must be a power of 2.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "__shfl_xor_sync") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque) .set_attr(tvm::backend::cuda::op_attr::kNeedWarpShuffle, true); OpDef("tirx.cuda.__activemask") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32)) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "__activemask") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::backend::cuda::op_attr::kNeedWarpShuffle, true); // clang-format on } diff --git a/tirx/src/backend/cuda/codegen/target_kind.cc b/tirx/src/backend/cuda/codegen/target_kind.cc index 9bd9e4c4f4..d0d2d99382 100644 --- a/tirx/src/backend/cuda/codegen/target_kind.cc +++ b/tirx/src/backend/cuda/codegen/target_kind.cc @@ -145,7 +145,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { TargetKindDef("cuda") .set_default_device_type(kDLCUDA) - .def_option<ffi::Map<ffi::String, ffi::Any>>(backend::attr::kBackendConfig.name) + .def_option<ffi::Map<ffi::String, ffi::Any>>(tvm::attr::kBackendConfig.name) .def_option<ffi::String>("mcpu") .def_option<ffi::String>("arch") .def_option<int64_t>("max_shared_memory_per_block") diff --git a/tirx/src/backend/cuda/op/iket.cc b/tirx/src/backend/cuda/op/iket.cc index f6a3c87763..9c6e329b6f 100644 --- a/tirx/src/backend/cuda/op/iket.cc +++ b/tirx/src/backend/cuda/op/iket.cc @@ -41,7 +41,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.cuda.iket_mark") .signature(sig::arg("name", "The name."), sig::var_args("args")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& iket_range_start_op() { @@ -53,7 +53,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.cuda.iket_range_start") .signature(sig::arg("name", "The name."), sig::var_args("args")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32)) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& iket_range_end_op() { @@ -65,7 +65,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.cuda.iket_range_end") .signature(sig::arg<IntExpr>("token", "The token."), sig::var_args("args")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& iket_range_push_op() { @@ -77,7 +77,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.cuda.iket_range_push") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) .signature(sig::arg("name", "The name."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& iket_range_pop_op() { @@ -88,7 +88,7 @@ const Op& iket_range_pop_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.cuda.iket_range_pop") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& iket_sentinel_token_op() { @@ -100,7 +100,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.cuda.iket_sentinel_token") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32)) .signature(sig::arg("name", "The name.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -108,7 +108,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32)) .signature(sig::arg("event_id", "The event identifier."), sig::arg("source_code", "The source code."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } } // namespace cuda diff --git a/tirx/src/backend/cuda/op/legacy.cc b/tirx/src/backend/cuda/op/legacy.cc index 42841e8804..57f54fa229 100644 --- a/tirx/src/backend/cuda/op/legacy.cc +++ b/tirx/src/backend/cuda/op/legacy.cc @@ -35,26 +35,26 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg("src_ptr", "The source pointer."), sig::arg<IntExpr>("src_offset", "The source offset."), sig::arg<IntExpr>("dst_stride", "The destination stride.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.cuda.mma_fill") .signature(sig::arg<IntExpr>("local_size", "The local allocation size."), sig::arg("local_ptr", "The local pointer."), sig::arg<IntExpr>("offset", "The offset.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.ptx_legacy.ldmatrix") .signature(sig::arg("trans"), sig::arg<IntExpr>("num"), sig::arg("dtype"), sig::arg("local_ptr"), sig::arg<IntExpr>("local_offset"), sig::arg("smem_ptr"), sig::arg<IntExpr>("smem_offset")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.ptx_legacy.mma") .signature(sig::arg("shape"), sig::arg("a_layout"), sig::arg("b_layout"), sig::arg("a_dtype"), sig::arg("b_dtype"), sig::arg("c_dtype"), sig::arg("a_ptr"), sig::arg<IntExpr>("a_offset"), sig::arg("b_ptr"), sig::arg<IntExpr>("b_offset"), sig::arg("acc_ptr"), sig::arg<IntExpr>("c_offset"), sig::arg("saturate"), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } } // namespace tirx diff --git a/tirx/src/backend/cuda/op/target_builtin.cc b/tirx/src/backend/cuda/op/target_builtin.cc index edbc2e0e54..9e6d3c8209 100644 --- a/tirx/src/backend/cuda/op/target_builtin.cc +++ b/tirx/src/backend/cuda/op/target_builtin.cc @@ -102,7 +102,7 @@ void RegisterCudaTargetBuiltins() { .set_validator(ffi::reflection::NativeFunctionView<void( const CallNode*)>::FromNative<&ValidateDeclTmem>(), true) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("tirx.cuda.bmma_sync") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) @@ -114,7 +114,7 @@ void RegisterCudaTargetBuiltins() { sig::arg<IntExpr>("index_b", "The B fragment index."), sig::arg("fragment_c", "The C fragment."), sig::arg<IntExpr>("index_c", "The C fragment index.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); // Raw legacy cp.async form emitted by InjectPTXAsyncCopy (and round-tripped by // the T.s_tir.cp_async_raw.legacy 6-arg surface). It carries the element dtype in Call.dtype @@ -125,7 +125,7 @@ void RegisterCudaTargetBuiltins() { sig::arg("src_ptr", "The source pointer."), sig::arg<IntExpr>("src_offset", "The source offset."), sig::arg<IntExpr>("cp_size", "The copy size."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.cuda.timer_init") .signature(sig::arg("profiler_buffer", "The profiler buffer."), @@ -133,7 +133,7 @@ void RegisterCudaTargetBuiltins() { sig::arg<IntExpr>("profiler_write_offset", "The profiler write offset."), sig::arg<IntExpr>("num_groups", "The number of groups."), sig::arg<IntExpr>("group_id", "The group identifier.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.cuda.timer_start") .signature(sig::arg("event_type", "The event type."), @@ -142,7 +142,7 @@ void RegisterCudaTargetBuiltins() { sig::arg<IntExpr>("profiler_write_offset", "The profiler write offset."), sig::arg<IntExpr>("profiler_write_stride", "The profiler write stride."), sig::arg("leader_cond", "The leader condition.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.cuda.timer_end") .signature(sig::arg("event_type", "The event type."), @@ -151,7 +151,7 @@ void RegisterCudaTargetBuiltins() { sig::arg<IntExpr>("profiler_write_offset", "The profiler write offset."), sig::arg<IntExpr>("profiler_write_stride", "The profiler write stride."), sig::arg("leader_cond", "The leader condition.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.cuda.timer_finalize") .signature(sig::arg("profiler_buffer", "The profiler buffer."), @@ -159,7 +159,7 @@ void RegisterCudaTargetBuiltins() { sig::arg<IntExpr>("profiler_write_offset", "The profiler write offset."), sig::arg<IntExpr>("profiler_write_stride", "The profiler write stride."), sig::arg("leader_cond", "The leader condition.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); RegisterDeviceIntrinsics(); } @@ -167,7 +167,7 @@ void RegisterCudaTargetBuiltins() { namespace { TVM_FFI_NO_INLINE void RegisterDeviceIntrinsicAttrs(OpDef& def, CallEffectKind effect_kind) { - def.set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(effect_kind)); + def.set_attr(tvm::op_attr::kCallEffectKind, effect_kind); } template <typename... Specs> diff --git a/tirx/src/backend/cuda/op/tcgen05.cc b/tirx/src/backend/cuda/op/tcgen05.cc index 8e3def9559..c50f6eddce 100644 --- a/tirx/src/backend/cuda/op/tcgen05.cc +++ b/tirx/src/backend/cuda/op/tcgen05.cc @@ -73,17 +73,17 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.cuda.tcgen05_encode_instr_descriptor") .signature(sig::arg("desc"), sig::call_attrs<backend::cuda::TCGen05InstrDescriptorAttrs>()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()); OpDef("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled") .signature(sig::arg("desc"), sig::call_attrs<backend::cuda::TCGen05InstrDescriptorBlockScaledAttrs>()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()); OpDef("tirx.cuda.tcgen05_encode_matrix_descriptor") .signature(sig::arg("desc"), sig::arg("addr"), sig::arg<IntExpr>("ldo"), sig::arg<IntExpr>("sdo"), sig::arg<IntExpr>("swizzle")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()); } diff --git a/tirx/src/backend/cuda/op/tensormap.cc b/tirx/src/backend/cuda/op/tensormap.cc index 3f138b27f8..f9ad36f95b 100644 --- a/tirx/src/backend/cuda/op/tensormap.cc +++ b/tirx/src/backend/cuda/op/tensormap.cc @@ -63,7 +63,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32)) .signature(sig::arg("descriptor", "The descriptor."), sig::arg("data", "The input data."), sig::var_args<PrimExpr>("args"), sig::call_attrs<TensorMapEncodeTiledAttr>()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } } // namespace cuda diff --git a/tirx/src/backend/cuda/tile/exec_context.cc b/tirx/src/backend/cuda/tile/exec_context.cc index 740f139ecf..53af5f5d4a 100644 --- a/tirx/src/backend/cuda/tile/exec_context.cc +++ b/tirx/src/backend/cuda/tile/exec_context.cc @@ -24,6 +24,7 @@ #include "exec_context.h" #include <tvm/ir/prim/expr.h> +#include <tvm/ir/prim/op.h> #include <tvm/runtime/logging.h> #include <tvm/sym/analyzer.h> diff --git a/tirx/src/backend/metal/codegen/intrin_rule_metal.cc b/tirx/src/backend/metal/codegen/intrin_rule_metal.cc index 766743b297..c92a9b7808 100644 --- a/tirx/src/backend/metal/codegen/intrin_rule_metal.cc +++ b/tirx/src/backend/metal/codegen/intrin_rule_metal.cc @@ -161,21 +161,21 @@ void RegisterMetalIntrinRules() { sig::arg<IntExpr>("lane", "The source thread id.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "simd_shuffle") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.metal.simd_shuffle_up") .signature(sig::arg("var", "The variable to sync."), sig::arg<IntExpr>("delta", "The source lane id offset to be added.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "simd_shuffle_up") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.metal.simd_shuffle_down") .signature(sig::arg("var", "The variable to sync."), sig::arg<IntExpr>("delta", "The source lane id offset to be subtracted.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "simd_shuffle_down") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); // clang-format on } diff --git a/tirx/src/backend/metal/op/cooperative_tensor.cc b/tirx/src/backend/metal/op/cooperative_tensor.cc index a3a202ef18..fcb6d07214 100644 --- a/tirx/src/backend/metal/op/cooperative_tensor.cc +++ b/tirx/src/backend/metal/op/cooperative_tensor.cc @@ -40,7 +40,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<PrimExpr>("value", "The value to use."), sig::arg<IntExpr>("rows", "The number of rows."), sig::arg<IntExpr>("cols", "The number of columns.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& cooperative_tensor_load_op() { @@ -60,7 +60,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("mma_N", "The N dimension of the matrix operation."), sig::arg<IntExpr>("mma_K", "The K dimension of the matrix operation."), sig::arg<IntExpr>("operand_role", "The matrix operand role.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& cooperative_tensor_store_op() { @@ -80,7 +80,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("mma_N", "The N dimension of the matrix operation."), sig::arg<IntExpr>("mma_K", "The K dimension of the matrix operation."), sig::arg<IntExpr>("operand_role", "The matrix operand role.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& cooperative_tensor_multiply_accumulate_op() { @@ -100,7 +100,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("K", "The K dimension."), sig::arg<PrimExpr>("transpose_a", "Whether to transpose A."), sig::arg<PrimExpr>("transpose_b", "Whether to transpose B.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } } // namespace metal diff --git a/tirx/src/backend/metal/op/target_builtin.cc b/tirx/src/backend/metal/op/target_builtin.cc index 2c9bb888b9..a2bc2afafb 100644 --- a/tirx/src/backend/metal/op/target_builtin.cc +++ b/tirx/src/backend/metal/op/target_builtin.cc @@ -44,7 +44,7 @@ void RegisterMetalTargetBuiltins() { sig::arg("value", "The value to use."), sig::arg<IntExpr>("col", "The column index."), sig::arg<IntExpr>("row", "The row index.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.metal.simdgroup_load") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) @@ -53,7 +53,7 @@ void RegisterMetalTargetBuiltins() { sig::arg<IntExpr>("col", "The column index."), sig::arg<IntExpr>("row", "The row index."), sig::arg("transpose_matrix", "Whether to transpose the matrix.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.metal.simdgroup_store") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) @@ -62,7 +62,7 @@ void RegisterMetalTargetBuiltins() { sig::arg<IntExpr>("col", "The column index."), sig::arg<IntExpr>("row", "The row index."), sig::arg("transpose_matrix", "Whether to transpose the matrix.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.metal.simdgroup_multiply_accumulate") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) @@ -71,7 +71,7 @@ void RegisterMetalTargetBuiltins() { sig::arg("a", "The A operand."), sig::arg<IntExpr>("index_a", "The A fragment index."), sig::arg("b", "The B operand."), sig::arg<IntExpr>("index_b", "The B fragment index."), sig::arg("c", "The C operand."), sig::arg<IntExpr>("index_c", "The C fragment index.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } TVM_FFI_STATIC_INIT_BLOCK() { RegisterMetalTargetBuiltins(); } diff --git a/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc b/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc index 99b749c1cc..1d0dcbe878 100644 --- a/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc +++ b/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc @@ -22,6 +22,7 @@ * \brief OpenCL intrinsic rules. */ #include <tvm/codegen/intrin_rule.h> +#include <tvm/ir/prim/op.h> #include <tvm/sym/analyzer.h> #include <tvm/tirx/op/abi.h> #include <tvm/tirx/op/attr.h> diff --git a/tirx/src/backend/opencl/op/memory.cc b/tirx/src/backend/opencl/op/memory.cc index e37d618e5e..6e8406984b 100644 --- a/tirx/src/backend/opencl/op/memory.cc +++ b/tirx/src/backend/opencl/op/memory.cc @@ -40,7 +40,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("channel_size", "The number of channels."), sig::arg<PrimExpr>("value", "The value to use.")) .set_attr(tvm::tirx::op_attr::kVectorizable, true) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& texture2d_load_op() { @@ -56,7 +56,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("channel_size", "The number of channels."), sig::arg<PrimExpr>("element_index", "The element index within a texture channel.")) .set_attr(tvm::tirx::op_attr::kVectorizable, true) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& nd_mem_alloc_with_scope_op() { @@ -69,7 +69,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg("storage_scope", "The storage scope."), sig::arg<IntExpr>("ndim", "The number of dimensions."), sig::arg("shape", "The shape."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } } // namespace opencl diff --git a/tirx/src/backend/trn/op/target_builtin.cc b/tirx/src/backend/trn/op/target_builtin.cc index 3b1674583c..51cb30d224 100644 --- a/tirx/src/backend/trn/op/target_builtin.cc +++ b/tirx/src/backend/trn/op/target_builtin.cc @@ -69,7 +69,7 @@ TVM_FFI_NO_INLINE std::string MakeNKIIntrinsicName(const char* op_name) { TVM_FFI_NO_INLINE void RegisterNKIIntrinsicAttrs(OpDef& def) { def.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } template <typename... Specs> diff --git a/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc b/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc index 012a02f3b0..b59539439c 100644 --- a/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc +++ b/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc @@ -187,21 +187,21 @@ void RegisterWebGPUIntrinRules() { sig::arg<IntExpr>("lane", "The source thread id.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "subgroupShuffle") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.webgpu.subgroup_shuffle_up") .signature(sig::arg("var", "The variable to sync."), sig::arg<IntExpr>("delta", "The source lane id offset to be added.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "subgroupShuffleUp") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); OpDef("tirx.webgpu.subgroup_shuffle_down") .signature(sig::arg("var", "The variable to sync."), sig::arg<IntExpr>("delta", "The source lane id offset to be subtracted.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>( tvm::op_attr::kInferType) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "subgroupShuffleDown") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); // clang-format on } diff --git a/tirx/src/ir/expr.cc b/tirx/src/ir/expr.cc index beebedaf04..a23aa87f58 100644 --- a/tirx/src/ir/expr.cc +++ b/tirx/src/ir/expr.cc @@ -39,6 +39,43 @@ 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> @@ -702,13 +739,13 @@ TensorLoad::TensorLoad(Expr source, ffi::Array<PrimExpr> indices, Location loc) << "TensorLoad requires a source with a defined type"; namespace refl = ffi::reflection; static const refl::TypeAttrColumn infer_column(type_attr::kTensorLoadInferType.name); - auto infer = infer_column[source->ty->type_index()]; - TVM_FFI_CHECK(infer != nullptr, TypeError) + auto infer = type_attr::kTensorLoadInferType.ValueFrom(infer_column[source->ty->type_index()]); + TVM_FFI_CHECK(infer.has_value(), TypeError) << "Type " << source->ty->GetTypeKey() << " does not support TensorLoad"; auto node = ffi::make_object<TensorLoadNode>(std::move(source)); node->indices = std::move(indices); node->loc = loc; - PrimType ty = infer.cast<FTensorLoadInferType>().CallExpected(node.get()).value(); + PrimType ty = infer->CallExpected(node.get()).value(); TVM_FFI_CHECK(ty.defined() && ty.as<PrimTypeNode>(), TypeError) << "TensorLoad inference requires a concrete PrimType"; node->ty = std::move(ty); diff --git a/tirx/src/ir/prim/builtin.cc b/tirx/src/ir/prim/builtin.cc index 5ef75f691b..de7069e05c 100644 --- a/tirx/src/ir/prim/builtin.cc +++ b/tirx/src/ir/prim/builtin.cc @@ -49,8 +49,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("prim.likely") .signature(sig::arg("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, - static_cast<int64_t>(CallEffectKind::kExprAnnotation)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kExprAnnotation) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.if_then_else") @@ -58,28 +57,28 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg("true_value", "The value when the condition is true."), sig::arg("false_value", "The value when the condition is false.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.vscale") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32)) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.ceil") .signature(sig::arg("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.log2") .signature(sig::arg("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.clz") .signature(sig::arg("x", "The input value.")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32)) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } } // namespace prim diff --git a/tirx/src/ir/prim/math.cc b/tirx/src/ir/prim/math.cc index 96a98d1fec..76e2cdda2d 100644 --- a/tirx/src/ir/prim/math.cc +++ b/tirx/src/ir/prim/math.cc @@ -387,206 +387,206 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg<PrimExpr>("x", "The input value."), sig::arg<PrimExpr>("y", "The second input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.fabs") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.fmod") .signature(sig::arg<PrimExpr>("x", "The input value."), sig::arg<PrimExpr>("y", "The second input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.floor") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.round") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.nearbyint") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.trunc") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.exp") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.exp2") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.exp10") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.erf") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.tanh") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.sigmoid") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.sqrt") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.rsqrt") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.log") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.log1p") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.log10") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.tan") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.cos") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.cosh") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.sin") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.sinh") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.asin") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.acos") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.atan") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.acosh") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.asinh") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.atanh") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.atan2") .signature(sig::arg<PrimExpr>("x1", "The first input value."), sig::arg<PrimExpr>("x2", "The second input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.nextafter") .signature(sig::arg<PrimExpr>("x1", "The first input value."), sig::arg<PrimExpr>("x2", "The second input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.hypot") .signature(sig::arg<PrimExpr>("x1", "The first input value."), sig::arg<PrimExpr>("x2", "The second input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.copysign") .signature(sig::arg<PrimExpr>("x1", "The first input value."), sig::arg<PrimExpr>("x2", "The second input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.ldexp") .signature(sig::arg<PrimExpr>("x1", "The first input value."), sig::arg<PrimExpr>("x2", "The second input value.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.isnan") .signature(sig::arg<PrimExpr>("x", "The input value.")) .set_native_fn_attr<&InferTypeIsNaN>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); OpDef("prim.popcount") .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) .signature(sig::arg<PrimExpr>("x", "The input value.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.fma") @@ -594,13 +594,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg<PrimExpr>("x", "The input value."), sig::arg<PrimExpr>("y", "The second input value."), sig::arg<PrimExpr>("z", "The third input value.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); OpDef("prim.assume") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Bool()) .signature(sig::arg<PrimExpr>("cond", "The condition.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kEmbedInfo)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kEmbedInfo); ffi::reflection::GlobalDef() .def("prim.infinity", prim::infinity) diff --git a/tirx/src/ir/side_effect.cc b/tirx/src/ir/side_effect.cc index 30eda9c617..1916c95576 100644 --- a/tirx/src/ir/side_effect.cc +++ b/tirx/src/ir/side_effect.cc @@ -35,8 +35,7 @@ CallEffectKind SideEffect(const Expr& expr) { expr, [&](const CallNode* node, ffi::StructuralVisitorObj* visitor) -> ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> { - auto effect = static_cast<CallEffectKind>( - effects.get(node->op, static_cast<TCallEffectKind>(CallEffectKind::kOpaque))); + auto effect = effects.get(node->op, CallEffectKind::kOpaque); kind = std::max(kind, std::min(effect, CallEffectKind::kUpdateState)); if (kind == CallEffectKind::kUpdateState) return ffi::VisitInterrupt(); return visitor->DefaultVisitExpected(node); diff --git a/tirx/src/ir/stmt.cc b/tirx/src/ir/stmt.cc index be6d5717a0..c315904326 100644 --- a/tirx/src/ir/stmt.cc +++ b/tirx/src/ir/stmt.cc @@ -649,8 +649,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorStoreMaybeInplace TVM_FFI_STATIC_INIT_BLOCK() { StmtNode::RegisterReflection(); - ffi::reflection::EnsureTypeAttrColumn(tvm::type_attr::kTensorStoreValidate.name); - ffi::reflection::EnsureTypeAttrColumn(tvm::type_attr::kEvaluateValidate.name); + ffi::reflection::EnsureTypeAttrColumn(tvm::type_attr::kTensorStoreValidateByTy.name); } // Bind @@ -1059,12 +1058,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { // Evaluate Evaluate::Evaluate(Expr value, Location loc) : Stmt(ffi::UnsafeInit{}) { TVM_FFI_ICHECK(value.defined()); - static const ffi::reflection::TypeAttrColumn validate_column( - tvm::type_attr::kEvaluateValidate.name); - if (auto validate = validate_column[value->ty->type_index()]; validate != nullptr) { - validate.cast<ffi::Function>()(value); - } - ffi::ObjectPtr<EvaluateNode> node = ffi::make_object<EvaluateNode>(std::move(value)); node->loc = loc; data_ = std::move(node); @@ -1088,11 +1081,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { TensorStore::TensorStore(Expr dest, ffi::Array<PrimExpr> indices, PrimExpr value, Location loc) : Stmt(ffi::UnsafeInit{}) { namespace refl = ffi::reflection; - static const refl::TypeAttrColumn validate_column(tvm::type_attr::kTensorStoreValidate.name); - auto validate = validate_column[dest->ty->type_index()]; - TVM_FFI_CHECK(validate != nullptr, TypeError) + static const refl::TypeAttrColumn validate_column(tvm::type_attr::kTensorStoreValidateByTy.name); + auto validate = + type_attr::kTensorStoreValidateByTy.ValueFrom(validate_column[dest->ty->type_index()]); + TVM_FFI_CHECK(validate.has_value(), TypeError) << "Type " << dest->ty->GetTypeKey() << " does not support TensorStore"; - validate.cast<ffi::Function>()(dest, indices, value); + validate->CallExpected(dest, indices, value).value(); auto node = ffi::make_object<TensorStoreNode>(std::move(dest), std::move(value)); node->indices = std::move(indices); node->loc = loc; diff --git a/tirx/src/ir/subscript_expr.cc b/tirx/src/ir/subscript_expr.cc index 7409d1a305..c7e0c3ecfc 100644 --- a/tirx/src/ir/subscript_expr.cc +++ b/tirx/src/ir/subscript_expr.cc @@ -28,15 +28,11 @@ namespace tvm { -using SubscriptSlice = ffi::Array<ffi::Variant< - ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>>, - PrimExpr>>; - TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::EnsureTypeAttrColumn(tvm::type_attr::kSubscriptExprRealize.name); - refl::TypeAttrDef<TupleTypeNode>().def( - tvm::type_attr::kSubscriptExprRealize.name, + refl::EnsureTypeAttrColumn(tvm::type_attr::kExprSubscriptRealizeByTy); + refl::TypeAttrDef<TupleTypeNode>().attr( + tvm::type_attr::kExprSubscriptRealizeByTy, [](Expr value, SubscriptSlice slice, Location loc) -> ffi::ObjectRef { TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError) << "A tuple expression requires exactly one index"; @@ -51,14 +47,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { "ir.SubscriptExprRealize", [](Expr value, SubscriptSlice slice, Location loc) -> ffi::ObjectRef { TVM_FFI_CHECK(value.defined(), TypeError) << "Cannot subscript an undefined expression"; - static refl::TypeAttrColumn realize_column(tvm::type_attr::kSubscriptExprRealize.name); - ffi::AnyView packed_realize = realize_column[value->ty->type_index()]; - TVM_FFI_CHECK(packed_realize != nullptr, TypeError) + static refl::TypeAttrColumn realize_column(tvm::type_attr::kExprSubscriptRealizeByTy); + auto realize = tvm::type_attr::kExprSubscriptRealizeByTy.ValueFrom( + realize_column[value->ty->type_index()]); + TVM_FFI_CHECK(realize.has_value(), TypeError) << "Type " << value->ty->GetTypeKey() << " does not support subscript"; - ffi::ObjectRef result = - packed_realize.cast<ffi::Function>()(value, slice, loc).cast<ffi::ObjectRef>(); + ffi::ObjectRef result = (*realize)(value, slice, loc); TVM_FFI_CHECK(result.defined(), TypeError) - << "__tvm_subscript_expr_realize__ for type " << value->ty->GetTypeKey() + << "__tvm_expr_subscript_realize__ for type " << value->ty->GetTypeKey() << " returned an undefined object"; return result; }); diff --git a/tirx/src/script/printer/ir/call.cc b/tirx/src/script/printer/ir/call.cc index a66749b247..4b4ddf7535 100644 --- a/tirx/src/script/printer/ir/call.cc +++ b/tirx/src/script/printer/ir/call.cc @@ -48,7 +48,10 @@ ffi::Optional<ExprDoc> CallDocTranslate(DocTranslatorObj* d, ffi::AnyView input, if (overrides.count(op.value())) { ffi::Any hook = overrides[op.value()]; if (hook.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) { - return hook.cast<FDocTranslate>().CallExpected(d, input, destination).value(); + return op_attr::kDocTranslateOpCall.ValueFrom(hook) + .value() + .CallExpected(d, input, destination) + .value(); } ffi::Any binding = nullptr; if (destination) binding = ffi::GetRef<ffi::ObjectRef>(destination); diff --git a/tirx/src/tirx/ir/stmt.cc b/tirx/src/tirx/ir/stmt.cc index d8f7e412eb..dd64e6471a 100644 --- a/tirx/src/tirx/ir/stmt.cc +++ b/tirx/src/tirx/ir/stmt.cc @@ -45,7 +45,8 @@ TVM_FFI_INLINE int GetLanesOrVScaleFactor(const PrimType& ty) { return ty.IsScalableVector() ? ty.VScaleFactor() : ty.lanes(); } -void ValidateTensorStore(Expr dest, ffi::Array<PrimExpr> indices, PrimExpr value) { +void ValidateTensorStore(const Expr& dest, const ffi::Array<PrimExpr>& indices, + const PrimExpr& value) { TensorType type = dest->ty.as_or_throw<TensorType>(); TVM_FFI_ICHECK_EQ(type->shape.size(), indices.size()) << "Destination " << dest << " is " << type->shape.size() @@ -95,18 +96,12 @@ void ValidateTensorStore(Expr dest, ffi::Array<PrimExpr> indices, PrimExpr value } } -void ValidateTensorEvaluate(Expr value) { - TVM_FFI_ICHECK(!value->IsInstance<VarNode>()) - << "A tensor variable cannot be used as a scalar Evaluate value; " - << "use tensor.data to evaluate its physical pointer"; -} } // namespace TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = ffi::reflection; - refl::TypeAttrDef<TensorTypeNode>() - .def(tvm::type_attr::kTensorStoreValidate.name, ValidateTensorStore) - .def(tvm::type_attr::kEvaluateValidate.name, ValidateTensorEvaluate); + refl::TypeAttrDef<TensorTypeNode>().native_fn_attr<&ValidateTensorStore>( + type_attr::kTensorStoreValidateByTy); } } // namespace tirx diff --git a/tirx/src/tirx/ir/tensor.cc b/tirx/src/tirx/ir/tensor.cc index 5890b18681..a21d8917c5 100644 --- a/tirx/src/tirx/ir/tensor.cc +++ b/tirx/src/tirx/ir/tensor.cc @@ -58,10 +58,6 @@ using namespace tvm::prim; namespace { -using SubscriptSlice = ffi::Array<ffi::Variant< - ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>>, - PrimExpr>>; - PrimExpr SliceExtent(const PrimExpr& start, const PrimExpr& stop, const sym::Analyzer& analyzer) { PrimExpr extent = Range::GetExtentFromBeginEnd(start, stop); if (const auto* add = stop.as<AddNode>(); @@ -71,13 +67,7 @@ PrimExpr SliceExtent(const PrimExpr& start, const PrimExpr& stop, const sym::Ana return analyzer->Simplify(extent); } -ffi::ObjectRef RealizeTensorSubscript( - Expr value, - ffi::Array<ffi::Variant< - ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>>, - PrimExpr>> - slice, - Location loc) { +ffi::ObjectRef RealizeTensorSubscript(Expr value, SubscriptSlice slice, Location loc) { TensorVar tensor = value.as_or_throw<TensorVar>(); TensorType tensor_ty = tensor.ty(); TVM_FFI_CHECK_LE(slice.size(), tensor_ty->shape.size(), IndexError) @@ -194,10 +184,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::TypeAttrDef<TensorTypeNode>().def(tvm::type_attr::kSubscriptExprRealize.name, - RealizeTensorSubscript); - refl::TypeAttrDef<TensorRegionTypeNode>().def(tvm::type_attr::kSubscriptExprRealize.name, - RealizeTensorRegionSubscript); + refl::TypeAttrDef<TensorTypeNode>().attr(tvm::type_attr::kExprSubscriptRealizeByTy, + RealizeTensorSubscript); + refl::TypeAttrDef<TensorRegionTypeNode>().attr(tvm::type_attr::kExprSubscriptRealizeByTy, + RealizeTensorRegionSubscript); } ffi::Array<PrimExpr> SimplifyArray(sym::AnalyzerObj* ana, ffi::Array<PrimExpr> array) { diff --git a/tirx/src/tirx/ir/tensor_load.cc b/tirx/src/tirx/ir/tensor_load.cc index 3129f59d16..4bc3bf1b84 100644 --- a/tirx/src/tirx/ir/tensor_load.cc +++ b/tirx/src/tirx/ir/tensor_load.cc @@ -70,9 +70,8 @@ PrimType InferTensorLoadType(const TensorLoadNode* load) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::TypeAttrDef<TensorTypeNode>().attr( - type_attr::kTensorLoadInferType.name, - FTensorLoadInferType::FromNative<&InferTensorLoadType>()); + refl::TypeAttrDef<TensorTypeNode>().native_fn_attr<&InferTensorLoadType>( + type_attr::kTensorLoadInferType); } } // namespace tirx diff --git a/tirx/src/tirx/op/abi.cc b/tirx/src/tirx/op/abi.cc index 9f25bd8ca3..d5419ceab1 100644 --- a/tirx/src/tirx/op/abi.cc +++ b/tirx/src/tirx/op/abi.cc @@ -61,7 +61,7 @@ const Op& call_extern_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.call_extern") .signature(sig::arg("func_name", "The function name."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& call_pure_extern_op() { @@ -72,7 +72,7 @@ const Op& call_pure_extern_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.call_pure_extern") .signature(sig::arg("func_name", "The function name."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& call_llvm_intrin_op() { @@ -83,7 +83,7 @@ const Op& call_llvm_intrin_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.call_llvm_intrin") .signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& call_llvm_pure_intrin_op() { @@ -94,7 +94,7 @@ const Op& call_llvm_pure_intrin_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.call_llvm_pure_intrin") .signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)) + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure) .set_attr(tvm::tirx::op_attr::kVectorizable, true); } @@ -107,7 +107,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.call_spirv_pure_glsl450") .signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."), sig::var_args<PrimExpr>("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& abi_field_get_op() { @@ -119,7 +119,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.abi_field_get") .signature(sig::arg("arr", "The array."), sig::arg<IntExpr>("index", "The index."), sig::arg<IntExpr>("field", "The field index.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kReadState)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState); } const Op& abi_field_set_op() { @@ -133,7 +133,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg("arr", "The array."), sig::arg<IntExpr>("index", "The index."), sig::arg<IntExpr>("field", "The field index."), sig::arg("value", "The value to use.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kUpdateState)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kUpdateState); } const Op& throw_last_error_op() { @@ -144,7 +144,7 @@ const Op& throw_last_error_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.throw_last_error") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& stack_alloca_op() { @@ -157,7 +157,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .set_native_fn_attr<&InferTypeStackAlloca>(tvm::op_attr::kInferType) .signature(sig::arg("dtype_str", "The data type name."), sig::arg<IntExpr>("num", "The number of entries.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& stack_make_shape_op() { @@ -169,7 +169,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.stack_make_shape") .set_attr(tvm::op_attr::kFixedReturnType, PtrType(PrimType::Int(64))) .signature(sig::var_args<IntExpr>("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& stack_make_dltensor_op() { @@ -185,7 +185,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("ndim", "The number of dimensions."), sig::arg("arr_dtype", "The array data type."), sig::arg<IntExpr>("elem_offset", "The element offset.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& call_packed_op() { @@ -197,7 +197,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.call_packed") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32)) .signature(sig::arg("func_name", "The function name."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -225,7 +225,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg("kernel", "The kernel."), sig::var_args("args"), sig::call_attrs<CallFFIKernelAttr>()) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32)) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& call_cpacked_op() { @@ -237,7 +237,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.call_cpacked") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32)) .signature(sig::arg("func_name", "The function name."), sig::var_args("args")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& call_packed_lowered_op() { @@ -252,7 +252,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg("args_stack", "The argument stack."), sig::arg<IntExpr>("begin", "The start index."), sig::arg<IntExpr>("end", "The end index.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& call_cpacked_lowered_op() { @@ -267,7 +267,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg("args_stack", "The argument stack."), sig::arg<IntExpr>("begin", "The start index."), sig::arg<IntExpr>("end", "The end index.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -279,7 +279,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("dtype_code_hint", "The data type code hint."), sig::arg<IntExpr>("dtype_bits_hint", "The data type bit-width hint.")) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "TVMBackendAllocWorkspace") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -289,7 +289,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("device_id", "The device index."), sig::arg("ptr", "The pointer.")) .set_attr(tvm::tirx::op_attr::kGlobalSymbol, "TVMBackendFreeWorkspace") - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } } // namespace tirx diff --git a/tirx/src/tirx/op/annotation.cc b/tirx/src/tirx/op/annotation.cc index 63d36d6350..a488797d51 100644 --- a/tirx/src/tirx/op/annotation.cc +++ b/tirx/src/tirx/op/annotation.cc @@ -39,7 +39,7 @@ const Op& kernel_replace_point_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.kernel_replace_point") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& ignore_loop_partition_op() { @@ -51,7 +51,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.ignore_loop_partition") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Bool()) .signature(sig::arg<PrimExpr>("predicate", "The predicate.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } } // namespace tirx diff --git a/tirx/src/tirx/op/cpu.cc b/tirx/src/tirx/op/cpu.cc index 2d1347601b..cc652174ef 100644 --- a/tirx/src/tirx/op/cpu.cc +++ b/tirx/src/tirx/op/cpu.cc @@ -40,7 +40,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.cpu_parallel_barrier") .signature() .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } } // namespace tirx diff --git a/tirx/src/tirx/op/gpu.cc b/tirx/src/tirx/op/gpu.cc index d943f39a86..fbe69f92e1 100644 --- a/tirx/src/tirx/op/gpu.cc +++ b/tirx/src/tirx/op/gpu.cc @@ -51,7 +51,7 @@ PrimExpr gpu_thread_return(Location loc) { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.gpu_thread_return") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kControlJump)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kControlJump); } const Op& gpu_thread_filter_op() { @@ -64,7 +64,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Bool()) .signature(sig::arg<PrimVar>("var", "The thread-axis variable."), sig::arg<PrimExpr>("pred", "The predicate.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& gpu_active_thread_selector_op() { @@ -77,7 +77,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg<PrimVar>("var", "The thread-axis variable."), sig::arg<PrimExpr>("pred", "The predicate.")) .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& gpu_thread_invariant_op() { @@ -89,7 +89,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.gpu_thread_invariant") .set_native_fn_attr<&InferTypeReturnArgType<0>>(tvm::op_attr::kInferType) .signature(sig::arg<PrimExpr>("cond", "The condition.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& gpu_storage_sync_op() { @@ -101,7 +101,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.gpu_storage_sync") .signature(sig::arg<StringImm>("storage_scope", "The storage scope.")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& gpu_warp_shuffle_op() { @@ -117,7 +117,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("width", "The width."), sig::arg<IntExpr>("warp_size", "The number of threads per warp.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& gpu_warp_shuffle_up_op() { @@ -132,7 +132,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("offset", "The offset."), sig::arg<IntExpr>("width", "The width."), sig::arg<IntExpr>("warp_size", "The number of threads per warp.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& gpu_warp_shuffle_down_op() { @@ -147,7 +147,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("offset", "The offset."), sig::arg<IntExpr>("width", "The width."), sig::arg<IntExpr>("warp_size", "The number of threads per warp.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& gpu_warp_shuffle_xor_op() { @@ -163,7 +163,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("width", "The width."), sig::arg<IntExpr>("warp_size", "The number of threads per warp.")) .set_native_fn_attr<&InferTypeReturnArgType<1>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& gpu_warp_activemask_op() { @@ -174,7 +174,7 @@ const Op& gpu_warp_activemask_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.gpu_warp_activemask") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::UInt(32)) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& gpu_thread_allreduce_op() { @@ -191,7 +191,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<Expr>("destinations", "The destination tensor loads."), sig::arg<Expr>("thread_axes", "The reduction thread axes.")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& gpu_dp4a_op() { @@ -205,7 +205,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<PrimExpr>("vec2", "The second input vector."), sig::arg<PrimExpr>("acc", "The accumulator.")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32)) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -216,7 +216,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("n", "The N dimension."), sig::arg<IntExpr>("k", "The K dimension."), sig::arg<IntExpr>("index", "The index."), sig::arg("buffer_ptr", "The buffer pointer."), sig::arg<IntExpr>("stride", "The stride."), sig::arg("layout", "The layout.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kReadState)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -230,7 +230,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("index_b", "The B fragment index."), sig::arg("fragment_c", "The C fragment."), sig::arg<IntExpr>("index_c", "The C fragment index.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -240,7 +240,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg("fragment", "The matrix fragment."), sig::arg<IntExpr>("m", "The M dimension."), sig::arg<IntExpr>("n", "The N dimension."), sig::arg<IntExpr>("k", "The K dimension."), sig::arg<IntExpr>("index", "The index."), sig::arg("value", "The value to use.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -251,7 +251,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<IntExpr>("n", "The N dimension."), sig::arg<IntExpr>("k", "The K dimension."), sig::arg<IntExpr>("index", "The index."), sig::arg("buffer_ptr", "The buffer pointer."), sig::arg<IntExpr>("stride", "The stride."), sig::arg("layout", "The layout.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } TVM_FFI_STATIC_INIT_BLOCK() { diff --git a/tirx/src/tirx/op/memory.cc b/tirx/src/tirx/op/memory.cc index a0373ff38f..5ee527db92 100644 --- a/tirx/src/tirx/op/memory.cc +++ b/tirx/src/tirx/op/memory.cc @@ -172,7 +172,7 @@ Expr reinterpret(Type target_ty, Expr value, Location loc) { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.reinterpret") .signature(sig::arg("x", "The input value.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& address_of_op() { @@ -184,7 +184,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.address_of") .signature(sig::arg("obj", "The referenced object.")) .set_native_fn_attr<&InferTypeAddressOf>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& isnullptr_op() { @@ -196,7 +196,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.isnullptr") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Bool()) .signature(sig::arg("x", "The input value.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& prefetch_op() { @@ -209,7 +209,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg("ptr", "The pointer."), sig::arg<IntExpr>("rw", "The read/write mode."), sig::arg<IntExpr>("locality", "The locality hint."), sig::arg<IntExpr>("cache_type", "The cache policy.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& ptr_byte_offset_op() { @@ -221,7 +221,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.ptr_byte_offset") .signature(sig::arg("data", "Base pointer."), sig::arg<IntExpr>("byte_offset", "Offset in bytes.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& static_handle_op() { @@ -231,8 +231,7 @@ const Op& static_handle_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.static_handle") - .set_attr(tvm::op_attr::kCallEffectKind, - static_cast<int64_t>(CallEffectKind::kSpecialCallArg)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kSpecialCallArg); } const Op& handle_add_byte_offset_op() { @@ -244,7 +243,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.handle_add_byte_offset") .signature(sig::arg("handle", "The handle."), sig::arg<IntExpr>("offset", "The offset.")) .set_native_fn_attr<&InferTypePointerOffset<0, -1>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& atomic_add_op() { @@ -255,7 +254,7 @@ const Op& atomic_add_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.atomic_add") .signature(sig::arg("ptr", "The pointer."), sig::arg("value", "The value to use.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& assume_aligned_op() { @@ -268,7 +267,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg<TensorVar>("tensor", "The tensor whose base address is aligned."), sig::arg<IntImm>("alignment_bytes", "The constant byte alignment.")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kEmbedInfo)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kEmbedInfo); } const Op& undef_op() { @@ -279,7 +278,7 @@ const Op& undef_op() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.undef") .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Int(32)) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kReadState)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState); } const Op& masked_load_op() { @@ -292,7 +291,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg<TensorVar>("tensor", "The tensor."), sig::arg<PrimExpr>("index", "The index."), sig::var_args<PrimExpr>("args")) .set_native_fn_attr<&InferTypeMaskedLoad>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kReadState)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState); } const Op& masked_store_op() { @@ -306,7 +305,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { sig::arg<PrimExpr>("value", "The value to use."), sig::arg<PrimExpr>("index", "The index."), sig::var_args<PrimExpr>("args")) .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void()) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kUpdateState)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kUpdateState); } const Op& alloc_tensor_op() { @@ -323,7 +322,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::var_args<tvm::Tuple>("placement", "Optional producer-owned placement."), sig::ty_arg<TensorType>("tensor_type")) .set_attrs_type_key(DictAttrsNode::_type_key) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque); } const Op& decl_tensor_op() { @@ -339,7 +338,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { true) .signature(sig::arg("data", "The existing data pointer."), sig::ty_arg<TensorType>("tensor_type")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& tensor_data_ptr_op() { @@ -351,7 +350,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.tensor_data_ptr") .signature(sig::arg("tensor", "The tensor variable.")) .set_native_fn_attr<&InferTypeTensorDataPtr>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } TVM_FFI_STATIC_INIT_BLOCK() { diff --git a/tirx/src/tirx/op/vector.cc b/tirx/src/tirx/op/vector.cc index 46c5f72fb5..225105a582 100644 --- a/tirx/src/tirx/op/vector.cc +++ b/tirx/src/tirx/op/vector.cc @@ -59,7 +59,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.vector_high") .signature(sig::arg<PrimExpr>("vec", "The input vector.")) .set_native_fn_attr<&InferTypeVectorPart<false>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& vector_low_op() { @@ -71,7 +71,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.vector_low") .signature(sig::arg<PrimExpr>("vec", "The input vector.")) .set_native_fn_attr<&InferTypeVectorPart<false>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& vector_combine_op() { @@ -84,7 +84,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .signature(sig::arg<PrimExpr>("vec1", "The first input vector."), sig::arg<PrimExpr>("vec2", "The second input vector.")) .set_native_fn_attr<&InferTypeVectorPart<true>>(tvm::op_attr::kInferType) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } const Op& get_active_lane_mask_op() { @@ -96,7 +96,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { OpDef("tirx.get_active_lane_mask") .signature(sig::arg<IntExpr>("base", "The base value."), sig::arg<IntExpr>("limit", "The limit value.")) - .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kPure)); + .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kPure); } } // namespace tirx diff --git a/tirx/tests/cpp/ir/tensor_load_test.cc b/tirx/tests/cpp/ir/tensor_load_test.cc index a74ad753d3..78accce96a 100644 --- a/tirx/tests/cpp/ir/tensor_load_test.cc +++ b/tirx/tests/cpp/ir/tensor_load_test.cc @@ -56,8 +56,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { if (load->indices.empty()) return Type::Missing(); return nullptr; })); - refl::TypeAttrDef<LoadSourceTypeNode>().attr(type_attr::kTensorLoadInferType.name, - FTensorLoadInferType::FromNative<&InferTestLoad>()); + refl::TypeAttrDef<LoadSourceTypeNode>().native_fn_attr<&InferTestLoad>( + type_attr::kTensorLoadInferType); } Var LoadSource() { diff --git a/tirx/tests/cpp/target/target_test.cc b/tirx/tests/cpp/target/target_test.cc index d8c3eb974e..d657a2fb2b 100644 --- a/tirx/tests/cpp/target/target_test.cc +++ b/tirx/tests/cpp/target/target_test.cc @@ -27,11 +27,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { TargetKindDef("TestTargetKind") .set_default_device_type(kDLCPU) .def_option<bool>("my_bool") - .def_option<ffi::Map<ffi::String, ffi::Any>>(backend::attr::kBackendConfig.name); + .def_option<ffi::Map<ffi::String, ffi::Any>>(attr::kBackendConfig.name); } TEST(TargetCreation, TypedBackendConfig) { - constexpr auto key = backend::attr::kBackendConfig; + constexpr auto key = attr::kBackendConfig; using Config = decltype(key)::ValueType; Config fallback{{"cuda", Config{{"arch", ffi::String("sm_80")}}}}; Target absent("TestTargetKind"); diff --git a/tirx/tests/python/tirx/script/test_error_handling.py b/tirx/tests/python/tirx/script/test_error_handling.py index e9a0dcdf60..69a822cec0 100644 --- a/tirx/tests/python/tirx/script/test_error_handling.py +++ b/tirx/tests/python/tirx/script/test_error_handling.py @@ -203,11 +203,12 @@ def test_tvm_exception_catch_from_scope_handler(): check_error(scope_handler_except, 2, TypeError) -def test_tvm_exception_catch_from_bare_intrin(): - def intrin_except_unassign(A: T.Tensor((16, 16), "float32")) -> None: - T.evaluate(A) # error +def test_evaluate_tensor(): + @T.function + def evaluate_tensor(A: T.Tensor((16, 16), "float32")) -> None: + T.evaluate(A) - check_error(intrin_except_unassign, 2, RuntimeError) + assert evaluate_tensor.body.seq[0].value.same_as(evaluate_tensor.params[0]) def test_tvm_exception_catch_from_assigned_intrin():
