You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/11 18:12:14 UTC
(tvm) branch main updated: [REFACTOR][TIRx] Carry tensor allocation types in type arguments (#20659)
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 83c79cf9de [REFACTOR][TIRx] Carry tensor allocation types in type arguments (#20659)
83c79cf9de is described below
commit 83c79cf9de06d902601ca342ad1e41aaec443700
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Sun Oct 11 14:12:08 2026 -0400
[REFACTOR][TIRx] Carry tensor allocation types in type arguments (#20659)
Represent tensor allocation and declaration with one complete TensorType
argument, retaining placement and backing pointers as value operands.
Rebuild type arguments through storage and layout lowering and print
producers through ordinary calls and bindings.
---
docs/deep_dive/relax/learning.rst | 2 +-
docs/deep_dive/tensor_ir/learning.rst | 6 +-
docs/deep_dive/tensor_ir/tutorials/tir_creation.py | 8 +-
docs/tirx/native_basics/cuda/data_types.rst | 4 +-
docs/tirx/native_basics/cuda/parser_utils.rst | 4 +-
docs/tirx/native_basics/cuda/tensors.rst | 52 +--
docs/tirx/tile_primitives/copy/ldstmatrix.rst | 8 +-
docs/tirx/tile_primitives/copy/reg.rst | 8 +-
docs/tirx/tile_primitives/copy_async/dsmem.rst | 28 +-
docs/tirx/tile_primitives/copy_async/ldgsts.rst | 2 +-
.../tirx/tile_primitives/copy_async/tcgen05_cp.rst | 13 +-
docs/tirx/tile_primitives/elementwise/reg.rst | 6 +-
docs/tirx/tile_primitives/elementwise/smem.rst | 4 +-
docs/tirx/tile_primitives/gemm.rst | 8 +-
docs/tirx/tile_primitives/gemm_async.rst | 22 +-
include/tvm/s_tir/stmt.h | 2 +-
include/tvm/s_tir/transform.h | 4 +-
include/tvm/tirx/op/memory.h | 26 +-
include/tvm/tirx/script/ir_builder/ir.h | 52 ---
python/tvm/backend/cuda/lang/alloc_pool.py | 30 +-
python/tvm/backend/cuda/lang/pipeline.py | 2 +-
.../cuda/tile_primitive/copy_async/dsmem.py | 20 +-
.../cuda/tile_primitive/copy_async/tcgen05_cp.py | 10 +-
.../backend/cuda/tile_primitive/copy_async/tma.py | 15 +-
.../backend/cuda/tile_primitive/elementwise/reg.py | 18 +-
.../cuda/tile_primitive/gemm_async/tcgen05.py | 27 +-
.../backend/trn/tile_primitive/private_alloc.py | 8 +-
.../tvm/backend/trn/transform/naive_allocator.py | 21 +-
python/tvm/relax/backend/gpu_generic/cumsum.py | 2 +-
.../relax/frontend/nn/llm/position_embedding.py | 36 +-
.../tvm/relax/frontend/tflite/tflite_frontend.py | 37 ++-
.../tvm/relax/transform/legalize_ops/inspect_op.py | 6 +-
python/tvm/script/printer/scriptable.py | 4 -
python/tvm/tirx/_tensor_view.py | 21 +-
python/tvm/tirx/bench.py | 8 +-
python/tvm/tirx/script/ir_builder/ir.py | 289 ++++++----------
.../tvm/tirx/script/ir_builder/parser_protocol.py | 10 +-
python/tvm/topi/gpu/scan.py | 28 +-
python/tvm/topi/gpu/sort.py | 42 ++-
python/tvm/topi/searchsorted.py | 4 +-
python/tvm/topi/sparse_reshape.py | 36 +-
python/tvm/topi/vision/nms.py | 30 +-
python/tvm/topi/vision/nms_util.py | 4 +-
src/backend/cuda/codegen/codegen_cuda.cc | 6 +-
src/backend/cuda/codegen/llvm/codegen_nvptx.cc | 6 +-
src/backend/metal/codegen/codegen_metal.cc | 6 +-
src/backend/opencl/codegen/codegen_opencl.cc | 4 +-
src/backend/rocm/codegen/llvm/codegen_amdgpu.cc | 6 +-
src/backend/trn/codegen/codegen_trn.cc | 22 +-
src/backend/trn/transform/lower_trainium_layout.cc | 30 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 8 +-
src/backend/webgpu/codegen/codegen_webgpu.cc | 6 +-
src/relax/op/tensor/inspect.cc | 13 +-
src/s_tir/analysis/calculate_allocated_memory.cc | 6 +-
src/s_tir/analysis/verify_gpu_code.cc | 6 +-
src/s_tir/backend/adreno/inject_texture_alloc.cc | 13 +-
src/s_tir/script/ir_builder/ir.cc | 20 +-
src/s_tir/transform/compact_tensor_region.cc | 7 +-
src/s_tir/transform/inject_virtual_thread.cc | 16 +-
src/s_tir/transform/lower_opaque_block.cc | 7 +-
src/s_tir/transform/lower_vtcm_alloc.cc | 10 +-
.../transform/merge_shared_memory_allocations.cc | 21 +-
src/script/printer/ir/call.cc | 28 +-
src/script/printer/ir/stmt.cc | 7 +-
src/target/llvm/codegen_llvm.cc | 8 +-
src/target/source/codegen_c.cc | 10 +-
src/tirx/ir/specialize.cc | 5 +-
src/tirx/ir/stmt_functor.cc | 3 +
src/tirx/ir/tile_dispatch.cc | 7 +-
src/tirx/op/memory.cc | 93 +++---
src/tirx/script/ir_builder/ir.cc | 75 -----
src/tirx/script/printer/stmt.cc | 37 +--
src/tirx/script/printer/tensor.cc | 122 +------
src/tirx/script/printer/utils.h | 12 -
src/tirx/transform/flatten_tensor.cc | 21 +-
src/tirx/transform/lower_thread_allreduce.h | 30 +-
src/tirx/transform/lower_tirx_cleanup.cc | 34 +-
src/tirx/transform/lower_tvm_builtin.cc | 20 +-
src/tirx/transform/lower_warp_memory.cc | 16 +-
src/tirx/transform/split_host_device.cc | 11 +-
src/tirx/transform/stmt_simplify.cc | 2 +-
src/tirx/transform/storage_rewrite.cc | 107 +++---
src/tirx/transform/tvm_ffi_binder.cc | 7 +-
src/tirx/transform/unsupported_dtype_legalize.cc | 22 +-
src/tirx/transform/update_pointer_storage_scope.cc | 8 +-
tests/cpp/ir_functor_test.cc | 35 +-
tests/python/codegen/test_gpu_codegen_allreduce.py | 4 +-
tests/python/codegen/test_target_codegen.py | 6 +-
tests/python/codegen/test_target_codegen_bool.py | 4 +-
tests/python/codegen/test_target_codegen_c_host.py | 6 +-
tests/python/codegen/test_target_codegen_cuda.py | 16 +-
.../python/codegen/test_target_codegen_cuda_fp4.py | 2 +-
.../python/codegen/test_target_codegen_cuda_fp8.py | 8 +-
tests/python/codegen/test_target_codegen_device.py | 2 +-
tests/python/codegen/test_target_codegen_llvm.py | 34 +-
tests/python/codegen/test_target_codegen_metal.py | 66 ++--
tests/python/codegen/test_target_codegen_rocm.py | 6 +-
tests/python/codegen/test_target_codegen_vulkan.py | 22 +-
tests/python/codegen/test_target_codegen_webgpu.py | 52 +--
tests/python/relax/backend/adreno/mod_utils.py | 6 +-
.../relax/script/test_relax_script_printer.py | 6 +-
tests/python/relax/test_relax_operators.py | 4 +
.../test_s_tir_analysis_estimate_tir_flops.py | 20 +-
.../test_s_tir_analysis_is_pure_function.py | 2 +-
.../s_tir/analysis/test_sblock_access_region.py | 2 +-
.../analysis/test_sblock_tensor_access_lca.py | 2 +-
.../test_meta_schedule_postproc_verify_gpu_code.py | 24 +-
.../s_tir/script/test_s_tir_script_basic_usage.py | 4 +-
.../script/test_s_tir_script_dynamic_shape.py | 10 +-
.../script/test_s_tir_script_error_handling.py | 2 +-
.../s_tir/script/test_s_tir_script_ir_builder.py | 4 +-
.../s_tir/script/test_s_tir_script_printer.py | 134 ++++----
tests/python/s_tir/test_s_tir_renew_defs.py | 2 +-
.../test_s_tir_transform_compact_tensor_region.py | 50 +--
.../test_s_tir_transform_inject_ptx_async_copy.py | 8 +-
.../test_s_tir_transform_inject_virtual_thread.py | 24 +-
.../test_s_tir_transform_loop_partition.py | 168 +++++++---
.../test_s_tir_transform_lower_opaque_block.py | 32 +-
...test_s_tir_transform_lower_thread_all_reduce.py | 134 ++++----
...form_merge_dynamic_shared_memory_allocations.py | 107 +++---
...sform_plan_update_tensor_allocation_location.py | 8 +-
.../test_s_tir_transform_rewrite_unsafe_select.py | 6 +-
.../transform/test_s_tir_transform_thread_sync.py | 86 ++---
.../test_tir_analysis_undefined_vars.py | 26 +-
.../test_tir_analysis_verify_well_formed.py | 8 +-
tests/python/tirx-base/test_tir_base.py | 2 +-
tests/python/tirx-base/test_tir_constructor.py | 8 +-
tests/python/tirx-base/test_tir_ptx_cp_async.py | 2 +-
tests/python/tirx-base/test_tir_ptx_ldmatrix.py | 4 +-
tests/python/tirx-base/test_tir_ptx_mma.py | 102 +++---
tests/python/tirx-base/test_tir_ptx_mma_sp.py | 52 +--
tests/python/tirx-base/test_tir_specialize.py | 32 +-
.../tirx-base/test_tir_stmt_functor_substitute.py | 12 +-
tests/python/tirx-base/test_tir_tensor.py | 9 +-
.../test_tir_transform_bf16_legalize.py | 96 +++---
.../test_tir_transform_convert_ssa.py | 23 +-
.../test_tir_transform_flatten_tensor.py | 125 ++++---
...test_tir_transform_force_narrow_index_to_i32.py | 2 +-
.../test_tir_transform_fp8_legalize.py | 26 +-
.../test_tir_transform_lower_tvm_builtin.py | 45 ++-
.../test_tir_transform_narrow_datatype.py | 8 +-
...est_tir_transform_pointer_value_type_rewrite.py | 18 +-
.../test_tir_transform_remove_no_op.py | 16 +-
.../tirx-transform/test_tir_transform_simplify.py | 38 ++-
.../test_tir_transform_split_host_device.py | 32 +-
.../test_tir_transform_storage_rewrite.py | 104 +++---
.../test_tir_transform_unroll_loop.py | 10 +-
.../tirx-transform/test_tir_transform_vectorize.py | 2 +-
.../python/tirx/codegen/test_codegen_blackwell.py | 58 ++--
tests/python/tirx/codegen/test_codegen_cuda.py | 37 +--
tests/python/tirx/codegen/test_codegen_dsmem.py | 2 +-
tests/python/tirx/codegen/test_codegen_hopper.py | 50 ++-
tests/python/tirx/codegen/test_codegen_nki.py | 32 +-
tests/python/tirx/codegen/test_cuda_cta_reduce.py | 12 +-
tests/python/tirx/codegen/test_ptx_addr.py | 10 +-
tests/python/tirx/codegen/test_ptx_dialect.py | 24 +-
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 8 +-
.../codegen/test_tcgen05_descriptor_encoders.py | 2 +-
tests/python/tirx/iket/test_iket_profiler.py | 2 +-
.../tile_primitive/cuda/copy/test_gmem_smem.py | 18 +-
.../tile_primitive/cuda/copy/test_ld_stmatrix.py | 60 +++-
.../operator/tile_primitive/cuda/copy/test_reg.py | 100 ++++--
.../cuda/copy/test_vec_forced_cache.py | 4 +-
.../tile_primitive/cuda/copy_async/test_dsmem.py | 12 +-
.../tile_primitive/cuda/copy_async/test_ldgsts.py | 6 +-
.../cuda/copy_async/test_tcgen05_cp.py | 38 ++-
.../tile_primitive/cuda/copy_async/test_tma.py | 76 +++--
.../tile_primitive/cuda/elementwise/test_binary.py | 116 +++++--
.../tile_primitive/cuda/elementwise/test_fma.py | 52 ++-
.../tile_primitive/cuda/elementwise/test_unary.py | 252 +++++++++-----
.../cuda/gemm/test_gemm_mma_m16n8k_.py | 68 ++--
.../cuda/gemm_async/test_gemm_async.py | 238 +++++++++----
.../operator/tile_primitive/trn/test_binary_trn.py | 116 ++++---
.../tile_primitive/trn/test_compose_op_trn.py | 368 +++++++++++++--------
.../operator/tile_primitive/trn/test_copy_trn.py | 202 +++++++----
.../operator/tile_primitive/trn/test_gemm_trn.py | 238 ++++++++-----
.../tile_primitive/trn/test_private_alloc_trn.py | 286 ++++++++++++----
.../tile_primitive/trn/test_reduction_trn.py | 80 +++--
.../operator/tile_primitive/trn/test_select_trn.py | 48 ++-
.../operator/tile_primitive/trn/test_unary_trn.py | 88 +++--
.../tirx/script/test_tirx_script_basic_usage.py | 18 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 2 +-
tests/python/tirx/script/test_tirx_script_gpu.py | 110 +++---
.../tirx/script/test_tirx_script_ir_builder.py | 2 +-
.../script/test_tirx_script_meta_programming.py | 6 +-
.../python/tirx/script/test_tirx_script_printer.py | 60 ++--
.../test_tirx_script_printer_structural_equal.py | 4 +-
.../script/test_tirx_script_source_locations.py | 10 +-
.../python/tirx/script/test_tirx_script_tensors.py | 349 ++++++++++++-------
tests/python/tirx/test_compilation_pipeline.py | 2 +-
tests/python/tirx/test_control_flow.py | 4 +-
tests/python/tirx/test_cuda_launch.py | 4 +-
tests/python/tirx/test_verifier.py | 18 +-
.../transform/test_transform_flatten_tensor.py | 12 +-
.../tirx/transform/test_transform_lower_tirx.py | 202 +++++++----
.../transform/test_transform_naive_allocator.py | 72 ++--
196 files changed, 4200 insertions(+), 3153 deletions(-)
diff --git a/docs/deep_dive/relax/learning.rst b/docs/deep_dive/relax/learning.rst
index be056a3ea8..2a26728681 100644
--- a/docs/deep_dive/relax/learning.rst
+++ b/docs/deep_dive/relax/learning.rst
@@ -139,7 +139,7 @@ for the end-to-end model execution. The code block below shows a TVMScript imple
- Y = T.alloc_tensor((M, N), "float32")
+ Y = T.alloc_tensor(ty_args=[T.Tensor((M, N), "float32")])
for i, j, k in T.grid(M, N, K):
with Ts.sblock("Y"):
v_i, v_j, v_k = Ts.axis.remap("SSR", [i, j, k])
diff --git a/docs/deep_dive/tensor_ir/learning.rst b/docs/deep_dive/tensor_ir/learning.rst
index f19f4d1ceb..4582dc1968 100644
--- a/docs/deep_dive/tensor_ir/learning.rst
+++ b/docs/deep_dive/tensor_ir/learning.rst
@@ -68,7 +68,7 @@ language called TVMScript, which is a domain-specific dialect embedded in python
def mm_relu(A: T.Tensor((128, 128), "float32"),
B: T.Tensor((128, 128), "float32"),
C: T.Tensor((128, 128), "float32")):
- Y = T.alloc_tensor((128, 128), dtype="float32")
+ Y = T.alloc_tensor(ty_args=[T.Tensor((128, 128), "float32")])
for i, j, k in T.grid(128, 128, 128):
with Ts.sblock("Y"):
vi = Ts.axis.spatial(128, i)
@@ -111,7 +111,7 @@ type.
.. code:: python
# TensorIR
- Y = T.alloc_tensor((128, 128), dtype="float32")
+ Y = T.alloc_tensor(ty_args=[T.Tensor((128, 128), "float32")])
# NumPy
Y = np.empty((128, 128), dtype="float32")
@@ -243,7 +243,7 @@ So we can also write the programs as follows.
def mm_relu(A: T.Tensor((128, 128), "float32"),
B: T.Tensor((128, 128), "float32"),
C: T.Tensor((128, 128), "float32")):
- Y = T.alloc_tensor((128, 128), dtype="float32")
+ Y = T.alloc_tensor(ty_args=[T.Tensor((128, 128), "float32")])
for i, j, k in T.grid(128, 128, 128):
with Ts.sblock("Y"):
vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
diff --git a/docs/deep_dive/tensor_ir/tutorials/tir_creation.py b/docs/deep_dive/tensor_ir/tutorials/tir_creation.py
index 166cc8bd35..a96f0cb45f 100644
--- a/docs/deep_dive/tensor_ir/tutorials/tir_creation.py
+++ b/docs/deep_dive/tensor_ir/tutorials/tir_creation.py
@@ -69,7 +69,7 @@ class MyModule:
B: T.Tensor((128, 128), "float32"),
C: T.Tensor((128, 128), "float32"),
):
- Y = T.alloc_tensor((128, 128), dtype="float32")
+ Y = T.alloc_tensor(ty_args=[T.Tensor((128, 128), "float32")])
for i in range(128):
for j in range(128):
for k in range(128):
@@ -112,7 +112,7 @@ class ConciseModule:
B: T.Tensor((128, 128), "float32"),
C: T.Tensor((128, 128), "float32"),
):
- Y = T.alloc_tensor((128, 128), dtype="float32")
+ Y = T.alloc_tensor(ty_args=[T.Tensor((128, 128), "float32")])
for i, j, k in T.grid(128, 128, 128):
with Ts.sblock("Y"):
vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
@@ -151,7 +151,7 @@ class ConciseModuleFromPython:
B: T.Tensor((K, N), dtype),
C: T.Tensor((M, N), dtype),
):
- Y = T.alloc_tensor((M, N), dtype)
+ Y = T.alloc_tensor(ty_args=[T.Tensor((M, N), dtype)])
for i, j, k in T.grid(M, N, K):
with Ts.sblock("Y"):
vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
@@ -188,7 +188,7 @@ class DynamicShapeModule:
def mm_relu(A: T.Tensor([M, K], dtype), B: T.Tensor([K, N], dtype), C: T.Tensor([M, N], dtype)):
# Bind the input buffers with the dynamic shapes
- Y = T.alloc_tensor((M, N), dtype)
+ Y = T.alloc_tensor(ty_args=[T.Tensor((M, N), dtype)])
for i, j, k in T.grid(M, N, K):
with Ts.sblock("Y"):
vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
diff --git a/docs/tirx/native_basics/cuda/data_types.rst b/docs/tirx/native_basics/cuda/data_types.rst
index 393b3d6c99..505ce62733 100644
--- a/docs/tirx/native_basics/cuda/data_types.rst
+++ b/docs/tirx/native_basics/cuda/data_types.rst
@@ -101,7 +101,7 @@ A tensor's ``data`` — its pointer — is a ``Var`` of pointer type, and it is
**immutable** (a pointer is never reassigned). That shapes how you obtain one:
- ``Tx.alloc_tensor(...)`` allocates storage **and** defines its ``data`` pointer.
-- ``Tx.decl_tensor(..., data=ptr)`` declares a tensor over an existing pointer
+- ``Tx.decl_tensor(ptr, ty_args=[Tx.Tensor(...)])`` declares a tensor over an existing pointer
``Var`` ``ptr``.
- To back a tensor with a pointer **expression** — e.g. ``Tx.ptx.mapa`` giving
another cluster CTA's shared address — convert the ``uint64`` address the
@@ -118,7 +118,7 @@ A tensor's ``data`` — its pointer — is a ``Var`` of pointer type, and it is
Tx.ptx.mapa.u64(mapped[0], mbar.ptr_to([0]), Tx.uint32(0))
ptr_ty = PtrType(PrimType("uint64"), "shared")
ptr = Tx.reinterpret(ptr_ty, mapped[0])
- remote_mbar = Tx.decl_tensor([1], "uint64", data=ptr, scope="shared")
+ remote_mbar = Tx.decl_tensor(ptr, ty_args=[Tx.Tensor([1], "uint64", scope="shared")])
Pointer bindings cannot be reassigned; use a new name for a different
pointer value.
diff --git a/docs/tirx/native_basics/cuda/parser_utils.rst b/docs/tirx/native_basics/cuda/parser_utils.rst
index 5a3856ce72..4172adcd91 100644
--- a/docs/tirx/native_basics/cuda/parser_utils.rst
+++ b/docs/tirx/native_basics/cuda/parser_utils.rst
@@ -64,7 +64,9 @@ allocations and state into one object and use it in the kernel body.
class State:
def __init__(self, smem):
self.acc = Tx.alloc_local([1], "float32")
- self.buf = Tx.decl_tensor([64], "float16", smem, scope="shared.dyn")
+ self.buf = Tx.decl_tensor(
+ smem, ty_args=[Tx.Tensor([64], "float16", scope="shared.dyn")]
+ )
s = State(smem.data)
s.acc[0] = Tx.float32(0.0) # use its fields like ordinary tensors
diff --git a/docs/tirx/native_basics/cuda/tensors.rst b/docs/tirx/native_basics/cuda/tensors.rst
index 61feebc759..b980b4bd07 100644
--- a/docs/tirx/native_basics/cuda/tensors.rst
+++ b/docs/tirx/native_basics/cuda/tensors.rst
@@ -28,22 +28,20 @@ pointer with ``A.ptr_to([i, j])`` or the raw data pointer ``A.data``.
Declaring tensors
-----------------
-The following APIs create a tensor variable:
-
-- ``Tx.alloc_tensor(shape, dtype, scope=..., ...)`` — **allocates new storage**
- (binds a ``tirx.alloc_tensor`` call) and returns an ``ir.Var``. ``Tx.alloc_shared`` /
- ``Tx.alloc_local`` are just ``alloc_tensor`` with ``scope="shared"`` /
- ``scope="local"``.
-- ``Tx.decl_tensor(shape, dtype, data=..., ...)`` — **declares a view** over an
- existing pointer ``data`` (no allocation); use it to alias or reinterpret
- storage — for example, a sub-region of a pool. With ``data=None`` it allocates,
- like ``alloc_tensor``.
+The following calls bind tensor variables when assigned in a script function:
+
+- ``Tx.alloc_tensor(ty_args=[Tx.Tensor(shape, dtype, scope=scope)])``
+ **allocates new storage**. ``Tx.alloc_shared`` and ``Tx.alloc_local`` remain
+ conveniences for constructing a tensor type with shared or local storage.
+- ``Tx.decl_tensor(data, ty_args=[Tx.Tensor(shape, dtype, scope=scope)])``
+ **declares a view** over an existing pointer ``data`` (no allocation); use it
+ to alias or reinterpret storage, for example a sub-region of a pool.
- ``Tx.cuda.decl_tmem(addr, ty_args=[Tx.Tensor(...)])`` — declares a tensor over
externally allocated tensor memory. The address is a value operand; the tensor
type supplies shape, dtype, scope and layout.
``A.data`` projects the physical pointer from the tensor variable.
-``alloc_tensor`` supplies new storage; ``decl_tensor`` with ``data`` binds an
+``alloc_tensor`` supplies new storage; ``decl_tensor`` binds an
existing pointer expression. See :doc:`data_types`.
The tensor descriptors share these parameters:
@@ -98,7 +96,9 @@ The ``scope`` argument selects the memory space:
def kernel(A: Tx.Tensor((M, K), "float16", align=16)):
As = Tx.alloc_shared((BM, BK), "float16") # new shared tile
acc = Tx.alloc_local((4,), "float32") # per-thread accumulator
- view = Tx.decl_tensor((BM, BK), "float16", data=As.data) # a view over As
+ view = Tx.decl_tensor(
+ As.data, ty_args=[Tx.Tensor((BM, BK), "float16", scope="shared")]
+ ) # a view over As
**A ptr-based tensor is just metadata over a pointer.** For any non-tmem tensor,
the declaration is a pointer plus a layout, and indexing resolves to an address::
@@ -175,14 +175,20 @@ Dynamic
**Dynamic** shared memory (``scope="shared.dyn"``) is sized per launch (the
``sharedMemBytes`` launch parameter), not at compile time. A kernel may have **only
one** dynamic-shared allocation — the *arena*. So you allocate it once and ``decl``
-each tensor as a view into it: ``Tx.decl_tensor`` with ``data=`` the arena pointer
-and an ``elem_offset``:
+each tensor as a view into it: ``Tx.decl_tensor`` with the arena pointer as its
+first argument and an ``elem_offset``:
.. code-block:: python
- arena = Tx.alloc_tensor((128,), "float32", scope="shared.dyn") # the one arena
- As = Tx.decl_tensor((64,), "float32", data=arena.data, scope="shared.dyn") # offset 0
- Bs = Tx.decl_tensor((64,), "float32", data=arena.data, elem_offset=64, scope="shared.dyn") # offset 64
+ arena = Tx.alloc_tensor(
+ ty_args=[Tx.Tensor((128,), "float32", scope="shared.dyn")]
+ ) # the one arena
+ As = Tx.decl_tensor(
+ arena.data, ty_args=[Tx.Tensor((64,), "float32", scope="shared.dyn")]
+ ) # offset 0
+ Bs = Tx.decl_tensor(
+ arena.data, ty_args=[Tx.Tensor((64,), "float32", elem_offset=64, scope="shared.dyn")]
+ ) # offset 64
As[tx] = A[tx]
Bs[tx] = B[tx]
Tx.cuda.cta_sync()
@@ -199,8 +205,8 @@ boilerplate elided; arena named ``smem`` for clarity):
__syncthreads();
C_ptr[tx] = smem[tx] + smem[tx + 64];
-(Two separate ``alloc_tensor(scope="shared.dyn")`` is an error — *only one dynamic
-shared memory allocation is allowed*.) So static shared memory is sized at compile
+(Two separate ``alloc_tensor`` calls with ``scope="shared.dyn"`` tensor types
+are an error — *only one dynamic shared memory allocation is allowed*.) So static shared memory is sized at compile
time (``__shared__ T x[N];``); dynamic shared memory is this one launch-sized arena
with views decl'd at offsets inside it.
@@ -509,7 +515,13 @@ no shape infers a one-dimensional shape from the logical storage size:
.. code-block:: python
- R = Tx.alloc_tensor((32, 8), "float32", scope="local", layout=TileLayout(S[(32, 8) : (1 @ laneid, 1)]))
+ R = Tx.alloc_tensor(
+ ty_args=[
+ Tx.Tensor(
+ (32, 8), "float32", scope="local", layout=TileLayout(S[(32, 8) : (1 @ laneid, 1)])
+ )
+ ]
+ )
R_flat = R.local() # this lane's 8 local elements, physical order
R_2d = R.local(2, 4) # the same elements, row-major 2x4 reshape
diff --git a/docs/tirx/tile_primitives/copy/ldstmatrix.rst b/docs/tirx/tile_primitives/copy/ldstmatrix.rst
index 7bdec6ee0b..8f429e292a 100644
--- a/docs/tirx/tile_primitives/copy/ldstmatrix.rst
+++ b/docs/tirx/tile_primitives/copy/ldstmatrix.rst
@@ -117,10 +117,14 @@ register, from ``test_ld_stmatrix.py`` (register layout = the m8n8 fragment,
_index = Tx.cuda.block_idx("x")
_lane = Tx.cuda.lane_id()
tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor((8, 4, num, 2), "float16", scope="shared", layout=s_layout)
+ A_smem = Tx.alloc_tensor(
+ ty_args=[Tx.Tensor((8, 4, num, 2), "float16", scope="shared", layout=s_layout)]
+ )
# ... stage A into A_smem (row = tid//4, cp = tid%4) ...
Tx.cuda.cta_sync()
- R = Tx.alloc_tensor((8, 4, num, 2), "float16", scope="local", layout=r_layout)
+ R = Tx.alloc_tensor(
+ ty_args=[Tx.Tensor((8, 4, num, 2), "float16", scope="local", layout=r_layout)]
+ )
Tx.cuda.tile.ldmatrix(R[full], A_smem[full], scope="warp") # shared -> register (ldmatrix)
# ... write R back out to B ...
diff --git a/docs/tirx/tile_primitives/copy/reg.rst b/docs/tirx/tile_primitives/copy/reg.rst
index e0724e2968..f994409706 100644
--- a/docs/tirx/tile_primitives/copy/reg.rst
+++ b/docs/tirx/tile_primitives/copy/reg.rst
@@ -94,11 +94,11 @@ contiguous elements). From ``test_reg.py``:
_index = Tx.cuda.block_idx("x")
_lane = Tx.cuda.lane_id()
tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
+ A_smem = Tx.alloc_tensor(ty_args=[Tx.Tensor(shape, dtype, scope="shared", layout=s_layout)])
for kk in range(8):
A_smem[tid, kk] = Tx.cast(tid * 100 + kk + 1, dtype)
Tx.cuda.cta_sync()
- R = Tx.alloc_tensor(shape, dtype, scope="local", layout=r_layout)
+ R = Tx.alloc_tensor(ty_args=[Tx.Tensor(shape, dtype, scope="local", layout=r_layout)])
Tx.cuda.tile.ld(R[fs], A_smem[fs], scope="warp") # shared -> register (this dispatch)
# ... clear A_smem, cta_sync ...
Tx.cuda.tile.st(A_smem[fs], R[fs], scope="warp") # register -> shared (this dispatch)
@@ -151,7 +151,9 @@ Generated TIRx IR
.. code-block:: python
- r_local = Tx.decl_tensor((8,), data=R.data, scope="local") # 8 fp32 elements / lane
+ r_local = Tx.decl_tensor(
+ R.data, ty_args=[Tx.Tensor((8,), scope="local")]
+ ) # 8 fp32 elements / lane
r_words = r_local.view("uint32")
for f in range(2): # outer = 8 / vec 4
s_ptr = pointer_offset(A_smem, ...) # this lane's row
diff --git a/docs/tirx/tile_primitives/copy_async/dsmem.rst b/docs/tirx/tile_primitives/copy_async/dsmem.rst
index a7ec6becef..97fe93976d 100644
--- a/docs/tirx/tile_primitives/copy_async/dsmem.rst
+++ b/docs/tirx/tile_primitives/copy_async/dsmem.rst
@@ -80,21 +80,29 @@ mbarrier and writes the result out (from ``test_dsmem.py``):
pool = Tx.SMEMPool()
src_raw = pool.alloc([8192], dtype, align=128)
src_smem = Tx.decl_tensor(
- list(shape),
- dtype,
src_raw.data,
- byte_offset=src_raw.byte_offset,
- scope="shared.dyn",
- layout=src_layout,
+ ty_args=[
+ Tx.Tensor(
+ list(shape),
+ dtype,
+ byte_offset=src_raw.byte_offset,
+ scope="shared.dyn",
+ layout=src_layout,
+ )
+ ],
)
dst_raw = pool.alloc([8192], dtype, align=128)
dst_smem = Tx.decl_tensor(
- list(shape),
- dtype,
dst_raw.data,
- byte_offset=dst_raw.byte_offset,
- scope="shared.dyn",
- layout=dst_layout,
+ ty_args=[
+ Tx.Tensor(
+ list(shape),
+ dtype,
+ byte_offset=dst_raw.byte_offset,
+ scope="shared.dyn",
+ layout=dst_layout,
+ )
+ ],
)
mbar = MBarrier(pool, 1)
pool.commit()
diff --git a/docs/tirx/tile_primitives/copy_async/ldgsts.rst b/docs/tirx/tile_primitives/copy_async/ldgsts.rst
index 95ed53bf85..8ccb773f59 100644
--- a/docs/tirx/tile_primitives/copy_async/ldgsts.rst
+++ b/docs/tirx/tile_primitives/copy_async/ldgsts.rst
@@ -101,7 +101,7 @@ shared, then commits and waits before reading it back (from ``test_ldgsts.py``):
_index = Tx.cuda.warp_id()
_lane = Tx.cuda.lane_id()
tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
+ A_smem = Tx.alloc_tensor(ty_args=[Tx.Tensor(shape, dtype, scope="shared", layout=s_layout)])
Tx.cuda.tile.cp_async(A_smem[full], A[full], scope="cta") # async global -> shared
Tx.ptx.cp.async_.commit_group() # caller commits ...
Tx.ptx.cp.async_.wait_group(0) # ... and waits
diff --git a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst b/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
index 2da2f9d744..58aa887804 100644
--- a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
+++ b/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
@@ -116,8 +116,17 @@ dealloc tail elided):
from tvm.tirx.layout import R, S, TCol, TileLayout, TLane
- A_smem = Tx.alloc_tensor([32, 16], "uint8", scope="shared",
- layout=TileLayout(S[(32, 16) : (16, 1)]), align=1024)
+ A_smem = Tx.alloc_tensor(
+ ty_args=[
+ Tx.Tensor(
+ [32, 16],
+ "uint8",
+ scope="shared",
+ layout=TileLayout(S[(32, 16) : (16, 1)]),
+ align=1024,
+ )
+ ]
+ )
tmem_addr = Tx.alloc_shared([1], "uint32")
cp_mbar = Tx.alloc_shared([1], "uint64")
if warp_id == 0:
diff --git a/docs/tirx/tile_primitives/elementwise/reg.rst b/docs/tirx/tile_primitives/elementwise/reg.rst
index b07e0ff4f8..868e32bf84 100644
--- a/docs/tirx/tile_primitives/elementwise/reg.rst
+++ b/docs/tirx/tile_primitives/elementwise/reg.rst
@@ -91,12 +91,14 @@ A warp takes the elementwise ``sqrt`` of a ``32×8`` ``float32`` local tile
_index = Tx.cuda.block_idx("x")
_lane = Tx.cuda.lane_id()
tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor((32, 8), "float32", scope="shared", layout=TileLayout(S[(32, 8)]))
+ A_smem = Tx.alloc_tensor(
+ ty_args=[Tx.Tensor((32, 8), "float32", scope="shared", layout=TileLayout(S[(32, 8)]))]
+ )
for k in Tx.serial(8):
value = A[tid, k]
A_smem[tid, k] = value
Tx.cuda.cta_sync()
- R = Tx.alloc_tensor((32, 8), "float32", scope="local", layout=r_layout)
+ R = Tx.alloc_tensor(ty_args=[Tx.Tensor((32, 8), "float32", scope="local", layout=r_layout)])
Tx.cuda.tile.ld(R[fs], A_smem[fs], scope="warp")
Tx.cuda.tile.sqrt(R[fs], R[fs], scope="warp") # elementwise reg dispatch
Tx.cuda.tile.st(A_smem[fs], R[fs], scope="warp")
diff --git a/docs/tirx/tile_primitives/elementwise/smem.rst b/docs/tirx/tile_primitives/elementwise/smem.rst
index 1e593cc28d..71a23806d0 100644
--- a/docs/tirx/tile_primitives/elementwise/smem.rst
+++ b/docs/tirx/tile_primitives/elementwise/smem.rst
@@ -85,7 +85,9 @@ round):
_index = Tx.cuda.warp_id()
_lane = Tx.cuda.lane_id()
tid = Tx.cuda.thread_idx("x")
- A_smem = Tx.alloc_tensor((32, 32), "float32", scope="shared", layout=s_layout)
+ A_smem = Tx.alloc_tensor(
+ ty_args=[Tx.Tensor((32, 32), "float32", scope="shared", layout=s_layout)]
+ )
for k in Tx.serial(4):
index = tid + k * 256
value = A[index // 32, index % 32]
diff --git a/docs/tirx/tile_primitives/gemm.rst b/docs/tirx/tile_primitives/gemm.rst
index 1a7f0739b1..5a29aea937 100644
--- a/docs/tirx/tile_primitives/gemm.rst
+++ b/docs/tirx/tile_primitives/gemm.rst
@@ -89,9 +89,11 @@ accumulate) — one ``m16n8k16`` atom (from ``test_gemm_mma_m16n8k_.py``):
_index = Tx.cuda.block_idx("x")
_index = Tx.cuda.warp_id()
lane = Tx.cuda.lane_id()
- A_f = Tx.alloc_tensor((16, 16), "float16", scope="local", layout=A_FRAG)
- B_f = Tx.alloc_tensor((16, 8), "float16", scope="local", layout=B_FRAG)
- D_f = Tx.alloc_tensor((16, 8), "float32", scope="local", layout=D_FRAG)
+ A_f = Tx.alloc_tensor(
+ ty_args=[Tx.Tensor((16, 16), "float16", scope="local", layout=A_FRAG)]
+ )
+ B_f = Tx.alloc_tensor(ty_args=[Tx.Tensor((16, 8), "float16", scope="local", layout=B_FRAG)])
+ D_f = Tx.alloc_tensor(ty_args=[Tx.Tensor((16, 8), "float32", scope="local", layout=D_FRAG)])
A_reg = A_f.local(8) # stage A into the lane's 8 regs
for s in Tx.unroll(8):
kp, rM, kHi = s % 2, (s // 2) % 2, s // 4
diff --git a/docs/tirx/tile_primitives/gemm_async.rst b/docs/tirx/tile_primitives/gemm_async.rst
index bae639a2d5..9775bb79b4 100644
--- a/docs/tirx/tile_primitives/gemm_async.rst
+++ b/docs/tirx/tile_primitives/gemm_async.rst
@@ -95,8 +95,26 @@ into shared (from
from tvm.tirx.layout import S, TCol, TLane, TileLayout, tid_in_wg as axis_tid_in_wg
from tvm.backend.cuda.tile_primitive.tma_utils import mma_shared_layout
- A_smem = Tx.alloc_tensor((3,128,64), "float16", scope="shared", layout=mma_shared_layout("float16", 3, (3,128,64)))
- B_smem = Tx.alloc_tensor((3,128,64), "float16", scope="shared", layout=mma_shared_layout("float16", 3, (3,128,64)))
+ A_smem = Tx.alloc_tensor(
+ ty_args=[
+ Tx.Tensor(
+ (3, 128, 64),
+ "float16",
+ scope="shared",
+ layout=mma_shared_layout("float16", 3, (3, 128, 64)),
+ )
+ ]
+ )
+ B_smem = Tx.alloc_tensor(
+ ty_args=[
+ Tx.Tensor(
+ (3, 128, 64),
+ "float16",
+ scope="shared",
+ layout=mma_shared_layout("float16", 3, (3, 128, 64)),
+ )
+ ]
+ )
tmem_addr = Tx.alloc_shared([1], "uint32"); mma_mbar = Tx.alloc_shared([1], "uint64")
# ... mbarrier.init, cta_sync ...
if warp_id == 0:
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 7d89f99b51..a30f365fd2 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -100,7 +100,7 @@ class MatchTensorRegion : public ffi::ObjectRef {
* T.reads([tensor0[start:end, ...], ...])
* T.writes([tensor1[start:end, ...], ...])
* T.where(predicate)
- * tensor2 = T.alloc_tensor(shape, dtype)
+ * tensor2 = T.alloc_tensor(ty_args=[T.Tensor(shape, dtype)])
* tensor3 = Ts.match_tensor(source_tensor[start:end, ...])
* Ts.sblock_attr({attr_key: attr_value, ...})
* with T.init():
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index be87e4c6c4..1513e12cab 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -99,7 +99,7 @@ TVM_DLL Pass LiftThreadBinding();
*
* for i in range(0, 16):
* with T.sblock():
- * B = T.alloc_tensor(16, 16)
+ * B = T.alloc_tensor(ty_args=[T.Tensor((16, 16), "float32")])
* for j in range(0, 16):
* B[i, j] = A[i, j] + 1
* for j in range(0, 16):
@@ -115,7 +115,7 @@ TVM_DLL Pass LiftThreadBinding();
*
* for i in range(0, 16):
* with T.sblock():
- * B = T.alloc_tensor(1, 16)
+ * B = T.alloc_tensor(ty_args=[T.Tensor((1, 16), "float32")])
* for j in range(0, 16):
* B[0, j] = A[i, j] + 1
* for j in range(0, 16):
diff --git a/include/tvm/tirx/op/memory.h b/include/tvm/tirx/op/memory.h
index 18e8dc69ae..92df1fe024 100644
--- a/include/tvm/tirx/op/memory.h
+++ b/include/tvm/tirx/op/memory.h
@@ -31,31 +31,19 @@ namespace tvm {
namespace tirx {
/*!
- * \brief Allocate a tensor: alloc_tensor(shape, dtype, scope) -> TensorType.
- *
- * Arguments, in order:
- * - args[0]: shape, Tuple of integer extents (IntImm or symbolic integer expressions).
- * - args[1]: dtype, DataTypeImm with a DLDataType payload for the element type.
- * - args[2]: scope, StringImm naming the storage scope.
- *
- * DictAttrs directly holds the allocation annotations, defaulting to an empty dictionary.
- * The TensorType result agrees with the operands and retains tensor access/storage metadata.
+ * \brief Allocate a tensor with exactly one TensorType in ty_args.
*
+ * The optional args[0] is a Tuple of producer-owned placement coordinates.
+ * DictAttrs directly holds allocation annotations. The TensorType owns all
+ * shape, dtype, storage scope and access/layout metadata, and is the result type.
*/
TVM_DLL const Op& alloc_tensor_op();
/*!
- * \brief Declare a tensor view: decl_tensor(data, shape, dtype, scope) -> TensorType.
- *
- * Arguments, in order:
- * - args[0]: data, Expr for the existing physical pointer backing the tensor view.
- * - args[1]: shape, Tuple of integer extents (IntImm or symbolic integer expressions).
- * - args[2]: dtype, DataTypeImm with a DLDataType payload for the element type.
- * - args[3]: scope, StringImm naming the storage scope.
- *
- * There are no attributes. The TensorType result agrees with the operands and retains
- * tensor access/storage metadata. The operation binds a view without allocating memory.
+ * \brief Declare a tensor view with exactly one TensorType in ty_args.
*
+ * args[0] is the existing physical pointer backing the view. There are no
+ * attributes. The TensorType is the result type; no memory is allocated.
*/
TVM_DLL const Op& decl_tensor_op();
diff --git a/include/tvm/tirx/script/ir_builder/ir.h b/include/tvm/tirx/script/ir_builder/ir.h
index d01017776e..9250b7d064 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -25,7 +25,6 @@
#include <tvm/runtime/tensor.h>
#include <tvm/script/ir_builder/ir.h>
#include <tvm/tirx/exec_scope.h>
-#include <tvm/tirx/layout.h>
#include <tvm/tirx/script/ir_builder/frame.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/type.h>
@@ -53,27 +52,8 @@ using tvm::ffi::Tuple;
using tvm::ffi::Variant;
using tvm::runtime::Tensor;
using tvm::tirx::ExecScope;
-using tvm::tirx::Layout;
using tvm::tirx::TensorVar;
-/*!
- * \brief The tensor declaration function.
- * \param shape The type of the tensor prior to flattening.
- * \param dtype The data type in the content of the tensor.
- * \param tensor_name The name of the tensor.
- * \param data The pointer to the head of the data.
- * \param strides The strides of each dimension.
- * \param elem_offset The offset in terms of number of dtype elements (including lanes).
- * \param storage_scope The optional storage scope of tensor data pointer.
- * \param align The alignment requirement of data pointer in bytes.
- * \param offset_factor The factor of elem_offset field.
- * \return The declared tensor.
- */
-TensorVar TensorDecl(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String tensor_name,
- ffi::Optional<Expr> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
- ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
- int offset_factor, ffi::Optional<Layout> layout = std::nullopt);
-
/*!
* \brief The primitive function statement.
* \return The FunctionFrame.
@@ -181,38 +161,6 @@ ForFrame Unroll(PrimExpr start, PrimExpr stop,
*/
ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread,
ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
-/*!
- * \brief Declare a tensor binding or allocate its storage.
- * \param shape The type of the tensor prior to flattening.
- * \param dtype The data type in the content of the tensor.
- * \param tensor_name The name of the tensor.
- * \param data The pointer to the head of the data.
- * \param strides The strides of each dimension.
- * \param elem_offset The offset in terms of number of dtype elements (including lanes).
- * \param storage_scope The optional storage scope of tensor data pointer.
- * \param align The alignment requirement of data pointer in bytes.
- * \param offset_factor The factor of elem_offset field.
- * \param layout The layout of the tensor.
- * \return The declared tensor.
- */
-TensorVar DeclTensor(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String tensor_name,
- ffi::Optional<Expr> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
- ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
- int offset_factor, ffi::Optional<Layout> layout = std::nullopt,
- ffi::Optional<PrimExpr> allocated_addr = std::nullopt);
-
-/*!
- * \brief Statement-level tensor allocation (binds a tensor-returning allocation Call).
- * \param shape The shape of the tensor to allocate.
- * \param dtype The data type of tensor elements.
- * \param storage_scope The storage scope (e.g., "global", "shared").
- * \param annotations Optional annotations for the allocation.
- * \return The allocated tensor.
- */
-TensorVar AllocTensor(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
- ffi::String storage_scope = "global",
- ffi::Optional<ffi::Map<ffi::String, ffi::Any>> annotations = std::nullopt);
-
/*!
* \brief Store data in a tensor destination.
* \param dest The destination expression.
diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py b/python/tvm/backend/cuda/lang/alloc_pool.py
index 33f202ef6b..7ef848976e 100644
--- a/python/tvm/backend/cuda/lang/alloc_pool.py
+++ b/python/tvm/backend/cuda/lang/alloc_pool.py
@@ -411,7 +411,7 @@ class SMEMPool:
Parameters
----------
ptr : Var or None, optional
- If omitted, an ``alloc_tensor([0], "uint8", scope="shared.dyn")`` is
+ If omitted, an ``alloc_tensor(ty_args=[Tensor([0], "uint8", scope="shared.dyn")])`` is
created automatically and ``commit()`` must be called after all
allocations to emit the size annotation.
If a ``Var`` is provided, the caller manages the backing buffer and
@@ -421,7 +421,9 @@ class SMEMPool:
def __init__(self, ptr=_POOL_UNSET):
ir = _get_ir()
if ptr is _POOL_UNSET:
- self.buf = ir.alloc_tensor([0], "uint8", scope="shared.dyn")
+ self.buf = ir._bind_tensor(
+ ir.alloc_tensor(ty_args=[ir.Tensor([0], "uint8", scope="shared.dyn")])
+ )
self.ptr = self.buf.data
self._owns_buffer = True
else:
@@ -443,15 +445,21 @@ class SMEMPool:
ir = _get_ir()
if align > 0:
self.offset = (self.offset + align - 1) // align * align
- res = ir.decl_tensor(
- shape,
- dtype,
- data=self.ptr,
- strides=strides,
- byte_offset=self.offset,
- scope=scope,
- align=align,
- layout=layout,
+ res = ir._bind_tensor(
+ ir.decl_tensor(
+ self.ptr,
+ ty_args=[
+ ir.Tensor(
+ shape,
+ dtype,
+ strides=strides,
+ byte_offset=self.offset,
+ scope=scope,
+ align=align,
+ layout=layout,
+ )
+ ],
+ )
)
# Advance in bits then round up to bytes so sub-byte dtypes (e.g.
# float4_e2m1fn = 4 bits) still bump the cursor instead of leaving it
diff --git a/python/tvm/backend/cuda/lang/pipeline.py b/python/tvm/backend/cuda/lang/pipeline.py
index 4260eb5337..3c6a3eca63 100644
--- a/python/tvm/backend/cuda/lang/pipeline.py
+++ b/python/tvm/backend/cuda/lang/pipeline.py
@@ -117,7 +117,7 @@ def _map_buffer_into_cta(ptr, rank, depth):
T.evaluate(T.ptx.mapa.u64(mapped[0], ptr, T.uint32(rank)))
remote_ptr = TIRVar("remote_mbar_ptr", ptr_ty)
T.bind(T.reinterpret(mapped[0], ty=ptr_ty), var=remote_ptr)
- return T.decl_tensor([depth], "uint64", data=remote_ptr, scope="shared")
+ return T.bind(T.decl_tensor(remote_ptr, ty_args=[T.Tensor([depth], "uint64", scope="shared")]))
def _mbarrier_arrive_remote(bar, pred=None, count=None):
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
index 5bac05b376..d62964a7e5 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
@@ -163,16 +163,20 @@ def copy_dsmem_impl(op_call: TensorCall, sctx: DispatchContext) -> Function:
src_elem_offset, dst_elem_offset = T.meta_var(compute_offsets(loop_vars))
src_buf_w = T.decl_tensor(
- src_buf.shape, src_buf.dtype, src_buf.data,
- elem_offset=src_buf.elem_offset + src_elem_offset,
- scope=src_buf.scope(),
- layout=src_tile,
+ src_buf.data,
+ ty_args=[T.Tensor(
+ src_buf.shape, src_buf.dtype,
+ elem_offset=src_buf.elem_offset + src_elem_offset,
+ scope=src_buf.scope(), layout=src_tile,
+ )],
)
dst_buf_w = T.decl_tensor(
- dst_buf.shape, dst_buf.dtype, dst_buf.data,
- elem_offset=dst_buf.elem_offset + dst_elem_offset,
- scope=dst_buf.scope(),
- layout=dst_tile,
+ dst_buf.data,
+ ty_args=[T.Tensor(
+ dst_buf.shape, dst_buf.dtype,
+ elem_offset=dst_buf.elem_offset + dst_elem_offset,
+ scope=dst_buf.scope(), layout=dst_tile,
+ )],
)
src_ptr = src_buf_w.ptr_to(src_st)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
index d1fe9003a4..9dcb53c10c 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
@@ -110,7 +110,7 @@ import operator
import tvm
from tvm.backend.cuda.op import _tmem_address
-from tvm.ir import Bind, Call, DataTypeImm, DictAttrs, Evaluate, SeqStmt, StringImm, Tuple
+from tvm.ir import Bind, Call, DictAttrs, Evaluate, SeqStmt, StringImm
from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.sym import Analyzer
@@ -688,13 +688,9 @@ def _get_or_create_desc(sctx, s_buf, ldo, sdo, swizzle):
desc_buf,
Call(
"tirx.alloc_tensor",
- [
- Tuple(desc_buf.ty.shape),
- DataTypeImm(desc_buf.ty.dtype.dtype),
- StringImm(desc_buf.scope()),
- ],
+ [],
attrs=DictAttrs({}),
- ty=desc_buf.ty,
+ ty_args=[desc_buf.ty],
),
),
Evaluate(encode_call),
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
index 9b4540a2e3..c959d95e87 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
@@ -2190,12 +2190,15 @@ def _emit_plan(
# elem_offset. The latter is part of a flat DeclTensor definition;
# after loop unrolling, CSE may otherwise lift an offset containing a
# locally bound coordinate above that coordinate's Bind statement.
- smem_view = T.decl_tensor(
- (1,),
- spec.smem_buffer.dtype,
- spec.smem_buffer.data,
- elem_offset=0,
- scope=spec.smem_buffer.scope(),
+ smem_view = T.bind(
+ T.decl_tensor(
+ spec.smem_buffer.data,
+ ty_args=[
+ T.Tensor(
+ (1,), spec.smem_buffer.dtype, elem_offset=0, scope=spec.smem_buffer.scope()
+ )
+ ],
+ )
)
return smem_view.ptr_to([spec.smem_base_offset + element_offset])
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
index e3cad591ff..95f7058b71 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
@@ -355,12 +355,18 @@ def _make_views_meta(per_op_carved, per_thread_total):
iter strides at codegen time.
"""
return {
- op_br: T.decl_tensor(
- (per_thread_total,),
- op_br.source.dtype,
- op_br.source.data,
- scope="local",
- layout=per_op_carved[op_br],
+ op_br: T.bind(
+ T.decl_tensor(
+ op_br.source.data,
+ ty_args=[
+ T.Tensor(
+ (per_thread_total,),
+ op_br.source.dtype,
+ scope="local",
+ layout=per_op_carved[op_br],
+ )
+ ],
+ )
)
for op_br in per_op_carved
}
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
index 3911ec0088..cd8e4d1fa1 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
@@ -25,13 +25,10 @@ from tvm.backend.cuda.op import _tmem_address
from tvm.ir import (
Bind,
Call,
- DataTypeImm,
DictAttrs,
Evaluate,
SeqStmt,
- StringImm,
TensorRegion,
- Tuple,
)
from tvm.runtime import DataType
from tvm.script import tirx as T
@@ -1106,26 +1103,18 @@ def gemm_async_tcgen05_impl(op_call: TensorCall, sctx: DispatchContext) -> Funct
desc_lo,
Call(
"tirx.alloc_tensor",
- [
- Tuple(desc_lo.ty.shape),
- DataTypeImm(desc_lo.ty.dtype.dtype),
- StringImm(desc_lo.scope()),
- ],
+ [],
attrs=DictAttrs({}),
- ty=desc_lo.ty,
+ ty_args=[desc_lo.ty],
),
),
Bind(
desc_hi,
Call(
"tirx.alloc_tensor",
- [
- Tuple(desc_hi.ty.shape),
- DataTypeImm(desc_hi.ty.dtype.dtype),
- StringImm(desc_hi.scope()),
- ],
+ [],
attrs=DictAttrs({}),
- ty=desc_hi.ty,
+ ty_args=[desc_hi.ty],
),
),
Evaluate(unpack),
@@ -1154,13 +1143,9 @@ def gemm_async_tcgen05_impl(op_call: TensorCall, sctx: DispatchContext) -> Funct
desc_buf,
Call(
"tirx.alloc_tensor",
- [
- Tuple(desc_buf.ty.shape),
- DataTypeImm(desc_buf.ty.dtype.dtype),
- StringImm(desc_buf.scope()),
- ],
+ [],
attrs=DictAttrs({}),
- ty=desc_buf.ty,
+ ty_args=[desc_buf.ty],
),
),
Evaluate(encode_call),
diff --git a/python/tvm/backend/trn/tile_primitive/private_alloc.py b/python/tvm/backend/trn/tile_primitive/private_alloc.py
index 95ad4f2cb1..67bd99d9f1 100644
--- a/python/tvm/backend/trn/tile_primitive/private_alloc.py
+++ b/python/tvm/backend/trn/tile_primitive/private_alloc.py
@@ -18,7 +18,7 @@
from typing import Any
from tvm.backend.trn.tile_primitive.common import nki_dim
-from tvm.ir import Bind, Call, DataTypeImm, DictAttrs, Stmt, StringImm, Tuple
+from tvm.ir import Bind, Call, DictAttrs, Stmt, Tuple
from tvm.script import tirx as T
from tvm.tirx import DispatchContext, FloatImm, IntImm, Var
from tvm.tirx.op.tile import (
@@ -36,8 +36,8 @@ from tvm.tirx.tensor_instruction import TensorCall
def _allocation(tensor: Var, allocated_addr=None) -> Bind:
"""Keep private workspace placement on its producing allocation."""
- args = [Tuple(tensor.ty.shape), DataTypeImm(tensor.ty.dtype.dtype), StringImm(tensor.scope())]
- if allocated_addr:
+ args = []
+ if allocated_addr is not None:
args.append(Tuple(allocated_addr))
return Bind(
tensor,
@@ -45,7 +45,7 @@ def _allocation(tensor: Var, allocated_addr=None) -> Bind:
"tirx.alloc_tensor",
args,
attrs=DictAttrs({}),
- ty=tensor.ty,
+ ty_args=[tensor.ty],
),
)
diff --git a/python/tvm/backend/trn/transform/naive_allocator.py b/python/tvm/backend/trn/transform/naive_allocator.py
index afc26d5267..f109ea8309 100644
--- a/python/tvm/backend/trn/transform/naive_allocator.py
+++ b/python/tvm/backend/trn/transform/naive_allocator.py
@@ -58,12 +58,13 @@ def _get_alloc_pool_start(stmt) -> int:
if not isinstance(op.value, Call) or op.value.op != Op.get("tirx.alloc_tensor"):
return
tensor = op.var
- allocated_addr = op.value.args[3].fields if len(op.value.args) == 4 else []
+ allocated_addr = op.value.args[0].fields if op.value.args else []
if len(allocated_addr) == 0:
return
- shape = op.value.args[0].fields
- dtype = op.value.args[1].value
- scope = op.value.args[2].value
+ tensor_type = op.value.ty_args[0]
+ shape = tensor_type.shape
+ dtype = tensor_type.dtype.dtype
+ scope = tensor_type.storage_scope
tensor_size = get_tensor_size(tensor, shape, dtype, scope)
if tensor_size is None:
return
@@ -81,14 +82,14 @@ def _allocate_missing_tensors(stmt, alloc_pool_start: int):
if not isinstance(op.value, Call) or op.value.op != Op.get("tirx.alloc_tensor"):
return op
tensor = op.var
- shape = op.value.args[0].fields
- dtype = op.value.args[1].value
- scope = op.value.args[2].value
+ tensor_type = op.value.ty_args[0]
+ shape = tensor_type.shape
+ dtype = tensor_type.dtype.dtype
+ scope = tensor_type.storage_scope
tensor_size = get_tensor_size(tensor, shape, dtype, scope)
- allocated_addr = op.value.args[3].fields if len(op.value.args) == 4 else []
+ allocated_addr = op.value.args[0].fields if op.value.args else []
if len(allocated_addr) == 0 and tensor_size is not None:
- args = list(op.value.args[:3])
- args.append(Tuple([IntImm("int32", int(alloc_offset))]))
+ args = [Tuple([IntImm("int32", int(alloc_offset))])]
alloc_offset += tensor_size
return Bind(
tensor,
diff --git a/python/tvm/relax/backend/gpu_generic/cumsum.py b/python/tvm/relax/backend/gpu_generic/cumsum.py
index cf57a2e182..0f6a39de05 100644
--- a/python/tvm/relax/backend/gpu_generic/cumsum.py
+++ b/python/tvm/relax/backend/gpu_generic/cumsum.py
@@ -181,7 +181,7 @@ def gpu_2d_continuous_cumsum(
def cumsum(A: T.Tensor([m, n], dtype=in_dtype), Out: T.Tensor([m, n], dtype=out_dtype)):
T.func_attr({"tirx.is_scheduled": True}) # prevent further scheduling
- Tmp = T.alloc_tensor([m, n], dtype=out_dtype)
+ Tmp = T.alloc_tensor(ty_args=[T.Tensor([m, n], out_dtype)])
# LowerIntrin may implement signed FloorDiv using a sign-bit shift. Keep
# hierarchy counting division-free so WebGPU can narrow indices to int32.
total_rounds: T.let[T.int64] = _get_total_rounds(n, log_block_n, index_bits)
diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py b/python/tvm/relax/frontend/nn/llm/position_embedding.py
index 880d3dd325..181db2a62f 100644
--- a/python/tvm/relax/frontend/nn/llm/position_embedding.py
+++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py
@@ -578,12 +578,20 @@ def llama_rope_with_position_map( # pylint: disable=too-many-arguments
)
# long factors is the first half, short factors is the second half
- long_factors = T.decl_tensor((rotary_dim // 2,), "float32", data=ext_factors.data)
+ long_factors = T.decl_tensor(
+ ext_factors.data,
+ ty_args=[T.Tensor((rotary_dim // 2,), "float32", data=ext_factors.data)],
+ )
short_factors = T.decl_tensor(
- (rotary_dim // 2,),
- "float32",
- data=ext_factors.data,
- elem_offset=(rotary_dim // 2),
+ ext_factors.data,
+ ty_args=[
+ T.Tensor(
+ (rotary_dim // 2,),
+ "float32",
+ elem_offset=rotary_dim // 2,
+ data=ext_factors.data,
+ )
+ ],
)
if seq_len > original_max_position_embeddings:
@@ -795,12 +803,20 @@ def llama4_rope_with_position_map( # pylint: disable=too-many-arguments
)
# long factors is the first half, short factors is the second half
- long_factors = T.decl_tensor((rotary_dim // 2,), "float32", data=ext_factors.data)
+ long_factors = T.decl_tensor(
+ ext_factors.data,
+ ty_args=[T.Tensor((rotary_dim // 2,), "float32", data=ext_factors.data)],
+ )
short_factors = T.decl_tensor(
- (rotary_dim // 2,),
- "float32",
- data=ext_factors.data,
- elem_offset=(rotary_dim // 2),
+ ext_factors.data,
+ ty_args=[
+ T.Tensor(
+ (rotary_dim // 2,),
+ "float32",
+ elem_offset=rotary_dim // 2,
+ data=ext_factors.data,
+ )
+ ],
)
if seq_len > original_max_position_embeddings:
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index 23f1a08911..461af96840 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -8343,8 +8343,13 @@ def _build_tflite_rfft2d_function(input_shape, output_pair_shape):
# Flat 1D aliases of the multi-dim buffers. The kernel is rank-agnostic
# over the leading batch dimensions, so collapsing the index space
# avoids special-casing 2D / 3D / 4D input shapes.
- data_flat = T.decl_tensor((input_total,), "float32", data=data.data)
- output_flat = T.decl_tensor((output_complex_total * 2,), "float32", data=output.data)
+ data_flat = T.decl_tensor(
+ data.data, ty_args=[T.Tensor((input_total,), "float32", data=data.data)]
+ )
+ output_flat = T.decl_tensor(
+ output.data,
+ ty_args=[T.Tensor((output_complex_total * 2,), "float32", data=output.data)],
+ )
neg_two_pi_const = T.float32(neg_two_pi)
for b_idx, out_y, out_x in T.grid(batch, height, out_width):
@@ -8515,10 +8520,10 @@ def _build_tflite_rfft2d_fft_function(input_shape, output_pair_shape):
f" data: T.Tensor({tuple(int(x) for x in input_shape)}, 'float32'),\n"
f" output: T.Tensor({tuple(int(x) for x in output_pair_shape)}, 'float32'),\n"
"):\n"
- f" data_flat = T.decl_tensor(({input_total},), 'float32', data=data.data)\n"
- f" output_flat = T.decl_tensor(({output_complex_total * 2},), 'float32', data=output.data)\n"
- f" scratch_real = T.decl_tensor(({input_total},), 'float32')\n"
- f" scratch_imag = T.decl_tensor(({input_total},), 'float32')\n"
+ f" data_flat = T.decl_tensor(data.data, ty_args=[T.Tensor(({input_total},), 'float32')])\n"
+ f" output_flat = T.decl_tensor(output.data, ty_args=[T.Tensor(({output_complex_total * 2},), 'float32')])\n"
+ f" scratch_real = T.alloc_tensor(ty_args=[T.Tensor(({input_total},), 'float32')])\n"
+ f" scratch_imag = T.alloc_tensor(ty_args=[T.Tensor(({input_total},), 'float32')])\n"
f" for b_idx in T.serial({batch}):\n"
f" with Ts.sblock('rfft2d_fft'):\n"
f" v_b = Ts.axis.remap('S', [b_idx])\n"
@@ -8635,10 +8640,12 @@ def _build_stablehlo_rng_bit_generator_function(algorithm, state_len, out_dtype,
key_1 = _u32(state_key >> T.uint64(32))
output_state[0] = state_key
output_state[1] = state_counter + T.uint64(num_blocks)
- out_flat = T.decl_tensor((total,), out_dtype, data=output.data)
- keys = T.decl_tensor((3,), "uint32", scope="local")
- rotations = T.decl_tensor((8,), "uint32", scope="local")
- ctr = T.decl_tensor((2,), "uint32", scope="local")
+ out_flat = T.decl_tensor(
+ output.data, ty_args=[T.Tensor((total,), out_dtype, data=output.data)]
+ )
+ keys = T.alloc_tensor(ty_args=[T.Tensor((3,), "uint32", scope="local")])
+ rotations = T.alloc_tensor(ty_args=[T.Tensor((8,), "uint32", scope="local")])
+ ctr = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint32", scope="local")])
keys[0] = key_0
keys[1] = key_1
keys[2] = key_0 ^ key_1 ^ T.uint32(parity)
@@ -8682,10 +8689,12 @@ def _build_stablehlo_rng_bit_generator_function(algorithm, state_len, out_dtype,
key_1 = _u32(state_key >> T.uint64(32))
output_state[0] = state_key
output_state[1] = state_counter + T.uint64(num_blocks)
- out_flat = T.decl_tensor((total,), out_dtype, data=output.data)
- ctr = T.decl_tensor((4,), "uint32", scope="local")
- keys = T.decl_tensor((2,), "uint32", scope="local")
- high_ctr = T.decl_tensor((2,), "uint32", scope="local")
+ out_flat = T.decl_tensor(
+ output.data, ty_args=[T.Tensor((total,), out_dtype, data=output.data)]
+ )
+ ctr = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint32", scope="local")])
+ keys = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint32", scope="local")])
+ high_ctr = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint32", scope="local")])
if state_len == 3:
# PHILOX u64[3]: the third state word feeds the high counter and
# is passed through to the output state unchanged.
diff --git a/python/tvm/relax/transform/legalize_ops/inspect_op.py b/python/tvm/relax/transform/legalize_ops/inspect_op.py
index 9794bcb0fd..94ef5ba4c9 100644
--- a/python/tvm/relax/transform/legalize_ops/inspect_op.py
+++ b/python/tvm/relax/transform/legalize_ops/inspect_op.py
@@ -73,9 +73,9 @@ def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
shape_ptr: T.let[T.handle("int64")] = T.abi_field_get(
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorShape), ty=T.handle("int64").ty
)
- shape = T.decl_tensor(ndim, "int64", data=shape_ptr)
+ shape = T.decl_tensor(shape_ptr, ty_args=[T.Tensor(ndim, "int64", data=shape_ptr)])
- product = T.decl_tensor([], "int64")
+ product = T.alloc_tensor(ty_args=[T.Tensor([], "int64")])
product[()] = 1
# TODO(Lunderberg): Add a TIR lowering pass to allow
@@ -87,7 +87,7 @@ def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
return product[()]
else:
- strides = T.decl_tensor(ndim, "int64", data=stride_ptr)
+ strides = T.decl_tensor(stride_ptr, ty_args=[T.Tensor(ndim, "int64", data=stride_ptr)])
stride: T.let[T.int64] = strides[axis]
return stride
diff --git a/python/tvm/script/printer/scriptable.py b/python/tvm/script/printer/scriptable.py
index e8e089a895..c3bf9f144a 100644
--- a/python/tvm/script/printer/scriptable.py
+++ b/python/tvm/script/printer/scriptable.py
@@ -35,10 +35,6 @@ class PrinterConfig(Object):
The ``int_dtype`` and ``float_dtype`` arguments populate the string-valued
``ir.int_dtype`` and ``ir.float_dtype`` entries in ``extra_config``.
Explicit entries in ``extra_config`` override those arguments.
-
- ``tirx.scalar_tensor_as_mutable_var`` in ``extra_config`` defaults to true.
- It prints eligible local scalar buffers as typed mutable variables. Set it
- to false to retain explicit buffer allocations and indexed loads/stores.
"""
binding_names: Sequence[str]
diff --git a/python/tvm/tirx/_tensor_view.py b/python/tvm/tirx/_tensor_view.py
index a02eb0c772..e810e5429a 100644
--- a/python/tvm/tirx/_tensor_view.py
+++ b/python/tvm/tirx/_tensor_view.py
@@ -53,18 +53,21 @@ def _redecl(buf: Var, shape, layout, *, dtype=None, elem_offset=None, addr_offse
layout=layout,
)
return bind_(tvm.ir.Call("tirx.cuda.decl_tmem", [addr], ty_args=[ty]))
- return tvm.tirx.script.ir_builder.decl_tensor(
+ from tvm.tirx.script.ir_builder.parser_protocol import ( # pylint: disable=import-outside-toplevel
+ bind_,
+ )
+
+ tensor_type = tvm.tirx.script.ir_builder.Tensor(
shape,
buf.dtype if dtype is None else dtype,
- buf.data,
- buf.strides,
- buf.elem_offset if elem_offset is None else elem_offset,
- None,
- buf.scope(),
- buf.data_alignment,
- buf.offset_factor,
- layout,
+ strides=buf.strides,
+ elem_offset=buf.elem_offset if elem_offset is None else elem_offset,
+ scope=buf.scope(),
+ align=buf.data_alignment,
+ offset_factor=buf.offset_factor,
+ layout=layout,
)
+ return bind_(tvm.tirx.script.ir_builder.decl_tensor(buf.data, ty_args=[tensor_type]))
def view(buf: Var, *args, **kwargs) -> Var:
diff --git a/python/tvm/tirx/bench.py b/python/tvm/tirx/bench.py
index ffc2686ee0..6cdf8776e5 100644
--- a/python/tvm/tirx/bench.py
+++ b/python/tvm/tirx/bench.py
@@ -940,8 +940,12 @@ class CudaProfiler:
# Assume Expr-like input; use as-is
self.profiler_enabled = profiler_enabled # type: ignore[assignment]
- self.profiler_tag = T.alloc_tensor([1], "uint64", scope="local", align=8)
- self.profiler_write_offset = T.alloc_tensor([1], "uint32", scope="local", align=8)
+ self.profiler_tag = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "uint64", scope="local", align=8)])
+ )
+ self.profiler_write_offset = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "uint32", scope="local", align=8)])
+ )
def _leader(self, leader: None | tvm.tirx.Expr | bool):
if leader is not None:
diff --git a/python/tvm/tirx/script/ir_builder/ir.py b/python/tvm/tirx/script/ir_builder/ir.py
index a683796180..bf4ae8dfd4 100644
--- a/python/tvm/tirx/script/ir_builder/ir.py
+++ b/python/tvm/tirx/script/ir_builder/ir.py
@@ -24,7 +24,6 @@ from collections.abc import Callable
from numbers import Integral
from typing import TYPE_CHECKING, Any, TypeVar
-import tvm
from tvm import ir as _ir
from tvm.script.ir_builder.base import annotation_constructor as _annotation_constructor
from tvm.script.ir_builder.base import at as _at
@@ -310,92 +309,72 @@ def elected():
)
-@_register_mutable_decl("tirx.alloc_tensor")
def alloc_tensor(
- shape: list[Expr] | tuple[Expr] | Expr | Integral,
- dtype: str = "float32",
- data: Var | None = None,
- strides: list[Expr] | None = None,
- elem_offset: Expr | None = None,
- byte_offset: Expr | None = None,
- scope: str = "global",
- align: int = -1,
- offset_factor: int = 0,
- layout: str | Layout | None = MISSING,
- allocated_addr: int | tuple[int, ...] | None = None,
+ placement=None,
+ *,
+ ty_args,
+ attrs=None,
annotations: dict[str, Any] | None = None,
-) -> Var:
- """Allocate a tensor and return its variable.
+ loc: Location = UNKNOWN_LOC,
+ ty=None,
+) -> ir.Call:
+ """Allocate the tensor described by one explicit ``Tensor`` type argument.
+
+ Optional placement is an ordinary tuple operand. Allocation annotations
+ belong to ``attrs``; ``annotations`` accepts a Python dictionary whose
+ scalar values are normalized to typed IR constants.
+ """
+ if attrs is not None and annotations is not None:
+ raise TypeError("alloc_tensor cannot mix attrs and annotations")
+
+ def normalize(value):
+ if isinstance(value, bool):
+ return tir.IntImm("bool", int(value))
+ if isinstance(value, int):
+ return tir.IntImm("int32", value)
+ if isinstance(value, float):
+ return tir.FloatImm("float32", value)
+ return value
- Emits a Bind statement with a ``tirx.alloc_tensor`` Call::
+ if attrs is None:
+ normalized = {key: normalize(value) for key, value in (annotations or {}).items()}
+ if "tensor_allocated_addr" in normalized:
+ if placement is not None:
+ raise ValueError("Allocation placement was specified twice")
+ placement = normalized.pop("tensor_allocated_addr")
+ attrs = ir.DictAttrs(normalized)
+ args = []
+ if placement is not None:
+ if not isinstance(placement, ir.Tuple):
+ placement = placement if isinstance(placement, list | tuple | ir.Array) else [placement]
+ placement = ir.Tuple([normalize(value) for value in placement])
+ args.append(placement)
+ return ir.Call("tirx.alloc_tensor", args, attrs=attrs, ty_args=ty_args, loc=loc, ty=ty)
- buf = T.alloc_tensor((128, 128))
- Parameters
- ----------
- shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
- The shape of the tensor to allocate.
- dtype : str
- The data type of the tensor elements.
- scope : str
- The storage scope of the tensor (e.g., "global", "shared").
- data : Optional[Var]
- Optional explicit data pointer.
- strides : Optional[List[Expr]]
- Optional strides.
- elem_offset : Optional[Expr]
- Optional element offset.
- byte_offset : Optional[Expr]
- Optional byte offset.
- align : int
- Alignment requirement in bytes.
- offset_factor : int
- Offset factor.
- layout : Optional[Union[str, Layout]]
- Optional layout.
- allocated_addr : Optional[Union[int, Tuple[int, ...]]]
- Optional pre-allocated address metadata.
- annotations : Optional[Dict[str, Any]]
- Optional annotations for the allocation.
+def _bind_tensor(value):
+ from tvm.script.ir_builder.stmt import bind # pylint: disable=import-outside-toplevel
- Returns
- -------
- res : Var
- The allocated tensor.
- """
- shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral) else shape
- buf = Var(
- "",
- _tensor_type(
- shape=shape,
- dtype=dtype,
- data=data,
- strides=strides,
- elem_offset=elem_offset,
- byte_offset=byte_offset,
- scope=scope,
- align=align,
- offset_factor=offset_factor,
- layout=layout,
- ),
- )
- _record_meta_resource(buf, skip_frames=2)
-
- # The allocation call annotations hold typed IR values. The C++ side stores
- # alignment / shape-like ints as ``IntImm(int32, ...)``; if the user
- # (or a parsed-source round-trip) passes a bare Python int, normalize
- # it so structural equality is preserved against the LowerOpaqueBlock
- # output. Booleans must stay as IntImm("bool", ...).
- def _normalize_ann_value(v):
- if isinstance(v, bool):
- return tir.IntImm("bool", int(v))
- if isinstance(v, int):
- return tir.IntImm("int32", v)
- if isinstance(v, float):
- return tir.FloatImm("float32", v)
- return v
-
- norm_annotations = {k: _normalize_ann_value(v) for k, v in (annotations or {}).items()}
+ tensor = bind(value)
+ _record_meta_resource(tensor, skip_frames=2)
+ return tensor
+
+
+def _alloc_tensor(
+ shape,
+ dtype="float32",
+ data=None,
+ strides=None,
+ elem_offset=None,
+ byte_offset=None,
+ scope="global",
+ align=-1,
+ offset_factor=0,
+ layout=MISSING,
+ allocated_addr=None,
+ annotations=None,
+):
+ """Implement separately named allocation conveniences using the canonical Op."""
if allocated_addr is not None:
if scope in ("", "global", "shared", "shared.dyn", "local"):
raise ValueError(f"{scope} does not support allocation placement")
@@ -403,24 +382,12 @@ def alloc_tensor(
raise ValueError(
"Allocation placement cannot be combined with pointer or offset metadata"
)
- addresses = allocated_addr if isinstance(allocated_addr, list | tuple) else [allocated_addr]
- allocated_addr = [_normalize_ann_value(v) for v in addresses]
- # Normalize the historical annotation spelling at the construction boundary.
- if "tensor_allocated_addr" in norm_annotations:
- if allocated_addr is not None:
- raise ValueError("Allocation placement was specified twice")
- allocated_addr = norm_annotations.pop("tensor_allocated_addr")
- args = [ir.Tuple(buf.shape), ir.DataTypeImm(DataType(buf.dtype)), ir.StringImm(buf.scope())]
- if allocated_addr is not None:
- args.append(ir.Tuple(allocated_addr))
- allocation = ir.Call(
- "tirx.alloc_tensor",
- args,
- attrs=ir.DictAttrs(norm_annotations),
- ty=buf.ty,
+ tensor_type = _tensor_type(
+ shape, dtype, data, strides, elem_offset, byte_offset, scope, align, offset_factor, layout
+ )
+ return _bind_tensor(
+ alloc_tensor(allocated_addr, ty_args=[tensor_type], annotations=annotations)
)
- _ffi_api.AddToParent(tvm.ir.Bind(buf, allocation))
- return buf
def wg_reg_tile(elem_per_thread: int, dtype: str = "float32") -> Var:
@@ -428,16 +395,15 @@ def wg_reg_tile(elem_per_thread: int, dtype: str = "float32") -> Var:
Sugar for the recurring pattern::
- T.alloc_tensor(
+ T.alloc_tensor(ty_args=[T.Tensor(
(128, elem_per_thread), dtype,
- layout=wg_local_layout(elem_per_thread),
- scope="local",
- )
+ layout=wg_local_layout(elem_per_thread), scope="local",
+ )])
Used to stage a tcgen05 load: each of the 128 threads in a warpgroup
owns one row of ``elem_per_thread`` contiguous elements.
"""
- return alloc_tensor(
+ return _alloc_tensor(
(128, elem_per_thread),
dtype,
layout=wg_local_layout(elem_per_thread),
@@ -532,96 +498,29 @@ class DtypeConstructor:
return f"DtypeConstructor({self._dtype_str!r})"
-@_register_mutable_decl("tirx.decl_tensor")
def decl_tensor(
- shape,
- dtype="float32",
- data=None,
- strides=None,
- elem_offset=None,
- byte_offset=None,
- scope="global",
- align=0,
- offset_factor=0,
- layout=MISSING,
- allocated_addr=None,
-) -> Var:
- """Declare a tensor backed by a pointer, or allocate its storage.
-
- With ``data``, bind a ``tirx.decl_tensor`` Call to the tensor variable.
- Without ``data``, bind a ``tirx.alloc_tensor`` Call instead. The ``tmem``
- scope uses ``allocated_addr`` to declare externally allocated tensor memory.
-
- Parameters
- ----------
- shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
- The type of the tensor prior to flattening.
-
- dtype : str
- The data type in the content of the tensor.
-
- data : Var
- The pointer to the head of the data.
-
- strides : List[Expr]
- The strides of each dimension.
-
- elem_offset : Expr
- The offset in terms of number of dtype elements (including lanes).
-
- byte_offset : Expr
- The offset in terms of number of bytes.
-
- scope : str
- The optional storage scope of tensor data pointer.
-
- align : int
- The alignment requirement of data pointer in bytes.
-
- offset_factor : int
- The factor of elem_offset field.
-
- layout : Layout
- The layout of the tensor.
-
- Returns
- -------
- res : Var
- The declared tensor.
- """
- shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral) else shape
- shape = tuple(shape)
- if strides is None:
- strides = []
- dtype = _normalize_prim_type(dtype)
- buf = _ffi_api.DeclTensor( # type: ignore[attr-defined] # pylint: disable=no-member
- shape,
- dtype,
- "",
- data,
- strides,
- _get_elem_offset(elem_offset, byte_offset, dtype),
- scope,
- align,
- offset_factor,
- _get_layout(layout, shape, scope),
- allocated_addr,
- )
- _record_meta_resource(buf, skip_frames=2)
- return buf
+ data,
+ *,
+ ty_args,
+ attrs=None,
+ loc: Location = UNKNOWN_LOC,
+ ty=None,
+) -> ir.Call:
+ """Declare a pointer-backed tensor with one explicit ``Tensor`` type argument."""
+ return ir.Call("tirx.decl_tensor", [data], attrs=attrs, ty_args=ty_args, loc=loc, ty=ty)
-alloc_shared = functools.partial(alloc_tensor, scope="shared")
+alloc_shared = functools.partial(_alloc_tensor, scope="shared")
_register_mutable_decl("tirx.alloc_shared")(alloc_shared)
-alloc_local = functools.partial(alloc_tensor, scope="local")
+alloc_local = functools.partial(_alloc_tensor, scope="local")
_register_mutable_decl("tirx.alloc_local")(alloc_local)
smem = _register_mutable_decl("tirx.smem")(alloc_shared)
-tmem = functools.partial(alloc_tensor, scope="tmem")
+tmem = functools.partial(_alloc_tensor, scope="tmem")
def alloc_tcgen05_ldst_frag(instr_shape, tensor_shape, dtype):
@@ -721,7 +620,7 @@ def alloc_scalar(
annotations: dict[str, Any] | None = None,
) -> _ir.TensorLoad:
"""Allocate a zero-dimensional tensor (scalar), with optional allocation annotations."""
- buf = alloc_tensor(
+ buf = _alloc_tensor(
shape=(1,), dtype=dtype, scope=scope, layout=TileLayout(S[1]), annotations=annotations
)
assert is_tensor_var(buf)
@@ -732,16 +631,20 @@ def alloc_scalar(
@_register_mutable_decl("tirx.decl_scalar")
def decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) -> _ir.TensorLoad:
"""Declare a zero-dimensional tensor (scalar) from a pointer."""
- buf = decl_tensor(
- shape=(1,),
- dtype=dtype,
- data=data,
- scope=scope,
- elem_offset=_get_elem_offset(elem_offset, byte_offset, dtype),
- strides=None,
- align=-1,
- offset_factor=0,
- layout=TileLayout(S[1]),
+ buf = _bind_tensor(
+ decl_tensor(
+ data,
+ ty_args=[
+ _tensor_type(
+ (1,),
+ dtype,
+ scope=scope,
+ elem_offset=_get_elem_offset(elem_offset, byte_offset, dtype),
+ align=-1,
+ layout=TileLayout(S[1]),
+ )
+ ],
+ )
)
assert is_tensor_var(buf)
scalar = buf[0]
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index 11950c102e..6b6ee83252 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -48,7 +48,7 @@ from tvm.script.ir_builder.stmt import (
assert_ as assert_,
)
from tvm.script.ir_builder.stmt import (
- bind as bind,
+ bind as _bind,
)
from tvm.script.ir_builder.stmt import (
break_ as break_,
@@ -235,6 +235,14 @@ def _name(value: Any, name: str | None, loc: _Loc) -> Any:
return _base.at_(loc, value)
+def bind(value, type_annotation=None, *, var=None):
+ """Bind an expression and track tensor resources created in meta classes."""
+ result = _bind(value, type_annotation, var=var)
+ if _tir.is_tensor_var(result):
+ _native._record_meta_resource(result)
+ return result
+
+
def bind_(
value: Any = _base.MISSING,
*,
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
index ebec0fa575..16b50a8cbf 100644
--- a/python/tvm/topi/gpu/scan.py
+++ b/python/tvm/topi/gpu/scan.py
@@ -132,9 +132,15 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
width = 2 << l2_width
blocks_per_batch = cast(ceil_div(scan_axis_size, max_threads * width), "int32")
- start_buf = T.decl_tensor([1], "int32", scope="local")
- middle_buf = T.decl_tensor([1], "int32", scope="local")
- end_buf = T.decl_tensor([1], "int32", scope="local")
+ start_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")])
+ )
+ middle_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")])
+ )
+ end_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")])
+ )
with T.frame_scope(
[
T.launch_thread("threadIdx.x", nthread_tx),
@@ -219,10 +225,18 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
width = 2 << (lim - l2_width - 1)
blocks_per_batch = cast(ceil_div(scan_axis_size, max_threads * width), "int32")
- start_buf = T.decl_tensor([1], "int32", scope="local")
- middle_buf = T.decl_tensor([1], "int32", scope="local")
- end_buf = T.decl_tensor([1], "int32", scope="local")
- tmp_buf = T.decl_tensor([1], out_dtype, scope="local")
+ start_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")])
+ )
+ middle_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")])
+ )
+ end_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")])
+ )
+ tmp_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], out_dtype, scope="local")])
+ )
with T.frame_scope(
[
T.launch_thread("threadIdx.x", nthread_tx),
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index a0296aea01..5f8bf69c44 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -111,13 +111,19 @@ def _odd_even_sort(
start = bx * block_size
# Tensor declarations without data allocate their storage.
- tmp_keys_swap = T.decl_tensor([block_size], keys_swap.dtype, scope="shared")
- temp_keys = T.decl_tensor([1], keys_swap.dtype, scope="local")
- temp_cond1 = T.decl_tensor([1], keys_swap.dtype, scope="local")
- temp_cond2 = T.decl_tensor([1], keys_swap.dtype, scope="local")
+ tmp_keys_swap = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([block_size], keys_swap.dtype, scope="shared")])
+ )
+ temp_keys = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], keys_swap.dtype, scope="local")]))
+ temp_cond1 = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], keys_swap.dtype, scope="local")]))
+ temp_cond2 = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], keys_swap.dtype, scope="local")]))
if values_swap is not None:
- tmp_values_swap = T.decl_tensor([block_size], values_swap.dtype, scope="shared")
- temp_values = T.decl_tensor([1], values_swap.dtype, scope="local")
+ tmp_values_swap = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([block_size], values_swap.dtype, scope="shared")])
+ )
+ temp_values = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], values_swap.dtype, scope="local")])
+ )
# Copy data to scratch space
base_idx = by_val * size * axis_mul_after + bz
@@ -399,10 +405,10 @@ def _sort_common(
step_count,
even,
):
- first_buf = T.decl_tensor([1], target_dtype, scope="local")
- last_buf = T.decl_tensor([1], target_dtype, scope="local")
- i_buf_buf = T.decl_tensor([1], target_dtype, scope="local")
- j_buf_buf = T.decl_tensor([1], target_dtype, scope="local")
+ first_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")]))
+ last_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")]))
+ i_buf_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")]))
+ j_buf_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")]))
first = first_buf
last = last_buf
i_buf = i_buf_buf
@@ -466,12 +472,16 @@ def _sort_common(
step_count,
even,
):
- outer_first_buf = T.decl_tensor([1], target_dtype, scope="local")
- outer_last_buf = T.decl_tensor([1], target_dtype, scope="local")
- first_buf = T.decl_tensor([1], target_dtype, scope="local")
- last_buf = T.decl_tensor([1], target_dtype, scope="local")
- i_buf_buf = T.decl_tensor([1], target_dtype, scope="local")
- j_buf_buf = T.decl_tensor([1], target_dtype, scope="local")
+ outer_first_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")])
+ )
+ outer_last_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")])
+ )
+ first_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")]))
+ last_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")]))
+ i_buf_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")]))
+ j_buf_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype, scope="local")]))
outer_first = outer_first_buf
outer_last = outer_last_buf
first = first_buf
diff --git a/python/tvm/topi/searchsorted.py b/python/tvm/topi/searchsorted.py
index 39b4b219a8..561bedd2d9 100644
--- a/python/tvm/topi/searchsorted.py
+++ b/python/tvm/topi/searchsorted.py
@@ -38,8 +38,8 @@ def binary_search(sequence_offset, search_range, sorted_sequence, value, right,
Note that we index N-D Buffer by 1-D linearlized indices.
"""
- lo_buf = T.decl_tensor([1], out_dtype, scope="local")
- hi_buf = T.decl_tensor([1], out_dtype, scope="local")
+ lo_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], out_dtype, scope="local")]))
+ hi_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], out_dtype, scope="local")]))
lo = lo_buf
hi = hi_buf
diff --git a/python/tvm/topi/sparse_reshape.py b/python/tvm/topi/sparse_reshape.py
index 97672dea8b..dd48ffde7b 100644
--- a/python/tvm/topi/sparse_reshape.py
+++ b/python/tvm/topi/sparse_reshape.py
@@ -89,19 +89,35 @@ def sparse_reshape(
prev_shape_size = prev_shape_ptr.shape[0]
new_shape_size = new_shape_ptr.shape[0]
- multipliers_buf = T.alloc_tensor([prev_shape_size], new_shape_ptr.dtype, scope="local")
+ multipliers_buf = T.bind(
+ T.alloc_tensor(
+ ty_args=[T.Tensor([prev_shape_size], new_shape_ptr.dtype, scope="local")]
+ )
+ )
multipliers = multipliers_buf
- dividers_buf = T.alloc_tensor([new_shape_size], new_shape_ptr.dtype, scope="local")
+ dividers_buf = T.bind(
+ T.alloc_tensor(
+ ty_args=[T.Tensor([new_shape_size], new_shape_ptr.dtype, scope="local")]
+ )
+ )
dividers = dividers_buf
- flattened_indices_buf = T.alloc_tensor(
- [sparse_indices_ptr.shape[0]], new_shape_ptr.dtype, scope="local"
+ flattened_indices_buf = T.bind(
+ T.alloc_tensor(
+ ty_args=[
+ T.Tensor([sparse_indices_ptr.shape[0]], new_shape_ptr.dtype, scope="local")
+ ]
+ )
)
flattened_indices = flattened_indices_buf
- total_ele_buf = T.alloc_tensor([1], new_shape_ptr.dtype, scope="local")
+ total_ele_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], new_shape_ptr.dtype, scope="local")])
+ )
total_ele = total_ele_buf
- division_total_ele_buf = T.alloc_tensor([1], new_shape_ptr.dtype, scope="local")
+ division_total_ele_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor([1], new_shape_ptr.dtype, scope="local")])
+ )
division_total_ele = division_total_ele_buf
- equal_shape_buf = T.alloc_tensor([1], "bool", scope="local")
+ equal_shape_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], "bool", scope="local")]))
equal_shape = equal_shape_buf
T.tensor_store(
@@ -227,8 +243,10 @@ def sparse_reshape(
)
with T.parallel(0, new_sparse_indices_ptr.shape[0]) as i:
- current_element_buf = T.alloc_tensor(
- [1], new_shape_ptr.dtype, scope="local"
+ current_element_buf = T.bind(
+ T.alloc_tensor(
+ ty_args=[T.Tensor([1], new_shape_ptr.dtype, scope="local")]
+ )
)
current_element = current_element_buf
T.tensor_store(
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
index eb18bcd937..4f72d434b9 100644
--- a/python/tvm/topi/vision/nms.py
+++ b/python/tvm/topi/vision/nms.py
@@ -212,7 +212,7 @@ def _classic_nms_ir(
with T.parallel(0, batch_size) as i:
# Step 1: Reorder data by sorted score
- nkeep_buf = T.alloc_tensor((1,), "int32", scope="local")
+ nkeep_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")]))
nkeep_local = nkeep_buf
T.tensor_store(
nkeep_local,
@@ -251,16 +251,22 @@ def _classic_nms_ir(
T.tensor_store(out_box_indices, (i, j), T.int32(-1))
# Step 2: Apply NMS - greedy suppression
- num_valid_boxes_buf = T.alloc_tensor((1,), "int32", scope="local")
+ num_valid_boxes_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
+ )
num_valid_boxes = num_valid_boxes_buf
T.tensor_store(num_valid_boxes, T.tensor_indices(num_valid_boxes, 0), T.int32(0))
- best_idx_buf = T.alloc_tensor((1,), "int32", scope="local")
+ best_idx_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")]))
best_idx = best_idx_buf
- best_score_buf = T.alloc_tensor((1,), data.dtype, scope="local")
+ best_score_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor((1,), data.dtype, scope="local")])
+ )
best_score = best_score_buf
- tmp_idx_buf = T.alloc_tensor((1,), "int32", scope="local")
+ tmp_idx_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")]))
tmp_idx = tmp_idx_buf
- tmp_val_buf = T.alloc_tensor((1,), data.dtype, scope="local")
+ tmp_val_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor((1,), data.dtype, scope="local")])
+ )
tmp_val = tmp_val_buf
zero = tvm.tirx.Cast(data.dtype, T.float32(0.0))
@@ -614,7 +620,9 @@ def _classic_nms_ir(
)
if return_indices:
- valid_idx_buf = T.alloc_tensor((1,), "int32", scope="local")
+ valid_idx_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
+ )
valid_idx = valid_idx_buf
T.tensor_store(valid_idx, T.tensor_indices(valid_idx, 0), T.int32(0))
@@ -851,7 +859,9 @@ def _rearrange_out(data, batch_size, num_anchors, box_data_length, score_index):
out = outs[0]
with T.parallel(0, batch_size) as i:
- valid_idx_buf = T.alloc_tensor((1,), "int32", scope="local")
+ valid_idx_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
+ )
valid_idx = valid_idx_buf
T.tensor_store(valid_idx, T.tensor_indices(valid_idx, 0), T.int32(0))
@@ -946,7 +956,9 @@ def _nms_loop(
with T.if_(tvm.tirx.all(iou_threshold > te.const(0), valid_count[i] > te.const(0))):
with T.then_():
- num_valid_boxes_local_buf = T.alloc_tensor((1,), "int32", scope="local")
+ num_valid_boxes_local_buf = T.bind(
+ T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
+ )
num_valid_boxes_local = num_valid_boxes_local_buf
T.tensor_store(
num_valid_boxes_local, T.tensor_indices(num_valid_boxes_local, 0), T.int32(0)
diff --git a/python/tvm/topi/vision/nms_util.py b/python/tvm/topi/vision/nms_util.py
index 684b0995c5..797f6cdbe0 100644
--- a/python/tvm/topi/vision/nms_util.py
+++ b/python/tvm/topi/vision/nms_util.py
@@ -67,8 +67,8 @@ def binary_search(y, num_boxes, scores, score_threshold, out):
Must be called within an IRBuilder context.
"""
- lo_buf = T.decl_tensor([1], "int32", scope="local")
- hi_buf = T.decl_tensor([1], "int32", scope="local")
+ lo_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")]))
+ hi_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")]))
lo = lo_buf
hi = hi_buf
T.tensor_store(lo, T.tensor_indices(lo, 0), T.int32(0))
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index fde9818cfc..d31cbec593 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1373,9 +1373,9 @@ void CodeGenCUDA::Dispatch_(const BindNode* op) {
}
void CodeGenCUDA::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.value().as_or_throw<DictAttrs>();
TVM_FFI_ICHECK(buffer.defined());
diff --git a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
index 3aeddc8122..f218ebeadf 100644
--- a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
+++ b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
@@ -90,9 +90,9 @@ class CodeGenNVPTX : public CodeGenLLVM {
}
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.value().as_or_throw<DictAttrs>();
llvm::Value* buf = nullptr;
diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc
index 4637c02e92..4e33677960 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -377,9 +377,9 @@ void CodeGenMetal::Dispatch_(const BindNode* op) {
}
void CodeGenMetal::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.value().as_or_throw<DictAttrs>();
TVM_FFI_ICHECK(buffer.defined());
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc b/src/backend/opencl/codegen/codegen_opencl.cc
index f6e96ef139..7d040e106f 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -464,8 +464,8 @@ void CodeGenOpenCL::Dispatch_(const BindNode* op) {
}
void CodeGenOpenCL::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
// Compute constant_size from buffer shape
size_t constant_size = 1;
diff --git a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
index c4036cd0bd..b6c637562f 100644
--- a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
+++ b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
@@ -107,9 +107,9 @@ class CodeGenAMDGPU : public CodeGenLLVM {
}
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.value().as_or_throw<DictAttrs>();
llvm::Value* buf = nullptr;
diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc
index 682709b716..87aa5e273e 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -229,9 +229,9 @@ void CodeGenTrainium::Dispatch_(const BindNode* op) {
}
void CodeGenTrainium::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
TVM_FFI_ICHECK(buffer.defined());
std::string vid = AllocVarID(buffer.get(), buffer.name() + "_ptr");
@@ -251,8 +251,8 @@ void CodeGenTrainium::DispatchAllocTensor(const BindNode* op, const CallNode* bu
<< ", dtype=" << dtype_str << ", buffer=";
}
Array<PrimExpr> addr;
- if (buffer_call->args.size() == 4) {
- addr = buffer_call->args[3].as_or_throw<tvm::Tuple>()->fields.as_or_throw<Array<PrimExpr>>();
+ if (!buffer_call->args.empty()) {
+ addr = buffer_call->args[0].as_or_throw<tvm::Tuple>()->fields.as_or_throw<Array<PrimExpr>>();
}
if (addr.empty()) {
stream << GetStorageScopeStr(scope) << ")\n";
@@ -260,9 +260,9 @@ void CodeGenTrainium::DispatchAllocTensor(const BindNode* op, const CallNode* bu
if (scope == "trn.psum") {
TVM_FFI_ICHECK(addr.size() == 2);
TVM_FFI_ICHECK(addr[0]->IsInstance<IntImmNode>())
- << "allocated_addr[0] must be a constant integer, got: " << addr[0];
+ << "placement[0] must be a constant integer, got: " << addr[0];
TVM_FFI_ICHECK(addr[1]->IsInstance<IntImmNode>())
- << "allocated_addr[1] must be a constant integer, got: " << addr[1];
+ << "placement[1] must be a constant integer, got: " << addr[1];
int64_t base_bank = static_cast<int64_t>(addr[0].as_or_throw<IntImm>()->value);
int64_t base_addr = static_cast<int64_t>(addr[1].as_or_throw<IntImm>()->value);
stream << "ncc.psum.mod_alloc(base_bank=" << base_bank << ", base_addr=" << base_addr;
@@ -270,7 +270,7 @@ void CodeGenTrainium::DispatchAllocTensor(const BindNode* op, const CallNode* bu
} else {
TVM_FFI_ICHECK(addr.size() == 1);
TVM_FFI_ICHECK(addr[0]->IsInstance<IntImmNode>())
- << "allocated_addr[0] must be a constant integer, got: " << addr[0];
+ << "placement[0] must be a constant integer, got: " << addr[0];
int64_t base_addr = static_cast<int64_t>(addr[0].as_or_throw<IntImm>()->value);
stream << "ncc.sbuf.mod_alloc(base_addr=" << base_addr << "))\n";
}
@@ -633,10 +633,10 @@ void CodeGenTrainium::Dispatch_(const prim::FloorModNode* op, std::ostream& os)
void CodeGenTrainium::DispatchDeclTensor(const BindNode* op, const CallNode* buffer_call) {
Expr data = buffer_call->args[0];
- tvm::Tuple shape = buffer_call->args[1].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[2].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[3].as_or_throw<StringImm>()->value;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
+ tvm::Tuple shape(buffer->shape);
+ DLDataType dtype = buffer->dtype->dtype;
+ ffi::String scope = buffer.scope();
if (scope == "trn.psum" || scope == "trn.sbuf") {
return;
}
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc
index fce99c3701..4da7cfbd66 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -80,11 +80,9 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
}
auto new_stmt = storage_lower->Mutate(stmt, InplaceMode::kDisallow).ValueOrUnchanged(stmt);
for (const auto& [buf, source] : param_flattened_tensors) {
- new_stmt = SeqStmt({Bind(buf, Call(buf.type(), tirx::decl_tensor_op(),
- {source.data(), tvm::Tuple(buf->shape),
- DataTypeImm(buf->dtype->dtype), StringImm(buf.scope())},
- {})),
- std::move(new_stmt)});
+ new_stmt = SeqStmt(
+ {Bind(buf, Call(buf.type(), tirx::decl_tensor_op(), {source.data()}, {}, {buf.type()})),
+ std::move(new_stmt)});
}
return std::make_pair(new_stmt, new_params);
}
@@ -113,8 +111,8 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
call && call->op.same_as(tirx::alloc_tensor_op())) {
TensorVar original_tensor = op->var.as_or_throw<TensorVar>();
ffi::Array<Expr> args = call->args;
- if (args.size() == 4) {
- args.Set(3, Mutate(args[3], inplace_mode).ValueOrUnchanged(args[3]));
+ if (!args.empty()) {
+ args.Set(0, Mutate(args[0], inplace_mode).ValueOrUnchanged(args[0]));
}
auto tensor = original_tensor->layout.has_value()
? GetFlattenedTensor(original_tensor, /*is_alloc=*/true)
@@ -122,13 +120,11 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
if (tensor.same_as(original_tensor) && args.same_as(call->args)) {
return ffi::Unchanged();
}
- args.Set(0, tvm::Tuple(tensor->shape, call->args[0]->loc));
- args.Set(1, DataTypeImm(tensor->dtype->dtype, call->args[1]->loc));
- args.Set(2, StringImm(tensor.scope(), call->args[2]->loc));
- return Bind(
- tensor.var(),
- Call(tensor.type(), tirx::alloc_tensor_op(), args, call->attrs, call->ty_args, call->loc),
- op->loc);
+
+ return Bind(tensor.var(),
+ Call(tensor.type(), tirx::alloc_tensor_op(), args, call->attrs, {tensor.type()},
+ call->loc),
+ op->loc);
}
if (const auto* call = op->value.as<CallNode>();
call && call->op.same_as(tirx::decl_tensor_op())) {
@@ -142,10 +138,8 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
return ffi::Unchanged();
}
return Bind(tensor,
- Call(tensor.type(), tirx::decl_tensor_op(),
- {std::move(data), tvm::Tuple(tensor->shape),
- DataTypeImm(tensor->dtype->dtype), StringImm(tensor.scope())},
- call->attrs, call->ty_args, call->loc),
+ Call(tensor.type(), tirx::decl_tensor_op(), {std::move(data)}, call->attrs,
+ {tensor.type()}, call->loc),
op->loc);
}
return IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc b/src/backend/vulkan/codegen/codegen_spirv.cc
index 2d10ab233c..f95d38d395 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -885,9 +885,9 @@ void CodeGenSPIRV::Dispatch_(const IfNode* op) {
}
void CodeGenSPIRV::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.value().as_or_throw<DictAttrs>();
TVM_FFI_ICHECK(!PrimType(dtype).IsVoid());
@@ -956,7 +956,7 @@ void CodeGenSPIRV::DispatchAllocTensor(const BindNode* op, const CallNode* buffe
void CodeGenSPIRV::DispatchDeclTensor(const BindNode* op, const CallNode* buffer_call) {
Expr data = buffer_call->args[0];
- DLDataType dtype = buffer_call->args[2].as_or_throw<DataTypeImm>()->value;
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
const VarNode* buffer_var = buffer.get();
TVM_FFI_ICHECK(!var_map_.count(buffer_var))
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 20f26360aa..2115e2b7e8 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -785,9 +785,9 @@ void CodeGenWebGPU::DispatchDeclTensor(const BindNode* op, const CallNode* buffe
}
void CodeGenWebGPU::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
TVM_FFI_ICHECK(buffer.defined());
std::string vid = AllocVarID(buffer.get());
diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc
index e3bc17f798..89274c7b00 100644
--- a/src/relax/op/tensor/inspect.cc
+++ b/src/relax/op/tensor/inspect.cc
@@ -281,14 +281,11 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const Call& call) {
{StringImm("Specified axis may not be larger than the tensor's dimensionality")}),
tvm::Bind(
shape_buffer,
- tvm::Call(
- shape_buffer.type(), tvm::tirx::decl_tensor_op(),
- {tvm::Call(shape_buffer.type()->DataPtrType(), tirx::abi_field_get_op(),
- {dlpack_handle, IntImm::Int32(0),
- IntImm::Int32(tirx::TVMStructFieldKind::kDLTensorShape)}),
- tvm::Tuple(shape_buffer->shape), tvm::DataTypeImm(shape_buffer->dtype->dtype),
- tvm::StringImm(shape_buffer.scope())},
- {})),
+ tvm::Call(shape_buffer.type(), tvm::tirx::decl_tensor_op(),
+ {tvm::Call(shape_buffer.type()->DataPtrType(), tirx::abi_field_get_op(),
+ {dlpack_handle, IntImm::Int32(0),
+ IntImm::Int32(tirx::TVMStructFieldKind::kDLTensorShape)})},
+ {}, {shape_buffer.type()})),
tvm::Bind(extent, tirx::MakeTensorLoad(shape_buffer, {axis.as_or_throw<PrimExpr>()})),
tvm::Return(extent)});
diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc b/src/s_tir/analysis/calculate_allocated_memory.cc
index b25ce84565..8f10899d9d 100644
--- a/src/s_tir/analysis/calculate_allocated_memory.cc
+++ b/src/s_tir/analysis/calculate_allocated_memory.cc
@@ -75,9 +75,9 @@ class AllocTensorCalculator : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const CallNode* call) {
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
+ ffi::String scope = call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
auto search = _current_size.find(scope);
if (search == _current_size.end()) {
_current_size[scope] = 0;
diff --git a/src/s_tir/analysis/verify_gpu_code.cc b/src/s_tir/analysis/verify_gpu_code.cc
index ade473011f..0bac812199 100644
--- a/src/s_tir/analysis/verify_gpu_code.cc
+++ b/src/s_tir/analysis/verify_gpu_code.cc
@@ -82,9 +82,9 @@ class GPUCodeVerifier : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const CallNode* call) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
+ ffi::String scope = call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
runtime::StorageScope storage_scope = runtime::StorageScope::Create(scope);
int64_t const_size = 1;
for (const Expr& e : shape->fields) {
diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc
index 526deac70c..48cea9be51 100644
--- a/src/s_tir/backend/adreno/inject_texture_alloc.cc
+++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc
@@ -75,7 +75,7 @@ class TextureAllocInjector : public s_tir::IRMutatorWithAnalyzer {
UnchangedOr<Stmt> Mutate_AllocTensor(const BindNode* op, const CallNode* call,
InplaceMode inplace_mode) {
Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ ffi::String scope = call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
if (IsTextureStorage(scope)) {
op = stmt.as<BindNode>();
if (const auto* call = op ? op->value.as<CallNode>() : nullptr;
@@ -83,8 +83,8 @@ class TextureAllocInjector : public s_tir::IRMutatorWithAnalyzer {
TVM_FFI_THROW(InternalError) << "Expected an allocation binding after tensor mutation";
}
const auto* allocation = op->value.as<CallNode>();
- tvm::Tuple shape = allocation->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = allocation->args[1].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(allocation->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = allocation->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
ffi::Array<PrimExpr> extents =
shape->fields.Map([](const Expr& e) { return e.as_or_throw<PrimExpr>(); });
TVM_FFI_ICHECK(extents.size() >= 3) << "Only 2D Array RGBA texture is currently supported";
@@ -105,11 +105,8 @@ class TextureAllocInjector : public s_tir::IRMutatorWithAnalyzer {
stmt = Bind(op->var.as_or_throw<TensorVar>(),
Call(op->var.as_or_throw<TensorVar>().type(), tirx::decl_tensor_op(),
{Call(op->var.as_or_throw<TensorVar>().type()->DataPtrType(),
- tvm::backend::opencl::nd_mem_alloc_with_scope_op(), args),
- tvm::Tuple(op->var.as_or_throw<TensorVar>()->shape),
- DataTypeImm(op->var.as_or_throw<TensorVar>()->dtype->dtype),
- StringImm(op->var.as_or_throw<TensorVar>().scope())},
- {}, allocation->ty_args, allocation->loc),
+ tvm::backend::opencl::nd_mem_alloc_with_scope_op(), args)},
+ {}, {op->var.as_or_throw<TensorVar>().type()}, allocation->loc),
op->loc);
}
return stmt;
diff --git a/src/s_tir/script/ir_builder/ir.cc b/src/s_tir/script/ir_builder/ir.cc
index f051098a5b..d3c96850b6 100644
--- a/src/s_tir/script/ir_builder/ir.cc
+++ b/src/s_tir/script/ir_builder/ir.cc
@@ -29,7 +29,6 @@ namespace ir_builder {
namespace s_tir {
using ir::ForFrameNode;
-using tirx::TensorDecl;
using tvm::s_tir::IterVar;
using tvm::s_tir::IterVarType;
@@ -52,8 +51,15 @@ TensorVar MatchTensor(ffi::ObjectRef param, ffi::Array<PrimExpr> shape, PrimType
ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
int offset_factor, ffi::Optional<Layout> layout,
ffi::Array<PrimExpr> allocated_addr) {
- TensorVar tensor = TensorDecl(shape, dtype, "", data, strides, elem_offset, storage_scope, align,
- offset_factor, layout);
+ if (data.has_value()) {
+ storage_scope = data.value()->ty.as_or_throw<PtrType>()->storage_scope;
+ }
+ if (!elem_offset.has_value() && offset_factor) {
+ PrimType shape_dtype = shape.empty() ? PrimType::Int(32) : shape[0].ty();
+ elem_offset = tvm::PrimVar("elem_offset", shape_dtype);
+ }
+ TensorVar tensor("", tvm::tirx::TensorType(storage_scope, dtype, shape, strides, elem_offset,
+ align, offset_factor, layout));
tvm::TensorRegion region{ffi::UnsafeInit{}};
if (auto load = param.as<TensorLoad>()) {
region = TensorRegionFromLoad(load.value());
@@ -210,8 +216,12 @@ TensorVar SBlockAllocTensor(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::Opt
<< "ValueError: For `" << scope
<< "` scope, Ts.alloc_tensor does not accept `allocated_addr`";
}
- TensorVar tensor = TensorDecl(shape, dtype, "", std::nullopt, strides, elem_offset, storage_scope,
- align, offset_factor, layout);
+ if (!elem_offset.has_value() && offset_factor) {
+ PrimType shape_dtype = shape.empty() ? PrimType::Int(32) : shape[0].ty();
+ elem_offset = tvm::PrimVar("elem_offset", shape_dtype);
+ }
+ TensorVar tensor("", tvm::tirx::TensorType(storage_scope, dtype, shape, strides, elem_offset,
+ align, offset_factor, layout));
IRBuilder builder = IRBuilder::Current();
auto opt_func_frame = builder->FindFrame<tirx::FunctionFrame>();
if (opt_func_frame.has_value()) {
diff --git a/src/s_tir/transform/compact_tensor_region.cc b/src/s_tir/transform/compact_tensor_region.cc
index 1b5cb950c3..6314be2296 100644
--- a/src/s_tir/transform/compact_tensor_region.cc
+++ b/src/s_tir/transform/compact_tensor_region.cc
@@ -669,18 +669,15 @@ class TensorCompactor : public StmtExprMutator {
TensorVar new_tensor = RewriteAllocTensor(tensor);
bool is_alloc = call->op.same_as(tirx::alloc_tensor_op());
if (new_tensor.same_as(tensor) ||
- (is_alloc &&
- PrimType(call->args[1].as_or_throw<DataTypeImm>()->value) != new_tensor->dtype)) {
+ (is_alloc && call->ty_args[0].as_or_throw<TensorType>()->dtype != new_tensor->dtype)) {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
// Update the producer before generic Bind mutation propagates its result type.
- size_t shape_index = is_alloc ? 0 : 1;
ffi::Array<Expr> args = call->args;
- args.Set(shape_index, tvm::Tuple(new_tensor->shape, args[shape_index]->loc));
auto rewritten = ffi::make_object<BindNode>(*op);
rewritten->value =
- Call(new_tensor.type(), call->op, args, call->attrs, call->ty_args, call->loc);
+ Call(new_tensor.type(), call->op, args, call->attrs, {new_tensor.type()}, call->loc);
tvm::Bind binding(std::move(rewritten));
return StmtExprMutator::Mutate_(binding.get(), inplace_mode).ValueOrUnchanged(binding);
}
diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc
index c651954280..df0b54537a 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -196,7 +196,7 @@ class VarTouchedAnalysis : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const CallNode* call) {
expr_touched_->Reset(false);
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
for (const Expr& extent : shape->fields) {
expr_touched_->Visit(extent);
}
@@ -583,7 +583,7 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
// AllocTensor
UnchangedOr<Stmt> Mutate_AllocTensor(const BindNode* op, const CallNode* call,
InplaceMode inplace_mode) {
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
ffi::Array<PrimExpr> original_shape = shape->fields.as_or_throw<ffi::Array<PrimExpr>>();
ffi::Array<PrimExpr> new_shape = original_shape.Map([this](const PrimExpr& s) {
// Keep the retained allocation and its tensor type unchanged.
@@ -591,8 +591,8 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
});
ffi::Array<Expr> args = call->args;
- if (args.size() == 4) {
- args.Set(3, Mutate(args[3], InplaceMode::kDisallow).ValueOrUnchanged(args[3]));
+ if (!args.empty()) {
+ args.Set(0, Mutate(args[0], InplaceMode::kDisallow).ValueOrUnchanged(args[0]));
}
if (visit_touched_var_ && !vt_loop_injected_) {
return InjectVTLoop(ffi::GetRef<Stmt>(op), true);
@@ -615,12 +615,10 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
type->shape = new_shape;
TensorVar new_tensor = RebuildTensorVar(op->var.as_or_throw<TensorVar>(), std::move(type));
VarRemapSet(op->var.as_or_throw<TensorVar>(), new_tensor);
- args.Set(0, tvm::Tuple(new_tensor->shape, call->args[0]->loc));
- args.Set(1, DataTypeImm(new_tensor->dtype->dtype, call->args[1]->loc));
- args.Set(2, StringImm(new_tensor.scope(), call->args[2]->loc));
+
return Bind(new_tensor.var(),
- Call(new_tensor.type(), tirx::alloc_tensor_op(), args, call->attrs, call->ty_args,
- call->loc),
+ Call(new_tensor.type(), tirx::alloc_tensor_op(), args, call->attrs,
+ {new_tensor.type()}, call->loc),
op->loc);
}
}
diff --git a/src/s_tir/transform/lower_opaque_block.cc b/src/s_tir/transform/lower_opaque_block.cc
index b82302604f..fed2fe5380 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -98,13 +98,12 @@ class OpaqueBlockLower : public StmtExprMutator {
}
allocate_annotations.Set(s_tir::attr::kTensorDimAlign, allocate_aligns);
}
- allocate_annotations.Set(tirx::attr::kTensorDataAlignment,
+ allocate_annotations.Set(tvm::tirx::attr::kTensorDataAlignment,
IntImm::Int32(tensor->data_alignment));
- ffi::Array<Expr> args{tvm::Tuple(tensor->shape), DataTypeImm(tensor->dtype->dtype),
- StringImm(tensor.scope())};
+ ffi::Array<Expr> args;
if (auto address = addresses.Get(tensor.var())) args.push_back(tvm::Tuple(address.value()));
body = SeqStmt({Bind(tensor.var(), Call(tensor.type(), tirx::alloc_tensor_op(), args,
- DictAttrs(allocate_annotations))),
+ DictAttrs(allocate_annotations), {tensor.type()})),
std::move(body)});
}
return body;
diff --git a/src/s_tir/transform/lower_vtcm_alloc.cc b/src/s_tir/transform/lower_vtcm_alloc.cc
index 4d6c5946ef..29545cf40d 100644
--- a/src/s_tir/transform/lower_vtcm_alloc.cc
+++ b/src/s_tir/transform/lower_vtcm_alloc.cc
@@ -51,9 +51,9 @@ class VtcmAllocator : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_AllocTensor(const BindNode* op, const CallNode* call,
InplaceMode inplace_mode) {
- ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ ffi::String scope = call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
if (IsVtcmStorage(scope)) {
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
ffi::Array<Expr> args;
args.push_back(StringImm(scope));
args.push_back(IntImm::Int64(shape->fields.size()));
@@ -62,10 +62,8 @@ class VtcmAllocator : public StmtExprMutator {
return Bind(tensor,
Call(tensor.type(), tirx::decl_tensor_op(),
{Call(tensor.type()->DataPtrType(),
- tvm::backend::opencl::nd_mem_alloc_with_scope_op(), args),
- tvm::Tuple(tensor->shape), DataTypeImm(tensor->dtype->dtype),
- StringImm(tensor.scope())},
- {}, call->ty_args, call->loc),
+ tvm::backend::opencl::nd_mem_alloc_with_scope_op(), args)},
+ {}, {tensor.type()}, call->loc),
op->loc);
}
return StmtExprMutator::Mutate_(op, inplace_mode);
diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc
index 87d03f1a25..35282988ef 100644
--- a/src/s_tir/transform/merge_shared_memory_allocations.cc
+++ b/src/s_tir/transform/merge_shared_memory_allocations.cc
@@ -124,7 +124,7 @@ class AllocateCollector : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const CallNode* call) {
- ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ ffi::String scope = call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
StorageScope storage_scope = StorageScope::Create(scope);
if (is_dynamic_ && storage_scope.rank == runtime::StorageRank::kShared &&
storage_scope.tag == ".dyn") {
@@ -497,9 +497,7 @@ class SharedMemoryRewriter : public StmtExprMutator {
// typed views; target codegen emits the required pointer cast.
visited_body = SeqStmt(
{Bind(remapped, Call(remapped.type(), tirx::decl_tensor_op(),
- {scope.merged_tensor.value().data(), tvm::Tuple(remapped->shape),
- DataTypeImm(remapped->dtype->dtype), StringImm(remapped.scope())},
- {})),
+ {scope.merged_tensor.value().data()}, {}, {remapped.type()})),
visited_body});
}
@@ -517,11 +515,8 @@ class SharedMemoryRewriter : public StmtExprMutator {
annotations.Set(tvm::tirx::attr::kVolatile, true);
}
Stmt alloc_stmt = Bind(scope.merged_tensor.value().var(),
- Call(scope.merged_tensor.value().type(), tirx::alloc_tensor_op(),
- {tvm::Tuple(scope.merged_tensor.value()->shape),
- DataTypeImm(scope.merged_tensor.value()->dtype->dtype),
- StringImm(scope.merged_tensor.value().scope())},
- DictAttrs(annotations)));
+ Call(scope.merged_tensor.value().type(), tirx::alloc_tensor_op(), {},
+ DictAttrs(annotations), {scope.merged_tensor.value().type()}));
Stmt new_body = SeqStmt({alloc_stmt, visited_body});
// 8. Pop the scope.
@@ -535,7 +530,7 @@ class SharedMemoryRewriter : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_AllocTensor(const BindNode* op, const CallNode* call,
InplaceMode inplace_mode) {
- ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ ffi::String scope = call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
StorageScope storage_scope = StorageScope::Create(scope);
if (storage_scope.rank == runtime::StorageRank::kShared &&
storage_scope.tag == (is_dynamic_ ? ".dyn" : "")) {
@@ -585,10 +580,8 @@ class SharedMemoryRewriter : public StmtExprMutator {
!new_buf.same_as(node->var)) {
const auto* new_call = node->value.as<CallNode>();
return Bind(new_buf,
- Call(new_buf.type(), tirx::decl_tensor_op(),
- {new_call->args[0], tvm::Tuple(new_buf->shape),
- DataTypeImm(new_buf->dtype->dtype), StringImm(new_buf.scope())},
- new_call->attrs, new_call->ty_args, new_call->loc),
+ Call(new_buf.type(), tirx::decl_tensor_op(), {new_call->args[0]}, new_call->attrs,
+ {new_buf.type()}, new_call->loc),
node->loc);
}
return node;
diff --git a/src/script/printer/ir/call.cc b/src/script/printer/ir/call.cc
index fe331c7dbc..eea3121c7b 100644
--- a/src/script/printer/ir/call.cc
+++ b/src/script/printer/ir/call.cc
@@ -47,11 +47,6 @@ ffi::Optional<ExprDoc> StandardCallDocTranslate(DocTranslatorObj* d, const CallN
: !op->attrs_type_key.empty()) {
return RawCall(d, call);
}
- try {
- op.Validate(call);
- } catch (const ffi::Error&) {
- return RawCall(d, call);
- }
ffi::Array<ExprDoc> args;
for (const Expr& arg : call->args) {
args.push_back(d->Translate(arg).value());
@@ -65,7 +60,13 @@ ffi::Optional<ExprDoc> StandardCallDocTranslate(DocTranslatorObj* d, const CallN
values.push_back(ListDoc(types));
}
bool default_attrs = false;
- if (call->attrs.defined()) {
+ if (call->attrs.defined() && call->attrs.value().as<DictAttrsNode>()) {
+ default_attrs = call->attrs.value().as_or_throw<DictAttrs>()->dict.empty();
+ if (!default_attrs) {
+ keys.push_back("attrs");
+ values.push_back(d->Translate(call->attrs.value()).value());
+ }
+ } else if (call->attrs.defined()) {
std::vector<std::pair<ffi::String, ffi::Any>> fields;
ffi::reflection::ForEachFieldInfo(
TVMFFIGetTypeInfo(call->attrs->type_index()), [&](const TVMFFIFieldInfo* field) {
@@ -94,14 +95,13 @@ ffi::Optional<ExprDoc> StandardCallDocTranslate(DocTranslatorObj* d, const CallN
values.push_back(AnyValue(d, value));
}
}
- bool omit_result = false;
- try {
- // Match construction without ty, rather than let inference read the stored result.
- Call inferred(std::nullopt, op, call->args, call->attrs, call->ty_args);
- omit_result = ffi::StructuralEqual()(inferred->ty, call->ty);
- } catch (const ffi::Error&) {
- // Missing information, unavailable inference, or errors require the exact stored type.
- }
+ // Well-formed calls reconstruct their result from registered inference or a fixed type.
+ // Builder-only inference is deliberately excluded: named construction cannot run it.
+ bool omit_result = call->ty.as<MissingType>().has_value() ||
+ (Op::HasAttrMap(tvm::op_attr::kInferType) &&
+ Op::GetAttrMap<FInferType>(tvm::op_attr::kInferType).count(op)) ||
+ (Op::HasAttrMap(tvm::op_attr::kFixedReturnType) &&
+ Op::GetAttrMap<TFixedReturnType>(tvm::op_attr::kFixedReturnType).count(op));
if (!omit_result) {
if (default_attrs) return RawCall(d, call);
keys.push_back("ty");
diff --git a/src/script/printer/ir/stmt.cc b/src/script/printer/ir/stmt.cc
index 5d4fd09ab8..3d99d77254 100644
--- a/src/script/printer/ir/stmt.cc
+++ b/src/script/printer/ir/stmt.cc
@@ -16,6 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ir/stmt.h>
#include <tvm/script/printer/doc_translator.h>
@@ -46,7 +47,11 @@ ffi::Optional<ExprDoc> BindDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
// an RHS, so this binding owns the one remaining assignment.
if (!rhs.has_value()) return std::nullopt;
ffi::Optional<ExprDoc> annotation = std::nullopt;
- if (!existing) {
+ // An explicit result type argument already exposes the binding type.
+ const auto* call = stmt->value.as<CallNode>();
+ bool explicit_type =
+ call && call->ty_args.size() == 1 && ffi::StructuralEqual()(stmt->var->ty, call->ty_args[0]);
+ if (!existing && !explicit_type) {
ExprDoc type = d->Translate(stmt->var->ty).value();
if (auto primitive = stmt->var->ty.as<PrimType>();
primitive && primitive.value().IsScalableVector()) {
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 2dc363cc48..658c3aa178 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -2198,8 +2198,8 @@ void CodeGenLLVM::Dispatch_(const IfNode* op) {
}
void CodeGenLLVM::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.as_or_throw<DictAttrs>();
EmitDebugLocation(op);
@@ -2332,8 +2332,8 @@ void CodeGenLLVM::Dispatch_(const SeqStmtNode* op) {
void CodeGenLLVM::DispatchDeclTensor(const BindNode* op, const CallNode* buffer_call) {
Expr data = buffer_call->args[0];
- DLDataType dtype = buffer_call->args[2].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[3].as_or_throw<StringImm>()->value;
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
EmitDebugLocation(op);
const VarNode* buffer_var = buffer.get();
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index ae2271e41d..294735afc7 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -937,8 +937,8 @@ void CodeGenC::PrintVecBinaryOp(const std::string& op, const PrimType& t, PrimEx
void CodeGenC::DispatchDeclTensor(const BindNode* op, const CallNode* buffer_call) {
Expr data = buffer_call->args[0];
- DLDataType dtype = buffer_call->args[2].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[3].as_or_throw<StringImm>()->value;
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
const VarNode* source = data.as<VarNode>();
if (const auto* call = data.as<CallNode>();
@@ -1276,9 +1276,9 @@ void CodeGenC::Dispatch_(const BindNode* op) {
}
void CodeGenC::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call) {
- tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = buffer_call->args[2].as_or_throw<StringImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
+ DLDataType dtype = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
+ ffi::String scope = buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.as_or_throw<DictAttrs>();
TVM_FFI_ICHECK(buffer.defined());
diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc
index 22e21f302c..b88c7176f4 100644
--- a/src/tirx/ir/specialize.cc
+++ b/src/tirx/ir/specialize.cc
@@ -175,8 +175,9 @@ class FunctionSpecializer : public StmtExprMutator {
(call->op.same_as(tirx::alloc_tensor_op()) || call->op.same_as(tirx::decl_tensor_op()))) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->WithDefRegionKind(
kTVMFFIDefRegionKindSimple, [&]() { return this->Visit(op->var); }));
- if (call->op.same_as(tirx::decl_tensor_op())) return Visit(call->args[0]);
- return std::nullopt;
+ // The definition above owns traversal of the complete tensor type.
+ // Value operands include declaration data or allocation placement.
+ return Visit(call->args);
}
return StmtExprVisitor::Visit_(op);
}
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index a265e64f11..64fa660b68 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -104,6 +104,9 @@ ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const CallNode* op) {
for (const auto& child : op->args) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(child));
}
+ for (const auto& type : op->ty_args) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(type));
+ }
return std::nullopt;
}
diff --git a/src/tirx/ir/tile_dispatch.cc b/src/tirx/ir/tile_dispatch.cc
index 417f5d3749..b7f14ff483 100644
--- a/src/tirx/ir/tile_dispatch.cc
+++ b/src/tirx/ir/tile_dispatch.cc
@@ -45,11 +45,10 @@ Value getOrSetDefault(ffi::Map<ffi::String, ffi::ObjectRef>& m, const Key& key,
void DispatchContextNode::AddAllocTensor(TensorVar tensor, ffi::Array<PrimExpr> allocated_addr,
ffi::Map<ffi::String, ffi::Any> annotations) {
auto tensors = getOrSetDefault(callbacks, callback::kPrivateAlloc, ffi::Array<Bind>());
- ffi::Array<Expr> args{tvm::Tuple(tensor->shape), DataTypeImm(tensor->dtype->dtype),
- StringImm(tensor.scope())};
+ ffi::Array<Expr> args;
if (!allocated_addr.empty()) args.push_back(tvm::Tuple(allocated_addr));
- tensors.push_back(
- Bind(tensor.var(), Call(tensor.type(), alloc_tensor_op(), args, DictAttrs(annotations))));
+ tensors.push_back(Bind(tensor.var(), Call(tensor.type(), alloc_tensor_op(), args,
+ DictAttrs(annotations), {tensor.type()})));
callbacks.Set(callback::kPrivateAlloc, tensors);
}
diff --git a/src/tirx/op/memory.cc b/src/tirx/op/memory.cc
index bfde7e8fe0..f61af69314 100644
--- a/src/tirx/op/memory.cc
+++ b/src/tirx/op/memory.cc
@@ -82,52 +82,37 @@ ffi::Expected<Type> InferTypeTensorDataPtr(const CallNode* call) noexcept try {
return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
-// Essential tensor properties follow operands; the supplied result type retains layout metadata.
-template <int shape_index>
-ffi::Expected<Type> InferTypeTensor(const CallNode* call) noexcept try {
- if constexpr (shape_index == 0) {
- TVM_FFI_CHECK(call->args.size() == 3 || call->args.size() == 4, ValueError);
- if (call->args.size() == 4) {
- auto placement = call->args[3].as_or_throw<tvm::Tuple>();
+// Allocation and declaration share the complete, explicit tensor type argument.
+template <bool is_alloc>
+Type InferTypeTensor(const CallNode* call) {
+ TVM_FFI_CHECK_EQ(call->ty_args.size(), 1U, ValueError)
+ << "Tensor allocation and declaration require exactly one TensorType argument";
+ auto type = call->ty_args[0].as_or_throw<TensorType>();
+ if constexpr (is_alloc) {
+ TVM_FFI_CHECK_LE(call->args.size(), 1U, ValueError)
+ << "alloc_tensor accepts at most one placement tuple";
+ if (!call->args.empty()) {
+ auto placement = call->args[0].as_or_throw<tvm::Tuple>();
for (const Expr& address : placement->fields) address.as_or_throw<PrimExpr>();
}
} else {
- TVM_FFI_CHECK_EQ(call->args.size(), 4U, ValueError);
+ TVM_FFI_CHECK_EQ(call->args.size(), 1U, ValueError)
+ << "decl_tensor requires one backing data operand";
}
- tvm::Tuple shape = call->args[shape_index].as_or_throw<tvm::Tuple>();
- DLDataType dtype = call->args[shape_index + 1].as_or_throw<DataTypeImm>()->value;
- ffi::String scope = call->args[shape_index + 2].as_or_throw<StringImm>()->value;
- auto original = call->ty.as_or_throw<TensorType>();
- if (ffi::StructuralEqual()(shape->fields, original->shape) && dtype == original->dtype->dtype &&
- scope == original->storage_scope) {
- return original;
- }
- auto inferred = ffi::make_object<TensorTypeNode>(*original.get());
- inferred->shape =
- shape->fields.Map([](const Expr& extent) { return extent.as_or_throw<PrimExpr>(); });
- inferred->dtype = PrimType(dtype);
- inferred->storage_scope = scope;
- return TensorType(std::move(inferred));
-} catch (const ffi::Error& error) {
- return ffi::Unexpected(error);
-} catch (const std::exception& error) {
- return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
+ return type;
}
-ffi::Expected<void> ValidateAllocTensor(const CallNode* call) noexcept {
- auto inferred = InferTypeTensor<0>(call);
- if (!inferred.has_value()) return ffi::Unexpected(inferred.error());
- return {};
-}
-
-ffi::Expected<void> ValidateDeclTensor(const CallNode* call) noexcept try {
- TVM_FFI_CHECK_EQ(call->args.size(), 4U, ValueError);
- call->ty.as_or_throw<TensorType>();
- return {};
-} catch (const ffi::Error& error) {
- return ffi::Unexpected(error);
-} catch (const std::exception& error) {
- return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
+template <bool is_alloc>
+void ValidateTensorProducer(const CallNode* call) {
+ Type inferred = InferTypeTensor<is_alloc>(call);
+ TVM_FFI_CHECK(ffi::StructuralEqual()(call->ty, inferred), TypeError)
+ << "Tensor producer result must match its TensorType argument";
+ if constexpr (is_alloc) {
+ TVM_FFI_CHECK(!call->attrs.defined() || call->attrs.as<DictAttrsNode>(), ValueError)
+ << "alloc_tensor annotations must be DictAttrs";
+ } else {
+ TVM_FFI_CHECK(!call->attrs.defined(), ValueError) << "decl_tensor does not accept attributes";
+ }
}
const Op& reinterpret_op() {
@@ -376,12 +361,16 @@ const Op& alloc_tensor_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.alloc_tensor")
.set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory, ffi::String("builtin"))
- .set_attr<FInferType>(tvm::op_attr::kInferType, FInferType::FromNative<&InferTypeTensor<0>>())
+ .set_attr<FInferType>(tvm::op_attr::kInferType,
+ FInferType::FromNative<&InferTypeTensor<true>>())
.set_validator(ffi::reflection::NativeFunctionView<void(
- const CallNode*)>::FromNative<&ValidateAllocTensor>())
- .add_arg("shape", "The tuple of tensor extents.")
- .add_arg("dtype", "The tensor data type.")
- .add_arg("scope", "The storage scope.")
+ const CallNode*)>::FromNative<&ValidateTensorProducer<true>>(),
+ true)
+ .signature(sig::var_args<tvm::Tuple>("placement", "Optional producer-owned placement."),
+ sig::ty_arg<TensorType>("tensor_type"))
+ .set_attrs_type_key(DictAttrsNode::_type_key)
+ .set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
+ ffi::String("tirx.alloc_tensor"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kOpaque));
}
@@ -394,13 +383,15 @@ const Op& decl_tensor_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.decl_tensor")
.set_attr<TIRxOpCategory>(tvm::tirx::op_attr::kOpCategory, ffi::String("builtin"))
- .set_attr<FInferType>(tvm::op_attr::kInferType, FInferType::FromNative<&InferTypeTensor<1>>())
+ .set_attr<FInferType>(tvm::op_attr::kInferType,
+ FInferType::FromNative<&InferTypeTensor<false>>())
.set_validator(ffi::reflection::NativeFunctionView<void(
- const CallNode*)>::FromNative<&ValidateDeclTensor>())
- .add_arg("data", "The existing data pointer.")
- .add_arg("shape", "The tuple of tensor extents.")
- .add_arg("dtype", "The tensor data type.")
- .add_arg("scope", "The storage scope.")
+ const CallNode*)>::FromNative<&ValidateTensorProducer<false>>(),
+ true)
+ .signature(sig::arg("data", "The existing data pointer."),
+ sig::ty_arg<TensorType>("tensor_type"))
+ .set_attr<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName,
+ ffi::String("tirx.decl_tensor"))
.set_attr<TCallEffectKind>(tvm::op_attr::kCallEffectKind,
static_cast<int64_t>(CallEffectKind::kPure));
}
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index 4a57c7fc31..bd9d97c97a 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -29,7 +29,6 @@
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
-#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/script/ir_builder/ir.h>
#include <tvm/tirx/stmt.h>
@@ -65,14 +64,6 @@ tvm::tirx::TensorType TensorTypeDecl(ffi::Array<PrimExpr> shape, PrimType dtype,
} // namespace
-TensorVar TensorDecl(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String tensor_name,
- ffi::Optional<Expr> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
- ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
- int offset_factor, ffi::Optional<Layout> layout) {
- return TensorVar(tensor_name, TensorTypeDecl(shape, dtype, data, strides, elem_offset,
- storage_scope, align, offset_factor, layout));
-}
-
FunctionFrame Function(bool is_private, bool persistent) {
ffi::ObjectPtr<FunctionFrameNode> n = ffi::make_object<FunctionFrameNode>();
n->name = std::nullopt;
@@ -284,70 +275,6 @@ tvm::Stmt TensorStore(Expr dest, ffi::Array<PrimExpr> indices, PrimExpr value) {
return store;
}
-TensorVar DeclTensor(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String tensor_name,
- ffi::Optional<Expr> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
- ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
- int offset_factor, ffi::Optional<Layout> layout,
- ffi::Optional<PrimExpr> allocated_addr) {
- std::string scope = static_cast<std::string>(storage_scope);
- if (scope.empty()) {
- scope = "global";
- }
-
- TensorVar tensor = TensorDecl(shape, dtype, tensor_name, data, strides, elem_offset,
- storage_scope, align, offset_factor, layout);
- if (scope == "tmem" && allocated_addr.has_value()) {
- TVM_FFI_CHECK(!data.has_value(), ValueError)
- << "A TMEM declaration cannot have both data and an address";
- Location loc = IRBuilder::Current()->GetCurrentLoc();
- AddToParent(tvm::Bind(tensor.var(),
- Call(std::nullopt, Op::Get("tirx.cuda.decl_tmem"),
- {allocated_addr.value()}, {}, {tensor.type()}, loc),
- loc));
- return tensor;
- }
- TVM_FFI_CHECK(!allocated_addr.has_value() || !data.has_value(), ValueError)
- << "Placement addresses apply to allocations, not pointer-backed declarations";
- TVM_FFI_CHECK(!allocated_addr.has_value() || (scope != "global" && scope != "shared" &&
- scope != "shared.dyn" && scope != "local"),
- ValueError)
- << "This storage scope does not support allocation placement";
- Location loc = IRBuilder::Current()->GetCurrentLoc();
- if (data.has_value()) {
- AddToParent(tvm::Bind(tensor.var(),
- Call(tensor.type(), tvm::tirx::decl_tensor_op(),
- {data.value(), tvm::Tuple(tensor->shape),
- DataTypeImm(tensor->dtype->dtype), StringImm(tensor.scope())},
- {}, {}, loc),
- loc));
- } else {
- // Without a backing pointer, declare and allocate the tensor together.
- ffi::Array<Expr> args{tvm::Tuple(tensor->shape), DataTypeImm(tensor->dtype->dtype),
- StringImm(tensor.scope())};
- if (allocated_addr.has_value()) args.push_back(tvm::Tuple({allocated_addr.value()}));
- AddToParent(tvm::Bind(
- tensor.var(), Call(tensor.type(), tvm::tirx::alloc_tensor_op(), args, DictAttrs(), {}, loc),
- loc));
- }
- return tensor;
-}
-
-TensorVar AllocTensor(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String storage_scope,
- ffi::Optional<ffi::Map<ffi::String, ffi::Any>> annotations) {
- TensorVar tensor = TensorDecl(shape, dtype, "", std::nullopt, std::nullopt, std::nullopt,
- storage_scope, 0, 0, std::nullopt);
- ffi::Array<Expr> args{tvm::Tuple(tensor->shape), DataTypeImm(tensor->dtype->dtype),
- StringImm(tensor.scope())};
- auto attrs = annotations.value_or(ffi::Map<ffi::String, ffi::Any>());
- if (auto placement = attrs.Get("tensor_allocated_addr")) {
- args.push_back(tvm::Tuple(placement.value().as_or_throw<ffi::Array<PrimExpr>>()));
- attrs.erase("tensor_allocated_addr");
- }
- AddToParent(tvm::Bind(tensor.var(), Call(tensor.type(), tvm::tirx::alloc_tensor_op(), args,
- DictAttrs(std::move(attrs)))));
- return tensor;
-}
-
Var Ptr(PrimType dtype, ffi::String storage_scope = "global") {
PtrType type_annotation(dtype, storage_scope);
return tvm::Var("", type_annotation);
@@ -408,13 +335,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("script.ir_builder.tirx.FuncName", FuncName)
.def("script.ir_builder.tirx.FuncAttrs", FuncAttrs)
.def("script.ir_builder.tirx.FuncRet", FuncRet)
- .def("script.ir_builder.tirx.AllocTensor", AllocTensor)
.def("script.ir_builder.tirx.Serial", Serial)
.def("script.ir_builder.tirx.Parallel", Parallel)
.def("script.ir_builder.tirx.Vectorized", Vectorized)
.def("script.ir_builder.tirx.Unroll", Unroll)
.def("script.ir_builder.tirx.ThreadBinding", ThreadBinding)
- .def("script.ir_builder.tirx.DeclTensor", DeclTensor)
.def("script.ir_builder.tirx.TensorStore", TensorStore)
.def("script.ir_builder.tirx.Ptr", Ptr);
}
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index a1a24b7c18..821e97c137 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -33,7 +33,6 @@
#include <vector>
#include "../../../script/printer/ir/utils.h"
-#include "utils.h"
namespace tvm {
namespace script {
@@ -71,40 +70,8 @@ ffi::Optional<ExprDoc> SeqStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView inp
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SeqStmtNode>(input);
TVM_FFI_CHECK(destination == nullptr, TypeError)
<< "printer statement-only node cannot fulfill a destination";
- for (size_t i = 0; i < stmt->seq.size(); ++i) {
- d->Translate(stmt->seq[i]);
- if (i + 1 == stmt->seq.size()) continue;
- const auto* alloc = stmt->seq[i].as<BindNode>();
- const auto* allocation = alloc ? alloc->value.as<CallNode>() : nullptr;
- const auto* store = stmt->seq[i + 1].as<TensorStoreNode>();
- auto docs = d->CurrentScopeDocs();
- if (!allocation || !allocation->op.same_as(tirx::alloc_tensor_op()) || !store ||
- !alloc->var.same_as(store->dest) || docs.empty())
- continue;
- auto scalar = docs.back().as<AssignDoc>();
- if (!IsScalarTensor(d, alloc->var) || !scalar.has_value() ||
- !std::all_of(store->indices.begin(), store->indices.end(), tvm::prim::IsZero))
- continue;
- bool reads_allocation = false;
- ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
- store->value, [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
- reads_allocation |= var.same_as(alloc->var);
- return ffi::WalkResult::Advance();
- });
- if (reads_allocation) continue;
- size_t before = docs.size();
- d->Translate(stmt->seq[++i]);
- // Only combine a plain adjacent store. Translation may emit prerequisites.
- if (docs.size() == before + 1) {
- if (auto initialization = docs.back().as<AssignDoc>()) {
- scalar.value()->rhs = initialization.value()->rhs;
- // Preserve both statement origins using ordinary annotation/assignment
- // occurrences while retaining the value's more precise child origin.
- d->RecordOrigin(scalar.value()->annotation.value(), ffi::GetRef<Bind>(alloc));
- d->RecordOrigin(scalar.value(), ffi::GetRef<TensorStore>(store));
- docs.pop_back();
- }
- }
+ for (const Stmt& child : stmt->seq) {
+ d->Translate(child);
}
return std::nullopt;
}
diff --git a/src/tirx/script/printer/tensor.cc b/src/tirx/script/printer/tensor.cc
index 5d8fe89f4d..589b71b4cf 100644
--- a/src/tirx/script/printer/tensor.cc
+++ b/src/tirx/script/printer/tensor.cc
@@ -22,7 +22,7 @@
#include <tvm/ir/prim/vector_expr.h>
#include <tvm/script/printer/doc_translator.h>
#include <tvm/tirx/layout.h>
-#include <tvm/tirx/op/memory.h>
+#include <tvm/tirx/type.h>
#include <algorithm>
#include <optional>
@@ -30,22 +30,12 @@
#include <vector>
#include "../../../script/printer/ir/utils.h"
-#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
namespace details {
-bool IsScalarTensor(DocTranslatorObj* d, const Expr& source) {
- auto var = source.as<Var>();
- if (!var.has_value()) return false;
- IdDoc id = d->VarGetOrAllocId(var.value(), false);
- auto column = d->GetOrCreateExtraState<ffi::Dict<IdDoc, bool>>("tirx.tensor_as_mutable_var");
- auto found = column.find(id);
- return found != column.end() && (*found).second;
-}
-
namespace {
ffi::Optional<ExprDoc> TensorVarDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
@@ -59,8 +49,7 @@ ffi::Optional<ExprDoc> TensorVarDocTranslate(DocTranslatorObj* d, ffi::AnyView i
EmitVarDefinition(d, var, rhs);
return std::nullopt;
}
- // Mutable scalar syntax binds a TensorLoad; resource uses need its allocation.
- return IsScalarTensor(d, var) ? IdDoc(id->name)->Attr("source") : ExprDoc(IdDoc(id->name));
+ return IdDoc(id->name);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -68,103 +57,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
type_attr::kDocTranslateVarByTy, FDocTranslate::FromNative<&TensorVarDocTranslate>());
}
-ffi::Optional<ExprDoc> TensorOperationDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
- const ffi::Object* destination) {
- const auto* call =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(input);
- // Surface tensor constructors emit a binding, so they cannot replace an
- // allocation expression nested inside another Call or statement.
- if (!destination) return RawCall(d, call);
- TVM_FFI_CHECK(destination->IsInstance<VarNode>(), TypeError)
- << "Tensor operation destination must be a variable";
- auto var = ffi::GetRef<Var>(static_cast<const VarNode*>(destination));
- if (!ffi::StructuralEqual()(var->ty, call->ty)) return RawCall(d, call);
- bool is_alloc = call->op.same_as(tirx::alloc_tensor_op());
- size_t shape_index = is_alloc ? 0 : 1;
- auto tensor = call->ty.as<tirx::TensorType>();
- if (!tensor || (call->args.size() != shape_index + 3 && !(is_alloc && call->args.size() == 4)) ||
- !call->ty_args.empty() || call->attrs.defined() != is_alloc ||
- (call->attrs.defined() &&
- !(call->attrs.has_value() ? call->attrs.value().as<DictAttrsNode>() : nullptr))) {
- return RawCall(d, call);
- }
- auto shape = call->args[shape_index].as<TupleNode>();
- auto dtype = call->args[shape_index + 1].as<DataTypeImmNode>();
- auto scope = call->args[shape_index + 2].as<StringImmNode>();
- if (!shape || !dtype || !scope || !ffi::StructuralEqual()(shape->fields, tensor.value()->shape) ||
- dtype->value != tensor.value()->dtype->dtype ||
- scope->value != tensor.value()->storage_scope || scope->value.empty() ||
- tensor.value()->data_alignment <= 0 || tensor.value()->offset_factor == 0) {
- return RawCall(d, call);
- }
- ffi::Map<ffi::String, ffi::Any> annotations;
- if (is_alloc) annotations = call->attrs.value().as_or_throw<DictAttrs>()->dict;
- for (const auto& [key, value] : annotations) {
- if (value.type_index() == ffi::TypeIndex::kTVMFFIInt ||
- value.type_index() == ffi::TypeIndex::kTVMFFIBool ||
- value.type_index() == ffi::TypeIndex::kTVMFFIFloat) {
- return RawCall(d, call);
- }
- }
- ffi::Optional<Expr> data = is_alloc ? std::nullopt : ffi::Optional<Expr>(call->args[0]);
- if (!is_alloc) {
- const auto* pointer = data.value()->ty.as<PtrTypeNode>();
- if (!pointer || pointer->storage_scope != scope->value) return RawCall(d, call);
- }
- CallDoc rhs = d->Translate(tensor.value()).value().as_or_throw<CallDoc>();
- if (rhs->callee.as_or_throw<AttrAccessDoc>()->name != "Tensor") return RawCall(d, call);
- ffi::String method = is_alloc ? "alloc_tensor" : "decl_tensor";
- if (is_alloc && (scope->value == "local" || scope->value == "shared")) {
- method = scope->value == "local" ? "alloc_local" : "alloc_shared";
- for (size_t i = 0; i < rhs->kwargs_keys.size(); ++i) {
- if (rhs->kwargs_keys[i] == "scope") {
- rhs->kwargs_keys.erase(rhs->kwargs_keys.begin() + i);
- rhs->kwargs_values.erase(rhs->kwargs_values.begin() + i);
- break;
- }
- }
- }
- rhs->callee = NamespaceDoc("tirx")->Attr(method);
- if (is_alloc && call->args.size() == 4) {
- auto placement = call->args[3].as<tvm::TupleNode>();
- if (!placement) return RawCall(d, call);
- rhs->kwargs_keys.push_back("allocated_addr");
- rhs->kwargs_values.push_back(AnyValue(d, placement->fields));
- }
- if (data.has_value()) {
- rhs->kwargs_keys.insert(rhs->kwargs_keys.begin(), "data");
- rhs->kwargs_values.insert(rhs->kwargs_values.begin(), d->Translate(data.value()).value());
- }
- if (!annotations.empty()) {
- CallDoc doc = rhs.as_or_throw<CallDoc>();
- doc->kwargs_keys.push_back("annotations");
- doc->kwargs_values.push_back(AnyValue(d, annotations));
- }
- IdDoc lhs = VarDoc(d, var);
- AssignDoc allocation(lhs, rhs, std::nullopt);
- if (is_alloc && d->GetExtraConfig<bool>("tirx.scalar_tensor_as_mutable_var", true) &&
- annotations.empty() && call->args.size() == 3 && scope->value == "local" &&
- tensor.value()->IsScalar(true)) {
- ExprDoc annotation = d->Translate(tensor.value()->dtype).value();
- // T.bool is a type annotation but has no mutable-declaration syntax.
- if (annotation.as<AttrAccessDoc>() && !tensor.value()->dtype.MatchesCode(kDLBool)) {
- auto scalars = d->GetOrCreateExtraState<ffi::Dict<IdDoc, bool>>("tirx.tensor_as_mutable_var");
- scalars.Set(d->VarGetOrAllocId(var, true), true);
- allocation->rhs = std::nullopt;
- allocation->annotation = annotation;
- }
- }
- d->Emit(allocation, ffi::GetRef<Call>(call));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- for (const char* name : {"tirx.alloc_tensor", "tirx.decl_tensor"}) {
- OpDef(name).set_attr<FDocTranslate>(tvm::script::printer::op_attr::kOpCallDocTranslate,
- FDocTranslate::FromNative<&TensorOperationDocTranslate>());
- }
-}
-
ffi::Optional<ExprDoc> TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
const ffi::Object*) {
const auto* tensor =
@@ -247,12 +139,9 @@ ffi::Optional<ExprDoc> TensorStoreDocTranslate(DocTranslatorObj* d, ffi::AnyView
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorStoreNode>(input);
TVM_FFI_CHECK(destination == nullptr, TypeError)
<< "printer statement-only node cannot fulfill a destination";
- auto tensor_var = store->dest.as<tvm::tirx::TensorVar>();
- bool scalar = tensor_var.has_value() && IsScalarTensor(d, tensor_var.value());
- ExprDoc tensor = scalar ? ExprDoc(VarDoc(d, store->dest.as<Var>().value(), false))
- : d->Translate(store->dest).value();
+ ExprDoc tensor = d->Translate(store->dest).value();
ExprDoc value = d->Translate(store->value).value();
- ExprDoc lhs = scalar ? tensor : ExprDoc(IndexDoc(tensor, TensorIndices(d, store->indices, true)));
+ ExprDoc lhs = IndexDoc(tensor, TensorIndices(d, store->indices, true));
d->Emit(AssignDoc(lhs, value, std::nullopt), ffi::GetRef<ffi::ObjectRef>(store));
return std::nullopt;
}
@@ -267,9 +156,6 @@ ffi::Optional<ExprDoc> TIRxTensorLoadDocTranslate(DocTranslatorObj* d, ffi::AnyV
const ffi::Object*) {
const auto* load =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorLoadNode>(input);
- if (IsScalarTensor(d, load->source)) {
- return VarDoc(d, load->source.as<Var>().value(), false);
- }
ExprDoc source = d->Translate(load->source).value();
return IndexDoc(source, TensorIndices(d, load->indices));
}
diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h
index f269100e92..0a376f83cf 100644
--- a/src/tirx/script/printer/utils.h
+++ b/src/tirx/script/printer/utils.h
@@ -27,16 +27,4 @@
#include "../../../script/printer/ir/utils.h"
-namespace tvm {
-namespace script {
-namespace printer {
-namespace details {
-
-bool IsScalarTensor(DocTranslatorObj* d, const Expr& source);
-
-} // namespace details
-} // namespace printer
-} // namespace script
-} // namespace tvm
-
#endif // SRC_TIRX_SCRIPT_PRINTER_UTILS_H_
diff --git a/src/tirx/transform/flatten_tensor.cc b/src/tirx/transform/flatten_tensor.cc
index 92f250a914..5e22c7c719 100644
--- a/src/tirx/transform/flatten_tensor.cc
+++ b/src/tirx/transform/flatten_tensor.cc
@@ -88,10 +88,7 @@ class TensorFlattener : public IRMutatorWithAnalyzer {
auto new_buf = pass->Lookup(old_buf.value()).flattened;
if (!old_buf.value().same_as(new_buf)) {
body = SeqStmt({Bind(new_buf, Call(new_buf.type(), decl_tensor_op(),
- {old_buf.value().data(), tvm::Tuple(new_buf->shape),
- DataTypeImm(new_buf->dtype->dtype),
- StringImm(new_buf.scope())},
- {})),
+ {old_buf.value().data()}, {}, {new_buf.type()})),
std::move(body).value()});
body_unchanged = false;
}
@@ -198,19 +195,17 @@ class TensorFlattener : public IRMutatorWithAnalyzer {
InplaceMode inplace_mode) {
const FlatInfo& info = Define(op->var.as_or_throw<TensorVar>());
ffi::Array<Expr> args = tensor_call->args;
- if (args.size() == 4) {
- args.Set(3, Mutate(args[3], inplace_mode).ValueOrUnchanged(args[3]));
+ if (!args.empty()) {
+ args.Set(0, Mutate(args[0], inplace_mode).ValueOrUnchanged(args[0]));
}
if (info.flattened.same_as(op->var.as_or_throw<TensorVar>()) &&
args.same_as(tensor_call->args)) {
return ffi::Unchanged();
}
- args.Set(0, tvm::Tuple(info.flattened->shape, tensor_call->args[0]->loc));
- args.Set(1, DataTypeImm(info.flattened->dtype->dtype, tensor_call->args[1]->loc));
- args.Set(2, StringImm(info.flattened.scope(), tensor_call->args[2]->loc));
+
return Bind(info.flattened.var(),
Call(info.flattened.type(), tirx::alloc_tensor_op(), args, tensor_call->attrs,
- tensor_call->ty_args, tensor_call->loc),
+ {info.flattened.type()}, tensor_call->loc),
op->loc);
}
@@ -234,10 +229,8 @@ class TensorFlattener : public IRMutatorWithAnalyzer {
return ffi::Unchanged();
}
return Bind(info.flattened,
- Call(info.flattened.type(), decl_tensor_op(),
- {std::move(data), tvm::Tuple(info.flattened->shape),
- DataTypeImm(info.flattened->dtype->dtype), StringImm(info.flattened.scope())},
- tensor_call->attrs, tensor_call->ty_args, tensor_call->loc),
+ Call(info.flattened.type(), decl_tensor_op(), {std::move(data)}, tensor_call->attrs,
+ {info.flattened.type()}, tensor_call->loc),
op->loc);
}
diff --git a/src/tirx/transform/lower_thread_allreduce.h b/src/tirx/transform/lower_thread_allreduce.h
index c8e2a2603b..6cb82207fc 100644
--- a/src/tirx/transform/lower_thread_allreduce.h
+++ b/src/tirx/transform/lower_thread_allreduce.h
@@ -137,12 +137,10 @@ class ThreadAllreduceBuilder final : public DialectMutator {
annotations.CopyOnWrite()->dict.Set(tvm::tirx::attr::kVolatile, true);
}
ffi::Array<Expr> args = call->args;
- args.Set(0, tvm::Tuple(replacement->shape, call->args[0]->loc));
- args.Set(1, DataTypeImm(replacement->dtype->dtype, call->args[1]->loc));
- args.Set(2, StringImm(replacement.scope(), call->args[2]->loc));
+
return Bind(replacement.var(),
- Call(replacement.type(), tirx::alloc_tensor_op(), args, annotations, call->ty_args,
- call->loc),
+ Call(replacement.type(), tirx::alloc_tensor_op(), args, annotations,
+ {replacement.type()}, call->loc),
node->loc);
}
@@ -530,10 +528,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
ffi::Array<Stmt> alloc_stmts;
for (TensorVar buf : new_alloc_bufs) {
alloc_stmts.push_back(Bind(
- buf.var(),
- Call(buf.type(), tirx::alloc_tensor_op(),
- {tvm::Tuple(buf->shape), DataTypeImm(buf->dtype->dtype), StringImm(buf.scope())},
- DictAttrs())));
+ buf.var(), Call(buf.type(), tirx::alloc_tensor_op(), {}, DictAttrs(), {buf.type()})));
}
// Prepend allocations before the sequence
for (const auto& s : seq) {
@@ -993,12 +988,10 @@ class DeferredRemapper : public DialectMutator {
annotations.CopyOnWrite()->dict.Set(tvm::tirx::attr::kVolatile, true);
}
ffi::Array<Expr> args = call->args;
- args.Set(0, tvm::Tuple(replacement->shape, call->args[0]->loc));
- args.Set(1, DataTypeImm(replacement->dtype->dtype, call->args[1]->loc));
- args.Set(2, StringImm(replacement.scope(), call->args[2]->loc));
+
return Bind(replacement.var(),
Call(replacement.type(), tirx::alloc_tensor_op(), args, annotations,
- call->ty_args, call->loc),
+ {replacement.type()}, call->loc),
node->loc);
}
}
@@ -1015,13 +1008,10 @@ class DeferredRemapper : public DialectMutator {
.template as_or_throw<Bind>();
if (auto new_buf = GetRemappedTensor(node->var.template as_or_throw<TensorVar>())) {
const CallNode* call = node->value.template as<CallNode>();
- return Bind(
- new_buf.value(),
- Call(new_buf.value().type(), tirx::decl_tensor_op(),
- {call->args[0], tvm::Tuple(new_buf.value()->shape),
- DataTypeImm(new_buf.value()->dtype->dtype), StringImm(new_buf.value().scope())},
- call->attrs, call->ty_args, call->loc),
- node->loc);
+ return Bind(new_buf.value(),
+ Call(new_buf.value().type(), tirx::decl_tensor_op(), {call->args[0]}, call->attrs,
+ {new_buf.value().type()}, call->loc),
+ node->loc);
}
return node;
}
diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc
index 6188cf482f..460c18b7ef 100644
--- a/src/tirx/transform/lower_tirx_cleanup.cc
+++ b/src/tirx/transform/lower_tirx_cleanup.cc
@@ -77,11 +77,9 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
}
auto new_stmt = storage_lower->Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(stmt);
for (const auto& [buf, source] : param_flattened_tensors) {
- new_stmt = SeqStmt({Bind(buf, Call(buf.type(), decl_tensor_op(),
- {source.data(), tvm::Tuple(buf->shape),
- DataTypeImm(buf->dtype->dtype), StringImm(buf.scope())},
- {})),
- std::move(new_stmt)});
+ new_stmt =
+ SeqStmt({Bind(buf, Call(buf.type(), decl_tensor_op(), {source.data()}, {}, {buf.type()})),
+ std::move(new_stmt)});
}
return std::make_pair(new_stmt, new_params);
}
@@ -127,9 +125,8 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode) final {
if (auto source = op->value.as<TensorVar>(); source.has_value()) {
auto tensor = op->var.as_or_throw<TensorVar>();
- Call declaration(tensor.type(), decl_tensor_op(),
- {source.value().data(), tvm::Tuple(tensor->shape),
- DataTypeImm(tensor->dtype->dtype), StringImm(tensor.scope())});
+ Call declaration(tensor.type(), decl_tensor_op(), {source.value().data()}, {},
+ {tensor.type()});
Bind binding(op->var, declaration, op->loc);
return MutateDeclTensor(binding.get(), declaration.get(), inplace_mode)
.ValueOrUnchanged(binding);
@@ -142,9 +139,8 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
// TMEM is already allocated by the pool. Retain an explicit backing pointer so
// ordinary data projections also work for views and function parameters.
Call declaration(tensor.type(), decl_tensor_op(),
- {Call(tensor->DataPtrType(), reinterpret_op(), {call->args[0]}),
- tvm::Tuple(tensor->shape), DataTypeImm(tensor->dtype->dtype),
- StringImm(tensor.scope())});
+ {Call(tensor->DataPtrType(), reinterpret_op(), {call->args[0]})}, {},
+ {tensor.type()});
Bind binding(op->var, declaration, op->loc);
return MutateDeclTensor(binding.get(), declaration.get(), inplace_mode)
.ValueOrUnchanged(binding);
@@ -164,15 +160,13 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
};
auto tensor = mutate(op->var.as_or_throw<TensorVar>());
ffi::Array<Expr> args = tensor_call->args;
- args.Set(0, tvm::Tuple(tensor->shape, tensor_call->args[0]->loc));
- args.Set(1, DataTypeImm(tensor->dtype->dtype, tensor_call->args[1]->loc));
- args.Set(2, StringImm(tensor.scope(), tensor_call->args[2]->loc));
- if (args.size() == 4) {
- args.Set(3, Mutate(args[3], inplace_mode).ValueOrUnchanged(args[3]));
+
+ if (!args.empty()) {
+ args.Set(0, Mutate(args[0], inplace_mode).ValueOrUnchanged(args[0]));
}
return Bind(tensor.var(),
Call(tensor.type(), tirx::alloc_tensor_op(), args, tensor_call->attrs,
- tensor_call->ty_args, tensor_call->loc),
+ {tensor.type()}, tensor_call->loc),
op->loc);
}
@@ -187,10 +181,8 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
return ffi::Unchanged();
}
return Bind(tensor,
- Call(tensor.type(), decl_tensor_op(),
- {std::move(data), tvm::Tuple(tensor->shape), DataTypeImm(tensor->dtype->dtype),
- StringImm(tensor.scope())},
- tensor_call->attrs, tensor_call->ty_args, tensor_call->loc),
+ Call(tensor.type(), decl_tensor_op(), {std::move(data)}, tensor_call->attrs,
+ {tensor.type()}, tensor_call->loc),
op->loc);
}
diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/src/tirx/transform/lower_tvm_builtin.cc
index b9f9432a19..14b9e11f8e 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -198,11 +198,8 @@ class BuiltinLower : public StmtExprMutator {
stmt = SeqStmt({Bind(scope.stack_shape.value(),
Call(scope.stack_shape.value().type(), decl_tensor_op(),
{StackAlloca(scope.stack_shape.value().type()->DataPtrType(),
- "shape", scope.max_sizes.shape_stack),
- tvm::Tuple(scope.stack_shape.value()->shape),
- DataTypeImm(scope.stack_shape.value()->dtype->dtype),
- StringImm(scope.stack_shape.value().scope())},
- {})),
+ "shape", scope.max_sizes.shape_stack)},
+ {}, {scope.stack_shape.value().type()})),
stmt});
}
@@ -274,10 +271,10 @@ class BuiltinLower : public StmtExprMutator {
Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
op = stmt.as<BindNode>();
const auto* tensor_call = op->value.as<CallNode>();
- tvm::Tuple shape = tensor_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = tensor_call->args[1].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(tensor_call->ty_args[0].as_or_throw<TensorType>()->shape);
+ DLDataType dtype = tensor_call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
- ffi::String scope = tensor_call->args[2].as_or_throw<StringImm>()->value;
+ ffi::String scope = tensor_call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
DictAttrs annotations = tensor_call->attrs.value().as_or_throw<DictAttrs>();
if (annotations->dict.count(transform::kDisableLowerTVMBuiltin)) {
if (annotations->dict[transform::kDisableLowerTVMBuiltin].as_or_throw<IntImm>()->value) {
@@ -330,11 +327,8 @@ class BuiltinLower : public StmtExprMutator {
{Call(op->var.as_or_throw<TensorVar>().type()->DataPtrType(), alloc_workspace_op,
{prim::cast(PrimType::Int(32), device_type_.value()),
prim::cast(PrimType::Int(32), device_id_.value()), total_bytes,
- IntImm::Int32(element_type.code()), IntImm::Int32(element_type.bits())}),
- tvm::Tuple(op->var.as_or_throw<TensorVar>()->shape),
- DataTypeImm(op->var.as_or_throw<TensorVar>()->dtype->dtype),
- StringImm(op->var.as_or_throw<TensorVar>().scope())},
- {}, tensor_call->ty_args, tensor_call->loc),
+ IntImm::Int32(element_type.code()), IntImm::Int32(element_type.bits())})},
+ {}, {op->var.as_or_throw<TensorVar>().type()}, tensor_call->loc),
op->loc);
return SeqStmt({alloc_bind, alloc_nullptr_check});
diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc
index d8adb7d9f3..d2b7e82294 100644
--- a/src/tirx/transform/lower_warp_memory.cc
+++ b/src/tirx/transform/lower_warp_memory.cc
@@ -343,8 +343,8 @@ class WarpAccessRewriter : public StmtExprMutator {
// \param tensor_call The matched allocation Call.
// \param body The remaining statements (siblings) that use this tensor.
Stmt Rewrite(const BindNode* op, const CallNode* tensor_call, Stmt body) {
- tvm::Tuple shape = tensor_call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = tensor_call->args[1].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(tensor_call->ty_args[0].as_or_throw<TensorType>()->shape);
+ DLDataType dtype = tensor_call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
tensor_ = op->var.get();
int64_t alloc_size = 1;
@@ -379,15 +379,13 @@ class WarpAccessRewriter : public StmtExprMutator {
new_tensor_ = new_buf;
Stmt rewritten_body = this->Mutate(body, InplaceMode::kDisallow).ValueOrUnchanged(body);
ffi::Array<Expr> args = tensor_call->args;
- if (args.size() == 4) {
- args.Set(3, Mutate(args[3], InplaceMode::kDisallow).ValueOrUnchanged(args[3]));
+ if (!args.empty()) {
+ args.Set(0, Mutate(args[0], InplaceMode::kDisallow).ValueOrUnchanged(args[0]));
}
- args.Set(0, tvm::Tuple(new_buf->shape, tensor_call->args[0]->loc));
- args.Set(1, DataTypeImm(new_buf->dtype->dtype, tensor_call->args[1]->loc));
- args.Set(2, StringImm(new_buf.scope(), tensor_call->args[2]->loc));
+
return SeqStmt({Bind(new_buf.var(),
Call(new_buf.type(), tirx::alloc_tensor_op(), args, tensor_call->attrs,
- tensor_call->ty_args, tensor_call->loc),
+ {new_buf.type()}, tensor_call->loc),
op->loc),
rewritten_body});
}
@@ -672,7 +670,7 @@ class WarpMemoryRewriter : public StmtExprMutator {
const auto* alloc = op->seq[i].as<BindNode>();
if (const auto* call = alloc ? alloc->value.as<CallNode>() : nullptr;
call && call->op.same_as(alloc_tensor_op()) &&
- call->args[2].as_or_throw<StringImm>()->value == "warp") {
+ call->ty_args[0].as_or_throw<TensorType>()->storage_scope == "warp") {
new_storage_scopes_[alloc->var] = "local";
// Gather remaining siblings as the "body" for rewriting.
ffi::Array<Stmt> remaining;
diff --git a/src/tirx/transform/split_host_device.cc b/src/tirx/transform/split_host_device.cc
index 85426d2d07..829a927c9b 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -247,10 +247,7 @@ class HostDeviceSplitter : public StmtExprMutator {
body = SeqStmt(
{Bind(kernel_tensor.as_or_throw<TensorVar>(),
Call(kernel_tensor.as_or_throw<TensorVar>().type(), decl_tensor_op(),
- {data_param.value(), tvm::Tuple(kernel_tensor.as_or_throw<TensorVar>()->shape),
- DataTypeImm(kernel_tensor.as_or_throw<TensorVar>()->dtype->dtype),
- StringImm(kernel_tensor.as_or_throw<TensorVar>().scope())},
- {})),
+ {data_param.value()}, {}, {kernel_tensor.as_or_throw<TensorVar>().type()})),
std::move(body)});
}
Function device_func(kernel_params, SeqStmt(body), kernel_ret_type);
@@ -491,7 +488,7 @@ class DeviceInfoCollector : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const CallNode* call) {
- ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ ffi::String scope = call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
auto storage_scope = runtime::StorageScope::Create(scope);
if (storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ".dyn") {
TVM_FFI_ICHECK(!saw_dyn_shared_alloc_)
@@ -500,8 +497,8 @@ class DeviceInfoCollector : public StmtExprVisitor {
// A zero extent is an extern placeholder; native launch regions supply
// its resource requirement independently of the backing allocation.
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
TVM_FFI_ICHECK_GT(shape->fields.size(), 0);
PrimExpr dyn_size = IntImm::Int32(1);
diff --git a/src/tirx/transform/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc
index 1ea72b3fac..36551a24c1 100644
--- a/src/tirx/transform/stmt_simplify.cc
+++ b/src/tirx/transform/stmt_simplify.cc
@@ -144,7 +144,7 @@ UnchangedOr<Stmt> StmtSimplifier::Mutate_(const ForNode* op, InplaceMode inplace
UnchangedOr<Stmt> StmtSimplifier::Mutate_(const BindNode* op, InplaceMode inplace_mode) {
if (const auto* call = op->value.as<CallNode>()) {
- // Preserve tensor metadata and shape operands; only declaration data is simplified.
+ // Preserve tensor type metadata; only declaration data is simplified.
if (call->op.same_as(tirx::alloc_tensor_op())) return ffi::Unchanged();
if (call->op.same_as(tirx::decl_tensor_op())) {
// The Call and its arguments may be shared even when the Bind is writable.
diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc
index c6eda604f5..27a71ed1d6 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -131,7 +131,7 @@ class LinearAccessPatternFinder final : public StmtExprVisitor {
};
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const CallNode* call) {
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
size_t level = scope_.size();
const VarNode* buf = op->var.get();
tensor_aliases_.Set(op->var, op->var);
@@ -726,13 +726,11 @@ class StoragePlanRewriter : public StmtExprMutator {
// This allocation was merged into another. Emit a DeclTensor
// aliasing the winner's data variable.
TensorVar buf = RemapTensor(op->var.as_or_throw<TensorVar>(), it->second->alloc_var);
- return Bind(
- buf,
- Call(buf.type(), tirx::decl_tensor_op(),
- {it->second->alloc_var.as_or_throw<TensorVar>().data(), tvm::Tuple(buf->shape),
- DataTypeImm(buf->dtype->dtype), StringImm(buf.scope())},
- {}, tensor_call->ty_args, tensor_call->loc),
- op->loc);
+ return Bind(buf,
+ Call(buf.type(), tirx::decl_tensor_op(),
+ {it->second->alloc_var.as_or_throw<TensorVar>().data()}, {}, {buf.type()},
+ tensor_call->loc),
+ op->loc);
}
// If not in alloc_map (e.g. unused), strip entirely.
return Evaluate(0);
@@ -747,13 +745,11 @@ class StoragePlanRewriter : public StmtExprMutator {
auto attrs = Mutate(tensor_call->attrs, InplaceMode::kDisallow)
.as_or_throw<UnchangedOr<ffi::Optional<Attrs>>>()
.ValueOrUnchanged(tensor_call->attrs);
- return Bind(
- tensor,
- Call(tensor.type(), tirx::decl_tensor_op(),
- {it->second->alloc_var.as_or_throw<TensorVar>().data(), tvm::Tuple(tensor->shape),
- DataTypeImm(tensor->dtype->dtype), StringImm(tensor.scope())},
- attrs, tensor_call->ty_args, tensor_call->loc),
- op->loc);
+ return Bind(tensor,
+ Call(tensor.type(), tirx::decl_tensor_op(),
+ {it->second->alloc_var.as_or_throw<TensorVar>().data()}, attrs,
+ {tensor.type()}, tensor_call->loc),
+ op->loc);
}
return StmtExprMutator::Mutate_(op, inplace_mode);
}
@@ -801,7 +797,7 @@ class StoragePlanRewriter : public StmtExprMutator {
};
static bool RequiresExactAllocation(const CallNode* call) {
- if (call->args.size() == 4) return true;
+ if (!call->args.empty()) return true;
const auto& annotations = call->attrs.value().as_or_throw<DictAttrs>()->dict;
// Rebuilt allocations only preserve volatile metadata. Other owners may attach
// constraints that cannot be transferred to merged or resized storage.
@@ -880,7 +876,7 @@ class StoragePlanRewriter : public StmtExprMutator {
continue;
}
const auto* first_call = e->allocs[0]->value.as<CallNode>();
- DLDataType first_dtype = first_call->args[1].as_or_throw<DataTypeImm>()->value;
+ DLDataType first_dtype = first_call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType first_type(first_dtype);
if (e->allocs.size() == 1 && first_type.IsScalableVector()) {
// Scalable vector lanes are runtime-dependent. Keep these allocations exact rather
@@ -891,11 +887,8 @@ class StoragePlanRewriter : public StmtExprMutator {
if (e->is_volatile) {
annotations.Set(tvm::tirx::attr::kVolatile, true);
}
- e->alloc_nest.push_back(Bind(
- buf.var(),
- Call(buf.type(), tirx::alloc_tensor_op(),
- {tvm::Tuple(buf->shape), DataTypeImm(buf->dtype->dtype), StringImm(buf.scope())},
- DictAttrs(annotations))));
+ e->alloc_nest.push_back(Bind(buf.var(), Call(buf.type(), tirx::alloc_tensor_op(), {},
+ DictAttrs(annotations), {buf.type()})));
continue;
}
// Get the allocation size;
@@ -903,7 +896,7 @@ class StoragePlanRewriter : public StmtExprMutator {
PrimType alloc_type = first_type;
for (const BindNode* op : e->allocs) {
const auto* call = op->value.as<CallNode>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
if (element_type.lanes() > alloc_type.lanes()) {
alloc_type = element_type;
@@ -913,12 +906,13 @@ class StoragePlanRewriter : public StmtExprMutator {
bool all_allocs_identical =
std::all_of(e->allocs.begin() + 1, e->allocs.end(), [&](const BindNode* op) -> bool {
const auto* call = op->value.as<CallNode>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
if (dtype != first_dtype) {
return false;
}
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- tvm::Tuple first_shape = first_call->args[0].as_or_throw<tvm::Tuple>();
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
+ tvm::Tuple first_shape =
+ tvm::Tuple(first_call->ty_args[0].as_or_throw<TensorType>()->shape);
if (shape->fields.size() != first_shape->fields.size()) {
return false;
}
@@ -939,18 +933,15 @@ class StoragePlanRewriter : public StmtExprMutator {
if (e->is_volatile) {
annotations.Set(tvm::tirx::attr::kVolatile, true);
}
- e->alloc_nest.push_back(Bind(
- buf.var(),
- Call(buf.type(), tirx::alloc_tensor_op(),
- {tvm::Tuple(buf->shape), DataTypeImm(buf->dtype->dtype), StringImm(buf.scope())},
- DictAttrs(annotations))));
+ e->alloc_nest.push_back(Bind(buf.var(), Call(buf.type(), tirx::alloc_tensor_op(), {},
+ DictAttrs(annotations), {buf.type()})));
} else {
// Build a merged allocation
ffi::Optional<PrimExpr> max_size;
for (const BindNode* op : e->allocs) {
const auto* call = op->value.as<CallNode>();
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
TVM_FFI_ICHECK_EQ(shape->fields.size(), 1)
<< "TensorVar var " << op->var.as_or_throw<TensorVar>().name()
@@ -994,11 +985,8 @@ class StoragePlanRewriter : public StmtExprMutator {
if (e->is_volatile) {
annotations.Set(tvm::tirx::attr::kVolatile, true);
}
- e->alloc_nest.push_back(Bind(
- buf.var(),
- Call(buf.type(), tirx::alloc_tensor_op(),
- {tvm::Tuple(buf->shape), DataTypeImm(buf->dtype->dtype), StringImm(buf.scope())},
- DictAttrs(annotations))));
+ e->alloc_nest.push_back(Bind(buf.var(), Call(buf.type(), tirx::alloc_tensor_op(), {},
+ DictAttrs(annotations), {buf.type()})));
}
}
}
@@ -1029,7 +1017,7 @@ class StoragePlanRewriter : public StmtExprMutator {
}
uint64_t type_bits = e->elem_type.bits() * e->elem_type.lanes();
const auto* call = e->allocs[0]->value.as<CallNode>();
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
PrimExpr alloc_size = MakeConst(shape->fields[0].as_or_throw<PrimExpr>().ty(),
(total_bits + type_bits - 1) / type_bits);
TensorVar buf(e->alloc_var->name, TensorType(e->scope.to_string(), e->elem_type, {alloc_size},
@@ -1046,11 +1034,8 @@ class StoragePlanRewriter : public StmtExprMutator {
if (any_volatile) {
annotations.Set(tvm::tirx::attr::kVolatile, true);
}
- e->alloc_nest.push_back(
- Bind(buf.var(),
- Call(buf.type(), tirx::alloc_tensor_op(),
- {tvm::Tuple(buf->shape), DataTypeImm(buf->dtype->dtype), StringImm(buf.scope())},
- DictAttrs(annotations))));
+ e->alloc_nest.push_back(Bind(buf.var(), Call(buf.type(), tirx::alloc_tensor_op(), {},
+ DictAttrs(annotations), {buf.type()})));
}
// Liveness analysis to find gen and kill point of each variable.
void LivenessAnalysis(const std::vector<StmtEntry>& seq) {
@@ -1129,12 +1114,12 @@ class StoragePlanRewriter : public StmtExprMutator {
const AllocEntry& entry = alloc_info.at(var);
const BindNode* alloc = entry.alloc;
const auto* call = alloc->value.as<CallNode>();
- ffi::String scope = call->args[2].as_or_throw<StringImm>()->value;
+ ffi::String scope = call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
auto storage_scope = StorageScope::Create(scope);
StorageEntry* dst_entry = nullptr;
// inplace detection
if (detect_inplace && !RequiresExactAllocation(call)) {
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
// only one inplace var for s.stmt
bool inplace_found = false;
@@ -1204,7 +1189,7 @@ class StoragePlanRewriter : public StmtExprMutator {
const StorageScope& scope, size_t const_nbits) {
TVM_FFI_ICHECK(op != nullptr);
const auto* call = op->value.as<CallNode>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
// Re-use not successful, allocate a new buffer.
auto entry = std::make_unique<StorageEntry>();
@@ -1223,7 +1208,7 @@ class StoragePlanRewriter : public StmtExprMutator {
bool enable_reuse, bool reuse_require_exact_matched_dtype) {
TVM_FFI_ICHECK(op != nullptr);
const auto* call = op->value.as<CallNode>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
// skip plan for local variable,
// compiler can do a better job with register allocation.
@@ -1312,7 +1297,7 @@ class StoragePlanRewriter : public StmtExprMutator {
// This rules only apply if we are using non special memory
if (e->scope.tag.length() == 0) {
const auto* call = e->allocs[0]->value.as<CallNode>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
// Disable sharing of local memory.
if (e->scope.rank >= StorageRank::kWarp || element_type.IsScalableVector()) {
@@ -1513,8 +1498,8 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const CallNode* call) {
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
tensor_aliases_.Set(op->var, op->var);
PrimExpr extent =
@@ -1527,8 +1512,8 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> DispatchDeclTensor(const BindNode* op,
const CallNode* tensor_call) {
RegisterTensorAlias(op->var.as_or_throw<TensorVar>(), tensor_call->args[0]);
- tvm::Tuple shape = tensor_call->args[1].as_or_throw<tvm::Tuple>();
- DLDataType dtype = tensor_call->args[2].as_or_throw<DataTypeImm>()->value;
+ tvm::Tuple shape = tvm::Tuple(tensor_call->ty_args[0].as_or_throw<TensorType>()->shape);
+ DLDataType dtype = tensor_call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
PrimExpr extent =
shape->fields.size() ? shape->fields.back().as_or_throw<PrimExpr>() : PrimExpr(0);
@@ -2016,8 +2001,8 @@ class VectorTypeRewriter : public StmtExprMutator {
InplaceMode inplace_mode) {
TensorVar new_buf = RemapTensor(op->var.as_or_throw<TensorVar>());
ffi::Array<Expr> args = tensor_call->args;
- if (args.size() == 4) {
- args.Set(3, Mutate(args[3], inplace_mode).ValueOrUnchanged(args[3]));
+ if (!args.empty()) {
+ args.Set(0, Mutate(args[0], inplace_mode).ValueOrUnchanged(args[0]));
}
auto attrs = Mutate(tensor_call->attrs, InplaceMode::kDisallow)
.as_or_throw<UnchangedOr<ffi::Optional<Attrs>>>()
@@ -2026,11 +2011,9 @@ class VectorTypeRewriter : public StmtExprMutator {
attrs.same_as(tensor_call->attrs)) {
return ffi::Unchanged();
}
- args.Set(0, tvm::Tuple(new_buf->shape, tensor_call->args[0]->loc));
- args.Set(1, DataTypeImm(new_buf->dtype->dtype, tensor_call->args[1]->loc));
- args.Set(2, StringImm(new_buf.scope(), tensor_call->args[2]->loc));
+
return Bind(new_buf.var(),
- Call(new_buf.type(), tirx::alloc_tensor_op(), args, attrs, tensor_call->ty_args,
+ Call(new_buf.type(), tirx::alloc_tensor_op(), args, attrs, {new_buf.type()},
tensor_call->loc),
op->loc);
}
@@ -2046,10 +2029,8 @@ class VectorTypeRewriter : public StmtExprMutator {
attrs.same_as(tensor_call->attrs))
return ffi::Unchanged();
return Bind(tensor,
- Call(tensor.type(), tirx::decl_tensor_op(),
- {data, tvm::Tuple(tensor->shape), DataTypeImm(tensor->dtype->dtype),
- StringImm(tensor.scope())},
- attrs, tensor_call->ty_args, tensor_call->loc),
+ Call(tensor.type(), tirx::decl_tensor_op(), {data}, attrs, {tensor.type()},
+ tensor_call->loc),
op->loc);
}
diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/src/tirx/transform/tvm_ffi_binder.cc
index bdee6892ad..6e8b976095 100644
--- a/src/tirx/transform/tvm_ffi_binder.cc
+++ b/src/tirx/transform/tvm_ffi_binder.cc
@@ -618,11 +618,8 @@ void TVMFFIABIBuilder::DecodeAllParams() {
Expr data = DecodeParamDLTensor(tensor.value(), device_type_, device_id_, handle,
func_name_ + "." + param->name, param_path);
decl_tensors_.push_back(
- Bind(tensor.value(),
- Call(tensor.value().type(), tirx::decl_tensor_op(),
- {data, tvm::Tuple(tensor.value()->shape),
- DataTypeImm(tensor.value()->dtype->dtype), StringImm(tensor.value().scope())},
- {})));
+ Bind(tensor.value(), Call(tensor.value().type(), tirx::decl_tensor_op(), {data}, {},
+ {tensor.value().type()})));
PrimExpr size = IntImm(PrimType(tensor.value()->DefaultIndexType()), 1);
for (const auto& extent : tensor.value()->shape) size *= extent;
Stmt alignment =
diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc b/src/tirx/transform/unsupported_dtype_legalize.cc
index 4b95a51c94..7c6280ab6f 100644
--- a/src/tirx/transform/unsupported_dtype_legalize.cc
+++ b/src/tirx/transform/unsupported_dtype_legalize.cc
@@ -92,7 +92,7 @@ class ComputeLegalizePlanner : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const CallNode* call) {
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
+ DLDataType dtype = call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType alloc_dtype(dtype);
// Select intermediate tensors with an unsupported element type.
if (MatchType(alloc_dtype)) {
@@ -239,11 +239,10 @@ class ComputeLegalizer : public StmtExprMutator {
.ValueOrUnchanged(ffi::GetRef<Expr>(op))
.as_or_throw<Call>();
if (op == allocation_to_promote_) {
- auto dtype = call->args[1].as_or_throw<DataTypeImm>();
- auto* node = call.CopyOnWrite();
- node->args.Set(1,
- DataTypeImm(promote_dtype_.WithLanes(PrimType(dtype->value).lanes())->dtype,
- dtype->loc));
+ auto type = call->ty_args[0].as_or_throw<TensorType>();
+ auto updated = ffi::make_object<TensorTypeNode>(*type.get());
+ updated->dtype = promote_dtype_.WithLanes(type->dtype.lanes());
+ call.CopyOnWrite()->ty_args.Set(0, TensorType(std::move(updated)));
}
return ReinferMutatedCallType(call, op, inplace_mode);
}
@@ -603,12 +602,11 @@ class StorageLegalizer : public StmtExprMutator {
Call call = StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Expr>(op))
.as_or_throw<Call>();
- int dtype_index = op->op.same_as(tirx::alloc_tensor_op()) ? 1 : 2;
- auto dtype = call->args[dtype_index].as_or_throw<DataTypeImm>();
- if (MatchType(PrimType(dtype->value))) {
- call.CopyOnWrite()->args.Set(
- dtype_index,
- DataTypeImm(GetStorageUIntDType(PrimType(dtype->value))->dtype, dtype->loc));
+ auto type = call->ty_args[0].as_or_throw<TensorType>();
+ if (MatchType(type->dtype)) {
+ auto updated = ffi::make_object<TensorTypeNode>(*type.get());
+ updated->dtype = GetStorageUIntDType(type->dtype);
+ call.CopyOnWrite()->ty_args.Set(0, TensorType(std::move(updated)));
}
return ReinferMutatedCallType(call, op, inplace_mode);
}
diff --git a/src/tirx/transform/update_pointer_storage_scope.cc b/src/tirx/transform/update_pointer_storage_scope.cc
index 5d1400e46b..be08c8a96c 100644
--- a/src/tirx/transform/update_pointer_storage_scope.cc
+++ b/src/tirx/transform/update_pointer_storage_scope.cc
@@ -81,10 +81,12 @@ UnchangedOr<Expr> UpdatePointerStorageScope::Mutate_(const CallNode* op, Inplace
if (auto it = tensor_scopes_.find(op); it != tensor_scopes_.end()) {
Expr value = std::move(result).ValueOrUnchanged(ffi::GetRef<Expr>(op));
auto call = value.as_or_throw<Call>();
- size_t scope_index = call->op.same_as(tirx::alloc_tensor_op()) ? 2 : 3;
- if (call->args[scope_index].as_or_throw<StringImm>()->value != it->second) {
+ auto type = call->ty_args[0].as_or_throw<TensorType>();
+ if (type->storage_scope != it->second) {
+ auto updated = ffi::make_object<TensorTypeNode>(*type.get());
+ updated->storage_scope = it->second;
auto copy = ffi::make_object<CallNode>(*call.get());
- copy->args.Set(scope_index, StringImm(it->second, call->args[scope_index]->loc));
+ copy->ty_args.Set(0, TensorType(std::move(updated)));
return ReinferMutatedCallType(Expr(std::move(copy)), op, inplace_mode);
}
return value;
diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc
index eb5a9a727c..5644ef9251 100644
--- a/tests/cpp/ir_functor_test.cc
+++ b/tests/cpp/ir_functor_test.cc
@@ -220,11 +220,9 @@ TEST(IRF, StmtVisitor) {
PrimType dtype = PrimType::Float(32);
TensorVar buf("b", TensorType("global", dtype, {z, z}, {}, std::nullopt, 0, 0));
// AllocTensor is flat (no body). Return as SeqStmt with eval.
- return SeqStmt({Bind(buf.var(), Call(buf.type(), tirx::alloc_tensor_op(),
- {tvm::Tuple(buf->shape), DataTypeImm(buf->dtype->dtype),
- StringImm(buf.scope())},
- DictAttrs())),
- eval_body});
+ return SeqStmt(
+ {Bind(buf.var(), Call(buf.type(), tirx::alloc_tensor_op(), {}, DictAttrs(), {buf.type()})),
+ eval_body});
};
v->Visit(fmaketest());
// AllocTensor visits buffer shape at its definition site.
@@ -237,11 +235,9 @@ TEST(IRF, StmtVisitor) {
PrimType dtype = PrimType::Float(32);
tvm::Var buf_var("b", PtrType(dtype));
TensorVar buffer = decl_tensor({16});
- body =
- SeqStmt({Bind(buffer, Call(buffer.type(), tvm::tirx::decl_tensor_op(),
- {buf_var, tvm::Tuple(buffer->shape),
- DataTypeImm(buffer->dtype->dtype), StringImm(buffer.scope())})),
- std::move(body)});
+ body = SeqStmt({Bind(buffer, Call(buffer.type(), tvm::tirx::decl_tensor_op(), {buf_var}, {},
+ {buffer.type()})),
+ std::move(body)});
TensorRegion buffer_region =
TensorRegion(buffer, {Range::FromMinExtent(x + 1, 1)}, TensorRegionType());
s_tir::MatchTensorRegion match_buffer_region(decl_tensor({1}), buffer_region);
@@ -275,10 +271,8 @@ TEST(IRF, StmtExprMutator) {
auto z = x + 1;
PrimType dtype = PrimType::Float(32);
TensorVar buf("b", TensorType("global", dtype, {1, z}, {}, std::nullopt, 0, 0));
- return Bind(buf.var(), Call(buf.type(), tirx::alloc_tensor_op(),
- {tvm::Tuple(buf->shape), DataTypeImm(buf->dtype->dtype),
- StringImm(buf.scope())},
- DictAttrs()));
+ return Bind(buf.var(),
+ Call(buf.type(), tirx::alloc_tensor_op(), {}, DictAttrs(), {buf.type()}));
};
auto fmakeif = [&]() {
@@ -373,9 +367,8 @@ TEST(IRF, StmtExprMutator) {
Stmt eval_body = Evaluate(x + 1);
TensorVar buffer = decl_tensor({16});
tvm::Var buffer_data("buffer_data", buffer.type()->DataPtrType());
- Stmt decl = Bind(buffer, Call(buffer.type(), tvm::tirx::decl_tensor_op(),
- {buffer_data, tvm::Tuple(buffer->shape),
- DataTypeImm(buffer->dtype->dtype), StringImm(buffer.scope())}));
+ Stmt decl = Bind(buffer, Call(buffer.type(), tvm::tirx::decl_tensor_op(), {buffer_data}, {},
+ {buffer.type()}));
Stmt alloc = fmakealloc();
// body is: DeclTensor, AllocTensor, Evaluate
Stmt body = SeqStmt({decl, alloc, eval_body});
@@ -704,11 +697,9 @@ TEST(IRF, StructuralMapBufferDefinition) {
PrimVar m("m", PrimType::Int(32));
TensorVar buffer = fmakebuffer();
Stmt store = TensorStore(buffer, {IntImm::Int32(0)}, FloatImm(dtype, 0));
- Stmt decl =
- SeqStmt({Bind(buffer, Call(buffer.type(), tvm::tirx::decl_tensor_op(),
- {x, tvm::Tuple(buffer->shape), DataTypeImm(buffer->dtype->dtype),
- StringImm(buffer.scope())})),
- store});
+ Stmt decl = SeqStmt(
+ {Bind(buffer, Call(buffer.type(), tvm::tirx::decl_tensor_op(), {x}, {}, {buffer.type()})),
+ store});
auto f_subst = [&](const tvm::Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
if (var.same_as(x)) return ffi::Any(y);
if (var.same_as(n)) return ffi::Any(m);
diff --git a/tests/python/codegen/test_gpu_codegen_allreduce.py b/tests/python/codegen/test_gpu_codegen_allreduce.py
index 9939dec38a..31d089ded5 100644
--- a/tests/python/codegen/test_gpu_codegen_allreduce.py
+++ b/tests/python/codegen/test_gpu_codegen_allreduce.py
@@ -40,7 +40,9 @@ def _reduce_module(d1, d2, d3, is_max=False):
for j in T.thread_binding(d1, thread="threadIdx.z"):
for k in T.thread_binding(d2, thread="threadIdx.y"):
for l in T.thread_binding(d3, thread="threadIdx.x"):
- reduced = T.alloc_tensor((1,), "float32", scope="local")
+ reduced = T.alloc_tensor(
+ ty_args=[T.Tensor((1,), "float32", scope="local")]
+ )
T.gpu_thread_allreduce(
combine,
(T.float32(-3.4028234663852886e38 if is_max else 0),),
diff --git a/tests/python/codegen/test_target_codegen.py b/tests/python/codegen/test_target_codegen.py
index dffc17d99c..e2d3e77423 100644
--- a/tests/python/codegen/test_target_codegen.py
+++ b/tests/python/codegen/test_target_codegen.py
@@ -146,8 +146,10 @@ def test_decl_buffer_offset_preserves_storage_scope(target, qualifier):
"tirx.noalias": True,
}
)
- A = T.decl_tensor((8,), "float32", data=A_ptr)
- B = T.decl_tensor((4,), "float32", data=T.address_of(A[4]))
+ A = T.decl_tensor(A_ptr, ty_args=[T.Tensor((8,), "float32", data=A_ptr)])
+ B = T.decl_tensor(
+ T.address_of(A[4]), ty_args=[T.Tensor((4,), "float32", data=T.address_of(A[4]))]
+ )
B[0] = T.float32(1)
mod = tvm.IRModule({"kernel": kernel})
diff --git a/tests/python/codegen/test_target_codegen_bool.py b/tests/python/codegen/test_target_codegen_bool.py
index 1665567b04..babf2ff3ab 100644
--- a/tests/python/codegen/test_target_codegen_bool.py
+++ b/tests/python/codegen/test_target_codegen_bool.py
@@ -42,7 +42,7 @@ def test_cmp_load_store(target):
T.func_attr({"tirx.noalias": True})
for bx in T.thread_binding(8, thread="blockIdx.x"):
for tx in T.thread_binding(4, thread="threadIdx.x"):
- C = T.alloc_tensor((1,), "bool", scope="local")
+ C = T.alloc_tensor(ty_args=[T.Tensor((1,), "bool", scope="local")])
C[0] = B[bx * 4 + tx] < A[bx * 4 + tx]
D[bx * 4 + tx] = T.Cast("float32", C[0] and T.float32(1.0) < A[bx * 4 + tx])
@@ -55,7 +55,7 @@ def test_cmp_load_store(target):
D: T.Tensor((32,), "float32"),
):
T.func_attr({"tirx.noalias": True})
- C = T.alloc_tensor((32,), "bool")
+ C = T.alloc_tensor(ty_args=[T.Tensor((32,), "bool")])
for i0 in range(32):
C[i0] = B[i0] < A[i0]
for i0 in range(32):
diff --git a/tests/python/codegen/test_target_codegen_c_host.py b/tests/python/codegen/test_target_codegen_c_host.py
index df8dd4f835..5b3b34cb48 100644
--- a/tests/python/codegen/test_target_codegen_c_host.py
+++ b/tests/python/codegen/test_target_codegen_c_host.py
@@ -198,7 +198,7 @@ def test_subroutine_call():
@T.function(private=True)
def subroutine(A_data: T.handle("float32")):
- A = T.decl_tensor(1, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(1, "float32", data=A_data)])
A[0] = 42.0
built = tvm.tirx.build(Module, target="c")
@@ -220,7 +220,7 @@ def test_workspace_allocation_cast():
class Module:
@T.function
def main(A: T.Tensor((256,), "float32")):
- workspace = T.alloc_tensor((256,), "float32", scope="global")
+ workspace = T.alloc_tensor(ty_args=[T.Tensor((256,), "float32", scope="global")])
for i in range(256):
workspace[i] = A[i]
for i in range(256):
@@ -238,7 +238,7 @@ def test_local_alloc_buffer_uses_plain_c_pointer():
class Module:
@T.function
def main(A: T.Tensor((1,), "float32")):
- B = T.alloc_tensor((1,), "float32", scope="local")
+ B = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
for i in range(1):
B[i] = A[i] + T.float32(1)
A[i] = B[i]
diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 31b23b0176..228b98beab 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -530,8 +530,8 @@ def test_crossthread_reduction1(target, *, backend_config):
for i in T.thread_binding(n, thread="blockIdx.x"):
for m_0 in T.thread_binding(nthd, thread="threadIdx.x"):
- partial = T.alloc_tensor((1,), "float32", scope="local")
- reduced = T.alloc_tensor((1,), "float32", scope="local")
+ partial = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ reduced = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
partial[0] = T.float32(0)
for m_1 in range((m + nthd - 1) // nthd):
if m_0 * ((m + nthd - 1) // nthd) + m_1 < m:
@@ -602,8 +602,12 @@ def test_crossthread_reduction2(target, *, backend_config):
for i in T.thread_binding(n, thread="blockIdx.x"):
for k0_0 in T.thread_binding(nthdx, thread="threadIdx.x"):
for k1_0 in T.thread_binding(nthdy, thread="threadIdx.y"):
- partial = T.alloc_tensor((1,), "float32", scope="local")
- reduced = T.alloc_tensor((1,), "float32", scope="local")
+ partial = T.alloc_tensor(
+ ty_args=[T.Tensor((1,), "float32", scope="local")]
+ )
+ reduced = T.alloc_tensor(
+ ty_args=[T.Tensor((1,), "float32", scope="local")]
+ )
partial[0] = T.float32(0)
for k0_1, k1_1 in T.grid(
(k0 + nthdx - 1) // nthdx, (k1 + nthdy - 1) // nthdy
@@ -1147,7 +1151,7 @@ def test_try_unaligned_vector_load(*, backend_config):
def test_cuda_thread_sync_inside_condition(*, backend_config):
@T.function
def func2(A: T.Tensor((4, 4), "float32")) -> None:
- A_shared = T.alloc_tensor((4, 4), "float32", scope="shared")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor((4, 4), "float32", scope="shared")])
for bx in T.thread_binding(1, "blockIdx.x"):
for tx in T.thread_binding(32, "threadIdx.x"):
if T.gpu_thread_invariant(A[0, 0] > 1.0):
@@ -1159,7 +1163,7 @@ def test_cuda_thread_sync_inside_condition(*, backend_config):
@T.function
def func3(A: T.Tensor((4, 4), "float32")) -> None:
- A_shared = T.alloc_tensor((4, 4), "float32", scope="shared")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor((4, 4), "float32", scope="shared")])
for bx in T.thread_binding(1, "blockIdx.x"):
for tx in T.thread_binding(32, "threadIdx.x"):
while T.gpu_thread_invariant(A[0, 0] > 1.0):
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp4.py b/tests/python/codegen/test_target_codegen_cuda_fp4.py
index db0352aed6..11cd49be06 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp4.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp4.py
@@ -234,7 +234,7 @@ def test_e2m1_scalar_buffer_offset():
@T.function
def func(A_raw: T.Tensor((n // 2,), "uint8"), B: T.Tensor((n,), "float16")):
T.func_attr({"tir.noalias": True})
- A = T.decl_tensor((n,), "float4_e2m1fn", data=A_raw.data)
+ A = T.decl_tensor(A_raw.data, ty_args=[T.Tensor((n,), "float4_e2m1fn", data=A_raw.data)])
for bx in T.thread_binding(n // 32, thread="blockIdx.x"):
for tx in T.thread_binding(32, thread="threadIdx.x"):
B[bx * 32 + tx] = T.Cast("float16", A[bx * 32 + tx])
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp8.py b/tests/python/codegen/test_target_codegen_cuda_fp8.py
index 0f0f748a88..d6e7616c4e 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp8.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp8.py
@@ -347,7 +347,7 @@ class BaseFP8E4M3QuantScaleOnly:
):
for row in T.thread_binding(rows, thread="blockIdx.x"):
for group in T.thread_binding(groups, thread="threadIdx.x"):
- maximum = T.alloc_tensor((1,), model_dtype, scope="local")
+ maximum = T.alloc_tensor(ty_args=[T.Tensor((1,), model_dtype, scope="local")])
maximum[0] = T.Cast(model_dtype, 0)
for k in range(group_size):
if group * group_size + k < columns:
@@ -498,7 +498,7 @@ class TestFP8e4x4QuantDequantScale(BaseFP8E4M3QuantScaleOnly):
def test_const(dtype):
@T.function
def func(A: T.Tensor((4,), dtype)) -> None:
- A_local = T.alloc_tensor((4,), dtype=dtype, scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor((4,), dtype, scope="local")])
for tx in T.thread_binding(0, 4, "threadIdx.x"):
for i in T.vectorized(4):
A_local[i] = T.float32(1.0).astype(dtype)
@@ -563,8 +563,8 @@ def test_moe_gemv_shfl_down_illegal_instr():
for block in T.thread_binding(spatial_size // 4, thread="blockIdx.x"):
for spatial in T.thread_binding(4, thread="threadIdx.y"):
for reduction in T.thread_binding(64, thread="threadIdx.x"):
- partial = T.alloc_tensor((1,), "float16", scope="local")
- reduced = T.alloc_tensor((1,), "float16", scope="local")
+ partial = T.alloc_tensor(ty_args=[T.Tensor((1,), "float16", scope="local")])
+ reduced = T.alloc_tensor(ty_args=[T.Tensor((1,), "float16", scope="local")])
partial[0] = T.float16(0)
for k in range(reduce_size // 64):
partial[0] = partial[0] + x[0, k * 64 + reduction] * (
diff --git a/tests/python/codegen/test_target_codegen_device.py b/tests/python/codegen/test_target_codegen_device.py
index 59982fe1e0..71fa043cfc 100644
--- a/tests/python/codegen/test_target_codegen_device.py
+++ b/tests/python/codegen/test_target_codegen_device.py
@@ -68,7 +68,7 @@ def test_add_pipeline():
def main(A: T.Tensor((n,)), B: T.Tensor((), "float32"), D: T.Tensor((n,))):
T.func_attr({"tirx.noalias": True})
- C = T.alloc_tensor((n,))
+ C = T.alloc_tensor(ty_args=[T.Tensor((n,))])
for i0_0 in T.thread_binding((n + 255) // 256, thread="blockIdx.x"):
for i0_1 in T.thread_binding(256, thread="threadIdx.x"):
if i0_0 * 256 + i0_1 < n:
diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py
index d309adc3eb..e5df6f7b7c 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -77,7 +77,7 @@ def test_llvm_intrin():
class Module:
@T.function
def main(A: T.handle("float32")):
- A_buf = T.decl_tensor((4,), "float32", data=A)
+ A_buf = T.decl_tensor(A, ty_args=[T.Tensor((4,), "float32", data=A)])
T.evaluate(T.Call("tirx.prefetch", [T.address_of(A_buf[0]), 0, 3, 1], ty="void"))
fcode = tvm.compile(Module)
@@ -128,7 +128,7 @@ def test_llvm_lookup_intrin():
class Module:
@T.function
def main(A: T.handle("uint8x8")):
- A_buf = T.decl_tensor((1,), "uint8x8", data=A)
+ A_buf = T.decl_tensor(A, ty_args=[T.Tensor((1,), "uint8x8", data=A)])
T.evaluate(T.call_llvm_pure_intrin("llvm.ctpop.v8i8", A_buf[0], ty="uint8x8"))
fcode = tvm.compile(Module, None)
@@ -160,7 +160,7 @@ def test_llvm_multi_parallel():
@T.function
def main(A: T.Tensor((128,), "float32"), C: T.Tensor((128,), "float32")):
T.func_attr({"tirx.noalias": True})
- B = T.alloc_tensor((128,))
+ B = T.alloc_tensor(ty_args=[T.Tensor((128,))])
for i0_0_0 in T.parallel(1):
for ax0 in range(128):
B[ax0] = A[ax0] + T.float32(1.0)
@@ -268,7 +268,7 @@ def test_llvm_temp_space():
@T.function
def main(A: T.Tensor((1024,), "float32"), C: T.Tensor((1024,), "float32")):
T.func_attr({"tirx.noalias": True})
- B = T.alloc_tensor((1024,))
+ B = T.alloc_tensor(ty_args=[T.Tensor((1024,))])
for i in range(1024):
B[i] = A[i] + T.float32(1.0)
for i in range(1024):
@@ -390,7 +390,7 @@ def test_rank_zero():
compute: T.Tensor((), "float32"),
):
T.func_attr({"tirx.noalias": True})
- C = T.alloc_tensor(())
+ C = T.alloc_tensor(ty_args=[T.Tensor(())])
C[()] = T.float32(0.0)
for k in range(64):
C[()] = C[()] + A[k] * scale[()]
@@ -826,9 +826,7 @@ def test_llvm_large_stack_allocation_uses_64bit_extent(extent):
@T.function
def main(A: T.Tensor((1,), "float32")):
B = T.alloc_tensor(
- (extent,),
- "float32",
- scope="global",
+ ty_args=[T.Tensor((extent,), "float32", scope="global")],
annotations={"disable_lower_builtin": True},
)
A[0] = B[extent - 1]
@@ -950,11 +948,11 @@ def test_llvm_assume():
@T.function
def main(A: T.Tensor((4, 4), "int32"), B: T.Tensor((14,), "int32")):
T.func_attr({"tirx.noalias": True})
- A_1 = T.decl_tensor((16,), "int32", data=A.data)
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((16,), "int32", data=A.data)])
for axis0, axis1 in T.grid(4, 4):
T.assume(axis0 < 3 or axis1 < 2 or A_1[axis0 * 4 + axis1] == 0)
for i in range(14):
- B_1 = T.decl_tensor((14,), "int32", data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((14,), "int32", data=B.data)])
B_1[i] = A_1[i] * 2
m = tvm.compile(Module, target="llvm")
@@ -974,8 +972,8 @@ def test_debug_symbol_for_float64():
@T.function
def main(a: T.handle("float64"), b: T.handle("float64"), n: T.int64):
T.func_attr({"calling_conv": 2})
- A = T.decl_tensor(16, "float64", data=a)
- B = T.decl_tensor(16, "float64", data=b)
+ A = T.decl_tensor(a, ty_args=[T.Tensor(16, "float64", data=a)])
+ B = T.decl_tensor(b, ty_args=[T.Tensor(16, "float64", data=b)])
for i in range(n):
B[i] = A[i]
@@ -990,7 +988,7 @@ def test_debug_symbol_for_buffer_var():
class Module:
@T.function
def main(A: T.Tensor((16,), "float32"), B: T.Tensor((16,), "float32")):
- C = T.alloc_tensor((16,), "float32")
+ C = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
for i in T.parallel(16):
C[i] = A[i]
B[i] = C[i]
@@ -1011,7 +1009,7 @@ def test_subroutine_call():
# The calling_conv parameter is to prevent MakePackedAPI
# from changing the call signature of the subroutine.
T.func_attr({"calling_conv": -1})
- A = T.decl_tensor(1, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(1, "float32", data=A_data)])
A[0] = 42.0
target = "llvm"
@@ -1093,7 +1091,7 @@ def test_invalid_volatile_masked_buffer_load():
class Module:
@T.function
def main(B: T.Tensor([4])):
- A = T.alloc_tensor((4,), annotations={"tirx.volatile": True})
+ A = T.alloc_tensor(ty_args=[T.Tensor((4,))], annotations={"tirx.volatile": True})
B[0:4] = T.call_intrin(
"tirx.masked_load", A, T.Ramp(0, 1, 4), T.Broadcast(T.bool(True), 4), ty="float32x4"
)
@@ -1109,8 +1107,8 @@ def test_invalid_volatile_masked_decl_buffer_load():
class Module:
@T.function
def main(B: T.Tensor([4])):
- A = T.alloc_tensor((4,), annotations={"tirx.volatile": True})
- A_alias = T.decl_tensor((4,), data=A.data)
+ A = T.alloc_tensor(ty_args=[T.Tensor((4,))], annotations={"tirx.volatile": True})
+ A_alias = T.decl_tensor(A.data, ty_args=[T.Tensor((4,), data=A.data)])
B[0:4] = T.call_intrin(
"tirx.masked_load",
A_alias,
@@ -1130,7 +1128,7 @@ def test_invalid_volatile_masked_tensor_store():
class Module:
@T.function
def main():
- A = T.alloc_tensor((4,), annotations={"tirx.volatile": True})
+ A = T.alloc_tensor(ty_args=[T.Tensor((4,))], annotations={"tirx.volatile": True})
T.evaluate(
T.call_intrin(
"tirx.masked_store",
diff --git a/tests/python/codegen/test_target_codegen_metal.py b/tests/python/codegen/test_target_codegen_metal.py
index 89091055f7..f0aa4f83cf 100644
--- a/tests/python/codegen/test_target_codegen_metal.py
+++ b/tests/python/codegen/test_target_codegen_metal.py
@@ -331,10 +331,10 @@ def test_codegen_simdgroup_buffer_data():
"tirx.kernel_launch_params": [],
}
)
- A = T.alloc_tensor((64,), "float16", scope="shared")
- A_frag = T.alloc_tensor((64,), "float16", scope="metal.simdgroup")
- B_frag = T.alloc_tensor((64,), "float16", scope="metal.simdgroup")
- C_frag = T.alloc_tensor((64,), "float16", scope="metal.simdgroup")
+ A = T.alloc_tensor(ty_args=[T.Tensor((64,), "float16", scope="shared")])
+ A_frag = T.alloc_tensor(ty_args=[T.Tensor((64,), "float16", scope="metal.simdgroup")])
+ B_frag = T.alloc_tensor(ty_args=[T.Tensor((64,), "float16", scope="metal.simdgroup")])
+ C_frag = T.alloc_tensor(ty_args=[T.Tensor((64,), "float16", scope="metal.simdgroup")])
T.metal.make_filled_simdgroup_matrix(C_frag.data, 0, T.float32(0), 8, 8)
T.metal.simdgroup_load(A_frag.data, 0, A.data, 8, 8, 8, T.bool(False))
T.metal.simdgroup_store(C_frag.data, 0, A.data, 8, 8, 8, T.bool(False))
@@ -371,7 +371,9 @@ def test_bounded_symbolic_stack_allocation():
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((T.min(n, 64), 2), "float32", scope="local")
+ scratch = T.alloc_tensor(
+ ty_args=[T.Tensor((T.min(n, 64), 2), "float32", scope="local")]
+ )
T.evaluate(scratch.data)
source = _build_metal(Module).inspect_source()
@@ -398,7 +400,7 @@ def test_bound_symbolic_stack_allocation(bounded):
# Common subexpression elimination can hoist the bounded extent.
extent: T.let[T.int32] = T.min(n, limit)
elements: T.let[T.int32] = extent * 2
- scratch = T.alloc_tensor((elements,), "float32", scope="local")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((elements,), "float32", scope="local")])
T.evaluate(scratch.data)
if bounded:
@@ -428,7 +430,7 @@ def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded):
"tirx.is_global_func": True,
}
)
- state = T.alloc_tensor((1,), "int32", scope="local")
+ state = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
state[0] = 0
snapshot: T.let[T.int32] = state[0]
state[0] = 32
@@ -436,9 +438,13 @@ def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded):
# The snapshot is immutable, but the buffer it read has changed.
# Substituting the load would incorrectly reduce this extent to 1.
scratch = T.alloc_tensor(
- (T.min(T.max(difference, 1), 32 if bounded else 2147483647),),
- "float32",
- scope=scope,
+ ty_args=[
+ T.Tensor(
+ (T.min(T.max(difference, 1), 32 if bounded else 2147483647),),
+ "float32",
+ scope=scope,
+ )
+ ]
)
scratch[31] = 1.0
@@ -469,7 +475,9 @@ def test_bounded_uint64_symbolic_stack_allocation():
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((T.min(n, T.uint64(64)),), "float32", scope="local")
+ scratch = T.alloc_tensor(
+ ty_args=[T.Tensor((T.min(n, T.uint64(64)),), "float32", scope="local")]
+ )
T.evaluate(scratch.data)
source = _build_metal(Module).inspect_source()
@@ -490,7 +498,7 @@ def test_unbounded_symbolic_stack_allocation_rejected():
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((n,), "float32", scope="local")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32", scope="local")])
scratch[0] = 1.0
T.evaluate(scratch[0])
@@ -515,7 +523,7 @@ def test_unbounded_uint64_symbolic_stack_allocation_rejected():
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((n,), "float32", scope="local")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32", scope="local")])
scratch[0] = 1.0
T.evaluate(scratch[0])
@@ -541,7 +549,7 @@ def test_nonpositive_stack_allocation_rejected(extent):
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((extent,), "float32", scope="local")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((extent,), "float32", scope="local")])
T.evaluate(scratch.data)
with pytest.raises(
@@ -566,9 +574,13 @@ def test_stack_allocation_element_count_overflow_rejected():
}
)
scratch = T.alloc_tensor(
- (T.min(n, 1 << 30), T.min(m, 1 << 30), T.min(k, 1 << 30)),
- "uint8",
- scope="local",
+ ty_args=[
+ T.Tensor(
+ (T.min(n, 1 << 30), T.min(m, 1 << 30), T.min(k, 1 << 30)),
+ "uint8",
+ scope="local",
+ )
+ ]
)
T.evaluate(scratch.data)
@@ -592,11 +604,15 @@ def test_codegen_pointer_byte_offsets_preserve_storage_scope():
"tirx.kernel_launch_params": [],
}
)
- shared = T.alloc_tensor((16,), "float16", scope="shared")
+ shared = T.alloc_tensor(ty_args=[T.Tensor((16,), "float16", scope="shared")])
typed_alias = T.ptr_byte_offset(shared.data, 4, ty=shared.data.ty)
- typed_buffer = T.decl_tensor((14,), "float16", data=typed_alias, scope="shared")
+ typed_buffer = T.decl_tensor(
+ typed_alias, ty_args=[T.Tensor((14,), "float16", scope="shared")]
+ )
void_alias = T.handle_add_byte_offset(shared.data, 8)
- void_buffer = T.decl_tensor((12,), "float16", data=void_alias, scope="shared")
+ void_buffer = T.decl_tensor(
+ void_alias, ty_args=[T.Tensor((12,), "float16", scope="shared")]
+ )
typed_buffer[0] = T.float16(1)
void_buffer[0] = T.float16(2)
@@ -620,11 +636,15 @@ def test_pointer_byte_offsets_execute_in_threadgroup_memory():
def main(A: T.Tensor((16,), "float32"), B: T.Tensor((16,), "float32")):
for bx in T.thread_binding(1, thread="blockIdx.x"):
for tx in T.thread_binding(1, thread="threadIdx.x"):
- shared = T.alloc_tensor((16,), "float32", scope="shared")
+ shared = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32", scope="shared")])
typed_alias = T.ptr_byte_offset(shared.data, 4, ty=shared.data.ty)
- typed_buffer = T.decl_tensor((15,), "float32", data=typed_alias, scope="shared")
+ typed_buffer = T.decl_tensor(
+ typed_alias, ty_args=[T.Tensor((15,), "float32", scope="shared")]
+ )
void_alias = T.handle_add_byte_offset(shared.data, 8)
- void_buffer = T.decl_tensor((14,), "float32", data=void_alias, scope="shared")
+ void_buffer = T.decl_tensor(
+ void_alias, ty_args=[T.Tensor((14,), "float32", scope="shared")]
+ )
shared[0] = A[0]
typed_buffer[0] = A[1]
void_buffer[0] = A[2]
diff --git a/tests/python/codegen/test_target_codegen_rocm.py b/tests/python/codegen/test_target_codegen_rocm.py
index 279fad75a5..916f8e3f2e 100644
--- a/tests/python/codegen/test_target_codegen_rocm.py
+++ b/tests/python/codegen/test_target_codegen_rocm.py
@@ -118,9 +118,9 @@ def test_rocm_warp_shuffle():
):
for bx in T.thread_binding(1, thread="blockIdx.x"):
for tx in T.thread_binding(32, thread="threadIdx.x"):
- A_local = T.alloc_tensor((1,), "float32", scope="local")
- mask = T.alloc_tensor((1,), "uint32", scope="local")
- t0 = T.alloc_tensor((1,), "float32", scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ mask = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="local")])
+ t0 = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
A_local[0] = A[tx]
A_local[0] = T.gpu_warp_shuffle(mask[0], A_local[0], 0, 32, 32)
A[tx] = A_local[0]
diff --git a/tests/python/codegen/test_target_codegen_vulkan.py b/tests/python/codegen/test_target_codegen_vulkan.py
index 04d875a595..32d6e947d2 100644
--- a/tests/python/codegen/test_target_codegen_vulkan.py
+++ b/tests/python/codegen/test_target_codegen_vulkan.py
@@ -278,7 +278,7 @@ def test_vulkan_while_if():
@T.function
def while_if_gpu(A: T.Tensor((1,), "int32"), B: T.Tensor((1,), "int32")):
for bx in T.thread_binding(1, thread="blockIdx.x"):
- iterations = T.decl_tensor((1,), "int32", scope="local")
+ iterations = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
iterations[0] = 0
B[0] = 0
while iterations[0] < T.if_then_else(A[0] > 0, 10, 20):
@@ -465,11 +465,11 @@ def test_cooperative_matrix(out_dtype):
@T.function
def main(X: T.Tensor((16, 32), "float16"), W: T.Tensor((32, 16), "float16"), compute: T.Tensor((16, 16), out_dtype)):
T.func_attr({"tirx.noalias": True})
- X_shared = T.alloc_tensor((16, 32), "float16", scope="shared")
- W_shared = T.alloc_tensor((32, 16), "float16", scope="shared")
- X_shared_wmma_matrix_a = T.alloc_tensor((16, 32), "float16", scope="wmma.matrix_a")
- W_shared_wmma_matrix_b = T.alloc_tensor((32, 16), "float16", scope="wmma.matrix_b")
- compute_wmma_accumulator = T.alloc_tensor((16, 16), out_dtype, scope="wmma.accumulator")
+ X_shared = T.alloc_tensor(ty_args=[T.Tensor((16, 32), "float16", scope="shared")])
+ W_shared = T.alloc_tensor(ty_args=[T.Tensor((32, 16), "float16", scope="shared")])
+ X_shared_wmma_matrix_a = T.alloc_tensor(ty_args=[T.Tensor((16, 32), "float16", scope="wmma.matrix_a")])
+ W_shared_wmma_matrix_b = T.alloc_tensor(ty_args=[T.Tensor((32, 16), "float16", scope="wmma.matrix_b")])
+ compute_wmma_accumulator = T.alloc_tensor(ty_args=[T.Tensor((16, 16), out_dtype, scope="wmma.accumulator")])
for i_0_j_0_fused in T.thread_binding(1, thread="blockIdx.x"):
T.gpu_fill_fragment(compute_wmma_accumulator.data, 16, 16, 16, 0, T.float32(0.0))
for k_0 in range(2):
@@ -521,8 +521,8 @@ def test_codegen_decl_buffer():
@T.function
def kernel():
T.func_attr({"calling_conv": 2, "global_symbol": "kernel", "tirx.noalias": True})
- A = T.alloc_tensor((256,), dtype="float32", scope="local")
- A_buf = T.decl_tensor([256], dtype="float32", scope="local", data=A.data)
+ A = T.alloc_tensor(ty_args=[T.Tensor((256,), "float32", scope="local")])
+ A_buf = T.decl_tensor(A.data, ty_args=[T.Tensor([256], "float32", scope="local")])
A_buf[0] = T.float32(1)
T.evaluate(A_buf[0])
@@ -534,8 +534,8 @@ def test_codegen_decl_buffer():
class ParameterBacked:
@T.function
def main(A: T.Tensor((1,), "float32"), B: T.Tensor((1,), "float32")):
- A_buf = T.decl_tensor([1], dtype="float32", data=A.data)
- B_buf = T.decl_tensor([1], dtype="float32", data=B.data)
+ A_buf = T.decl_tensor(A.data, ty_args=[T.Tensor([1], "float32", data=A.data)])
+ B_buf = T.decl_tensor(B.data, ty_args=[T.Tensor([1], "float32", data=B.data)])
for tx in T.thread_binding(1, thread="threadIdx.x"):
B_buf[tx] = A_buf[tx]
@@ -551,7 +551,7 @@ def test_codegen_static_shared_memory():
class Module:
@T.function
def main(A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")):
- A_shared = T.alloc_tensor((128,), dtype="float32", scope="shared")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared")])
for bx in T.thread_binding(1, thread="blockIdx.x"):
for tx in T.thread_binding(128, thread="threadIdx.x"):
diff --git a/tests/python/codegen/test_target_codegen_webgpu.py b/tests/python/codegen/test_target_codegen_webgpu.py
index 80024df01e..e20f259c2f 100644
--- a/tests/python/codegen/test_target_codegen_webgpu.py
+++ b/tests/python/codegen/test_target_codegen_webgpu.py
@@ -60,7 +60,9 @@ def test_bounded_symbolic_stack_allocation():
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((T.min(n, 64), 2), "float32", scope="local")
+ scratch = T.alloc_tensor(
+ ty_args=[T.Tensor((T.min(n, 64), 2), "float32", scope="local")]
+ )
T.evaluate(scratch.data)
source = _build_webgpu(Module).inspect_source()
@@ -86,9 +88,9 @@ def test_bound_symbolic_allocation(scope, bounded):
)
# Common subexpression elimination can hoist the bounded extent.
extent: T.let[T.int32] = T.min(n, limit)
- first = T.alloc_tensor((extent * 2,), "float32", scope=scope)
+ first = T.alloc_tensor(ty_args=[T.Tensor((extent * 2,), "float32", scope=scope)])
elements: T.let[T.int32] = extent * 2
- second = T.alloc_tensor((elements,), "float32", scope=scope)
+ second = T.alloc_tensor(ty_args=[T.Tensor((elements,), "float32", scope=scope)])
first[0] = 1.0
second[0] = first[0]
@@ -119,7 +121,7 @@ def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded):
"tirx.is_global_func": True,
}
)
- state = T.alloc_tensor((1,), "int32", scope="local")
+ state = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
state[0] = 0
snapshot: T.let[T.int32] = state[0]
state[0] = 32
@@ -127,9 +129,13 @@ def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded):
# The snapshot is immutable, but the buffer it read has changed.
# Substituting the load would incorrectly reduce this extent to 1.
scratch = T.alloc_tensor(
- (T.min(T.max(difference, 1), 32 if bounded else 2147483647),),
- "float32",
- scope=scope,
+ ty_args=[
+ T.Tensor(
+ (T.min(T.max(difference, 1), 32 if bounded else 2147483647),),
+ "float32",
+ scope=scope,
+ )
+ ]
)
scratch[31] = 1.0
@@ -162,7 +168,7 @@ def test_bound_symbolic_workgroup_allocation_respects_target_limit(target_limit)
)
extent: T.let[T.int32] = T.min(n, 64)
elements: T.let[T.int32] = extent * 2
- scratch = T.alloc_tensor((elements,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((elements,), "float32", scope="shared")])
scratch[0] = 1.0
target = {"kind": "webgpu", "max_shared_memory_per_block": target_limit}
@@ -190,7 +196,7 @@ def test_unbounded_symbolic_stack_allocation_rejected():
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((n,), "float32", scope="local")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32", scope="local")])
scratch[0] = 1.0
T.evaluate(scratch[0])
@@ -215,7 +221,7 @@ def test_nonpositive_stack_allocation_rejected(extent):
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((extent,), "float32", scope="local")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((extent,), "float32", scope="local")])
T.evaluate(scratch.data)
with pytest.raises(
@@ -239,9 +245,13 @@ def test_stack_allocation_element_count_overflow_rejected():
}
)
scratch = T.alloc_tensor(
- (T.min(n, 1 << 30), T.min(m, 1 << 30), T.min(k, 1 << 30)),
- "uint8",
- scope="local",
+ ty_args=[
+ T.Tensor(
+ (T.min(n, 1 << 30), T.min(m, 1 << 30), T.min(k, 1 << 30)),
+ "uint8",
+ scope="local",
+ )
+ ]
)
T.evaluate(scratch.data)
@@ -265,7 +275,9 @@ def test_stack_allocation_byte_size_overflow_rejected():
}
)
scratch = T.alloc_tensor(
- (T.min(n, 1 << 30), T.min(m, 1 << 30), 4), "float32", scope="local"
+ ty_args=[
+ T.Tensor((T.min(n, 1 << 30), T.min(m, 1 << 30), 4), "float32", scope="local")
+ ]
)
T.evaluate(scratch.data)
@@ -288,7 +300,7 @@ def test_workgroup_allocation_at_target_limit():
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((8192,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((8192,), "float32", scope="shared")])
scratch[0] = 1.0
source = _build_webgpu(Module).inspect_source()
@@ -308,8 +320,8 @@ def test_total_workgroup_allocation_above_target_limit_rejected():
"tirx.is_global_func": True,
}
)
- first = T.alloc_tensor((4096,), "float32", scope="shared")
- second = T.alloc_tensor((4097,), "float32", scope="shared")
+ first = T.alloc_tensor(ty_args=[T.Tensor((4096,), "float32", scope="shared")])
+ second = T.alloc_tensor(ty_args=[T.Tensor((4097,), "float32", scope="shared")])
first[0] = 1.0
second[0] = 2.0
@@ -333,8 +345,8 @@ def test_workgroup_allocation_accounts_for_declaration_alignment():
"tirx.is_global_func": True,
}
)
- first = T.alloc_tensor((1,), "float32", scope="shared")
- second = T.alloc_tensor((1,), "float32", scope="shared")
+ first = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="shared")])
+ second = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="shared")])
first[0] = 1.0
second[0] = 2.0
@@ -358,7 +370,7 @@ def test_workgroup_allocation_uses_target_limit():
"tirx.is_global_func": True,
}
)
- scratch = T.alloc_tensor((16384,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((16384,), "float32", scope="shared")])
scratch[0] = 1.0
_build_webgpu(Module, {"kind": "webgpu", "max_shared_memory_per_block": 65536})
diff --git a/tests/python/relax/backend/adreno/mod_utils.py b/tests/python/relax/backend/adreno/mod_utils.py
index 5cfdb5484b..9efa163677 100644
--- a/tests/python/relax/backend/adreno/mod_utils.py
+++ b/tests/python/relax/backend/adreno/mod_utils.py
@@ -759,7 +759,7 @@ def get_dequant_matmul_module(K, N):
T.func_attr({"tirx.noalias": T.bool(True)})
# with Ts.sblock("root"):
- compute = T.alloc_tensor((T.int64(K), T.int64(N)), "float16")
+ compute = T.alloc_tensor(ty_args=[T.Tensor((T.int64(K), T.int64(N)), "float16")])
for i0, i1 in T.grid(T.int64(K), T.int64(N)):
with Ts.sblock("compute"):
v_i0, v_i1 = Ts.axis.remap("SS", [i0, i1])
@@ -821,7 +821,9 @@ def get_dequant_vec_matmul_module(K, N):
T.func_attr({"tirx.noalias": T.bool(True)})
# with Ts.sblock("root"):
- compute = T.alloc_tensor((T.int64(K), vocab_size_dequantize), "float16")
+ compute = T.alloc_tensor(
+ ty_args=[T.Tensor((T.int64(K), vocab_size_dequantize), "float16")]
+ )
for i0, i1 in T.grid(T.int64(K), vocab_size_dequantize):
with Ts.sblock("compute"):
v_i0, v_i1 = Ts.axis.remap("SS", [i0, i1])
diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py
index 4a19046895..b24523ea7e 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -723,7 +723,7 @@ def test_call():
"""
x = I.dynamic("x", dtype="int64")
a = I.Var("a", R.Tensor((1, x, 3), dtype="float32"))
-I.Call("relax.call_tir", [Module.tir_func, (a, x)], ty_args=[R.Tensor((1, x, 3), dtype="float32")], ty=R.Tensor((1, x, 3), dtype="float32"))
+R.call_tir(Module.tir_func, (a, x), ty_args=[R.Tensor((1, x, 3), dtype="float32")])
""",
)
_assert_print_lines(
@@ -752,7 +752,7 @@ def test_call_tir_with_grad():
"""
v0 = I.Var("v0", R.Tensor((54, 96), dtype="float32"))
x = I.dynamic("x", dtype="int64")
-I.Call("relax.call_tir_with_grad", [Module.tir_func, (v0,)], attrs=I.make_node("relax.attrs.CallTIRWithGradAttrs", te_grad_kwargs={"k": 1.0, "x": x}, te_grad_name="grad_func"), ty_args=[R.Tensor((54, 96), dtype="float32")], ty=R.Tensor((54, 96), dtype="float32"))
+R.call_tir_with_grad(Module.tir_func, (v0,), ty_args=[R.Tensor((54, 96), dtype="float32")], te_grad_kwargs={"k": 1.0, "x": x}, te_grad_name="grad_func")
""",
)
@@ -779,7 +779,7 @@ def test_call_tir_inplace():
x = I.Var("x", R.Tensor((32, 32), dtype="int32"))
y = I.Var("y", R.Tensor((32, 32), dtype="int32"))
t = I.dynamic("t", dtype="int64")
-I.Call("relax.call_tir_inplace", [Module.tir_func, (x, y, t)], attrs=I.make_node("relax.attrs.CallTIRInplaceAttrs", inplace_indices=[-1, 0]), ty_args=[R.Tuple(R.Tensor((32, 32), dtype="int32"), R.Tensor((32, 32), dtype="int32"))], ty=R.Tuple(R.Tensor((32, 32), dtype="int32"), R.Tensor((32, 32), dtype="int32")))
+R.call_tir_inplace(Module.tir_func, (x, y, t), ty_args=[R.Tuple(R.Tensor((32, 32), dtype="int32"), R.Tensor((32, 32), dtype="int32"))], inplace_indices=[-1, 0])
""",
)
diff --git a/tests/python/relax/test_relax_operators.py b/tests/python/relax/test_relax_operators.py
index 5436ee5633..e105bc31a1 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/tests/python/relax/test_relax_operators.py
@@ -446,6 +446,8 @@ def test_call_tir_canonical_roundtrip(name, tuple_result, result_kind, namespace
callee, call.args[1], ty_args=[output_type], attrs=call.attrs, ty=result_type, loc=loc
)
tvm.ir.assert_structural_equal(call, with_attrs)
+ # Roundtrip printing requires a type-consistent Call after constructor checks.
+ call = tvm.ir.Call(call.op, call.args, attrs=call.attrs, ty_args=call.ty_args, loc=call.loc)
printed = call.script()
assert printed.startswith(f"R.{name}(")
assert "ty_args=" in printed and "out_ty=" not in printed
@@ -552,6 +554,8 @@ def test_call_py_func_canonical_roundtrip(tuple_size, result_kind, builder):
assert len(call.args[1].fields) == 1 and call.args[1].fields[0].same_as(x)
tvm.ir.assert_structural_equal(call.ty_args[0], output_type)
tvm.ir.assert_structural_equal(call.ty, output_type if result_type is None else result_type)
+ # Roundtrip printing requires a type-consistent Call after constructor checks.
+ call = tvm.ir.Call(call.op, call.args, attrs=call.attrs, ty_args=call.ty_args, loc=call.loc)
printed = call.script()
assert printed.startswith("R.call_py_func(")
assert "ty_args=" in printed and "out_ty=" not in printed
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py b/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
index d5352cdc4a..bc5f21db7a 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
@@ -105,20 +105,14 @@ def test_estimate_flops_forloop_as_expression():
def test_estimate_flops_with_decl_tensor():
- def make_func(use_decl_tensor):
- tensor_func = T.decl_tensor if use_decl_tensor else T.Tensor
-
- @Ts.function
- def func(A_data: T.handle("float32")):
- A = tensor_func(16, "float32", data=A_data)
- for i in range(16):
- A[0] = A[0] + 1
-
- return func
+ @Ts.function
+ def func(A_data: T.handle("float32")):
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
+ for i in range(16):
+ A[0] = A[0] + 1
- flops_with_decl_tensor = estimate_tir_flops(IRModule.from_expr(make_func(True)))
- flops_without_decl_tensor = estimate_tir_flops(IRModule.from_expr(make_func(True)))
- assert flops_with_decl_tensor == flops_without_decl_tensor
+ flops = estimate_tir_flops(IRModule.from_expr(func))
+ assert flops == 16
@Ts.function
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py b/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
index 0e9178170e..6fd0cc078e 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
@@ -74,7 +74,7 @@ class TestWriteToTensorArgument(CheckImpureFunction):
class TestWriteToInternalAllocation(CheckPureFunction):
@Ts.function
def func(A: T.Tensor([16, 16], "float32")) -> T.float32:
- Sum = T.decl_tensor([], "float32")
+ Sum = T.alloc_tensor(ty_args=[T.Tensor([], "float32")])
Sum[()] = 0.0
for i, j in T.grid(16, 16):
Sum[()] = Sum[()] + A[i, j]
diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py
index c7a86b1c4e..8cad304c2c 100644
--- a/tests/python/s_tir/analysis/test_sblock_access_region.py
+++ b/tests/python/s_tir/analysis/test_sblock_access_region.py
@@ -141,7 +141,7 @@ def decl_tensor_alias_func(
with Ts.sblock("alias"):
Ts.reads(A[0])
Ts.writes(B[0])
- A_view = T.decl_tensor((16,), "float32", data=A.data)
+ A_view = T.decl_tensor(A.data, ty_args=[T.Tensor((16,), "float32", data=A.data)])
B[0] = A[0] + A_view[0]
diff --git a/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py b/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
index f460bd8703..924d24e979 100644
--- a/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
+++ b/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
@@ -52,7 +52,7 @@ def tensor_opaque_access(
with Ts.sblock():
Ts.reads([])
Ts.writes(B[0:16, 0:16])
- A = T.decl_tensor([256], "float32")
+ A = T.alloc_tensor(ty_args=[T.Tensor([256], "float32")])
for i, j in T.grid(16, 16):
A[i * 16 + j] = 1
for i in range(0, 16):
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 1e38dd10f4..4921d19c28 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
@@ -57,9 +57,9 @@ class Conv2dCuda0:
# body
blockIdx_z = T.launch_thread("blockIdx.z", 196)
- B_local = T.decl_tensor([64], "float32", scope="local")
- Apad_shared = T.decl_tensor([512], "float32", scope="shared")
- Apad_shared_local = T.decl_tensor([8], "float32", scope="local")
+ B_local = T.alloc_tensor(ty_args=[T.Tensor([64], "float32", scope="local")])
+ Apad_shared = T.alloc_tensor(ty_args=[T.Tensor([512], "float32", scope="shared")])
+ Apad_shared_local = T.alloc_tensor(ty_args=[T.Tensor([8], "float32", scope="local")])
blockIdx_y = T.launch_thread("blockIdx.y", 8)
blockIdx_x = T.launch_thread("blockIdx.x", 4)
threadIdx_y = T.launch_thread("threadIdx.y", 8)
@@ -92,9 +92,9 @@ class Conv2dCuda1:
# body
blockIdx_z = T.launch_thread("blockIdx.z", 196)
- B_local = T.decl_tensor([6400000], "float32", scope="local")
- Apad_shared = T.decl_tensor([512], "float32", scope="shared")
- Apad_shared_local = T.decl_tensor([8], "float32", scope="local")
+ B_local = T.alloc_tensor(ty_args=[T.Tensor([6400000], "float32", scope="local")])
+ Apad_shared = T.alloc_tensor(ty_args=[T.Tensor([512], "float32", scope="shared")])
+ Apad_shared_local = T.alloc_tensor(ty_args=[T.Tensor([8], "float32", scope="local")])
blockIdx_y = T.launch_thread("blockIdx.y", 8)
blockIdx_x = T.launch_thread("blockIdx.x", 4)
threadIdx_y = T.launch_thread("threadIdx.y", 8)
@@ -131,9 +131,9 @@ class Conv2dCuda2:
# body
blockIdx_z = T.launch_thread("blockIdx.z", 196)
- B_local = T.decl_tensor([64], "float32", scope="local")
- Apad_shared = T.decl_tensor([512000], "float32", scope="shared")
- Apad_shared_local = T.decl_tensor([8], "float32", scope="local")
+ B_local = T.alloc_tensor(ty_args=[T.Tensor([64], "float32", scope="local")])
+ Apad_shared = T.alloc_tensor(ty_args=[T.Tensor([512000], "float32", scope="shared")])
+ Apad_shared_local = T.alloc_tensor(ty_args=[T.Tensor([8], "float32", scope="local")])
blockIdx_y = T.launch_thread("blockIdx.y", 8)
blockIdx_x = T.launch_thread("blockIdx.x", 4)
threadIdx_y = T.launch_thread("threadIdx.y", 8)
@@ -170,9 +170,9 @@ class Conv2dCuda3:
# body
blockIdx_z = T.launch_thread("blockIdx.z", 196)
- B_local = T.decl_tensor([64], "float32", scope="local")
- Apad_shared = T.decl_tensor([512], "float32", scope="shared")
- Apad_shared_local = T.decl_tensor([8], "float32", scope="local")
+ B_local = T.alloc_tensor(ty_args=[T.Tensor([64], "float32", scope="local")])
+ Apad_shared = T.alloc_tensor(ty_args=[T.Tensor([512], "float32", scope="shared")])
+ Apad_shared_local = T.alloc_tensor(ty_args=[T.Tensor([8], "float32", scope="local")])
blockIdx_y = T.launch_thread("blockIdx.y", 8)
blockIdx_x = T.launch_thread("blockIdx.x", 4)
threadIdx_y = T.launch_thread("threadIdx.y", 8)
diff --git a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
index d3508eb972..34700a8bb9 100644
--- a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
+++ b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
@@ -242,7 +242,7 @@ def test_inferred_ty_with_internal_allocation():
@Ts.function
def func(A: T.Tensor([16, 16], "float32")) -> T.float32:
- Sum = T.decl_tensor([], "float32")
+ Sum = T.alloc_tensor(ty_args=[T.Tensor([], "float32")])
Sum[()] = 0.0
for i, j in T.grid(16, 16):
Sum[()] = Sum[()] + A[i, j]
@@ -432,7 +432,7 @@ def vthread_func():
i0 = T.launch_thread("blockIdx.x", 4)
i1 = T.launch_thread("threadIdx.x", 2)
i2 = T.launch_thread("vthread", 2)
- B = T.alloc_tensor((16,), scope="local")
+ B = T.alloc_tensor(ty_args=[T.Tensor((16,), scope="local")])
for j in range(16):
B[j] = A[i0 * 64 + i1 * 32 + i2 * 16 + j] + T.float32(1)
for j in range(16):
diff --git a/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py b/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
index 2064e39591..60a4021b09 100644
--- a/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
+++ b/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
@@ -218,7 +218,7 @@ def undefined_shape_in_decl_tensor():
@Ts.function(check_well_formed=False)
def func():
- buf = T.decl_tensor(shape=[size], dtype="float32")
+ buf = T.alloc_tensor(ty_args=[T.Tensor([size], "float32")])
T.evaluate(buf[0])
return func
@@ -231,7 +231,9 @@ def undefined_stride_in_decl_tensor():
@Ts.function(check_well_formed=False)
def func():
data_ptr = T.handle("float32")
- buf = T.decl_tensor(shape=[1], dtype="float32", data=data_ptr, strides=[stride])
+ buf = T.decl_tensor(
+ data_ptr, ty_args=[T.Tensor([1], "float32", strides=[stride], data=data_ptr)]
+ )
T.evaluate(buf[0])
return func
@@ -244,7 +246,9 @@ def undefined_elem_offset_in_decl_tensor():
@Ts.function(check_well_formed=False)
def func():
data_ptr = T.handle("float32")
- buf = T.decl_tensor(shape=[1], dtype="float32", data=data_ptr, elem_offset=elem_offset)
+ buf = T.decl_tensor(
+ data_ptr, ty_args=[T.Tensor([1], "float32", elem_offset=elem_offset, data=data_ptr)]
+ )
T.evaluate(buf[0])
return func
diff --git a/tests/python/s_tir/script/test_s_tir_script_error_handling.py b/tests/python/s_tir/script/test_s_tir_script_error_handling.py
index afc77a7ddb..133106633c 100644
--- a/tests/python/s_tir/script/test_s_tir_script_error_handling.py
+++ b/tests/python/s_tir/script/test_s_tir_script_error_handling.py
@@ -373,7 +373,7 @@ def test_match_tensor_shape_mismatch():
def test_high_dim_store():
def high_dim_store() -> None:
with Ts.sblock("root"):
- B = T.alloc_tensor((256,), "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor((256,), "float32")])
for i, j in T.grid(16, 16):
B[i, j] = 1.0 # error: Store is only allowed with one index
diff --git a/tests/python/s_tir/script/test_s_tir_script_ir_builder.py b/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
index 58bd5a75ce..a5d444c401 100644
--- a/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
+++ b/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
@@ -265,7 +265,7 @@ def test_ir_builder_tir_allocate():
with IRBuilder() as ib:
with build_function():
T.func_name_("test")
- buf = T.alloc_tensor([10], "float32", scope="local")
+ buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([10], "float32", scope="local")]))
T.evaluate(1)
# the allocate generated by IRBuilder
@@ -284,7 +284,7 @@ def test_ir_builder_tir_decl_tensor():
with IRBuilder() as ib:
with build_function():
T.func_name_("test")
- buf = T.decl_tensor([128, 128], "float32")
+ buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([128, 128], "float32")]))
T.evaluate(1)
# the decl_tensor generated by IRBuilder
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 d96a08250c..2fb3790285 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
@@ -92,13 +92,9 @@ def test_function_symbolic_alloc_tensor_roundtrip():
buf,
tvm.ir.Call(
"tirx.alloc_tensor",
- [
- tvm.ir.Tuple(buf.shape),
- tvm.ir.DataTypeImm(tvm.DataType(buf.dtype)),
- tvm.ir.StringImm(buf.scope()),
- ],
+ [],
attrs=tvm.ir.DictAttrs({}),
- ty=buf.ty,
+ ty_args=[buf.ty],
),
),
tvm.ir.Evaluate(tirx.TensorLoad(buf, [0])),
@@ -107,7 +103,7 @@ def test_function_symbolic_alloc_tensor_roundtrip():
).with_attr({"s_tir": True, "script.namespace": "s_tir"})
source = func.script()
- assert 'T.alloc_tensor((size,), "float32", layout=None)' in source
+ assert 'T.alloc_tensor(ty_args=[T.Tensor((size,), "float32", layout=None)])' in source
tvm.ir.assert_structural_equal(
tvm.script.from_source(
source,
@@ -628,7 +624,7 @@ def test_masked_load_prevents_scalar_allocation_init_fusion():
@Ts.function
def main():
- A = TB.alloc_tensor((1,), "float32x4")
+ A = TB.alloc_tensor(ty_args=[TB.Tensor((1,), "float32x4")])
A[0] = TB.masked_load(A, 0, TB.Broadcast(TB.bool(True), 4), ty="float32x4")
source = main.script()
@@ -731,12 +727,12 @@ def opt_gemm_lower():
# function attr dict
T.func_attr({"tirx.noalias": True})
# body
- packedB = T.alloc_tensor((32768,))
+ packedB = T.alloc_tensor(ty_args=[T.Tensor((32768,))])
for x in T.parallel(0, 32):
for y in T.serial(0, 1024):
packedB[T.ramp(((x * 32768) + (y * 32)), 1, 32)] = B_1[y, T.ramp(x * 32, 1, 32)]
for x_outer in T.parallel(0, 32):
- C_global = T.alloc_tensor((1024,))
+ C_global = T.alloc_tensor(ty_args=[T.Tensor((1024,))])
for y_outer in T.serial(0, 32):
for x_c_init in T.serial(0, 32):
C_global[T.ramp((x_c_init * 32), 1, 32)] = T.broadcast(T.float32(0), 32)
@@ -788,15 +784,21 @@ def opt_conv_tensorcore_lower():
# function attr dict
T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
# body
- A_1 = T.decl_tensor([12845056], dtype="float16", data=A.data)
- W_1 = T.decl_tensor([1179648], dtype="float16", data=W.data)
- Conv_1 = T.decl_tensor([25690112], data=Conv.data)
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor([12845056], "float16", data=A.data)])
+ W_1 = T.decl_tensor(W.data, ty_args=[T.Tensor([1179648], "float16", data=W.data)])
+ Conv_1 = T.decl_tensor(Conv.data, ty_args=[T.Tensor([25690112], data=Conv.data)])
bz = T.launch_thread("blockIdx.z", 196)
- Conv_wmma_accumulator = T.alloc_tensor((2048,), scope="wmma.accumulator")
- Apad_shared = T.alloc_tensor((12288,), "float16", scope="shared")
- W_shared = T.alloc_tensor((12288,), "float16", scope="shared")
- Apad_shared_wmma_matrix_a = T.alloc_tensor((512,), "float16", scope="wmma.matrix_a")
- W_shared_wmma_matrix_b = T.alloc_tensor((1024,), "float16", scope="wmma.matrix_b")
+ Conv_wmma_accumulator = T.alloc_tensor(
+ ty_args=[T.Tensor((2048,), scope="wmma.accumulator")]
+ )
+ Apad_shared = T.alloc_tensor(ty_args=[T.Tensor((12288,), "float16", scope="shared")])
+ W_shared = T.alloc_tensor(ty_args=[T.Tensor((12288,), "float16", scope="shared")])
+ Apad_shared_wmma_matrix_a = T.alloc_tensor(
+ ty_args=[T.Tensor((512,), "float16", scope="wmma.matrix_a")]
+ )
+ W_shared_wmma_matrix_b = T.alloc_tensor(
+ ty_args=[T.Tensor((1024,), "float16", scope="wmma.matrix_b")]
+ )
bx = T.launch_thread("blockIdx.x", 2)
by = T.launch_thread("blockIdx.y", 4)
ty = T.launch_thread("threadIdx.y", 4)
@@ -1083,7 +1085,9 @@ def opt_conv_tensorcore_mod_host():
)
# body
stack_tcode_data: T.let[T.handle("int32")] = T.stack_alloca("arg_tcode", 10)
- stack_tcode = T.decl_tensor([9], "int32", data=stack_tcode_data)
+ stack_tcode = T.decl_tensor(
+ stack_tcode_data, ty_args=[T.Tensor([9], "int32", data=stack_tcode_data)]
+ )
stack_value: T.let[T.handle] = T.stack_alloca("arg_value", 10)
assert num_args == 3, "default_function: num_args should be 3"
arg0: T.let[T.handle] = T.abi_field_get(args, 0, 12, ty="handle")
@@ -1092,25 +1096,29 @@ def opt_conv_tensorcore_mod_host():
arg2: T.let[T.handle] = T.abi_field_get(args, 2, 12, ty="handle")
A: T.let[T.handle] = T.abi_field_get(arg0, 0, 1, ty="handle")
- A_tensor = T.decl_tensor([1], "float16", data=A)
+ A_tensor = T.decl_tensor(A, ty_args=[T.Tensor([1], "float16", data=A)])
T.assume_aligned(A_tensor, 128)
arg0_shape_data: T.let[T.handle("int64")] = T.abi_field_get(
arg0, 0, 2, ty=T.handle("int64").ty
)
- arg0_shape = T.decl_tensor([6], "int64", data=arg0_shape_data)
+ arg0_shape = T.decl_tensor(
+ arg0_shape_data, ty_args=[T.Tensor([6], "int64", data=arg0_shape_data)]
+ )
arg0_strides_data: T.let[T.handle("int64")] = T.abi_field_get(
arg0, 0, 3, ty=T.handle("int64").ty
)
- arg0_strides = T.decl_tensor([6], "int64", data=arg0_strides_data)
+ arg0_strides = T.decl_tensor(
+ arg0_strides_data, ty_args=[T.Tensor([6], "int64", data=arg0_strides_data)]
+ )
dev_id: T.let[T.int32] = T.abi_field_get(arg0, 0, 9, ty="int32")
W: T.let[T.handle] = T.abi_field_get(arg1, 0, 1, ty="handle")
- W_tensor = T.decl_tensor([1], "float16", data=W)
+ W_tensor = T.decl_tensor(W, ty_args=[T.Tensor([1], "float16", data=W)])
T.assume_aligned(W_tensor, 128)
Conv: T.let[T.handle] = T.abi_field_get(arg2, 0, 1, ty="handle")
- Conv_tensor = T.decl_tensor([1], "float32", data=Conv)
+ Conv_tensor = T.decl_tensor(Conv, ty_args=[T.Tensor([1], "float32", data=Conv)])
T.assume_aligned(Conv_tensor, 128)
assert (((arg0_code == 3) or (arg0_code == 13)) or (arg0_code == 7)) or (arg0_code == 4), (
@@ -1286,7 +1294,9 @@ def function_with_allocate_annotations():
)
# body
- tensor_2 = T.alloc_tensor((200704,), "uint8", annotations={"attr1_key": "attr1_value"})
+ tensor_2 = T.alloc_tensor(
+ ty_args=[T.Tensor((200704,), "uint8")], annotations={"attr1_key": "attr1_value"}
+ )
for ax0_ax1_fused_4 in T.serial(0, 56):
for ax2_4 in T.serial(0, 56):
for ax3_init in T.serial(0, 64):
@@ -1335,7 +1345,7 @@ def comm_reducer_single_reduce_group():
for i in T.serial(0, 128):
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- reduce_temp0 = T.alloc_tensor((1,), scope="local")
+ reduce_temp0 = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
T.evaluate(
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x, y: (x + y,)),
@@ -1359,8 +1369,8 @@ def comm_reducer_multiple_reduce_groups():
for i in T.serial(0, 128):
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- reduce_temp0 = T.alloc_tensor((1,), "int32", scope="local")
- reduce_temp1 = T.alloc_tensor((1,), "float32", scope="local")
+ reduce_temp0 = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
+ reduce_temp1 = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
T.evaluate(
T.gpu_thread_allreduce(
T.Lambda(
@@ -1506,7 +1516,9 @@ def func_T_ptr_let_statement():
) -> None:
# The T.Ptr declaration in the parameter list should parse
# correctly, and should be usable as the data pointer in a tensor.
- arg_type_ids = T.decl_tensor([2], dtype="int32", data=arg_type_ids_handle)
+ arg_type_ids = T.decl_tensor(
+ arg_type_ids_handle, ty_args=[T.Tensor([2], "int32", data=arg_type_ids_handle)]
+ )
arg0: T.let[T.handle] = T.abi_field_get(args, 0, 12, ty="handle")
arg1: T.let[T.handle] = T.abi_field_get(args, 1, 12, ty="handle")
@@ -1521,11 +1533,11 @@ def func_T_ptr_let_statement():
# this function. It should only be defined after the data pointer
# has been defined, and should not be hoisted into the header of
# the function as other tensor_decl statements can be.
- A = T.decl_tensor([1024], dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor([1024], "float32", data=A_data)])
B_data: T.let[T.handle("float32")] = T.reinterpret(
T.abi_field_get(arg1, 0, 1, ty="handle"), ty=T.handle("float32").ty
)
- B = T.decl_tensor([1024], dtype="float32", data=B_data)
+ B = T.decl_tensor(B_data, ty_args=[T.Tensor([1024], "float32", data=B_data)])
B[0] = A[0]
@@ -1535,7 +1547,7 @@ def func_T_ptr_let_statement():
def func_T_ptr_allocate():
@Ts.function
def func_T_ptr_allocate() -> None:
- A = T.alloc_tensor((1024,))
+ A = T.alloc_tensor(ty_args=[T.Tensor((1024,))])
A[0] = 0.0
return func_T_ptr_allocate
@@ -1582,8 +1594,8 @@ def string_annotation_escaping():
def pointer_type():
@Ts.function
def func_with_ptr_type_annotations(x: T.handle("int32"), y: T.handle("int32", "shared")):
- xx = T.alloc_tensor((16,), "int32")
- yy = T.alloc_tensor((16,), "int32", scope="shared")
+ xx = T.alloc_tensor(ty_args=[T.Tensor((16,), "int32")])
+ yy = T.alloc_tensor(ty_args=[T.Tensor((16,), "int32", scope="shared")])
a: T.let[T.handle("int32")] = T.address_of(xx[0])
b: T.let[T.handle("int32", "shared")] = T.address_of(yy[0])
T.evaluate(T.call_extern("copy", a, b, ty=""))
@@ -1723,9 +1735,11 @@ def void_ptr():
def decl_tensor():
@Ts.function
def func(A: T.Tensor((16, 16), "float32"), B: T.Tensor((16, 16), "float32")) -> None:
- A_flattened = T.decl_tensor(data=A.data, shape=(256,), dtype="float32")
- B_flattened = T.decl_tensor(data=B.data, shape=(256,), dtype="float32")
- C_alias = T.decl_tensor(data=A_flattened.data, shape=(256,), dtype="float32")
+ A_flattened = T.decl_tensor(A.data, ty_args=[T.Tensor((256,), "float32", data=A.data)])
+ B_flattened = T.decl_tensor(B.data, ty_args=[T.Tensor((256,), "float32", data=B.data)])
+ C_alias = T.decl_tensor(
+ A_flattened.data, ty_args=[T.Tensor((256,), "float32", data=A_flattened.data)]
+ )
for i in range(256):
B_flattened[i] = A_flattened[i] + C_alias[i] + T.float32(1.0)
@@ -1735,9 +1749,9 @@ def decl_tensor():
def allocate_and_decl_tensor():
@Ts.function
def func(A: T.Tensor((16,), "float32"), B: T.Tensor((16,), "float32")) -> None:
- D = T.alloc_tensor((16,))
+ D = T.alloc_tensor(ty_args=[T.Tensor((16,))])
for i in range(4):
- C = T.alloc_tensor((4,))
+ C = T.alloc_tensor(ty_args=[T.Tensor((4,))])
for j in range(4):
C[j] = A[i * 4 + j] + T.float32(1.0)
for j in range(4):
@@ -1751,7 +1765,7 @@ def allocate_and_decl_tensor():
def alloc_tensor_example():
@Ts.function
def func(A: T.Tensor((128,), "float32"), C: T.Tensor((128,), "float32")):
- B = T.alloc_tensor((128,), "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32")])
for i in range(128):
B[i] = A[i] * T.float32(2)
for i in range(128):
@@ -1877,24 +1891,24 @@ def tvm_shfl_builtins():
):
blockIdx_x = T.launch_thread("blockIdx.x", 1)
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- A_warp = T.alloc_tensor((1,), scope="local")
- B_warp = T.alloc_tensor((1,), scope="local")
- red_buf0 = T.alloc_tensor((1,), scope="local")
- A_warp_1 = T.decl_tensor((32,), data=A_warp.data, scope="local")
- A_1 = T.decl_tensor((32,), data=A) # A is a handle param
+ A_warp = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ B_warp = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ red_buf0 = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ A_warp_1 = T.decl_tensor(A_warp.data, ty_args=[T.Tensor((32,), scope="local")])
+ A_1 = T.decl_tensor(A, ty_args=[T.Tensor((32,), data=A)]) # A is a handle param
A_warp_1[0] = A_1[threadIdx_x]
- B_warp_1 = T.decl_tensor((32,), data=B_warp.data, scope="local")
+ B_warp_1 = T.decl_tensor(B_warp.data, ty_args=[T.Tensor((32,), scope="local")])
T.gpu_storage_sync("warp")
B_warp_1[0] = T.gpu_warp_shuffle(
T.gpu_warp_activemask(), A_warp_1[0], threadIdx_x % 4 * 8 + threadIdx_x // 4, 32, 32
) + T.float32(1)
- red_buf0_1 = T.decl_tensor((1,), data=red_buf0.data, scope="local")
- mask = T.alloc_tensor((1,), "uint32", scope="local")
- t0 = T.alloc_tensor((1,), scope="local")
+ red_buf0_1 = T.decl_tensor(red_buf0.data, ty_args=[T.Tensor((1,), scope="local")])
+ mask = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="local")])
+ t0 = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
red_buf0_1[0] = A_warp_1[0]
- mask_1 = T.decl_tensor((1,), "uint32", data=mask.data, scope="local")
+ mask_1 = T.decl_tensor(mask.data, ty_args=[T.Tensor((1,), "uint32", scope="local")])
mask_1[0] = T.gpu_warp_activemask()
- t0_1 = T.decl_tensor((1,), data=t0.data, scope="local")
+ t0_1 = T.decl_tensor(t0.data, ty_args=[T.Tensor((1,), scope="local")])
t0_1[0] = T.gpu_warp_shuffle_down(mask_1[0], red_buf0_1[0], 16, 32, 32)
red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
t0_1[0] = T.gpu_warp_shuffle_down(mask_1[0], red_buf0_1[0], 8, 32, 32)
@@ -1909,9 +1923,9 @@ def tvm_shfl_builtins():
# NOTE(Zihao): test gpu_warp_shuffle_up
red_buf0_1[0] = T.gpu_warp_shuffle_up(mask_1[0], red_buf0_1[0], 0, 32, 32)
if threadIdx_x == 0:
- C_1 = T.decl_tensor((1,), data=C)
+ C_1 = T.decl_tensor(C, ty_args=[T.Tensor((1,), data=C)])
C_1[0] = red_buf0_1[0]
- B_1 = T.decl_tensor((32,), data=B)
+ B_1 = T.decl_tensor(B, ty_args=[T.Tensor((32,), data=B)])
B_1[threadIdx_x] = B_warp_1[0]
return func
@@ -1970,7 +1984,7 @@ def undefined_data_ptr_in_decl_tensor():
@Ts.function(check_well_formed=False)
def func():
data_ptr = T.handle("float32")
- buf = T.decl_tensor(shape=[1], dtype="float32", data=data_ptr)
+ buf = T.decl_tensor(data_ptr, ty_args=[T.Tensor([1], "float32", data=data_ptr)])
T.evaluate(buf[0])
return func
@@ -2107,8 +2121,8 @@ def test_disable_concise_scoping_when_scope_annotated():
y: T.int32 = x + 1
T.evaluate(y - 1)
- # Explicit scalar declarations lower to AllocTensor + TensorStore (local_scalar).
- # The printer fuses each pair into one line; annotate the allocation for y.
+ # Explicit scalar declarations lower to allocation and store statements.
+ # Keep the annotation attached to the allocation for y.
result = _func.with_attr("global_symbol", "main").script(
obj_to_annotate={
_func.body.seq[2]: "annotation 1",
@@ -2122,9 +2136,11 @@ def test_disable_concise_scoping_when_scope_annotated():
@Ts.function
def main():
T.func_attr({"s_tir": True, "script.namespace": "s_tir"})
- x: T.int32 = 1
- y: T.int32 = x + 1 # annotation 1
- T.evaluate(y - 1)"""
+ x = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local", layout="default")])
+ x[0] = 1
+ y = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local", layout="default")]) # annotation 1
+ y[0] = x[0] + 1
+ T.evaluate(y[0] - 1)"""
)
diff --git a/tests/python/s_tir/test_s_tir_renew_defs.py b/tests/python/s_tir/test_s_tir_renew_defs.py
index 459365ad29..924003f4f1 100644
--- a/tests/python/s_tir/test_s_tir_renew_defs.py
+++ b/tests/python/s_tir/test_s_tir_renew_defs.py
@@ -139,7 +139,7 @@ def test_undefined_tensor():
@Ts.function
def access_alloc():
# Var A should be remapped
- A = T.alloc_tensor((128,), "float16")
+ A = T.alloc_tensor(ty_args=[T.Tensor((128,), "float16")])
T.evaluate(A.data)
for i in range(128):
A[i] = A[i] + T.float16(1.0)
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py b/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
index d6e3ec4911..599a574661 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
@@ -661,9 +661,9 @@ class TestDataDependentRegion(BaseCompactTest):
p1: T.Tensor((1,), "int32"),
hybrid_nms: T.Tensor((30,), "float32"),
):
- argsort_nms_cpu = T.decl_tensor([5], "int32", scope="global")
+ argsort_nms_cpu = T.alloc_tensor(ty_args=[T.Tensor([5], "int32", scope="global")])
for i in range(1):
- nkeep = T.decl_tensor([1], "int32", scope="global")
+ nkeep = T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="global")])
if 0 < p1[i]:
nkeep[0] = p1[i]
if 2 < nkeep[0]:
@@ -1005,9 +1005,9 @@ class TestTileAwareCompaction(BaseCompactTest):
):
for i_0 in range(5, annotations={"loop_partition_hint": 1}):
for j_0 in range(5, annotations={"loop_partition_hint": 1}):
- A_local = T.decl_tensor((26, 128), scope="local")
- B_local = T.decl_tensor((128, 26), scope="local")
- C_local = T.decl_tensor((26, 26), scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor((26, 128), scope="local")])
+ B_local = T.alloc_tensor(ty_args=[T.Tensor((128, 26), scope="local")])
+ C_local = T.alloc_tensor(ty_args=[T.Tensor((26, 26), scope="local")])
for ax0, ax1 in T.grid(26, 128):
if i_0 * 26 + ax0 < 128:
A_local[ax0, ax1] = A[i_0 * 26 + ax0, ax1]
@@ -1043,9 +1043,9 @@ class TestTileAwareCompaction(BaseCompactTest):
):
for i_0 in range(4):
for j_0 in range(4):
- A_local_tile0 = T.decl_tensor((26, 128), scope="local")
- B_local_tile0 = T.decl_tensor((128, 26), scope="local")
- C_local_tile0 = T.decl_tensor((26, 26), scope="local")
+ A_local_tile0 = T.alloc_tensor(ty_args=[T.Tensor((26, 128), scope="local")])
+ B_local_tile0 = T.alloc_tensor(ty_args=[T.Tensor((128, 26), scope="local")])
+ C_local_tile0 = T.alloc_tensor(ty_args=[T.Tensor((26, 26), scope="local")])
for ax0, ax1 in T.grid(26, 128):
A_local_tile0[ax0, ax1] = A[i_0 * 26 + ax0, ax1]
for ax0, ax1 in T.grid(128, 26):
@@ -1059,9 +1059,9 @@ class TestTileAwareCompaction(BaseCompactTest):
for ax0, ax1 in T.grid(26, 26):
C[i_0 * 26 + ax0, j_0 * 26 + ax1] = C_local_tile0[ax0, ax1]
- A_local_tile1 = T.decl_tensor((26, 128), scope="local")
- B_local_tile1 = T.decl_tensor((128, 24), scope="local")
- C_local_tile1 = T.decl_tensor((26, 24), scope="local")
+ A_local_tile1 = T.alloc_tensor(ty_args=[T.Tensor((26, 128), scope="local")])
+ B_local_tile1 = T.alloc_tensor(ty_args=[T.Tensor((128, 24), scope="local")])
+ C_local_tile1 = T.alloc_tensor(ty_args=[T.Tensor((26, 24), scope="local")])
for ax0, ax1 in T.grid(26, 128):
A_local_tile1[ax0, ax1] = A[i_0 * 26 + ax0, ax1]
for ax0, ax1 in T.grid(128, 26):
@@ -1079,9 +1079,9 @@ class TestTileAwareCompaction(BaseCompactTest):
C[i_0 * 26 + ax0, ax1 + 104] = C_local_tile1[ax0, ax1]
for j_0 in range(4):
- A_local_tile2 = T.decl_tensor((24, 128), scope="local")
- B_local_tile2 = T.decl_tensor((128, 26), scope="local")
- C_local_tile2 = T.decl_tensor((24, 26), scope="local")
+ A_local_tile2 = T.alloc_tensor(ty_args=[T.Tensor((24, 128), scope="local")])
+ B_local_tile2 = T.alloc_tensor(ty_args=[T.Tensor((128, 26), scope="local")])
+ C_local_tile2 = T.alloc_tensor(ty_args=[T.Tensor((24, 26), scope="local")])
for ax0, ax1 in T.grid(26, 128):
if ax0 < 24:
A_local_tile2[ax0, ax1] = A[ax0 + 104, ax1]
@@ -1098,9 +1098,9 @@ class TestTileAwareCompaction(BaseCompactTest):
if ax0 < 24:
C[ax0 + 104, j_0 * 26 + ax1] = C_local_tile2[ax0, ax1]
- A_local_tile3 = T.decl_tensor((24, 128), scope="local")
- B_local_tile3 = T.decl_tensor((128, 24), scope="local")
- C_local_tile3 = T.decl_tensor((24, 24), scope="local")
+ A_local_tile3 = T.alloc_tensor(ty_args=[T.Tensor((24, 128), scope="local")])
+ B_local_tile3 = T.alloc_tensor(ty_args=[T.Tensor((128, 24), scope="local")])
+ C_local_tile3 = T.alloc_tensor(ty_args=[T.Tensor((24, 24), scope="local")])
for ax0, ax1 in T.grid(26, 128):
if ax0 < 24:
A_local_tile3[ax0, ax1] = A[ax0 + 104, ax1]
@@ -1203,9 +1203,9 @@ class TestNotCompactAliasTensor(BaseCompactTest):
def before():
"""Partially accessed tensor, but should not compact
because existence of aliasing tensor B."""
- data = T.alloc_tensor((1024,), "int8")
- A = T.decl_tensor([1024], "int8", data=data.data)
- B = T.decl_tensor([512], "float16", data=data.data)
+ data = T.alloc_tensor(ty_args=[T.Tensor((1024,), "int8")])
+ A = T.decl_tensor(data.data, ty_args=[T.Tensor([1024], "int8", data=data.data)])
+ B = T.decl_tensor(data.data, ty_args=[T.Tensor([512], "float16", data=data.data)])
for i in range(10):
A[i] = A[i] + T.int8(1)
for i in range(10):
@@ -1222,8 +1222,8 @@ class TestNotCompactTensorWithDifferentDtype(BaseCompactTest):
def before():
"""Partially accessed tensor, but should not compact
because existence of aliasing tensor B."""
- data = T.alloc_tensor((1024,), "int8")
- A = T.decl_tensor([256], "int32", data=data.data)
+ data = T.alloc_tensor(ty_args=[T.Tensor((1024,), "int8")])
+ A = T.decl_tensor(data.data, ty_args=[T.Tensor([256], "int32", data=data.data)])
for i in range(10):
A[i] = A[i] + 1
@@ -1236,14 +1236,14 @@ class TestNonBoolCondition(BaseCompactTest):
@Ts.function
def before():
- A = T.decl_tensor([12], "int32")
+ A = T.alloc_tensor(ty_args=[T.Tensor([12], "int32")])
for i in range(10):
if i:
A[i] = A[i] + 1
@Ts.function
def expected():
- A = T.decl_tensor((9,), "int32")
+ A = T.alloc_tensor(ty_args=[T.Tensor((9,), "int32")])
for i in range(10):
if i:
A[i - 1] = A[i - 1] + 1
@@ -1254,7 +1254,7 @@ def test_loop_var_does_not_escape_compacted_tensor_extent():
@Ts.function(private=True)
def before(A: T.Tensor((n,), "int32")):
- tmp = T.alloc_tensor((n,), "int32")
+ tmp = T.alloc_tensor(ty_args=[T.Tensor((n,), "int32")])
for i in range(n):
length: T.let[T.int64] = T.ceildiv(n, T.shift_left(T.int64(1), i + 1))
for j in range(length):
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 ce6522338c..b091df9b12 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
@@ -879,8 +879,8 @@ def test_multiplication_nodes_are_inlined():
@Ts.function
def main(A: T.Tensor((32, 128), "float16")):
tx = T.launch_thread("threadIdx.x", T.int64(32))
- A_flattened = T.decl_tensor((4096,), "float16", data=A.data)
- A_shared = T.decl_tensor([4096], "float16", scope="shared")
+ A_flattened = T.decl_tensor(A.data, ty_args=[T.Tensor((4096,), "float16", data=A.data)])
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor([4096], "float16", scope="shared")])
with Ts.async_copy_scope():
for i in range(16):
@@ -896,8 +896,8 @@ def test_multiplication_nodes_are_inlined():
@Ts.function
def main(A: T.Tensor((32, 128), "float16")):
tx = T.launch_thread("threadIdx.x", T.int64(32))
- A_flattened = T.decl_tensor((4096,), "float16", data=A.data)
- A_shared = T.decl_tensor((4096,), "float16", scope="shared")
+ A_flattened = T.decl_tensor(A.data, ty_args=[T.Tensor((4096,), "float16", data=A.data)])
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor((4096,), "float16", scope="shared")])
for i in range(16):
cse_v1: T.int64 = T.Cast("int64", i)
T.s_tir.cp_async_raw(
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
index 65a7baa217..c89253ff92 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
@@ -43,12 +43,12 @@ def test_vthread():
class Module:
@Ts.function
def main(A: T.handle("float32"), C: T.handle("float32")):
- A_buf = T.decl_tensor((n * nthread,), "float32", data=A)
- C_buf = T.decl_tensor((n * nthread,), "float32", data=C)
+ A_buf = T.decl_tensor(A, ty_args=[T.Tensor((n * nthread,), "float32", data=A)])
+ C_buf = T.decl_tensor(C, ty_args=[T.Tensor((n * nthread,), "float32", data=C)])
for i in range(n):
vt_x = T.launch_thread("vthread", nthread)
vt_y = T.launch_thread("vthread", nthread)
- B = T.alloc_tensor((m,), scope="shared")
+ B = T.alloc_tensor(ty_args=[T.Tensor((m,), scope="shared")])
B[i] = A_buf[i * nthread + vt_x]
T.evaluate(T.call_extern("Run", B.ptr_to([0] * len(B.shape)), vt_x, ty="int32"))
C_buf[i * nthread + vt_x] = B[i] + T.float32(1)
@@ -84,9 +84,9 @@ def test_vthread_extern():
for i in range(n):
vt_x = T.launch_thread("vthread", nthread)
vt_y = T.launch_thread("vthread", nthread)
- A = T.alloc_tensor((m,), scope="shared")
- B = T.alloc_tensor((m,), scope="shared")
- C = T.alloc_tensor((m,), scope="shared")
+ A = T.alloc_tensor(ty_args=[T.Tensor((m,), scope="shared")])
+ B = T.alloc_tensor(ty_args=[T.Tensor((m,), scope="shared")])
+ C = T.alloc_tensor(ty_args=[T.Tensor((m,), scope="shared")])
A[vt_x] = T.Cast("float32", vt_x) + T.float32(1)
B[vt_y] = T.Cast("float32", vt_y) + T.float32(1)
T.evaluate(
@@ -128,10 +128,10 @@ def test_vthread_if_then_else():
@Ts.function
def main(A: T.handle("float32")):
T.func_attr({"global_symbol": "main"})
- A_buf = T.decl_tensor((100 * nthread,), "float32", data=A)
+ A_buf = T.decl_tensor(A, ty_args=[T.Tensor((100 * nthread,), "float32", data=A)])
for i in range(100):
vt = T.launch_thread("vthread", nthread)
- B = T.alloc_tensor((128,), scope="shared")
+ B = T.alloc_tensor(ty_args=[T.Tensor((128,), scope="shared")])
if i == 0:
B[i] = A_buf[i * nthread + vt]
else:
@@ -166,12 +166,12 @@ def test_vthread_simplified():
@Ts.function
def before_func():
vthread = T.launch_thread("vthread", 4)
- B = T.alloc_tensor((4,), "int32", scope="shared")
+ B = T.alloc_tensor(ty_args=[T.Tensor((4,), "int32", scope="shared")])
B[T.ramp(0, 1, 4)] = T.broadcast(vthread, 4)
@Ts.function
def expected_func():
- B = T.alloc_tensor((16,), "int32", scope="shared")
+ B = T.alloc_tensor(ty_args=[T.Tensor((16,), "int32", scope="shared")])
# The indices for B should each be a single Ramp node, and
# should not be the sum of a Ramp and Broadcast node.
B[T.ramp(T.Mul(0, 4), 1, 4)] = T.broadcast(0, 4)
@@ -192,7 +192,7 @@ def test_vthread_vectorized():
@Ts.function
def before_func():
vthread = T.launch_thread("vthread", 4)
- B = T.alloc_tensor((4,), "int32", scope="shared")
+ B = T.alloc_tensor(ty_args=[T.Tensor((4,), "int32", scope="shared")])
B[T.ramp(0, 1, 4)] = T.broadcast(vthread, 4)
before_mod = tvm.IRModule.from_expr(before_func.with_attr("global_symbol", "main"))
@@ -218,7 +218,7 @@ def test_vthread_rewrites_masked_accesses():
@Ts.function
def before_func():
vthread = T.launch_thread("vthread", 2)
- B = T.alloc_tensor((4,), "float32", scope="shared")
+ B = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32", scope="shared")])
mask = T.meta_var(T.Broadcast(T.bool(True), 4))
loaded = T.meta_var(T.masked_load(B, T.Ramp(0, 1, 4), mask, ty="float32x4"))
value = T.meta_var(loaded + T.Broadcast(T.Cast("float32", vthread), 4))
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
index cdc8694809..3e3895c844 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
@@ -119,8 +119,8 @@ def test_everything_during_deduction():
def test_oneD_pool():
@Ts.function
def func(m: T.int64, data: T.handle("float32"), out: T.handle("float32")):
- data_ptr = T.decl_tensor((16,), "float32", data=data)
- out_ptr = T.decl_tensor((16,), "float32", data=out)
+ data_ptr = T.decl_tensor(data, ty_args=[T.Tensor((16,), "float32", data=data)])
+ out_ptr = T.decl_tensor(out, ty_args=[T.Tensor((16,), "float32", data=out)])
for ow in range(16):
for kw in range(3):
if T.likely(ow > 0):
@@ -240,10 +240,18 @@ def partitioned_concat_3(
placeholder_2: T.Tensor((1, 32, 28, 28), "int8"),
T_concat: T.Tensor((1, 128, 28, 28), "int8"),
) -> None:
- placeholder_flat = T.decl_tensor([50176], "int8", data=placeholder.data)
- placeholder_1_flat = T.decl_tensor([25088], "int8", data=placeholder_1.data)
- placeholder_2_flat = T.decl_tensor([25088], "int8", data=placeholder_2.data)
- T_concat_flat = T.decl_tensor([100352], "int8", data=T_concat.data)
+ placeholder_flat = T.decl_tensor(
+ placeholder.data, ty_args=[T.Tensor([50176], "int8", data=placeholder.data)]
+ )
+ placeholder_1_flat = T.decl_tensor(
+ placeholder_1.data, ty_args=[T.Tensor([25088], "int8", data=placeholder_1.data)]
+ )
+ placeholder_2_flat = T.decl_tensor(
+ placeholder_2.data, ty_args=[T.Tensor([25088], "int8", data=placeholder_2.data)]
+ )
+ T_concat_flat = T.decl_tensor(
+ T_concat.data, ty_args=[T.Tensor([100352], "int8", data=T_concat.data)]
+ )
for i1, i2, i3 in T.grid(64, 28, 28):
T_concat_flat[i1 * 784 + i2 * 28 + i3] = placeholder_flat[i1 * 784 + i2 * 28 + i3]
for i1, i2, i3 in T.grid(32, 28, 28):
@@ -259,10 +267,18 @@ def concat_func_3(
placeholder_2: T.Tensor((1, 32, 28, 28), "int8"),
T_concat: T.Tensor((1, 128, 28, 28), "int8"),
) -> None:
- placeholder_flat = T.decl_tensor([50176], "int8", data=placeholder.data)
- placeholder_1_flat = T.decl_tensor([25088], "int8", data=placeholder_1.data)
- placeholder_2_flat = T.decl_tensor([25088], "int8", data=placeholder_2.data)
- T_concat_flat = T.decl_tensor([100352], "int8", data=T_concat.data)
+ placeholder_flat = T.decl_tensor(
+ placeholder.data, ty_args=[T.Tensor([50176], "int8", data=placeholder.data)]
+ )
+ placeholder_1_flat = T.decl_tensor(
+ placeholder_1.data, ty_args=[T.Tensor([25088], "int8", data=placeholder_1.data)]
+ )
+ placeholder_2_flat = T.decl_tensor(
+ placeholder_2.data, ty_args=[T.Tensor([25088], "int8", data=placeholder_2.data)]
+ )
+ T_concat_flat = T.decl_tensor(
+ T_concat.data, ty_args=[T.Tensor([100352], "int8", data=T_concat.data)]
+ )
for i1 in T.serial(128, annotations={"loop_partition_hint": 1}):
for i2, i3 in T.grid(28, 28):
if 96 <= i1:
@@ -291,8 +307,8 @@ def test_loop_partition_unroll_hint():
def main(
A_arg: T.Tensor((1, 3, 224, 224), "int8"), B_arg: T.Tensor((1, 224, 7, 16), "int8")
) -> None:
- A = T.decl_tensor(150528, "int8", data=A_arg.data)
- B = T.decl_tensor(25088, "int8", data=B_arg.data)
+ A = T.decl_tensor(A_arg.data, ty_args=[T.Tensor(150528, "int8", data=A_arg.data)])
+ B = T.decl_tensor(B_arg.data, ty_args=[T.Tensor(25088, "int8", data=B_arg.data)])
for ax0 in T.serial(
112,
annotations={"loop_partition_hint": True},
@@ -305,8 +321,8 @@ def test_loop_partition_unroll_hint():
def partitioned_main(
A_arg: T.Tensor((1, 3, 224, 224), "int8"), B_arg: T.Tensor((1, 224, 7, 16), "int8")
) -> None:
- A = T.decl_tensor(150528, dtype="int8", data=A_arg.data)
- B = T.decl_tensor(25088, dtype="int8", data=B_arg.data)
+ A = T.decl_tensor(A_arg.data, ty_args=[T.Tensor(150528, "int8", data=A_arg.data)])
+ B = T.decl_tensor(B_arg.data, ty_args=[T.Tensor(25088, "int8", data=B_arg.data)])
# body
for ax1, ax2, ax3 in T.grid(224, 7, 16):
if 3 <= ax2 and ax3 < 3:
@@ -339,10 +355,10 @@ def test_loop_partition_unroll_hint():
def test_loop_partition_recursive_unroll_hint():
@Ts.function
def main():
- placeholder_0_dm = T.decl_tensor([1, 32, 32, 16], dtype="int8")
+ placeholder_0_dm = T.alloc_tensor(ty_args=[T.Tensor([1, 32, 32, 16], "int8")])
for i3_0 in T.serial(5, annotations={"loop_partition_hint": 1}):
for i2_0 in T.serial(2, annotations={"loop_partition_hint": 1}):
- pad_temp = T.decl_tensor([1, 16, 16, 16], dtype="int8")
+ pad_temp = T.alloc_tensor(ty_args=[T.Tensor([1, 16, 16, 16], "int8")])
for ax0, ax1, ax2 in T.grid(16, 16, 16):
if (
6 <= i2_0 * 4 + ax0
@@ -364,17 +380,17 @@ def test_loop_partition_recursive_unroll_hint():
@Ts.function
def partitioned_main():
- placeholder_0_dm = T.decl_tensor((16384,), "int8")
+ placeholder_0_dm = T.alloc_tensor(ty_args=[T.Tensor((16384,), "int8")])
for i3_0 in T.unroll(2):
for i2_0 in T.unroll(2):
- pad_temp = T.decl_tensor((4096,), "int8")
+ pad_temp = T.alloc_tensor(ty_args=[T.Tensor((4096,), "int8")])
for ax0, ax1, ax2 in T.grid(16, 16, 16):
if 6 <= i2_0 * 4 + ax0 and 6 <= i3_0 * 4 + ax1:
pad_temp[ax0 * 256 + ax1 * 16 + ax2] = placeholder_0_dm[
i2_0 * 2048 + ax0 * 512 + i3_0 * 64 + ax1 * 16 + ax2
]
for i2_0 in T.unroll(2):
- pad_temp = T.decl_tensor((4096,), "int8")
+ pad_temp = T.alloc_tensor(ty_args=[T.Tensor((4096,), "int8")])
for ax0, ax1, ax2 in T.grid(16, 16, 16):
if 6 <= i2_0 * 4 + ax0:
pad_temp[ax0 * 256 + ax1 * 16 + ax2] = placeholder_0_dm[
@@ -382,7 +398,7 @@ def test_loop_partition_recursive_unroll_hint():
]
for i3_0 in T.unroll(2):
for i2_0 in T.unroll(2):
- pad_temp = T.decl_tensor((4096,), "int8")
+ pad_temp = T.alloc_tensor(ty_args=[T.Tensor((4096,), "int8")])
for ax0, ax1, ax2 in T.grid(16, 16, 16):
if 6 <= i2_0 * 4 + ax0 and i3_0 * 4 + ax1 < 14:
pad_temp[ax0 * 256 + ax1 * 16 + ax2] = placeholder_0_dm[
@@ -417,8 +433,8 @@ def test_loop_partition_keep_loop_annotations():
@Ts.function
def after(A: T.Tensor(160, "int32"), B: T.Tensor(160, "int32")) -> None:
- A_1 = T.decl_tensor((160,), "int32", data=A.data)
- B_1 = T.decl_tensor((160,), "int32", data=B.data)
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((160,), "int32", data=A.data)])
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((160,), "int32", data=B.data)])
for i in T.serial(10, annotations={"key": "value"}):
B_1[i] = A_1[i] + 1
for i in T.serial(140, annotations={"key": "value"}):
@@ -466,10 +482,18 @@ def test_loop_partition_with_unit_loop_in_condition():
placeholder_2: T.Tensor(25088, "int8"),
T_concat: T.Tensor(100352, "int8"),
) -> None:
- placeholder_3 = T.decl_tensor((50176,), "int8", data=placeholder.data)
- placeholder_1_1 = T.decl_tensor((25088,), "int8", data=placeholder_1.data)
- placeholder_2_1 = T.decl_tensor((25088,), "int8", data=placeholder_2.data)
- T_concat_1 = T.decl_tensor((100352,), "int8", data=T_concat.data)
+ placeholder_3 = T.decl_tensor(
+ placeholder.data, ty_args=[T.Tensor((50176,), "int8", data=placeholder.data)]
+ )
+ placeholder_1_1 = T.decl_tensor(
+ placeholder_1.data, ty_args=[T.Tensor((25088,), "int8", data=placeholder_1.data)]
+ )
+ placeholder_2_1 = T.decl_tensor(
+ placeholder_2.data, ty_args=[T.Tensor((25088,), "int8", data=placeholder_2.data)]
+ )
+ T_concat_1 = T.decl_tensor(
+ T_concat.data, ty_args=[T.Tensor((100352,), "int8", data=T_concat.data)]
+ )
for k in T.serial(1, annotations={"preserve_unit_loop": True}):
for i1, i2, i3 in T.grid(64, 28, 28):
T_concat_1[i1 * 784 + i2 * 28 + i3] = placeholder_3[i1 * 784 + i2 * 28 + i3]
@@ -515,10 +539,18 @@ def expected_partitioned_concat_single_point(
placeholder_2: T.Tensor((28, 63), "int8"),
T_concat: T.Tensor((28, 128), "int8"),
):
- placeholder_3 = T.decl_tensor((1792,), "int8", data=placeholder.data)
- placeholder_1_1 = T.decl_tensor((28,), "int8", data=placeholder_1.data)
- placeholder_2_1 = T.decl_tensor((1764,), "int8", data=placeholder_2.data)
- T_concat_1 = T.decl_tensor((3584,), "int8", data=T_concat.data)
+ placeholder_3 = T.decl_tensor(
+ placeholder.data, ty_args=[T.Tensor((1792,), "int8", data=placeholder.data)]
+ )
+ placeholder_1_1 = T.decl_tensor(
+ placeholder_1.data, ty_args=[T.Tensor((28,), "int8", data=placeholder_1.data)]
+ )
+ placeholder_2_1 = T.decl_tensor(
+ placeholder_2.data, ty_args=[T.Tensor((1764,), "int8", data=placeholder_2.data)]
+ )
+ T_concat_1 = T.decl_tensor(
+ T_concat.data, ty_args=[T.Tensor((3584,), "int8", data=T_concat.data)]
+ )
for i0 in range(28):
for i1 in range(63):
T_concat_1[i0 * 128 + i1] = placeholder_2_1[i0 * 63 + i1]
@@ -554,10 +586,18 @@ def concat_func_start_point_equality_expected(
placeholder_2: T.Tensor((28, 63), "int8"),
T_concat: T.Tensor((28, 128), "int8"),
):
- placeholder_3 = T.decl_tensor((1792,), "int8", data=placeholder.data)
- placeholder_1_1 = T.decl_tensor((28,), "int8", data=placeholder_1.data)
- placeholder_2_1 = T.decl_tensor((1764,), "int8", data=placeholder_2.data)
- T_concat_1 = T.decl_tensor((3584,), "int8", data=T_concat.data)
+ placeholder_3 = T.decl_tensor(
+ placeholder.data, ty_args=[T.Tensor((1792,), "int8", data=placeholder.data)]
+ )
+ placeholder_1_1 = T.decl_tensor(
+ placeholder_1.data, ty_args=[T.Tensor((28,), "int8", data=placeholder_1.data)]
+ )
+ placeholder_2_1 = T.decl_tensor(
+ placeholder_2.data, ty_args=[T.Tensor((1764,), "int8", data=placeholder_2.data)]
+ )
+ T_concat_1 = T.decl_tensor(
+ T_concat.data, ty_args=[T.Tensor((3584,), "int8", data=T_concat.data)]
+ )
for i0 in range(28):
T_concat_1[i0 * 128] = placeholder_1_1[i0]
for i1 in range(63):
@@ -593,10 +633,18 @@ def concat_func_end_point_equality_expected(
placeholder_2: T.Tensor((28, 63), "int8"),
T_concat: T.Tensor((28, 128), "int8"),
):
- placeholder_3 = T.decl_tensor((1792,), "int8", data=placeholder.data)
- placeholder_1_1 = T.decl_tensor((28,), "int8", data=placeholder_1.data)
- placeholder_2_1 = T.decl_tensor((1764,), "int8", data=placeholder_2.data)
- T_concat_1 = T.decl_tensor((3584,), "int8", data=T_concat.data)
+ placeholder_3 = T.decl_tensor(
+ placeholder.data, ty_args=[T.Tensor((1792,), "int8", data=placeholder.data)]
+ )
+ placeholder_1_1 = T.decl_tensor(
+ placeholder_1.data, ty_args=[T.Tensor((28,), "int8", data=placeholder_1.data)]
+ )
+ placeholder_2_1 = T.decl_tensor(
+ placeholder_2.data, ty_args=[T.Tensor((1764,), "int8", data=placeholder_2.data)]
+ )
+ T_concat_1 = T.decl_tensor(
+ T_concat.data, ty_args=[T.Tensor((3584,), "int8", data=T_concat.data)]
+ )
for i0 in range(28):
for i1 in range(64):
T_concat_1[i0 * 128 + i1] = placeholder_2_1[i0 * 63 + i1]
@@ -632,10 +680,18 @@ def concat_func_edge_equalities_expected(
placeholder_2: T.Tensor((28, 1), "int8"),
T_concat: T.Tensor((28, 66), "int8"),
):
- placeholder_3 = T.decl_tensor((1792,), "int8", data=placeholder.data)
- placeholder_1_1 = T.decl_tensor((28,), "int8", data=placeholder_1.data)
- placeholder_2_1 = T.decl_tensor((28,), "int8", data=placeholder_2.data)
- T_concat_1 = T.decl_tensor((1848,), "int8", data=T_concat.data)
+ placeholder_3 = T.decl_tensor(
+ placeholder.data, ty_args=[T.Tensor((1792,), "int8", data=placeholder.data)]
+ )
+ placeholder_1_1 = T.decl_tensor(
+ placeholder_1.data, ty_args=[T.Tensor((28,), "int8", data=placeholder_1.data)]
+ )
+ placeholder_2_1 = T.decl_tensor(
+ placeholder_2.data, ty_args=[T.Tensor((28,), "int8", data=placeholder_2.data)]
+ )
+ T_concat_1 = T.decl_tensor(
+ T_concat.data, ty_args=[T.Tensor((1848,), "int8", data=T_concat.data)]
+ )
for i0 in range(28):
T_concat_1[i0 * 66] = placeholder_2_1[i0]
for i1 in range(64):
@@ -675,12 +731,18 @@ def concat_five_tensors_with_equalities_expected(
tensor_e: T.Tensor((28, 1), "int8"), # Used for i1 == 129
T_concat: T.Tensor((28, 129), "int8"),
):
- tensor_a_1 = T.decl_tensor((28,), "int8", data=tensor_a.data)
- tensor_b_1 = T.decl_tensor((1764,), "int8", data=tensor_b.data)
- tensor_c_1 = T.decl_tensor((28,), "int8", data=tensor_c.data)
- tensor_d_1 = T.decl_tensor((1764,), "int8", data=tensor_d.data)
- tensor_e_1 = T.decl_tensor((28,), "int8", data=tensor_e.data)
- T_concat_1 = T.decl_tensor((3612,), "int8", data=T_concat.data)
+ tensor_a_1 = T.decl_tensor(tensor_a.data, ty_args=[T.Tensor((28,), "int8", data=tensor_a.data)])
+ tensor_b_1 = T.decl_tensor(
+ tensor_b.data, ty_args=[T.Tensor((1764,), "int8", data=tensor_b.data)]
+ )
+ tensor_c_1 = T.decl_tensor(tensor_c.data, ty_args=[T.Tensor((28,), "int8", data=tensor_c.data)])
+ tensor_d_1 = T.decl_tensor(
+ tensor_d.data, ty_args=[T.Tensor((1764,), "int8", data=tensor_d.data)]
+ )
+ tensor_e_1 = T.decl_tensor(tensor_e.data, ty_args=[T.Tensor((28,), "int8", data=tensor_e.data)])
+ T_concat_1 = T.decl_tensor(
+ T_concat.data, ty_args=[T.Tensor((3612,), "int8", data=T_concat.data)]
+ )
for i0 in range(28):
T_concat_1[i0 * 129] = tensor_a_1[i0]
for i1 in range(63):
@@ -706,7 +768,7 @@ def nested_partition_with_single_points(A: T.Tensor((25,), "int32")):
@Ts.function
def nested_partition_with_single_points_expected(A: T.Tensor((25,), "int32")):
- A_1 = T.decl_tensor((25,), "int32", data=A.data)
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((25,), "int32", data=A.data)])
for j in range(2):
A_1[j + 3] = j + 3
for j in range(2):
@@ -767,8 +829,8 @@ def test_ignore_loop_partition_hint():
@Ts.function
def before(A: T.Tensor((10), "float32"), D: T.Tensor((10), "float32")):
- B = T.decl_tensor([2], "float32")
- C = T.decl_tensor([2], "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor([2], "float32")])
+ C = T.alloc_tensor(ty_args=[T.Tensor([2], "float32")])
for i in T.serial(12, annotations={"loop_partition_hint": 1}):
if T.ignore_loop_partition(i < 10):
B[i % 2] = A[i] + 1.0
@@ -779,8 +841,8 @@ def test_ignore_loop_partition_hint():
@Ts.function
def expected(A: T.Tensor((10), "float32"), D: T.Tensor((10), "float32")):
- B = T.decl_tensor([2], "float32")
- C = T.decl_tensor([2], "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor([2], "float32")])
+ C = T.alloc_tensor(ty_args=[T.Tensor([2], "float32")])
for i in range(2):
B[i] = A[i] + 1.0
if i == 1:
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
index d9adef349d..67c5a6f6d4 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
@@ -62,9 +62,7 @@ def transformed_elementwise_func(
) -> None:
for i in T.serial(0, 16):
B_new = T.alloc_tensor(
- [1, 16],
- "float32",
- annotations={"tensor_data_alignment": 64},
+ ty_args=[T.Tensor([1, 16], "float32")], annotations={"tensor_data_alignment": 64}
)
for j in T.serial(0, 16):
B_new[0, j] = A[i, j] + 1.0
@@ -101,9 +99,7 @@ def transformed_gpu_func(
i1 = T.launch_thread("threadIdx.x", 2)
i2 = T.launch_thread("vthread", 2)
B = T.alloc_tensor(
- [1, 16],
- "float32",
- scope="local",
+ ty_args=[T.Tensor([1, 16], "float32", scope="local")],
annotations={"tensor_data_alignment": 64},
)
for j in range(0, 16):
@@ -145,9 +141,7 @@ def transformed_symbolic_func(
) -> None:
for i in range(0, n):
B = T.alloc_tensor(
- [m],
- "float32",
- annotations={"tensor_data_alignment": 64},
+ ty_args=[T.Tensor([m], "float32")], annotations={"tensor_data_alignment": 64}
)
for j in range(0, m):
B[j] = A[i, j] + 1.0
@@ -204,14 +198,10 @@ def compacted_multi_alloc_func(A: T.Tensor(32, "float32"), D: T.Tensor(32, "floa
def transformed_multi_alloc_func(A: T.Tensor(32, "float32"), D: T.Tensor(32, "float32")) -> None:
for i in range(0, 32):
B = T.alloc_tensor(
- (32,),
- "float32",
- annotations={"tensor_data_alignment": 64},
+ ty_args=[T.Tensor((32,), "float32")], annotations={"tensor_data_alignment": 64}
)
C = T.alloc_tensor(
- (32,),
- "float32",
- annotations={"tensor_data_alignment": 64},
+ ty_args=[T.Tensor((32,), "float32")], annotations={"tensor_data_alignment": 64}
)
B[i] = A[i] + 1.0
C[i] = A[i] + B[i]
@@ -248,9 +238,7 @@ def transformed_strided_tensor_func(
# body
for i0 in T.serial(4):
B = T.alloc_tensor(
- [4, 16],
- "float32",
- strides=[17, 1],
+ ty_args=[T.Tensor([4, 16], "float32", strides=[17, 1])],
annotations={"tensor_data_alignment": 64},
)
for i1, j in T.grid(4, 16):
@@ -289,9 +277,11 @@ def transformed_symbolic_strided_tensor_func(A: T.Tensor((1, n, 10240))):
padded_size = T.min((n + 63) // 64 * 64, 96)
for i, j, k in T.grid(((n + 63) // 64 * 4 + 7) // 8, 2, 160):
A_pad_shared_dyn = T.alloc_tensor(
- (1, padded_size, 64),
- strides=(72 * padded_size, 72, 1),
- scope="shared.dyn",
+ ty_args=[
+ T.Tensor(
+ (1, padded_size, 64), strides=(72 * padded_size, 72, 1), scope="shared.dyn"
+ )
+ ],
annotations={"tensor_data_alignment": 64},
)
for ax0, ax1 in T.grid(96, 64):
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
index 9a354cb8b7..6e59f8491b 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
@@ -55,13 +55,13 @@ def test_basic():
@Ts.function(private=True)
def main(A: T.Tensor((128, 32), "float32"), B: T.Tensor(128, "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
- A_flat = T.decl_tensor(4096, data=A.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor(4096, data=A.data)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce = T.alloc_tensor((1,), scope="local")
- reduce_1 = T.decl_tensor(1, data=reduce.data, scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ reduce_1 = T.decl_tensor(reduce.data, ty_args=[T.Tensor(1, scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x, y: x + y),
@@ -91,12 +91,12 @@ def test_basic_with_decl_tensor():
@Ts.function(private=True)
def main(A: T.Tensor((128, 32), "float32"), B: T.Tensor(128, "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
- A_flat = T.decl_tensor(4096, data=A.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor(4096, data=A.data)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce = T.decl_tensor(1, dtype="float32", scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor(1, "float32", scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x, y: (x + y,)),
@@ -124,16 +124,18 @@ def test_reduce_summation():
@Ts.function(private=True)
def main(A: T.Tensor((128, 128), "float32"), B: T.Tensor(128, "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
- A_flat = T.decl_tensor(16384, data=A.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor(16384, data=A.data)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- normal_reduce = T.alloc_tensor((1,), scope="local")
- normal_reduce_1 = T.decl_tensor(1, data=normal_reduce.data, scope="local")
+ normal_reduce = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ normal_reduce_1 = T.decl_tensor(
+ normal_reduce.data, ty_args=[T.Tensor(1, scope="local")]
+ )
- reduce = T.alloc_tensor((1,), scope="local")
- reduce_1 = T.decl_tensor(1, data=reduce.data, scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ reduce_1 = T.decl_tensor(reduce.data, ty_args=[T.Tensor(1, scope="local")])
normal_reduce_1[0] = T.float32(0)
@@ -168,10 +170,12 @@ def test_multi_group_reduction():
def main(A: T.Tensor((32, 32), "float32"), B: T.Tensor((32,), "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
threadIdx_y = T.launch_thread("threadIdx.y", 32)
- cross_thread_B = T.alloc_tensor((1,), scope="local")
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- cross_thread_B_1 = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
- A_1 = T.decl_tensor((1024,), data=A.data)
+ cross_thread_B_1 = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1024,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -181,7 +185,7 @@ def test_multi_group_reduction():
(threadIdx_x,),
)
if threadIdx_x == 0:
- B_1 = T.decl_tensor((32,), data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((32,), data=B.data)])
B_1[threadIdx_y] = cross_thread_B_1[0]
After = transform(Before)
@@ -199,10 +203,12 @@ def test_multi_group_reduction_consumed_through_alias():
def main(A: T.Tensor((4, 128), "float32"), B: T.Tensor((4,), "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
threadIdx_y = T.launch_thread("threadIdx.y", 4)
- cross_thread_B = T.alloc_tensor((1,), scope="local")
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- cross_thread_B_alias = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
- A_flat = T.decl_tensor((512,), data=A.data)
+ cross_thread_B_alias = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor((512,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -213,7 +219,7 @@ def test_multi_group_reduction_consumed_through_alias():
)
cross_thread_B_alias[0] = cross_thread_B[0]
if threadIdx_x == 0:
- B_flat = T.decl_tensor((4,), data=B.data)
+ B_flat = T.decl_tensor(B.data, ty_args=[T.Tensor((4,), data=B.data)])
B_flat[threadIdx_y] = cross_thread_B_alias[0]
After = transform(Before)
@@ -235,9 +241,9 @@ def test_multi_group_reduction_with_alias_declared_after_allreduce():
def main(A: T.Tensor((4, 128), "float32"), B: T.Tensor((4,), "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
threadIdx_y = T.launch_thread("threadIdx.y", 4)
- cross_thread_B = T.alloc_tensor((1,), scope="local")
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- A_flat = T.decl_tensor((512,), data=A.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor((512,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -246,10 +252,12 @@ def test_multi_group_reduction_with_alias_declared_after_allreduce():
(cross_thread_B[0],),
(threadIdx_x,),
)
- cross_thread_B_alias = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
+ cross_thread_B_alias = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
cross_thread_B[0] = cross_thread_B_alias[0]
if threadIdx_x == 0:
- B_flat = T.decl_tensor((4,), data=B.data)
+ B_flat = T.decl_tensor(B.data, ty_args=[T.Tensor((4,), data=B.data)])
B_flat[threadIdx_y] = cross_thread_B_alias[0]
After = transform(Before)
@@ -271,10 +279,12 @@ def test_multi_group_mask1():
def main(A: T.Tensor((32, 8), "float32"), B: T.Tensor((32,), "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
threadIdx_y = T.launch_thread("threadIdx.y", 32)
- cross_thread_B = T.alloc_tensor((1,), scope="local")
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 8)
- cross_thread_B_1 = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
- A_1 = T.decl_tensor((256,), data=A.data)
+ cross_thread_B_1 = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((256,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -284,7 +294,7 @@ def test_multi_group_mask1():
(threadIdx_x,),
)
if threadIdx_x == 0:
- B_1 = T.decl_tensor((32,), data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((32,), data=B.data)])
B_1[threadIdx_y] = cross_thread_B_1[0]
After = transform(Before)
@@ -303,9 +313,11 @@ def test_multi_warp_reduce1():
T.func_attr({"target": T.target("cuda", host="llvm")})
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- cross_thread_B = T.alloc_tensor((1,), scope="local")
- cross_thread_B_1 = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
- A_1 = T.decl_tensor((16384,), data=A.data)
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ cross_thread_B_1 = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((16384,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -315,7 +327,7 @@ def test_multi_warp_reduce1():
(threadIdx_x,),
)
if threadIdx_x == 0:
- B_1 = T.decl_tensor((128,), data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((128,), data=B.data)])
B_1[i] = cross_thread_B_1[0]
After = transform(Before)
@@ -334,9 +346,11 @@ def test_multi_warp_reduce2():
def main(A: T.Tensor((1, 1024), "float32"), B: T.Tensor((1,), "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
threadIdx_x = T.launch_thread("threadIdx.x", 1024)
- cross_thread_B = T.alloc_tensor((1,), scope="local")
- cross_thread_B_1 = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
- A_1 = T.decl_tensor((1024,), data=A.data)
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ cross_thread_B_1 = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((1024,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -346,7 +360,7 @@ def test_multi_warp_reduce2():
(threadIdx_x,),
)
if threadIdx_x == 0:
- B_1 = T.decl_tensor((1,), data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((1,), data=B.data)])
B_1[0] = cross_thread_B_1[0]
After = transform(Before)
@@ -365,10 +379,12 @@ def test_multi_group_multi_warp_reduction():
def main(A: T.Tensor((4, 128), "float32"), B: T.Tensor((4,), "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
threadIdx_y = T.launch_thread("threadIdx.y", 4)
- cross_thread_B = T.alloc_tensor((1,), scope="local")
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- cross_thread_B_1 = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
- A_1 = T.decl_tensor((512,), data=A.data)
+ cross_thread_B_1 = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((512,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -378,7 +394,7 @@ def test_multi_group_multi_warp_reduction():
(threadIdx_x,),
)
if threadIdx_x == 0:
- B_1 = T.decl_tensor((4,), data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((4,), data=B.data)])
B_1[threadIdx_y] = cross_thread_B_1[0]
After = transform(Before)
@@ -397,15 +413,17 @@ def test_multi_group_multi_warp_predicated_reduction():
def main(A: T.Tensor((2, 70), "float32"), B: T.Tensor((2,), "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
threadIdx_y = T.launch_thread("threadIdx.y", 2)
- in_thread_B = T.alloc_tensor((1,), scope="local")
- cross_thread_B = T.alloc_tensor((1,), scope="local")
+ in_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 512)
- in_thread_B_1 = T.decl_tensor((1,), data=in_thread_B.data, scope="local")
+ in_thread_B_1 = T.decl_tensor(in_thread_B.data, ty_args=[T.Tensor((1,), scope="local")])
in_thread_B_1[0] = T.float32(0)
if threadIdx_x < 70:
- A_1 = T.decl_tensor((140,), data=A.data)
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((140,), data=A.data)])
in_thread_B_1[0] = in_thread_B_1[0] + A_1[threadIdx_y * 70 + threadIdx_x]
- cross_thread_B_1 = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
+ cross_thread_B_1 = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -415,7 +433,7 @@ def test_multi_group_multi_warp_predicated_reduction():
(threadIdx_x,),
)
if threadIdx_x == 0:
- B_1 = T.decl_tensor((2,), data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((2,), data=B.data)])
B_1[threadIdx_y] = cross_thread_B_1[0]
After = transform(Before)
@@ -445,12 +463,14 @@ def test_metal_no_mask():
}
)
blockIdx_x = T.launch_thread("blockIdx.x", 1)
- cross_thread_B = T.alloc_tensor((1,), scope="local")
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), scope="local")])
threadIdx_z = T.launch_thread("threadIdx.z", 1)
threadIdx_y = T.launch_thread("threadIdx.y", 2)
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- cross_thread_B_1 = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
- A_1 = T.decl_tensor((256,), data=A.data)
+ cross_thread_B_1 = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((256,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -460,7 +480,7 @@ def test_metal_no_mask():
(threadIdx_x,),
)
if threadIdx_x == 0:
- B_1 = T.decl_tensor((2,), data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((2,), data=B.data)])
B_1[threadIdx_y] = cross_thread_B_1[0]
After = transform(Before)
@@ -489,14 +509,14 @@ def test_webgpu_warp_reduce():
),
}
)
- A_flat = T.decl_tensor(4096, data=A.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor(4096, data=A.data)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce_data = T.alloc_tensor((1,), "float32", scope="local")
- reduce = T.decl_tensor(1, data=reduce_data.data, scope="local")
- reduce_alias = T.decl_tensor(1, data=reduce.data, scope="local")
+ reduce_data = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ reduce = T.decl_tensor(reduce_data.data, ty_args=[T.Tensor(1, scope="local")])
+ reduce_alias = T.decl_tensor(reduce.data, ty_args=[T.Tensor(1, scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x, y: (x + y,)),
@@ -539,12 +559,14 @@ def test_webgpu_multi_warp_reduce():
}
)
blockIdx_x = T.launch_thread("blockIdx.x", 1)
- cross_thread_B = T.alloc_tensor((1,), "float32", scope="local")
+ cross_thread_B = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
threadIdx_z = T.launch_thread("threadIdx.z", 1)
threadIdx_y = T.launch_thread("threadIdx.y", 2)
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- cross_thread_B_1 = T.decl_tensor((1,), data=cross_thread_B.data, scope="local")
- A_1 = T.decl_tensor((256,), data=A.data)
+ cross_thread_B_1 = T.decl_tensor(
+ cross_thread_B.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((256,), data=A.data)])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -554,7 +576,7 @@ def test_webgpu_multi_warp_reduce():
(threadIdx_x,),
)
if threadIdx_x == 0:
- B_1 = T.decl_tensor((2,), data=B.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((2,), data=B.data)])
B_1[threadIdx_y] = cross_thread_B_1[0]
After = transform(Before)
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py b/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
index e042e7dceb..fe348d17b7 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
@@ -44,15 +44,15 @@ def test_matmul_t_tensor():
B: T.Tensor((1024, 1024), "float16"),
matmul: T.Tensor((1024, 1024), "float32"),
):
- A_flat = T.decl_tensor(1048576, "float16", data=A.data)
- B_flat = T.decl_tensor(1048576, "float16", data=B.data)
- matmul_flat = T.decl_tensor(1048576, data=matmul.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor(1048576, "float16", data=A.data)])
+ B_flat = T.decl_tensor(B.data, ty_args=[T.Tensor(1048576, "float16", data=B.data)])
+ matmul_flat = T.decl_tensor(matmul.data, ty_args=[T.Tensor(1048576, data=matmul.data)])
threadIdx_x = T.launch_thread("threadIdx.x", 16)
- C_local = T.alloc_tensor((1,), "float32", scope="local")
- A_sh = T.alloc_tensor((256,), "float16", scope="shared.dyn")
- B_sh = T.alloc_tensor((256,), "float16", scope="shared.dyn")
- C_sh = T.alloc_tensor((256,), "float32", scope="shared.dyn")
+ C_local = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ A_sh = T.alloc_tensor(ty_args=[T.Tensor((256,), "float16", scope="shared.dyn")])
+ B_sh = T.alloc_tensor(ty_args=[T.Tensor((256,), "float16", scope="shared.dyn")])
+ C_sh = T.alloc_tensor(ty_args=[T.Tensor((256,), "float32", scope="shared.dyn")])
threadIdx_y = T.launch_thread("threadIdx.y", 16)
blockIdx_x = T.launch_thread("blockIdx.x", 64)
blockIdx_y = T.launch_thread("blockIdx.y", 64)
@@ -89,18 +89,24 @@ def test_matmul_t_tensor():
B: T.Tensor((1024, 1024), "float16"),
matmul: T.Tensor((1024, 1024), "float32"),
):
- A_flat = T.decl_tensor(1048576, "float16", data=A.data)
- B_flat = T.decl_tensor(1048576, "float16", data=B.data)
- matmul_flat = T.decl_tensor(1048576, data=matmul.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor(1048576, "float16", data=A.data)])
+ B_flat = T.decl_tensor(B.data, ty_args=[T.Tensor(1048576, "float16", data=B.data)])
+ matmul_flat = T.decl_tensor(matmul.data, ty_args=[T.Tensor(1048576, data=matmul.data)])
threadIdx_x = T.launch_thread("threadIdx.x", 16)
- buf_dyn_shmem = T.alloc_tensor((1024,), "uint8", scope="shared.dyn")
+ buf_dyn_shmem = T.alloc_tensor(ty_args=[T.Tensor((1024,), "uint8", scope="shared.dyn")])
- C_local = T.alloc_tensor((1,), "float32", scope="local")
- A_sh = T.decl_tensor(256, "float16", data=buf_dyn_shmem.data, scope="shared.dyn")
- B_sh = T.decl_tensor(256, "float16", data=buf_dyn_shmem.data, scope="shared.dyn")
- C_sh = T.decl_tensor(256, "float32", data=buf_dyn_shmem.data, scope="shared.dyn")
+ C_local = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ A_sh = T.decl_tensor(
+ buf_dyn_shmem.data, ty_args=[T.Tensor(256, "float16", scope="shared.dyn")]
+ )
+ B_sh = T.decl_tensor(
+ buf_dyn_shmem.data, ty_args=[T.Tensor(256, "float16", scope="shared.dyn")]
+ )
+ C_sh = T.decl_tensor(
+ buf_dyn_shmem.data, ty_args=[T.Tensor(256, "float32", scope="shared.dyn")]
+ )
threadIdx_y = T.launch_thread("threadIdx.y", 16)
blockIdx_x = T.launch_thread("blockIdx.x", 64)
@@ -128,7 +134,7 @@ def test_matmul_t_tensor():
After = transform(Before)
script = After["main"].script()
# Verify merged allocation: one shared.dyn tensor of 1024 bytes (256*2 float16 + 256 float32)
- assert "alloc_tensor((1024,)" in script
+ assert "alloc_tensor(ty_args=[T.Tensor((1024,)" in script
assert '"uint8"' in script
assert '"shared.dyn"' in script
# Verify storage sync calls preserved
@@ -155,15 +161,15 @@ def test_matmul_decl_tensor():
B: T.Tensor((1024, 1024), "float16"),
matmul: T.Tensor((1024, 1024), "float32"),
):
- A_flat = T.decl_tensor(1048576, "float16", data=A.data)
- B_flat = T.decl_tensor(1048576, "float16", data=B.data)
- matmul_flat = T.decl_tensor(1048576, data=matmul.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor(1048576, "float16", data=A.data)])
+ B_flat = T.decl_tensor(B.data, ty_args=[T.Tensor(1048576, "float16", data=B.data)])
+ matmul_flat = T.decl_tensor(matmul.data, ty_args=[T.Tensor(1048576, data=matmul.data)])
threadIdx_x = T.launch_thread("threadIdx.x", 16)
- C_local = T.alloc_tensor((1,), "float32", scope="local")
- A_sh = T.alloc_tensor((256,), "float16", scope="shared.dyn")
- B_sh = T.alloc_tensor((256,), "float16", scope="shared.dyn")
- C_sh = T.alloc_tensor((256,), "float32", scope="shared.dyn")
+ C_local = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ A_sh = T.alloc_tensor(ty_args=[T.Tensor((256,), "float16", scope="shared.dyn")])
+ B_sh = T.alloc_tensor(ty_args=[T.Tensor((256,), "float16", scope="shared.dyn")])
+ C_sh = T.alloc_tensor(ty_args=[T.Tensor((256,), "float32", scope="shared.dyn")])
threadIdx_y = T.launch_thread("threadIdx.y", 16)
blockIdx_x = T.launch_thread("blockIdx.x", 64)
blockIdx_y = T.launch_thread("blockIdx.y", 64)
@@ -195,7 +201,7 @@ def test_matmul_decl_tensor():
After = transform(Before)
script = After["main"].script()
# Verify merged allocation: one shared.dyn tensor of 1024 bytes
- assert "alloc_tensor((1024,)" in script
+ assert "alloc_tensor(ty_args=[T.Tensor((1024,)" in script
assert '"uint8"' in script
assert '"shared.dyn"' in script
assert "gpu_storage_sync" in script
@@ -211,14 +217,14 @@ def test_simple_alloc_no_reuse():
@Ts.function
def main():
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- A_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
- B_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
+ A_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
+ B_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
B_sh[threadIdx_x] = A_sh[threadIdx_x]
After = transform(Before)
script = After["main"].script()
# Verify merged allocation: 1024 bytes (128*4 + 128*4)
- assert "alloc_tensor((1024,)" in script
+ assert "alloc_tensor(ty_args=[T.Tensor((1024,)" in script
assert '"uint8"' in script
assert '"shared.dyn"' in script
# Verify offset indexing
@@ -234,15 +240,15 @@ def test_simple_alloc_reuse():
@Ts.function
def main():
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- A_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
- B_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
+ A_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
+ B_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
A_sh[threadIdx_x] = 0
B_sh[threadIdx_x] = 0
After = transform(Before)
script = After["main"].script()
# Verify merged allocation: 512 bytes (128*4, reusable)
- assert "alloc_tensor((512,)" in script
+ assert "alloc_tensor(ty_args=[T.Tensor((512,)" in script
assert '"uint8"' in script
assert '"shared.dyn"' in script
@@ -256,8 +262,8 @@ def test_async_copy():
@Ts.function
def main(A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")):
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- A_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
- B_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
+ A_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
+ B_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
T.s_tir.cp_async_raw(A_sh.data, threadIdx_x, A.data, threadIdx_x, 512, ty="float32")
T.s_tir.cp_async_raw(B_sh.data, threadIdx_x, B.data, threadIdx_x, 512, ty="float32")
@@ -273,12 +279,12 @@ def test_async_copy():
# offsets remain in float32 elements and are scaled by the intrinsic
# lowering, rather than being pre-scaled as byte offsets here.
assert (
- 'A_sh = T.decl_tensor((128,), "float32", data=buf_dyn_shmem.data, '
- 'scope="shared.dyn", layout=None)' in script
+ "A_sh = T.decl_tensor(buf_dyn_shmem.data, ty_args=[T.Tensor((128,), "
+ '"float32", scope="shared.dyn", layout=None)])' in script
)
assert (
- 'B_sh = T.decl_tensor((128,), "float32", data=buf_dyn_shmem.data, '
- 'scope="shared.dyn", layout=None)' in script
+ "B_sh = T.decl_tensor(buf_dyn_shmem.data, ty_args=[T.Tensor((128,), "
+ '"float32", scope="shared.dyn", layout=None)])' in script
)
assert (
"T.s_tir.cp_async_raw(A_sh.data, threadIdx_x, "
@@ -299,17 +305,24 @@ def test_decl_tensor_alias_extends_allocation_lifetime():
@Ts.function
def main(C: T.Tensor((128,), "float32")):
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- A_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
- B_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
- A_view = T.decl_tensor((128,), "float32", data=A_sh.data, scope="shared.dyn")
- B_view = T.decl_tensor((128,), "float32", data=B_sh.data, scope="shared.dyn")
+ A_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
+ B_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
+ A_view = T.decl_tensor(
+ A_sh.data, ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")]
+ )
+ B_view = T.decl_tensor(
+ B_sh.data, ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")]
+ )
A_view[threadIdx_x] = T.float32(1)
B_view[threadIdx_x] = T.float32(2)
C[threadIdx_x] = A_view[threadIdx_x] + B_view[threadIdx_x]
After = transform(Before)
script = After["main"].script()
- assert 'alloc_tensor((1024,), "uint8", scope="shared.dyn", layout=None)' in script
+ assert (
+ 'alloc_tensor(ty_args=[T.Tensor((1024,), "uint8", scope="shared.dyn", layout=None)])'
+ in script
+ )
assert "B_view[threadIdx_x + 128]" in script
assert "A_view[threadIdx_x]" in script
@@ -331,21 +344,21 @@ def test_multi_thread_extent_blocks():
X: T.Tensor((128,), "float32"),
Y: T.Tensor((128,), "float32"),
):
- X_flat = T.decl_tensor(128, data=X.data)
- Y_flat = T.decl_tensor(128, data=Y.data)
+ X_flat = T.decl_tensor(X.data, ty_args=[T.Tensor(128, data=X.data)])
+ Y_flat = T.decl_tensor(Y.data, ty_args=[T.Tensor(128, data=Y.data)])
# First kernel launch
with T.launch_thread("threadIdx.x", 128) as tx0:
- A_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
- B_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
+ A_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
+ B_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
A_sh[tx0] = X_flat[tx0]
B_sh[tx0] = A_sh[tx0]
X_flat[tx0] = B_sh[tx0]
# Second kernel launch — must NOT see kernel #0's merged tensor.
with T.launch_thread("threadIdx.x", 128) as tx1:
- C_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
- D_sh = T.alloc_tensor((128,), "float32", scope="shared.dyn")
+ C_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
+ D_sh = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])
C_sh[tx1] = Y_flat[tx1]
D_sh[tx1] = C_sh[tx1]
Y_flat[tx1] = D_sh[tx1]
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py b/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
index 252040f788..6b3bfb603b 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
@@ -357,7 +357,7 @@ def test_tensor_conditional_lowering():
def before(A: T.handle("float32")):
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
for i in range(1):
- A_1 = T.decl_tensor((1,), data=A)
+ A_1 = T.decl_tensor(A, ty_args=[T.Tensor((1,), data=A)])
A_1[i] = 0
after = before
@@ -383,14 +383,14 @@ def test_dltensor_tensor_is_unlowered():
shape_ptr: T.let[T.handle("int64")] = T.abi_field_get(
dlpack_handle, 0, 3, ty=T.handle("int64").ty
)
- shape = T.decl_tensor(ndim, "int64", data=shape_ptr)
- product = T.decl_tensor([], "int64")
+ shape = T.decl_tensor(shape_ptr, ty_args=[T.Tensor(ndim, "int64", data=shape_ptr)])
+ product = T.alloc_tensor(ty_args=[T.Tensor([], "int64")])
product[()] = 1
for dim in range(axis + 1, ndim):
product[()] = product[()] * shape[dim]
return product[()]
else:
- strides = T.decl_tensor(ndim, "int64", data=stride_ptr)
+ strides = T.decl_tensor(stride_ptr, ty_args=[T.Tensor(ndim, "int64", data=stride_ptr)])
stride: T.int64 = strides[axis]
return stride
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py b/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
index bb99ed1c20..9ff7beba3f 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
@@ -27,7 +27,7 @@ def test_rewrite_Select():
class ModuleY:
@Ts.function
def main(i: T.int32):
- A = T.alloc_tensor((100,))
+ A = T.alloc_tensor(ty_args=[T.Tensor((100,))])
T.evaluate(T.Select(i > 1, A[i - 1], T.float32(1.0)))
yy = tvm.s_tir.transform.RewriteUnsafeSelect()(ModuleY)["main"].body.seq[-1].value
@@ -36,7 +36,7 @@ def test_rewrite_Select():
class ModuleZ:
@Ts.function
def main(i: T.int32):
- A = T.alloc_tensor((100,))
+ A = T.alloc_tensor(ty_args=[T.Tensor((100,))])
T.evaluate(
T.Select(
T.Select(i > 1, A[i - 1], T.float32(1.0)) > T.float32(0.0), A[i], T.float32(0.1)
@@ -49,7 +49,7 @@ def test_rewrite_Select():
class ModuleA:
@Ts.function
def main(i: T.int32):
- A = T.alloc_tensor((100,))
+ A = T.alloc_tensor(ty_args=[T.Tensor((100,))])
# Inline y and z to avoid Let bindings - outer Select condition is safe (no tensor access)
T.evaluate(
T.Select(
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
index 0d4767512f..32cb7ed146 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
@@ -58,13 +58,13 @@ def test_sync_read_thread_id_independent_location():
result_local[0] = result_local[0] + temp_shared[0] * p1[1]
mod = run_passes(func)
- assert 'T.gpu_storage_sync("shared", ty="int32")' in str(mod)
+ assert 'T.gpu_storage_sync("shared")' in str(mod)
def test_sync_inside_condition():
@Ts.function
def func1(A: T.Tensor((4, 4), "float32")) -> None:
- A_shared = T.alloc_tensor((4, 4), "float32", scope="shared")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor((4, 4), "float32", scope="shared")])
bx = T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
if A[0, 0] > 1.0:
@@ -75,7 +75,7 @@ def test_sync_inside_condition():
@Ts.function
def func2(A: T.Tensor((4, 4), "float32")) -> None:
- A_shared = T.alloc_tensor((4, 4), "float32", scope="shared")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor((4, 4), "float32", scope="shared")])
bx = T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
if T.gpu_thread_invariant(A[0, 0] > 1.0):
@@ -86,7 +86,7 @@ def test_sync_inside_condition():
@Ts.function
def func3(A: T.Tensor((4, 4), "float32")) -> None:
- A_shared = T.alloc_tensor((4, 4), "float32", scope="shared")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor((4, 4), "float32", scope="shared")])
bx = T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
while T.gpu_thread_invariant(A[0, 0] > 1.0):
@@ -100,43 +100,43 @@ def test_sync_inside_condition():
for func in (func2, func3):
mod = s_tir.transform.ThreadSync("shared")(tvm.IRModule.from_expr(func))
- assert 'T.gpu_storage_sync("shared", ty="int32")' in mod.script()
+ assert 'T.gpu_storage_sync("shared")' in mod.script()
def test_sync_shared_dyn():
@Ts.function(private=True)
def func(A: T.Tensor((4, 4), "float32"), E: T.Tensor((4, 4), "float32")):
blockIdx_x = T.launch_thread("blockIdx.x", 1)
- B = T.alloc_tensor((24,), "float32", scope="shared.dyn")
- C = T.alloc_tensor((1,), "float32", scope="local")
- D = T.alloc_tensor((16,), "float32", scope="shared.dyn")
+ B = T.alloc_tensor(ty_args=[T.Tensor((24,), "float32", scope="shared.dyn")])
+ C = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ D = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32", scope="shared.dyn")])
threadIdx_x = T.launch_thread("threadIdx.x", 16)
- B_1 = T.decl_tensor((24,), data=B.data, scope="shared.dyn")
- A_1 = T.decl_tensor((16,), data=A.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((24,), scope="shared.dyn")])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((16,), data=A.data)])
B_1[threadIdx_x // 4 * 6 + threadIdx_x % 4] = A_1[threadIdx_x]
- C_1 = T.decl_tensor((1,), data=C.data, scope="local")
+ C_1 = T.decl_tensor(C.data, ty_args=[T.Tensor((1,), scope="local")])
C_1[0] = B_1[threadIdx_x // 4 * 6 + threadIdx_x % 4]
- D_1 = T.decl_tensor((16,), data=D.data, scope="shared.dyn")
+ D_1 = T.decl_tensor(D.data, ty_args=[T.Tensor((16,), scope="shared.dyn")])
D_1[threadIdx_x] = C_1[0]
- E_1 = T.decl_tensor((16,), data=E.data)
+ E_1 = T.decl_tensor(E.data, ty_args=[T.Tensor((16,), data=E.data)])
E_1[threadIdx_x] = D_1[threadIdx_x]
@Ts.function(private=True)
def expected(A: T.Tensor((4, 4), "float32"), E: T.Tensor((4, 4), "float32")):
blockIdx_x = T.launch_thread("blockIdx.x", 1)
- B_1 = T.alloc_tensor((24,), "float32", scope="shared.dyn")
- C_1 = T.alloc_tensor((1,), "float32", scope="local")
- D_1 = T.alloc_tensor((16,), "float32", scope="shared.dyn")
+ B_1 = T.alloc_tensor(ty_args=[T.Tensor((24,), "float32", scope="shared.dyn")])
+ C_1 = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ D_1 = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32", scope="shared.dyn")])
threadIdx_x = T.launch_thread("threadIdx.x", 16)
- B_1_1 = T.decl_tensor((24,), data=B_1.data, scope="shared.dyn")
- A_1 = T.decl_tensor((16,), data=A.data)
+ B_1_1 = T.decl_tensor(B_1.data, ty_args=[T.Tensor((24,), scope="shared.dyn")])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((16,), data=A.data)])
B_1_1[threadIdx_x // 4 * 6 + threadIdx_x % 4] = A_1[threadIdx_x]
- C_1_1 = T.decl_tensor((1,), data=C_1.data, scope="local")
+ C_1_1 = T.decl_tensor(C_1.data, ty_args=[T.Tensor((1,), scope="local")])
C_1_1[0] = B_1_1[threadIdx_x // 4 * 6 + threadIdx_x % 4]
- D_1_1 = T.decl_tensor((16,), data=D_1.data, scope="shared.dyn")
+ D_1_1 = T.decl_tensor(D_1.data, ty_args=[T.Tensor((16,), scope="shared.dyn")])
T.evaluate(T.call_intrin("tirx.gpu_storage_sync", "shared.dyn", ty="int32"))
D_1_1[threadIdx_x] = C_1_1[0]
- E_1 = T.decl_tensor((16,), data=E.data)
+ E_1 = T.decl_tensor(E.data, ty_args=[T.Tensor((16,), data=E.data)])
E_1[threadIdx_x] = D_1_1[threadIdx_x]
mod = tvm.IRModule({"main": func})
@@ -148,11 +148,15 @@ def test_sync_shared_aliasing_tensor_views():
@Ts.function(private=True)
def func(A: T.Tensor((64,), "float32")):
blockIdx_x = T.launch_thread("blockIdx.x", 1)
- shared_storage = T.alloc_tensor((32,), "float16", scope="shared")
- local = T.alloc_tensor((1,), "float32", scope="local")
+ shared_storage = T.alloc_tensor(ty_args=[T.Tensor((32,), "float16", scope="shared")])
+ local = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- shared_half = T.decl_tensor((32,), "float16", data=shared_storage.data, scope="shared")
- shared_float = T.decl_tensor((16,), "float32", data=shared_storage.data, scope="shared")
+ shared_half = T.decl_tensor(
+ shared_storage.data, ty_args=[T.Tensor((32,), "float16", scope="shared")]
+ )
+ shared_float = T.decl_tensor(
+ shared_storage.data, ty_args=[T.Tensor((16,), "float32", scope="shared")]
+ )
for i in range(2):
shared_half[threadIdx_x] = T.Cast("float16", A[i * 32 + threadIdx_x])
T.gpu_storage_sync("shared")
@@ -173,14 +177,16 @@ def test_sync_bind():
@Ts.function(private=True)
def func(A: T.Tensor((16 * 512), "float32")):
blockIdx_x = T.launch_thread("blockIdx.x", 16)
- A_shared = T.alloc_tensor((512,), "float32", scope="shared")
- in_thread_A_temp = T.alloc_tensor((1,), "float32", scope="local")
- cross_thread_A_temp = T.alloc_tensor((1,), "float32", scope="local")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor((512,), "float32", scope="shared")])
+ in_thread_A_temp = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ cross_thread_A_temp = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- A_shared_1 = T.decl_tensor((512,), data=A_shared.data, scope="shared")
+ A_shared_1 = T.decl_tensor(A_shared.data, ty_args=[T.Tensor((512,), scope="shared")])
for ax0 in range(512):
A_shared_1[ax0] = A[blockIdx_x * 512 + ax0]
- in_thread_A_temp_1 = T.decl_tensor((1,), data=in_thread_A_temp.data, scope="local")
+ in_thread_A_temp_1 = T.decl_tensor(
+ in_thread_A_temp.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
in_thread_A_temp_1[0] = T.float32(0)
A_temp_1 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x])
in_thread_A_temp_1[0] = A_temp_1
@@ -190,7 +196,9 @@ def test_sync_bind():
in_thread_A_temp_1[0] = A_temp_3
A_temp_4 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x + 384])
in_thread_A_temp_1[0] = A_temp_4
- cross_thread_A_temp_1 = T.decl_tensor((1,), data=cross_thread_A_temp.data, scope="local")
+ cross_thread_A_temp_1 = T.decl_tensor(
+ cross_thread_A_temp.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
(T.float32(0),),
@@ -203,14 +211,16 @@ def test_sync_bind():
@Ts.function(private=True)
def expected(A: T.Tensor((8192,), "float32")):
blockIdx_x = T.launch_thread("blockIdx.x", 16)
- A_shared_1 = T.alloc_tensor((512,), "float32", scope="shared")
- in_thread_A_temp_1 = T.alloc_tensor((1,), "float32", scope="local")
- cross_thread_A_temp_1 = T.alloc_tensor((1,), "float32", scope="local")
+ A_shared_1 = T.alloc_tensor(ty_args=[T.Tensor((512,), "float32", scope="shared")])
+ in_thread_A_temp_1 = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
+ cross_thread_A_temp_1 = T.alloc_tensor(ty_args=[T.Tensor((1,), "float32", scope="local")])
threadIdx_x = T.launch_thread("threadIdx.x", 128)
- A_shared_1_1 = T.decl_tensor((512,), data=A_shared_1.data, scope="shared")
+ A_shared_1_1 = T.decl_tensor(A_shared_1.data, ty_args=[T.Tensor((512,), scope="shared")])
for ax0 in range(512):
A_shared_1_1[ax0] = A[blockIdx_x * 512 + ax0]
- in_thread_A_temp_1_1 = T.decl_tensor((1,), data=in_thread_A_temp_1.data, scope="local")
+ in_thread_A_temp_1_1 = T.decl_tensor(
+ in_thread_A_temp_1.data, ty_args=[T.Tensor((1,), scope="local")]
+ )
in_thread_A_temp_1_1[0] = T.float32(0)
T.evaluate(T.call_intrin("tirx.gpu_storage_sync", "shared", ty="int32"))
A_temp_1 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x])
@@ -222,7 +232,7 @@ def test_sync_bind():
A_temp_4 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x + 384])
in_thread_A_temp_1_1[0] = A_temp_4
cross_thread_A_temp_1_1 = T.decl_tensor(
- (1,), data=cross_thread_A_temp_1.data, scope="local"
+ cross_thread_A_temp_1.data, ty_args=[T.Tensor((1,), scope="local")]
)
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x0, y0: (x0 + y0,)),
@@ -244,7 +254,7 @@ def test_hoist_if_preserves_sync_placement(loop):
class Before:
@T.function
def main(A: T.Tensor((4,), "int32")):
- B = T.alloc_tensor((4,), "int32", scope="shared")
+ B = T.alloc_tensor(ty_args=[T.Tensor((4,), "int32", scope="shared")])
tx = T.launch_thread("threadIdx.x", 4)
B[tx] = A[tx]
for i in loop(1, 4):
diff --git a/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py b/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py
index d6e96b2642..e2cd51d75c 100644
--- a/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py
+++ b/tests/python/tirx-analysis/test_tir_analysis_undefined_vars.py
@@ -39,13 +39,8 @@ def test_decl_tensor_data_is_use():
buf,
tvm.ir.Call(
"tirx.decl_tensor",
- [
- data_ptr,
- tvm.ir.Tuple(buf.shape),
- tvm.ir.DataTypeImm(tvm.DataType(buf.dtype)),
- tvm.ir.StringImm(buf.scope()),
- ],
- ty=buf.ty,
+ [data_ptr],
+ ty_args=[buf.ty],
),
)
stmt = tvm.ir.SeqStmt([decl, body])
@@ -74,13 +69,8 @@ def test_decl_tensor_elem_offset_is_use():
buf,
tvm.ir.Call(
"tirx.decl_tensor",
- [
- data_ptr,
- tvm.ir.Tuple(buf.shape),
- tvm.ir.DataTypeImm(tvm.DataType(buf.dtype)),
- tvm.ir.StringImm(buf.scope()),
- ],
- ty=buf.ty,
+ [data_ptr],
+ ty_args=[buf.ty],
),
)
stmt = tvm.ir.SeqStmt([decl, body])
@@ -107,13 +97,9 @@ def test_alloc_tensor_data_is_def():
buf,
tvm.ir.Call(
"tirx.alloc_tensor",
- [
- tvm.ir.Tuple(buf.shape),
- tvm.ir.DataTypeImm(tvm.DataType(buf.dtype)),
- tvm.ir.StringImm(buf.scope()),
- ],
+ [],
attrs=tvm.ir.DictAttrs({}),
- ty=buf.ty,
+ ty_args=[buf.ty],
),
)
stmt = tvm.ir.SeqStmt([alloc, body])
diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
index 9774d280e4..a4a207af27 100644
--- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
+++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
@@ -34,7 +34,7 @@ def test_pass_simple():
A: T.Tensor((128, 128), "float32"),
C: T.Tensor((128, 128), "float32"),
):
- B = T.alloc_tensor((128, 128), "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor((128, 128), "float32")])
for i, j in T.grid(128, 128):
B[i, j] = A[i, j] * 2.0
for i, j in T.grid(128, 128):
@@ -303,7 +303,7 @@ def test_decl_tensor_is_well_formed():
@T.function
def func(A: T.Tensor((128,), "float32")):
- B = T.alloc_tensor((128,), "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32")])
for i in T.grid(128):
B[i] = A[i] * 2.0
@@ -317,7 +317,7 @@ def test_alloc_tensor_is_well_formed():
class mod:
@T.function
def func(A: T.Tensor((128,), "float32")):
- B = T.alloc_tensor([128], "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor([128], "float32")])
for i in T.grid(128):
B[i] = A[i] * 2.0
@@ -327,7 +327,7 @@ def test_alloc_tensor_is_well_formed():
def test_tensor_load_asserted_type_matches_source_and_indices():
@T.function
def func():
- tensor = T.alloc_tensor((4,), "float32")
+ tensor = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32")])
T.evaluate(tensor[0])
serialized = tvm.ir.save_json(func)
diff --git a/tests/python/tirx-base/test_tir_base.py b/tests/python/tirx-base/test_tir_base.py
index 19130609ed..c9027efaeb 100644
--- a/tests/python/tirx-base/test_tir_base.py
+++ b/tests/python/tirx-base/test_tir_base.py
@@ -198,7 +198,7 @@ def test_continue_statement():
if (i * 10 + j) % 3 != 0:
continue
Out[0] = Out[0] + 1
- k = T.decl_tensor([], "int32")
+ k = T.alloc_tensor(ty_args=[T.Tensor([], "int32")])
k[()] = 0
while k[()] < Out[0]:
k[()] = k[()] + 1
diff --git a/tests/python/tirx-base/test_tir_constructor.py b/tests/python/tirx-base/test_tir_constructor.py
index ea256d0607..0ef66ac0a4 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -333,13 +333,9 @@ def test_stmt_constructor():
buf,
tvm.ir.Call(
"tirx.alloc_tensor",
- [
- tvm.ir.Tuple(buf.shape),
- tvm.ir.DataTypeImm(tvm.DataType(buf.dtype)),
- tvm.ir.StringImm(buf.scope()),
- ],
+ [],
attrs=tvm.ir.DictAttrs({}),
- ty=buf.ty,
+ ty_args=[buf.ty],
),
)
assert _is_tensor_binding(x, "tirx.alloc_tensor")
diff --git a/tests/python/tirx-base/test_tir_ptx_cp_async.py b/tests/python/tirx-base/test_tir_ptx_cp_async.py
index e69ca4bb86..8af3437ad7 100644
--- a/tests/python/tirx-base/test_tir_ptx_cp_async.py
+++ b/tests/python/tirx-base/test_tir_ptx_cp_async.py
@@ -29,7 +29,7 @@ def ptx_cp_async(A: T.Tensor((32, 128), "float16"), B: T.Tensor((32, 128), "floa
T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- A_shared = T.alloc_tensor([32, 128], "float16", scope="shared")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor([32, 128], "float16", scope="shared")])
for i in range(16):
T.evaluate(
T.s_tir.cp_async_raw.legacy(
diff --git a/tests/python/tirx-base/test_tir_ptx_ldmatrix.py b/tests/python/tirx-base/test_tir_ptx_ldmatrix.py
index a0b27f31c7..c94a962689 100644
--- a/tests/python/tirx-base/test_tir_ptx_ldmatrix.py
+++ b/tests/python/tirx-base/test_tir_ptx_ldmatrix.py
@@ -31,8 +31,8 @@ def ptx_ldmatrix(
T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- A_shared = T.alloc_tensor([16, 16], "float16", scope="shared")
- A_local = T.alloc_tensor([8], "float16", scope="local")
+ A_shared = T.alloc_tensor(ty_args=[T.Tensor([16, 16], "float16", scope="shared")])
+ A_local = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
for i in range(8):
A_shared[i * 2 + tx // 16, tx % 16] = A[i * 2 + tx // 16, tx % 16]
T.evaluate(
diff --git a/tests/python/tirx-base/test_tir_ptx_mma.py b/tests/python/tirx-base/test_tir_ptx_mma.py
index a85bf25c2a..5f2fa83f63 100644
--- a/tests/python/tirx-base/test_tir_ptx_mma.py
+++ b/tests/python/tirx-base/test_tir_ptx_mma.py
@@ -35,9 +35,9 @@ def gemm_mma_m8n8k4_row_col_fp64pf64fp64(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([1], "float64", scope="local")
- MultiB = T.decl_tensor([1], "float64", scope="local")
- Accum = T.decl_tensor([2], "float64", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([1], "float64", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([1], "float64", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "float64", scope="local")])
for i in range(2):
Accum[i] = T.float64(0)
@@ -100,9 +100,9 @@ def gemm_mma_m8n8k4_row_row_fp16fp16fp16(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([4], "float16", scope="local")
- MultiB = T.decl_tensor([4], "float16", scope="local")
- Accum = T.decl_tensor([8], "float16", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
for i in range(8):
Accum[i] = T.float32(0)
@@ -176,9 +176,9 @@ def gemm_mma_m8n8k4_row_row_fp16fp16fp32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([4], "float16", scope="local")
- MultiB = T.decl_tensor([4], "float16", scope="local")
- Accum = T.decl_tensor([8], "float32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([8], "float32", scope="local")])
for i in range(8):
Accum[i] = T.float32(0)
@@ -259,9 +259,9 @@ def gemm_mma_m8n8k16_row_col_s8s8s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([4], "int8", scope="local")
- MultiB = T.decl_tensor([4], "int8", scope="local")
- Accum = T.decl_tensor([2], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "int8", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "int8", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "int32", scope="local")])
for i in range(2):
Accum[i] = T.int32(0)
@@ -330,9 +330,9 @@ def gemm_mma_m8n8k16_row_col_s8u8s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([4], "int8", scope="local")
- MultiB = T.decl_tensor([4], "uint8", scope="local")
- Accum = T.decl_tensor([2], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "int8", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "uint8", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "int32", scope="local")])
for i in range(2):
Accum[i] = T.int32(0)
@@ -401,9 +401,9 @@ def gemm_mma_m8n8k32_row_col_s4s4s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([8], "int4", scope="local")
- MultiB = T.decl_tensor([8], "int4", scope="local")
- Accum = T.decl_tensor([2], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "int4", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([8], "int4", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "int32", scope="local")])
for i in range(2):
Accum[i] = T.int32(0)
@@ -466,9 +466,9 @@ def gemm_mma_m8n8k32_row_col_s4u4s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([8], "int4", scope="local")
- MultiB = T.decl_tensor([8], "uint4", scope="local")
- Accum = T.decl_tensor([2], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "int4", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([8], "uint4", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "int32", scope="local")])
for i in range(2):
Accum[i] = T.int32(0)
@@ -531,9 +531,9 @@ def gemm_mma_m16n8k8_row_col_fp16fp16fp32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([4], "float16", scope="local")
- MultiB = T.decl_tensor([2], "float16", scope="local")
- Accum = T.decl_tensor([4], "float32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([2], "float16", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float32", scope="local")])
for i in range(4):
Accum[i] = T.float32(0)
@@ -604,9 +604,9 @@ def gemm_mma_m16n8k16_row_col_fp16fp16fp16(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([8], "float16", scope="local")
- MultiB = T.decl_tensor([4], "float16", scope="local")
- Accum = T.decl_tensor([4], "float16", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
for i in range(4):
Accum[i] = T.float32(0)
@@ -680,9 +680,9 @@ def gemm_mma_m16n8k16_row_col_fp16fp16fp32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([8], "float16", scope="local")
- MultiB = T.decl_tensor([4], "float16", scope="local")
- Accum = T.decl_tensor([4], "float32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float32", scope="local")])
for i in range(4):
Accum[i] = T.float32(0)
@@ -756,9 +756,9 @@ def gemm_mma_m16n8k16_row_col_s8s8s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([8], "int8", scope="local")
- MultiB = T.decl_tensor([4], "int8", scope="local")
- Accum = T.decl_tensor([4], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "int8", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "int8", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
for i in range(4):
Accum[i] = T.int32(0)
@@ -832,9 +832,9 @@ def gemm_mma_m16n8k16_row_col_s8u8s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([8], "int8", scope="local")
- MultiB = T.decl_tensor([4], "uint8", scope="local")
- Accum = T.decl_tensor([4], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "int8", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "uint8", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
for i in range(4):
Accum[i] = T.int32(0)
@@ -908,9 +908,9 @@ def gemm_mma_m16n8k32_row_col_s8s8s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([16], "int8", scope="local")
- MultiB = T.decl_tensor([8], "int8", scope="local")
- Accum = T.decl_tensor([4], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([16], "int8", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([8], "int8", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
for i in range(4):
Accum[i] = T.int32(0)
@@ -984,9 +984,9 @@ def gemm_mma_m16n8k32_row_col_s8u8s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([16], "int8", scope="local")
- MultiB = T.decl_tensor([8], "uint8", scope="local")
- Accum = T.decl_tensor([4], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([16], "int8", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([8], "uint8", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
for i in range(4):
Accum[i] = T.int32(0)
@@ -1060,9 +1060,9 @@ def gemm_mma_m16n8k64_row_col_s4s4s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([32], "int4", scope="local")
- MultiB = T.decl_tensor([16], "int4", scope="local")
- Accum = T.decl_tensor([4], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([32], "int4", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([16], "int4", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
for i in range(4):
Accum[i] = T.int32(0)
@@ -1130,9 +1130,9 @@ def gemm_mma_m16n8k64_row_col_s4u4s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([32], "int4", scope="local")
- MultiB = T.decl_tensor([16], "uint4", scope="local")
- Accum = T.decl_tensor([4], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([32], "int4", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([16], "uint4", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
for i in range(4):
Accum[i] = T.int32(0)
@@ -1200,9 +1200,9 @@ def gemm_mma_m16n8k256_row_col_b1b1s32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.decl_tensor([128], "int1", scope="local")
- MultiB = T.decl_tensor([64], "int1", scope="local")
- Accum = T.decl_tensor([4], "int32", scope="local")
+ MultiA = T.alloc_tensor(ty_args=[T.Tensor([128], "int1", scope="local")])
+ MultiB = T.alloc_tensor(ty_args=[T.Tensor([64], "int1", scope="local")])
+ Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
for i in range(4):
Accum[i] = T.int32(0)
diff --git a/tests/python/tirx-base/test_tir_ptx_mma_sp.py b/tests/python/tirx-base/test_tir_ptx_mma_sp.py
index 05dfe0a2f9..5635220847 100644
--- a/tests/python/tirx-base/test_tir_ptx_mma_sp.py
+++ b/tests/python/tirx-base/test_tir_ptx_mma_sp.py
@@ -54,10 +54,10 @@ def mma_sp_m16n8k16_f16f16f16(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- multi_a = T.decl_tensor([4], "float16", scope="local")
- multi_b = T.decl_tensor([4], "float16", scope="local")
- accum = T.decl_tensor([4], "float16", scope="local")
- meta_local = T.decl_tensor([1], "uint32", scope="local")
+ multi_a = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ multi_b = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ meta_local = T.alloc_tensor(ty_args=[T.Tensor([1], "uint32", scope="local")])
for i in range(4):
accum[i] = T.float16(0)
@@ -69,9 +69,9 @@ def mma_sp_m16n8k16_f16f16f16(
meta_local[0] = metadata[tx // 4]
- a_words = T.decl_tensor([2], "uint32", data=multi_a.data, scope="local")
- b_words = T.decl_tensor([2], "uint32", data=multi_b.data, scope="local")
- acc_words = T.decl_tensor([2], "uint32", data=accum.data, scope="local")
+ a_words = T.decl_tensor(multi_a.data, ty_args=[T.Tensor([2], "uint32", scope="local")])
+ b_words = T.decl_tensor(multi_b.data, ty_args=[T.Tensor([2], "uint32", scope="local")])
+ acc_words = T.decl_tensor(accum.data, ty_args=[T.Tensor([2], "uint32", scope="local")])
T.ptx.mma.sp.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16(
acc_words[0],
acc_words[1],
@@ -101,10 +101,10 @@ def mma_sp_m16n8k16_f16f16f32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- multi_a = T.decl_tensor([4], "float16", scope="local")
- multi_b = T.decl_tensor([4], "float16", scope="local")
- accum = T.decl_tensor([4], "float32", scope="local")
- meta_local = T.decl_tensor([1], "uint32", scope="local")
+ multi_a = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ multi_b = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float32", scope="local")])
+ meta_local = T.alloc_tensor(ty_args=[T.Tensor([1], "uint32", scope="local")])
for i in range(4):
accum[i] = T.float16(0)
@@ -116,8 +116,8 @@ def mma_sp_m16n8k16_f16f16f32(
meta_local[0] = metadata[tx // 4]
- a_words = T.decl_tensor([2], "uint32", data=multi_a.data, scope="local")
- b_words = T.decl_tensor([2], "uint32", data=multi_b.data, scope="local")
+ a_words = T.decl_tensor(multi_a.data, ty_args=[T.Tensor([2], "uint32", scope="local")])
+ b_words = T.decl_tensor(multi_b.data, ty_args=[T.Tensor([2], "uint32", scope="local")])
T.ptx.mma.sp.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32(
accum[0],
accum[1],
@@ -151,10 +151,10 @@ def mma_sp_m16n8k32_f16f16f16(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- multi_a = T.decl_tensor([8], "float16", scope="local")
- multi_b = T.decl_tensor([8], "float16", scope="local")
- accum = T.decl_tensor([4], "float16", scope="local")
- meta_local = T.decl_tensor([1], "uint32", scope="local")
+ multi_a = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
+ multi_b = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
+ accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
+ meta_local = T.alloc_tensor(ty_args=[T.Tensor([1], "uint32", scope="local")])
for i in range(4):
accum[i] = T.float16(0)
@@ -166,9 +166,9 @@ def mma_sp_m16n8k32_f16f16f16(
meta_local[0] = metadata[tx // 4 * 2 + tx % 2]
- a_words = T.decl_tensor([4], "uint32", data=multi_a.data, scope="local")
- b_words = T.decl_tensor([4], "uint32", data=multi_b.data, scope="local")
- acc_words = T.decl_tensor([2], "uint32", data=accum.data, scope="local")
+ a_words = T.decl_tensor(multi_a.data, ty_args=[T.Tensor([4], "uint32", scope="local")])
+ b_words = T.decl_tensor(multi_b.data, ty_args=[T.Tensor([4], "uint32", scope="local")])
+ acc_words = T.decl_tensor(accum.data, ty_args=[T.Tensor([2], "uint32", scope="local")])
T.ptx.mma.sp.sync.aligned.m16n8k32.row.col.f16.f16.f16.f16(
acc_words[0],
acc_words[1],
@@ -202,10 +202,10 @@ def mma_sp_m16n8k32_f16f16f32(
T.launch_thread("blockIdx.y", 1)
T.launch_thread("blockIdx.x", 1)
tx = T.launch_thread("threadIdx.x", 32)
- multi_a = T.decl_tensor([8], "float16", scope="local")
- multi_b = T.decl_tensor([8], "float16", scope="local")
- accum = T.decl_tensor([4], "float32", scope="local")
- meta_local = T.decl_tensor([1], "uint32", scope="local")
+ multi_a = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
+ multi_b = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
+ accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float32", scope="local")])
+ meta_local = T.alloc_tensor(ty_args=[T.Tensor([1], "uint32", scope="local")])
for i in range(4):
accum[i] = T.float16(0)
@@ -217,8 +217,8 @@ def mma_sp_m16n8k32_f16f16f32(
meta_local[0] = metadata[tx // 4 * 2 + tx % 2]
- a_words = T.decl_tensor([4], "uint32", data=multi_a.data, scope="local")
- b_words = T.decl_tensor([4], "uint32", data=multi_b.data, scope="local")
+ a_words = T.decl_tensor(multi_a.data, ty_args=[T.Tensor([4], "uint32", scope="local")])
+ b_words = T.decl_tensor(multi_b.data, ty_args=[T.Tensor([4], "uint32", scope="local")])
T.ptx.mma.sp.sync.aligned.m16n8k32.row.col.f32.f16.f16.f32(
accum[0],
accum[1],
diff --git a/tests/python/tirx-base/test_tir_specialize.py b/tests/python/tirx-base/test_tir_specialize.py
index 2c675da380..bac573da10 100644
--- a/tests/python/tirx-base/test_tir_specialize.py
+++ b/tests/python/tirx-base/test_tir_specialize.py
@@ -80,7 +80,7 @@ n = T.dynamic("n", "int32")
@T.function
def element_wise(A: T.Tensor((m, n), "float32"), C: T.Tensor((m, n), "float32")) -> None:
- B = T.alloc_tensor((m, n), "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor((m, n), "float32")])
for i, j in T.grid(m, n):
B[i, j] = A[i, j] * 2.0
@@ -93,7 +93,7 @@ def element_wise(A: T.Tensor((m, n), "float32"), C: T.Tensor((m, n), "float32"))
def element_wise_128_64(
A: T.Tensor((128, 64), "float32"), C: T.Tensor((128, 64), "float32")
) -> None:
- B = T.alloc_tensor((128, 64), "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor((128, 64), "float32")])
for i, j in T.grid(128, 64):
B[i, j] = A[i, j] * 2.0
@@ -107,7 +107,7 @@ n = T.dynamic("n", "int32")
@T.function
def element_wise_128_n(A: T.Tensor((128, n), "float32"), C: T.Tensor((128, n), "float32")) -> None:
- B = T.alloc_tensor((128, n), "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor((128, n), "float32")])
for i, j in T.grid(128, n):
B[i, j] = A[i, j] * 2.0
@@ -239,13 +239,13 @@ def test_specialize_decl_tensor():
@T.function(private=True)
def before(A_data: T.handle("float32"), A_size: T.int32):
- A_buf = T.decl_tensor(A_size, "float32", data=A_data)
+ A_buf = T.decl_tensor(A_data, ty_args=[T.Tensor(A_size, "float32", data=A_data)])
for i in range(A_size):
A_buf[i] = A_buf[i] * 2.0
@T.function(private=True)
def expected(A_data: T.handle("float32")):
- A_buf = T.decl_tensor(16, "float32", data=A_data)
+ A_buf = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
for i in range(16):
A_buf[i] = A_buf[i] * 2.0
@@ -260,12 +260,12 @@ def test_specialize_preserves_decl_tensor_alias():
@T.function(private=True)
def before(A: T.Tensor((before_n,), "int32"), n: before_n):
- A_flat = T.decl_tensor((n,), "int32", data=A.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor((n,), "int32", data=A.data)])
A_flat[n - 1] = 42
@T.function(private=True)
def expected(A: T.Tensor((8,), "int32")):
- A_flat = T.decl_tensor((8,), "int32", data=A.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor((8,), "int32", data=A.data)])
A_flat[7] = 42
after = before.specialize({before.params[1]: 8})
@@ -282,15 +282,15 @@ def test_specialize_tensor_var_to_var():
@T.function(private=True)
def before(A: T.Tensor([16, 16], "float32"), B: T.Tensor([16, 16], "float32")):
- A_flat = T.decl_tensor([256], "float32", data=A.data)
- B_flat = T.decl_tensor([256], "float32", data=B.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor([256], "float32", data=A.data)])
+ B_flat = T.decl_tensor(B.data, ty_args=[T.Tensor([256], "float32", data=B.data)])
for i in range(256):
B_flat[i] = A_flat[i] * 2.0
@T.function(private=True)
def expected(A: T.Tensor([16, 16], "float32")):
- A_flat = T.decl_tensor([256], "float32", data=A.data)
- B_flat = T.decl_tensor([256], "float32", data=A.data)
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor([256], "float32", data=A.data)])
+ B_flat = T.decl_tensor(A.data, ty_args=[T.Tensor([256], "float32", data=A.data)])
for i in range(256):
B_flat[i] = A_flat[i] * 2.0
@@ -307,15 +307,17 @@ def test_specialize_tensor_var_to_expr():
@T.function(private=True)
def before(A_data: T.handle("float32"), B_data: T.handle("float32")):
- A_buf = T.decl_tensor(32, "float32", data=A_data)
- B_buf = T.decl_tensor(16, "float32", data=B_data)
+ A_buf = T.decl_tensor(A_data, ty_args=[T.Tensor(32, "float32", data=A_data)])
+ B_buf = T.decl_tensor(B_data, ty_args=[T.Tensor(16, "float32", data=B_data)])
for i in range(16):
B_buf[i] = A_buf[i] * 2.0
@T.function(private=True)
def expected(A_data: T.handle("float32")):
- A_buf = T.decl_tensor(32, "float32", data=A_data)
- B_buf = T.decl_tensor(16, "float32", data=T.address_of(A_buf[16]))
+ A_buf = T.decl_tensor(A_data, ty_args=[T.Tensor(32, "float32", data=A_data)])
+ B_buf = T.decl_tensor(
+ T.address_of(A_buf[16]), ty_args=[T.Tensor(16, "float32", data=T.address_of(A_buf[16]))]
+ )
for i in range(16):
B_buf[i] = A_buf[i] * 2.0
diff --git a/tests/python/tirx-base/test_tir_stmt_functor_substitute.py b/tests/python/tirx-base/test_tir_stmt_functor_substitute.py
index e48f905c44..1c79bc9f67 100644
--- a/tests/python/tirx-base/test_tir_stmt_functor_substitute.py
+++ b/tests/python/tirx-base/test_tir_stmt_functor_substitute.py
@@ -62,14 +62,14 @@ def test_substitute_allocate():
class Before:
@T.function
def main(n: T.int32):
- A = T.alloc_tensor((n,), "float32")
+ A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
T.evaluate(A.data)
@I.ir_module
class Expected:
@T.function
def main():
- A = T.alloc_tensor((16,), "float32")
+ A = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
T.evaluate(A.data)
After = _apply_substitute(Before)
@@ -81,7 +81,7 @@ def test_substitute_tensor_load():
class Before:
@T.function
def main(n: T.int32):
- A = T.alloc_tensor((n,), "float32")
+ A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
for i in range(n):
T.evaluate(A[i])
@@ -89,7 +89,7 @@ def test_substitute_tensor_load():
class Expected:
@T.function
def main():
- A = T.alloc_tensor((16,), "float32")
+ A = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
for i in range(16):
T.evaluate(A[i])
@@ -102,14 +102,14 @@ def test_substitute_decl_tensor():
class Before:
@T.function
def main(n: T.int32):
- A = T.alloc_tensor((n,), "float32")
+ A = T.alloc_tensor(ty_args=[T.Tensor((n,), "float32")])
T.evaluate(A.data)
@I.ir_module
class Expected:
@T.function
def main():
- A = T.alloc_tensor((16,), "float32")
+ A = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
T.evaluate(A.data)
After = _apply_substitute(Before)
diff --git a/tests/python/tirx-base/test_tir_tensor.py b/tests/python/tirx-base/test_tir_tensor.py
index 06009c88ec..fbb5043aeb 100644
--- a/tests/python/tirx-base/test_tir_tensor.py
+++ b/tests/python/tirx-base/test_tir_tensor.py
@@ -90,13 +90,8 @@ def test_decl_tensor_physical_data_binding():
tensor,
tvm.ir.Call(
"tirx.decl_tensor",
- [
- data,
- tvm.ir.Tuple(tensor.shape),
- tvm.ir.DataTypeImm(tvm.DataType(tensor.dtype)),
- tvm.ir.StringImm(tensor.scope()),
- ],
- ty=tensor.ty,
+ [data],
+ ty_args=[tensor.ty],
),
)
assert decl.var.same_as(tensor)
diff --git a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py b/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
index 11e3f61b4b..6fa3ec726e 100644
--- a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
+++ b/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
@@ -64,9 +64,9 @@ def test_bf16_simple_store_will_legalize():
Cptr: T.handle("bfloat16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "bfloat16", data=Aptr)
- B = T.decl_tensor((100,), "bfloat16")
- C = T.decl_tensor((100,), "bfloat16", data=Cptr)
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "bfloat16", data=Aptr)])
+ B = T.alloc_tensor(ty_args=[T.Tensor((100,), "bfloat16")])
+ C = T.decl_tensor(Cptr, ty_args=[T.Tensor((100,), "bfloat16", data=Cptr)])
for i in T.grid(100):
B[i] = A[i]
C[i] = T.exp(B[i])
@@ -82,9 +82,9 @@ def test_bf16_simple_store_will_legalize():
Cptr: T.handle("bfloat16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "bfloat16", data=Aptr)
- B = T.decl_tensor((100,), "float32")
- C = T.decl_tensor((100,), "bfloat16", data=Cptr)
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "bfloat16", data=Aptr)])
+ B = T.alloc_tensor(ty_args=[T.Tensor((100,), "float32")])
+ C = T.decl_tensor(Cptr, ty_args=[T.Tensor((100,), "bfloat16", data=Cptr)])
for i in T.grid(100):
B[i] = bf16tof32(A[i])
C[i] = f32tobf16(T.exp(B[i]))
@@ -100,9 +100,9 @@ def test_bf16_simple_store_will_legalize():
Cptr: T.handle("uint16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "uint16", data=Aptr)
- B = T.decl_tensor((100,), "float32")
- C = T.decl_tensor((100,), "uint16", data=Cptr)
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "uint16", data=Aptr)])
+ B = T.alloc_tensor(ty_args=[T.Tensor((100,), "float32")])
+ C = T.decl_tensor(Cptr, ty_args=[T.Tensor((100,), "uint16", data=Cptr)])
for i in T.grid(100):
B[i] = u16tof32(A[i])
C[i] = f32tou16(T.exp(B[i]))
@@ -124,9 +124,9 @@ def test_bf16_masked_load_store_will_legalize():
@T.function
def main(Aptr: T.handle("bfloat16"), Cptr: T.handle("bfloat16")):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((16,), "bfloat16", data=Aptr)
- B = T.decl_tensor((16,), "bfloat16")
- C = T.decl_tensor((16,), "bfloat16", data=Cptr)
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((16,), "bfloat16", data=Aptr)])
+ B = T.alloc_tensor(ty_args=[T.Tensor((16,), "bfloat16")])
+ C = T.decl_tensor(Cptr, ty_args=[T.Tensor((16,), "bfloat16", data=Cptr)])
mask = T.local_scalar("boolx4")
mask = T.Broadcast(T.bool(True), 4)
T.evaluate(
@@ -221,10 +221,10 @@ def test_bf16_storage_compute_scope_will_legalize():
Dptr: T.handle("bfloat16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "bfloat16", data=Aptr)
- B = T.decl_tensor((100,), "bfloat16", data=Bptr)
- D = T.decl_tensor((100,), "bfloat16", data=Dptr)
- C = T.decl_tensor((100,), "bfloat16")
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "bfloat16", data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), "bfloat16", data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), "bfloat16", data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), "bfloat16")])
for i in T.grid(100):
C[i] = A[i] + B[i]
D[i] = T.exp(C[i])
@@ -241,10 +241,10 @@ def test_bf16_storage_compute_scope_will_legalize():
Dptr: T.handle("bfloat16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "bfloat16", data=Aptr)
- B = T.decl_tensor((100,), "bfloat16", data=Bptr)
- D = T.decl_tensor((100,), "bfloat16", data=Dptr)
- C = T.decl_tensor((100,), "float32")
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "bfloat16", data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), "bfloat16", data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), "bfloat16", data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), "float32")])
for i in T.grid(100):
C[i] = bf16tof32(A[i]) + bf16tof32(B[i])
D[i] = f32tobf16(T.exp(C[i]))
@@ -261,10 +261,10 @@ def test_bf16_storage_compute_scope_will_legalize():
Dptr: T.handle("uint16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "uint16", data=Aptr)
- B = T.decl_tensor((100,), "uint16", data=Bptr)
- D = T.decl_tensor((100,), "uint16", data=Dptr)
- C = T.decl_tensor((100,), "float32")
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "uint16", data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), "uint16", data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), "uint16", data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), "float32")])
for i in T.grid(100):
C[i] = u16tof32(A[i]) + u16tof32(B[i])
D[i] = f32tou16(T.exp(C[i]))
@@ -290,10 +290,10 @@ def test_bf16_storage_compute_scope_wont_legalize():
Dptr: T.handle("bfloat16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "bfloat16", data=Aptr)
- B = T.decl_tensor((100,), "bfloat16", data=Bptr)
- D = T.decl_tensor((100,), "bfloat16", data=Dptr)
- C = T.decl_tensor((100,), "bfloat16")
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "bfloat16", data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), "bfloat16", data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), "bfloat16", data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), "bfloat16")])
for i in T.grid(100):
C[i] = A[i] + B[i]
D[i] = T.exp(C[i])
@@ -310,10 +310,10 @@ def test_bf16_storage_compute_scope_wont_legalize():
Dptr: T.handle("bfloat16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "bfloat16", data=Aptr)
- B = T.decl_tensor((100,), "bfloat16", data=Bptr)
- D = T.decl_tensor((100,), "bfloat16", data=Dptr)
- C = T.decl_tensor((100,), "bfloat16")
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "bfloat16", data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), "bfloat16", data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), "bfloat16", data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), "bfloat16")])
for i in T.grid(100):
C[i] = A[i] + B[i]
D[i] = T.exp(C[i])
@@ -330,10 +330,10 @@ def test_bf16_storage_compute_scope_wont_legalize():
Dptr: T.handle("bfloat16"),
):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "bfloat16", data=Aptr)
- B = T.decl_tensor((100,), "bfloat16", data=Bptr)
- D = T.decl_tensor((100,), "bfloat16", data=Dptr)
- C = T.decl_tensor((100,), "bfloat16")
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "bfloat16", data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), "bfloat16", data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), "bfloat16", data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), "bfloat16")])
for i in T.grid(100):
C[i] = A[i] + B[i]
D[i] = T.exp(C[i])
@@ -356,12 +356,12 @@ def test_bf16_reduce_will_legalize():
def main(
Aptr: T.handle("bfloat16", storage_scope="shared"),
):
- A_flat = T.decl_tensor(4096, "bfloat16", data=Aptr)
+ A_flat = T.decl_tensor(Aptr, ty_args=[T.Tensor(4096, "bfloat16", data=Aptr)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce = T.decl_tensor(1, dtype="bfloat16", scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor(1, "bfloat16", scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.bfloat16, T.bfloat16], lambda x, y: (x + y,)),
@@ -381,12 +381,12 @@ def test_bf16_reduce_will_legalize():
def main(
Aptr: T.handle("bfloat16", storage_scope="shared"),
):
- A_flat_1 = T.decl_tensor(4096, "bfloat16", data=Aptr)
+ A_flat_1 = T.decl_tensor(Aptr, ty_args=[T.Tensor(4096, "bfloat16", data=Aptr)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce = T.decl_tensor(1, dtype="float32", scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor(1, "float32", scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x, y: (x + y,)),
@@ -414,12 +414,12 @@ def test_bf16_reduce_will_legalize():
def main(
Aptr: T.handle("uint16", storage_scope="shared"),
):
- A_flat = T.decl_tensor(4096, "uint16", data=Aptr)
+ A_flat = T.decl_tensor(Aptr, ty_args=[T.Tensor(4096, "uint16", data=Aptr)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce = T.decl_tensor(1, dtype="float32", scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor(1, "float32", scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.float32, T.float32], lambda x, y: (x + y,)),
@@ -456,12 +456,12 @@ def test_bf16_reduce_wont_legalize():
def main(
Aptr: T.handle("bfloat16", storage_scope="shared"),
):
- A_flat = T.decl_tensor(4096, "bfloat16", data=Aptr)
+ A_flat = T.decl_tensor(Aptr, ty_args=[T.Tensor(4096, "bfloat16", data=Aptr)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce = T.decl_tensor(1, dtype="bfloat16", scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor(1, "bfloat16", scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.bfloat16, T.bfloat16], lambda x, y: (x + y,)),
@@ -481,12 +481,12 @@ def test_bf16_reduce_wont_legalize():
def main(
Aptr: T.handle("bfloat16", storage_scope="shared"),
):
- A_flat = T.decl_tensor(4096, "bfloat16", data=Aptr)
+ A_flat = T.decl_tensor(Aptr, ty_args=[T.Tensor(4096, "bfloat16", data=Aptr)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce = T.decl_tensor(1, dtype="bfloat16", scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor(1, "bfloat16", scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.bfloat16, T.bfloat16], lambda x, y: (x + y,)),
@@ -506,12 +506,12 @@ def test_bf16_reduce_wont_legalize():
def main(
Aptr: T.handle("bfloat16", storage_scope="shared"),
):
- A_flat = T.decl_tensor(4096, "bfloat16", data=Aptr)
+ A_flat = T.decl_tensor(Aptr, ty_args=[T.Tensor(4096, "bfloat16", data=Aptr)])
for i in range(128):
threadIdx_x = T.launch_thread("threadIdx.x", 32)
- reduce = T.decl_tensor(1, dtype="bfloat16", scope="local")
+ reduce = T.alloc_tensor(ty_args=[T.Tensor(1, "bfloat16", scope="local")])
T.gpu_thread_allreduce(
T.Lambda([T.bfloat16, T.bfloat16], lambda x, y: (x + y,)),
diff --git a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py b/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
index 1cb669e261..b07fb89220 100644
--- a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
+++ b/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
@@ -171,7 +171,7 @@ def test_reused_tensor_obj():
@T.function(private=True)
def func(a: T.handle("float32")):
- A = T.decl_tensor(shape=1, dtype="float32", data=a)
+ A = T.decl_tensor(a, ty_args=[T.Tensor(1, "float32", data=a)])
T.evaluate(A[0])
before = tvm.IRModule(
@@ -185,12 +185,12 @@ def test_reused_tensor_obj():
class expected:
@T.function
def func_a(a: T.handle("float32")):
- A = T.decl_tensor(shape=1, dtype="float32", data=a)
+ A = T.decl_tensor(a, ty_args=[T.Tensor(1, "float32", data=a)])
T.evaluate(A[0])
@T.function
def func_b(a: T.handle("float32")):
- A = T.decl_tensor(shape=1, dtype="float32", data=a)
+ A = T.decl_tensor(a, ty_args=[T.Tensor(1, "float32", data=a)])
T.evaluate(A[0])
after = tvm.tirx.transform.ConvertSSA()(before)
@@ -474,13 +474,9 @@ def test_shared_shape_var_in_tensor_params_and_alloc_tensor():
C,
tvm.ir.Call(
"tirx.alloc_tensor",
- [
- tvm.ir.Tuple(C.shape),
- tvm.ir.DataTypeImm(tvm.DataType(C.dtype)),
- tvm.ir.StringImm(C.scope()),
- ],
+ [],
attrs=tvm.ir.DictAttrs({}),
- ty=C.ty,
+ ty_args=[C.ty],
),
),
tvm.ir.Evaluate(1),
@@ -517,13 +513,8 @@ def test_reused_loop_var_in_decl_tensor_elem_offset():
tensor,
tvm.ir.Call(
"tirx.decl_tensor",
- [
- buffer_data,
- tvm.ir.Tuple(tensor.shape),
- tvm.ir.DataTypeImm(tvm.DataType(tensor.dtype)),
- tvm.ir.StringImm(tensor.scope()),
- ],
- ty=tensor.ty,
+ [buffer_data],
+ ty_args=[tensor.ty],
),
),
tvm.ir.Evaluate(tirx.TensorLoad(tensor, [0])),
diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py b/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
index 3e7783d11a..8d75dce878 100644
--- a/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
+++ b/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
@@ -39,7 +39,7 @@ def test_elementwise():
@T.function
def main(A: T.Tensor((16, 16), "float32"), C: T.Tensor((16, 16), "float32")):
for i in T.serial(0, 16):
- B_new = T.decl_tensor([1, 16], "float32")
+ B_new = T.alloc_tensor(ty_args=[T.Tensor([1, 16], "float32")])
for j in T.serial(0, 16):
B_new[0, j] = A[i, j] + 1.0
for j in T.serial(0, 16):
@@ -49,10 +49,14 @@ def test_elementwise():
class Expected:
@T.function
def main(A: T.Tensor((16, 16), "float32"), C: T.Tensor((16, 16), "float32")):
- A_1 = T.decl_tensor(256, dtype="float32", data=A.data, layout=None)
- C_1 = T.decl_tensor(256, dtype="float32", data=C.data, layout=None)
+ A_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor(256, "float32", layout=None, data=A.data)]
+ )
+ C_1 = T.decl_tensor(
+ C.data, ty_args=[T.Tensor(256, "float32", layout=None, data=C.data)]
+ )
for i in T.serial(0, 16):
- B_new = T.decl_tensor([16], "float32", layout=None)
+ B_new = T.alloc_tensor(ty_args=[T.Tensor([16], "float32", layout=None)])
for j in T.serial(0, 16):
B_new[j] = A_1[((i * 16) + j)] + 1.0
for j in T.serial(0, 16):
@@ -77,8 +81,10 @@ def test_elementwise_without_decl_tensor():
@T.function
def main(A: T.Tensor((16, 16), "float32"), C: T.Tensor((16, 16), "float32")):
for i in T.serial(0, 16):
- B_new_buf = T.alloc_tensor((1, 16), "float32")
- B_new = T.decl_tensor([1, 16], "float32", data=B_new_buf.data)
+ B_new_buf = T.alloc_tensor(ty_args=[T.Tensor((1, 16), "float32")])
+ B_new = T.decl_tensor(
+ B_new_buf.data, ty_args=[T.Tensor([1, 16], "float32", data=B_new_buf.data)]
+ )
for j in T.serial(0, 16):
B_new[0, j] = A[i, j] + 1.0
for j in T.serial(0, 16):
@@ -88,11 +94,18 @@ def test_elementwise_without_decl_tensor():
class Expected:
@T.function
def main(input_A: T.Tensor((16, 16), "float32"), input_C: T.Tensor((16, 16), "float32")):
- A = T.decl_tensor(256, dtype="float32", data=input_A.data, layout=None)
- C = T.decl_tensor(256, dtype="float32", data=input_C.data, layout=None)
+ A = T.decl_tensor(
+ input_A.data, ty_args=[T.Tensor(256, "float32", layout=None, data=input_A.data)]
+ )
+ C = T.decl_tensor(
+ input_C.data, ty_args=[T.Tensor(256, "float32", layout=None, data=input_C.data)]
+ )
for i in T.serial(0, 16):
- B_new_buf = T.alloc_tensor((16,), "float32", layout=None)
- B_new = T.decl_tensor(16, "float32", data=B_new_buf.data, layout=None)
+ B_new_buf = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32", layout=None)])
+ B_new = T.decl_tensor(
+ B_new_buf.data,
+ ty_args=[T.Tensor(16, "float32", layout=None, data=B_new_buf.data)],
+ )
for j in T.serial(0, 16):
B_new[j] = A[((i * 16) + j)] + 1.0
for j in T.serial(0, 16):
@@ -112,7 +125,7 @@ def test_gpu():
i0 = T.launch_thread("blockIdx.x", 4)
i1 = T.launch_thread("threadIdx.x", 2)
i2 = T.launch_thread("vthread", 2)
- B = T.decl_tensor([1, 16], "float32", scope="local")
+ B = T.alloc_tensor(ty_args=[T.Tensor([1, 16], "float32", scope="local")])
for j in range(0, 16):
B[0, j] = A[i0 * 4 + i1 * 2 + i2, j] + 1.0
for j in range(0, 16):
@@ -122,13 +135,17 @@ def test_gpu():
class Expected:
@T.function
def main(A: T.Tensor((16, 16), "float32"), C: T.Tensor((16, 16), "float32")):
- A_1 = T.decl_tensor(256, dtype="float32", data=A.data, layout=None)
- C_1 = T.decl_tensor(256, dtype="float32", data=C.data, layout=None)
+ A_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor(256, "float32", layout=None, data=A.data)]
+ )
+ C_1 = T.decl_tensor(
+ C.data, ty_args=[T.Tensor(256, "float32", layout=None, data=C.data)]
+ )
i0 = T.launch_thread("blockIdx.x", 4)
i1 = T.launch_thread("threadIdx.x", 2)
i2 = T.launch_thread("vthread", 2)
- B = T.decl_tensor([16], "float32", scope="local", layout=None)
+ B = T.alloc_tensor(ty_args=[T.Tensor([16], "float32", scope="local", layout=None)])
for j in range(0, 16):
B[j] = A_1[i0 * 64 + i1 * 32 + i2 * 16 + j] + 1.0
for j in range(0, 16):
@@ -151,7 +168,7 @@ def test_symbolic():
m: T.int32,
) -> None:
for i in range(0, n):
- B = T.decl_tensor([m], "float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor([m], "float32")])
for j in range(0, m):
B[j] = A[i, j] + 1.0
for j in range(0, m):
@@ -166,11 +183,15 @@ def test_symbolic():
n: T.int32,
m: T.int32,
) -> None:
- A_1 = T.decl_tensor(n * m, "float32", data=A.data, layout=None)
- C_1 = T.decl_tensor(n * m, "float32", data=C.data, layout=None)
+ A_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor(n * m, "float32", layout=None, data=A.data)]
+ )
+ C_1 = T.decl_tensor(
+ C.data, ty_args=[T.Tensor(n * m, "float32", layout=None, data=C.data)]
+ )
for i in range(0, n):
- B = T.decl_tensor([m], "float32", layout=None)
+ B = T.alloc_tensor(ty_args=[T.Tensor([m], "float32", layout=None)])
for j in range(0, m):
B[j] = A_1[i * m + j] + 1.0
for j in range(0, m):
@@ -204,8 +225,14 @@ def test_fused_symbolic():
input_B: T.Tensor((32, n, n), "float32"), # noqa: F821
n: T.int32,
) -> None:
- A = T.decl_tensor(n * n * 32, "float32", data=input_A.data, layout=None)
- B = T.decl_tensor(n * n * 32, "float32", data=input_B.data, layout=None)
+ A = T.decl_tensor(
+ input_A.data,
+ ty_args=[T.Tensor(n * n * 32, "float32", layout=None, data=input_A.data)],
+ )
+ B = T.decl_tensor(
+ input_B.data,
+ ty_args=[T.Tensor(n * n * 32, "float32", layout=None, data=input_B.data)],
+ )
for i in range(0, n * n * 32):
B[i] = A[i]
@@ -245,8 +272,14 @@ def test_fused_symbolic_with_predicate():
input_B: T.Tensor((32, n, n), "float32"), # noqa: F821
n: T.int32,
) -> None:
- A = T.decl_tensor(n * n * 32, "float32", data=input_A.data, layout=None)
- B = T.decl_tensor(n * n * 32, "float32", data=input_B.data, layout=None)
+ A = T.decl_tensor(
+ input_A.data,
+ ty_args=[T.Tensor(n * n * 32, "float32", layout=None, data=input_A.data)],
+ )
+ B = T.decl_tensor(
+ input_B.data,
+ ty_args=[T.Tensor(n * n * 32, "float32", layout=None, data=input_B.data)],
+ )
for bx, tx in T.grid((n * n + 1) // 2, 64):
if bx * 64 + tx < n * n * 32:
@@ -264,8 +297,8 @@ def test_multi_alloc():
@T.function
def main(A: T.Tensor((4, 32), "float32"), D: T.Tensor((4, 32), "float32")):
for i, j in T.grid(4, 32):
- B = T.decl_tensor((4, 32), "float32", scope="global")
- C = T.decl_tensor((4, 32), "float32", scope="global")
+ B = T.alloc_tensor(ty_args=[T.Tensor((4, 32), "float32", scope="global")])
+ C = T.alloc_tensor(ty_args=[T.Tensor((4, 32), "float32", scope="global")])
B[i, j] = A[i, j] + 1.0
C[i, j] = A[i, j] + B[i, j]
D[i, j] = C[i, j] * 2.0
@@ -274,12 +307,16 @@ def test_multi_alloc():
class Expected:
@T.function
def main(A: T.Tensor((4, 32), "float32"), D: T.Tensor((4, 32), "float32")):
- A_1 = T.decl_tensor(128, "float32", data=A.data, layout=None)
- D_1 = T.decl_tensor(128, "float32", data=D.data, layout=None)
+ A_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor(128, "float32", layout=None, data=A.data)]
+ )
+ D_1 = T.decl_tensor(
+ D.data, ty_args=[T.Tensor(128, "float32", layout=None, data=D.data)]
+ )
for i, j in T.grid(4, 32):
- B = T.decl_tensor([128], "float32", layout=None)
- C = T.decl_tensor([128], "float32", layout=None)
+ B = T.alloc_tensor(ty_args=[T.Tensor([128], "float32", layout=None)])
+ C = T.alloc_tensor(ty_args=[T.Tensor([128], "float32", layout=None)])
B[i * 32 + j] = A_1[i * 32 + j] + 1.0
C[i * 32 + j] = A_1[i * 32 + j] + B[i * 32 + j]
D_1[i * 32 + j] = C[i * 32 + j] * 2.0
@@ -296,8 +333,10 @@ def test_strided():
@T.function
def main(A: T.Tensor((16, 16), "float32"), C: T.Tensor((16, 16), "float32")):
for i0 in T.serial(4):
- B = T.decl_tensor([4, 17], "float32")
- B_1 = T.decl_tensor([4, 16], dtype="float32", data=B.data, strides=[17, 1])
+ B = T.alloc_tensor(ty_args=[T.Tensor([4, 17], "float32")])
+ B_1 = T.decl_tensor(
+ B.data, ty_args=[T.Tensor([4, 16], "float32", strides=[17, 1], data=B.data)]
+ )
for i1, j in T.grid(4, 16):
B_1[i1, j] = A[i0 * 4 + i1, j] + 1.0
for i1, j in T.grid(4, 16):
@@ -307,11 +346,17 @@ def test_strided():
class Expected:
@T.function
def main(A: T.Tensor((16, 16), "float32"), C: T.Tensor((16, 16), "float32")):
- A_1 = T.decl_tensor(256, dtype="float32", data=A.data, layout=None)
- C_1 = T.decl_tensor(256, dtype="float32", data=C.data, layout=None)
+ A_1 = T.decl_tensor(
+ A.data, ty_args=[T.Tensor(256, "float32", layout=None, data=A.data)]
+ )
+ C_1 = T.decl_tensor(
+ C.data, ty_args=[T.Tensor(256, "float32", layout=None, data=C.data)]
+ )
for i0 in T.serial(0, 4):
- B = T.decl_tensor([68], "float32", layout=None)
- B_1 = T.decl_tensor([68], "float32", data=B.data, layout=None)
+ B = T.alloc_tensor(ty_args=[T.Tensor([68], "float32", layout=None)])
+ B_1 = T.decl_tensor(
+ B.data, ty_args=[T.Tensor([68], "float32", layout=None, data=B.data)]
+ )
for i1 in T.serial(0, 4):
for j in T.serial(0, 16):
B_1[i1 * 17 + j] = A_1[i0 * 64 + i1 * 16 + j] + 1.0
@@ -337,8 +382,12 @@ def test_boolean():
class Expected:
@T.function
def main(input_A: T.Tensor(10, "bool"), input_B: T.Tensor(10, "bool")) -> None:
- A = T.decl_tensor(10, dtype="bool", data=input_A.data, layout=None)
- B = T.decl_tensor(10, dtype="bool", data=input_B.data, layout=None)
+ A = T.decl_tensor(
+ input_A.data, ty_args=[T.Tensor(10, "bool", layout=None, data=input_A.data)]
+ )
+ B = T.decl_tensor(
+ input_B.data, ty_args=[T.Tensor(10, "bool", layout=None, data=input_B.data)]
+ )
# body
for i0 in T.serial(10):
B[i0] = A[i0]
@@ -354,7 +403,7 @@ def test_flatten_inside_block():
class Before:
@T.function
def main():
- A = T.alloc_tensor([32, 32])
+ A = T.alloc_tensor(ty_args=[T.Tensor([32, 32])])
for i, j in T.grid(32, 32):
T.evaluate(A[i, j])
@@ -362,7 +411,7 @@ def test_flatten_inside_block():
class Expected:
@T.function
def main():
- A = T.alloc_tensor([1024], layout=None)
+ A = T.alloc_tensor(ty_args=[T.Tensor([1024], layout=None)])
for i, j in T.grid(32, 32):
T.evaluate(A[i * 32 + j])
diff --git a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
index a3377cdbe4..9ea8144605 100644
--- a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
+++ b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
@@ -279,7 +279,7 @@ def test_fail_on_tensor_param():
def test_fail_on_internal_tensor():
@T.function(private=True)
def func(A: T.Tensor((128,), "int32"), B: T.Tensor((128,), "int32")):
- C = T.alloc_tensor((128,), "int64")
+ C = T.alloc_tensor(ty_args=[T.Tensor((128,), "int64")])
for i in T.serial(0, 16):
for j in T.serial(0, 8):
C[i * 8 + j] = T.cast(A[i * 8 + j], "int64") + T.int64(1)
diff --git a/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py b/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py
index a7e3db4485..4d8897b641 100644
--- a/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py
+++ b/tests/python/tirx-transform/test_tir_transform_fp8_legalize.py
@@ -30,10 +30,10 @@ def get_before(dtype: str):
@T.function
def main(Aptr: T.handle(dtype), Bptr: T.handle(dtype), Dptr: T.handle(dtype)):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), dtype, data=Aptr)
- B = T.decl_tensor((100,), dtype, data=Bptr)
- D = T.decl_tensor((100,), dtype, data=Dptr)
- C = T.decl_tensor((100,), dtype)
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), dtype, data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), dtype, data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), dtype, data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), dtype)])
for i in T.grid(100):
C[i] = A[i] + B[i]
D[i] = T.exp(C[i])
@@ -55,10 +55,10 @@ def get_after_compute_legalize(dtype: str, promote_dtype: str):
@T.function
def main(Aptr: T.handle(dtype), Bptr: T.handle(dtype), Dptr: T.handle(dtype)):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), dtype, data=Aptr)
- B = T.decl_tensor((100,), dtype, data=Bptr)
- D = T.decl_tensor((100,), dtype, data=Dptr)
- C = T.decl_tensor((100,), promote_dtype)
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), dtype, data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), dtype, data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), dtype, data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), promote_dtype)])
for i in T.grid(100):
C[i] = promote_f8(dtype, promote_dtype, A[i]) + promote_f8(
dtype, promote_dtype, B[i]
@@ -188,10 +188,10 @@ def get_after_storage_legalize(dtype: str, promote_dtype: str):
@T.function
def main(Aptr: T.handle("uint8"), Bptr: T.handle("uint8"), Dptr: T.handle("uint8")):
T.func_attr({"global_symbol": "main"})
- A = T.decl_tensor((100,), "uint8", data=Aptr)
- B = T.decl_tensor((100,), "uint8", data=Bptr)
- D = T.decl_tensor((100,), "uint8", data=Dptr)
- C = T.decl_tensor((100,), promote_dtype)
+ A = T.decl_tensor(Aptr, ty_args=[T.Tensor((100,), "uint8", data=Aptr)])
+ B = T.decl_tensor(Bptr, ty_args=[T.Tensor((100,), "uint8", data=Bptr)])
+ D = T.decl_tensor(Dptr, ty_args=[T.Tensor((100,), "uint8", data=Dptr)])
+ C = T.alloc_tensor(ty_args=[T.Tensor((100,), promote_dtype)])
for i in T.grid(100):
C[i] = promote_uint8(dtype, promote_dtype, A[i]) + promote_uint8(
dtype, promote_dtype, B[i]
@@ -219,7 +219,7 @@ def test_fp8_compute_legalize(dtype, promote_dtype):
def test_fp8_compute_legalize_preserves_opaque_tensor_access(dtype, promote_dtype):
@T.function
def before():
- tensor = T.alloc_tensor((16,), dtype)
+ tensor = T.alloc_tensor(ty_args=[T.Tensor((16,), dtype)])
T.evaluate(T.call_extern("consume", tensor.data, ty="void"))
before_mod = tvm.IRModule.from_expr(before)
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
index 8d5fed120a..691c88eb92 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
@@ -55,7 +55,10 @@ def test_lower_call_packed():
stack_ffi_any: T.let[T.handle] = T.stack_alloca("tvm_ffi_any", 4)
stack_array: T.let[T.handle] = T.stack_alloca("array", 3)
stack_shape = T.decl_tensor(
- (T.int64(6),), "int64", data=T.stack_alloca("shape", 6), layout=None
+ T.stack_alloca("shape", 6),
+ ty_args=[
+ T.Tensor((T.int64(6),), "int64", layout=None, data=T.stack_alloca("shape", 6))
+ ],
)
stack_shape[0] = T.int64(64)
stack_shape[1] = T.int64(64)
@@ -189,11 +192,17 @@ def test_lower_overflow_int32():
@T.function(check_well_formed=False)
def variance4(rxplaceholder: T.Tensor((T.int64(1), T.int64(32), T.int64(25690112)), "float32")):
T.func_attr({"global_symbol": "variance4", "tirx.noalias": True})
- rxplaceholder_red = T.alloc_tensor((32,), "float32")
- T_subtract = T.alloc_tensor((822083584,), "float32")
- rxplaceholder_red_1 = T.decl_tensor((T.int64(32),), data=rxplaceholder_red.data)
- rxplaceholder_1 = T.decl_tensor((T.int64(822083584),), data=rxplaceholder.data)
- T_subtract_1 = T.decl_tensor((T.int64(822083584),), data=T_subtract.data)
+ rxplaceholder_red = T.alloc_tensor(ty_args=[T.Tensor((32,), "float32")])
+ T_subtract = T.alloc_tensor(ty_args=[T.Tensor((822083584,), "float32")])
+ rxplaceholder_red_1 = T.decl_tensor(
+ rxplaceholder_red.data, ty_args=[T.Tensor((T.int64(32),), data=rxplaceholder_red.data)]
+ )
+ rxplaceholder_1 = T.decl_tensor(
+ rxplaceholder.data, ty_args=[T.Tensor((T.int64(822083584),), data=rxplaceholder.data)]
+ )
+ T_subtract_1 = T.decl_tensor(
+ T_subtract.data, ty_args=[T.Tensor((T.int64(822083584),), data=T_subtract.data)]
+ )
for ax1, ax2 in T.grid(32, 25690112):
cse_v1: T.let[T.int32] = ax1 * 25690112 + ax2
T_subtract_1[cse_v1] = rxplaceholder_1[cse_v1] - rxplaceholder_red_1[ax1]
@@ -219,8 +228,8 @@ def test_lower_device_allocate():
def main():
T.func_attr({"target": T.target("llvm")})
with T.device_context(2, 0): # kDLCuda
- ptr = T.alloc_tensor((16,), "float32")
- buf = T.decl_tensor(16, "float32", data=ptr.data)
+ ptr = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
+ buf = T.decl_tensor(ptr.data, ty_args=[T.Tensor(16, "float32", data=ptr.data)])
buf[0] = 0.0
After = tvm.tirx.transform.LowerTVMBuiltin()(Before)
@@ -242,8 +251,8 @@ def test_lower_cpu_allocation():
def main():
T.func_attr({"target": T.target("llvm")})
with T.device_context(1, 0): # kDLCPU
- ptr = T.alloc_tensor((16,), "float32")
- buf = T.decl_tensor(16, "float32", data=ptr.data)
+ ptr = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
+ buf = T.decl_tensor(ptr.data, ty_args=[T.Tensor(16, "float32", data=ptr.data)])
buf[0] = 0.0
@I.ir_module
@@ -251,8 +260,8 @@ def test_lower_cpu_allocation():
@T.function
def main():
T.func_attr({"target": T.target("llvm")})
- ptr = T.alloc_tensor((16,), "float32")
- buf = T.decl_tensor(16, "float32", data=ptr.data)
+ ptr = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
+ buf = T.decl_tensor(ptr.data, ty_args=[T.Tensor(16, "float32", data=ptr.data)])
buf[0] = 0.0
After = tvm.tirx.transform.LowerTVMBuiltin()(Before)
@@ -267,8 +276,8 @@ def test_lower_allocate_requires_device_id():
@T.function
def main():
T.func_attr({"target": T.target("cuda"), "tirx.is_host_func": True})
- ptr = T.alloc_tensor((16,), "float32")
- buf = T.decl_tensor(16, "float32", data=ptr.data)
+ ptr = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
+ buf = T.decl_tensor(ptr.data, ty_args=[T.Tensor(16, "float32", data=ptr.data)])
buf[0] = 0.0
with pytest.raises(RuntimeError):
@@ -290,8 +299,8 @@ def test_lower_allocate_requires_device_type():
@T.function
def main():
T.func_attr({"tirx.is_host_func": True})
- ptr = T.alloc_tensor((1024 * 1024,), "float32")
- buf = T.decl_tensor(1024 * 1024, "float32", data=ptr.data)
+ ptr = T.alloc_tensor(ty_args=[T.Tensor((1024 * 1024,), "float32")])
+ buf = T.decl_tensor(ptr.data, ty_args=[T.Tensor(1024 * 1024, "float32", data=ptr.data)])
buf[0] = 0.0
with pytest.raises(RuntimeError):
@@ -312,8 +321,8 @@ def test_lower_cpu_alloc_with_function_attr():
@T.function
def main():
T.func_attr({"target": T.target("llvm")})
- ptr = T.alloc_tensor((16,), "float32")
- buf = T.decl_tensor(16, "float32", data=ptr.data)
+ ptr = T.alloc_tensor(ty_args=[T.Tensor((16,), "float32")])
+ buf = T.decl_tensor(ptr.data, ty_args=[T.Tensor(16, "float32", data=ptr.data)])
buf[0] = 0.0
# Expected is same as before for this transform
diff --git a/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py b/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py
index 0bbe0d5049..e3d58a0701 100644
--- a/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py
+++ b/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py
@@ -63,8 +63,8 @@ def test_basic():
@T.function
def func(A: T.handle("float32"), B: T.handle("float32"), m: T.int32, n: T.int32):
- A_buf = T.decl_tensor((m * n,), "float32", data=A)
- B_buf = T.decl_tensor((m * n,), "float32", data=B)
+ A_buf = T.decl_tensor(A, ty_args=[T.Tensor((m * n,), "float32", data=A)])
+ B_buf = T.decl_tensor(B, ty_args=[T.Tensor((m * n,), "float32", data=B)])
for i in T.serial(m):
for j in T.serial(n):
B_buf[i * n + j] = A_buf[i * n + j] + T.float32(1)
@@ -73,8 +73,8 @@ def test_basic():
@T.function
def func(A: T.handle("float32"), B: T.handle("float32"), m: T.int64, n: T.int64):
- A_buf = T.decl_tensor((m * n,), "float32", data=A)
- B_buf = T.decl_tensor((m * n,), "float32", data=B)
+ A_buf = T.decl_tensor(A, ty_args=[T.Tensor((m * n,), "float32", data=A)])
+ B_buf = T.decl_tensor(B, ty_args=[T.Tensor((m * n,), "float32", data=B)])
for i in T.serial(m):
for j in T.serial(n):
B_buf[i * n + j] = A_buf[i * n + j] + T.float32(1)
diff --git a/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py b/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py
index f44815ba1f..a174ff1173 100644
--- a/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py
+++ b/tests/python/tirx-transform/test_tir_transform_pointer_value_type_rewrite.py
@@ -40,7 +40,7 @@ def test_rewrite_to_shuffle_0():
class Before:
@T.function
def main(A: T.Tensor((16,), "float32"), B: T.Tensor((4,), "float32")):
- A_local = T.alloc_tensor((16,), scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor((16,), scope="local")])
for i in range(4):
A_local[T.ramp(i * 4, 1, 4)] = A[T.ramp(i * 4, 1, 4)]
for i in range(4):
@@ -50,7 +50,9 @@ def test_rewrite_to_shuffle_0():
class Expected:
@T.function
def main(A: T.Tensor((4,), "float32x4", layout=None), B: T.Tensor((4,), "float32")):
- A_local = T.alloc_tensor((4,), "float32x4", scope="local", layout=None)
+ A_local = T.alloc_tensor(
+ ty_args=[T.Tensor((4,), "float32x4", scope="local", layout=None)]
+ )
for i in range(4):
A_local[T.Div(i * 4, 4)] = A[T.Div(i * 4, 4)]
for i in range(4):
@@ -72,7 +74,7 @@ def test_rewrite_to_shuffle_1():
class Before:
@T.function
def main(A: T.Tensor((8,), "float32"), B: T.Tensor((1,), "float32")):
- A_local = T.alloc_tensor((8,), scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor((8,), scope="local")])
A_local[T.ramp(0, 1, 4)] = A[T.ramp(0, 1, 4)]
A_local[T.ramp(4, 1, 4)] = A[T.ramp(4, 1, 4)]
B[0] = (
@@ -90,7 +92,9 @@ def test_rewrite_to_shuffle_1():
class Expected:
@T.function
def main(A: T.Tensor((2,), "float32x4", layout=None), B: T.Tensor((1,), "float32")):
- A_local = T.alloc_tensor((2,), "float32x4", scope="local", layout=None)
+ A_local = T.alloc_tensor(
+ ty_args=[T.Tensor((2,), "float32x4", scope="local", layout=None)]
+ )
A_local[0] = A[0]
A_local[1] = A[1]
B[0] = (
@@ -160,8 +164,10 @@ def test_decl_tensor_alias_chain_uses_flat_root_map():
class Before:
@T.function
def main(A: T.Tensor((16,), "float32")):
- A_view = T.decl_tensor((16,), "float32", data=A.data)
- A_view_2 = T.decl_tensor((16,), "float32", data=A_view.data)
+ A_view = T.decl_tensor(A.data, ty_args=[T.Tensor((16,), "float32", data=A.data)])
+ A_view_2 = T.decl_tensor(
+ A_view.data, ty_args=[T.Tensor((16,), "float32", data=A_view.data)]
+ )
for i in range(4):
A_view_2[T.ramp(i * 4, 1, 4)] = T.broadcast(T.float32(1), 4)
diff --git a/tests/python/tirx-transform/test_tir_transform_remove_no_op.py b/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
index 4174e38514..4a356d9674 100644
--- a/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
+++ b/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
@@ -379,7 +379,7 @@ def test_remove_empty_temporary():
@T.function(private=True)
def before():
- A = T.alloc_tensor((16,), "int32", scope="local")
+ A = T.alloc_tensor(ty_args=[T.Tensor((16,), "int32", scope="local")])
T.evaluate(0)
@T.function(private=True)
@@ -401,8 +401,8 @@ def test_remove_empty_temporary_with_decl_tensor():
@T.function(private=True)
def before():
- A = T.decl_tensor([4, 4], "int32", scope="local")
- A_flat = T.decl_tensor(16, "int32", scope="local", data=A.data)
+ A = T.alloc_tensor(ty_args=[T.Tensor([4, 4], "int32", scope="local")])
+ A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "int32", scope="local")])
T.evaluate(0)
@T.function(private=True)
@@ -420,7 +420,7 @@ def test_remove_unused_temporary():
@T.function(private=True)
def before(A: T.Tensor(16, "int32")):
- B = T.alloc_tensor((16,), "int32", scope="local")
+ B = T.alloc_tensor(ty_args=[T.Tensor((16,), "int32", scope="local")])
for i in T.serial(16):
A[i] = 1
@@ -440,7 +440,7 @@ def test_remove_unused_write_into_temporary():
@T.function(private=True)
def before():
- A = T.decl_tensor([16], "int32", scope="local")
+ A = T.alloc_tensor(ty_args=[T.Tensor([16], "int32", scope="local")])
for i in T.serial(16):
A[i] = 0
@@ -458,7 +458,7 @@ def test_keep_used_write_into_temporary():
@T.function(private=True)
def before(B: T.Tensor(16, "int32")):
- A = T.decl_tensor([16], "int32", scope="local")
+ A = T.alloc_tensor(ty_args=[T.Tensor([16], "int32", scope="local")])
for i in T.serial(16):
A[i] = 0
@@ -476,7 +476,7 @@ def test_remove_write_into_temporary():
@T.function(private=True)
def before(A: T.Tensor(16, "int32"), C: T.Tensor(1, "int32")):
- B = T.decl_tensor([16], "int32", scope="local")
+ B = T.alloc_tensor(ty_args=[T.Tensor([16], "int32", scope="local")])
for i in T.serial(16):
B[i] = A[i]
@@ -489,7 +489,7 @@ def test_remove_write_into_temporary():
@T.function(private=True)
def expected(A: T.Tensor(16, "int32"), C: T.Tensor(1, "int32")):
- B = T.decl_tensor([16], "int32", scope="local")
+ B = T.alloc_tensor(ty_args=[T.Tensor([16], "int32", scope="local")])
for i in T.serial(16):
B[i] = A[i]
diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py b/tests/python/tirx-transform/test_tir_transform_simplify.py
index 586d39b7f1..cf0377033b 100644
--- a/tests/python/tirx-transform/test_tir_transform_simplify.py
+++ b/tests/python/tirx-transform/test_tir_transform_simplify.py
@@ -26,8 +26,8 @@ from tvm.script import tirx as T
def test_stmt_simplify():
@T.function(private=True)
def func(A: T.handle("float32"), C: T.handle("float32"), n: T.int32):
- A_ptr = T.decl_tensor((10,), "float32", data=A)
- C_ptr = T.decl_tensor((10,), "float32", data=C)
+ A_ptr = T.decl_tensor(A, ty_args=[T.Tensor((10,), "float32", data=A)])
+ C_ptr = T.decl_tensor(C, ty_args=[T.Tensor((10,), "float32", data=C)])
n_val: T.let[T.int32] = 10
for i in T.serial(n_val):
if i < 12:
@@ -48,8 +48,8 @@ def test_stmt_simplify():
def test_thread_extent_simplify():
@T.function(private=True)
def func(A: T.handle("float32"), C: T.handle("float32"), n: T.int32):
- A_ptr = T.decl_tensor((10,), "float32", data=A)
- C_ptr = T.decl_tensor((10,), "float32", data=C)
+ A_ptr = T.decl_tensor(A, ty_args=[T.Tensor((10,), "float32", data=A)])
+ C_ptr = T.decl_tensor(C, ty_args=[T.Tensor((10,), "float32", data=C)])
n_val: T.let[T.int32] = 10
for tx in T.thread_binding(n_val, thread="threadIdx.x"):
for ty in T.thread_binding(1, thread="threadIdx.y"):
@@ -73,8 +73,8 @@ def test_thread_extent_simplify():
def test_if_likely():
@T.function(private=True)
def func(A: T.handle("float32"), C: T.handle("float32"), n: T.int32):
- A_ptr = T.decl_tensor((32,), "float32", data=A)
- C_ptr = T.decl_tensor((1024,), "float32", data=C)
+ A_ptr = T.decl_tensor(A, ty_args=[T.Tensor((32,), "float32", data=A)])
+ C_ptr = T.decl_tensor(C, ty_args=[T.Tensor((1024,), "float32", data=C)])
for tx in T.thread_binding(32, thread="threadIdx.x"):
for ty in T.thread_binding(32, thread="threadIdx.y"):
if T.likely(tx * 32 + ty < n):
@@ -1161,7 +1161,7 @@ def test_simplify_trivial_let_tensor_var():
@T.function(private=True)
def before(A_ptr: T.handle("float32")):
A_ptr_redef: T.let[T.handle("float32")] = A_ptr
- A = T.decl_tensor(1, "float32", data=A_ptr_redef)
+ A = T.decl_tensor(A_ptr_redef, ty_args=[T.Tensor(1, "float32", data=A_ptr_redef)])
A[0] = 42.0
expected = before
@@ -1176,13 +1176,17 @@ def test_simplify_trivial_let_elem_offset():
@T.function(private=True)
def before(A_ptr: T.handle("float32"), A_offset: T.int32):
A_offset_redef = A_offset
- A = T.decl_tensor(1, "float32", elem_offset=A_offset_redef, data=A_ptr)
+ A = T.decl_tensor(
+ A_ptr, ty_args=[T.Tensor(1, "float32", elem_offset=A_offset_redef, data=A_ptr)]
+ )
A[0] = 42.0
@T.function(private=True)
def expected(A_ptr: T.handle("float32"), A_offset: T.int32):
A_offset_redef = A_offset
- A = T.decl_tensor(1, "float32", elem_offset=A_offset_redef, data=A_ptr)
+ A = T.decl_tensor(
+ A_ptr, ty_args=[T.Tensor(1, "float32", elem_offset=A_offset_redef, data=A_ptr)]
+ )
A[0] = 42.0
after = _apply_simplify(before)
@@ -1195,13 +1199,13 @@ def test_simplify_trivial_let_shape():
@T.function(private=True)
def before(A_ptr: T.handle("float32"), A_size: T.int32):
A_size_redef = A_size
- A = T.decl_tensor([A_size_redef], "float32", data=A_ptr)
+ A = T.decl_tensor(A_ptr, ty_args=[T.Tensor([A_size_redef], "float32", data=A_ptr)])
A[0] = 42.0
@T.function(private=True)
def expected(A_ptr: T.handle("float32"), A_size: T.int32):
A_size_redef = A_size
- A = T.decl_tensor([A_size_redef], "float32", data=A_ptr)
+ A = T.decl_tensor(A_ptr, ty_args=[T.Tensor([A_size_redef], "float32", data=A_ptr)])
A[0] = 42.0
after = _apply_simplify(before)
@@ -1214,13 +1218,17 @@ def test_simplify_trivial_let_stride():
@T.function(private=True)
def before(A_ptr: T.handle("float32"), A_stride: T.int32):
A_stride_redef = A_stride
- A = T.decl_tensor(1, "float32", strides=[A_stride_redef], data=A_ptr)
+ A = T.decl_tensor(
+ A_ptr, ty_args=[T.Tensor(1, "float32", strides=[A_stride_redef], data=A_ptr)]
+ )
A[0] = 42.0
@T.function(private=True)
def expected(A_ptr: T.handle("float32"), A_stride: T.int32):
A_stride_redef = A_stride
- A = T.decl_tensor(1, "float32", strides=[A_stride_redef], data=A_ptr)
+ A = T.decl_tensor(
+ A_ptr, ty_args=[T.Tensor(1, "float32", strides=[A_stride_redef], data=A_ptr)]
+ )
A[0] = 42.0
after = _apply_simplify(before)
@@ -1240,8 +1248,8 @@ def test_simplify_tensor_identity_well_formed():
@T.function(private=True)
def before(A_ptr: T.handle("float32"), B_ptr: T.handle("float32"), n: T.int32):
n_val = n
- A = T.decl_tensor([n_val], "float32", data=A_ptr)
- B = T.decl_tensor([n_val], "float32", data=B_ptr)
+ A = T.decl_tensor(A_ptr, ty_args=[T.Tensor([n_val], "float32", data=A_ptr)])
+ B = T.decl_tensor(B_ptr, ty_args=[T.Tensor([n_val], "float32", data=B_ptr)])
B[0] = A[0]
after = _apply_simplify(before)
diff --git a/tests/python/tirx-transform/test_tir_transform_split_host_device.py b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
index 05af7357d1..58528e6828 100644
--- a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
+++ b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
@@ -329,8 +329,8 @@ def test_dynamic_launch_thread():
"tirx.noalias": True,
}
)
- A = T.decl_tensor(seq_len, "int32", data=A_data)
- B = T.decl_tensor(seq_len, "int32", data=B_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(seq_len, "int32", data=A_data)])
+ B = T.decl_tensor(B_data, ty_args=[T.Tensor(seq_len, "int32", data=B_data)])
blockIdx_x = T.launch_thread("blockIdx.x", num_blocks)
threadIdx_x = T.launch_thread("threadIdx.x", 128)
if blockIdx_x * 128 + threadIdx_x < seq_len:
@@ -353,8 +353,8 @@ def test_symbolic_var_parameter():
T.region("tirx.device_scope", [], attrs={"target": T.target("cuda")})
blockIdx_x = T.launch_thread("blockIdx.x", m)
- B_1 = T.decl_tensor((m,), data=B.data)
- A_1 = T.decl_tensor((m,), data=A.data)
+ B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((m,), data=B.data)])
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((m,), data=A.data)])
B_1[blockIdx_x] = A_1[blockIdx_x]
after = tvm.tirx.transform.SplitHostDevice()(Module)
@@ -414,7 +414,7 @@ def test_thread_extent_region_extracted_as_device_kernel():
"tirx.is_global_func": True,
}
)
- A = T.decl_tensor(16, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
i = T.launch_thread("threadIdx.x", 16)
A[i] = 0.0
@@ -570,7 +570,7 @@ def test_device_scope_region_extracted_as_device_kernel():
"tirx.is_global_func": True,
}
)
- A = T.decl_tensor(1, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(1, "float32", data=A_data)])
A[0] = 0.0
After = tvm.tirx.transform.SplitHostDevice()(Before)
@@ -590,7 +590,7 @@ def test_lower_device_kernel_launch():
@T.function
def kernel(A_data: T.handle("float32")):
T.func_attr({"target": T.target("cuda")})
- A = T.decl_tensor(1, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(1, "float32", data=A_data)])
A[0] = 0.0
@I.ir_module
@@ -611,7 +611,7 @@ def test_lower_device_kernel_launch():
"tirx.is_global_func": True,
}
)
- A = T.decl_tensor(1, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(1, "float32", data=A_data)])
A[0] = 0.0
After = tvm.tirx.transform.SplitHostDevice()(Before)
@@ -631,7 +631,7 @@ def test_externally_visible_kernel_launch():
@T.function
def kernel(A_data: T.handle("float32")):
T.func_attr({"target": T.target("cuda"), "global_symbol": "kernel_by_another_name"})
- A = T.decl_tensor(1, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(1, "float32", data=A_data)])
A[0] = 0.0
@I.ir_module
@@ -652,7 +652,7 @@ def test_externally_visible_kernel_launch():
"tirx.is_global_func": True,
}
)
- A = T.decl_tensor(1, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(1, "float32", data=A_data)])
A[0] = 0.0
After = tvm.tirx.transform.SplitHostDevice()(Before)
@@ -677,7 +677,7 @@ def test_collect_launch_parameter():
"global_symbol": "kernel",
}
)
- A = T.decl_tensor(16, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
i = T.launch_thread("threadIdx.x", 16)
A[i] = 0.0
@@ -699,7 +699,7 @@ def test_collect_launch_parameter():
"tirx.is_global_func": True,
}
)
- A = T.decl_tensor(16, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
i = T.launch_thread("threadIdx.x", 16)
A[i] = 0.0
@@ -720,7 +720,7 @@ def test_same_device_different_target():
@T.function
def kernel(A_data: T.handle("float32")):
T.func_attr({"target": T.target("c")})
- A = T.decl_tensor(16, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
A[0] = 0.0
@I.ir_module
@@ -739,7 +739,7 @@ def test_same_device_different_target():
"tirx.is_global_func": True,
}
)
- A = T.decl_tensor(16, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
A[0] = 0.0
After = tvm.tirx.transform.SplitHostDevice()(Before)
@@ -759,7 +759,7 @@ def test_bind_before_thread_extent():
@T.function
def kernel(A_data: T.handle("float32"), n: T.int32):
T.func_attr({"target": T.target("cuda"), "global_symbol": "kernel"})
- A = T.decl_tensor(16, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
v: T.let[T.int32] = n + 1
i = T.launch_thread("threadIdx.x", v)
A[i] = 0.0
@@ -782,7 +782,7 @@ def test_bind_before_thread_extent():
"tirx.is_global_func": True,
}
)
- A = T.decl_tensor(16, dtype="float32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32", data=A_data)])
v: T.let[T.int32] = n + 1
i = T.launch_thread("threadIdx.x", v)
A[i] = 0.0
diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
index 88ed638604..df2a18f6de 100644
--- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
+++ b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
@@ -42,10 +42,10 @@ def test_alloc_seq():
def func(n: T.int32):
for i in T.serial(n):
for j in range(10):
- A = T.alloc_tensor((200,), scope=scope_tb)
+ A = T.alloc_tensor(ty_args=[T.Tensor((200,), scope=scope_tb)])
A[j] = T.float32(1.2)
for j in range(10):
- B = T.alloc_tensor((200,), scope=scope_tb)
+ B = T.alloc_tensor(ty_args=[T.Tensor((200,), scope=scope_tb)])
B[j] = T.float32(1.3)
mod = tvm.IRModule.from_expr(func)
@@ -70,11 +70,11 @@ def test_alloc_different_dtypes():
@T.function
def func():
# Allocate all tensors in parent scope (before any loops)
- A = T.alloc_tensor((length,), dtype_list[0], scope="local.L0A")
- B = T.alloc_tensor((length,), dtype_list[1], scope="local.L0A")
- C = T.alloc_tensor((length,), dtype_list[2], scope="local.L0A")
- D = T.alloc_tensor((length,), dtype_list[3], scope="local.L0A")
- E = T.alloc_tensor((length,), "int8", scope="local.L0A")
+ A = T.alloc_tensor(ty_args=[T.Tensor((length,), dtype_list[0], scope="local.L0A")])
+ B = T.alloc_tensor(ty_args=[T.Tensor((length,), dtype_list[1], scope="local.L0A")])
+ C = T.alloc_tensor(ty_args=[T.Tensor((length,), dtype_list[2], scope="local.L0A")])
+ D = T.alloc_tensor(ty_args=[T.Tensor((length,), dtype_list[3], scope="local.L0A")])
+ E = T.alloc_tensor(ty_args=[T.Tensor((length,), "int8", scope="local.L0A")])
for j in range(length):
A[j] = T.Cast(dtype_list[0], 1)
@@ -137,21 +137,21 @@ def test_address_of():
# combine tensors B and D, but not C
@T.function
def before(A: T.Tensor(8, "float32"), E: T.Tensor(8, "float32")):
- B = T.alloc_tensor((8,))
+ B = T.alloc_tensor(ty_args=[T.Tensor((8,))])
for i in range(8):
B[i] = (
T.call_extern("deref", T.address_of(A[i]), ty="float32")
+ T.call_extern("deref", T.address_of(A[0]), ty="float32")
+ T.float32(1)
)
- C = T.alloc_tensor((8,))
+ C = T.alloc_tensor(ty_args=[T.Tensor((8,))])
for i in range(8):
C[i] = (
T.call_extern("deref", T.address_of(B[i]), ty="float32")
+ T.call_extern("deref", T.address_of(B[0]), ty="float32")
+ T.float32(2)
)
- D = T.alloc_tensor((8,))
+ D = T.alloc_tensor(ty_args=[T.Tensor((8,))])
for i in range(8):
D[i] = (
T.call_extern("deref", T.address_of(C[i]), ty="float32")
@@ -185,7 +185,7 @@ def test_parallel_alloc():
def func1(n: T.int32):
for i in T.parallel(n):
for j in range(10):
- A = T.alloc_tensor((n,))
+ A = T.alloc_tensor(ty_args=[T.Tensor((n,))])
A[j] = A[j] + T.float32(2)
mod = tvm.IRModule.from_expr(func1)
@@ -200,7 +200,7 @@ def test_parallel_alloc():
with T.parallel_launch():
for i in T.parallel(n):
for j in range(10):
- A = T.alloc_tensor((n,))
+ A = T.alloc_tensor(ty_args=[T.Tensor((n,))])
A[j] = A[j] + T.float32(2)
mod = tvm.IRModule.from_expr(func2)
@@ -214,20 +214,20 @@ def test_while_alloc():
@T.function
def func_parallel(n: T.int32):
for i in T.parallel(n):
- j = T.alloc_tensor((1,), "int32")
+ j = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32")])
j[0] = 0
while j[0] < 10:
- A = T.alloc_tensor((n,))
+ A = T.alloc_tensor(ty_args=[T.Tensor((n,))])
A[j[0]] = A[j[0]] + T.float32(2)
j[0] = j[0] + j[0] + 1
@T.function
def func_serial(n: T.int32):
for i in T.serial(n):
- j = T.alloc_tensor((1,), "int32")
+ j = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32")])
j[0] = 0
while j[0] < 10:
- A = T.alloc_tensor((n,))
+ A = T.alloc_tensor(ty_args=[T.Tensor((n,))])
A[j[0]] = A[j[0]] + T.float32(2)
j[0] = j[0] + j[0] + 1
@@ -270,17 +270,17 @@ def test_alloc_seq_type():
def func(n: T.int32):
for i in T.serial(n):
for j in range(10):
- A = T.alloc_tensor((200,), scope="local.L0A")
- A1 = T.alloc_tensor((200,), scope="local.L0A")
+ A = T.alloc_tensor(ty_args=[T.Tensor((200,), scope="local.L0A")])
+ A1 = T.alloc_tensor(ty_args=[T.Tensor((200,), scope="local.L0A")])
A[j] = T.float32(1.2)
A1[j] = T.float32(1.3)
- B = T.alloc_tensor((200,), "int16", scope="local.L0A")
+ B = T.alloc_tensor(ty_args=[T.Tensor((200,), "int16", scope="local.L0A")])
B[j] = T.int16(1)
- C = T.alloc_tensor((200,), "int16", scope="local.L0A")
+ C = T.alloc_tensor(ty_args=[T.Tensor((200,), "int16", scope="local.L0A")])
C[j] = T.int16(1)
- D = T.alloc_tensor((200,), "int16", scope="local.L0A")
+ D = T.alloc_tensor(ty_args=[T.Tensor((200,), "int16", scope="local.L0A")])
D[j] = B[j] + C[j]
- A2 = T.alloc_tensor((200,), scope="local.L0A")
+ A2 = T.alloc_tensor(ty_args=[T.Tensor((200,), scope="local.L0A")])
A2[j] = A[j]
mod = tvm.IRModule.from_expr(func)
@@ -304,13 +304,13 @@ def test_alloc_seq_type2():
def func(n: T.int32):
for i in T.serial(n):
for j in range(10):
- A = T.alloc_tensor((200,), scope=scope_tb)
+ A = T.alloc_tensor(ty_args=[T.Tensor((200,), scope=scope_tb)])
A[j] = T.float32(1.2)
for j in range(20):
- B = T.alloc_tensor((400,), "int16", scope=scope_tb)
+ B = T.alloc_tensor(ty_args=[T.Tensor((400,), "int16", scope=scope_tb)])
B[j] = T.int16(1)
for j in range(10):
- C = T.alloc_tensor((200,), scope=scope_tb)
+ C = T.alloc_tensor(ty_args=[T.Tensor((200,), scope=scope_tb)])
C[j] = T.float32(1.2)
mod = tvm.IRModule.from_expr(func)
@@ -332,17 +332,17 @@ def test_reuse_small_tensor():
def func(n: T.int32):
for i in T.serial(n):
for j in range(10):
- A = T.alloc_tensor((200,), "int16", scope="local.L0A")
+ A = T.alloc_tensor(ty_args=[T.Tensor((200,), "int16", scope="local.L0A")])
A[j] = T.int16(1)
- B = T.alloc_tensor((200,), "int16", scope="local.L0A")
+ B = T.alloc_tensor(ty_args=[T.Tensor((200,), "int16", scope="local.L0A")])
B[j] = T.int16(1)
- B1 = T.alloc_tensor((200,), "int16", scope="local.L0A")
+ B1 = T.alloc_tensor(ty_args=[T.Tensor((200,), "int16", scope="local.L0A")])
B1[j] = A[j] + B[j]
- C = T.alloc_tensor((400,), "int16", scope="local.L0A")
+ C = T.alloc_tensor(ty_args=[T.Tensor((400,), "int16", scope="local.L0A")])
C[j] = T.int16(1)
- D = T.alloc_tensor((400,), "int16", scope="local.L0A")
+ D = T.alloc_tensor(ty_args=[T.Tensor((400,), "int16", scope="local.L0A")])
D[j] = T.int16(1)
- E = T.alloc_tensor((400,), "int16", scope="local.L0A")
+ E = T.alloc_tensor(ty_args=[T.Tensor((400,), "int16", scope="local.L0A")])
E[j] = C[j]
mod = tvm.IRModule.from_expr(func)
@@ -363,14 +363,14 @@ def test_access_in_let_value():
@T.function
def func(A: T.Tensor((8,), "float32")):
for i in range(8):
- B = T.alloc_tensor((1,))
+ B = T.alloc_tensor(ty_args=[T.Tensor((1,))])
B[0] = 3.14
x: T.let[T.float32] = T.exp(B[0])
A[i] = (x + 1.0) / (x - 1.0)
@T.function
def func_rewritten(A: T.Tensor((8,), "float32")) -> None:
- B = T.alloc_tensor((1,))
+ B = T.alloc_tensor(ty_args=[T.Tensor((1,))])
for i in range(8):
B[0] = 3.14
x: T.let[T.float32] = T.exp(B[0])
@@ -394,7 +394,7 @@ def test_decl_tensor_is_not_vectorized():
@T.function
def main() -> None:
A_data: T.let[T.handle("int32")] = T.call_extern("dummy_func", ty=T.handle("int32").ty)
- A = T.decl_tensor([8], "int32", data=A_data)
+ A = T.decl_tensor(A_data, ty_args=[T.Tensor([8], "int32", data=A_data)])
A[T.ramp(0, 1, 8)] = T.broadcast(42, 8)
After = tvm.tirx.transform.StorageRewrite()(Before)
@@ -408,8 +408,8 @@ def test_rewrite_decl_tensor():
class Before:
@T.function
def main(A: T.Tensor(16, "float32"), D: T.Tensor(16, "float32")):
- B = T.decl_tensor(16, dtype="float32")
- C = T.decl_tensor(16, dtype="float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
+ C = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
for i in range(16):
B[i] = A[i]
@@ -424,8 +424,8 @@ def test_rewrite_decl_tensor():
class Expected:
@T.function
def main(A: T.Tensor(16, "float32"), D: T.Tensor(16, "float32")):
- B = T.decl_tensor(16, dtype="float32")
- C = T.decl_tensor(16, dtype="float32", data=B.data)
+ B = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
+ C = T.decl_tensor(B.data, ty_args=[T.Tensor(16, "float32", data=B.data)])
for i in range(16):
B[i] = A[i]
@@ -448,9 +448,9 @@ def test_decl_tensor_alias_chain_uses_flat_root():
class Before:
@T.function
def main(D: T.Tensor(1, "float32")):
- A = T.decl_tensor(16, dtype="float32")
- B = T.decl_tensor(16, dtype="float32", data=A.data)
- C = T.decl_tensor(16, dtype="float32", data=B.data)
+ A = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
+ B = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "float32", data=A.data)])
+ C = T.decl_tensor(B.data, ty_args=[T.Tensor(16, "float32", data=B.data)])
A[0] = 1.0
D[0] = C[0]
@@ -458,9 +458,9 @@ def test_decl_tensor_alias_chain_uses_flat_root():
class Expected:
@T.function
def main(D: T.Tensor(1, "float32")):
- A = T.decl_tensor(16, dtype="float32")
- B = T.decl_tensor(16, dtype="float32", data=A.data)
- C = T.decl_tensor(16, dtype="float32", data=A.data)
+ A = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
+ B = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "float32", data=A.data)])
+ C = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "float32", data=A.data)])
A[0] = 1.0
D[0] = C[0]
@@ -474,11 +474,11 @@ def test_decl_tensor_alias_extends_source_lifetime():
@T.function
def func(D: T.Tensor(1, "float32")):
- A = T.decl_tensor(16, dtype="float32")
- B = T.decl_tensor(16, dtype="float32", data=A.data)
+ A = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
+ B = T.decl_tensor(A.data, ty_args=[T.Tensor(16, "float32", data=A.data)])
A[0] = 1.0
- C = T.decl_tensor(16, dtype="float32")
+ C = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
C[0] = 2.0
D[0] = B[0] + C[0]
@@ -505,9 +505,9 @@ def test_no_orphaned_decl_tensor():
class Before:
@T.function
def main(A: T.Tensor(16, "float32"), D: T.Tensor(16, "float32")):
- B = T.decl_tensor(16, dtype="float32")
- C = T.decl_tensor(16, dtype="float32")
- Unused = T.decl_tensor(16, dtype="float32")
+ B = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
+ C = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
+ Unused = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
for i in range(16):
B[i] = A[i]
@@ -522,8 +522,8 @@ def test_no_orphaned_decl_tensor():
class Expected:
@T.function
def main(A: T.Tensor(16, "float32"), D: T.Tensor(16, "float32")):
- B = T.decl_tensor(16, dtype="float32")
- C = T.decl_tensor(16, dtype="float32", data=B.data)
+ B = T.alloc_tensor(ty_args=[T.Tensor(16, "float32")])
+ C = T.decl_tensor(B.data, ty_args=[T.Tensor(16, "float32", data=B.data)])
for i in range(16):
B[i] = A[i]
diff --git a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
index e1ef7d81ec..29c3a764ec 100644
--- a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
+++ b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
@@ -96,16 +96,16 @@ def test_unroll_allocations():
@T.function
def main():
for i in T.unroll(2):
- buf = T.alloc_tensor([16], "float32")
+ buf = T.alloc_tensor(ty_args=[T.Tensor([16], "float32")])
buf[0] = 0.0
@I.ir_module
class Expected:
@T.function
def main():
- buf1 = T.alloc_tensor([16], "float32")
+ buf1 = T.alloc_tensor(ty_args=[T.Tensor([16], "float32")])
buf1[0] = 0.0
- buf2 = T.alloc_tensor([16], "float32")
+ buf2 = T.alloc_tensor(ty_args=[T.Tensor([16], "float32")])
buf2[0] = 0.0
after = tvm.tirx.transform.UnrollLoop()(Before)
@@ -120,7 +120,7 @@ def test_unroll_local_access():
def main(B: T.Tensor((64,), "float32")):
for bx in T.thread_binding(4, thread="blockIdx.x"):
for tx in T.thread_binding(4, thread="threadIdx.x"):
- A_local = T.alloc_tensor((4,), scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor((4,), scope="local")])
for i in T.serial(4):
A_local[i] = T.float32(i)
@@ -130,7 +130,7 @@ def test_unroll_local_access():
def main(B: T.Tensor((64,), "float32")):
for bx in T.thread_binding(4, thread="blockIdx.x"):
for tx in T.thread_binding(4, thread="threadIdx.x"):
- A_local = T.alloc_tensor((4,), scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor((4,), scope="local")])
A_local[0] = T.float32(0)
A_local[1] = T.float32(1)
A_local[2] = T.float32(2)
diff --git a/tests/python/tirx-transform/test_tir_transform_vectorize.py b/tests/python/tirx-transform/test_tir_transform_vectorize.py
index 2d880fc673..38a68580a8 100644
--- a/tests/python/tirx-transform/test_tir_transform_vectorize.py
+++ b/tests/python/tirx-transform/test_tir_transform_vectorize.py
@@ -353,7 +353,7 @@ def test_vectorize_while_fail():
C[j] = T.float32(0)
# While loop inside vectorized loop (should fail)
- i = T.decl_tensor((1,), "int32", scope="local")
+ i = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
i[0] = 0
for j in T.vectorized(64):
while i[0] < 10:
diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py b/tests/python/tirx/codegen/test_codegen_blackwell.py
index 4fc633b6fa..8ffbf0fdf7 100644
--- a/tests/python/tirx/codegen/test_codegen_blackwell.py
+++ b/tests/python/tirx/codegen/test_codegen_blackwell.py
@@ -101,7 +101,7 @@ def test_tmem_alloc_dealloc_relinquish():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- # tmem_addr = T.alloc_tensor((1,), "uint32", scope="shared", align=8)
+ # tmem_addr = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="shared", align=8)])
tmem_addr = T.shared_scalar("uint32")
# alloc TMEM
@@ -341,8 +341,8 @@ def test_tcgen05_ld_st_roundtrip():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tx = T.cuda.thread_idx('x')
- reg = T.alloc_tensor((WIDTH,), "float32", scope="local")
- # tmem_addr = T.alloc_tensor((1,), "uint32", scope="shared", align=8)
+ reg = T.alloc_tensor(ty_args=[T.Tensor((WIDTH,), "float32", scope="local")])
+ # tmem_addr = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="shared", align=8)])
tmem_addr = T.shared_scalar("uint32")
# alloc TMEM
@@ -421,13 +421,15 @@ def test_tcgen05_cp_ld_roundtrip():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tx = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor((HEIGHT, WIDTH), dtype, scope="shared", layout=A_layout)
- reg = T.alloc_tensor((WIDTH,), dtype, scope="local")
- # tmem_addr = T.alloc_tensor((1,), "uint32", scope="shared", align=8)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor((HEIGHT, WIDTH), dtype, scope="shared", layout=A_layout)]
+ )
+ reg = T.alloc_tensor(ty_args=[T.Tensor((WIDTH,), dtype, scope="local")])
+ # tmem_addr = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="shared", align=8)])
tmem_addr = T.shared_scalar("uint32")
- descA = T.alloc_tensor((1,), "uint64", scope="local")
- bar = T.alloc_tensor((1,), "uint64", scope="shared", align=8)
- phase = T.alloc_tensor((1,), "int32", scope="local")
+ descA = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="local")])
+ bar = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="shared", align=8)])
+ phase = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
# alloc TMEM
if warp_id == 0:
@@ -565,17 +567,27 @@ def test_tcgen05_mma_ss_no_tma(swizzle):
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tx = T.cuda.thread_idx('x')
- dyn = T.alloc_tensor((dyn_smem_bytes,), "uint8", scope="shared")
+ dyn = T.alloc_tensor(ty_args=[T.Tensor((dyn_smem_bytes,), "uint8", scope="shared")])
tmem_addr = T.decl_scalar("uint32", dyn.data, scope="shared", elem_offset=0)
- A_smem = T.decl_tensor((M, K), a_type, dyn.data, elem_offset=256, layout=A_layout)
- B_smem = T.decl_tensor((N, K), b_type, dyn.data, elem_offset=256 + M*K, layout=B_layout)
- bar = T.decl_tensor((1,), "uint64", dyn.data, scope="shared", elem_offset=8)
+ A_smem = T.decl_tensor(
+ dyn.data,
+ ty_args=[T.Tensor((M, K), a_type, elem_offset=256, layout=A_layout, data=dyn.data)],
+ )
+ B_smem = T.decl_tensor(
+ dyn.data,
+ ty_args=[
+ T.Tensor((N, K), b_type, elem_offset=256 + M * K, layout=B_layout, data=dyn.data)
+ ],
+ )
+ bar = T.decl_tensor(
+ dyn.data, ty_args=[T.Tensor((1,), "uint64", scope="shared", elem_offset=8)]
+ )
- reg = T.alloc_tensor((N,), d_type, scope="local")
- descA = T.alloc_tensor((1,), "uint64", scope="local")
- descB = T.alloc_tensor((1,), "uint64", scope="local")
- descI = T.alloc_tensor((1,), "uint32", scope="local")
- phase = T.alloc_tensor((1,), "int32", scope="local")
+ reg = T.alloc_tensor(ty_args=[T.Tensor((N,), d_type, scope="local")])
+ descA = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="local")])
+ descB = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="local")])
+ descI = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="local")])
+ phase = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
# alloc TMEM
if warp_id == 0:
@@ -719,11 +731,11 @@ def test_tcgen05_mma_pred_codegen():
@T.function
def test_mma_pred():
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(1,)))
- tmem_addr = T.alloc_tensor((1,), "uint32", scope="local")
- desc_a = T.alloc_tensor((1,), "uint64", scope="local")
- desc_b = T.alloc_tensor((1,), "uint64", scope="local")
- desc_i = T.alloc_tensor((1,), "uint32", scope="local")
- pred = T.alloc_tensor((1,), "uint32", scope="local")
+ tmem_addr = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="local")])
+ desc_a = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="local")])
+ desc_b = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="local")])
+ desc_i = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="local")])
+ pred = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="local")])
tmem_addr[0] = T.uint32(0)
desc_a[0] = T.uint64(0)
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py
index 31e72be560..fbdd0e1b13 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -156,13 +156,13 @@ def test_typed_buffer_bindings(monkeypatch):
@T.function
def main(data: T.handle("float32"), output: T.handle):
T.func_attr({"global_symbol": "main", "calling_conv": 2, "tirx.kernel_launch_params": []})
- A = T.decl_tensor((1,), "float32", data=data)
- B = T.decl_tensor((1,), "float32", data=output)
- storage = T.alloc_tensor((1,), "uint32", scope="shared")
- S = T.decl_tensor((1,), "float32", data=storage.data, scope="shared")
- pool = T.alloc_tensor((4,), "uint8", scope="shared.dyn")
- D = T.decl_tensor((1,), "float32", data=pool.data, scope="shared.dyn")
- alias = T.decl_tensor((1,), "float32", data=D.data, scope="shared.dyn")
+ A = T.decl_tensor(data, ty_args=[T.Tensor((1,), "float32")])
+ B = T.decl_tensor(output, ty_args=[T.Tensor((1,), "float32")])
+ storage = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="shared")])
+ S = T.decl_tensor(storage.data, ty_args=[T.Tensor((1,), "float32", scope="shared")])
+ pool = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint8", scope="shared.dyn")])
+ D = T.decl_tensor(pool.data, ty_args=[T.Tensor((1,), "float32", scope="shared.dyn")])
+ alias = T.decl_tensor(D.data, ty_args=[T.Tensor((1,), "float32", scope="shared.dyn")])
S[0] = A[0]
D[0] = S[0]
B[0] = alias[0]
@@ -195,13 +195,8 @@ def test_vector_pointer_preserves_packed_offset(monkeypatch):
tensor,
tvm.ir.Call(
"tirx.decl_tensor",
- [
- data,
- tvm.ir.Tuple(tensor.shape),
- tvm.ir.DataTypeImm(tvm.DataType(tensor.dtype)),
- tvm.ir.StringImm(tensor.scope()),
- ],
- ty=tensor.ty,
+ [data],
+ ty_args=[tensor.ty],
),
),
tvm.ir.Evaluate(tvm.tirx.call_extern("void", "consume", pointer)),
@@ -682,8 +677,8 @@ def test_ptx_sync_and_clc_codegen():
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(32,)))
tx = T.cuda.thread_idx("x")
if tx == 0:
- bar = T.alloc_tensor((5,), "uint64", scope="shared", align=16)
- response = T.alloc_tensor((4,), "uint32", scope="shared", align=16)
+ bar = T.alloc_tensor(ty_args=[T.Tensor((5,), "uint64", scope="shared", align=16)])
+ response = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint32", scope="shared", align=16)])
T.ptx.cp.async_.mbarrier.arrive.shared.b64(bar.ptr_to([0]))
T.ptx.cp.async_.mbarrier.arrive.noinc.shared__cta.b64(bar.ptr_to([0]))
T.cuda.mbarrier_wait(bar.ptr_to([0]), T.int32(0))
@@ -754,7 +749,7 @@ def test_ptx_mbarrier_arrive_new_forms_codegen():
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(32,)))
tx = T.cuda.thread_idx("x")
if tx == 0:
- bar = T.alloc_tensor((6,), "uint64", scope="shared", align=16)
+ bar = T.alloc_tensor(ty_args=[T.Tensor((6,), "uint64", scope="shared", align=16)])
state = T.local_scalar("uint64")
T.ptx.mbarrier.arrive.relaxed.cta.shared__cta.b64(bar.ptr_to([0]))
T.ptx.mbarrier.arrive.relaxed.cluster.shared__cluster.b64(bar.ptr_to([1]))
@@ -820,7 +815,7 @@ def test_tma_cache_policy_operand_codegen():
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(32,)))
tx = T.cuda.thread_idx("x")
if tx == 0:
- smem = T.alloc_tensor((128,), "float32", scope="shared", align=128)
+ smem = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared", align=128)])
bar = T.shared_scalar("uint64")
T.ptx[_TMA_G2S_CG2_CACHE](
smem.data, T.address_of(A_map), 0, 0, T.address_of(bar), Cache[0]
@@ -973,8 +968,8 @@ def test_warp_shuffle_xor_sync():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
- A_local = T.alloc_tensor([1], "float32", scope="local")
- i = T.alloc_tensor([1], "int32", scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor([1], "float32", scope="local")])
+ i = T.alloc_tensor(ty_args=[T.Tensor([1], "int32", scope="local")])
A_local[0] = T.float32(31 - lane_id)
i[0] = 16
@@ -1172,7 +1167,7 @@ def test_uint32_loop_var_runs_correctly():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
_ = T.cuda.block_idx("x")
tx = T.cast(T.cuda.thread_idx("x"), "uint32")
- acc = T.alloc_tensor((1,), "int32", scope="local")
+ acc = T.alloc_tensor(ty_args=[T.Tensor((1,), "int32", scope="local")])
acc[0] = 0
for k in T.serial(4, dtype="uint32"):
acc[0] = acc[0] + A[tx] + T.int32(k)
diff --git a/tests/python/tirx/codegen/test_codegen_dsmem.py b/tests/python/tirx/codegen/test_codegen_dsmem.py
index 98d79fc38d..0983523a47 100644
--- a/tests/python/tirx/codegen/test_codegen_dsmem.py
+++ b/tests/python/tirx/codegen/test_codegen_dsmem.py
@@ -118,7 +118,7 @@ def test_mapa_pointer_bind_codegen():
mapped = T.alloc_local([1], "uint64")
T.ptx.mapa.u64(mapped[0], mbar.ptr_to([0]), T.uint32(0))
remote_ptr = T.reinterpret( mapped[0], ty=ptr_ty)
- remote_mbar = T.decl_tensor([1], "uint64", data=remote_ptr, scope="shared")
+ remote_mbar = T.decl_tensor(remote_ptr, ty_args=[T.Tensor([1], "uint64", scope="shared")])
A[0] = remote_mbar[0]
# fmt: on
diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py b/tests/python/tirx/codegen/test_codegen_hopper.py
index b58ff5a70f..edc842c5aa 100644
--- a/tests/python/tirx/codegen/test_codegen_hopper.py
+++ b/tests/python/tirx/codegen/test_codegen_hopper.py
@@ -89,8 +89,8 @@ def test_stmatrix_sync_aligned(trans):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
cta_id = T.cuda.block_idx('x')
tx = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor((16, 16), "float16", scope="shared", align=16)
- reg = T.alloc_tensor((8,), "float16", scope="local")
+ A_smem = T.alloc_tensor(ty_args=[T.Tensor((16, 16), "float16", scope="shared", align=16)])
+ reg = T.alloc_tensor(ty_args=[T.Tensor((8,), "float16", scope="local")])
for i in range(8):
reg[i] = tx * 8 + i
# stmatrix stores 4 b32 registers; reg is fp16, so they ride a uint32
@@ -454,7 +454,9 @@ def test_cp_async_bulk_tensor_global_to_shared_unicast(dtype, inputs):
for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
bar = T.shared_scalar("uint64")
phase: T.int32
- A_smem = T.alloc_tensor(shape, dtype, scope="shared", align=128)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="shared", align=128)]
+ )
phase = 0
if threadIdx == 0:
@@ -654,7 +656,9 @@ def test_cp_async_bulk_tensor_global_to_shared_swizzle(swizzle, dtype):
T.device_entry()
for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
- A_smem = T.alloc_tensor((total_elems,), dtype, scope="shared", align=128)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor((total_elems,), dtype, scope="shared", align=128)]
+ )
bar = T.shared_scalar("uint64")
phase: T.int32
@@ -759,7 +763,9 @@ def test_cp_async_bulk_tensor_global_to_shared_multicast1(inputs):
for tx in T.thread_binding(128, thread="threadIdx.x"):
bar = T.shared_scalar("uint64")
phase: T.int32
- A_smem = T.alloc_tensor(shape[::-1], "float32", scope="shared", align=128)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape[::-1], "float32", scope="shared", align=128)]
+ )
phase = 0
if tx == 0:
@@ -870,7 +876,9 @@ def test_cp_async_bulk_tensor_global_to_shared_multicast2(inputs):
for tx in T.thread_binding(128, thread="threadIdx.x"):
bar = T.shared_scalar("uint64")
phase: T.int32
- A_smem = T.alloc_tensor(shape[::-1], "float32", scope="shared", align=128)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape[::-1], "float32", scope="shared", align=128)]
+ )
phase = 0
if tx == 0:
@@ -987,7 +995,7 @@ def test_cp_async_bulk_tensor_shared_to_global(inputs):
cta_id = T.cuda.block_idx('x')
tx = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor(elems, "float32", scope="shared", align=128)
+ A_smem = T.alloc_tensor(ty_args=[T.Tensor(elems, "float32", scope="shared", align=128)])
if tx == 0:
for i in T.serial(0, elems):
@@ -1079,14 +1087,18 @@ def test_wgmma_ss_nt():
cta_id = T.cuda.block_idx('x')
tx = T.cuda.thread_idx('x') # A warpgroup is 128 threads
- A_smem = T.alloc_tensor(shapeA, in_dtype, scope="shared", align=1024)
- B_smem = T.alloc_tensor(shapeB, in_dtype, scope="shared", align=1024)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shapeA, in_dtype, scope="shared", align=1024)]
+ )
+ B_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shapeB, in_dtype, scope="shared", align=1024)]
+ )
bar = T.shared_scalar("uint64")
phase: T.int32
descA: T.uint64
descB: T.uint64
- C_local = T.alloc_tensor((C_elems,), out_dtype, scope="local")
+ C_local = T.alloc_tensor(ty_args=[T.Tensor((C_elems,), out_dtype, scope="local")])
# init phase and bar
phase = 0
@@ -1250,17 +1262,21 @@ def test_wgmma_rs_nt():
cta_id = T.cuda.block_idx('x')
tx = T.cuda.thread_idx('x') # A warpgroup is 128 threads
- B_smem = T.alloc_tensor(shapeB, in_dtype, scope="shared", align=1024)
- # bar = T.alloc_tensor((1,), "uint64", scope="shared", align=8)
+ B_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shapeB, in_dtype, scope="shared", align=1024)]
+ )
+ # bar = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="shared", align=8)])
bar = T.shared_scalar("uint64")
- # descB = T.alloc_tensor((1,), "uint64", scope="local")
+ # descB = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="local")])
descB: T.uint64
- A_local = T.alloc_tensor((A_elems,), in_dtype, scope="local")
- C_local = T.alloc_tensor((C_elems,), out_dtype, scope="local")
+ A_local = T.alloc_tensor(ty_args=[T.Tensor((A_elems,), in_dtype, scope="local")])
+ C_local = T.alloc_tensor(ty_args=[T.Tensor((C_elems,), out_dtype, scope="local")])
A_elems_b32 = T.meta_var(A_elems // (32 // in_dtype_bits))
- A_local_b32 = T.decl_tensor((A_elems_b32,), "uint32", data=A_local.data)
+ A_local_b32 = T.decl_tensor(
+ A_local.data, ty_args=[T.Tensor((A_elems_b32,), "uint32", data=A_local.data)]
+ )
# load A to regs
for i in T.serial(0, A_elems // 4):
@@ -1388,7 +1404,7 @@ def test_mapa():
cbx = T.cuda.cluster_cta_id("x")
cta_id = T.cuda.block_idx("x")
tx = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor([1], "uint32", scope="shared")
+ A_smem = T.alloc_tensor(ty_args=[T.Tensor([1], "uint32", scope="shared")])
mapped = T.alloc_local([1], "uint64")
if cbx == 0 and tx == 0:
T.ptx.mapa.u64(mapped[0], A_smem.data, T.uint32(cbx))
diff --git a/tests/python/tirx/codegen/test_codegen_nki.py b/tests/python/tirx/codegen/test_codegen_nki.py
index a6a2592e39..b854f69a11 100644
--- a/tests/python/tirx/codegen/test_codegen_nki.py
+++ b/tests/python/tirx/codegen/test_codegen_nki.py
@@ -42,8 +42,8 @@ def test_nki_add_1():
def func(A: T.Tensor((128, 512)), B: T.Tensor((128, 512))):
T.func_attr({"num_inputs": 1})
T.device_entry()
- A_sbuf = T.alloc_tensor((128, 512), "float32", scope="trn.sbuf",)
- B_sbuf = T.alloc_tensor((128, 512), "float32", scope="trn.sbuf",)
+ A_sbuf = T.alloc_tensor(ty_args=[T.Tensor((128, 512), "float32", scope="trn.sbuf")])
+ B_sbuf = T.alloc_tensor(ty_args=[T.Tensor((128, 512), "float32", scope="trn.sbuf")])
with T.nki.tensorized_instruction():
for i in range(0, 128):
for j in range(0, 512):
@@ -96,8 +96,8 @@ def test_nki_add_2():
def func(A: T.Tensor((128, 2048)), B: T.Tensor((128, 2048))):
T.func_attr({"num_inputs": 1})
T.device_entry()
- A_sbuf = T.alloc_tensor((128, 512), "float32", scope="trn.sbuf",)
- B_sbuf = T.alloc_tensor((128, 512), "float32", scope="trn.sbuf",)
+ A_sbuf = T.alloc_tensor(ty_args=[T.Tensor((128, 512), "float32", scope="trn.sbuf")])
+ B_sbuf = T.alloc_tensor(ty_args=[T.Tensor((128, 512), "float32", scope="trn.sbuf")])
for k in range(0, 4):
with T.nki.tensorized_instruction():
for i in range(0, 128):
@@ -176,16 +176,26 @@ def test_nki_matmul_1():
):
T.func_attr({"num_inputs": 2})
result_tiles = T.alloc_tensor(
- (TILE_M, NUM_BLOCK_M, TILES_IN_BLOCK_M, TILES_IN_BLOCK_N, TILE_N),
- "float32",
- scope="trn.sbuf",
+ ty_args=[
+ T.Tensor(
+ (TILE_M, NUM_BLOCK_M, TILES_IN_BLOCK_M, TILES_IN_BLOCK_N, TILE_N),
+ "float32",
+ scope="trn.sbuf",
+ )
+ ]
+ )
+ rhs_tiles = T.alloc_tensor(
+ ty_args=[T.Tensor((TILE_K, TILES_IN_BLOCK_K, BLOCK_N), "float16", scope="trn.sbuf")]
)
- rhs_tiles = T.alloc_tensor((TILE_K, TILES_IN_BLOCK_K, BLOCK_N), "float16", scope="trn.sbuf")
lhsT_tiles = T.alloc_tensor(
- (TILE_K, TILES_IN_BLOCK_K, BLOCK_M), "float16", scope="trn.sbuf"
+ ty_args=[T.Tensor((TILE_K, TILES_IN_BLOCK_K, BLOCK_M), "float16", scope="trn.sbuf")]
+ )
+ res_tile = T.alloc_tensor(
+ ty_args=[T.Tensor((1, TILE_M, TILE_N), "float32", scope="trn.psum")]
+ )
+ result_packed = T.alloc_tensor(
+ ty_args=[T.Tensor((TILE_K, BLOCK_N), "float32", scope="trn.sbuf")]
)
- res_tile = T.alloc_tensor((1, TILE_M, TILE_N), "float32", scope="trn.psum")
- result_packed = T.alloc_tensor((TILE_K, BLOCK_N), "float32", scope="trn.sbuf")
for n in range(NUM_BLOCK_N):
with T.nki.tensorized_instruction():
for i0 in range(TILE_M):
diff --git a/tests/python/tirx/codegen/test_cuda_cta_reduce.py b/tests/python/tirx/codegen/test_cuda_cta_reduce.py
index 3fdd2c67f7..450021c585 100644
--- a/tests/python/tirx/codegen/test_cuda_cta_reduce.py
+++ b/tests/python/tirx/codegen/test_cuda_cta_reduce.py
@@ -56,7 +56,7 @@ def test_cta_sum_4_warps():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- scratch = T.alloc_tensor((NUM_WARPS,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((NUM_WARPS,), "float32", scope="shared")])
val: T.f32 = T.float32(tid + 1)
val = T.cuda.cta_sum(val, NUM_WARPS, scratch.ptr_to([0]))
out[tid] = val
@@ -84,7 +84,7 @@ def test_cta_sum_8_warps():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- scratch = T.alloc_tensor((NUM_WARPS,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((NUM_WARPS,), "float32", scope="shared")])
val: T.f32 = T.float32(tid + 1)
val = T.cuda.cta_sum(val, NUM_WARPS, scratch.ptr_to([0]))
out[tid] = val
@@ -111,7 +111,7 @@ def test_cta_max_4_warps():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- scratch = T.alloc_tensor((NUM_WARPS,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((NUM_WARPS,), "float32", scope="shared")])
val: T.f32 = T.float32(tid + 1)
val = T.cuda.cta_max(val, NUM_WARPS, scratch.ptr_to([0]))
out[tid] = val
@@ -137,7 +137,7 @@ def test_cta_min_4_warps():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- scratch = T.alloc_tensor((NUM_WARPS,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((NUM_WARPS,), "float32", scope="shared")])
val: T.f32 = T.float32(tid + 1)
val = T.cuda.cta_min(val, NUM_WARPS, scratch.ptr_to([0]))
out[tid] = val
@@ -163,7 +163,7 @@ def test_cta_sum_1_warp():
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- scratch = T.alloc_tensor((NUM_WARPS,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((NUM_WARPS,), "float32", scope="shared")])
val: T.f32 = T.float32(tid + 1)
val = T.cuda.cta_sum(val, NUM_WARPS, scratch.ptr_to([0]))
out[tid] = val
@@ -190,7 +190,7 @@ def test_cta_sum_all_warp_counts(num_warps):
warp_id = T.cuda.warp_id()
lane_id = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- scratch = T.alloc_tensor((num_warps,), "float32", scope="shared")
+ scratch = T.alloc_tensor(ty_args=[T.Tensor((num_warps,), "float32", scope="shared")])
val: T.f32 = T.float32(tid + 1)
val = T.cuda.cta_sum(val, num_warps, scratch.ptr_to([0]))
out[tid] = val
diff --git a/tests/python/tirx/codegen/test_ptx_addr.py b/tests/python/tirx/codegen/test_ptx_addr.py
index a268bcdf58..a2b8c97bd7 100644
--- a/tests/python/tirx/codegen/test_ptx_addr.py
+++ b/tests/python/tirx/codegen/test_ptx_addr.py
@@ -89,7 +89,7 @@ def test_ptx_addr_coercion_ir_order_and_shared_codegen():
def kernel(global_buf: T.Tensor((8,), "uint64"), raw_shared: T.uint32, raw_global: T.uint64):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(32,)))
tx = T.cuda.thread_idx("x")
- shared_buf = T.alloc_tensor((8,), "uint64", scope="shared")
+ shared_buf = T.alloc_tensor(ty_args=[T.Tensor((8,), "uint64", scope="shared")])
value = T.local_scalar("uint64")
if tx == 0:
T.ptx.ld.shared.b64(value, T.ptx.addr(shared_buf.data, 4))
@@ -122,8 +122,8 @@ def test_ptx_addr_scalar_vector_cache_predicate_and_multi_address_codegen():
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(32,)))
tx = T.cuda.thread_idx("x")
- shared_buf = T.alloc_tensor((64,), "uint32", scope="shared")
- barrier = T.alloc_tensor((1,), "uint64", scope="shared")
+ shared_buf = T.alloc_tensor(ty_args=[T.Tensor((64,), "uint32", scope="shared")])
+ barrier = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="shared")])
values = T.alloc_local((2,), "uint32")
T.ptx.ld.global_.b32(values[0], T.ptx.addr(src.data, 16))
T.ptx.ld.global_.L2__cache_hint.b32(values[1], T.ptx.addr(src.data, -16), policy[0])
@@ -265,8 +265,8 @@ def test_ptx_addr_tma_tmem_and_independent_immediate_rejections():
@T.function
def tma(tmap: T.Tensor((8,), "uint64")):
T.device_entry()
- shared_buf = T.alloc_tensor((16,), "uint32", scope="shared")
- barrier = T.alloc_tensor((1,), "uint64", scope="shared")
+ shared_buf = T.alloc_tensor(ty_args=[T.Tensor((16,), "uint32", scope="shared")])
+ barrier = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint64", scope="shared")])
T.ptx["cp.async.bulk.tensor.1d.shared::cta.global.mbarrier::complete_tx::bytes"](
shared_buf.data, T.ptx.addr(tmap.data, 16), T.int32(0), barrier.data
)
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py b/tests/python/tirx/codegen/test_ptx_dialect.py
index 59d6c9ec7c..89a881fdf3 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -137,7 +137,7 @@ def test_ptx_st_shared_coercion():
def kernel(out: T.Tensor((1,), "uint32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
tx = T.cuda.thread_idx("x")
- smem = T.alloc_tensor((4,), "uint32", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint32", scope="shared")])
if tx == 0:
# Shared-space slot fed a shared-scope pointer: engine must
# auto-wrap with cvta_generic_to_shared.
@@ -155,7 +155,7 @@ def test_ptx_explicit_cvta():
def kernel(out: T.Tensor((1,), "uint64")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
tx = T.cuda.thread_idx("x")
- smem = T.alloc_tensor((4,), "uint32", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint32", scope="shared")])
smem[tx % 4] = T.uint32(0)
if tx == 0:
T.ptx.cvta.to.shared.u64(out[0], smem.data)
@@ -626,8 +626,8 @@ def test_ptx_92_cp_bulk_roundtrip():
@T.function
def kernel(src: T.Tensor((64,), "uint32"), dst: T.Tensor((64,), "uint32")):
T.device_entry()
- smem = T.alloc_tensor((64,), "uint32", scope="shared")
- mbar = T.alloc_tensor((2,), "uint64", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((64,), "uint32", scope="shared")])
+ mbar = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint64", scope="shared")])
T.ptx[
"cp.async.bulk.shared::cta.global."
"mbarrier::complete_tx::bytes.L2::cache_hint.ignore_oob"
@@ -839,7 +839,7 @@ def test_ptx_optional_operand_arity_dispatch():
def kernel(A: T.Tensor((4,), "uint32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
tx = T.cuda.thread_idx("x")
- bar = T.alloc_tensor((2,), "uint64", scope="shared")
+ bar = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint64", scope="shared")])
T.ptx.bar.sync(T.uint32(0))
T.ptx.bar.sync(T.uint32(0), T.uint32(64))
T.ptx.mbarrier.arrive.shared.b64(bar.ptr_to([0]))
@@ -870,7 +870,7 @@ def test_ptx_mbarrier_92_shapes_render_and_roundtrip():
def kernel(out: T.Tensor((4,), "uint32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
tx = T.cuda.thread_idx("x")
- bar = T.alloc_tensor((2,), "uint64", scope="shared")
+ bar = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint64", scope="shared")])
state = T.local_scalar("uint64")
pending = T.local_scalar("uint32")
wait_complete = T.local_scalar("uint32")
@@ -1906,7 +1906,7 @@ def test_ptx_data_movement_dispatch():
def kernel(A: T.Tensor((8,), "uint32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
tx = T.cuda.thread_idx("x")
- smem = T.alloc_tensor((4,), "uint32", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint32", scope="shared")])
d = T.local_scalar("uint32")
p = T.local_scalar("uint32")
v = T.local_scalar("uint32")
@@ -2123,7 +2123,7 @@ def test_ptx_lazy_subscript_operands_realize():
@T.function
def kernel(A: T.Tensor((1,), "uint32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
- regs = T.alloc_tensor((2,), "uint32", scope="local")
+ regs = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint32", scope="local")])
full = T.uint32(0xFFFFFFFF)
T.ptx.elect_sync(regs[0], regs[1], full)
T.ptx.vote_sync.all.pred(regs[0], T.ptx.pred(regs[1]), full)
@@ -2143,7 +2143,7 @@ def test_ptx_parser_roundtrip():
def kernel(A: T.Tensor((32,), "uint32"), B: T.Tensor((32,), "uint32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
tx = T.cuda.thread_idx("x")
- smem = T.alloc_tensor((4,), "uint32", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint32", scope="shared")])
if tx == 0:
val = T.local_scalar("uint32")
smem_addr = T.local_scalar("uint64")
@@ -2938,8 +2938,8 @@ def test_ptx_94_cp_bulk_semantics_roundtrip():
@T.function
def kernel(src: T.Tensor((64,), "uint32")):
T.device_entry()
- smem = T.alloc_tensor((64,), "uint32", scope="shared")
- mbar = T.alloc_tensor((2,), "uint64", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((64,), "uint32", scope="shared")])
+ mbar = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint64", scope="shared")])
T.ptx[
"cp.async.bulk.relaxed.sys.shared::cta.global."
"mbarrier::complete_tx::bytes."
@@ -4515,7 +4515,7 @@ def test_ptx_nvcc_smoke():
def kernel(A: T.Tensor((32,), "uint32"), B: T.Tensor((32,), "uint32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
tx = T.cuda.thread_idx("x")
- smem = T.alloc_tensor((4,), "uint32", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint32", scope="shared")])
if tx == 0:
val = T.local_scalar("uint32")
T.ptx.ld.global_.acquire.gpu.b32(val, A.ptr_to([0]))
diff --git a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
index 9ec4b24032..e33a19b379 100644
--- a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
+++ b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
@@ -78,8 +78,8 @@ def _shared_scratch_copy_kernel(num_bytes: int):
def func(out: T.Tensor((nelems,), smem_dtype)):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=1 * 32))
lane = T.cuda.lane_id()
- src_buf = T.alloc_tensor((nelems,), smem_dtype, scope="shared")
- dst_buf = T.alloc_tensor((nelems,), smem_dtype, scope="shared")
+ src_buf = T.alloc_tensor(ty_args=[T.Tensor((nelems,), smem_dtype, scope="shared")])
+ dst_buf = T.alloc_tensor(ty_args=[T.Tensor((nelems,), smem_dtype, scope="shared")])
tmp = T.alloc_local((lanes,), reg_dtype)
if T.constexpr(fill_offset is not None):
if lane < nelems:
@@ -112,7 +112,7 @@ def test_ptx_ld_st_codegen_emits_shared_asm():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=4 * 32))
tid_in_wg = T.cuda.thread_in_warpgroup()
- smem = T.alloc_tensor((4,), "uint32", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint32", scope="shared")])
reg = T.alloc_local((4,), "uint32")
if tid_in_wg == 0:
T.ptx.st.shared.v4.u32(smem.ptr_to([0]), reg[0], reg[1], reg[2], reg[3])
@@ -137,7 +137,7 @@ def test_ptx_ld_st_raw_shared_address_codegen():
def main(out: T.Tensor((2,), "uint64")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(32,)))
tx = T.cuda.thread_idx("x")
- smem = T.alloc_tensor((2,), "uint64", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((2,), "uint64", scope="shared")])
values = T.alloc_local((4,), "uint32")
if tx == 0:
raw_addr: T.uint32 = T.cuda.cvta_generic_to_shared(smem.data)
diff --git a/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py b/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
index 76dd8f4af2..0c373c324f 100644
--- a/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
+++ b/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
@@ -72,7 +72,7 @@ def test_smem_descriptor_matches_runtime_encoder(ldo, sdo, swizzle):
def kernel(out: T.Tensor((2,), "uint64")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
tx = T.cuda.thread_idx("x")
- smem = T.alloc_tensor((64,), "uint32", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((64,), "uint32", scope="shared")])
smem[tx] = T.uint32(0)
smem[tx + 32] = T.uint32(0)
if tx == 0:
diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py
index 0cde9aac47..15e046463b 100644
--- a/tests/python/tirx/iket/test_iket_profiler.py
+++ b/tests/python/tirx/iket/test_iket_profiler.py
@@ -393,7 +393,7 @@ def test_public_interface_is_official_only():
payload_script = payload_types.script()
assert 'T.cuda.iket_mark("i8", T.int8(-8))' in payload_script
assert 'T.cuda.iket_range_start("token_payload", -7)' in payload_script
- assert "T.cuda.iket_range_end(token, 9)" in payload_script
+ assert "T.cuda.iket_range_end(token[0], 9)" in payload_script
assert (
tvm.script.from_source(
payload_script, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}
diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
index bfa28ce070..a28a3ded91 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
@@ -45,7 +45,9 @@ def _build_kernel(scope, n_threads, shape, dtype):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(n_threads,)))
_lane = T.cuda.lane_id()
copy_thread_id = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="shared", layout=s_layout)]
+ )
copy_src_6 = T.meta_var(A[full_slices])
copy_src_tensor_6 = T.meta_var(copy_src_6.source)
copy_dst_6 = T.meta_var(A_smem[full_slices])
@@ -119,7 +121,9 @@ def _build_kernel(scope, n_threads, shape, dtype):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(n_threads,)))
_lane = T.cuda.lane_id()
copy_thread_id = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="shared", layout=s_layout)]
+ )
copy_src_8 = T.meta_var(A[full_slices])
copy_src_tensor_8 = T.meta_var(copy_src_8.source)
copy_dst_8 = T.meta_var(A_smem[full_slices])
@@ -193,7 +197,9 @@ def _build_kernel(scope, n_threads, shape, dtype):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(n_threads,)))
_lane = T.cuda.lane_id()
copy_thread_id = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="shared", layout=s_layout)]
+ )
copy_src_10 = T.meta_var(A[full_slices])
copy_src_tensor_10 = T.meta_var(copy_src_10.source)
copy_dst_10 = T.meta_var(A_smem[full_slices])
@@ -392,7 +398,7 @@ def test_copy_g2s_s2g(task, dtype, scope):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(2,), block=(thread_cnt,)))
copy_thread_id = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(s_shape, dtype, scope="shared", layout=layoutS)
+ A_smem = T.alloc_tensor(ty_args=[T.Tensor(s_shape, dtype, scope="shared", layout=layoutS)])
# `scope` is parametrized at runtime; select the scope namespace
# dynamically (T.cta / T.thread) instead of a literal prefix.
copy_src_1 = T.meta_var(A[r_gmem])
@@ -595,7 +601,9 @@ def test_swizzled_smem_emit_must_be_swizzle_aware():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
_lane = T.cuda.lane_id()
copy_thread_id = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(shape, "float16", scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, "float16", scope="shared", layout=s_layout)]
+ )
copy_src_3 = T.meta_var(A[0:128, 0:32])
copy_src_tensor_3 = T.meta_var(copy_src_3.source)
copy_dst_3 = T.meta_var(A_smem[0:128, 0:32])
diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
index d25bb3a079..0837eb0c88 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
@@ -117,7 +117,9 @@ def _build_warp_kernel(num, direction, trans, swizzle=False):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor(s_shape, "float16", scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="shared", layout=s_layout)]
+ )
row = tid // 4
cp = tid % 4
for t in range(num):
@@ -125,7 +127,9 @@ def _build_warp_kernel(num, direction, trans, swizzle=False):
gr, gc = _coord(row, cp, t, w)
A_smem[row, cp, t, w] = A[gr, gc]
T.cuda.cta_sync()
- R_local = T.alloc_tensor(s_shape, "float16", scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="local", layout=r_layout)]
+ )
T.cuda.tile.ldmatrix(R_local[full], A_smem[full], scope='warp')
r_view = R_local.local()
for t in range(num):
@@ -139,10 +143,14 @@ def _build_warp_kernel(num, direction, trans, swizzle=False):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor(s_shape, "float16", scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="shared", layout=s_layout)]
+ )
row = tid // 4
cp = tid % 4
- R_local = T.alloc_tensor(s_shape, "float16", scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="local", layout=r_layout)]
+ )
r_view = R_local.local()
for t in range(num):
for w in range(2):
@@ -181,7 +189,9 @@ def _build_warpgroup_kernel(num, direction, trans, swizzle=False):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor(s_shape, "float16", scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="shared", layout=s_layout)]
+ )
wid = tid // 32
lid = tid % 32
row = lid // 4
@@ -191,7 +201,9 @@ def _build_warpgroup_kernel(num, direction, trans, swizzle=False):
gr, gc = _coord(wid, row, cp, t, w)
A_smem[wid, row, cp, t, w] = A[gr, gc]
T.cuda.cta_sync()
- R_local = T.alloc_tensor(s_shape, "float16", scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="local", layout=r_layout)]
+ )
T.cuda.tile.ldmatrix(R_local[full], A_smem[full], scope='warpgroup')
r_view = R_local.local()
for t in range(num):
@@ -205,12 +217,16 @@ def _build_warpgroup_kernel(num, direction, trans, swizzle=False):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor(s_shape, "float16", scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="shared", layout=s_layout)]
+ )
wid = tid // 32
lid = tid % 32
row = lid // 4
cp = lid % 4
- R_local = T.alloc_tensor(s_shape, "float16", scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="local", layout=r_layout)]
+ )
r_view = R_local.local()
for t in range(num):
for w in range(2):
@@ -249,7 +265,9 @@ def _build_cta_kernel(num, direction, trans, swizzle=False):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor(s_shape, "float16", scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="shared", layout=s_layout)]
+ )
wid = tid // 32
lid = tid % 32
row = lid // 4
@@ -259,7 +277,9 @@ def _build_cta_kernel(num, direction, trans, swizzle=False):
gr, gc = _coord(wid, row, cp, t, w)
A_smem[wid, row, cp, t, w] = A[gr, gc]
T.cuda.cta_sync()
- R_local = T.alloc_tensor(s_shape, "float16", scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="local", layout=r_layout)]
+ )
T.cuda.tile.ldmatrix(R_local[full], A_smem[full], scope='cta')
r_view = R_local.local()
for t in range(num):
@@ -273,12 +293,16 @@ def _build_cta_kernel(num, direction, trans, swizzle=False):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx('x')
- A_smem = T.alloc_tensor(s_shape, "float16", scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="shared", layout=s_layout)]
+ )
wid = tid // 32
lid = tid % 32
row = lid // 4
cp = lid % 4
- R_local = T.alloc_tensor(s_shape, "float16", scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(s_shape, "float16", scope="local", layout=r_layout)]
+ )
r_view = R_local.local()
for t in range(num):
for w in range(2):
@@ -383,14 +407,18 @@ def _build_multi_iter_kernel(outer_ext: int):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(shape, "float16", scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, "float16", scope="shared", layout=s_layout)]
+ )
for a in range(outer_ext):
for c in range(2):
for d in range(4):
for e in range(2):
A_smem[a, tid // 4, c, d, tid % 4, e] = A[a, tid // 4, c, d, tid % 4, e]
T.cuda.cta_sync()
- R_local = T.alloc_tensor(shape, "float16", scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, "float16", scope="local", layout=r_layout)]
+ )
T.cuda.tile.ldmatrix(R_local[full], A_smem[full], scope="warp")
r_view = R_local.local()
for a in range(outer_ext):
@@ -462,7 +490,9 @@ def test_ldstmatrix_tcgen05_warpgroup_atom_emits_ldmatrix():
def kernel(smem: T.Tensor((m, k), "bfloat16", scope="shared", layout=smem_layout)) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=1 * 128))
_lane = T.cuda.lane_id()
- a_reg = T.alloc_tensor((m, k), "bfloat16", scope="local", layout=reg_layout)
+ a_reg = T.alloc_tensor(
+ ty_args=[T.Tensor((m, k), "bfloat16", scope="local", layout=reg_layout)]
+ )
T.cuda.tile.ldmatrix(a_reg, smem, scope="warpgroup")
_, src = _compile_src(kernel)
diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
index 616b30e42d..931f1e0914 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
@@ -72,11 +72,15 @@ def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(n_threads,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="shared", layout=s_layout)]
+ )
for kk in range(k):
A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
T.cuda.cta_sync()
- R_local = T.alloc_tensor(shape, dtype, scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="local", layout=r_layout)]
+ )
T.cuda.tile.ld(R_local[full_slices], A_smem[full_slices], scope="warpgroup")
for kk in range(k):
A_smem[tid, kk] = T.cast(0, dtype)
@@ -93,11 +97,15 @@ def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(n_threads,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="shared", layout=s_layout)]
+ )
for kk in range(k):
A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
T.cuda.cta_sync()
- R_local = T.alloc_tensor(shape, dtype, scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="local", layout=r_layout)]
+ )
T.cuda.tile.ld(R_local[full_slices], A_smem[full_slices], scope="warp")
for kk in range(k):
A_smem[tid, kk] = T.cast(0, dtype)
@@ -114,11 +122,15 @@ def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope):
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(n_threads,)))
_lane = T.cuda.lane_id()
tid = T.cuda.thread_idx("x")
- A_smem = T.alloc_tensor(shape, dtype, scope="shared", layout=s_layout)
+ A_smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="shared", layout=s_layout)]
+ )
for kk in range(k):
A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
T.cuda.cta_sync()
- R_local = T.alloc_tensor(shape, dtype, scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="local", layout=r_layout)]
+ )
T.cuda.tile.ld(R_local[full_slices], A_smem[full_slices], scope="cta")
for kk in range(k):
A_smem[tid, kk] = T.cast(0, dtype)
@@ -141,7 +153,9 @@ def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope):
for kk in range(k):
A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
T.cuda.cta_sync()
- R_local = T.alloc_tensor(shape, dtype, scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="local", layout=r_layout)]
+ )
T.cuda.tile.ld(R_local[full_slices], A[full_slices], scope="warpgroup")
for kk in range(k):
A[tid, kk] = T.cast(0, dtype)
@@ -161,7 +175,9 @@ def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope):
for kk in range(k):
A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
T.cuda.cta_sync()
- R_local = T.alloc_tensor(shape, dtype, scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="local", layout=r_layout)]
+ )
T.cuda.tile.ld(R_local[full_slices], A[full_slices], scope="warp")
for kk in range(k):
A[tid, kk] = T.cast(0, dtype)
@@ -181,7 +197,9 @@ def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope):
for kk in range(k):
A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
T.cuda.cta_sync()
- R_local = T.alloc_tensor(shape, dtype, scope="local", layout=r_layout)
+ R_local = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, dtype, scope="local", layout=r_layout)]
+ )
T.cuda.tile.ld(R_local[full_slices], A[full_slices], scope="cta")
for kk in range(k):
A[tid, kk] = T.cast(0, dtype)
@@ -266,7 +284,7 @@ def test_reg_roundtrip_gapped_permuted_storage():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(32,)))
_lane = T.cuda.lane_id()
- reg = T.alloc_tensor(shape, "float32", scope="local", layout=r_layout)
+ reg = T.alloc_tensor(ty_args=[T.Tensor(shape, "float32", scope="local", layout=r_layout)])
T.cuda.tile.ld(reg, A, scope='warp')
T.cuda.tile.st(B, reg, scope='warp')
# fmt: on
@@ -331,7 +349,9 @@ def test_copy_g2l_l2g_vec_load(task, dtype):
A: T.Tensor(g_shape, dtype, layout=layoutA), B: T.Tensor(g_shape, dtype, layout=layoutB)
) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(2,), block=(thread_cnt,)))
- A_local = T.alloc_tensor(l_shape, dtype, scope="local", layout=layoutLocal)
+ A_local = T.alloc_tensor(
+ ty_args=[T.Tensor(l_shape, dtype, scope="local", layout=layoutLocal)]
+ )
T.cuda.tile.ld(A_local[r_lmem], A[r_gmem])
T.cuda.tile.st(B[r_gmem], A_local[r_lmem])
@@ -425,8 +445,12 @@ def test_reg_copy_linear_shared_hoists_thread_base():
def kernel(A: T.Tensor(shape, "float16", layout=linear_layout)) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(n_threads,)))
tid = T.cuda.thread_in_warpgroup()
- reg = T.alloc_tensor(shape, "float16", scope="local", layout=wg_local_layout(width))
- smem = T.alloc_tensor(shape, "float16", scope="shared", layout=linear_layout)
+ reg = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, "float16", scope="local", layout=wg_local_layout(width))]
+ )
+ smem = T.alloc_tensor(
+ ty_args=[T.Tensor(shape, "float16", scope="shared", layout=linear_layout)]
+ )
reg_local = reg.local(width)
for i in T.serial(width):
@@ -507,8 +531,12 @@ def test_reg_copy_wg_local_to_swizzled_shared_uses_structured_compose_apply():
) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(N_THREADS,)))
tid = T.cuda.thread_in_warpgroup()
- reg = T.alloc_tensor(g_shape, "float16", scope="local", layout=wg_local_layout(EPI_N))
- smem = T.alloc_tensor(g_shape, "float16", scope="shared", layout=smem_layout)
+ reg = T.alloc_tensor(
+ ty_args=[T.Tensor(g_shape, "float16", scope="local", layout=wg_local_layout(EPI_N))]
+ )
+ smem = T.alloc_tensor(
+ ty_args=[T.Tensor(g_shape, "float16", scope="shared", layout=smem_layout)]
+ )
# Populate the per-thread slice via .local() (decomposes the wg
# thread-axis layout into a per-thread 1D view).
@@ -555,7 +583,7 @@ def test_ptx_st_from_src_f32_vector_preserves_values():
@T.function
def kernel(B: T.Tensor((4,), "float32")) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(1,)))
- smem = T.alloc_tensor((4,), "float32", scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32", scope="shared")])
reg = T.alloc_local((4,), "float32")
out = T.alloc_local((4,), "float32")
for i in range(4):
@@ -615,7 +643,7 @@ def test_copy_forced_vec_width_codegen(variant, dtype, n_elements, expected_st,
@T.function
def kernel(B: T.Tensor((n_elements,), dtype)) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(1,)))
- smem = T.alloc_tensor((n_elements,), dtype, scope="shared")
+ smem = T.alloc_tensor(ty_args=[T.Tensor((n_elements,), dtype, scope="shared")])
reg = T.alloc_local((n_elements,), dtype)
out = T.alloc_local((n_elements,), dtype)
for i in range(n_elements):
@@ -657,7 +685,9 @@ def test_copy_forced_vec_dynamic_swizzled_shared_uses_vector_ptx():
def kernel(B: T.Tensor((128, 4), "float32")) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
tid = T.cuda.thread_idx("x")
- smem = T.alloc_tensor((64, 256), "float32", scope="shared", layout=smem_layout)
+ smem = T.alloc_tensor(
+ ty_args=[T.Tensor((64, 256), "float32", scope="shared", layout=smem_layout)]
+ )
reg = T.alloc_local((4,), "float32")
out = T.alloc_local((4,), "float32")
for i in range(4):
@@ -693,9 +723,13 @@ def test_copy_explicit_vec_auto_uses_auto_family():
@T.function
def kernel(B: T.Tensor((4,), "float32")) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(1,)))
- smem = T.alloc_tensor((4,), "float32", scope="shared")
- reg = T.alloc_tensor((4,), "float32", scope="local", layout=TileLayout(S[4]))
- out = T.alloc_tensor((4,), "float32", scope="local", layout=TileLayout(S[4]))
+ smem = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32", scope="shared")])
+ reg = T.alloc_tensor(
+ ty_args=[T.Tensor((4,), "float32", scope="local", layout=TileLayout(S[4]))]
+ )
... 6304 lines suppressed ...