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 e17fcc98f9 [REFACTOR][TIR] Replace type annotation markers with
explicit types (#20602)
e17fcc98f9 is described below
commit e17fcc98f9375fb7cbf0f7030d01e75db0e5ade4
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 8 11:54:15 2026 -0400
[REFACTOR][TIR] Replace type annotation markers with explicit types (#20602)
Replace synthetic `TypeAnnotation` expressions with explicit pointer
type information. `tvm_access_ptr` carries its independent access
element type in `ty_args`, while `ptr_byte_offset` uses its result
`PointerType`.
Update pointer construction, lowering, storage analyses and rewrites,
backend code generation, and script reconstruction to preserve access
units, address spaces, and requested pointer results.
---
include/tvm/tirx/op.h | 6 +--
include/tvm/tirx/stmt.h | 7 ----
python/tvm/backend/cuda/op.py | 24 +++++------
python/tvm/s_tir/tensor_intrin/riscv_cpu.py | 7 ++--
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/op/__init__.py | 45 +++++---------------
python/tvm/tirx/script/ir_builder/ir.py | 3 +-
python/tvm/tirx/script/ir_builder/op.py | 4 --
src/backend/metal/codegen/codegen_metal.cc | 3 +-
.../analysis/sblock_access_region_detector.cc | 6 +--
src/s_tir/transform/inject_permuted_layout.cc | 12 +++---
src/s_tir/transform/inject_software_pipeline.cc | 8 ++--
src/s_tir/transform/inject_virtual_thread.cc | 25 +++++------
.../transform/memhammer_tensorcore_rewrite.cc | 24 +++++------
.../transform/merge_shared_memory_allocations.cc | 14 +++----
src/s_tir/transform/storage_access.cc | 14 +++----
src/target/source/codegen_c.cc | 4 +-
src/tirx/ir/buffer.cc | 10 ++---
src/tirx/ir/stmt.cc | 12 ------
src/tirx/op/builtin.cc | 17 +++++---
src/tirx/op/op.cc | 21 ----------
src/tirx/script/printer/expr.cc | 15 +++++++
src/tirx/transform/lower_intrin.cc | 20 ++++-----
src/tirx/transform/storage_rewrite.cc | 48 ++++++++++------------
src/tirx/transform/vectorize_loop.cc | 2 +-
tests/python/codegen/test_target_codegen_riscv.py | 4 +-
tests/python/codegen/test_target_codegen_vulkan.py | 6 +--
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 36 ++++++++--------
.../test_meta_schedule_postproc_verify_gpu_code.py | 6 +--
.../test_meta_schedule_trace_apply.py | 6 +--
.../schedule/test_tir_schedule_cache_read_write.py | 12 +++---
.../s_tir/script/test_s_tir_script_printer.py | 12 +++---
.../test_s_tir_transform_inject_permuted_layout.py | 20 ++++-----
.../test_s_tir_transform_inject_ptx_async_copy.py | 6 +--
...st_s_tir_transform_memhammer_lower_auto_copy.py | 8 ++--
tests/python/tirx-base/test_tir_buffer.py | 30 +++++++-------
tests/python/tirx-base/test_tir_op_types.py | 11 ++---
.../test_tir_transform_lower_intrin.py | 6 +--
.../tirx/transform/test_transform_lower_tirx.py | 4 +-
39 files changed, 228 insertions(+), 292 deletions(-)
diff --git a/include/tvm/tirx/op.h b/include/tvm/tirx/op.h
index c88bd88076..7017b70965 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/tirx/op.h
@@ -255,10 +255,10 @@ TVM_DLL const Op& prefetch_op();
* rw_mask is a bit_mask setting whether the access is a read(1) or write(2).
* The access is assume to happen in the current expression.
*
- * PtrType tvm_access_ptr(Expr dtype, DType* data,
+ * PtrType tvm_access_ptr<DType>(DType* data,
* int offset, int extent,
* int rw_mask) {
- * // DType == dtype.type();
+ * // DType is the independent access type in ty_args[0].
* return &data[offset];
* }
*/
@@ -267,7 +267,7 @@ TVM_DLL const Op& tvm_access_ptr_op();
/*!
* \brief Cast a handle to a typed pointer after adding a byte offset.
*
- * DType* ptr_byte_offset(void* data, int byte_offset, Expr dtype) {
+ * DType* ptr_byte_offset(void* data, int byte_offset) {
* return reinterpret_cast<DType*>(reinterpret_cast<char*>(data) +
byte_offset);
* }
*/
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 2a2df56114..8a2d730a08 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -720,13 +720,6 @@ constexpr const char* buffer_allocated_addr =
"buffer_allocated_addr";
constexpr const char* kPersistentKernel = "tirx.persistent_kernel";
} // namespace attr
-/*!
- * \brief Create a type annotation expression
- * \param dtype The data type
- * \param span The location of this object in the source code.
- * \return Expr a expression with dtype.
- */
-TVM_DLL PrimExpr TypeAnnotation(PrimType dtype, Span span = Span());
// overload printing of for type.
TVM_DLL std::ostream& operator<<(std::ostream& os, ForKind kind);
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index 971979eac8..c60f788576 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -529,7 +529,7 @@ def _wrap_or_fold_access_ptr(ptr, offset, elem_dtype):
Several s_tir tensor intrinsics already pass ``buffer.access_ptr(...)``
(an ``tvm_access_ptr`` Call) for the pointer argument. Naively wrapping
that again yields a nested ``tvm_access_ptr(... access_ptr(...) ...)``
- whose ``args[1]`` is a Call rather than a Var, which crashes the
+ whose ``args[0]`` is a Call rather than a Var, which crashes the
lowering rule (Downcast<Var> at intrin_rule.cc) and several s_tir
passes that assume a raw buffer var. Detect that case and fold the
outer offset into the inner one.
@@ -540,22 +540,20 @@ def _wrap_or_fold_access_ptr(ptr, offset, elem_dtype):
)
if is_access_ptr_call:
# Inner Call already wraps the buffer var. Reuse its inner var and
- # inner element dtype (the marker type_annotation), and add the
+ # inner access element type, and add the
# outer offset (which is in `elem_dtype` units, same convention as
# the inner since both come from the same buffer).
inner_args = ptr.args
- inner_marker = inner_args[0]
- inner_var = inner_args[1]
- inner_offset = inner_args[2]
- rw_mask = inner_args[4]
- return call_intrin(
- ptr.ty,
+ inner_var = inner_args[0]
+ inner_offset = inner_args[1]
+ rw_mask = inner_args[3]
+ return Call(
"tirx.tvm_access_ptr",
- inner_marker,
- inner_var,
- inner_offset + offset,
- 1,
- rw_mask,
+ [inner_var, inner_offset + offset, 1, rw_mask],
+ ty=ptr.ty,
+ attrs=ptr.attrs,
+ ty_args=ptr.ty_args,
+ span=ptr.span,
)
return tvm_access_ptr(elem_dtype, ptr, offset, 1, 1)
diff --git a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
index 14fbe1fadd..3e35ecc531 100644
--- a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
+++ b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
@@ -15,7 +15,6 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name,line-too-long
-# ruff: noqa: E501
"""Intrinsics for RISCV tensorization"""
import logging
@@ -119,7 +118,7 @@ def rvv_vec_dot_product_kernels(
vec_A = T.call_llvm_intrin(
"llvm.riscv.vle",
T.broadcast(T.Cast(data_dtype, 0), T.vscale() * d_dtype_lanes),
- T.tvm_access_ptr(T.type_annotation(ty=data_dtype), A.data, 0,
n_elems, 1),
+ T.tvm_access_ptr(data_dtype, A.data, 0, n_elems, 1),
T.int64(n_elems), ty=f"{data_dtype}xvscalex{d_dtype_lanes}")
for i in range(n_lanes):
@@ -130,7 +129,7 @@ def rvv_vec_dot_product_kernels(
vec_B_row = T.call_llvm_intrin(
"llvm.riscv.vle",
T.broadcast(T.Cast(data_dtype, 0), T.vscale() *
w_dtype_lanes),
- T.tvm_access_ptr(T.type_annotation(ty=weight_dtype),
B.data, i * n_elems, n_elems, 1),
+ T.tvm_access_ptr(weight_dtype, B.data, i * n_elems,
n_elems, 1),
T.int64(n_elems),
ty=f"{weight_dtype}xvscalex{w_dtype_lanes}")
product = T.call_llvm_intrin(
@@ -146,7 +145,7 @@ def rvv_vec_dot_product_kernels(
ini_acc = T.call_llvm_intrin(
"llvm.riscv.vle",
T.broadcast(T.Cast(out_dtype, 0), T.vscale() *
o_dtype_lanes),
- T.tvm_access_ptr(T.type_annotation(ty=out_dtype),
C.data, i, 1, 1),
+ T.tvm_access_ptr(out_dtype, C.data, i, 1, 1),
T.int64(1), ty=f"{out_dtype}xvscalex{o_dtype_lanes}")
red_sum = T.call_llvm_intrin(
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 6e7eda8b05..f69f0e9886 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -65,7 +65,7 @@ from .op import call_llvm_intrin, call_llvm_pure_intrin, all,
any, min_value, ma
from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
from .op import handle_add_byte_offset, tvm_struct_get, tvm_struct_set
from .op import address_of, assume, assume_aligned, undef
-from .op import tvm_thread_allreduce, type_annotation, tvm_access_ptr,
ptr_byte_offset
+from .op import tvm_thread_allreduce, tvm_access_ptr, ptr_byte_offset
from .op import tvm_throw_last_error, cpu_parallel_barrier
from .op import (
tvm_load_matrix_sync,
diff --git a/python/tvm/tirx/op/__init__.py b/python/tvm/tirx/op/__init__.py
index 710b2edb92..f31082fdb3 100644
--- a/python/tvm/tirx/op/__init__.py
+++ b/python/tvm/tirx/op/__init__.py
@@ -1107,32 +1107,13 @@ def tvm_warp_activemask(*, ty=None, span=None):
return call_intrin(ty, "tirx.tvm_warp_activemask", span=span)
-def type_annotation(dtype):
- """Create a type annotation expression
-
- Parameters
- ----------
- dtype : Expr
- The data type.
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin(dtype, "tirx.type_annotation")
-
-
def tvm_access_ptr(ptype, data, offset, extent, rw_mask, *, ty=None,
span=None):
"""Get head access address with memory access pattern info
Parameters
----------
- ptype : Expr, PrimType, or str
- The data type of pointer. If a ``PrimType`` or ``str``, it is wrapped
- via :func:`type_annotation` so that the lowering rule (which reads
- ``args[0].dtype()`` for the cast type) sees the intended dtype instead
- of StringType from a string literal.
+ ptype : PrimType or str
+ The accessed element type. Offsets and extents are in units of this
type.
data : DType*
The data of pointer.
@@ -1151,16 +1132,12 @@ def tvm_access_ptr(ptype, data, offset, extent,
rw_mask, *, ty=None, span=None):
call : Expr
The call expression.
"""
- if isinstance(ptype, str | PrimType):
- ptype = type_annotation(ptype)
- return call_intrin(
- ty,
+ ptype = PrimType(ptype) if isinstance(ptype, str) else ptype
+ return Call(
"tirx.tvm_access_ptr",
- ptype,
- data,
- offset,
- extent,
- rw_mask,
+ [data, offset, extent, rw_mask],
+ ty=ty,
+ ty_args=[ptype],
span=span,
)
@@ -1171,14 +1148,14 @@ def ptr_byte_offset(data, byte_offset, dtype, *,
ty=None, span=None):
``byte_offset`` is always in bytes. Use this when the source CUDA shape
needs an explicitly typed local pointer derived from a byte-addressed base.
"""
- if isinstance(dtype, str | PrimType):
- dtype = type_annotation(dtype)
+ dtype = PrimType(dtype) if isinstance(dtype, str) else dtype
+ data_type = getattr(data, "ty", None)
+ storage_scope = data_type.storage_scope if isinstance(data_type,
PointerType) else "global"
return call_intrin(
- ty,
+ PointerType(dtype, storage_scope) if ty is None else ty,
"tirx.ptr_byte_offset",
data,
byte_offset,
- dtype,
span=span,
)
diff --git a/python/tvm/tirx/script/ir_builder/ir.py
b/python/tvm/tirx/script/ir_builder/ir.py
index 0bde7d5538..b2b97e4f77 100644
--- a/python/tvm/tirx/script/ir_builder/ir.py
+++ b/python/tvm/tirx/script/ir_builder/ir.py
@@ -47,7 +47,7 @@ from tvm.script.parser.protocol_registry import (
from tvm.target import Target
# pylint: disable=unused-import
-from tvm.tirx import Expr, IndexMap, Var, is_tensor_var, type_annotation
+from tvm.tirx import Expr, IndexMap, Var, is_tensor_var
from tvm.tirx.exec_scope import ExecScope, ScopeIdDef
# import tirx.expr for direct ir construction to pass structural_equal
comparison
@@ -1758,7 +1758,6 @@ __all__ = [
"thread_id",
"thread_id_in_wg",
"tmem",
- "type_annotation",
"u8",
"u16",
"u32",
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index 25fdbe6e68..423aa1849c 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -595,9 +595,6 @@ vscale = _tir_op.vscale
ignore_loop_partition = _tir_op.ignore_loop_partition
-type_annotation = _ir.op._make_op_api(_ir.Op.get("tirx.type_annotation"),
__name__)
-
-
reinterpret = _ir.op._make_op_api(_ir.Op.get("tirx.reinterpret"), __name__)
@@ -961,7 +958,6 @@ __all__ = [
"tvm_warp_shuffle_down",
"tvm_warp_shuffle_up",
"tvm_warp_shuffle_xor",
- "type_annotation",
"undef",
"vectorcombine",
"vectorhigh",
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index 245661bec1..30f1a148cf 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -511,8 +511,7 @@ void CodeGenMetal::Dispatch_(const CallNode* op,
std::ostream& os) { // NOLINT(
<< PrintExpr(c) << "[" << PrintExpr(op->args[7]) << "])";
} else if (op->op.same_as(tirx::ptr_byte_offset_op()) ||
op->op.same_as(tirx::handle_add_byte_offset_op())) {
- bool is_typed_offset = op->op.same_as(tirx::ptr_byte_offset_op());
- TVM_FFI_ICHECK_EQ(op->args.size(), is_typed_offset ? 3U : 2U);
+ TVM_FFI_ICHECK_EQ(op->args.size(), 2U);
const auto* pointer_type = op->ty.as<PointerTypeNode>();
TVM_FFI_ICHECK(pointer_type)
<< "Metal pointer byte offsets must have a pointer result type, but
got " << op->ty;
diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc
b/src/s_tir/analysis/sblock_access_region_detector.cc
index e1de6780b8..27368ffac5 100644
--- a/src/s_tir/analysis/sblock_access_region_detector.cc
+++ b/src/s_tir/analysis/sblock_access_region_detector.cc
@@ -283,12 +283,12 @@ ffi::Optional<VisitInterrupt>
BlockReadWriteDetector::Visit_(const CallNode* op)
return std::nullopt;
}
if (op->op.same_as(tirx::tvm_access_ptr_op())) {
- const VarNode* buffer_var = op->args[1].as<VarNode>();
- if (const auto* data = op->args[1].as<CallNode>();
+ const VarNode* buffer_var = op->args[0].as<VarNode>();
+ if (const auto* data = op->args[0].as<CallNode>();
data && data->op.same_as(tirx::buffer_data_op())) {
buffer_var = data->args[0].as<VarNode>();
}
- const IntImmNode* access_mask = op->args[4].as<IntImmNode>();
+ const IntImmNode* access_mask = op->args[3].as<IntImmNode>();
if (buffer_var && access_mask) {
auto it = buffer_var_map_.find(ffi::GetRef<Var>(buffer_var));
if (it != buffer_var_map_.end()) {
diff --git a/src/s_tir/transform/inject_permuted_layout.cc
b/src/s_tir/transform/inject_permuted_layout.cc
index 8f773ef7f8..cb9419519d 100644
--- a/src/s_tir/transform/inject_permuted_layout.cc
+++ b/src/s_tir/transform/inject_permuted_layout.cc
@@ -249,7 +249,7 @@ class PermutedLayoutInjector : public IRMutatorWithAnalyzer
{
}
Expr HandleAccessPtrAndOffset(Expr access_ptr, ffi::Optional<PrimExpr>
offset = std::nullopt) {
- // The 2th arg of T.tvm_access_ptr call is offset, we set it to 0 and
accumulate it to
+ // The second arg of T.tvm_access_ptr call is offset, we set it to 0 and
accumulate it to
// smem_offset
TVM_FFI_ICHECK(access_ptr->IsInstance<CallNode>())
<< "Invalid access ptr for permuted layout: " << access_ptr;
@@ -257,15 +257,15 @@ class PermutedLayoutInjector : public
IRMutatorWithAnalyzer {
TVM_FFI_ICHECK(access_ptr_call->op.same_as(tirx::tvm_access_ptr_op()))
<< "Invalid access ptr for permuted layout: " << access_ptr;
- auto data_var = GetBufferDataVar(access_ptr_call->args[1]);
+ auto data_var = GetBufferDataVar(access_ptr_call->args[0]);
TVM_FFI_ICHECK(data_var.has_value())
- << "Expected a buffer data expression, but received " <<
access_ptr_call->args[1];
+ << "Expected a buffer data expression, but received " <<
access_ptr_call->args[0];
auto buffer_map_iter = buffer_map_.find(data_var.value());
TVM_FFI_ICHECK(buffer_map_iter != buffer_map_.end())
- << "The buffer corresponding to data Var " << access_ptr_call->args[1]
<< " is not found";
+ << "The buffer corresponding to data Var " << access_ptr_call->args[0]
<< " is not found";
int buffer_row_size = CheckAndGetBufferRowSize(buffer_map_iter->second);
- PrimExpr smem_offset = access_ptr_call->args[2].as_or_throw<PrimExpr>() +
+ PrimExpr smem_offset = access_ptr_call->args[1].as_or_throw<PrimExpr>() +
(offset.has_value() ? offset.value() : 0);
// Convert offset to 2-dimension, reindex it and convert it back
@@ -276,7 +276,7 @@ class PermutedLayoutInjector : public IRMutatorWithAnalyzer
{
auto new_offset = analyzer_->Simplify(new_indices[0] * buffer_row_size +
new_indices[1]);
auto new_access_ptr = access_ptr_call.CopyOnWrite();
- new_access_ptr->args.Set(2, new_offset);
+ new_access_ptr->args.Set(1, new_offset);
return access_ptr_call;
}
diff --git a/src/s_tir/transform/inject_software_pipeline.cc
b/src/s_tir/transform/inject_software_pipeline.cc
index 2c7e77ae72..473662e82f 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/src/s_tir/transform/inject_software_pipeline.cc
@@ -141,7 +141,7 @@ class PipelineOpaqueAccessRewriter {
new_args.Set(0, new_buffer.data());
new_args.Set(
4, RewriteWmmaFragmentIndex(buffer, new_buffer,
call->args[4].as_or_throw<PrimExpr>()));
- return Call(call->ty, call->op, new_args, call->attrs, {}, call->span);
+ return Call(call->ty, call->op, new_args, call->attrs, call->ty_args,
call->span);
}
} else if (call->op.same_as(mma_sync)) {
ffi::Array<Expr> new_args = call->args;
@@ -157,9 +157,9 @@ class PipelineOpaqueAccessRewriter {
new_args.Set(i * 2 + 1, new_index);
}
}
- return Call(call->ty, call->op, new_args, call->attrs, {}, call->span);
+ return Call(call->ty, call->op, new_args, call->attrs, call->ty_args,
call->span);
} else if (call->op.same_as(access_ptr)) {
- return RewriteBufferAccess(call, {1});
+ return RewriteBufferAccess(call, {0});
} else if (call->op.same_as(ptx_mma_legacy)) {
return RewriteBufferAccess(call, {6, 8, 10});
} else if (call->op.same_as(ptx_ldmatrix_legacy)) {
@@ -217,7 +217,7 @@ class PipelineOpaqueAccessRewriter {
new_args.Set(i + 1, new_index);
}
}
- return Call(call->ty, call->op, new_args, call->attrs, {}, call->span);
+ return Call(call->ty, call->op, new_args, call->attrs, call->ty_args,
call->span);
}
const ffi::Map<Var, TensorVar>& buffer_data_to_buffer_;
diff --git a/src/s_tir/transform/inject_virtual_thread.cc
b/src/s_tir/transform/inject_virtual_thread.cc
index 4357be5ecc..a33f0c132f 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -97,14 +97,14 @@ class ExprTouched final : public StmtExprVisitor {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->args[i]));
}
} else if (op->op.same_as(tirx::tvm_access_ptr_op())) {
- const auto* rw_mask = op->args[4].as<IntImmNode>();
- auto buffer = GetBufferDataVar(op->args[1]);
+ const auto* rw_mask = op->args[3].as<IntImmNode>();
+ auto buffer = GetBufferDataVar(op->args[0]);
if (!buffer.has_value()) {
// Nested access pointers are valid pointer expressions. Visit the
// inner pointer and this access's offset instead of assuming a raw
// buffer Var at every level.
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->args[1]));
- return this->Visit(op->args[2].as_or_throw<PrimExpr>());
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->args[0]));
+ return this->Visit(op->args[1].as_or_throw<PrimExpr>());
}
const VarNode* buffer_var = buffer.value().get();
TVM_FFI_ICHECK(rw_mask);
@@ -115,7 +115,7 @@ class ExprTouched final : public StmtExprVisitor {
if (rw_mask->value & 2) {
HandleWriteVar(buffer_var);
}
-
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->args[2].as_or_throw<PrimExpr>()));
+
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->args[1].as_or_throw<PrimExpr>()));
} else {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
}
@@ -339,25 +339,26 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
visit_touched_var_ = true;
return GetRemappedBuffer(buffer.as_or_throw<TensorVar>(),
it->second).data();
} else if (op->op.same_as(tirx::tvm_access_ptr_op())) {
- TVM_FFI_ICHECK_EQ(op->args.size(), 5U);
- PrimType dtype = op->args[0].as_or_throw<PrimExpr>().ty();
- auto buffer = GetBufferDataVar(op->args[1]);
+ TVM_FFI_ICHECK_EQ(op->args.size(), 4U);
+ PrimType dtype = op->ty_args[0].as_or_throw<PrimType>();
+ auto buffer = GetBufferDataVar(op->args[0]);
if (!buffer.has_value()) {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
auto it = alloc_remap_.find(buffer.value().get());
if (it == alloc_remap_.end()) return StmtExprMutator::Mutate_(op,
inplace_mode);
visit_touched_var_ = true;
- PrimExpr offset =
Mutate(op->args[2]).ValueOrUnchanged(op->args[2]).as_or_throw<PrimExpr>();
- PrimExpr extent =
Mutate(op->args[3]).ValueOrUnchanged(op->args[3]).as_or_throw<PrimExpr>();
+ PrimExpr offset =
Mutate(op->args[1]).ValueOrUnchanged(op->args[1]).as_or_throw<PrimExpr>();
+ PrimExpr extent =
Mutate(op->args[2]).ValueOrUnchanged(op->args[2]).as_or_throw<PrimExpr>();
PrimExpr stride = it->second / prim::MakeConst(offset.ty(),
dtype.lanes());
offset = RewriteIndex(offset, stride);
Expr data =
buffer.value()->ty.as<TensorTypeNode>()
? GetRemappedBuffer(buffer.value().as_or_throw<TensorVar>(),
it->second).data()
- : op->args[1];
+ : op->args[0];
- return Call(op->ty, op->op, {op->args[0], data, offset, extent,
op->args[4]});
+ return Call(op->ty, op->op, {data, offset, extent, op->args[3]},
op->attrs, op->ty_args,
+ op->span);
} else {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
index 121bc1e509..9fa81b516c 100644
--- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
+++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
@@ -187,12 +187,12 @@ Stmt RewriteWmmaLoad(Stmt stmt) {
/*op=*/tirx::tvm_access_ptr_op(),
/*args=*/
ffi::Array<Expr>{
- /*0:*/ TypeAnnotation(new_src_buffer->dtype),
- /*1:*/ new_src_buffer.data(),
- /*2:*/ new_src_buffer->elem_offset,
- /*3:*/
new_src_buffer->strides[new_src_buffer->strides.size() - 2] * 16,
- /*4:*/ PrimExpr(1),
- }),
+ /*0:*/ new_src_buffer.data(),
+ /*1:*/ new_src_buffer->elem_offset,
+ /*2:*/
new_src_buffer->strides[new_src_buffer->strides.size() - 2] * 16,
+ /*3:*/ PrimExpr(1),
+ },
+ {}, {new_src_buffer->dtype}),
/*6:*/
new_src_buffer->strides[new_src_buffer->strides.size() - 2],
/*7:*/ StringImm(layout),
})),
@@ -285,12 +285,12 @@ Stmt RewriteWmmaStore(Stmt stmt) {
/*data=*/new_tgt_buffer.data()->ty,
/*op=*/tirx::tvm_access_ptr_op(),
ffi::Array<Expr>{
- /*0:*/
TypeAnnotation(new_tgt_buffer->dtype),
- /*1:*/ new_tgt_buffer.data(),
- /*2:*/ new_tgt_buffer->elem_offset,
- /*3:*/ new_tgt_buffer->strides[0] *
16,
- /*4:*/ PrimExpr(2),
- }),
+ /*0:*/ new_tgt_buffer.data(),
+ /*1:*/ new_tgt_buffer->elem_offset,
+ /*2:*/ new_tgt_buffer->strides[0] *
16,
+ /*3:*/ PrimExpr(2),
+ },
+ {}, {new_tgt_buffer->dtype}),
/*6:*/ new_tgt_buffer->strides[0],
/*7:*/ StringImm("row_major")})),
/*init=*/std::nullopt,
diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc
b/src/s_tir/transform/merge_shared_memory_allocations.cc
index 56c353f1c1..39ed921f6d 100644
--- a/src/s_tir/transform/merge_shared_memory_allocations.cc
+++ b/src/s_tir/transform/merge_shared_memory_allocations.cc
@@ -665,9 +665,9 @@ class SharedMemoryRewriter : public StmtExprMutator {
UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
final {
static const Op ptx_cp_async_op = Op::Get("tirx.s_tir.cp_async_raw");
if (op->op.same_as(tirx::tvm_access_ptr_op())) {
- TVM_FFI_ICHECK_EQ(op->args.size(), 5U);
- DLDataType dtype = op->args[0].as_or_throw<PrimExpr>().ty()->dtype;
- auto buffer_opt = GetBufferDataVar(op->args[1]);
+ TVM_FFI_ICHECK_EQ(op->args.size(), 4U);
+ DLDataType dtype = op->ty_args[0].as_or_throw<PrimType>()->dtype;
+ auto buffer_opt = GetBufferDataVar(op->args[0]);
if (!buffer_opt.has_value()) {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
@@ -684,10 +684,10 @@ class SharedMemoryRewriter : public StmtExprMutator {
?
GetUpdatedBuffer(buffer.as_or_throw<TensorVar>()).data()
:
scope_stack_.back().merged_buffer.value().data();
- PrimExpr offset =
Mutate(op->args[2]).ValueOrUnchanged(op->args[2]).as_or_throw<PrimExpr>();
- PrimExpr extent =
Mutate(op->args[3]).ValueOrUnchanged(op->args[3]).as_or_throw<PrimExpr>();
- return Call(op->ty, op->op,
- {op->args[0], merged_data, extra_offset + offset, extent,
op->args[4]});
+ PrimExpr offset =
Mutate(op->args[1]).ValueOrUnchanged(op->args[1]).as_or_throw<PrimExpr>();
+ PrimExpr extent =
Mutate(op->args[2]).ValueOrUnchanged(op->args[2]).as_or_throw<PrimExpr>();
+ return Call(op->ty, op->op, {merged_data, extra_offset + offset, extent,
op->args[3]},
+ op->attrs, op->ty_args, op->span);
} else if (op->op.same_as(ptx_cp_async_op)) {
TVM_FFI_ICHECK((op->args.size() == 5U) || (op->args.size() == 6U));
auto buffer_opt = GetBufferDataVar(op->args[0]);
diff --git a/src/s_tir/transform/storage_access.cc
b/src/s_tir/transform/storage_access.cc
index bba32badd7..f5913de776 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -295,20 +295,20 @@ ffi::Optional<VisitInterrupt>
StorageAccessVisitor::Visit_(const CallNode* op) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
}
} else if (op->op.same_as(tirx::tvm_access_ptr_op())) {
- TVM_FFI_ICHECK_EQ(op->args.size(), 5U);
- PrimType dtype = op->args[0].as_or_throw<PrimExpr>().ty();
- auto buffer_var = GetBufferDataVar(op->args[1]);
+ TVM_FFI_ICHECK_EQ(op->args.size(), 4U);
+ PrimType dtype = op->ty_args[0].as_or_throw<PrimType>();
+ auto buffer_var = GetBufferDataVar(op->args[0]);
if (!buffer_var.has_value()) {
- // args[1] is not a raw Var — e.g. a nested tvm_access_ptr or some
+ // args[0] is not a raw Var — e.g. a nested tvm_access_ptr or some
// other PrimExpr. Recurse into sub-exprs so any inner buffer var
// refs still get visited, but don't try to record an access entry
// here (GetScope on a null Var would dereference a null pointer).
return StmtExprVisitor::Visit_(op);
}
Var buffer = ResolveBuffer(buffer_var.value());
- PrimExpr offset = op->args[2].as_or_throw<PrimExpr>();
- PrimExpr extent = op->args[3].as_or_throw<PrimExpr>();
- const IntImmNode* flag = op->args[4].as<IntImmNode>();
+ PrimExpr offset = op->args[1].as_or_throw<PrimExpr>();
+ PrimExpr extent = op->args[2].as_or_throw<PrimExpr>();
+ const IntImmNode* flag = op->args[3].as<IntImmNode>();
StorageScope scope = GetScope(buffer_var.value());
// The buffer scope.
if (Enabled(buffer.get(), scope)) {
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 7604fb2653..b10c648863 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -838,9 +838,9 @@ void CodeGenC::Dispatch_(const CallNode* op, std::ostream&
os) { // NOLINT(*)
this->PrintExpr(op->args[0], os);
os << " == NULL)";
} else if (op->op.same_as(tirx::ptr_byte_offset_op())) {
- TVM_FFI_ICHECK_EQ(op->args.size(), 3U);
+ TVM_FFI_ICHECK_EQ(op->args.size(), 2U);
os << "((";
- PrintType(op->args[2].as_or_throw<PrimExpr>().ty(), os);
+ PrintType(op->ty.as_or_throw<PointerType>()->element_type, os);
os << "*)(((char*)";
this->PrintExpr(op->args[0], os);
os << ") + ";
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 475fe0d736..514f3352bb 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -621,7 +621,7 @@ Expr TensorVar::access_ptr(int access_mask, PointerType
ptr_type, int content_la
// requested type controls its pointee, while the buffer controls its address
// space (for example, shared or local memory).
ptr_type = PointerType(ptr_type->element_type, self->storage_scope);
- PrimExpr e_dtype{ffi::UnsafeInit{}};
+ PrimType access_dtype = self->dtype;
PrimExpr extent{ffi::UnsafeInit{}};
if (self->shape.size() == 0) {
extent = IntImm(PrimType(self->DefaultIndexType()), 1);
@@ -635,18 +635,16 @@ Expr TensorVar::access_ptr(int access_mask, PointerType
ptr_type, int content_la
}
PrimExpr elem_offset = self->elem_offset + offset;
if (content_lanes > 1) {
- e_dtype =
tirx::TypeAnnotation(PrimType(self->dtype).WithLanes(content_lanes));
+ access_dtype = PrimType(self->dtype).WithLanes(content_lanes);
extent = extent / MakeConst(self->elem_offset.ty(), content_lanes);
elem_offset = self->elem_offset / MakeConst(self->elem_offset.ty(),
content_lanes);
- } else {
- e_dtype = tirx::TypeAnnotation(self->dtype);
}
if (input_extent.has_value()) {
extent = input_extent.value();
}
- ffi::Array<Expr> acc_args{e_dtype, data(), elem_offset, extent,
IntImm::Int32(access_mask)};
- return Call(ptr_type, tirx::tvm_access_ptr_op(), acc_args);
+ ffi::Array<Expr> acc_args{data(), elem_offset, extent,
IntImm::Int32(access_mask)};
+ return Call(ptr_type, tirx::tvm_access_ptr_op(), acc_args, {},
{access_dtype});
}
TensorVar::TensorVar(ffi::String name, TensorType type, Span span)
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 270ffd5674..269b57cf4e 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -1253,17 +1253,5 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](ScopeIdDef def, Span span) { return
ScopeIdDefStmt(def, span); });
}
-PrimExpr TypeAnnotation(PrimType dtype, Span span) {
- static const Op type_annotation_op = Op::Get("tirx.type_annotation");
- return Call(dtype, type_annotation_op, {}, {}, {},
span).as_or_throw<PrimExpr>();
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- OpDef("tirx.type_annotation")
- .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.type_annotation"))
- .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
-}
-
} // namespace tirx
} // namespace tvm
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 8ae57c9318..ee755d37ef 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -96,6 +96,13 @@ Type InferTypePointerOffset(const CallNode* call) {
return PointerType(element_type, pointer ? pointer.value()->storage_scope :
"global");
}
+Type InferTypeAccessPointer(const CallNode* call) {
+ TVM_FFI_CHECK_EQ(call->ty_args.size(), 1U, ValueError);
+ TVM_FFI_CHECK_EQ(call->args.size(), 4U, ValueError);
+ auto pointer = call->args[0]->ty.as<PointerType>();
+ return PointerType(call->ty_args[0], pointer ?
pointer.value()->storage_scope : "global");
+}
+
Type InferTypePrintBuffer(const CallNode* call) {
TVM_FFI_CHECK_GE(call->args.size(), 2U, ValueError);
return
PrimType(ffi::StringToDLDataType(call->args[1].as_or_throw<StringImm>()->value));
@@ -470,22 +477,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.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<IntExpr>("offset", "The offset."),
+ .signature(sig::ty_arg<PrimType>("access_dtype", "The accessed element
type."),
+ sig::arg("data", "The input data."),
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<FInferType>("FInferType",
FInferType::FromNative<&InferTypePointerOffset<1, 0>>())
+ .set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeAccessPointer>())
.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<IntExpr>("byte_offset", "Offset in bytes."),
- sig::arg("dtype", "Type annotation for pointed-to elements."))
+ sig::arg<IntExpr>("byte_offset", "Offset in bytes."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.ptr_byte_offset"))
- .set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypePointerOffset<0, 2>>())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure));
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index 40a5540056..d6914151ac 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -68,27 +68,6 @@ Type GetType(const PrimExpr& expr) {
}
}
- static const Op type_annotation_op = Op::Get("tirx.type_annotation");
- if (auto* access = expr.as<CallNode>()) {
- if (access->op.same_as(tirx::tvm_access_ptr_op())) {
- TVM_FFI_ICHECK(access->args.size())
- << "Builtin tvm_access_ptr() may not have empty arguments";
- auto type_annotation = access->args[0].as_or_throw<Call>();
- TVM_FFI_ICHECK(type_annotation->op.same_as(type_annotation_op))
- << "Expected the first argument of builtin tvm_access_ptr() "
- << "to be a type annotation, but found " << type_annotation->op;
- return PointerType(type_annotation->ty.as_or_throw<PrimType>());
- }
- if (access->op.same_as(tirx::ptr_byte_offset_op())) {
- TVM_FFI_ICHECK_EQ(access->args.size(), 3U);
- auto type_annotation = access->args[2].as_or_throw<Call>();
- TVM_FFI_ICHECK(type_annotation->op.same_as(type_annotation_op))
- << "Expected the third argument of builtin ptr_byte_offset() "
- << "to be a type annotation, but found " << type_annotation->op;
- return PointerType(type_annotation->ty.as_or_throw<PrimType>());
- }
- }
-
if (auto* address_of = expr.as<CallNode>()) {
if (address_of->op.same_as(tirx::address_of_op())) {
TVM_FFI_ICHECK_EQ(address_of->args.size(), 1)
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index 62c17398ef..949043cbc3 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -324,6 +324,21 @@ ffi::Optional<ExprDoc>
ConsumedPTXAddressDocTranslate(DocTranslatorObj* d, const
namespace {
+// Pointer helpers have semantic type parameters outside the stored value
arguments.
+ffi::Optional<ExprDoc> PointerCallDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* call =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
+ return RawCall(d, call);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ for (const char* name : {"tirx.tvm_access_ptr", "tirx.ptr_byte_offset"}) {
+ OpDef(name).set_attr<FDocTranslate>(kOpCallDocTranslate,
+
FDocTranslate::FromNative<&PointerCallDocTranslate>());
+ }
+}
+
// The isnan convenience helper folds constants and widens float16 inputs.
ffi::Optional<ExprDoc> IsNaNDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object*) {
diff --git a/src/tirx/transform/lower_intrin.cc
b/src/tirx/transform/lower_intrin.cc
index 3afa62725a..37d79e2ec9 100644
--- a/src/tirx/transform/lower_intrin.cc
+++ b/src/tirx/transform/lower_intrin.cc
@@ -50,27 +50,27 @@ struct AccessPtrBufferAlias {
static Expr LowerAccessPtr(const CallNode* call,
std::vector<AccessPtrBufferAlias>* buffer_aliases) {
- TVM_FFI_ICHECK_EQ(call->args.size(), 5U);
- PrimType dtype = call->args[0].as_or_throw<PrimExpr>().ty();
- PrimExpr offset = call->args[2].as_or_throw<PrimExpr>();
+ TVM_FFI_ICHECK_EQ(call->args.size(), 4U);
+ PrimType dtype = call->ty_args[0].as_or_throw<PrimType>();
+ PrimExpr offset = call->args[1].as_or_throw<PrimExpr>();
TVM_FFI_ICHECK(call->ty.as<PointerTypeNode>());
// An access pointer may itself be used as the base of another access
// pointer. Fold those offsets before constructing the synthetic
- // TensorLoad so lowering never assumes that args[1] is immediately a Var.
- Expr buffer = call->args[1];
+ // TensorLoad so lowering never assumes that args[0] is immediately a Var.
+ Expr buffer = call->args[0];
while (const auto* inner = buffer.as<CallNode>()) {
if (!inner->op.same_as(tvm_access_ptr_op())) break;
- TVM_FFI_ICHECK_EQ(inner->args.size(), 5U);
- PrimType inner_dtype = inner->args[0].as_or_throw<PrimExpr>().ty();
+ TVM_FFI_ICHECK_EQ(inner->args.size(), 4U);
+ PrimType inner_dtype = inner->ty_args[0].as_or_throw<PrimType>();
TVM_FFI_ICHECK_EQ(inner_dtype, dtype)
<< "Nested tvm_access_ptr calls must use the same element type";
- PrimExpr inner_offset = inner->args[2].as_or_throw<PrimExpr>();
+ PrimExpr inner_offset = inner->args[1].as_or_throw<PrimExpr>();
if (inner_offset.ty() != offset.ty()) {
inner_offset = prim::Cast(offset.ty(), inner_offset);
}
offset = inner_offset + offset;
- buffer = inner->args[1];
+ buffer = inner->args[0];
}
const auto* buffer_data = buffer.as<CallNode>();
@@ -81,7 +81,7 @@ static Expr LowerAccessPtr(const CallNode* call,
const auto* buffer_node = buffer.as<VarNode>();
TVM_FFI_ICHECK(buffer_node)
- << "tvm_access_ptr expects a buffer Var or nested tvm_access_ptr as
args[1], but got "
+ << "tvm_access_ptr expects a buffer Var or nested tvm_access_ptr as
args[0], but got "
<< buffer;
Var buffer_var = ffi::GetRef<Var>(buffer_node);
PrimExpr scalar_extent = offset + IntImm(offset.ty(), 1);
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index a4023ec519..eb61cdb76d 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -632,11 +632,10 @@ class StoragePlanRewriter : public StmtExprMutator {
.as_or_throw<Expr>());
return Call(PrimType::Void(), op->op, args, op->attrs, op->ty_args,
op->span);
}
- } else if (op->op.same_as(tirx::tvm_access_ptr_op())) {
- TVM_FFI_ICHECK_EQ(op->args.size(), 5U);
- PrimExpr dtype_marker = op->args[0].as_or_throw<PrimExpr>();
- PrimType dtype = dtype_marker.ty();
- auto buffer_var = GetBufferDataVar(op->args[1]);
+ } else if (op->op.same_as(tvm_access_ptr_op())) {
+ TVM_FFI_ICHECK_EQ(op->args.size(), 4U);
+ PrimType dtype = op->ty_args[0].as_or_throw<PrimType>();
+ auto buffer_var = GetBufferDataVar(op->args[0]);
if (!buffer_var.has_value()) {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
@@ -651,18 +650,17 @@ class StoragePlanRewriter : public StmtExprMutator {
}
const StorageEntry* se = it->second;
PrimExpr offset =
-
this->Mutate(op->args[2]).ValueOrUnchanged(op->args[2]).as_or_throw<PrimExpr>();
+
this->Mutate(op->args[1]).ValueOrUnchanged(op->args[1]).as_or_throw<PrimExpr>();
PrimExpr extent =
-
this->Mutate(op->args[3]).ValueOrUnchanged(op->args[3]).as_or_throw<PrimExpr>();
+
this->Mutate(op->args[2]).ValueOrUnchanged(op->args[2]).as_or_throw<PrimExpr>();
uint64_t elem_bits = dtype.bits() * dtype.lanes();
TVM_FFI_ICHECK_EQ(se->bits_offset % elem_bits, 0U);
if (se->bits_offset != 0) {
offset = MakeConst(offset.ty(), se->bits_offset / elem_bits) + offset;
}
- return Call(
- op->ty, op->op,
- {dtype_marker, se->alloc_var, offset, extent,
op->args[4].as_or_throw<PrimExpr>()},
- op->attrs, {}, op->span);
+ return Call(op->ty, op->op,
+ {se->alloc_var, offset, extent,
op->args[3].as_or_throw<PrimExpr>()}, op->attrs,
+ op->ty_args, op->span);
} else {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
@@ -1481,11 +1479,11 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
indices.push_back(op->args[i].as_or_throw<PrimExpr>());
}
OnArrayAccess(dtype, buffer.get(), indices, is_load);
- } else if (op->op.same_as(tirx::tvm_access_ptr_op())) {
- PrimType dtype = op->args[0].as_or_throw<PrimExpr>().ty();
- auto buffer_var = GetBufferDataVar(op->args[1]);
- PrimExpr index = op->args[2].as_or_throw<PrimExpr>();
- // args[1] may be a nested Call (e.g. another tvm_access_ptr) rather
+ } else if (op->op.same_as(tvm_access_ptr_op())) {
+ PrimType dtype = op->ty_args[0].as_or_throw<PrimType>();
+ auto buffer_var = GetBufferDataVar(op->args[0]);
+ PrimExpr index = op->args[1].as_or_throw<PrimExpr>();
+ // args[0] may be a nested Call (e.g. another tvm_access_ptr) rather
// than a raw Var; OnArrayAccess derefs `buffer` so skip the record
// here and let the recursive visit handle any inner buffer var.
if (buffer_var.has_value()) {
@@ -2066,8 +2064,8 @@ class VectorTypeRewriter : public StmtExprMutator {
return RemapBuffer(var.value().as_or_throw<TensorVar>()).data();
}
}
- if (op->op.same_as(tirx::tvm_access_ptr_op())) {
- auto buffer = GetBufferDataVar(op->args[1]);
+ if (op->op.same_as(tvm_access_ptr_op())) {
+ auto buffer = GetBufferDataVar(op->args[0]);
Expr expr =
StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Expr>(op));
op = expr.as<CallNode>();
@@ -2087,21 +2085,19 @@ class VectorTypeRewriter : public StmtExprMutator {
}
const auto& info = it->second;
- PrimExpr index = op->args[2].as_or_throw<PrimExpr>();
- PrimExpr extent = op->args[3].as_or_throw<PrimExpr>();
- PrimExpr flag = op->args[4].as_or_throw<PrimExpr>();
+ PrimExpr index = op->args[1].as_or_throw<PrimExpr>();
+ PrimExpr extent = op->args[2].as_or_throw<PrimExpr>();
+ PrimExpr flag = op->args[3].as_or_throw<PrimExpr>();
- PrimExpr e_dtype = tirx::TypeAnnotation(info.new_element_dtype);
int factor = info.factor();
extent = extent / MakeConst(extent.ty(), factor);
index = index / MakeConst(index.ty(), factor);
Expr data = info.new_buffer_var->ty.as<TensorTypeNode>()
? info.new_buffer_var.as_or_throw<TensorVar>().data()
: Expr(info.new_buffer_var);
- ffi::Array<Expr> acc_args{e_dtype, data, index, extent, flag};
- auto old_pointer_type = op->ty.as_or_throw<PointerType>();
- Type new_pointer_type = PointerType(info.new_element_dtype,
old_pointer_type->storage_scope);
- return Call(new_pointer_type, tirx::tvm_access_ptr_op(), acc_args);
+ ffi::Array<Expr> acc_args{data, index, extent, flag};
+ return Call(op->ty, tvm_access_ptr_op(), acc_args, op->attrs,
{info.new_element_dtype},
+ op->span);
} else {
return StmtExprMutator::Mutate_(op, inplace_mode);
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index 933c6a2b29..9d716ad571 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -798,7 +798,7 @@ class Vectorizer : public StmtExprMutator {
if (op->args.same_as(new_args)) {
return ffi::Unchanged();
}
- return Call(op->ty, op->op, new_args, op->attrs, {}, op->span);
+ return Call(op->ty, op->op, new_args, op->attrs, op->ty_args, op->span);
}
PrimType ret_ty = optional_ret_ty.value();
if (op->op.same_as(prim::if_then_else_op())) {
diff --git a/tests/python/codegen/test_target_codegen_riscv.py
b/tests/python/codegen/test_target_codegen_riscv.py
index a3d97af717..5b1e83c17a 100644
--- a/tests/python/codegen/test_target_codegen_riscv.py
+++ b/tests/python/codegen/test_target_codegen_riscv.py
@@ -112,8 +112,8 @@ def test_rvv_vscale_llvm_dbginfo(target):
def rvv_with_vscale(A: T.Tensor((8,), dtype='float32', align=4,
offset_factor=1), B: T.Tensor((4, 8), dtype='float32', align=4,
offset_factor=1, strides=[8, 1]), C: T.Tensor((4,), dtype='float32', align=4,
offset_factor=1)):
zero = T.call_llvm_intrin( 'llvm.riscv.vfmv.v.f',
T.Broadcast(T.float32(0.0), T.vscale() * 2), C[0], T.uint64(1),
ty='float32xvscalex2')
- vec_A = T.call_llvm_intrin( 'llvm.riscv.vle',
T.Broadcast(T.float32(0.0), T.vscale() * 4),
T.tvm_access_ptr(T.type_annotation(ty='float32'), A.data, 0, 8, 1), T.int64(8),
ty='float32xvscalex4')
- vec_B = T.call_llvm_intrin( 'llvm.riscv.vle',
T.Broadcast(T.float32(0.0), T.vscale() * 4),
T.tvm_access_ptr(T.type_annotation(ty='float32'), B.data, 0 * 8, 8, 1),
T.int64(8), ty='float32xvscalex4')
+ vec_A = T.call_llvm_intrin( 'llvm.riscv.vle',
T.Broadcast(T.float32(0.0), T.vscale() * 4), T.tvm_access_ptr('float32',
A.data, 0, 8, 1), T.int64(8), ty='float32xvscalex4')
+ vec_B = T.call_llvm_intrin( 'llvm.riscv.vle',
T.Broadcast(T.float32(0.0), T.vscale() * 4), T.tvm_access_ptr('float32',
B.data, 0 * 8, 8, 1), T.int64(8), ty='float32xvscalex4')
prod = T.call_llvm_intrin( 'llvm.riscv.vfmul',
T.Broadcast(T.float32(0.0), T.vscale() * 4), vec_A, vec_B, T.uint64(7),
T.uint64(8), ty='float32xvscalex4')
redsum = T.call_llvm_intrin( 'llvm.riscv.vfredusum',
T.Broadcast(T.float32(0.0), T.vscale() * 2), prod, zero, T.uint64(7),
T.uint64(8), ty='float32xvscalex2')
# fmt: on
diff --git a/tests/python/codegen/test_target_codegen_vulkan.py
b/tests/python/codegen/test_target_codegen_vulkan.py
index d4e28379fa..d59021231c 100644
--- a/tests/python/codegen/test_target_codegen_vulkan.py
+++ b/tests/python/codegen/test_target_codegen_vulkan.py
@@ -481,14 +481,14 @@ def test_cooperative_matrix(out_dtype):
W_shared[k_0 * 16 + index // 16, index % 16] =
W[k_0 * 16 + index // 16, index % 16]
T.tvm_storage_sync("shared")
T.tvm_load_matrix_sync(X_shared_wmma_matrix_a.data, 16,
16, 16, k_0,
- T.tvm_access_ptr(T.type_annotation(ty="float16"),
X_shared.data, k_0 * 16, 512, 1), 32, "row_major")
+ T.tvm_access_ptr("float16", X_shared.data, k_0 * 16,
512, 1), 32, "row_major")
T.tvm_load_matrix_sync(W_shared_wmma_matrix_b.data, 16,
16, 16, k_0,
- T.tvm_access_ptr(T.type_annotation(ty="float16"),
W_shared.data, k_0 * 256, 256, 1), 16, "row_major")
+ T.tvm_access_ptr("float16", W_shared.data, k_0 * 256,
256, 1), 16, "row_major")
T.tvm_mma_sync(compute_wmma_accumulator.data, 0,
X_shared_wmma_matrix_a.data, k_0,
W_shared_wmma_matrix_b.data, k_0,
compute_wmma_accumulator.data, 0)
T.tvm_store_matrix_sync(compute_wmma_accumulator.data, 16, 16,
16, 0,
- T.tvm_access_ptr(T.type_annotation(ty=out_dtype),
compute.data, 0, 256, 2), 16, "row_major")
+ T.tvm_access_ptr(out_dtype, compute.data, 0, 256, 2), 16,
"row_major")
# fmt: on
target = {"kind": "vulkan", "from_device": 0}
diff --git a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
index cfb379cfcf..dd1c6b15c0 100644
--- a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
@@ -125,7 +125,7 @@ def test_matmul_tensorize():
Ts.writes(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16,
v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o *
16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_s0, A_s1),
scope="shared.dyn", offset_factor=16)
C =
Ts.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 +
16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_1_s0, C_1_s1),
scope="wmma.matrix_a", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="float16"), A.data, A.elem_offset,
A.strides[0] * 16, 1), A.strides[0], "row_major")
+ T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr("float16", A.data,
A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
for ax0_0 in T.unroll(2):
for ax1_0 in T.unroll(1):
with
Ts.sblock("W_reindex_shared.dyn_wmma.matrix_b_o"):
@@ -136,7 +136,7 @@ def test_matmul_tensorize():
Ts.writes(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16,
v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o *
16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_1_s0, A_1_s1),
scope="shared.dyn", offset_factor=16)
C =
Ts.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 +
16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_2_s0, C_2_s1),
scope="wmma.matrix_b", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="float16"), A.data, A.elem_offset,
A.strides[0] * 16, 1), A.strides[0], "col_major")
+ T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr("float16", A.data,
A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
for ax1_0_3, ax2_0_3 in T.grid(2, 2):
with Ts.sblock("compute_o_update"):
v0_o = Ts.axis.spatial(1, ax0)
@@ -164,7 +164,7 @@ def test_matmul_tensorize():
Ts.writes(compute_reindex_shared_dyn[v0_o,
v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o *
16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16",
strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
C =
Ts.match_buffer(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o
* 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_4_s0, C_4_s1),
scope="shared.dyn", offset_factor=16)
- T.tvm_store_matrix_sync(A.data, 16, 16, 16,
A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset %
A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(ty="float16"), C.data,
C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
+ T.tvm_store_matrix_sync(A.data, 16, 16, 16,
A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset %
A.strides[0] // 16, T.tvm_access_ptr("float16", C.data, C.elem_offset,
C.strides[0] * 16, 2), C.strides[0], "row_major")
for ax0_ax1_fused_0 in range(8):
for ax0_ax1_fused_1 in T.thread_binding(32,
thread="threadIdx.x"):
for ax0_ax1_fused_2 in T.vectorized(4):
@@ -407,7 +407,7 @@ def test_matmul_tensorize_epilogue():
Ts.writes(lv42_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 +
16, v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(lv42_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16,
v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_s0, A_s1),
scope="shared.dyn", offset_factor=16)
C =
Ts.match_buffer(lv42_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o
* 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_1_s0,
C_1_s1), scope="wmma.matrix_a", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="float16"), A.data, A.elem_offset,
A.strides[0] * 16, 1), A.strides[0], "row_major")
+ T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr("float16", A.data,
A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
for ax0_0 in T.unroll(2):
for ax1_0 in T.unroll(1):
with
Ts.sblock("p_output0_intermediate_1_reindex_shared.dyn_wmma.matrix_b_o"):
@@ -418,7 +418,7 @@ def test_matmul_tensorize_epilogue():
Ts.writes(p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o
* 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(p_output0_intermediate_1_reindex_shared_dyn[v0_o, v1_o *
16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16",
strides=(A_1_s0, A_1_s1), scope="shared.dyn", offset_factor=16)
C =
Ts.match_buffer(p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[v0_o,
v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16",
strides=(C_2_s0, C_2_s1), scope="wmma.matrix_b", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="float16"), A.data, A.elem_offset,
A.strides[0] * 16, 1), A.strides[0], "col_major")
+ T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr("float16", A.data,
A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
for ax1_0_3, ax2_0_3 in T.grid(2, 2):
with Ts.sblock("NT_matmul_o_update"):
v0_o = Ts.axis.spatial(1, ax0)
@@ -446,7 +446,7 @@ def test_matmul_tensorize_epilogue():
Ts.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0_o, v1_o *
16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[v0_o,
v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16",
strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
C =
Ts.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0_o, v1_o *
16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16",
strides=(C_4_s0, C_4_s1), scope="shared.dyn", offset_factor=16)
- T.tvm_store_matrix_sync(A.data, 16, 16, 16,
A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset %
A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(ty="float16"), C.data,
C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
+ T.tvm_store_matrix_sync(A.data, 16, 16, 16,
A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset %
A.strides[0] // 16, T.tvm_access_ptr("float16", C.data, C.elem_offset,
C.strides[0] * 16, 2), C.strides[0], "row_major")
for ax0_ax1_fused_0 in range(8):
for ax0_ax1_fused_1 in T.thread_binding(32,
thread="threadIdx.x"):
for ax0_ax1_fused_2 in T.vectorized(4):
@@ -567,7 +567,7 @@ def test_matmul_int8_tensorize():
Ts.writes(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16,
v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o *
16:v2_o * 16 + 16], (16, 16), "int8", strides=(A_s0, A_s1), scope="shared.dyn",
offset_factor=16)
C =
Ts.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 +
16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(C_1_s0, C_1_s1),
scope="wmma.matrix_a", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="int8"), A.data, A.elem_offset,
A.strides[0] * 16, 1), A.strides[0], "row_major")
+ T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr("int8", A.data,
A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
for ax0_0 in T.unroll(2):
for ax1_0 in T.unroll(1):
with
Ts.sblock("W_reindex_shared.dyn_wmma.matrix_b_o"):
@@ -578,7 +578,7 @@ def test_matmul_int8_tensorize():
Ts.writes(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16,
v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o *
16:v2_o * 16 + 16], (16, 16), "int8", strides=(A_1_s0, A_1_s1),
scope="shared.dyn", offset_factor=16)
C =
Ts.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 +
16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(C_2_s0, C_2_s1),
scope="wmma.matrix_b", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="int8"), A.data, A.elem_offset,
A.strides[0] * 16, 1), A.strides[0], "col_major")
+ T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr("int8", A.data,
A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
for ax1_0_3, ax2_0_3 in T.grid(2, 2):
with Ts.sblock("compute_o_update"):
v0_o = Ts.axis.spatial(1, ax0)
@@ -606,7 +606,7 @@ def test_matmul_int8_tensorize():
Ts.writes(compute_reindex_shared_dyn[v0_o,
v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
A =
Ts.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o *
16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32",
strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
C =
Ts.match_buffer(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o
* 16:v2_o * 16 + 16], (16, 16), "int32", strides=(C_4_s0, C_4_s1),
scope="shared.dyn", offset_factor=16)
- T.tvm_store_matrix_sync(A.data, 16, 16, 16,
A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset %
A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(ty="int32"), C.data,
C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
+ T.tvm_store_matrix_sync(A.data, 16, 16, 16,
A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset %
A.strides[0] // 16, T.tvm_access_ptr("int32", C.data, C.elem_offset,
C.strides[0] * 16, 2), C.strides[0], "row_major")
for ax0_ax1_fused_0 in range(8):
for ax0_ax1_fused_1 in T.thread_binding(32,
thread="threadIdx.x"):
for ax0_ax1_fused_2 in T.vectorized(4):
@@ -731,7 +731,7 @@ def test_matmul_int8_tensorize_3d2d_dyn():
Ts.writes(A_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 +
16, v2_o * 16:v2_o * 16 + 16])
A_1 =
Ts.match_buffer(A_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o *
16:v2_o * 16 + 16], (16, 16), "int8", strides=(A_s0, A_s1), scope="shared.dyn",
offset_factor=16)
C =
Ts.match_buffer(A_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o *
16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(C_1_s0, C_1_s1),
scope="wmma.matrix_a", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="int8"), A_1.data, A_1.elem_offset,
A_1.strides[0] * 16, 1), A_1.strides[0], "row_major")
+ T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr("int8", A_1.data,
A_1.elem_offset, A_1.strides[0] * 16, 1), A_1.strides[0], "row_major")
for ax0_0 in T.unroll(2):
for ax1_0 in T.unroll(1):
with
Ts.sblock("B_reindex_shared.dyn_wmma.matrix_b_o"):
@@ -742,7 +742,7 @@ def test_matmul_int8_tensorize_3d2d_dyn():
Ts.writes(B_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16,
v2_o * 16:v2_o * 16 + 16])
A_1 =
Ts.match_buffer(B_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o *
16:v2_o * 16 + 16], (16, 16), "int8", strides=(A_1_s0, A_1_s1),
scope="shared.dyn", offset_factor=16)
C =
Ts.match_buffer(B_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 +
16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(C_2_s0, C_2_s1),
scope="wmma.matrix_b", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="int8"), A_1.data, A_1.elem_offset,
A_1.strides[0] * 16, 1), A_1.strides[0], "col_major")
+ T.tvm_load_matrix_sync(C.data, 16,
16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) +
C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr("int8", A_1.data,
A_1.elem_offset, A_1.strides[0] * 16, 1), A_1.strides[0], "col_major")
for ax1_0_3, ax2_0_3 in T.grid(2, 2):
with Ts.sblock("matmul_o_update"):
v0_o = Ts.axis.spatial(1, ax0)
@@ -770,7 +770,7 @@ def test_matmul_int8_tensorize_3d2d_dyn():
Ts.writes(matmul_1_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o
* 16:v2_o * 16 + 16])
A_1 =
Ts.match_buffer(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o *
16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32",
strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
C =
Ts.match_buffer(matmul_1_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16,
v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(C_4_s0, C_4_s1),
scope="shared.dyn", offset_factor=16)
- T.tvm_store_matrix_sync(A_1.data, 16, 16, 16,
A_1.elem_offset // A_1.strides[0] // 16 * (A_1.strides[0] // 16) +
A_1.elem_offset % A_1.strides[0] // 16,
T.tvm_access_ptr(T.type_annotation(ty="int32"), C.data, C.elem_offset,
C.strides[0] * 16, 2), C.strides[0], "row_major")
+ T.tvm_store_matrix_sync(A_1.data, 16, 16, 16,
A_1.elem_offset // A_1.strides[0] // 16 * (A_1.strides[0] // 16) +
A_1.elem_offset % A_1.strides[0] // 16, T.tvm_access_ptr("int32", C.data,
C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
for ax0_ax1_fused_0 in range(8):
for ax0_ax1_fused_1 in T.thread_binding(32,
thread="threadIdx.x"):
for ax0_ax1_fused_2 in T.vectorized(4):
@@ -892,7 +892,7 @@ def test_matmul_metal():
Ts.writes(A_reindex_pad_shared_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8,
v2_o * 8:v2_o * 8 + 8])
A_1 =
Ts.match_buffer(A_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o
* 8 + 8], (8, 8), "float16", strides=(A_1_s0, A_1_s1), scope="shared",
offset_factor=1)
C_1 =
Ts.match_buffer(A_reindex_pad_shared_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 +
8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=(C_s0, C_s1),
scope="metal.simdgroup", offset_factor=1)
- T.metal.simdgroup_load(C_1.data,
C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) +
C_1.elem_offset % C_1.strides[0] // 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), A_1.data, A_1.elem_offset,
A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(False))
+ T.metal.simdgroup_load(C_1.data,
C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) +
C_1.elem_offset % C_1.strides[0] // 8, T.tvm_access_ptr("float16", A_1.data,
A_1.elem_offset, A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(False))
for ax0_0, ax1_0_1 in T.grid(2, 1):
with
Ts.sblock("B_reindex_shared_metal.simdgroup_o"):
v0_o = Ts.axis.spatial(1, 0)
@@ -902,7 +902,7 @@ def test_matmul_metal():
Ts.writes(B_reindex_shared_metal_simdgroup[v0_o, v2_o * 8:v2_o * 8 + 8, v1_o *
8:v1_o * 8 + 8])
A_1 =
Ts.match_buffer(B_reindex_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8
+ 8], (8, 8), "float16", strides=(A_2_s0, A_2_s1), scope="shared",
offset_factor=1)
C_1 =
Ts.match_buffer(B_reindex_shared_metal_simdgroup[v0_o, v2_o * 8:v2_o * 8 + 8,
v1_o * 8:v1_o * 8 + 8], (8, 8), "float16", strides=(C_1_s0, C_1_s1),
scope="metal.simdgroup", offset_factor=1)
- T.metal.simdgroup_load(C_1.data,
C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) +
C_1.elem_offset % C_1.strides[0] // 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), A_1.data, A_1.elem_offset,
A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(True))
+ T.metal.simdgroup_load(C_1.data,
C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) +
C_1.elem_offset % C_1.strides[0] // 8, T.tvm_access_ptr("float16", A_1.data,
A_1.elem_offset, A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(True))
for ax1_2, ax2_2 in T.grid(2, 2):
with Ts.sblock("C_update_o"):
v0_o = Ts.axis.spatial(1, ax0)
@@ -924,7 +924,7 @@ def test_matmul_metal():
Ts.writes(C_reindex_pad_shared[v0_o, v1_o
* 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
A_1 =
Ts.match_buffer(C_reindex_pad_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o
* 8:v2_o * 8 + 8], (8, 8), "float16", strides=(A_4_s0, A_4_s1),
scope="metal.simdgroup", offset_factor=1)
C_1 =
Ts.match_buffer(C_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o
* 8 + 8], (8, 8), "float16", strides=(C_3_s0, C_3_s1), scope="shared",
offset_factor=1)
- T.metal.simdgroup_store(A_1.data,
A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) +
A_1.elem_offset % A_1.strides[0] // 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), C_1.data, C_1.elem_offset,
C_1.strides[0] * 8, 2), C_1.strides[0], 8, 8, T.bool(False))
+ T.metal.simdgroup_store(A_1.data,
A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) +
A_1.elem_offset % A_1.strides[0] // 8, T.tvm_access_ptr("float16", C_1.data,
C_1.elem_offset, C_1.strides[0] * 8, 2), C_1.strides[0], 8, 8, T.bool(False))
for ax0_1, ax1_ax2_fused_0 in T.grid(1, 2):
for ax1_ax2_fused_1 in T.thread_binding(4,
thread="threadIdx.z"):
for ax1_ax2_fused_2 in T.thread_binding(1,
thread="threadIdx.y"):
@@ -1057,7 +1057,7 @@ def test_matmul_metal_int4_quant():
Ts.writes(A_reindex_pad_shared_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8,
v2_o * 8:v2_o * 8 + 8])
A_1 =
Ts.match_buffer(A_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o
* 8 + 8], (8, 8), "float16", strides=(A_1_s0, A_1_s1), scope="shared",
offset_factor=1)
C_1 =
Ts.match_buffer(A_reindex_pad_shared_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 +
8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=(C_s0, C_s1),
scope="metal.simdgroup", offset_factor=1)
- T.metal.simdgroup_load(C_1.data,
C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) +
C_1.elem_offset % C_1.strides[0] // 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), A_1.data, A_1.elem_offset,
A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(False))
+ T.metal.simdgroup_load(C_1.data,
C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) +
C_1.elem_offset % C_1.strides[0] // 8, T.tvm_access_ptr("float16", A_1.data,
A_1.elem_offset, A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(False))
for ax0_0, ax1_0_1 in T.grid(2, 1):
with
Ts.sblock("B_reindex_shared_metal.simdgroup_o"):
v0_o = Ts.axis.spatial(1, 0)
@@ -1067,7 +1067,7 @@ def test_matmul_metal_int4_quant():
Ts.writes(B_reindex_shared_metal_simdgroup[v0_o, v2_o * 8:v2_o * 8 + 8, v1_o *
8:v1_o * 8 + 8])
A_1 =
Ts.match_buffer(B_reindex_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8
+ 8], (8, 8), "float16", strides=(A_2_s0, A_2_s1), scope="shared",
offset_factor=1)
C_1 =
Ts.match_buffer(B_reindex_shared_metal_simdgroup[v0_o, v2_o * 8:v2_o * 8 + 8,
v1_o * 8:v1_o * 8 + 8], (8, 8), "float16", strides=(C_1_s0, C_1_s1),
scope="metal.simdgroup", offset_factor=1)
- T.metal.simdgroup_load(C_1.data,
C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) +
C_1.elem_offset % C_1.strides[0] // 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), A_1.data, A_1.elem_offset,
A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(True))
+ T.metal.simdgroup_load(C_1.data,
C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) +
C_1.elem_offset % C_1.strides[0] // 8, T.tvm_access_ptr("float16", A_1.data,
A_1.elem_offset, A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(True))
for ax1_2, ax2_2 in T.grid(2, 2):
with Ts.sblock("NT_matmul_update_o"):
v0_o = Ts.axis.spatial(1, ax0)
@@ -1089,7 +1089,7 @@ def test_matmul_metal_int4_quant():
Ts.writes(C_reindex_pad_shared[v0_o, v1_o
* 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
A_1 =
Ts.match_buffer(C_reindex_pad_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o
* 8:v2_o * 8 + 8], (8, 8), "float16", strides=(A_4_s0, A_4_s1),
scope="metal.simdgroup", offset_factor=1)
C_1 =
Ts.match_buffer(C_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o
* 8 + 8], (8, 8), "float16", strides=(C_3_s0, C_3_s1), scope="shared",
offset_factor=1)
- T.metal.simdgroup_store(A_1.data,
A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) +
A_1.elem_offset % A_1.strides[0] // 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), C_1.data, C_1.elem_offset,
C_1.strides[0] * 8, 2), C_1.strides[0], 8, 8, T.bool(False))
+ T.metal.simdgroup_store(A_1.data,
A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) +
A_1.elem_offset % A_1.strides[0] // 8, T.tvm_access_ptr("float16", C_1.data,
C_1.elem_offset, C_1.strides[0] * 8, 2), C_1.strides[0], 8, 8, T.bool(False))
for ax0_1, ax1_ax2_fused_0 in T.grid(1, 2):
for ax1_ax2_fused_1 in T.thread_binding(4,
thread="threadIdx.z"):
for ax1_ax2_fused_2 in T.thread_binding(1,
thread="threadIdx.y"):
diff --git
a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
index 045ed9ecfd..3c43c111fe 100644
---
a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
+++
b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
@@ -549,7 +549,7 @@ def GMMCUDATensorCore(
16,
C_1.elem_offset // 256 +
C_1.elem_offset % 256 // 16,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
A.data,
A.elem_offset,
s1 * 16,
@@ -602,7 +602,7 @@ def GMMCUDATensorCore(
16,
C_2.elem_offset // 256 +
C_2.elem_offset % 256 // 16,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
A_1.data,
A_1.elem_offset,
s1_1 * 16,
@@ -736,7 +736,7 @@ def GMMCUDATensorCore(
16,
A_3.elem_offset // 256 + A_3.elem_offset % 256
// 16,
T.tvm_access_ptr(
- T.type_annotation(ty="float32"),
+ "float32",
C_4.data,
C_4.elem_offset,
s1_2 * 16,
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
index 3928fc428e..6ac69eff7e 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
@@ -705,7 +705,7 @@ class Conv2dInt8_tensorcore_scheduled:
Ts.writes(pad_temp_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o
* 16:v1_o * 16 + 16])
A =
Ts.match_buffer(pad_temp_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o *
16:v1_o * 16 + 16], (16, 16), "int8", strides=(A_s0, A_s1), scope="shared",
offset_factor=16)
C =
Ts.match_buffer(pad_temp_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16,
v1_o * 16:v1_o * 16 + 16], (16, 16), "int8", strides=(C_1_s0, C_1_s1),
scope="wmma.matrix_a", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16, 16, 16,
C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset %
C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(ty="int8"), A.data,
A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
+ T.tvm_load_matrix_sync(C.data, 16, 16, 16,
C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset %
C.strides[0] // 16, T.tvm_access_ptr("int8", A.data, A.elem_offset,
A.strides[0] * 16, 1), A.strides[0], "row_major")
for ax0, ax1, ax2_0, ax3_0 in T.grid(1, 1, 1, 2):
with
Ts.sblock("p1_reindex_shared_wmma.matrix_b_o"):
v0_o, v1_o = Ts.axis.remap("SS", [ax0,
ax1])
@@ -715,7 +715,7 @@ class Conv2dInt8_tensorcore_scheduled:
Ts.writes(p1_reindex_shared_wmma_matrix_b[v0_o, v1_o, v2_o * 16:v2_o * 16 + 16,
v3_o * 16:v3_o * 16 + 16])
A =
Ts.match_buffer(p1_reindex_shared[v0_o, v1_o, v2_o * 16:v2_o * 16 + 16, v3_o *
16:v3_o * 16 + 16], (16, 16), "int8", strides=(A_1_s0, A_1_s1), scope="shared",
offset_factor=16)
C =
Ts.match_buffer(p1_reindex_shared_wmma_matrix_b[v0_o, v1_o, v2_o * 16:v2_o * 16
+ 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=(C_2_s0, C_2_s1),
scope="wmma.matrix_b", offset_factor=16)
- T.tvm_load_matrix_sync(C.data, 16, 16, 16,
C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset %
C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(ty="int8"), A.data,
A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
+ T.tvm_load_matrix_sync(C.data, 16, 16, 16,
C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset %
C.strides[0] // 16, T.tvm_access_ptr("int8", A.data, A.elem_offset,
A.strides[0] * 16, 1), A.strides[0], "col_major")
for ax2_0_3, ax3_0_3, ax0_2, ax1_2, ax4_0_2,
ax2_0_4, ax3_0_4 in T.grid(1, 1, 1, 1, 2, 1, 1):
with Ts.sblock("conv2d_nhwc_o_update"):
v0_o = Ts.axis.spatial(1, ax0_0 + ax0_1 +
ax0_2)
@@ -738,7 +738,7 @@ class Conv2dInt8_tensorcore_scheduled:
Ts.writes(conv2d_nhwc_reindex_shared[v0_o *
16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
A =
Ts.match_buffer(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o * 16:v0_o * 16
+ 16, v1_o * 16:v1_o * 16 + 16], (16, 16), "int32", strides=(A_3_s0, A_3_s1),
scope="wmma.accumulator", offset_factor=16)
C =
Ts.match_buffer(conv2d_nhwc_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o *
16:v1_o * 16 + 16], (16, 16), "int32", strides=(C_4_s0, C_4_s1),
scope="shared", offset_factor=16)
- T.tvm_store_matrix_sync(A.data, 16, 16, 16,
A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset %
A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(ty="int32"), C.data,
C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
+ T.tvm_store_matrix_sync(A.data, 16, 16, 16,
A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset %
A.strides[0] // 16, T.tvm_access_ptr("int32", C.data, C.elem_offset,
C.strides[0] * 16, 2), C.strides[0], "row_major")
for ax0, ax1_0 in T.grid(128, 2):
for ax1_1 in T.thread_binding(16, thread="threadIdx.x"):
with Ts.sblock("conv2d_nhwc_reindex_shared"):
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
b/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
index ba64a1a553..1a5608c4e9 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
@@ -275,7 +275,7 @@ def opaque_access(
16,
vi * 8 + vj,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
A.data,
vi * 2048 + vj * 16,
128,
@@ -318,7 +318,7 @@ def opaque_access(
16,
vi * 8 + vj,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
A0.data,
A0.elem_offset,
A0.strides[0],
@@ -592,7 +592,7 @@ def cache_read_opaque_access(
16,
vi * 8 + vj,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
A_global.data,
vi * 2048 + vj * 16,
128,
@@ -635,7 +635,7 @@ def cache_read_opaque_access(
16,
vi * 8 + vj,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
A0.data,
A0.elem_offset,
A0.strides[0],
@@ -953,7 +953,7 @@ def cache_write_opaque_access(
16,
vi * 8 + vj,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
A.data,
vi * 2048 + vj * 16,
128,
@@ -996,7 +996,7 @@ def cache_write_opaque_access(
16,
vi * 8 + vj,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
A0.data,
A0.elem_offset,
A0.strides[0],
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py
b/tests/python/s_tir/script/test_s_tir_script_printer.py
index 2c1b47d803..b78478d981 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -909,7 +909,7 @@ def opt_conv_tensorcore_lower():
16,
0,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
Apad_shared.data,
(((ty * 3072) + (kw * 512)) + (ic_inner *
256)),
256,
@@ -927,7 +927,7 @@ def opt_conv_tensorcore_lower():
16,
1,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
Apad_shared.data,
((((ty * 3072) + (kw * 512)) + (ic_inner *
256)) + 1536),
256,
@@ -945,7 +945,7 @@ def opt_conv_tensorcore_lower():
16,
0,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
W_shared.data,
(((kw * 4096) + (ic_inner * 2048)) + (tz *
1024)),
256,
@@ -963,7 +963,7 @@ def opt_conv_tensorcore_lower():
16,
3,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
W_shared.data,
((((kw * 4096) + (ic_inner * 2048)) + (tz
* 1024)) + 768),
256,
@@ -1005,7 +1005,7 @@ def opt_conv_tensorcore_lower():
16,
0,
T.tvm_access_ptr(
- T.type_annotation(ty="float32"),
+ "float32",
Conv_1.data,
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) +
(by * 2048))
@@ -1026,7 +1026,7 @@ def opt_conv_tensorcore_lower():
16,
7,
T.tvm_access_ptr(
- T.type_annotation(ty="float32"),
+ "float32",
Conv_1.data,
(
(
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
index 6341bd235c..dc80659206 100644
---
a/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
+++
b/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
@@ -68,7 +68,7 @@ def test_backward_compatibility_shared_a():
Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y //
2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
Ts.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
Ts.sblock_attr({"permuted_layout": "s2l_A"})
-
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), X_reindex_shared_dyn.data,
threadIdx_y // 2 * 2048 + ax0_0 * 1024 + ax2_0_1 * 8, 1024, 1), threadIdx_x *
32, ty="float16")
+
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8,
T.tvm_access_ptr("float16", X_reindex_shared_dyn.data, threadIdx_y // 2 * 2048
+ ax0_0 * 1024 + ax2_0_1 * 8, 1024, 1), threadIdx_x * 32, ty="float16")
@Ts.function
def expected(X: T.Tensor((4096, 4096), "float16")):
@@ -93,7 +93,7 @@ def test_backward_compatibility_shared_a():
with
Ts.sblock("X_reindex_shared.dyn_m16n8k8.matrixA_o"):
Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y //
2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
Ts.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
-
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), X_reindex_shared_dyn.data,
threadIdx_y // 2 * 2048 + ax0_0 * 1024 + threadIdx_x * 32 +
T.bitwise_xor(ax2_0_1, threadIdx_x % 8 // 2) * 8, 1024, 1), 0, ty="float16")
+
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8,
T.tvm_access_ptr("float16", X_reindex_shared_dyn.data, threadIdx_y // 2 * 2048
+ ax0_0 * 1024 + threadIdx_x * 32 + T.bitwise_xor(ax2_0_1, threadIdx_x % 8 //
2) * 8, 1024, 1), 0, ty="float16")
# fmt: on
_check_function_transform(before, expected)
@@ -130,13 +130,13 @@ def test_backward_compatibility_shared_a_and_b():
Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y //
2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
Ts.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
Ts.sblock_attr({"permuted_layout": "s2l_A"})
-
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), X_reindex_shared_dyn.data,
threadIdx_y // 2 * 2048 + ax0_0 * 1024 + ax2_0_1 * 8, 1024, 1), threadIdx_x *
32, ty="float16")
+
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8,
T.tvm_access_ptr("float16", X_reindex_shared_dyn.data, threadIdx_y // 2 * 2048
+ ax0_0 * 1024 + ax2_0_1 * 8, 1024, 1), threadIdx_x * 32, ty="float16")
for ax0_0, ax1_0 in T.grid(1, 2):
with
Ts.sblock("Y_reindex_shared.dyn_m16n8k8.matrixB_o"):
Ts.reads(Y_reindex_shared_dyn[ax2_0_1 * 8:ax2_0_1 * 8 + 8, threadIdx_y % 2 * 64
+ ax1_0 * 32:threadIdx_y % 2 * 64 + ax1_0 * 32 + 32])
Ts.writes(Y_reindex_shared_dyn_m16n8k8_matrixB[0:8, ax1_0 * 32:ax1_0 * 32 + 32])
Ts.sblock_attr({"permuted_layout": "s2l_B"})
-
T.ptx_legacy.ldmatrix(T.bool(True), 4, ".b16",
Y_reindex_shared_dyn_m16n8k8_matrixB.data, ax1_0 * 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), Y_reindex_shared_dyn.data,
ax2_0_1 * 1024 + threadIdx_y % 2 * 64 + ax1_0 * 32, 1024, 1), threadIdx_x % 8 *
128 + threadIdx_x // 8 * 8, ty="float16")
+
T.ptx_legacy.ldmatrix(T.bool(True), 4, ".b16",
Y_reindex_shared_dyn_m16n8k8_matrixB.data, ax1_0 * 8,
T.tvm_access_ptr("float16", Y_reindex_shared_dyn.data, ax2_0_1 * 1024 +
threadIdx_y % 2 * 64 + ax1_0 * 32, 1024, 1), threadIdx_x % 8 * 128 +
threadIdx_x // 8 * 8, ty="float16")
@Ts.function
def expected(X: T.Tensor((4096, 4096), "float16"), Y: T.Tensor((4096,
4096), "float16")):
@@ -173,12 +173,12 @@ def test_backward_compatibility_shared_a_and_b():
with
Ts.sblock("X_reindex_shared.dyn_m16n8k8.matrixA_o"):
Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y //
2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
Ts.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
-
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), X_reindex_shared_dyn.data,
threadIdx_y // 2 * 2048 + ax0_0 * 1024 + threadIdx_x * 32 +
T.bitwise_xor(ax2_0_1, threadIdx_x % 8 // 2) * 8, 1024, 1), 0, ty="float16")
+
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8,
T.tvm_access_ptr("float16", X_reindex_shared_dyn.data, threadIdx_y // 2 * 2048
+ ax0_0 * 1024 + threadIdx_x * 32 + T.bitwise_xor(ax2_0_1, threadIdx_x % 8 //
2) * 8, 1024, 1), 0, ty="float16")
for ax0_0, ax1_0 in T.grid(1, 2):
with
Ts.sblock("Y_reindex_shared.dyn_m16n8k8.matrixB_o"):
Ts.reads(Y_reindex_shared_dyn[ax2_0_1 * 8:ax2_0_1 * 8 + 8, threadIdx_y % 2 * 64
+ ax1_0 * 32:threadIdx_y % 2 * 64 + ax1_0 * 32 + 32])
Ts.writes(Y_reindex_shared_dyn_m16n8k8_matrixB[0:8, ax1_0 * 32:ax1_0 * 32 + 32])
-
T.ptx_legacy.ldmatrix(T.bool(True), 4, ".b16",
Y_reindex_shared_dyn_m16n8k8_matrixB.data, ax1_0 * 8,
T.tvm_access_ptr(T.type_annotation(ty="float16"), Y_reindex_shared_dyn.data,
ax2_0_1 * 1024 + threadIdx_x % 8 * 128 + T.bitwise_xor(threadIdx_y % 2 * 8 +
ax1_0 * 4 + threadIdx_x // 8, threadIdx_x % 8) * 8, 1024, 1), 0, ty="float16")
+
T.ptx_legacy.ldmatrix(T.bool(True), 4, ".b16",
Y_reindex_shared_dyn_m16n8k8_matrixB.data, ax1_0 * 8,
T.tvm_access_ptr("float16", Y_reindex_shared_dyn.data, ax2_0_1 * 1024 +
threadIdx_x % 8 * 128 + T.bitwise_xor(threadIdx_y % 2 * 8 + ax1_0 * 4 +
threadIdx_x // 8, threadIdx_x % 8) * 8, 1024, 1), 0, ty="float16")
# fmt: on
_check_function_transform(before, expected)
@@ -213,7 +213,7 @@ def test_buffer_a():
T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
A_warp.data,
v1 * T.int64(256) + threadIdx_x *
T.int64(8),
-
T.tvm_access_ptr(T.type_annotation(ty="float16"),
+ T.tvm_access_ptr("float16",
A_shared_dyn.data,
threadIdx_z * T.int64(2048) + v1 *
T.int64(512) + v0 * T.int64(16), T.int64(512),
1
@@ -241,7 +241,7 @@ def test_buffer_a():
with Ts.sblock("A_reindex_shared.dyn_warp_o"):
Ts.reads(A_shared_dyn[threadIdx_z *
T.int64(64) + v1 * T.int64(16):threadIdx_z * T.int64(64) + v1 * T.int64(16) +
T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
Ts.writes(A_warp[v1, T.int64(0),
T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- T.ptx_legacy.ldmatrix(T.bool(False), 4,
".b16", A_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8),
T.tvm_access_ptr(T.type_annotation(ty="float16"), A_shared_dyn.data,
threadIdx_z * T.int64(2048) + v1 * T.int64(512) + threadIdx_x % T.int64(16) *
T.int64(32) + T.bitwise_xor(v0 * T.int64(2) + threadIdx_x // T.int64(16),
threadIdx_x % T.int64(8) // T.int64(2)) * T.int64(8), T.int64(512), 1),
T.int64(0), ty="float16")
+ T.ptx_legacy.ldmatrix(T.bool(False), 4,
".b16", A_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8),
T.tvm_access_ptr("float16", A_shared_dyn.data, threadIdx_z * T.int64(2048) + v1
* T.int64(512) + threadIdx_x % T.int64(16) * T.int64(32) + T.bitwise_xor(v0 *
T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(8) //
T.int64(2)) * T.int64(8), T.int64(512), 1), T.int64(0), ty="float16")
# fmt: on
_check_function_transform(before, expected)
@@ -269,7 +269,7 @@ def test_buffer_b():
with
Ts.sblock("B_reindex_shared.dyn_warp_o"):
Ts.reads(B_shared_dyn[threadIdx_y *
T.int64(64) + v1 * T.int64(16):threadIdx_y * T.int64(64) + v1 * T.int64(16) +
T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
Ts.writes(B_warp[v1, T.int64(0),
T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- T.ptx_legacy.ldmatrix(T.bool(False),
4, ".b16", B_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8),
T.tvm_access_ptr(T.type_annotation(ty="float16"), B_shared_dyn.data,
threadIdx_y * T.int64(2048) + v1 * T.int64(512) + v0 * T.int64(16),
T.int64(512), 1), threadIdx_x // T.int64(16) * T.int64(256) + threadIdx_x %
T.int64(8) * T.int64(32) + threadIdx_x % T.int64(16) // T.int64(8) *
T.int64(8), ty="float16")
+ T.ptx_legacy.ldmatrix(T.bool(False),
4, ".b16", B_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8),
T.tvm_access_ptr("float16", B_shared_dyn.data, threadIdx_y * T.int64(2048) + v1
* T.int64(512) + v0 * T.int64(16), T.int64(512), 1), threadIdx_x // T.int64(16)
* T.int64(256) + threadIdx_x % T.int64(8) * T.int64(32) + threadIdx_x %
T.int64(16) // T.int64(8) * T.int64(8), ty="float16")
@Ts.function
def expected(B: T.Tensor((T.int64(128), T.int64(32)), "float16")):
@@ -293,7 +293,7 @@ def test_buffer_b():
with
Ts.sblock("B_reindex_shared.dyn_warp_o"):
Ts.reads(B_shared_dyn[threadIdx_y *
T.int64(64) + v1 * T.int64(16):threadIdx_y * T.int64(64) + v1 * T.int64(16) +
T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
Ts.writes(B_warp[v1, T.int64(0),
T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- T.ptx_legacy.ldmatrix(T.bool(False),
4, ".b16", B_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8),
T.tvm_access_ptr(T.type_annotation(ty="float16"), B_shared_dyn.data,
threadIdx_y * T.int64(2048) + v1 * T.int64(512) + threadIdx_x // T.int64(16) *
T.int64(256) + threadIdx_x % T.int64(8) * T.int64(32) + T.bitwise_xor(v0 *
T.int64(2) + threadIdx_x % T.int64(16) // T.int64(8), threadIdx_x % T.int64(8)
// T.int64(2)) * T.int64(8), T.int64(512), [...]
+ T.ptx_legacy.ldmatrix(T.bool(False),
4, ".b16", B_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8),
T.tvm_access_ptr("float16", B_shared_dyn.data, threadIdx_y * T.int64(2048) + v1
* T.int64(512) + threadIdx_x // T.int64(16) * T.int64(256) + threadIdx_x %
T.int64(8) * T.int64(32) + T.bitwise_xor(v0 * T.int64(2) + threadIdx_x %
T.int64(16) // T.int64(8), threadIdx_x % T.int64(8) // T.int64(2)) *
T.int64(8), T.int64(512), 1), T.int64(0), ty="float16")
# fmt: on
_check_function_transform(before, expected)
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
index 07542b9c03..1853391e42 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
@@ -698,7 +698,7 @@ def
test_vectorize_cp_async_in_if_then_else(postproc_if_missing_async_support):
C.ty.elem_offset // C_s0_1 // 16 *
(C_s0_1 // 16)
+ C.ty.elem_offset % C_s0_1 // 16,
T.tvm_access_ptr(
-
T.type_annotation(ty="float16"),
+ "float16",
A_1.data,
A_1.ty.elem_offset,
A_s0_0 * 16,
@@ -755,7 +755,7 @@ def
test_vectorize_cp_async_in_if_then_else(postproc_if_missing_async_support):
C.ty.elem_offset // C_s0_2 // 16 *
(C_s0_2 // 16)
+ C.ty.elem_offset % C_s0_2 // 16,
T.tvm_access_ptr(
-
T.type_annotation(ty="float16"),
+ "float16",
A_1.data,
A_1.ty.elem_offset,
A_s0_1 * 16,
@@ -873,7 +873,7 @@ def
test_vectorize_cp_async_in_if_then_else(postproc_if_missing_async_support):
A_1.ty.elem_offset // A_s0_3 // 16 *
(A_s0_3 // 16)
+ A_1.ty.elem_offset % A_s0_3 // 16,
T.tvm_access_ptr(
- T.type_annotation(ty="float16"),
+ "float16",
C.data,
C.ty.elem_offset,
C_s0_4 * 16,
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
index 56b7d2311e..6d6b096084 100644
---
a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
+++
b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
@@ -481,7 +481,7 @@ class TransformedSharedToWmma:
tgt.elem_offset // 256
+ tgt.elem_offset % 256 // 16,
T.tvm_access_ptr(
-
T.type_annotation(ty="float16"),
+ "float16",
src.data,
src.elem_offset,
s1 * 16,
@@ -562,7 +562,7 @@ class TransformedWmmaToShared:
src.elem_offset // 256
+ src.elem_offset % 256 // 16,
T.tvm_access_ptr(
-
T.type_annotation(ty="float32"),
+ "float32",
tgt.data,
tgt.elem_offset,
s1 * 16,
@@ -632,7 +632,7 @@ class TransformedWmmaToGlobal:
src.elem_offset // 256
+ src.elem_offset % 256 // 16,
T.tvm_access_ptr(
-
T.type_annotation(ty="float32"),
+ "float32",
tgt.data,
tgt.elem_offset,
s1 * 16,
@@ -836,7 +836,7 @@ class TransformedWmmaToGlobalWithFusion:
src.elem_offset // 256
+ src.elem_offset % 256 // 16,
T.tvm_access_ptr(
-
T.type_annotation(ty="float32"),
+ "float32",
tgt.data,
tgt.elem_offset,
s1_1 * 16,
diff --git a/tests/python/tirx-base/test_tir_buffer.py
b/tests/python/tirx-base/test_tir_buffer.py
index f6e78c7464..b6b1562a72 100644
--- a/tests/python/tirx-base/test_tir_buffer.py
+++ b/tests/python/tirx-base/test_tir_buffer.py
@@ -114,9 +114,9 @@ def test_buffer_access_ptr():
aptr = Ab.access_ptr("rw")
assert isinstance(aptr.ty, tvm.ir.PointerType)
assert aptr.ty.element_type == tvm.ir.PrimType("void")
- tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m)
- assert aptr.args[0].ty == Ab.ty.dtype
- assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
+ tvm.ir.assert_structural_equal(aptr.args[2], Ab.ty.strides[0] * m)
+ assert aptr.ty_args[0] == Ab.ty.dtype
+ assert aptr.args[3].value == BufferAccessKind.READ | BufferAccessKind.WRITE
typed_ptr = Ab.access_ptr("r", ptr_type="uint8")
assert typed_ptr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"))
shared = tvm.tirx.decl_tensor((m, n), "float32", scope="shared")
@@ -125,7 +125,7 @@ def test_buffer_access_ptr():
tvm.ir.PrimType("uint8"), "shared"
)
aptr = Ab.access_ptr("w")
- assert aptr.args[4].value == BufferAccessKind.WRITE
+ assert aptr.args[3].value == BufferAccessKind.WRITE
def test_buffer_access_ptr_offset():
@@ -133,17 +133,17 @@ def test_buffer_access_ptr_offset():
n = tvm.tirx.Var("n", "int32")
Ab = tvm.tirx.decl_tensor((m, n), "float32")
aptr = Ab.access_ptr("rw", offset=100)
- tvm.testing.assert_prim_expr_equal(aptr.args[2], 100)
- assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
+ tvm.testing.assert_prim_expr_equal(aptr.args[1], 100)
+ assert aptr.args[3].value == BufferAccessKind.READ | BufferAccessKind.WRITE
v = tvm.tirx.Var("int32", "int32")
aptr = Ab.access_ptr("rw", offset=100 + 100 + v)
- tvm.testing.assert_prim_expr_equal(aptr.args[2], 200 + v)
- assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
+ tvm.testing.assert_prim_expr_equal(aptr.args[1], 200 + v)
+ assert aptr.args[3].value == BufferAccessKind.READ | BufferAccessKind.WRITE
aptr = Ab.access_ptr("rw", offset=tvm.tirx.call_extern("int32",
"test_call", 100 + 100 + v))
tvm.testing.assert_prim_expr_equal(
- aptr.args[2], tvm.tirx.call_extern("int32", "test_call", 200 + v)
+ aptr.args[1], tvm.tirx.call_extern("int32", "test_call", 200 + v)
)
- assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
+ assert aptr.args[3].value == BufferAccessKind.READ | BufferAccessKind.WRITE
def test_buffer_access_ptr_extent():
@@ -151,18 +151,18 @@ def test_buffer_access_ptr_extent():
n = tvm.tirx.Var("n", "int32")
Ab = tvm.tirx.decl_tensor((m, n), "float32")
aptr = Ab.access_ptr("rw")
- tvm.ir.assert_structural_equal(aptr.args[3], m * n)
+ tvm.ir.assert_structural_equal(aptr.args[2], m * n)
aptr = Ab.access_ptr("rw", offset=100)
- tvm.ir.assert_structural_equal(aptr.args[3], m * n - 100)
+ tvm.ir.assert_structural_equal(aptr.args[2], m * n - 100)
Ab = tvm.tirx.decl_tensor((m, n), "float32", strides=[n + 1, 1])
aptr = Ab.access_ptr("rw", offset=100)
- tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m - 100)
+ tvm.ir.assert_structural_equal(aptr.args[2], Ab.ty.strides[0] * m - 100)
# Test extent from input params
aptr = Ab.access_ptr("rw", extent=200)
- tvm.ir.assert_structural_equal(aptr.args[3], T.int32(200))
+ tvm.ir.assert_structural_equal(aptr.args[2], T.int32(200))
aptr = Ab.access_ptr("rw", offset=100, extent=100)
- tvm.ir.assert_structural_equal(aptr.args[3], T.int32(100))
+ tvm.ir.assert_structural_equal(aptr.args[2], T.int32(100))
def test_buffer_vload():
diff --git a/tests/python/tirx-base/test_tir_op_types.py
b/tests/python/tirx-base/test_tir_op_types.py
index 8a6648006e..734f751bdf 100644
--- a/tests/python/tirx-base/test_tir_op_types.py
+++ b/tests/python/tirx-base/test_tir_op_types.py
@@ -103,11 +103,6 @@ def test_tir_op_tvm_thread_allreduce():
assert expr.op.name == "tirx.tvm_thread_allreduce"
-def test_tir_op_type_annotation():
- expr = tirx.type_annotation("int32")
- assert expr.op.name == "tirx.type_annotation"
-
-
def test_tir_op_tvm_access_ptr():
buffer = tirx.decl_tensor((128), "float32")
for ptype in ("float32", tvm.ir.PrimType("float32")):
@@ -224,9 +219,9 @@ def test_op_ptx_cp_async():
# Raw-form layout: (dst, dst_off, src, src_off, cp_size), offsets folded.
for access_ptr, expected_offset in zip((expr.args[0], expr.args[2]), [5,
9]):
assert access_ptr.op.name == "tirx.tvm_access_ptr"
- assert access_ptr.args[1].op.name == "tirx.buffer_data"
- assert isinstance(access_ptr.args[1].args[0], tirx.Var)
- simplified_offset = tvm.sym.Analyzer().simplify(access_ptr.args[2])
+ assert access_ptr.args[0].op.name == "tirx.buffer_data"
+ assert isinstance(access_ptr.args[0].args[0], tirx.Var)
+ simplified_offset = tvm.sym.Analyzer().simplify(access_ptr.args[1])
assert int(simplified_offset) == expected_offset
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
index 2b451a768b..3f2e0c3272 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
@@ -142,9 +142,9 @@ def test_lower_vector_access_ptr():
access_ptr = buffer.access_ptr(access_mask=3, offset=2, extent=4)
assert access_ptr.op.name == "tirx.tvm_access_ptr"
- assert int(access_ptr.args[2]) == 2
- assert int(access_ptr.args[3]) == 4
- assert int(access_ptr.args[4]) == 3
+ assert int(access_ptr.args[1]) == 2
+ assert int(access_ptr.args[2]) == 4
+ assert int(access_ptr.args[3]) == 3
mod = tvm.IRModule.from_expr(
tvm.tirx.Function([buffer], tvm.tirx.Evaluate(access_ptr)).with_attr(
diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py
b/tests/python/tirx/transform/test_transform_lower_tirx.py
index 97f3a5267e..e0c60d951e 100644
--- a/tests/python/tirx/transform/test_transform_lower_tirx.py
+++ b/tests/python/tirx/transform/test_transform_lower_tirx.py
@@ -741,9 +741,7 @@ def test_lower_decl_buffer_access_ptr():
A = T.decl_tensor(
(128,), "float16", data=buf.data, elem_offset=32,
scope="shared.dyn", layout=None
)
- T.tvm_access_ptr(
- T.type_annotation(ty="float16"), buf.data, T.Add(32, 64),
T.Sub(128, 64), 3
- )
+ T.tvm_access_ptr("float16", buf.data, T.Add(32, 64), T.Sub(128, 64), 3)
compare(before, after, LowerTIRx)
