This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 4d98064b4f [REFACTOR][IR] Colocate typed attribute contracts (#20795)
4d98064b4f is described below
commit 4d98064b4f5b5d6595451f1dc40abbcb0d0c6e4d
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.
---
src/relax/ir/dependent_type.cc | 29 +-
src/relax/utils.cc | 2 +-
src/s_tir/analysis/is_pure_function.cc | 2 +-
.../feature_extractor/per_store_feature.cc | 3 +-
src/s_tir/stmt.cc | 4 +-
src/s_tir/transform/rewrite_unsafe_select.cc | 5 +-
src/s_tir/transform/storage_access.cc | 2 +-
src/te/tensor.cc | 2 +-
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 +-
49 files changed, 680 insertions(+), 669 deletions(-)
diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc
index 4430f7e527..2134064bc3 100644
--- a/src/relax/ir/dependent_type.cc
+++ b/src/relax/ir/dependent_type.cc
@@ -272,23 +272,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::FStructuralMutate::FromNative<&TensorTypeMutate>())
.attr(ffi::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&TensorTypeMaybeInplaceMutate>())
- .def(tvm::type_attr::kSubscriptExprRealize.name,
- [](Expr value,
- ffi::Array<ffi::Variant<ffi::Tuple<ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>,
- ffi::Optional<PrimExpr>>,
- PrimExpr>>
- slice,
- Location loc) -> ffi::ObjectRef {
- TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError)
- << "A Relax expression requires exactly one index";
- auto index = slice[0].as<PrimExpr>();
- TVM_FFI_CHECK(index.has_value(), TypeError)
- << "A Relax expression requires a point index";
- const auto* imm = index.value().as<IntImmNode>();
- TVM_FFI_CHECK(imm != nullptr, TypeError)
- << "A Relax expression requires a constant integer index";
- return TupleGetItem(value, imm->value.as<int>().value(), loc);
- });
+ .attr(tvm::type_attr::kExprSubscriptRealizeByTy,
+ [](Expr value, SubscriptSlice slice, Location loc) ->
ffi::ObjectRef {
+ TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError)
+ << "A Relax expression requires exactly one index";
+ auto index = slice[0].as<PrimExpr>();
+ TVM_FFI_CHECK(index.has_value(), TypeError)
+ << "A Relax expression requires a point index";
+ const auto* imm = index.value().as<IntImmNode>();
+ TVM_FFI_CHECK(imm != nullptr, TypeError)
+ << "A Relax expression requires a constant integer index";
+ return TupleGetItem(value, imm->value.as<int>().value(), loc);
+ });
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/utils.cc b/src/relax/utils.cc
index 4e1c6bbe6d..e0cf657da4 100644
--- a/src/relax/utils.cc
+++ b/src/relax/utils.cc
@@ -247,7 +247,7 @@ bool IsImpureCall(const Call& call) {
static auto effect_map = Op::GetAttrMap(tvm::op_attr::kCallEffectKind);
TVM_FFI_ICHECK(effect_map.count(op))
<< "Cannot find the registered purity or call effect of this op: " <<
op->name;
- auto effect = static_cast<CallEffectKind>(effect_map[op]);
+ auto effect = effect_map[op];
return effect > CallEffectKind::kPure;
}
TVM_FFI_CHECK(!call->op->ty.as<tvm::FuncTypeNode>(), TypeError)
diff --git a/src/s_tir/analysis/is_pure_function.cc
b/src/s_tir/analysis/is_pure_function.cc
index a4084d69f5..873a7a224c 100644
--- a/src/s_tir/analysis/is_pure_function.cc
+++ b/src/s_tir/analysis/is_pure_function.cc
@@ -105,7 +105,7 @@ class PurityChecker : public StmtExprVisitor {
static auto op_call_effect = Op::GetAttrMap(tvm::op_attr::kCallEffectKind);
CallEffectKind effect = [&]() {
if (auto opt = call->op.as<Op>()) {
- return static_cast<CallEffectKind>(op_call_effect[opt.value()]);
+ return op_call_effect[opt.value()];
} else {
return CallEffectKind::kOpaque;
}
diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
index 07b4990769..e42aaa826d 100644
--- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
@@ -694,8 +694,7 @@ Feature::ArithOps::ArithOps(const TensorStoreNode* store,
int64_t prod_loop_exte
return StmtExprVisitor::Visit_(op);
}
static auto op_call_effect_ =
Op::GetAttrMap(tvm::op_attr::kCallEffectKind);
- CallEffectKind effect_kind =
-
static_cast<CallEffectKind>(op_call_effect_[op->op.as_or_throw<Op>()]);
+ CallEffectKind effect_kind = op_call_effect_[op->op.as_or_throw<Op>()];
bool is_pure =
effect_kind == CallEffectKind::kPure || effect_kind ==
CallEffectKind::kExprAnnotation;
if (is_pure) {
diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc
index 0b196a9c0f..979adb4456 100644
--- a/src/s_tir/stmt.cc
+++ b/src/s_tir/stmt.cc
@@ -69,11 +69,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("s_tir.async_commit", "Commit asynchronous copies to a queue.")
.signature(sig::arg<IntImm>("queue_id"))
.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);
OpDef("s_tir.async_wait", "Wait for committed asynchronous copies.")
.signature(sig::arg<IntImm>("queue_id"),
sig::arg<PrimExpr>("inflight_count"))
.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);
}
namespace {
diff --git a/src/s_tir/transform/rewrite_unsafe_select.cc
b/src/s_tir/transform/rewrite_unsafe_select.cc
index ff2e90f4f5..0ee073315a 100644
--- a/src/s_tir/transform/rewrite_unsafe_select.cc
+++ b/src/s_tir/transform/rewrite_unsafe_select.cc
@@ -56,7 +56,7 @@ class UnsafeExprDetector : public tvm::ExprFunctor<bool(const
Expr& n)> {
}
return Dispatch(op->args[0]);
} else if (auto opt = op->op.as<Op>()) {
- auto effect_kind =
static_cast<CallEffectKind>(op_call_effect_[opt.value()]);
+ auto effect_kind = op_call_effect_[opt.value()];
if (effect_kind == CallEffectKind::kPure || effect_kind ==
CallEffectKind::kExprAnnotation) {
for (const Expr& arg : op->args) {
if (Dispatch(arg)) return true;
@@ -122,8 +122,7 @@ class UnsafeExprDetector : public
tvm::ExprFunctor<bool(const Expr& n)> {
return Dispatch(op->a) || Dispatch(op->b);
}
- decltype(Op::GetAttrMap(tvm::op_attr::kCallEffectKind)) op_call_effect_ =
- Op::GetAttrMap(tvm::op_attr::kCallEffectKind);
+ OpAttrMap<CallEffectKind> op_call_effect_ =
Op::GetAttrMap(tvm::op_attr::kCallEffectKind);
};
class UnsafeSelectRewriter : public StmtExprMutator {
diff --git a/src/s_tir/transform/storage_access.cc
b/src/s_tir/transform/storage_access.cc
index 5ecc9878f1..e47bd8a5fa 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -396,7 +396,7 @@ ffi::Optional<VisitInterrupt>
StorageAccessVisitor::Visit_(const CallNode* op) {
auto effect_map = Op::GetAttrMap(tvm::op_attr::kCallEffectKind);
auto callee = op->op.as<Op>();
if (!callee.has_value() || !effect_map.count(callee.value()) ||
- effect_map[callee.value()] > static_cast<int>(CallEffectKind::kPure)) {
+ effect_map[callee.value()] > CallEffectKind::kPure) {
in_opaque_call_ = true;
}
auto result = StmtExprVisitor::Visit_(op);
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index d1a41975ed..01168cf97b 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -104,7 +104,7 @@ PrimVar var(std::string name_hint, PrimType t) { return
PrimVar(name_hint, t); }
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("te.tensor_load")
.signature(sig::arg("tensor", "The input tensor."),
sig::var_args("args"))
- .set_attr(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kReadState));
+ .set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kReadState);
}
// Tensor
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():
