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 c7931c6200 [REFACTOR][IR] Add scalar integer expression signatures
(#20521)
c7931c6200 is described below
commit c7931c6200924b1d049bf6ff5dd5d89613b537ca
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 2 13:49:42 2026 -0400
[REFACTOR][IR] Add scalar integer expression signatures (#20521)
Introduce IntExpr as a scalar signed or unsigned integer view for
operator signatures while retaining PrimExpr implementation APIs and
integer literal conversions. Update tvm-ffi and simplify expression
traits to use inherited casting with specialized semantic checks.
---
3rdparty/tvm-ffi | 2 +-
include/tvm/ir/base_expr.h | 112 ++++++------
src/backend/cuda/codegen/intrin_rule_cuda.cc | 24 +--
src/backend/cuda/op/iket.cc | 2 +-
src/backend/cuda/op/target_builtin.cc | 222 +++++++++++++----------
src/backend/metal/codegen/intrin_rule_metal.cc | 6 +-
src/backend/metal/op/target_builtin.cc | 30 +--
src/backend/webgpu/codegen/intrin_rule_webgpu.cc | 6 +-
src/relax/op/op.cc | 46 ++---
src/relax/op/tensor/create.cc | 8 +-
src/relax/op/tensor/inspect.cc | 4 +-
src/relax/op/tensor/manipulate.cc | 10 +-
src/tirx/op/builtin.cc | 155 +++++++++-------
src/tirx/op/op.cc | 15 +-
tests/cpp/expr_test.cc | 21 +++
tests/python/tirx-base/test_tir_op_types.py | 11 ++
16 files changed, 378 insertions(+), 296 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index c96f9e31b2..df463f9bd2 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit c96f9e31b279ddfecf543539263c0f3e082feeeb
+Subproject commit df463f9bd269867fc7f4d578d62d80cc2a11d004
diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
index 87855537b4..4302a570b2 100644
--- a/include/tvm/ir/base_expr.h
+++ b/include/tvm/ir/base_expr.h
@@ -427,6 +427,18 @@ class PrimExpr : public TypedExpr<PrimType> {
static constexpr bool _type_container_is_exact = false;
};
+/*!
+ * \brief View over a scalar signed or unsigned integer expression of any bit
width.
+ *
+ * IntExpr validates integer operands in operator signatures without changing
the
+ * underlying expression node or its PrimExpr implementation API.
+ */
+class IntExpr : public PrimExpr {
+ public:
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntExpr, PrimExpr, ExprNode);
+ static constexpr bool _type_container_is_exact = false;
+};
+
/*!
* \brief Base class for other IR constructs that can be converted to PrimExpr.
* This is useful for the FFI to convert the expressions to PrimExpr.
@@ -464,15 +476,6 @@ inline constexpr bool
use_default_type_traits_v<TypedExpr<ExpectedType>> = false
template <typename ExpectedType>
struct TypeTraits<TypedExpr<ExpectedType>>
: public ObjectRefTypeTraitsBase<TypedExpr<ExpectedType>> {
- using Base = ObjectRefTypeTraitsBase<TypedExpr<ExpectedType>>;
- using Base::CopyFromAnyViewAfterCheck;
- using Base::CopyToAnyView;
- using Base::GetMismatchTypeInfo;
- using Base::MoveFromAnyAfterCheck;
- using Base::MoveToAny;
- using Base::TypeSchema;
- using Base::TypeStr;
-
TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
if (src->type_index == TypeIndex::kTVMFFINone) {
return TypedExpr<ExpectedType>::_type_is_nullable;
@@ -484,58 +487,25 @@ struct TypeTraits<TypedExpr<ExpectedType>>
// Non-owning: this only reads `ty`, and the owning form's incref/decref
pair costs two
// atomics per check on a path every typed field assignment takes.
const auto* expr =
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
- return
details::AnyUnsafe::CheckAnyViewStrict<ExpectedType>(AnyView(expr->ty));
- }
-
- TVM_FFI_INLINE static std::optional<TypedExpr<ExpectedType>>
TryCastFromAnyView(
- const TVMFFIAny* src) {
- if (CheckAnyStrict(src)) {
- if (src->type_index == TypeIndex::kTVMFFINone) {
- return
details::ObjectUnsafe::ObjectRefFromObjectPtr<TypedExpr<ExpectedType>>(nullptr);
- }
- return
details::ObjectUnsafe::ObjectRefFromObjectPtr<TypedExpr<ExpectedType>>(
- details::ObjectUnsafe::ObjectPtrFromUnowned<ExprNode>(src->v_obj));
- }
- return std::nullopt;
+ return details::AnyUnsafe::CheckAnyViewStrict<ExpectedType>(expr->ty);
}
};
template <>
inline constexpr bool use_default_type_traits_v<PrimExpr> = false;
-template <typename ObjectRefType, typename ExpectedType, typename...
FallbackTypes>
-struct TypedExprWithFallbackTraitsBase
- : public ObjectRefWithFallbackTraitsBase<ObjectRefType, FallbackTypes...> {
- using Base = ObjectRefWithFallbackTraitsBase<ObjectRefType,
FallbackTypes...>;
-
- TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
- return TypeTraits<TypedExpr<ExpectedType>>::CheckAnyStrict(src);
- }
-
- TVM_FFI_INLINE static std::optional<ObjectRefType> TryCastFromAnyView(const
TVMFFIAny* src) {
- if (TypeTraits<TypedExpr<ExpectedType>>::TryCastFromAnyView(src)) {
- return details::ObjectUnsafe::ObjectRefFromObjectPtr<ObjectRefType>(
- details::ObjectUnsafe::ObjectPtrFromUnowned<ExprNode>(src->v_obj));
- }
- return Base::template TryFallbackTypes<FallbackTypes...>(src);
- }
-};
-
// define automatic conversion from bool, int64_t, double to PrimExpr
// These functions are declared early to avoid circular dependency
template <>
-struct TypeTraits<PrimExpr>
- : public TypedExprWithFallbackTraitsBase<PrimExpr, PrimType, StrictBool,
int64_t, double,
- PrimExprConvertible> {
- using Base = TypedExprWithFallbackTraitsBase<PrimExpr, PrimType, StrictBool,
int64_t, double,
- PrimExprConvertible>;
- using Base::CheckAnyStrict;
- using Base::CopyFromAnyViewAfterCheck;
- using Base::CopyToAnyView;
- using Base::MoveFromAnyAfterCheck;
- using Base::MoveToAny;
- using Base::TryCastFromAnyView;
- using Base::TypeSchema;
+struct TypeTraits<PrimExpr> : public ObjectRefWithFallbackTraitsBase<PrimExpr,
StrictBool, int64_t,
+ double,
PrimExprConvertible> {
+ using Base =
+ ObjectRefWithFallbackTraitsBase<PrimExpr, StrictBool, int64_t, double,
PrimExprConvertible>;
+
+ 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::string TypeStr() { return "ir.PrimExpr"; }
@@ -543,8 +513,7 @@ struct TypeTraits<PrimExpr>
if (source->type_index >= TypeIndex::kTVMFFIStaticObjectBegin &&
source->v_obj != nullptr &&
details::IsObjectInstance<ExprNode>(source->type_index)) {
const auto* expr =
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(source->v_obj);
- if (expr->ty.defined() &&
-
!details::AnyUnsafe::CheckAnyViewStrict<PrimType>(AnyView(expr->ty))) {
+ if (expr->ty.defined() &&
!details::AnyUnsafe::CheckAnyViewStrict<PrimType>(expr->ty)) {
return TypeIndexToTypeKey(source->type_index) + "[ty=" +
expr->ty.GetTypeKey() + "]";
}
}
@@ -559,6 +528,41 @@ struct TypeTraits<PrimExpr>
}
};
+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) {
+ if (src->type_index >= TypeIndex::kTVMFFIStaticObjectBegin &&
+ details::IsObjectInstance<ExprNode>(src->type_index)) {
+ const auto* expr =
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
+ if (const auto* ty = expr->ty.as<PrimTypeNode>()) {
+ return TypeIndexToTypeKey(src->type_index) +
+ "[dtype=" + ffi::DLDataTypeToString(ty->dtype) + "]";
+ }
+ }
+ return TypeTraits<PrimExpr>::GetMismatchTypeInfo(src);
+ }
+};
+
template <>
inline constexpr bool use_default_type_traits_v<Expr> = false;
diff --git a/src/backend/cuda/codegen/intrin_rule_cuda.cc
b/src/backend/cuda/codegen/intrin_rule_cuda.cc
index bbbbcca4f7..8b9a96b929 100644
--- a/src/backend/cuda/codegen/intrin_rule_cuda.cc
+++ b/src/backend/cuda/codegen/intrin_rule_cuda.cc
@@ -281,10 +281,10 @@ void RegisterCudaIntrinRules() {
// Register low-level CUDA device intrinsics.
// TODO(tvm-team): consider make CUDA its own subfolder and create a file
for low-level builtins.
OpDef("tirx.cuda.__shfl_sync")
- .signature(sig::arg("mask", "The thread mask."),
+ .signature(sig::arg<IntExpr>("mask", "The thread mask."),
sig::arg("var", "The variable to sync."),
- sig::arg("lane", "The source thread id."),
- sig::arg("width", "The warp thread width, must be a power of 2."))
+ sig::arg<IntExpr>("lane", "The source thread id."),
+ sig::arg<IntExpr>("width", "The warp thread width, must be a power
of 2."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("cuda"))
@@ -295,10 +295,10 @@ void RegisterCudaIntrinRules() {
.set_attr<bool>("cuda.need_warp_shuffle", true);
OpDef("tirx.cuda.__shfl_up_sync")
- .signature(sig::arg("mask", "The thread mask."),
+ .signature(sig::arg<IntExpr>("mask", "The thread mask."),
sig::arg("var", "The variable to sync."),
- sig::arg("delta", "The source lane id offset to be added."),
- sig::arg("width", "The warp thread width, must be a power of 2."))
+ sig::arg<IntExpr>("delta", "The source lane id offset to be added."),
+ sig::arg<IntExpr>("width", "The warp thread width, must be a power
of 2."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("cuda"))
@@ -309,10 +309,10 @@ void RegisterCudaIntrinRules() {
.set_attr<bool>("cuda.need_warp_shuffle", true);
OpDef("tirx.cuda.__shfl_down_sync")
- .signature(sig::arg("mask", "The thread mask."),
+ .signature(sig::arg<IntExpr>("mask", "The thread mask."),
sig::arg("var", "The variable to sync."),
- sig::arg("delta", "The source lane id offset to be subtracted."),
- sig::arg("width", "The warp thread width, must be a power of 2."))
+ sig::arg<IntExpr>("delta", "The source lane id offset to be
subtracted."),
+ sig::arg<IntExpr>("width", "The warp thread width, must be a power
of 2."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("cuda"))
@@ -323,10 +323,10 @@ void RegisterCudaIntrinRules() {
.set_attr<bool>("cuda.need_warp_shuffle", true);
OpDef("tirx.cuda.__shfl_xor_sync")
- .signature(sig::arg("mask", "The thread mask."),
+ .signature(sig::arg<IntExpr>("mask", "The thread mask."),
sig::arg("var", "The variable to sync."),
- sig::arg("lane_mask", "The lane mask."),
- sig::arg("width", "The warp thread width, must be a power of 2."))
+ sig::arg<IntExpr>("lane_mask", "The lane mask."),
+ sig::arg<IntExpr>("width", "The warp thread width, must be a power
of 2."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("cuda"))
diff --git a/src/backend/cuda/op/iket.cc b/src/backend/cuda/op/iket.cc
index 0186402fa4..2e8c4fdcb4 100644
--- a/src/backend/cuda/op/iket.cc
+++ b/src/backend/cuda/op/iket.cc
@@ -45,7 +45,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_range_end")
- .signature(sig::arg("token", "The token."), sig::var_args("args"))
+ .signature(sig::arg<IntExpr>("token", "The token."),
sig::var_args("args"))
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("cuda"))
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index 73b511e96a..4ac2b7b8ee 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -61,48 +61,57 @@ void RegisterCudaTargetBuiltins() {
registered = true;
OpDef("tirx.tvm_load_matrix_sync")
- .signature(sig::arg("fragment", "The matrix fragment."), sig::arg("m",
"The M dimension."),
- sig::arg("n", "The N dimension."), sig::arg("k", "The K
dimension."),
- sig::arg("index", "The index."), sig::arg("buffer_ptr", "The
buffer pointer."),
- sig::arg("stride", "The stride."), sig::arg("layout", "The
layout."))
+ .signature(
+ 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("buffer_ptr",
"The buffer pointer."),
+ sig::arg<IntExpr>("stride", "The stride."), sig::arg("layout", "The
layout."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_load_matrix_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState));
OpDef("tirx.tvm_mma_sync")
- .signature(
- sig::arg("fragment_d", "The D fragment."), sig::arg("index_d", "The
D fragment index."),
- sig::arg("fragment_a", "The A fragment."), sig::arg("index_a", "The
A fragment index."),
- sig::arg("fragment_b", "The B fragment."), sig::arg("index_b", "The
B fragment index."),
- sig::arg("fragment_c", "The C fragment."), sig::arg("index_c", "The
C fragment index."))
+ .signature(sig::arg("fragment_d", "The D fragment."),
+ sig::arg<IntExpr>("index_d", "The D fragment index."),
+ sig::arg("fragment_a", "The A fragment."),
+ sig::arg<IntExpr>("index_a", "The A fragment index."),
+ sig::arg("fragment_b", "The B fragment."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_mma_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_bmma_sync")
- .signature(
- sig::arg("fragment_d", "The D fragment."), sig::arg("index_d", "The
D fragment index."),
- sig::arg("fragment_a", "The A fragment."), sig::arg("index_a", "The
A fragment index."),
- sig::arg("fragment_b", "The B fragment."), sig::arg("index_b", "The
B fragment index."),
- sig::arg("fragment_c", "The C fragment."), sig::arg("index_c", "The
C fragment index."))
+ .signature(sig::arg("fragment_d", "The D fragment."),
+ sig::arg<IntExpr>("index_d", "The D fragment index."),
+ sig::arg("fragment_a", "The A fragment."),
+ sig::arg<IntExpr>("index_a", "The A fragment index."),
+ sig::arg("fragment_b", "The B fragment."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_bmma_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_fill_fragment")
- .signature(sig::arg("fragment", "The matrix fragment."), sig::arg("m",
"The M dimension."),
- sig::arg("n", "The N dimension."), sig::arg("k", "The K
dimension."),
- sig::arg("index", "The index."), sig::arg("value", "The value
to use."))
+ .signature(
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_fill_fragment"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_store_matrix_sync")
- .signature(sig::arg("fragment", "The matrix fragment."), sig::arg("m",
"The M dimension."),
- sig::arg("n", "The N dimension."), sig::arg("k", "The K
dimension."),
- sig::arg("index", "The index."), sig::arg("buffer_ptr", "The
buffer pointer."),
- sig::arg("stride", "The stride."), sig::arg("layout", "The
layout."))
+ .signature(
+ 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("buffer_ptr",
"The buffer pointer."),
+ sig::arg<IntExpr>("stride", "The stride."), sig::arg("layout", "The
layout."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_store_matrix_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -113,18 +122,20 @@ void RegisterCudaTargetBuiltins() {
// to its thread-local index. Used by the s_tir tensor_intrin tensorize
// path so per-thread fragment offsets stay element-accurate.
OpDef("tirx.mma_store_legacy")
- .signature(sig::arg("m", "The M dimension."), sig::arg("n", "The N
dimension."),
+ .signature(sig::arg<IntExpr>("m", "The M dimension."),
+ sig::arg<IntExpr>("n", "The N dimension."),
sig::arg("dst_ptr", "The destination pointer."),
sig::arg("src_ptr", "The source pointer."),
- sig::arg("src_offset", "The source offset."),
- sig::arg("dst_stride", "The destination stride."))
+ sig::arg<IntExpr>("src_offset", "The source offset."),
+ sig::arg<IntExpr>("dst_stride", "The destination stride."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.mma_store_legacy"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.mma_fill_legacy")
- .signature(sig::arg("local_size", "The local allocation size."),
- sig::arg("local_ptr", "The local pointer."),
sig::arg("offset", "The offset."))
+ .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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.mma_fill_legacy"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -134,10 +145,10 @@ void RegisterCudaTargetBuiltins() {
// and prints it dtype-first; user-issued copies go through T.ptx instead.
OpDef("tirx.s_tir.cp_async_raw")
.signature(sig::arg("dst_ptr", "The destination pointer."),
- sig::arg("dst_offset", "The destination offset."),
+ sig::arg<IntExpr>("dst_offset", "The destination offset."),
sig::arg("src_ptr", "The source pointer."),
- sig::arg("src_offset", "The source offset."),
- sig::arg("cp_size", "The copy size."), sig::var_args("args"))
+ sig::arg<IntExpr>("src_offset", "The source offset."),
+ sig::arg<IntExpr>("cp_size", "The copy size."),
sig::var_args("args"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("s_tir"))
@@ -146,11 +157,12 @@ void RegisterCudaTargetBuiltins() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.mma_store")
- .signature(sig::arg("m", "The M dimension."), sig::arg("n", "The N
dimension."),
+ .signature(sig::arg<IntExpr>("m", "The M dimension."),
+ sig::arg<IntExpr>("n", "The N dimension."),
sig::arg("dst_ptr", "The destination pointer."),
sig::arg("src_ptr", "The source pointer."),
- sig::arg("src_offset", "The source offset."),
- sig::arg("dst_stride", "The destination stride."))
+ sig::arg<IntExpr>("src_offset", "The source offset."),
+ sig::arg<IntExpr>("dst_stride", "The destination stride."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.mma_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
@@ -158,8 +170,9 @@ void RegisterCudaTargetBuiltins() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.mma_fill")
- .signature(sig::arg("local_size", "The local allocation size."),
- sig::arg("local_ptr", "The local pointer."),
sig::arg("offset", "The offset."))
+ .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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.mma_fill"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
@@ -169,9 +182,9 @@ void RegisterCudaTargetBuiltins() {
OpDef("tirx.timer_init_cuda")
.signature(sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
- sig::arg("profiler_write_offset", "The profiler write
offset."),
- sig::arg("num_groups", "The number of groups."),
- sig::arg("group_id", "The group identifier."))
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.timer_init"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -180,8 +193,8 @@ void RegisterCudaTargetBuiltins() {
.signature(sig::arg("event_type", "The event type."),
sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
- sig::arg("profiler_write_offset", "The profiler write
offset."),
- sig::arg("profiler_write_stride", "The profiler write
stride."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.timer_start"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
@@ -191,8 +204,8 @@ void RegisterCudaTargetBuiltins() {
.signature(sig::arg("event_type", "The event type."),
sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
- sig::arg("profiler_write_offset", "The profiler write
offset."),
- sig::arg("profiler_write_stride", "The profiler write
stride."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.timer_end"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
@@ -201,8 +214,8 @@ void RegisterCudaTargetBuiltins() {
OpDef("tirx.timer_finalize_cuda")
.signature(sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
- sig::arg("profiler_write_offset", "The profiler write
offset."),
- sig::arg("profiler_write_stride", "The profiler write
stride."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.timer_finalize"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
@@ -261,7 +274,7 @@ void RegisterDeviceIntrinsic(const char* op_name, const
char* op_namespace,
}
void RegisterDeviceIntrinsicAliases() {
- RegisterDeviceIntrinsic("cuda_any_sync", "cuda", CallEffectKind::kPure,
sig::arg("mask"),
+ RegisterDeviceIntrinsic("cuda_any_sync", "cuda", CallEffectKind::kPure,
sig::arg<IntExpr>("mask"),
sig::arg("pred"));
RegisterDeviceIntrinsic("cuda_atomic_add", "cuda", CallEffectKind::kOpaque,
sig::arg("res_addr"),
sig::arg("value"));
@@ -269,16 +282,16 @@ void RegisterDeviceIntrinsicAliases() {
sig::arg("old_val"), sig::arg("new_val"));
RegisterDeviceIntrinsic("cuda_wait_until", "cuda", CallEffectKind::kOpaque,
sig::arg("dst"),
sig::arg("ptr"), sig::arg("condition"),
sig::arg("scope"),
- sig::arg("space"), sig::arg("ptx_type"),
sig::arg("backoff_ns"));
- RegisterDeviceIntrinsic("cuda_ballot_sync", "cuda", CallEffectKind::kOpaque,
sig::arg("mask"),
- sig::arg("pred"));
+ sig::arg("space"), sig::arg("ptx_type"),
sig::arg<IntExpr>("backoff_ns"));
+ RegisterDeviceIntrinsic("cuda_ballot_sync", "cuda", CallEffectKind::kOpaque,
+ sig::arg<IntExpr>("mask"), sig::arg("pred"));
RegisterDeviceIntrinsic("cuda_bfloat1622float2", "cuda",
CallEffectKind::kOpaque,
sig::arg("packed"));
RegisterDeviceIntrinsic("cuda_bfloat162float", "cuda",
CallEffectKind::kOpaque, sig::arg("src"));
RegisterDeviceIntrinsic("cuda_clock64", "cuda", CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("cuda_cluster_sync", "cuda",
CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("cuda_cta_reduce", "cuda", CallEffectKind::kOpaque,
sig::arg("value"),
- sig::arg("op"), sig::arg("num_warps"),
sig::arg("scratch"));
+ sig::arg("op"), sig::arg<IntExpr>("num_warps"),
sig::arg("scratch"));
RegisterDeviceIntrinsic("cuda_cta_sync", "cuda", CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("cuda_cvta_generic_to_shared", "cuda",
CallEffectKind::kOpaque,
sig::arg("ptr"));
@@ -287,7 +300,8 @@ void RegisterDeviceIntrinsicAliases() {
sig::arg("b"));
RegisterDeviceIntrinsic("cuda_fdividef", "cuda", CallEffectKind::kPure,
sig::arg("x"),
sig::arg("y"));
- RegisterDeviceIntrinsic("cuda_ffs_u32", "cuda", CallEffectKind::kOpaque,
sig::arg("value"));
+ RegisterDeviceIntrinsic("cuda_ffs_u32", "cuda", CallEffectKind::kOpaque,
+ sig::arg<IntExpr>("value"));
RegisterDeviceIntrinsic("cuda_float22bfloat162_rn", "cuda",
CallEffectKind::kOpaque,
sig::arg("v0"), sig::arg("v1"));
RegisterDeviceIntrinsic("cuda_float22bfloat162_rn_from_float2", "cuda",
CallEffectKind::kOpaque,
@@ -306,7 +320,7 @@ void RegisterDeviceIntrinsicAliases() {
RegisterDeviceIntrinsic("cuda_func_call", "cuda", CallEffectKind::kOpaque,
sig::arg("func_name"),
sig::var_args("args"));
RegisterDeviceIntrinsic("cuda_get_tmem_addr", "cuda",
CallEffectKind::kOpaque, sig::arg("addr"),
- sig::arg("row_offset"), sig::arg("col_offset"));
+ sig::arg<IntExpr>("row_offset"),
sig::arg<IntExpr>("col_offset"));
RegisterDeviceIntrinsic("cuda_grid_sync", "cuda", CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("cuda_half2float", "cuda", CallEffectKind::kOpaque,
sig::arg("src"));
RegisterDeviceIntrinsic("cuda_half8tofloat8", "cuda",
CallEffectKind::kOpaque,
@@ -319,102 +333,110 @@ void RegisterDeviceIntrinsicAliases() {
RegisterDeviceIntrinsic("cuda_make_float2", "cuda", CallEffectKind::kOpaque,
sig::arg("x"),
sig::arg("y"));
RegisterDeviceIntrinsic("cuda_mbarrier_wait", "cuda",
CallEffectKind::kOpaque, sig::arg("bar"),
- sig::arg("phase"));
+ sig::arg<IntExpr>("phase"));
RegisterDeviceIntrinsic("cuda_mbarrier_wait_acquire_cluster", "cuda",
CallEffectKind::kOpaque,
- sig::arg("bar"), sig::arg("phase"));
- RegisterDeviceIntrinsic("cuda_mov_sreg", "cuda", CallEffectKind::kPure,
sig::arg("bits"),
+ sig::arg("bar"), sig::arg<IntExpr>("phase"));
+ RegisterDeviceIntrinsic("cuda_mov_sreg", "cuda", CallEffectKind::kPure,
sig::arg<IntExpr>("bits"),
sig::arg("reg_name"));
- RegisterDeviceIntrinsic("cuda_nano_sleep", "cuda", CallEffectKind::kOpaque,
sig::arg("time"));
+ RegisterDeviceIntrinsic("cuda_nano_sleep", "cuda", CallEffectKind::kOpaque,
+ sig::arg<IntExpr>("time"));
RegisterDeviceIntrinsic("cuda_printf", "cuda", CallEffectKind::kOpaque,
sig::arg("fmt"),
sig::var_args("args"));
RegisterDeviceIntrinsic("cuda_reduce_add_sync_u32", "cuda",
CallEffectKind::kOpaque,
- sig::arg("mask"), sig::arg("value"));
+ sig::arg<IntExpr>("mask"),
sig::arg<IntExpr>("value"));
RegisterDeviceIntrinsic("cuda_reduce_min_sync_u32", "cuda",
CallEffectKind::kOpaque,
- sig::arg("mask"), sig::arg("value"));
+ sig::arg<IntExpr>("mask"),
sig::arg<IntExpr>("value"));
RegisterDeviceIntrinsic("cuda_runtime_instr_desc", "cuda",
CallEffectKind::kOpaque,
- sig::arg("desc"), sig::arg("sf_id"));
+ sig::arg("desc"), sig::arg<IntExpr>("sf_id"));
RegisterDeviceIntrinsic("cuda_sm100_2sm_leader_smem_addr", "cuda",
CallEffectKind::kOpaque,
sig::arg("ptr"));
RegisterDeviceIntrinsic("cuda_smem_addr_from_uint64", "cuda",
CallEffectKind::kOpaque,
- sig::arg("cluster_addr"));
+ sig::arg<IntExpr>("cluster_addr"));
RegisterDeviceIntrinsic("cuda_syncthreads_and", "cuda",
CallEffectKind::kOpaque,
sig::arg("cond"));
RegisterDeviceIntrinsic("cuda_syncthreads_or", "cuda",
CallEffectKind::kOpaque, sig::arg("cond"));
- RegisterDeviceIntrinsic(
- "cuda_tcgen05_encode_instr_descriptor", "cuda", CallEffectKind::kOpaque,
sig::arg("desc"),
- sig::arg("d_dtype"), sig::arg("a_dtype"), sig::arg("b_dtype"),
sig::arg("M"), sig::arg("N"),
- sig::arg("K"), sig::arg("trans_a"), sig::arg("trans_b"),
sig::arg("n_cta_groups"),
- sig::arg("neg_a"), sig::arg("neg_b"), sig::arg("sat_d"),
sig::arg("is_sparse"));
+ RegisterDeviceIntrinsic("cuda_tcgen05_encode_instr_descriptor", "cuda",
CallEffectKind::kOpaque,
+ sig::arg("desc"), sig::arg("d_dtype"),
sig::arg("a_dtype"),
+ sig::arg("b_dtype"), sig::arg<IntExpr>("M"),
sig::arg<IntExpr>("N"),
+ sig::arg<IntExpr>("K"), sig::arg("trans_a"),
sig::arg("trans_b"),
+ sig::arg<IntExpr>("n_cta_groups"),
sig::arg("neg_a"), sig::arg("neg_b"),
+ sig::arg("sat_d"), sig::arg("is_sparse"));
RegisterDeviceIntrinsic("cuda_tcgen05_encode_instr_descriptor_block_scaled",
"cuda",
CallEffectKind::kOpaque, sig::arg("desc"),
sig::arg("d_dtype"),
sig::arg("a_dtype"), sig::arg("b_dtype"),
sig::arg("sfa_dtype"),
sig::arg("sfb_dtype"), sig::arg("sfa_tmem_addr"),
- sig::arg("sfb_tmem_addr"), sig::arg("M"),
sig::arg("N"), sig::arg("K"),
- sig::arg("trans_a"), sig::arg("trans_b"),
sig::arg("n_cta_groups"),
- sig::arg("neg_a"), sig::arg("neg_b"),
sig::arg("is_sparse"));
+ sig::arg("sfb_tmem_addr"), sig::arg<IntExpr>("M"),
sig::arg<IntExpr>("N"),
+ sig::arg<IntExpr>("K"), sig::arg("trans_a"),
sig::arg("trans_b"),
+ sig::arg<IntExpr>("n_cta_groups"),
sig::arg("neg_a"), sig::arg("neg_b"),
+ sig::arg("is_sparse"));
RegisterDeviceIntrinsic("cuda_tcgen05_encode_matrix_descriptor", "cuda",
CallEffectKind::kOpaque,
- sig::arg("desc"), sig::arg("addr"), sig::arg("ldo"),
sig::arg("sdo"),
- sig::arg("swizzle"));
+ sig::arg("desc"), sig::arg("addr"),
sig::arg<IntExpr>("ldo"),
+ sig::arg<IntExpr>("sdo"),
sig::arg<IntExpr>("swizzle"));
RegisterDeviceIntrinsic("cuda_thread_fence", "cuda",
CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("cuda_thread_rank", "cuda", CallEffectKind::kPure);
RegisterDeviceIntrinsic("cuda_trap_when_assert_failed", "cuda",
CallEffectKind::kOpaque,
sig::arg("cond"));
- RegisterDeviceIntrinsic("cuda_uint_as_float", "cuda",
CallEffectKind::kOpaque, sig::arg("bits"));
+ RegisterDeviceIntrinsic("cuda_uint_as_float", "cuda",
CallEffectKind::kOpaque,
+ sig::arg<IntExpr>("bits"));
RegisterDeviceIntrinsic("cuda_warp_reduce", "cuda", CallEffectKind::kOpaque,
sig::arg("value"),
- sig::arg("op"), sig::arg("width"));
+ sig::arg("op"), sig::arg<IntExpr>("width"));
RegisterDeviceIntrinsic("cuda_warp_sync", "cuda", CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("cuda_warpgroup_sync", "cuda",
CallEffectKind::kOpaque,
- sig::arg("bar_no"));
+ sig::arg<IntExpr>("bar_no"));
RegisterDeviceIntrinsic("cuda_wgmma_encode_matrix_descriptor", "cuda",
CallEffectKind::kOpaque,
- sig::arg("desc"), sig::arg("addr"), sig::arg("ldo"),
sig::arg("sdo"),
- sig::arg("swizzle"));
+ sig::arg("desc"), sig::arg("addr"),
sig::arg<IntExpr>("ldo"),
+ sig::arg<IntExpr>("sdo"),
sig::arg<IntExpr>("swizzle"));
RegisterDeviceIntrinsic("cuda_wgmma_noop_barrier", "cuda",
CallEffectKind::kOpaque,
sig::arg("reg"));
RegisterDeviceIntrinsic("nvshmem_barrier_all", "nvshmem",
CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("nvshmem_fence", "nvshmem", CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("nvshmem_getmem_nbi", "nvshmem",
CallEffectKind::kOpaque, sig::arg("dst"),
- sig::arg("src"), sig::arg("nelems"), sig::arg("pe"));
+ sig::arg("src"), sig::arg<IntExpr>("nelems"),
sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_getmem_nbi_block", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("dst"), sig::arg("src"),
sig::arg("nelems"), sig::arg("pe"));
+ sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
+ sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_getmem_nbi_warp", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("dst"), sig::arg("src"),
sig::arg("nelems"), sig::arg("pe"));
+ sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
+ sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_my_pe", "nvshmem", CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("nvshmem_n_pes", "nvshmem", CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("nvshmem_putmem_nbi", "nvshmem",
CallEffectKind::kOpaque, sig::arg("dst"),
- sig::arg("src"), sig::arg("nelems"), sig::arg("pe"));
+ sig::arg("src"), sig::arg<IntExpr>("nelems"),
sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_putmem_nbi_block", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("dst"), sig::arg("src"),
sig::arg("nelems"), sig::arg("pe"));
+ sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
+ sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_putmem_nbi_warp", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("dst"), sig::arg("src"),
sig::arg("nelems"), sig::arg("pe"));
+ sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
+ sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_putmem_signal_nbi", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("dst"), sig::arg("src"), sig::arg("nelems"),
- sig::arg("sig_addr"), sig::arg("signal"),
sig::arg("sig_op"),
- sig::arg("pe"));
+ sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
+ sig::arg("sig_addr"), sig::arg<IntExpr>("signal"),
sig::arg("sig_op"),
+ sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_putmem_signal_nbi_block", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("dst"), sig::arg("src"), sig::arg("nelems"),
- sig::arg("sig_addr"), sig::arg("signal"),
sig::arg("sig_op"),
- sig::arg("pe"));
+ sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
+ sig::arg("sig_addr"), sig::arg<IntExpr>("signal"),
sig::arg("sig_op"),
+ sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_putmem_signal_nbi_warp", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("dst"), sig::arg("src"), sig::arg("nelems"),
- sig::arg("sig_addr"), sig::arg("signal"),
sig::arg("sig_op"),
- sig::arg("pe"));
+ sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
+ sig::arg("sig_addr"), sig::arg<IntExpr>("signal"),
sig::arg("sig_op"),
+ sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_quiet", "nvshmem", CallEffectKind::kOpaque);
RegisterDeviceIntrinsic("nvshmem_signal_op", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("sig_addr"), sig::arg("signal"),
sig::arg("sig_op"),
- sig::arg("pe"));
+ sig::arg("sig_addr"), sig::arg<IntExpr>("signal"),
sig::arg("sig_op"),
+ sig::arg<IntExpr>("pe"));
RegisterDeviceIntrinsic("nvshmem_wait_until", "nvshmem",
CallEffectKind::kOpaque,
- sig::arg("ivar"), sig::arg("cmp"),
sig::arg("cmp_value"),
+ sig::arg("ivar"), sig::arg("cmp"),
sig::arg<IntExpr>("cmp_value"),
sig::arg("type"));
RegisterDeviceIntrinsic("ptx_legacy_ldmatrix", "ptx_legacy",
CallEffectKind::kOpaque,
- sig::arg("trans"), sig::arg("num"),
sig::arg("dtype"),
- sig::arg("local_ptr"), sig::arg("local_offset"),
sig::arg("smem_ptr"),
- sig::arg("smem_offset"));
- RegisterDeviceIntrinsic("ptx_legacy_mma", "ptx_legacy",
CallEffectKind::kOpaque,
- 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("a_offset"),
sig::arg("b_ptr"),
- sig::arg("b_offset"), sig::arg("acc_ptr"),
sig::arg("c_offset"),
- sig::arg("saturate"), sig::var_args("args"));
+ 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"));
+ RegisterDeviceIntrinsic(
+ "ptx_legacy_mma", "ptx_legacy", CallEffectKind::kOpaque,
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"));
for (const char* name : {
"tirx.cuda.tcgen05_encode_matrix_descriptor",
diff --git a/src/backend/metal/codegen/intrin_rule_metal.cc
b/src/backend/metal/codegen/intrin_rule_metal.cc
index 2f3d559ff2..402ebb6293 100644
--- a/src/backend/metal/codegen/intrin_rule_metal.cc
+++ b/src/backend/metal/codegen/intrin_rule_metal.cc
@@ -155,7 +155,7 @@ void RegisterMetalIntrinRules() {
// Register low-level Metal device intrinsics.
OpDef("tirx.metal.simd_shuffle")
.signature(sig::arg("var", "The variable to sync."),
- sig::arg("lane", "The source thread id."))
+ sig::arg<IntExpr>("lane", "The source thread id."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("metal"))
@@ -166,7 +166,7 @@ void RegisterMetalIntrinRules() {
OpDef("tirx.metal.simd_shuffle_up")
.signature(sig::arg("var", "The variable to sync."),
- sig::arg("delta", "The source lane id offset to be added."))
+ sig::arg<IntExpr>("delta", "The source lane id offset to be added."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("metal"))
@@ -177,7 +177,7 @@ void RegisterMetalIntrinRules() {
OpDef("tirx.metal.simd_shuffle_down")
.signature(sig::arg("var", "The variable to sync."),
- sig::arg("delta", "The source lane id offset to be subtracted."))
+ sig::arg<IntExpr>("delta", "The source lane id offset to be
subtracted."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("metal"))
diff --git a/src/backend/metal/op/target_builtin.cc
b/src/backend/metal/op/target_builtin.cc
index d9f3be63c0..794c802019 100644
--- a/src/backend/metal/op/target_builtin.cc
+++ b/src/backend/metal/op/target_builtin.cc
@@ -37,37 +37,41 @@ void RegisterMetalTargetBuiltins() {
registered = true;
OpDef("tirx.make_filled_simdgroup_matrix")
- .signature(sig::arg("d", "The D operand."), sig::arg("index", "The
index."),
- sig::arg("value", "The value to use."), sig::arg("col", "The
column index."),
- sig::arg("row", "The row index."))
+ .signature(sig::arg("d", "The D operand."), sig::arg<IntExpr>("index",
"The index."),
+ sig::arg("value", "The value to use."),
+ sig::arg<IntExpr>("col", "The column index."),
+ sig::arg<IntExpr>("row", "The row index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.metal.make_filled_simdgroup_matrix"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.simdgroup_load")
- .signature(sig::arg("d", "The D operand."), sig::arg("index", "The
index."),
- sig::arg("ptr", "The pointer."), sig::arg("stride", "The
stride."),
- sig::arg("col", "The column index."), sig::arg("row", "The
row index."),
+ .signature(sig::arg("d", "The D operand."), sig::arg<IntExpr>("index",
"The index."),
+ sig::arg("ptr", "The pointer."), sig::arg<IntExpr>("stride",
"The stride."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.metal.simdgroup_load"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.simdgroup_store")
- .signature(sig::arg("d", "The D operand."), sig::arg("index", "The
index."),
- sig::arg("ptr", "The pointer."), sig::arg("stride", "The
stride."),
- sig::arg("col", "The column index."), sig::arg("row", "The
row index."),
+ .signature(sig::arg("d", "The D operand."), sig::arg<IntExpr>("index",
"The index."),
+ sig::arg("ptr", "The pointer."), sig::arg<IntExpr>("stride",
"The stride."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.metal.simdgroup_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.simdgroup_multiply_accumulate")
- .signature(sig::arg("d", "The D operand."), sig::arg("index_d", "The D
fragment index."),
- sig::arg("a", "The A operand."), sig::arg("index_a", "The A
fragment index."),
- sig::arg("b", "The B operand."), sig::arg("index_b", "The B
fragment index."),
- sig::arg("c", "The C operand."), sig::arg("index_c", "The C
fragment index."))
+ .signature(
+ sig::arg("d", "The D operand."), sig::arg<IntExpr>("index_d", "The D
fragment index."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.metal.simdgroup_multiply_accumulate"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
diff --git a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
index b55c517000..6178cc6d8b 100644
--- a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
+++ b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
@@ -182,7 +182,7 @@ void RegisterWebGPUIntrinRules() {
// Register low-level WebGPU device intrinsics.
OpDef("tirx.webgpu.subgroup_shuffle")
.signature(sig::arg("var", "The variable to sync."),
- sig::arg("lane", "The source thread id."))
+ sig::arg<IntExpr>("lane", "The source thread id."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("webgpu"))
@@ -193,7 +193,7 @@ void RegisterWebGPUIntrinRules() {
OpDef("tirx.webgpu.subgroup_shuffle_up")
.signature(sig::arg("var", "The variable to sync."),
- sig::arg("delta", "The source lane id offset to be added."))
+ sig::arg<IntExpr>("delta", "The source lane id offset to be added."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("webgpu"))
@@ -204,7 +204,7 @@ void RegisterWebGPUIntrinRules() {
OpDef("tirx.webgpu.subgroup_shuffle_down")
.signature(sig::arg("var", "The variable to sync."),
- sig::arg("delta", "The source lane id offset to be subtracted."))
+ sig::arg<IntExpr>("delta", "The source lane id offset to be
subtracted."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory",
ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("webgpu"))
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index fb776eda1e..5b4185fe90 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -1195,14 +1195,15 @@ Type InferTypeAllocateTensor(const Call& call, const
BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.builtin.alloc_tensor")
- .signature(sig::arg("shape", "The shape of the tensor to allocate."),
- sig::arg("dtype", "The dtype of the tensor to allocate."),
- sig::arg("runtime_device_index",
- "The device index indicating on which device the
tensor is to be "
- "allocated at runtime. Index -1 is reserved for the
host device."),
- sig::arg("storage_scope",
- "The storage scope of the storage to allocate.
Default is global."),
- sig::var_ty_args("out_type", "Optional output type used by
allocation rewrites."))
+ .signature(
+ sig::arg("shape", "The shape of the tensor to allocate."),
+ sig::arg("dtype", "The dtype of the tensor to allocate."),
+ sig::arg<IntExpr>("runtime_device_index",
+ "The device index indicating on which device the
tensor is to be "
+ "allocated at runtime. Index -1 is reserved for
the host device."),
+ sig::arg("storage_scope",
+ "The storage scope of the storage to allocate. Default is
global."),
+ sig::var_ty_args("out_type", "Optional output type used by
allocation rewrites."))
.set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder",
InferTypeAllocateTensor)
// memory allocation isn't considered a "visible effect" as far as
purity is concerned
.set_attr<bool>("FPurity", true)
@@ -1277,13 +1278,14 @@ Type InferTypeMemAllocTensor(const Call& call, const
BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.memory.alloc_tensor")
- .signature(sig::arg("storage", "The storage to allocate the tensor to."),
- sig::arg("offset", "Storage offset to allocate the tensor."),
- sig::arg("shape", "The shape of the tensor to allocate."),
- sig::arg("dtype", "The dtype of the tensor to allocate."),
- sig::arg("runtime_device_index",
- "The device index indicating on which device the
tensor is to be "
- "allocated at runtime. Index -1 is reserved for the
host device."))
+ .signature(
+ sig::arg("storage", "The storage to allocate the tensor to."),
+ sig::arg<IntExpr>("offset", "Storage offset to allocate the
tensor."),
+ sig::arg("shape", "The shape of the tensor to allocate."),
+ sig::arg("dtype", "The dtype of the tensor to allocate."),
+ sig::arg<IntExpr>("runtime_device_index",
+ "The device index indicating on which device the
tensor is to be "
+ "allocated at runtime. Index -1 is reserved for
the host device."))
.set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder",
InferTypeMemAllocTensor)
// memory allocation isn't considered a "visible effect" as far as
purity is concerned
.set_attr<bool>("FPurity", true)
@@ -1352,9 +1354,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vm.alloc_storage")
.signature(sig::arg("size", "The size of the storage to allocate."),
- sig::arg("runtime_device_index",
- "The device index indicating on which device the
tensor is "
- "to be allocated at runtime."),
+ sig::arg<IntExpr>("runtime_device_index",
+ "The device index indicating on which
device the tensor is "
+ "to be allocated at runtime."),
sig::arg("dtype", "The dtype of the tensor to allocate."),
sig::arg("storage_scope",
"The storage scope of the storage to allocate.
Default is global."))
@@ -1405,12 +1407,12 @@ Type InferTypeVMAllocTensor(const Call& call, const
BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vm.alloc_tensor")
.signature(sig::arg("storage", "The storage to allocate the tensor to."),
- sig::arg("offset", "Storage offset to allocate the tensor."),
+ sig::arg<IntExpr>("offset", "Storage offset to allocate the
tensor."),
sig::arg("shape", "The shape of the tensor to allocate."),
sig::arg("dtype", "The dtype of the tensor to allocate."),
- sig::arg("runtime_device_index",
- "The device index indicating on which device the
tensor is "
- "to be allocated at runtime."))
+ sig::arg<IntExpr>("runtime_device_index",
+ "The device index indicating on which
device the tensor is "
+ "to be allocated at runtime."))
.set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder",
InferTypeVMAllocTensor)
// memory allocation isn't considered a "visible effect" as far as
purity is concerned
.set_attr<bool>("FPurity", true)
diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc
index e48c62f32e..76c887c6bd 100644
--- a/src/relax/op/tensor/create.cc
+++ b/src/relax/op/tensor/create.cc
@@ -450,7 +450,7 @@ Type InferTypeHammingWindow(const CallNode* call_node) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.hamming_window")
.signature(
- sig::arg("window_size", "The size of the window"),
+ sig::arg<IntExpr>("window_size", "The size of the window"),
sig::arg("periodic",
"If True, returns a window to be used as periodic function.
If False, return a "
"symmetric window"),
@@ -495,12 +495,14 @@ Type InferTypeTrilTriu(const CallNode* call_node) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.tril")
- .signature(sig::arg("x", "The input tensor."), sig::arg("k", "The offset
of the diagonal."))
+ .signature(sig::arg("x", "The input tensor."),
+ sig::arg<IntExpr>("k", "The offset of the diagonal."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeTrilTriu>())
.set_attr<bool>("FPurity", true);
OpDef("relax.triu")
- .signature(sig::arg("x", "The input tensor."), sig::arg("k", "The offset
of the diagonal."))
+ .signature(sig::arg("x", "The input tensor."),
+ sig::arg<IntExpr>("k", "The offset of the diagonal."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeTrilTriu>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc
index 1e62cf350e..ca60f3e135 100644
--- a/src/relax/op/tensor/inspect.cc
+++ b/src/relax/op/tensor/inspect.cc
@@ -313,7 +313,7 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const
Call& call) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_shape_i")
.signature(sig::arg("tensor", "The tensor to be inspected"),
- sig::arg("axis", "The axis whose extent should be returned"))
+ sig::arg<IntExpr>("axis", "The axis whose extent should be
returned"))
.set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder",
InferTypeTensorShape)
.set_attr<FLegalize>("FLegalize", LegalizeTensorShape)
.set_attr<bool>("RequiresArgumentShapes", false)
@@ -362,7 +362,7 @@ Type InferTypeTensorStride(const Call& call, const
BlockBuilder&) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_stride_i")
.signature(sig::arg("tensor", "The tensor to be inspected"),
- sig::arg("axis", "The axis whose extent should be returned"))
+ sig::arg<IntExpr>("axis", "The axis whose extent should be
returned"))
.set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder",
InferTypeTensorStride)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
diff --git a/src/relax/op/tensor/manipulate.cc
b/src/relax/op/tensor/manipulate.cc
index 38eb54a039..24a46d5105 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -3105,11 +3105,11 @@ Type InferTypeSliceScatter(const Call& call, const
BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.slice_scatter")
- .signature(sig::arg("input", "The input tensor."),
- sig::arg("src", "The source tensor to scatter."),
- sig::arg("start", "The starting index of the slice
(inclusive)."),
- sig::arg("end", "The ending index of the slice (exclusive)."),
- sig::arg("step", "The step of the slice."),
sig::call_attrs<SliceScatterAttrs>())
+ .signature(
+ sig::arg("input", "The input tensor."), sig::arg("src", "The source
tensor to scatter."),
+ sig::arg<IntExpr>("start", "The starting index of the slice
(inclusive)."),
+ sig::arg<IntExpr>("end", "The ending index of the slice
(exclusive)."),
+ sig::arg<IntExpr>("step", "The step of the slice."),
sig::call_attrs<SliceScatterAttrs>())
.set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder",
InferTypeSliceScatter)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 1b757ffda2..eaa0465f58 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -365,7 +365,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.call_llvm_intrin")
- .signature(sig::arg("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
+ .signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.call_llvm_intrin"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
@@ -373,7 +373,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.call_llvm_pure_intrin")
- .signature(sig::arg("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
+ .signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.call_llvm_pure_intrin"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure))
@@ -382,28 +382,30 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.call_spirv_pure_glsl450")
- .signature(sig::arg("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
+ .signature(sig::arg<IntExpr>("intrin_id", "The intrinsic identifier."),
sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.prefetch")
- .signature(sig::arg("ptr", "The pointer."), sig::arg("rw", "The
read/write mode."),
- sig::arg("locality", "The locality hint."),
- sig::arg("cache_type", "The cache policy."))
+ .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<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_access_ptr")
.signature(sig::arg("ptype", "The pointer type."), sig::arg("data", "The
input data."),
- sig::arg("offset", "The offset."), sig::arg("extent", "The
extent."),
- sig::arg("rw_mask", "The read/write mask."))
+ sig::arg<IntExpr>("offset", "The offset."),
+ sig::arg<IntExpr>("extent", "The extent."),
+ sig::arg<IntExpr>("rw_mask", "The read/write mask."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_access_ptr"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kSpecialCallArg));
OpDef("tirx.ptr_byte_offset")
- .signature(sig::arg("data", "Base pointer."), sig::arg("byte_offset",
"Offset in bytes."),
+ .signature(sig::arg("data", "Base pointer."),
+ sig::arg<IntExpr>("byte_offset", "Offset in bytes."),
sig::arg("dtype", "Type annotation for pointed-to elements."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.ptr_byte_offset"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
@@ -415,15 +417,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(CallEffectKind::kSpecialCallArg));
OpDef("tirx.handle_add_byte_offset")
- .signature(sig::arg("handle", "The handle."), sig::arg("offset", "The
offset."))
+ .signature(sig::arg("handle", "The handle."),
sig::arg<IntExpr>("offset", "The offset."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.handle_add_byte_offset"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.tvm_struct_get")
- .signature(sig::arg("arr", "The array."), sig::arg("index", "The
index."),
- sig::arg("field", "The field index."))
+ .signature(sig::arg("arr", "The array."), sig::arg<IntExpr>("index",
"The index."),
+ sig::arg<IntExpr>("field", "The field index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_struct_get"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
@@ -432,8 +434,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kLast));
OpDef("tirx.tvm_struct_set")
- .signature(sig::arg("arr", "The array."), sig::arg("index", "The
index."),
- sig::arg("field", "The field index."), sig::arg("value", "The
value to use."))
+ .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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_struct_set"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
@@ -446,22 +449,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.tvm_stack_alloca")
.signature(sig::arg("dtype_str", "The data type name."),
- sig::arg("num", "The number of entries."))
+ sig::arg<IntExpr>("num", "The number of entries."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_stack_alloca"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_stack_make_shape")
- .signature(sig::var_args("args"))
+ .signature(sig::var_args<IntExpr>("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_stack_make_shape"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_stack_make_array")
.signature(sig::arg("data", "The input data."), sig::arg("shape", "The
shape."),
- sig::arg("strides", "The strides."), sig::arg("ndim", "The
number of dimensions."),
+ sig::arg("strides", "The strides."),
+ sig::arg<IntExpr>("ndim", "The number of dimensions."),
sig::arg("arr_dtype", "The array data type."),
- sig::arg("elem_offset", "The element offset."))
+ sig::arg<IntExpr>("elem_offset", "The element offset."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_stack_make_array"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -501,7 +505,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.tvm_call_packed_lowered")
.signature(sig::arg("func_name", "The function name."),
sig::arg("args_stack", "The argument stack."),
- sig::arg("begin", "The start index."), sig::arg("end", "The
end index."))
+ sig::arg<IntExpr>("begin", "The start index."),
+ sig::arg<IntExpr>("end", "The end index."))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.call_packed_lowered"));
@@ -509,7 +514,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.tvm_call_cpacked_lowered")
.signature(sig::arg("func_name", "The function name."),
sig::arg("args_stack", "The argument stack."),
- sig::arg("begin", "The start index."), sig::arg("end", "The
end index."))
+ sig::arg<IntExpr>("begin", "The start index."),
+ sig::arg<IntExpr>("end", "The end index."))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.call_cpacked_lowered"));
@@ -531,36 +537,40 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_warp_shuffle")
- .signature(sig::arg("mask", "The mask."), sig::arg("value", "The value
to use."),
- sig::arg("warp_id", "The warp identifier."),
sig::arg("width", "The width."),
- sig::arg("warp_size", "The number of threads per warp."))
+ .signature(sig::arg<IntExpr>("mask", "The mask."), sig::arg("value",
"The value to use."),
+ sig::arg<IntExpr>("warp_id", "The warp identifier."),
+ sig::arg<IntExpr>("width", "The width."),
+ sig::arg<IntExpr>("warp_size", "The number of threads per
warp."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_warp_shuffle"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_warp_shuffle_up")
- .signature(sig::arg("mask", "The mask."), sig::arg("value", "The value
to use."),
- sig::arg("offset", "The offset."), sig::arg("width", "The
width."),
- sig::arg("warp_size", "The number of threads per warp."))
+ .signature(sig::arg<IntExpr>("mask", "The mask."), sig::arg("value",
"The value to use."),
+ sig::arg<IntExpr>("offset", "The offset."),
+ sig::arg<IntExpr>("width", "The width."),
+ sig::arg<IntExpr>("warp_size", "The number of threads per
warp."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_warp_shuffle_up"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_warp_shuffle_down")
- .signature(sig::arg("mask", "The mask."), sig::arg("value", "The value
to use."),
- sig::arg("offset", "The offset."), sig::arg("width", "The
width."),
- sig::arg("warp_size", "The number of threads per warp."))
+ .signature(sig::arg<IntExpr>("mask", "The mask."), sig::arg("value",
"The value to use."),
+ sig::arg<IntExpr>("offset", "The offset."),
+ sig::arg<IntExpr>("width", "The width."),
+ sig::arg<IntExpr>("warp_size", "The number of threads per
warp."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_warp_shuffle_down"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_warp_shuffle_xor")
- .signature(sig::arg("mask", "The mask."), sig::arg("value", "The value
to use."),
- sig::arg("lane_mask", "The lane mask."), sig::arg("width",
"The width."),
- sig::arg("warp_size", "The number of threads per warp."))
+ .signature(sig::arg<IntExpr>("mask", "The mask."), sig::arg("value",
"The value to use."),
+ sig::arg<IntExpr>("lane_mask", "The lane mask."),
+ sig::arg<IntExpr>("width", "The width."),
+ sig::arg<IntExpr>("warp_size", "The number of threads per
warp."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_warp_shuffle_xor"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
@@ -573,59 +583,61 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_thread_allreduce")
- .signature(sig::arg("size", "The size."), sig::var_args("args"))
+ .signature(sig::arg<IntExpr>("size", "The size."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tvm_thread_allreduce"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cooperative_tensor_fill")
- .signature(sig::arg("d", "The D operand."), sig::arg("index", "The
index."),
- sig::arg("value", "The value to use."), sig::arg("rows", "The
number of rows."),
- sig::arg("cols", "The number of columns."))
+ .signature(sig::arg("d", "The D operand."), sig::arg<IntExpr>("index",
"The index."),
+ sig::arg("value", "The value to use."),
+ sig::arg<IntExpr>("rows", "The number of rows."),
+ sig::arg<IntExpr>("cols", "The number of columns."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cooperative_tensor_fill"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cooperative_tensor_load")
- .signature(sig::arg("d", "The D operand."), sig::arg("index", "The
index."),
- sig::arg("ptr", "The pointer."), sig::arg("stride", "The
stride."),
- sig::arg("rows", "The number of rows."),
- sig::arg("cols", "The number of columns."),
+ .signature(sig::arg("d", "The D operand."), sig::arg<IntExpr>("index",
"The index."),
+ sig::arg("ptr", "The pointer."), sig::arg<IntExpr>("stride",
"The stride."),
+ sig::arg<IntExpr>("rows", "The number of rows."),
+ sig::arg<IntExpr>("cols", "The number of columns."),
sig::arg("transpose_matrix", "Whether to transpose the
matrix."),
- sig::arg("mma_M", "The M dimension of the matrix operation."),
- sig::arg("mma_N", "The N dimension of the matrix operation."),
- sig::arg("mma_K", "The K dimension of the matrix operation."),
- sig::arg("operand_role", "The matrix operand role."))
+ sig::arg<IntExpr>("mma_M", "The M dimension of the matrix
operation."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cooperative_tensor_load"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cooperative_tensor_store")
- .signature(sig::arg("d", "The D operand."), sig::arg("index", "The
index."),
- sig::arg("ptr", "The pointer."), sig::arg("stride", "The
stride."),
- sig::arg("rows", "The number of rows."),
- sig::arg("cols", "The number of columns."),
+ .signature(sig::arg("d", "The D operand."), sig::arg<IntExpr>("index",
"The index."),
+ sig::arg("ptr", "The pointer."), sig::arg<IntExpr>("stride",
"The stride."),
+ sig::arg<IntExpr>("rows", "The number of rows."),
+ sig::arg<IntExpr>("cols", "The number of columns."),
sig::arg("transpose_matrix", "Whether to transpose the
matrix."),
- sig::arg("mma_M", "The M dimension of the matrix operation."),
- sig::arg("mma_N", "The N dimension of the matrix operation."),
- sig::arg("mma_K", "The K dimension of the matrix operation."),
- sig::arg("operand_role", "The matrix operand role."))
+ sig::arg<IntExpr>("mma_M", "The M dimension of the matrix
operation."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cooperative_tensor_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cooperative_tensor_multiply_accumulate")
- .signature(sig::arg("d", "The D operand."), sig::arg("index_d", "The D
fragment index."),
- sig::arg("a", "The A operand."), sig::arg("index_a", "The A
fragment index."),
- sig::arg("b", "The B operand."), sig::arg("index_b", "The B
fragment index."),
- sig::arg("c", "The C operand."), sig::arg("index_c", "The C
fragment index."),
- sig::arg("M", "The M dimension."), sig::arg("N", "The N
dimension."),
- sig::arg("K", "The K dimension."),
- sig::arg("transpose_a", "Whether to transpose A."),
- sig::arg("transpose_b", "Whether to transpose B."))
+ .signature(
+ sig::arg("d", "The D operand."), sig::arg<IntExpr>("index_d", "The D
fragment index."),
+ 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."),
+ sig::arg<IntExpr>("M", "The M dimension."), sig::arg<IntExpr>("N",
"The N dimension."),
+ sig::arg<IntExpr>("K", "The K dimension."),
+ sig::arg("transpose_a", "Whether to transpose A."),
+ sig::arg("transpose_b", "Whether to transpose B."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cooperative_tensor_multiply_accumulate"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
@@ -672,24 +684,26 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.nd_mem_alloc_with_scope")
.signature(sig::arg("storage_scope", "The storage scope."),
- sig::arg("ndim", "The number of dimensions."),
sig::arg("shape", "The shape."),
- sig::var_args("args"))
+ sig::arg<IntExpr>("ndim", "The number of dimensions."),
+ sig::arg("shape", "The shape."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.texture2d_store")
- .signature(sig::arg("texture", "The texture."), sig::arg("x", "The input
value."),
- sig::arg("y", "The second input value."), sig::arg("z", "The
third input value."),
- sig::arg("channel_size", "The number of channels."),
+ .signature(sig::arg("texture", "The texture."), sig::arg<IntExpr>("x",
"The input value."),
+ sig::arg<IntExpr>("y", "The second input value."),
+ sig::arg<IntExpr>("z", "The third input value."),
+ sig::arg<IntExpr>("channel_size", "The number of channels."),
sig::arg("value", "The value to use."))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TVectorizable>("TVectorizable", true)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.texture2d_load")
- .signature(sig::arg("texture", "The texture."), sig::arg("x", "The input
value."),
- sig::arg("y", "The second input value."), sig::arg("z", "The
third input value."),
- sig::arg("channel_size", "The number of channels."),
+ .signature(sig::arg("texture", "The texture."), sig::arg<IntExpr>("x",
"The input value."),
+ sig::arg<IntExpr>("y", "The second input value."),
+ sig::arg<IntExpr>("z", "The third input value."),
+ sig::arg<IntExpr>("channel_size", "The number of channels."),
sig::arg("element_index", "The element index within a texture
channel."))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TVectorizable>("TVectorizable", true)
@@ -709,7 +723,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(CallEffectKind::kReadState));
OpDef("tirx.get_active_lane_mask")
- .signature(sig::arg("base", "The base value."), sig::arg("limit", "The
limit value."))
+ .signature(sig::arg<IntExpr>("base", "The base value."),
+ sig::arg<IntExpr>("limit", "The limit value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.get_active_lane_mask"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure))
@@ -778,7 +793,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "The input data."), sig::arg("dtype", "The
data type."),
sig::arg("is_string", "Whether to print as a string."),
sig::arg("is_scalar", "Whether to print as a scalar."),
- sig::arg("ndim", "The number of dimensions."),
sig::var_args("args"))
+ sig::arg<IntExpr>("ndim", "The number of dimensions."),
sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.print_buffer"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index aa662a0aac..877e5f583e 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -699,11 +699,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.TVMBackendAllocWorkspace")
- .signature(sig::arg("device_type", "The device type."),
- sig::arg("device_id", "The device index."),
- sig::arg("nbytes", "The number of bytes."),
- sig::arg("dtype_code_hint", "The data type code hint."),
- sig::arg("dtype_bits_hint", "The data type bit-width hint."))
+ .signature(sig::arg<IntExpr>("device_type", "The device type."),
+ sig::arg<IntExpr>("device_id", "The device index."),
+ sig::arg<IntExpr>("nbytes", "The number of bytes."),
+ 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<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.TVMBackendAllocWorkspace"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
@@ -711,8 +711,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.TVMBackendFreeWorkspace")
- .signature(sig::arg("device_type", "The device type."),
- sig::arg("device_id", "The device index."), sig::arg("ptr",
"The pointer."))
+ .signature(sig::arg<IntExpr>("device_type", "The device type."),
+ sig::arg<IntExpr>("device_id", "The device index."),
+ sig::arg("ptr", "The pointer."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.TVMBackendFreeWorkspace"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
diff --git a/tests/cpp/expr_test.cc b/tests/cpp/expr_test.cc
index 04bedfce99..067962a6e3 100644
--- a/tests/cpp/expr_test.cc
+++ b/tests/cpp/expr_test.cc
@@ -90,6 +90,27 @@ TEST(Expr, PrimTypeBoolLanes) {
TVM_FFI_ICHECK(boolx4.MatchesElementType(DLDataTypeCode::kDLBool, 8));
}
+TEST(Expr, IntegerView) {
+ using namespace tvm;
+ for (PrimType ty : {PrimType::Int(8), PrimType::UInt(64),
PrimType::Int(128)}) {
+ ffi::Any value = Var("i", ty);
+ EXPECT_TRUE(value.cast<IntExpr>().same_as(value.cast<PrimExpr>()));
+ EXPECT_TRUE(value.cast<TypedExpr<PrimType>>().ty() == ty);
+ }
+ for (PrimType ty : {PrimType::Float(32), PrimType::Bool(), PrimType::Int(32,
4),
+ PrimType::ScalableVector(kDLInt, 32, 4)}) {
+ ffi::Any value = Var("x", ty);
+ EXPECT_FALSE(value.try_cast<IntExpr>());
+ EXPECT_TRUE(value.try_cast<PrimExpr>());
+ }
+ ffi::Any nonprimitive = Var("x", AnyType());
+ EXPECT_FALSE(nonprimitive.try_cast<IntExpr>());
+ EXPECT_FALSE(nonprimitive.try_cast<PrimExpr>());
+ EXPECT_FALSE(nonprimitive.try_cast<TypedExpr<PrimType>>());
+ EXPECT_EQ(ffi::Any(1).cast<IntExpr>().ty().bits(), 32);
+ EXPECT_EQ(ffi::Any(int64_t{1} << 40).cast<IntExpr>().ty().bits(), 64);
+}
+
TEST(ExprNodeRef, Basic) {
using namespace tvm;
using namespace tvm::tirx;
diff --git a/tests/python/tirx-base/test_tir_op_types.py
b/tests/python/tirx-base/test_tir_op_types.py
index 6ff7c05e6e..085b0c69f6 100644
--- a/tests/python/tirx-base/test_tir_op_types.py
+++ b/tests/python/tirx-base/test_tir_op_types.py
@@ -29,6 +29,17 @@ def test_tir_op_tvm_struct_get():
assert expr.op.name == "tirx.tvm_struct_get"
+def test_scalar_integer_signature():
+ ptr = tirx.Var("ptr", ty="handle")
+ for dtype in ("int8", "uint32", "int64", "uint64"):
+ index = tirx.Var("index", dtype)
+ assert tirx.tvm_struct_get(ptr, index, 2,
dtype="int32").args[1].same_as(index)
+ assert tirx.tvm_struct_get(ptr, 1 << 40, 2,
dtype="int32").args[1].ty.dtype == "int64"
+ for dtype in ("float32", "bool", "int32x4", "int32xvscalex4"):
+ with pytest.raises(TypeError, match=r"index.*expected
`ir.IntExpr`.*dtype="):
+ tirx.tvm_struct_get(ptr, tirx.Var("index", dtype), 2,
dtype="int32")
+
+
def test_tir_op_tvm_struct_set():
x = tirx.Var("x", ty="handle")
expr = tirx.tvm_struct_set(x, 1, 2, 3)
