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 4e03bc537d [TIRx] Canonicalize CUDA and vector operator identities
(#20630)
4e03bc537d is described below
commit 4e03bc537db3ded29e29d41d2bd8c4b13426a807
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 9 19:17:37 2026 -0400
[TIRx] Canonicalize CUDA and vector operator identities (#20630)
Canonicalize the remaining CUDA MMA/timer and vector operator
identities, migrating their direct consumers to a single registration.
Consolidate MMA fill/store lowering while preserving raw-pointer offsets
and fragment layouts, attach Hexagon population count lowering to the
canonical primitive, and retire unused finiteness registrations.
---
docs/tirx/api/cuda.rst | 3 +-
include/tvm/tirx/op/vector.h | 6 +--
python/tvm/backend/cuda/cpp/instrument.py | 12 +++--
python/tvm/backend/cuda/op.py | 45 ++++-------------
python/tvm/backend/cuda/script.py | 14 +++--
python/tvm/s_tir/tensor_intrin/arm_cpu.py | 4 +-
python/tvm/s_tir/tensor_intrin/cuda.py | 6 +--
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/op/__init__.py | 12 ++---
python/tvm/tirx/script/ir_builder/op.py | 12 ++---
src/backend/cuda/codegen/codegen_cuda.cc | 59 +---------------------
src/backend/cuda/op/target_builtin.cc | 36 +++----------
.../hexagon/codegen/llvm/intrin_rule_hexagon.cc | 4 +-
src/s_tir/transform/inject_permuted_layout.cc | 2 +-
src/target/intrin_rule.cc | 17 -------
src/target/llvm/codegen_llvm.cc | 6 +--
src/tirx/op/vector.cc | 26 +++++-----
src/tirx/transform/lower_warp_memory.cc | 23 ++-------
src/tirx/transform/remove_no_op.cc | 12 ++---
tests/python/tirx-base/test_tir_op_types.py | 22 ++++----
20 files changed, 92 insertions(+), 231 deletions(-)
diff --git a/docs/tirx/api/cuda.rst b/docs/tirx/api/cuda.rst
index e507a93851..3f897ee31f 100644
--- a/docs/tirx/api/cuda.rst
+++ b/docs/tirx/api/cuda.rst
@@ -74,8 +74,7 @@ The current ``Tx.cuda`` surface includes:
- ``iket.mark``, ``iket.range_start``, ``iket.range_end``,
``iket.range_push``, ``iket.range_pop``, ``iket.sentinel_token``,
``iket.official_event``, ``timer_init``, ``timer_start``, ``timer_end``,
- ``timer_finalize``, ``mma_store``, ``mma_fill``, ``mma_store_legacy``,
- and ``mma_fill_legacy``
+ ``timer_finalize``, ``mma_store``, and ``mma_fill``
NVSHMEM namespace
-----------------
diff --git a/include/tvm/tirx/op/vector.h b/include/tvm/tirx/op/vector.h
index 7b55ad4acc..4a5e8036c4 100644
--- a/include/tvm/tirx/op/vector.h
+++ b/include/tvm/tirx/op/vector.h
@@ -36,7 +36,7 @@ namespace tirx {
* Arguments, in order:
* - args[0]: vec, The input vector.
*/
-TVM_DLL const Op& vectorhigh_op();
+TVM_DLL const Op& vector_high_op();
/*!
* \brief Get the low-level half of the vector.
@@ -44,7 +44,7 @@ TVM_DLL const Op& vectorhigh_op();
* Arguments, in order:
* - args[0]: vec, The input vector.
*/
-TVM_DLL const Op& vectorlow_op();
+TVM_DLL const Op& vector_low_op();
/*!
* \brief Concat two vectors.
@@ -53,7 +53,7 @@ TVM_DLL const Op& vectorlow_op();
* - args[0]: vec1, The first input vector.
* - args[1]: vec2, The second input vector.
*/
-TVM_DLL const Op& vectorcombine_op();
+TVM_DLL const Op& vector_combine_op();
/*!
* \brief Calculate a predicate mask given an upper bound (limit) and a
current value (base).
diff --git a/python/tvm/backend/cuda/cpp/instrument.py
b/python/tvm/backend/cuda/cpp/instrument.py
index 2552f3531f..0a74899f47 100644
--- a/python/tvm/backend/cuda/cpp/instrument.py
+++ b/python/tvm/backend/cuda/cpp/instrument.py
@@ -53,7 +53,8 @@ def _write_event(event_bits: str) -> str:
device_intrinsic(
- "timer_init_cuda",
+ "cuda.timer_init",
+ helper_name="tvm_builtin_timer_init_cuda",
c_signature=(
"(uint64_t* profiler_buffer, uint64_t* profiler_tag, "
"uint32_t* profiler_write_offset, int num_groups, int group_id)"
@@ -75,7 +76,8 @@ device_intrinsic(
)
device_intrinsic(
- "timer_start_cuda",
+ "cuda.timer_start",
+ helper_name="tvm_builtin_timer_start_cuda",
c_signature=f"({_EVENT_PARAMS})",
body=(
f" if (leader_cond) {{\n {_write_event('(uint32_t)event_type
<< 2 | 0x0')}\n }}\n"
@@ -85,7 +87,8 @@ device_intrinsic(
)
device_intrinsic(
- "timer_end_cuda",
+ "cuda.timer_end",
+ helper_name="tvm_builtin_timer_end_cuda",
c_signature=f"({_EVENT_PARAMS})",
body=(
" __threadfence_block();\n"
@@ -95,7 +98,8 @@ device_intrinsic(
)
device_intrinsic(
- "timer_finalize_cuda",
+ "cuda.timer_finalize",
+ helper_name="tvm_builtin_timer_finalize_cuda",
c_signature=f"({_COMMON_PARAMS})",
body=(
f" __threadfence_block();\n if (leader_cond) {{\n
{_write_event('0x3')}\n }}"
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index d745075c2f..cd21f3bdda 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -519,38 +519,13 @@ def bmma_sync(
def mma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
"""Store the result of PTX MMA into a destination pointer."""
- return call_intrin(dtype, "tirx.mma_store", m, n, dst_ptr, src_ptr,
src_offset, dst_stride)
-
-
-_mma_store_legacy = _make_op_api(Op.get("tirx.mma_store_legacy"), __name__)
-
-
-def mma_store_legacy(m, n, dst_ptr, src_ptr, src_offset, dst_stride, *,
ty=None, span=None):
- """Store MMA registers using explicit pointer and element-offset
operands."""
- return _mma_store_legacy(
- m,
- n,
- dst_ptr,
- src_ptr,
- src_offset,
- dst_stride,
- ty=ty,
- span=span,
- )
+ return call_intrin(dtype, "tirx.cuda.mma_store", m, n, dst_ptr, src_ptr,
src_offset, dst_stride)
def mma_fill(dtype, local_size, local_ptr, offset):
"""Zero-initialize an MMA accumulation register."""
- return call_intrin(dtype, "tirx.mma_fill", local_size, local_ptr, offset)
-
-
-_mma_fill_legacy = _make_op_api(Op.get("tirx.mma_fill_legacy"), __name__)
-
-
-def mma_fill_legacy(local_size, local_ptr, offset, *, ty=None, span=None):
- """Initialize MMA registers using an explicit element offset."""
- return _mma_fill_legacy(local_size, local_ptr, offset, ty=ty, span=span)
+ return call_intrin(dtype, "tirx.cuda.mma_fill", local_size, local_ptr,
offset)
_PTX_TO_NUMPY_DTYPE = {
@@ -873,7 +848,7 @@ def _static_str(value):
# shape tokens in backend/cuda/ptx/table.py.
-def timer_init_cuda(
+def timer_init(
profiler_buffer,
profiler_tag,
profiler_write_offset,
@@ -910,14 +885,14 @@ def timer_init_cuda(
"""
return Call(
- "tirx.timer_init_cuda",
+ "tirx.cuda.timer_init",
[profiler_buffer, profiler_tag, profiler_write_offset, num_groups,
group_id],
ty=ty,
span=span,
)
-def timer_start_cuda(
+def timer_start(
event_type,
profiler_buffer,
profiler_tag,
@@ -958,7 +933,7 @@ def timer_start_cuda(
""" # noqa: E501
return Call(
- "tirx.timer_start_cuda",
+ "tirx.cuda.timer_start",
[
event_type.value if isinstance(event_type, Enum) else event_type,
profiler_buffer,
@@ -972,7 +947,7 @@ def timer_start_cuda(
)
-def timer_end_cuda(
+def timer_end(
event_type,
profiler_buffer,
profiler_tag,
@@ -1013,7 +988,7 @@ def timer_end_cuda(
""" # noqa: E501
return Call(
- "tirx.timer_end_cuda",
+ "tirx.cuda.timer_end",
[
event_type.value if isinstance(event_type, Enum) else event_type,
profiler_buffer,
@@ -1027,7 +1002,7 @@ def timer_end_cuda(
)
-def timer_finalize_cuda(
+def timer_finalize(
profiler_buffer,
profiler_tag,
profiler_write_offset,
@@ -1064,7 +1039,7 @@ def timer_finalize_cuda(
"""
return Call(
- "tirx.timer_finalize_cuda",
+ "tirx.cuda.timer_finalize",
[profiler_buffer, profiler_tag, profiler_write_offset,
profiler_write_stride, leader_cond],
ty=ty,
span=span,
diff --git a/python/tvm/backend/cuda/script.py
b/python/tvm/backend/cuda/script.py
index 26f05b09b5..1f1dc441f8 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -119,15 +119,13 @@ cta_max = _cuda_op.cuda_cta_max
cta_min = _cuda_op.cuda_cta_min
func_call = _cuda_op.cuda_func_call
sm100_2sm_leader_smem_addr = _cuda_op.cuda_sm100_2sm_leader_smem_addr
-timer_init = _cuda_op.timer_init_cuda
-timer_start = _cuda_op.timer_start_cuda
-timer_end = _cuda_op.timer_end_cuda
-timer_finalize = _cuda_op.timer_finalize_cuda
-mma_store = _ir.op._make_op_api(_ir.Op.get("tirx.mma_store"), __name__)
-mma_fill = _ir.op._make_op_api(_ir.Op.get("tirx.mma_fill"), __name__)
+timer_init = _cuda_op.timer_init
+timer_start = _cuda_op.timer_start
+timer_end = _cuda_op.timer_end
+timer_finalize = _cuda_op.timer_finalize
+mma_store = _ir.op._make_op_api(_ir.Op.get("tirx.cuda.mma_store"), __name__)
+mma_fill = _ir.op._make_op_api(_ir.Op.get("tirx.cuda.mma_fill"), __name__)
bmma_sync = _cuda_op.bmma_sync
-mma_store_legacy = _cuda_op.mma_store_legacy
-mma_fill_legacy = _cuda_op.mma_fill_legacy
atomic_add = _cuda_op.cuda_atomic_add
atomic_cas = _cuda_op.cuda_atomic_cas
cta_reduce = _cuda_op.cuda_cta_reduce
diff --git a/python/tvm/s_tir/tensor_intrin/arm_cpu.py
b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
index 86ba0dc7e8..841448dd2f 100644
--- a/python/tvm/s_tir/tensor_intrin/arm_cpu.py
+++ b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
@@ -71,7 +71,7 @@ def neon_4x4_i8i8i32_impl(
vec_b = B.vload([0, 0], dtype="int8x16")
# TODO(masahi): Remove duplication when inlined function call is
supported
- vec_b_low = T.vectorlow(vec_b, ty="int8x8")
+ vec_b_low = T.vector_low(vec_b, ty="int8x8")
multiply_low = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.smull.v8i16"),
@@ -86,7 +86,7 @@ def neon_4x4_i8i8i32_impl(
ty="int32x4",
)
- vec_b_high = T.vectorhigh(vec_b, ty="int8x8")
+ vec_b_high = T.vector_high(vec_b, ty="int8x8")
multiply_high = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.smull.v8i16"),
diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py
b/python/tvm/s_tir/tensor_intrin/cuda.py
index f985bb93ea..d77e19ad22 100644
--- a/python/tvm/s_tir/tensor_intrin/cuda.py
+++ b/python/tvm/s_tir/tensor_intrin/cuda.py
@@ -544,9 +544,7 @@ def get_mma_fill_intrin(dtype, local_size):
Ts.writes(C_warp[0:WARP_SIZE, 0:local_size])
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
- T.evaluate(
- T.cuda.mma_fill_legacy(local_size, C_warp.data,
C_warp.elem_offset, ty=dtype)
- )
+ T.evaluate(T.cuda.mma_fill(local_size, C_warp.data,
C_warp.elem_offset, ty=dtype))
return mma_fill_desc, mma_fill_impl
@@ -599,7 +597,7 @@ def get_mma_store_intrin(dtype, local_size, scope="global",
use_mma_store_intrin
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(
- T.cuda.mma_store_legacy(
+ T.cuda.mma_store(
M_DIM,
N_DIM,
C.access_ptr("w"),
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index cced97669a..f747862bef 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -65,7 +65,7 @@ from .op import (
gpu_mma_sync,
gpu_fill_fragment,
)
-from .op import vectorlow, vectorhigh, vectorcombine
+from .op import vector_low, vector_high, vector_combine
from .op import infinity, reinterpret
from .op import exp, exp2, exp10, log, log2, log10, log1p, ldexp, clz
from .op import sin, sinh, asin, asinh
diff --git a/python/tvm/tirx/op/__init__.py b/python/tvm/tirx/op/__init__.py
index 241dc97b04..1681df0ea4 100644
--- a/python/tvm/tirx/op/__init__.py
+++ b/python/tvm/tirx/op/__init__.py
@@ -1112,7 +1112,7 @@ def throw_last_error(*, ty=None, span=None):
return call_intrin(ty, "tirx.throw_last_error", span=span)
-def vectorlow(dtype, vec):
+def vector_low(dtype, vec):
"""Get the low level half of the vector
Parameters
@@ -1128,10 +1128,10 @@ def vectorlow(dtype, vec):
call : Expr
The call expression.
"""
- return call_intrin(dtype, "tirx.vectorlow", vec)
+ return call_intrin(dtype, "tirx.vector_low", vec)
-def vectorhigh(dtype, vec):
+def vector_high(dtype, vec):
"""Get the high level half of the vector
Parameters
@@ -1147,10 +1147,10 @@ def vectorhigh(dtype, vec):
call : Expr
The call expression.
"""
- return call_intrin(dtype, "tirx.vectorhigh", vec)
+ return call_intrin(dtype, "tirx.vector_high", vec)
-def vectorcombine(dtype, vec1, vec2):
+def vector_combine(dtype, vec1, vec2):
"""Concat two vectors
Parameters
@@ -1166,7 +1166,7 @@ def vectorcombine(dtype, vec1, vec2):
call : Expr
The call expression.
"""
- return call_intrin(dtype, "tirx.vectorcombine", vec1, vec2)
+ return call_intrin(dtype, "tirx.vector_combine", vec1, vec2)
def gpu_dp4a(vec1, vec2, acc=0, **kwargs):
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index bacc176903..e217c76616 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -558,13 +558,13 @@ call_llvm_pure_intrin =
_llvm_result_type(_tir_op.call_llvm_pure_intrin)
call_pure_extern = _ir.op._make_op_api(_ir.Op.get("tirx.call_pure_extern"),
__name__)
-vectorlow = _ir.op._make_op_api(_ir.Op.get("tirx.vectorlow"), __name__)
+vector_low = _ir.op._make_op_api(_ir.Op.get("tirx.vector_low"), __name__)
-vectorhigh = _ir.op._make_op_api(_ir.Op.get("tirx.vectorhigh"), __name__)
+vector_high = _ir.op._make_op_api(_ir.Op.get("tirx.vector_high"), __name__)
-vectorcombine = _ir.op._make_op_api(_ir.Op.get("tirx.vectorcombine"), __name__)
+vector_combine = _ir.op._make_op_api(_ir.Op.get("tirx.vector_combine"),
__name__)
get_active_lane_mask =
_ir.op._make_op_api(_ir.Op.get("tirx.get_active_lane_mask"), __name__)
@@ -862,9 +862,9 @@ __all__ = [
"truncdiv",
"truncmod",
"undef",
- "vectorcombine",
- "vectorhigh",
- "vectorlow",
+ "vector_combine",
+ "vector_high",
+ "vector_low",
"vscale",
"webgpu",
]
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 3f746a2e4f..1be72070c7 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1083,12 +1083,10 @@ void CodeGenCUDA::Dispatch_(const CallNode* op,
std::ostream& os) {
}
}
- static const Op mma_store_op = Op::Get("tirx.mma_store");
- static const Op mma_fill_op = Op::Get("tirx.mma_fill");
+ static const Op mma_store_op = Op::Get("tirx.cuda.mma_store");
+ static const Op mma_fill_op = Op::Get("tirx.cuda.mma_fill");
static const Op ptx_mma_legacy_op = Op::Get("tirx.ptx_legacy.mma");
static const Op ptx_ldmatrix_legacy_op = Op::Get("tirx.ptx_legacy.ldmatrix");
- static const Op mma_store_legacy_op = Op::Get("tirx.mma_store_legacy");
- static const Op mma_fill_legacy_op = Op::Get("tirx.mma_fill_legacy");
static const Op cuda_func_call_op = Op::Get("tirx.cuda.func_call");
if (op->op.same_as(mma_store_op)) {
@@ -1204,59 +1202,6 @@ void CodeGenCUDA::Dispatch_(const CallNode* op,
std::ostream& os) {
this->stream << PrintLoadMatrixAssembly(trans, num, type_str, local_ptr,
local_offset,
smem_ptr, smem_offset);
}
- } else if (op->op.same_as(mma_store_legacy_op)) {
- // args: m, n, dst_ptr, src_ptr_var, src_offset, dst_stride
- // (dst_ptr is typically an access_ptr Call that already encodes
- // dst.elem_offset and the global pointer cast.)
- int m = op->args[0].as_or_throw<IntImm>()->value.as<int>().value();
- int n = op->args[1].as_or_throw<IntImm>()->value.as<int>().value();
- std::string dst = this->PrintExpr(op->args[2]);
- std::string src = this->PrintExpr(op->args[3]);
- std::string src_offset = this->PrintExpr(op->args[4]);
- PrimExpr stride = op->args[5].as_or_throw<PrimExpr>();
-
- TVM_FFI_ICHECK(m == 16 && n == 16) << "Only m == 16 && n == 16 case
supported for now";
-
- const auto index_map_func =
-
tvm::ffi::Function::GetGlobal("tirx.index_map.shared_16x16_to_ldmatrix_32x8_layout");
- TVM_FFI_ICHECK(index_map_func.has_value());
-
- sym::Analyzer analyzer;
- auto inverse_index_map =
- IndexMap::FromFunc(2, *index_map_func).Inverse({Range(0, m), Range(0,
n)}, analyzer);
- auto indices_16x16 = inverse_index_map->final_indices;
-
- class LowerFloorDivMod : public tirx::StmtExprMutator {
- public:
- UnchangedOr<PrimExpr> Mutate_(const prim::FloorDivNode* op, InplaceMode
inplace_mode) {
- return prim::Div(this->Mutate(op->a,
inplace_mode).ValueOrUnchanged(op->a),
- this->Mutate(op->b,
inplace_mode).ValueOrUnchanged(op->b));
- }
- UnchangedOr<PrimExpr> Mutate_(const prim::FloorModNode* op, InplaceMode
inplace_mode) {
- return prim::Mod(this->Mutate(op->a,
inplace_mode).ValueOrUnchanged(op->a),
- this->Mutate(op->b,
inplace_mode).ValueOrUnchanged(op->b));
- }
- };
-
- PrimExpr dst_expr = indices_16x16[0] * stride + indices_16x16[1];
- auto dst_ind =
-
ffi::make_object<LowerFloorDivMod>()->Mutate(dst_expr).ValueOrUnchanged(dst_expr);
-
- var_idmap_[inverse_index_map->initial_indices[0].get()] = "threadIdx.x";
- var_idmap_[inverse_index_map->initial_indices[1].get()] = "local_id";
-
- os << "for (int local_id = 0; local_id < 8; ++local_id) {\n";
- os << dst << "[" << this->PrintExpr(dst_ind) << "] = " << src << "[" <<
src_offset
- << " + local_id];\n";
- os << "}\n";
- } else if (op->op.same_as(mma_fill_legacy_op)) {
- // args: local_size, local_ptr_var, offset
- std::string num_elem = this->PrintExpr(op->args[0]);
- std::string dst = this->PrintExpr(op->args[1]);
- std::string dst_offset = this->PrintExpr(op->args[2]);
- os << "for (int i = 0; i < " << num_elem << "; ++i) {\n";
- os << dst << "[" << dst_offset << " + i] = 0.0;";
- os << "}\n";
} else if (op->op.same_as(tirx::reinterpret_op())) {
// Compile-time pointer reinterpret of a literal (e.g. a tcgen05 descriptor
// template encoded at address 0): emit C++-style reinterpret_cast<T*>(...)
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index ac109a82f7..39a5a6dc67 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -110,30 +110,6 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
- // Siblings of mma_store / mma_fill that accept
- // (ptr_var, offset) pairs. Codegen emits `ptr + offset` C-pointer
- // arithmetic and lower_warp_memory rewrites the offset's group component
- // 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<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<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<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));
-
// 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
// and prints it with ty=; user-issued copies go through T.ptx instead.
@@ -148,7 +124,7 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("s_tir"))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.s_tir.cp_async_raw"));
- OpDef("tirx.mma_store")
+ OpDef("tirx.cuda.mma_store")
.signature(sig::arg<IntExpr>("m", "The M dimension."),
sig::arg<IntExpr>("n", "The N dimension."),
sig::arg("dst_ptr", "The destination pointer."),
@@ -159,7 +135,7 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
- OpDef("tirx.mma_fill")
+ 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."))
@@ -167,7 +143,7 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
- OpDef("tirx.timer_init_cuda")
+ OpDef("tirx.cuda.timer_init")
.signature(sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
sig::arg<IntExpr>("profiler_write_offset", "The profiler
write offset."),
@@ -177,7 +153,7 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
- OpDef("tirx.timer_start_cuda")
+ OpDef("tirx.cuda.timer_start")
.signature(sig::arg("event_type", "The event type."),
sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
@@ -188,7 +164,7 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
- OpDef("tirx.timer_end_cuda")
+ OpDef("tirx.cuda.timer_end")
.signature(sig::arg("event_type", "The event type."),
sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
@@ -199,7 +175,7 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
- OpDef("tirx.timer_finalize_cuda")
+ OpDef("tirx.cuda.timer_finalize")
.signature(sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
sig::arg<IntExpr>("profiler_write_offset", "The profiler
write offset."),
diff --git a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
b/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
index 1bcb32c8e0..182a9eea58 100644
--- a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
@@ -109,9 +109,7 @@ void RegisterHexagonIntrinRules() {
.set_attr<FLowerIntrinsic>("hexagon.FLowerIntrinsic",
DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
- OpDef("tirx.ctpop")
- .signature(sig::arg("x", "The input value."))
- .set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
+ OpDef("prim.popcount")
.set_attr<FLowerIntrinsic>("hexagon.FLowerIntrinsic",
DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
diff --git a/src/s_tir/transform/inject_permuted_layout.cc
b/src/s_tir/transform/inject_permuted_layout.cc
index 9a954d2b0b..c720616b26 100644
--- a/src/s_tir/transform/inject_permuted_layout.cc
+++ b/src/s_tir/transform/inject_permuted_layout.cc
@@ -293,7 +293,7 @@ class PermutedLayoutInjector : public IRMutatorWithAnalyzer
{
}
static const Op ptx_ldmatrix_op = Op::Get("tirx.ptx_legacy.ldmatrix");
- static const Op mma_store_op = Op::Get("tirx.mma_store_legacy");
+ static const Op mma_store_op = Op::Get("tirx.cuda.mma_store");
if (!call->op.same_as(ptx_ldmatrix_op) && !call->op.same_as(mma_store_op))
{
return call;
}
diff --git a/src/target/intrin_rule.cc b/src/target/intrin_rule.cc
index 1df6952fdb..68b7355cec 100644
--- a/src/target/intrin_rule.cc
+++ b/src/target/intrin_rule.cc
@@ -210,23 +210,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return one / (one + exp(-arg));
});
- OpDef("tirx.isfinite")
- .signature(sig::arg("x", "The input value."))
- .set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
- .set_attr<FLegalize>("default.FLegalize", [](const PrimExpr& e) ->
PrimExpr {
- const CallNode* call = e.as<CallNode>();
- TVM_FFI_ICHECK(call != nullptr);
- return isfinite(call->args[0].as_or_throw<PrimExpr>());
- });
-
- OpDef("tirx.isinf")
- .signature(sig::arg("x", "The input value."))
- .set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
- .set_attr<FLegalize>("default.FLegalize", [](const PrimExpr& e) ->
PrimExpr {
- const CallNode* call = e.as<CallNode>();
- TVM_FFI_ICHECK(call != nullptr);
- return isinf(call->args[0].as_or_throw<PrimExpr>());
- });
// clang-format on
}
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 9aa1cf54fd..65fc4439af 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -1441,15 +1441,15 @@ llvm::Value* CodeGenLLVM::CreateIntrinsic(const
CallNode* op) {
return llvm::Constant::getNullValue(GetLLVMType(op->ty));
}
return builder_->CreateFCmpUNO(a, a);
- } else if (op->op.same_as(tirx::vectorlow_op())) {
+ } else if (op->op.same_as(tirx::vector_low_op())) {
llvm::Value* v = MakeValue(args[0]);
int l = GetVectorNumElements(v);
return CreateVecSlice(v, 0, l / 2);
- } else if (op->op.same_as(tirx::vectorhigh_op())) {
+ } else if (op->op.same_as(tirx::vector_high_op())) {
llvm::Value* v = MakeValue(args[0]);
int l = GetVectorNumElements(v);
return CreateVecSlice(v, l / 2, l / 2);
- } else if (op->op.same_as(tirx::vectorcombine_op())) {
+ } else if (op->op.same_as(tirx::vector_combine_op())) {
llvm::Value* v0 = MakeValue(args[0]);
llvm::Value* v1 = MakeValue(args[1]);
int num_elems = GetVectorNumElements(v0) * 2;
diff --git a/src/tirx/op/vector.cc b/src/tirx/op/vector.cc
index 10c1135898..4796c2b903 100644
--- a/src/tirx/op/vector.cc
+++ b/src/tirx/op/vector.cc
@@ -35,7 +35,7 @@ Type InferTypeVectorPart(const CallNode* call) {
PrimType input = call->args[0]->ty.as_or_throw<PrimType>();
if constexpr (Combine) {
TVM_FFI_CHECK(ffi::StructuralEqual()(input, call->args[1]->ty), TypeError)
- << "vectorcombine requires equal input vector types";
+ << "vector_combine requires equal input vector types";
}
int lanes = input.IsScalableVector() ? input.VScaleFactor() : input.lanes();
if constexpr (!Combine) {
@@ -48,45 +48,45 @@ Type InferTypeVectorPart(const CallNode* call) {
: input.WithLanes(result_lanes);
}
-const Op& vectorhigh_op() {
- static const Op op = Op::Get("tirx.vectorhigh");
+const Op& vector_high_op() {
+ static const Op op = Op::Get("tirx.vector_high");
return op;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- OpDef("tirx.vectorhigh")
+ OpDef("tirx.vector_high")
.signature(sig::arg<PrimExpr>("vec", "The input vector."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeVectorPart<false>>())
- .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.vectorhigh"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.vector_high"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
}
-const Op& vectorlow_op() {
- static const Op op = Op::Get("tirx.vectorlow");
+const Op& vector_low_op() {
+ static const Op op = Op::Get("tirx.vector_low");
return op;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- OpDef("tirx.vectorlow")
+ OpDef("tirx.vector_low")
.signature(sig::arg<PrimExpr>("vec", "The input vector."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeVectorPart<false>>())
- .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.vectorlow"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.vector_low"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
}
-const Op& vectorcombine_op() {
- static const Op op = Op::Get("tirx.vectorcombine");
+const Op& vector_combine_op() {
+ static const Op op = Op::Get("tirx.vector_combine");
return op;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- OpDef("tirx.vectorcombine")
+ OpDef("tirx.vector_combine")
.signature(sig::arg<PrimExpr>("vec1", "The first input vector."),
sig::arg<PrimExpr>("vec2", "The second input vector."))
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeVectorPart<true>>())
- .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.vectorcombine"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.vector_combine"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
}
diff --git a/src/tirx/transform/lower_warp_memory.cc
b/src/tirx/transform/lower_warp_memory.cc
index 00c47aadc7..8456d8c9b7 100644
--- a/src/tirx/transform/lower_warp_memory.cc
+++ b/src/tirx/transform/lower_warp_memory.cc
@@ -203,9 +203,8 @@ class WarpStoreCoeffFinder : public StmtExprVisitor {
/// Visitor implementation
ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
- static const Op mma_fill_op = Op::Get("tirx.mma_fill");
+ static const Op mma_fill_op = Op::Get("tirx.cuda.mma_fill");
static const Op ptx_ldmatrix_legacy_op =
Op::Get("tirx.ptx_legacy.ldmatrix");
- static const Op mma_fill_legacy_op = Op::Get("tirx.mma_fill_legacy");
if (op->op.same_as(mma_fill_op) && GetTensorVar(op->args[1]) == buffer_) {
auto* local_size = op->args[0].as<IntImmNode>();
TVM_FFI_ICHECK(local_size) << "Integer expected for the first argument
of mma_fill";
@@ -215,12 +214,8 @@ class WarpStoreCoeffFinder : public StmtExprVisitor {
// ``... + lift(local_size) * tx`` from which the warp coefficient
// is derived.
UpdatePattern(op->args[4].as_or_throw<PrimExpr>());
- } else if (op->op.same_as(mma_fill_legacy_op) && GetTensorVar(op->args[1])
== buffer_) {
- auto* local_size = op->args[0].as<IntImmNode>();
- TVM_FFI_ICHECK(local_size) << "Integer expected for the first argument
of mma_fill_legacy";
- UpdateCoefficient(local_size->value.as<int>().value());
}
- // mma_store_legacy/ptx_mma_legacy only *use* the warp buffer
+ // mma_store/ptx_mma_legacy only *use* the warp buffer
// (read+rewrite); WarpStoreCoeffFinder relies on ldmatrix/mma_fill
// (the actual stores) for the warp coefficient.
@@ -430,12 +425,10 @@ class WarpAccessRewriter : public StmtExprMutator {
}
UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
override {
- static const Op mma_store_op = Op::Get("tirx.mma_store");
- static const Op mma_fill_op = Op::Get("tirx.mma_fill");
+ static const Op mma_store_op = Op::Get("tirx.cuda.mma_store");
+ static const Op mma_fill_op = Op::Get("tirx.cuda.mma_fill");
static const Op ptx_mma_legacy_op = Op::Get("tirx.ptx_legacy.mma");
static const Op ptx_ldmatrix_legacy_op =
Op::Get("tirx.ptx_legacy.ldmatrix");
- static const Op mma_store_legacy_op = Op::Get("tirx.mma_store_legacy");
- static const Op mma_fill_legacy_op = Op::Get("tirx.mma_fill_legacy");
if (op->op.same_as(mma_store_op)) {
return RewriteIndicesAt(op, {3});
}
@@ -454,14 +447,6 @@ class WarpAccessRewriter : public StmtExprMutator {
// access_ptr Call wrapping a shared-scope var.
return RewriteIndicesAt(op, {3});
}
- if (op->op.same_as(mma_store_legacy_op)) {
- // args: m, n, dst_ptr, src_ptr, src_offset, dst_stride
- return RewriteIndicesAt(op, {3});
- }
- if (op->op.same_as(mma_fill_legacy_op)) {
- // args: local_size, local_ptr, offset
- return RewriteIndicesAt(op, {1});
- }
return StmtExprMutator::Mutate_(op, inplace_mode);
}
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index 3d65ee9c27..1a1790d21c 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -209,12 +209,12 @@ class NoOpRemover : public IRMutatorWithAnalyzer {
bool HasSideEffect(const PrimExpr& value) {
if (ignore_profiler_call_) {
if (const CallNode* call = value.as<CallNode>()) {
- static const Op timer_init_cuda_op = Op::Get("tirx.timer_init_cuda");
- static const Op timer_start_cuda_op = Op::Get("tirx.timer_start_cuda");
- static const Op timer_end_cuda_op = Op::Get("tirx.timer_end_cuda");
- static const Op timer_finalize_cuda_op =
Op::Get("tirx.timer_finalize_cuda");
- if (call->op.same_as(timer_init_cuda_op) ||
call->op.same_as(timer_start_cuda_op) ||
- call->op.same_as(timer_end_cuda_op) ||
call->op.same_as(timer_finalize_cuda_op)) {
+ static const Op timer_init_op = Op::Get("tirx.cuda.timer_init");
+ static const Op timer_start_op = Op::Get("tirx.cuda.timer_start");
+ static const Op timer_end_op = Op::Get("tirx.cuda.timer_end");
+ static const Op timer_finalize_op =
Op::Get("tirx.cuda.timer_finalize");
+ if (call->op.same_as(timer_init_op) ||
call->op.same_as(timer_start_op) ||
+ call->op.same_as(timer_end_op) ||
call->op.same_as(timer_finalize_op)) {
return false;
}
}
diff --git a/tests/python/tirx-base/test_tir_op_types.py
b/tests/python/tirx-base/test_tir_op_types.py
index 2c38ddc078..73dd64606d 100644
--- a/tests/python/tirx-base/test_tir_op_types.py
+++ b/tests/python/tirx-base/test_tir_op_types.py
@@ -200,13 +200,13 @@ def test_tir_op_mma_store():
buffer_w.ty.elem_offset,
x,
)
- assert expr.op.name == "tirx.mma_store"
+ assert expr.op.name == "tirx.cuda.mma_store"
def test_tir_op_mma_fill():
buffer_w = tirx.decl_tensor([16, 8], dtype="int32", scope="warp",
offset_factor=1)
expr = _cuda_op.mma_fill("int32", 8, buffer_w.data,
buffer_w.ty.elem_offset)
- assert expr.op.name == "tirx.mma_fill"
+ assert expr.op.name == "tirx.cuda.mma_fill"
def test_op_ptx_cp_async():
@@ -227,18 +227,18 @@ def test_op_ptx_cp_async():
assert int(simplified_offset) == expected_offset
-def test_tir_op_vectorlow():
+def test_tir_op_vector_low():
buffer = tirx.decl_tensor((4, 4), "int8", offset_factor=1)
vec = buffer.vload([0, 0], dtype="int8x16")
- expr = tirx.vectorlow("int8x8", vec)
- assert expr.op.name == "tirx.vectorlow"
+ expr = tirx.vector_low("int8x8", vec)
+ assert expr.op.name == "tirx.vector_low"
-def test_tir_op_vectorhigh():
+def test_tir_op_vector_high():
buffer = tirx.decl_tensor((4, 4), "int8", offset_factor=1)
vec = buffer.vload([0, 0], dtype="int8x16")
- expr = tirx.vectorhigh("int8x8", vec)
- assert expr.op.name == "tirx.vectorhigh"
+ expr = tirx.vector_high("int8x8", vec)
+ assert expr.op.name == "tirx.vector_high"
def test_tir_op_dp4a():
@@ -249,11 +249,11 @@ def test_tir_op_dp4a():
assert expr.op.name == "tirx.gpu_dp4a"
-def test_tir_op_vectorcombine():
+def test_tir_op_vector_combine():
buffer = tirx.decl_tensor((4, 4), "int8", offset_factor=1)
vec = buffer.vload([0, 0], dtype="int8x16")
- expr = tirx.vectorcombine("int8x8", vec, vec)
- assert expr.op.name == "tirx.vectorcombine"
+ expr = tirx.vector_combine("int8x8", vec, vec)
+ assert expr.op.name == "tirx.vector_combine"
def test_tir_op_shift_left():
