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 d4759d513e [REFACTOR][TIRx] Remove redundant helpers and construct
tensor variables directly (#20776)
d4759d513e is described below
commit d4759d513e376bf7f6ca25f9925d472a4c19aae7
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Oct 17 15:35:59 2026 -0400
[REFACTOR][TIRx] Remove redundant helpers and construct tensor variables
directly (#20776)
---
python/tvm/relax/transform/legalize_ops/grad.py | 2 +-
python/tvm/te/operation.py | 24 +--
python/tvm/topi/gpu/scan.py | 50 ++++--
python/tvm/topi/gpu/scatter_elements.py | 2 +-
python/tvm/topi/gpu/scatter_nd.py | 2 +-
python/tvm/topi/gpu/sort.py | 90 +++++++----
python/tvm/topi/index_put.py | 2 +-
python/tvm/topi/scan.py | 4 +-
python/tvm/topi/scatter.py | 2 +-
python/tvm/topi/scatter_elements.py | 2 +-
python/tvm/topi/signal.py | 2 +-
python/tvm/topi/sort.py | 46 +++---
python/tvm/topi/sparse_reshape.py | 8 +-
python/tvm/topi/vision/nms.py | 84 ++++++----
python/tvm/topi/vision/nms_util.py | 12 +-
tests/python/contrib/test_sort.py | 10 +-
.../test_s_tir_analysis_verify_block_scope.py | 6 +-
.../s_tir/analysis/test_sblock_access_region.py | 4 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 8 +-
.../schedule/test_tir_schedule_transform_layout.py | 2 +-
.../s_tir/script/test_s_tir_script_basic_usage.py | 14 +-
.../script/test_s_tir_script_dynamic_shape.py | 8 +-
.../s_tir/script/test_s_tir_script_ir_builder.py | 18 ++-
.../s_tir/script/test_s_tir_script_printer.py | 32 ++--
tests/python/s_tir/test_s_tir_renew_defs.py | 4 +-
.../test_s_tir_transform_mma_tensor_layout.py | 2 +-
tests/python/te/test_te_tensor.py | 6 +
tests/python/tirx/base/test_tirx_constructor.py | 16 +-
.../tirx/script/test_tirx_script_ir_builder.py | 13 +-
.../python/tirx/script/test_tirx_script_printer.py | 12 +-
tirx/python/tvm_tirx/backend/cuda/op.py | 176 ---------------------
tirx/python/tvm_tirx/script/ir_builder/base.py | 18 ---
tirx/python/tvm_tirx/tirx/__init__.py | 1 -
tirx/python/tvm_tirx/tirx/compilation_pipeline.py | 6 -
tirx/python/tvm_tirx/tirx/function.py | 3 +-
tirx/python/tvm_tirx/tirx/tensor.py | 49 +-----
tirx/src/tirx/ir/specialize.cc | 2 +-
tirx/src/tirx/ir/tensor.cc | 4 +-
tirx/tests/python/codegen/c/test_basic.py | 2 +-
tirx/tests/python/sym/test_deduce_bound.py | 4 +-
tirx/tests/python/sym/test_iter_affine_map.py | 2 +-
tirx/tests/python/sym/test_rewrite_simplify.py | 2 +-
tirx/tests/python/sym/test_z3.py | 2 +-
.../tirx/analysis/test_check_well_defined.py | 2 +-
.../python/tirx/analysis/test_undefined_vars.py | 8 +-
tirx/tests/python/tirx/base/test_nodes.py | 22 +--
tirx/tests/python/tirx/base/test_op_types.py | 63 ++++----
tirx/tests/python/tirx/base/test_specialize.py | 25 ++-
tirx/tests/python/tirx/base/test_tensor_type.py | 110 ++++++++++---
tirx/tests/python/tirx/codegen/test_cuda.py | 2 +-
tirx/tests/python/tirx/codegen/test_ptx_dialect.py | 6 +-
tirx/tests/python/tirx/script/test_tensors.py | 38 +++--
tirx/tests/python/tirx/test_op.py | 16 +-
.../tests/python/tirx/test_op_namespace_cleanup.py | 2 +-
.../python/tirx/test_tensor_data_reinfer_type.py | 4 +-
tirx/tests/python/tirx/test_thread_scope.py | 8 +-
tirx/tests/python/tirx/test_tile_namespace.py | 20 +--
.../python/tirx/tile/cuda/test_cp_async_bulk.py | 12 +-
.../tirx/tile/cuda/test_cp_async_bulk_tensor.py | 115 ++++++++------
tirx/tests/python/tirx/tile/test_dispatcher.py | 10 +-
tirx/tests/python/tirx/tile/trn/test_copy.py | 4 +-
61 files changed, 617 insertions(+), 608 deletions(-)
diff --git a/python/tvm/relax/transform/legalize_ops/grad.py
b/python/tvm/relax/transform/legalize_ops/grad.py
index 8b63a39f55..cda3f6b456 100644
--- a/python/tvm/relax/transform/legalize_ops/grad.py
+++ b/python/tvm/relax/transform/legalize_ops/grad.py
@@ -229,7 +229,7 @@ def _grad_take_backward(bb: BlockBuilder, call: Call) ->
Expr:
return ib.get()
shape = x.shape
- out_buf = tirx.decl_tensor(shape, x.dtype, "out_buf", layout=None)
+ out_buf = tirx.Var("out_buf", tirx.TensorType(shape, x.dtype,
layout=None))
return te.extern(
[shape],
diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py
index 82e791b714..8aa676445b 100644
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -312,12 +312,14 @@ def extern(
raise ValueError("expect inputs to be tensor")
if in_tensors is None:
input_placeholders.append(
- tvm_tirx.tirx.decl_tensor(
- t.shape,
- t.dtype,
+ tvm_tirx.tirx.Var(
t.op.name,
- elem_offset=tvm_tirx.tirx.Var("elem_offset", "int32"),
- layout=None,
+ tvm_tirx.tirx.TensorType(
+ t.shape,
+ t.dtype,
+ elem_offset=tvm_tirx.tirx.Var("elem_offset", "int32"),
+ layout=None,
+ ),
)
)
types.add(t.dtype)
@@ -333,12 +335,14 @@ def extern(
for shp, dt in zip(shape, dtype):
output_placeholders.append(
- tvm_tirx.tirx.decl_tensor(
- shp,
- dt,
+ tvm_tirx.tirx.Var(
name,
- elem_offset=tvm_tirx.tirx.Var("elem_offset", "int32"),
- layout=None,
+ tvm_tirx.tirx.TensorType(
+ shp,
+ "float32" if dt is None else dt,
+ elem_offset=tvm_tirx.tirx.Var("elem_offset", "int32"),
+ layout=None,
+ ),
)
)
body = fcompute(input_placeholders, output_placeholders)
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
index 937ae36b74..bd5bbc613b 100644
--- a/python/tvm/topi/gpu/scan.py
+++ b/python/tvm/topi/gpu/scan.py
@@ -411,15 +411,15 @@ def get_reduction_from_exclusive_scan(data,
ex_scan_output, binop=operator.add):
return ib.get()
- data_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "valid_indices_buf", data_alignment=8,
layout=None
+ data_buf = tvm_tirx.tirx.Var(
+ "valid_indices_buf",
+ tvm_tirx.tirx.TensorType(data.shape, data.dtype, data_alignment=8,
layout=None),
)
- ex_scan_output_buf = tvm_tirx.tirx.decl_tensor(
- ex_scan_output.shape,
- ex_scan_output.dtype,
+ ex_scan_output_buf = tvm_tirx.tirx.Var(
"ex_scan_output_buf",
- data_alignment=8,
- layout=None,
+ tvm_tirx.tirx.TensorType(
+ ex_scan_output.shape, ex_scan_output.dtype, data_alignment=8,
layout=None
+ ),
)
reduction = te.extern(
@@ -483,16 +483,25 @@ def scan_thrust(
(N-1)-D tensor storing the reduction of each scan axis.
Returned if return_reduction is True.
"""
- data_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+ data_buf = tvm_tirx.tirx.Var(
+ "data_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
- output_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, output_dtype, "output_buf", data_alignment=8, layout=None
+ output_buf = tvm_tirx.tirx.Var(
+ "output_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape,
+ "float32" if output_dtype is None else output_dtype,
+ data_alignment=8,
+ layout=None,
+ ),
)
workspace_buf = (
- tvm_tirx.tirx.decl_tensor(
- workspace.shape, workspace.dtype, "workspace_buf",
data_alignment=8, layout=None
+ tvm_tirx.tirx.Var(
+ "workspace_buf",
+ tvm_tirx.tirx.TensorType(
+ workspace.shape, workspace.dtype, data_alignment=8, layout=None
+ ),
)
if workspace is not None
else None
@@ -592,11 +601,18 @@ def exclusive_scan(
# TIR exclusive scan accepts only 2D or higher-rank inputs.
data = expand_dims(data, axis=0)
- data_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+ data_buf = tvm_tirx.tirx.Var(
+ "data_buf",
+ tvm_tirx.tirx.TensorType(data.shape, data.dtype, data_alignment=8,
layout=None),
)
- output_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, output_dtype, "output_buf", data_alignment=8,
layout=None
+ output_buf = tvm_tirx.tirx.Var(
+ "output_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape,
+ "float32" if output_dtype is None else output_dtype,
+ data_alignment=8,
+ layout=None,
+ ),
)
if return_reduction:
diff --git a/python/tvm/topi/gpu/scatter_elements.py
b/python/tvm/topi/gpu/scatter_elements.py
index 7fb95ad036..8a0c4b8e36 100644
--- a/python/tvm/topi/gpu/scatter_elements.py
+++ b/python/tvm/topi/gpu/scatter_elements.py
@@ -170,7 +170,7 @@ def scatter_elements(data, indices, updates, axis=0,
reduction="update"):
"scatter_elements reduction not in [update, add, mul, mean, min,
max]:", reduction
)
- out_buf = tirx.decl_tensor(data.shape, data.dtype, "out_buf", layout=None)
+ out_buf = tirx.Var("out_buf", tirx.TensorType(data.shape, data.dtype,
layout=None))
return te.extern(
[data.shape],
[data, indices, updates],
diff --git a/python/tvm/topi/gpu/scatter_nd.py
b/python/tvm/topi/gpu/scatter_nd.py
index 67f780979c..0ea9f04e34 100644
--- a/python/tvm/topi/gpu/scatter_nd.py
+++ b/python/tvm/topi/gpu/scatter_nd.py
@@ -168,7 +168,7 @@ def scatter_nd(data, indices, updates, mode):
return ib.get()
- out_buf = tirx.decl_tensor(data.shape, data.dtype, "out_buf", layout=None)
+ out_buf = tirx.Var("out_buf", tirx.TensorType(data.shape, data.dtype,
layout=None))
return te.extern(
[data.shape],
[data, indices, updates],
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index 7ebf2d9499..4c3df93f16 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -788,11 +788,12 @@ def sort(data, axis=-1, is_ascend=1):
axes = swap(list(range(ndim)), axis)
data = transpose(data, axes)
- value_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "value_buf", data_alignment=8, layout=None
+ value_buf = tvm_tirx.tirx.Var(
+ "value_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
- value_buf_swap = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "value_buf_swap", data_alignment=8, layout=None
+ value_buf_swap = tvm_tirx.tirx.Var(
+ "value_buf_swap",
+ tvm_tirx.tirx.TensorType(data.shape, data.dtype, data_alignment=8,
layout=None),
)
out = te.extern(
@@ -846,11 +847,14 @@ def sort_thrust(data, axis=-1, is_ascend=1,
workspace=None):
axes = swap(list(range(ndim)), axis)
data = transpose(data, axes)
- value_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "value_buf", data_alignment=8, layout=None
+ value_buf = tvm_tirx.tirx.Var(
+ "value_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
- indices_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, dtype, "out_buf", data_alignment=8, layout=None
+ indices_buf = tvm_tirx.tirx.Var(
+ "out_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape, "float32" if dtype is None else dtype,
data_alignment=8, layout=None
+ ),
)
def f_compute(ins, outs):
@@ -912,17 +916,24 @@ def argsort(data, axis=-1, is_ascend=1, dtype="float32",
ret_type="indices"):
axes = swap(list(range(ndim)), axis)
data = transpose(data, axes)
- value_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "value_buf", data_alignment=8, layout=None
+ value_buf = tvm_tirx.tirx.Var(
+ "value_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
- value_swap_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "value_swap_buf", data_alignment=8, layout=None
+ value_swap_buf = tvm_tirx.tirx.Var(
+ "value_swap_buf",
+ tvm_tirx.tirx.TensorType(data.shape, data.dtype, data_alignment=8,
layout=None),
)
- indices_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, dtype, "out_buf", data_alignment=8, layout=None
+ indices_buf = tvm_tirx.tirx.Var(
+ "out_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape, "float32" if dtype is None else dtype,
data_alignment=8, layout=None
+ ),
)
- indices_swap_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, dtype, "out_swap_buf", data_alignment=8, layout=None
+ indices_swap_buf = tvm_tirx.tirx.Var(
+ "out_swap_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape, "float32" if dtype is None else dtype,
data_alignment=8, layout=None
+ ),
)
outs = te.extern(
@@ -1028,17 +1039,25 @@ def topk(data, k=1, axis=-1, ret_type="both",
is_ascend=False, dtype="int64"):
axes = swap(list(range(ndim)), axis)
data = transpose(data, axes)
- values_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "values_buf", data_alignment=8, layout=None
+ values_buf = tvm_tirx.tirx.Var(
+ "values_buf",
+ tvm_tirx.tirx.TensorType(data.shape, data.dtype, data_alignment=8,
layout=None),
)
- values_swap_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "values_swap_buf", data_alignment=8,
layout=None
+ values_swap_buf = tvm_tirx.tirx.Var(
+ "values_swap_buf",
+ tvm_tirx.tirx.TensorType(data.shape, data.dtype, data_alignment=8,
layout=None),
)
- indices_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, dtype, "indices_buf", data_alignment=8, layout=None
+ indices_buf = tvm_tirx.tirx.Var(
+ "indices_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape, "float32" if dtype is None else dtype,
data_alignment=8, layout=None
+ ),
)
- indices_swap_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, dtype, "indies_swap_buf", data_alignment=8, layout=None
+ indices_swap_buf = tvm_tirx.tirx.Var(
+ "indies_swap_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape, "float32" if dtype is None else dtype,
data_alignment=8, layout=None
+ ),
)
if ret_type == "values":
@@ -1139,20 +1158,29 @@ def topk_thrust(
axes = swap(list(range(ndim)), axis)
data = transpose(data, axes)
- data_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+ data_buf = tvm_tirx.tirx.Var(
+ "data_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
if workspace is not None:
- workspace_buf = tvm_tirx.tirx.decl_tensor(
- workspace.shape, workspace.dtype, "workspace_buf",
data_alignment=8, layout=None
+ workspace_buf = tvm_tirx.tirx.Var(
+ "workspace_buf",
+ tvm_tirx.tirx.TensorType(
+ workspace.shape, workspace.dtype, data_alignment=8, layout=None
+ ),
)
else:
workspace_buf = None
out_bufs = [
- tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "value_buf", data_alignment=8, layout=None
+ tvm_tirx.tirx.Var(
+ "value_buf",
+ tvm_tirx.tirx.TensorType(data.shape, data.dtype, data_alignment=8,
layout=None),
+ ),
+ tvm_tirx.tirx.Var(
+ "indices_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape, "float32" if dtype is None else dtype,
data_alignment=8, layout=None
+ ),
),
- tvm_tirx.tirx.decl_tensor(data.shape, dtype, "indices_buf",
data_alignment=8, layout=None),
]
def f_compute(ins, outs):
diff --git a/python/tvm/topi/index_put.py b/python/tvm/topi/index_put.py
index 03948f9870..1b45515f7f 100644
--- a/python/tvm/topi/index_put.py
+++ b/python/tvm/topi/index_put.py
@@ -157,7 +157,7 @@ def index_put(data, indices, values, accumulate=False):
in_tensors.extend(indices)
in_tensors.append(values)
- out_buf = tirx.decl_tensor(data.shape, data.dtype, "out_buf", layout=None)
+ out_buf = tirx.Var("out_buf", tirx.TensorType(data.shape, data.dtype,
layout=None))
return te.extern(
[data.shape],
in_tensors,
diff --git a/python/tvm/topi/scan.py b/python/tvm/topi/scan.py
index 3f737c4a20..0260a8610b 100644
--- a/python/tvm/topi/scan.py
+++ b/python/tvm/topi/scan.py
@@ -22,7 +22,7 @@ from collections.abc import Callable
import tvm
from tvm_tirx.script.ir_builder import IRBuilder
-from tvm_tirx.tirx import decl_tensor
+from tvm_tirx.tirx import TensorType, Var
from tvm_tirx.tirx.script import ir_builder as T
from ..te import extern
@@ -156,7 +156,7 @@ def scanop(
return ib.get()
- out_buf = decl_tensor(shape, dtype, "out_buf")
+ out_buf = Var("out_buf", TensorType(shape, dtype))
return extern(
[shape],
diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py
index 91b9bda8b2..d38f46090e 100644
--- a/python/tvm/topi/scatter.py
+++ b/python/tvm/topi/scatter.py
@@ -179,7 +179,7 @@ def scatter_nd(data, indices, updates, mode):
return ib.get()
- out_buf = tirx.decl_tensor(data.shape, data.dtype, "out_buf", layout=None)
+ out_buf = tirx.Var("out_buf", tirx.TensorType(data.shape, data.dtype,
layout=None))
return te.extern(
[data.shape],
[data, indices, updates],
diff --git a/python/tvm/topi/scatter_elements.py
b/python/tvm/topi/scatter_elements.py
index 4dcfbac76d..425597dfa8 100644
--- a/python/tvm/topi/scatter_elements.py
+++ b/python/tvm/topi/scatter_elements.py
@@ -182,7 +182,7 @@ def scatter_elements(data, indices, updates, axis=0,
reduction="update"):
"scatter_elements reduction not in [update, add, mul, mean, min,
max]:", reduction
)
- out_buf = tirx.decl_tensor(data.shape, data.dtype, "out_buf", layout=None)
+ out_buf = tirx.Var("out_buf", tirx.TensorType(data.shape, data.dtype,
layout=None))
return te.extern(
[data.shape],
[data, indices, updates],
diff --git a/python/tvm/topi/signal.py b/python/tvm/topi/signal.py
index 2ff84a1b05..9d16c4c22a 100644
--- a/python/tvm/topi/signal.py
+++ b/python/tvm/topi/signal.py
@@ -140,7 +140,7 @@ def stft(
return ib.get()
- output_buf = tirx.decl_tensor(output_shape, data.dtype, "output_buf",
layout=None)
+ output_buf = tirx.Var("output_buf", tirx.TensorType(output_shape,
data.dtype, layout=None))
loop_kind = "vectorize"
if ir.is_prim_var(output_shape[2]): # any_dim
loop_kind = "serial"
diff --git a/python/tvm/topi/sort.py b/python/tvm/topi/sort.py
index af99d445cc..7424b09849 100644
--- a/python/tvm/topi/sort.py
+++ b/python/tvm/topi/sort.py
@@ -49,11 +49,11 @@ def sort(data, axis=-1, is_ascend=1):
Sorted index tensor.
"""
- data_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+ data_buf = tvm_tirx.tirx.Var(
+ "data_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
- out_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "out_buf", data_alignment=8, layout=None
+ out_buf = tvm_tirx.tirx.Var(
+ "out_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
out = te.extern(
data.shape,
@@ -116,15 +116,18 @@ def argsort(data, valid_count=None, axis=-1, is_ascend=1,
dtype="float32"):
tvm_out = tvm_tirx.runtime.tensor(np.zeros(dshape,
dtype=data.dtype.dtype), dev)
f(tvm_data, tvm_out)
"""
- data_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+ data_buf = tvm_tirx.tirx.Var(
+ "data_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
if valid_count is not None:
- valid_count_buf = tvm_tirx.tirx.decl_tensor(
- valid_count.shape, valid_count.dtype, "valid_count_buf",
data_alignment=4, layout=None
+ valid_count_buf = tvm_tirx.tirx.Var(
+ "valid_count_buf",
+ tvm_tirx.tirx.TensorType(
+ valid_count.shape, valid_count.dtype, data_alignment=4,
layout=None
+ ),
)
- out_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, "int32", "out_buf", data_alignment=8, layout=None
+ out_buf = tvm_tirx.tirx.Var(
+ "out_buf", tvm_tirx.tirx.TensorType(data.shape, "int32",
data_alignment=8, layout=None)
)
out = te.extern(
data.shape,
@@ -139,8 +142,11 @@ def argsort(data, valid_count=None, axis=-1, is_ascend=1,
dtype="float32"):
tag="argsort_nms_cpu",
)
else:
- out_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, dtype, "out_buf", data_alignment=8, layout=None
+ out_buf = tvm_tirx.tirx.Var(
+ "out_buf",
+ tvm_tirx.tirx.TensorType(
+ data.shape, "float32" if dtype is None else dtype,
data_alignment=8, layout=None
+ ),
)
out = te.extern(
data.shape,
@@ -189,8 +195,8 @@ def topk(data, k=1, axis=-1, ret_type="both",
is_ascend=False, dtype="int64"):
The computed result.
"""
assert ret_type in ["both", "values", "indices"]
- data_buf = tvm_tirx.tirx.decl_tensor(
- data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+ data_buf = tvm_tirx.tirx.Var(
+ "data_buf", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
data_alignment=8, layout=None)
)
out_shape = list(get_const_tuple(data.shape))
kvar = tvm.te.var("k")
@@ -201,14 +207,18 @@ def topk(data, k=1, axis=-1, ret_type="both",
is_ascend=False, dtype="int64"):
out_bufs = []
if ret_type in ["both", "values"]:
out_bufs.append(
- tvm_tirx.tirx.decl_tensor(
- out_shape, data.dtype, "value_buf", data_alignment=8,
layout=None
+ tvm_tirx.tirx.Var(
+ "value_buf",
+ tvm_tirx.tirx.TensorType(out_shape, data.dtype,
data_alignment=8, layout=None),
)
)
if ret_type in ["both", "indices"]:
out_bufs.append(
- tvm_tirx.tirx.decl_tensor(
- out_shape, dtype, "indices_buf", data_alignment=8, layout=None
+ tvm_tirx.tirx.Var(
+ "indices_buf",
+ tvm_tirx.tirx.TensorType(
+ out_shape, "float32" if dtype is None else dtype,
data_alignment=8, layout=None
+ ),
)
)
out_shapes = [out_shape] * len(out_bufs)
diff --git a/python/tvm/topi/sparse_reshape.py
b/python/tvm/topi/sparse_reshape.py
index e35e174ba8..5542af272b 100644
--- a/python/tvm/topi/sparse_reshape.py
+++ b/python/tvm/topi/sparse_reshape.py
@@ -19,7 +19,7 @@
from tvm.te import div, extern, floordiv, floormod
from tvm_tirx.script.ir_builder import IRBuilder
-from tvm_tirx.tirx import Cast, decl_tensor
+from tvm_tirx.tirx import Cast, TensorType, Var
from tvm_tirx.tirx.script import ir_builder as T
@@ -278,10 +278,10 @@ def sparse_reshape(
return ib.get()
- new_sparse_indices_buf = decl_tensor(
- new_sparse_indices_shape, sparse_indices.dtype,
"new_sparse_indices_buf"
+ new_sparse_indices_buf = Var(
+ "new_sparse_indices_buf", TensorType(new_sparse_indices_shape,
sparse_indices.dtype)
)
- new_shape_buf = decl_tensor(new_shape_shape, prev_shape.dtype,
"new_shape_buf")
+ new_shape_buf = Var("new_shape_buf", TensorType(new_shape_shape,
prev_shape.dtype))
return extern(
[new_sparse_indices_shape, new_shape_shape],
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
index 3c817d1b8d..6d0dc136cc 100644
--- a/python/tvm/topi/vision/nms.py
+++ b/python/tvm/topi/vision/nms.py
@@ -125,17 +125,23 @@ def get_valid_counts(data, score_threshold=0, id_index=0,
score_index=1):
id_index_const = tvm_tirx.tirx.const(id_index, "int32")
score_index_const = tvm_tirx.tirx.const(score_index, "int32")
- valid_count_buf = tvm_tirx.tirx.decl_tensor((batch_size,), "int32",
"valid_count", layout=None)
- out_tensor_buf = tvm_tirx.tirx.decl_tensor(
- (batch_size, num_anchors, box_data_length), data.dtype, "out_tensor",
layout=None
+ valid_count_buf = tvm_tirx.tirx.Var(
+ "valid_count", tvm_tirx.tirx.TensorType((batch_size,), "int32",
layout=None)
)
- out_indices_buf = tvm_tirx.tirx.decl_tensor(
- (batch_size, num_anchors), "int32", "out_indices", layout=None
+ out_tensor_buf = tvm_tirx.tirx.Var(
+ "out_tensor",
+ tvm_tirx.tirx.TensorType(
+ (batch_size, num_anchors, box_data_length), data.dtype, layout=None
+ ),
+ )
+ out_indices_buf = tvm_tirx.tirx.Var(
+ "out_indices", tvm_tirx.tirx.TensorType((batch_size, num_anchors),
"int32", layout=None)
)
if is_score_threshold_tensor:
- score_thresh_buf = tvm_tirx.tirx.decl_tensor(
- score_threshold.shape, score_threshold.dtype, "score_threshold",
layout=None
+ score_thresh_buf = tvm_tirx.tirx.Var(
+ "score_threshold",
+ tvm_tirx.tirx.TensorType(score_threshold.shape,
score_threshold.dtype, layout=None),
)
valid_count, out_tensor, out_indices = te.extern(
[(batch_size,), (batch_size, num_anchors, box_data_length),
(batch_size, num_anchors)],
@@ -152,7 +158,9 @@ def get_valid_counts(data, score_threshold=0, id_index=0,
score_index=1):
dtype=["int32", data.dtype, "int32"],
out_tensors=[valid_count_buf, out_tensor_buf, out_indices_buf],
in_tensors=[
- tvm_tirx.tirx.decl_tensor(data.shape, data.dtype, "data",
layout=None),
+ tvm_tirx.tirx.Var(
+ "data", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
layout=None)
+ ),
score_thresh_buf,
],
name="get_valid_counts",
@@ -177,7 +185,11 @@ def get_valid_counts(data, score_threshold=0, id_index=0,
score_index=1):
_ir_with_const_threshold,
dtype=["int32", data.dtype, "int32"],
out_tensors=[valid_count_buf, out_tensor_buf, out_indices_buf],
- in_tensors=[tvm_tirx.tirx.decl_tensor(data.shape, data.dtype,
"data", layout=None)],
+ in_tensors=[
+ tvm_tirx.tirx.Var(
+ "data", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
layout=None)
+ )
+ ],
name="get_valid_counts",
tag="get_valid_counts",
)
@@ -757,23 +769,29 @@ def non_max_suppression(
)
sort_tensor = argsort(score_tensor, valid_count=valid_count, axis=1,
is_ascend=False)
- data_buf = tvm_tirx.tirx.decl_tensor(data.shape, data.dtype, "data",
layout=None)
- sort_buf = tvm_tirx.tirx.decl_tensor(
- sort_tensor.shape, sort_tensor.dtype, "sorted_index", layout=None
+ data_buf = tvm_tirx.tirx.Var(
+ "data", tvm_tirx.tirx.TensorType(data.shape, data.dtype, layout=None)
)
- valid_count_buf = tvm_tirx.tirx.decl_tensor(
- valid_count.shape, valid_count.dtype, "valid_count", layout=None
+ sort_buf = tvm_tirx.tirx.Var(
+ "sorted_index", tvm_tirx.tirx.TensorType(sort_tensor.shape,
sort_tensor.dtype, layout=None)
+ )
+ valid_count_buf = tvm_tirx.tirx.Var(
+ "valid_count", tvm_tirx.tirx.TensorType(valid_count.shape,
valid_count.dtype, layout=None)
+ )
+ indices_buf = tvm_tirx.tirx.Var(
+ "indices", tvm_tirx.tirx.TensorType(indices.shape, indices.dtype,
layout=None)
)
- indices_buf = tvm_tirx.tirx.decl_tensor(indices.shape, indices.dtype,
"indices", layout=None)
- out_data_buf = tvm_tirx.tirx.decl_tensor(data.shape, data.dtype,
"out_data", layout=None)
- out_box_indices_buf = tvm_tirx.tirx.decl_tensor(
- (batch_size, num_anchors), "int32", "out_box_indices", layout=None
+ out_data_buf = tvm_tirx.tirx.Var(
+ "out_data", tvm_tirx.tirx.TensorType(data.shape, data.dtype,
layout=None)
+ )
+ out_box_indices_buf = tvm_tirx.tirx.Var(
+ "out_box_indices", tvm_tirx.tirx.TensorType((batch_size, num_anchors),
"int32", layout=None)
)
if return_indices:
- out_valid_box_count_buf = tvm_tirx.tirx.decl_tensor(
- (batch_size, 1), "int32", "out_valid_box_count", layout=None
+ out_valid_box_count_buf = tvm_tirx.tirx.Var(
+ "out_valid_box_count", tvm_tirx.tirx.TensorType((batch_size, 1),
"int32", layout=None)
)
out_data, out_box_indices, out_valid_box_count = te.extern(
@@ -852,8 +870,11 @@ def non_max_suppression(
def _rearrange_out(data, batch_size, num_anchors, box_data_length,
score_index):
"""Move valid boxes (score >= 0) to the top of output."""
- out_buf = tvm_tirx.tirx.decl_tensor(
- (batch_size, num_anchors, box_data_length), data.dtype, "rearranged",
layout=None
+ out_buf = tvm_tirx.tirx.Var(
+ "rearranged",
+ tvm_tirx.tirx.TensorType(
+ (batch_size, num_anchors, box_data_length), data.dtype, layout=None
+ ),
)
def _rearrange_ir(ins, outs):
@@ -1014,20 +1035,21 @@ def _get_valid_box_count(scores, score_threshold):
return ib.get()
- scores_buf = tvm_tirx.tirx.decl_tensor(
- scores.shape, scores.dtype, "scores_buf", data_alignment=8, layout=None
+ scores_buf = tvm_tirx.tirx.Var(
+ "scores_buf",
+ tvm_tirx.tirx.TensorType(scores.shape, scores.dtype, data_alignment=8,
layout=None),
)
- searchsorted_buf = tvm_tirx.tirx.decl_tensor(
- (batch_classes,), "int32", "searchsorted", data_alignment=8,
layout=None
+ searchsorted_buf = tvm_tirx.tirx.Var(
+ "searchsorted",
+ tvm_tirx.tirx.TensorType((batch_classes,), "int32", data_alignment=8,
layout=None),
)
if hasattr(score_threshold, "shape"):
- score_thresh_buf = tvm_tirx.tirx.decl_tensor(
- score_threshold.shape,
- score_threshold.dtype,
+ score_thresh_buf = tvm_tirx.tirx.Var(
"score_thresh_buf",
- data_alignment=8,
- layout=None,
+ tvm_tirx.tirx.TensorType(
+ score_threshold.shape, score_threshold.dtype,
data_alignment=8, layout=None
+ ),
)
return te.extern(
[(batch_classes,)],
diff --git a/python/tvm/topi/vision/nms_util.py
b/python/tvm/topi/vision/nms_util.py
index a574c99620..26fc0b32f0 100644
--- a/python/tvm/topi/vision/nms_util.py
+++ b/python/tvm/topi/vision/nms_util.py
@@ -424,11 +424,15 @@ def run_all_class_nms(
num_class = batch_class // batch
if return_scores is False:
- all_class_num0_buf = tvm_tirx.tirx.decl_tensor(
- (batch_class, num_boxes), "int32", "all_class_nms0",
data_alignment=8, layout=None
+ all_class_num0_buf = tvm_tirx.tirx.Var(
+ "all_class_nms0",
+ tvm_tirx.tirx.TensorType(
+ (batch_class, num_boxes), "int32", data_alignment=8,
layout=None
+ ),
)
- all_class_num1_buf = tvm_tirx.tirx.decl_tensor(
- (batch_class,), "int32", "all_class_nms1", data_alignment=8,
layout=None
+ all_class_num1_buf = tvm_tirx.tirx.Var(
+ "all_class_nms1",
+ tvm_tirx.tirx.TensorType((batch_class,), "int32",
data_alignment=8, layout=None),
)
extern_inputs = [boxes, sorted_scores, sorted_indices, valid_count]
if score_threshold is not None:
diff --git a/tests/python/contrib/test_sort.py
b/tests/python/contrib/test_sort.py
index 8f5364aef4..04535b335a 100644
--- a/tests/python/contrib/test_sort.py
+++ b/tests/python/contrib/test_sort.py
@@ -23,7 +23,7 @@ import numpy as np
import tvm
import tvm_tirx
import tvm_tirx.testing
-from tvm import te
+from tvm import te, topi
def test_sort():
@@ -98,6 +98,14 @@ def test_sort_np():
tvm_tirx.testing.assert_allclose(c.numpy(), np_out, rtol=1e-5)
+def test_sort_none_index_dtype_defaults_to_float32():
+ data = te.placeholder((8,), dtype="int32")
+ indices = topi.argsort(data, dtype=None)
+ top_indices = topi.topk(data, k=2, ret_type="indices", dtype=None)
+ assert indices.dtype == "float32"
+ assert top_indices.dtype == "float32"
+
+
if __name__ == "__main__":
test_sort()
test_sort_np()
diff --git
a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
index dda027f75e..be58e2b90c 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
@@ -29,7 +29,7 @@ def test_tensor_region_bounds_are_visited():
data = tvm_tirx.tirx.Var(
"data", tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("int32"),
scope="global")
)
- tensor = tvm_tirx.tirx.decl_tensor([4], "int32", data=data)
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([4],
"int32", data=data))
undefined = tvm_tirx.tirx.Var("undefined", "int32")
region = TensorRegion(
tensor, [tvm_tirx.ir.Range.from_min_extent(undefined, 4)],
ty=TensorRegionType()
@@ -121,13 +121,13 @@ def test_error_undeclared_tensor_in_schedulable_tir():
# Manually construct a TensorStore that uses a tensor without any
declaration
# inside a block context.
n = tvm_tirx.tirx.Var("n", "int32")
- A = tvm_tirx.tirx.decl_tensor([n], "float32", name="A")
+ A = tvm_tirx.tirx.Var("A", tvm_tirx.tirx.TensorType([n], "float32"))
i = tvm_tirx.tirx.Var("i", "int32")
# Create an undeclared tensor using an explicit data pointer that is NOT
# a function parameter and NOT wrapped with DeclTensor.
B_data = tvm_tirx.tirx.Var("B_data",
tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32")))
- B = tvm_tirx.tirx.decl_tensor([n], "float32", name="B", data=B_data)
+ B = tvm_tirx.tirx.Var("B", tvm_tirx.tirx.TensorType([n], "float32",
data=B_data))
# Build a block that writes to B without any declaration of B.
bi = tvm_tirx.tirx.Var("bi", "int32")
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 2f3b0d4015..764d1d1da2 100644
--- a/tests/python/s_tir/analysis/test_sblock_access_region.py
+++ b/tests/python/s_tir/analysis/test_sblock_access_region.py
@@ -501,8 +501,8 @@ def test_conditional_inequality_access_regions(case):
"rounded": ([x], [(-20, 41)], tirx.all(x * 3 >= -7, x * 2 <= 9), [(-2,
7)]),
}
variables, domains, condition, expected = cases[case]
- inside = tirx.decl_tensor([256] * len(variables), name="inside")
- outside = tirx.decl_tensor([256] * len(variables), name="outside")
+ inside = tirx.Var("inside", tirx.TensorType([256] * len(variables),
"float32"))
+ outside = tirx.Var("outside", tirx.TensorType([256] * len(variables),
"float32"))
body = tvm_tirx.ir.SeqStmt(
[
tvm_tirx.ir.If(condition,
tvm_tirx.ir.Evaluate(inside[tuple(variables)]), None),
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
index 2cbb44826f..68ea9a4d63 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
@@ -38,7 +38,7 @@ from tvm.te import create_function
from tvm_tirx.ir import Evaluate, For, ForKind
from tvm_tirx.ir.prim import expr_deep_equal
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import IndexMap, Var, decl_tensor, floordiv, floormod
+from tvm_tirx.tirx import IndexMap, TensorType, Var, floordiv, floormod
def _make_vars(*args: str) -> list[Var]:
@@ -62,7 +62,7 @@ def _make_loops(loop_vars: list[Var], extents: list[int]) ->
list[For]:
def test_suggest_index_map_simple():
i, j = _make_vars("i", "j")
index_map = suggest_index_map(
- tensor=decl_tensor(shape=[8, 256]),
+ tensor=Var("tensor", TensorType([8, 256], "float32")),
indices=[
floordiv(i, 16) * 4 + floordiv(j, 16),
floormod(i, 16) * 16 + floormod(j, 16),
@@ -87,7 +87,7 @@ def test_suggest_index_map_simple():
def test_suggest_index_map_bijective():
i, j = _make_vars("i", "j")
index_map = suggest_index_map(
- tensor=decl_tensor(shape=[8]),
+ tensor=Var("tensor", TensorType([8], "float32")),
indices=[floormod(j, 4) * 2 + i],
loops=_make_loops(
loop_vars=[i, j],
@@ -111,7 +111,7 @@ def test_suggest_index_map_winograd():
nu = floordiv(floormod(fused_outer, 336), 112) * 2 +
floordiv(floormod(fused_outer, 8), 4)
co = floormod(fused_outer, 4) * 32 + i3_3_fused
ci = (i4_0 * 32) + i4_1
- tensor = decl_tensor(shape=[6, 6, 128, 128])
+ tensor = Var("tensor", TensorType([6, 6, 128, 128], "float32"))
index_map = suggest_index_map(
tensor=tensor,
indices=[eps, nu, co, ci],
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
index 6b22646383..54a485b932 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
@@ -1041,7 +1041,7 @@ def test_pad_value_may_not_reference_other_tensor():
sch = tvm.s_tir.Schedule(Before)
A = sch.get(sch.get_sblock("block")).reads[0].source
- other = tirx.decl_tensor(1, A.ty.dtype, name="other")
+ other = tirx.Var("other", tirx.TensorType(1, A.ty.dtype))
with pytest.raises(tvm.s_tir.schedule.schedule.ScheduleError):
sch.transform_layout(
"block",
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 b9a868dce4..4859a11264 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,8 +242,10 @@ def test_inferred_ty_with_tensor_args():
expected = tvm_tirx.ir.FuncType(
[
- tvm_tirx.tirx.decl_tensor([16, 16], "float32", layout=None).ty,
- tvm_tirx.tirx.decl_tensor([256], "int32", layout=None).ty,
+ tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType([16, 16], "float32",
layout=None)
+ ).ty,
+ tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([256],
"int32", layout=None)).ty,
],
tvm_tirx.ir.PrimType("float32"),
)
@@ -264,7 +266,9 @@ def test_inferred_ty_with_internal_allocation():
expected = tvm_tirx.ir.FuncType(
[
- tvm_tirx.tirx.decl_tensor([16, 16], "float32", layout=None).ty,
+ tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType([16, 16], "float32",
layout=None)
+ ).ty,
],
tvm_tirx.ir.PrimType("float32"),
)
@@ -281,8 +285,8 @@ def test_inferred_ty_with_output_tensor():
expected = tvm_tirx.ir.FuncType(
[
- tvm_tirx.tirx.decl_tensor([16], "float32", layout=None).ty,
- tvm_tirx.tirx.decl_tensor([16], "float32", layout=None).ty,
+ tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([16],
"float32", layout=None)).ty,
+ tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([16],
"float32", layout=None)).ty,
],
tvm_tirx.ir.TupleType([]),
)
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 09f3e47847..8dd7dd2572 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
@@ -61,8 +61,12 @@ def test_inferred_ty_with_dynamic_tensor():
N = tvm_tirx.tirx.Var("N", "int64")
expected = tvm_tirx.ir.FuncType(
[
- tvm_tirx.tirx.decl_tensor([M, N], "float32", layout=None).ty,
- tvm_tirx.tirx.decl_tensor([M * N], "float32", layout=None).ty,
+ tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType([M, N], "float32",
layout=None)
+ ).ty,
+ tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType([M * N], "float32",
layout=None)
+ ).ty,
],
tvm_tirx.ir.TupleType([]),
)
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 8a72782e9a..ea8c3149c6 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
@@ -93,9 +93,9 @@ def test_ir_builder_tir_function_complete():
function_actual = ib.get()
# the expected function
- c_tensor = tirx.decl_tensor((128, 128), "float32", name="c", layout=None)
- d_tensor = tirx.decl_tensor((64, 64), "int64", name="d", layout=None)
- e_tensor = tirx.decl_tensor((1024,), "int8", name="e", layout=None)
+ c_tensor = tirx.Var("c", tirx.TensorType((128, 128), "float32",
layout=None))
+ d_tensor = tirx.Var("d", tirx.TensorType((64, 64), "int64", layout=None))
+ e_tensor = tirx.Var("e", tirx.TensorType((1024,), "int8", layout=None))
function_expected = tirx.Function(
params=[
tirx.Var("a", tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("void"))),
@@ -167,10 +167,10 @@ def test_ir_builder_tir_block_complete():
# the expected block
var_a = tirx.Var("a", "int64")
- tensor_b = tirx.decl_tensor((128, 128), "float32", name="b")
- tensor_c = tirx.decl_tensor((128, 128), "float32", name="c")
+ tensor_b = tirx.Var("b", tirx.TensorType((128, 128), "float32"))
+ tensor_c = tirx.Var("c", tirx.TensorType((128, 128), "float32"))
var_d = tirx.Var("d", "int32")
- tensor_e = tirx.decl_tensor((128, 128), "float32", name="c")
+ tensor_e = tirx.Var("c", tirx.TensorType((128, 128), "float32"))
var_f = tirx.Var("f", "int32")
block_expected = s_tir.SBlock(
iter_vars=[s_tir.IterVar((0, 128), tirx.Var("", "int32"),
iter_type=s_tir.IterVar.DataPar)],
@@ -178,9 +178,11 @@ def test_ir_builder_tir_block_complete():
writes=[tensor_c[var_d:128, var_d:128]],
name_hint="block",
body=tvm_tirx.ir.Evaluate(0),
- alloc_tensors=[tirx.decl_tensor((128, 128), "float32")],
+ alloc_tensors=[tirx.Var("tensor", tirx.TensorType((128, 128),
"float32"))],
match_tensors=[
- s_tir.MatchTensorRegion(tirx.decl_tensor((32, 32), "float32"),
tensor_e[0:32, 0:32])
+ s_tir.MatchTensorRegion(
+ tirx.Var("tensor", tirx.TensorType((32, 32), "float32")),
tensor_e[0:32, 0:32]
+ )
],
annotations={"key": "value"},
)
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 6dce514f71..84a85d65f9 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
@@ -45,7 +45,7 @@ def _assert_print(obj, expected):
@pytest.mark.skipif(sys.version_info < (3, 12), reason="PEP 695 requires
Python 3.12")
def test_function_symbolic_tensor_param_roundtrip():
n = tirx.Var("n", "int32")
- A = tirx.decl_tensor(shape=[n + 1, n], dtype="float32", name="A",
layout=None)
+ A = tirx.Var("A", tirx.TensorType([n + 1, n], "float32", layout=None))
func = (
tirx.Function(params=[A], body=tvm_tirx.ir.Evaluate(n))
.with_attr("global_symbol", "main")
@@ -67,7 +67,7 @@ def test_function_symbolic_tensor_param_roundtrip():
@pytest.mark.skipif(sys.version_info < (3, 12), reason="PEP 695 requires
Python 3.12")
def test_function_compound_tensor_shape_first_use_roundtrip():
n = tirx.Var("n", "int32")
- A = tirx.decl_tensor(shape=[tirx.max(n, 1)], dtype="float32", name="A",
layout=None)
+ A = tirx.Var("A", tirx.TensorType([tirx.max(n, 1)], "float32",
layout=None))
func = (
tirx.Function(params=[A], body=tvm_tirx.ir.Evaluate(n))
.with_attr("global_symbol", "main")
@@ -88,7 +88,7 @@ def test_function_compound_tensor_shape_first_use_roundtrip():
def test_function_symbolic_alloc_tensor_roundtrip():
size = tirx.Var("size", "int32")
- buf = tirx.decl_tensor(shape=[size], dtype="float32", name="buf",
layout=None)
+ buf = tirx.Var("buf", tirx.TensorType([size], "float32", layout=None))
func = tirx.Function(
params=[],
body=tvm_tirx.ir.SeqStmt(
@@ -120,8 +120,8 @@ def test_function_symbolic_alloc_tensor_roundtrip():
def test_function():
- A = tirx.decl_tensor(shape=[128, 128], dtype="float32", name="A")
- B = tirx.decl_tensor(shape=[256, 256], dtype="float32", name="B")
+ A = tirx.Var("A", tirx.TensorType([128, 128], "float32"))
+ B = tirx.Var("B", tirx.TensorType([256, 256], "float32"))
func = (
tirx.Function(
params=[A, B],
@@ -147,8 +147,8 @@ def main(A: T.Tensor((128, 128), T.float32), B:
T.Tensor((256, 256), T.float32))
def test_function_tensor_data_use():
- A = tirx.decl_tensor(shape=[128, 128], dtype="float32", name="A")
- B = tirx.decl_tensor(shape=[256, 256], dtype="float32", name="B")
+ A = tirx.Var("A", tirx.TensorType([128, 128], "float32"))
+ B = tirx.Var("B", tirx.TensorType([256, 256], "float32"))
func = (
tirx.Function(
params=[A, B],
@@ -175,9 +175,9 @@ def main(A: T.Tensor((128, 128), T.float32), B:
T.Tensor((256, 256), T.float32))
def test_function_tensor_data_argument_is_scope_hint():
- buffer_data = tirx.decl_tensor(shape=[128, 128], dtype="float32",
name="A").data
- A = tirx.decl_tensor(shape=[128, 128], dtype="float32", name="A",
data=buffer_data)
- B = tirx.decl_tensor(shape=[256, 256], dtype="float32", name="B",
data=buffer_data)
+ buffer_data = tirx.Var("A", tirx.TensorType([128, 128], "float32")).data
+ A = tirx.Var("A", tirx.TensorType([128, 128], "float32", data=buffer_data))
+ B = tirx.Var("B", tirx.TensorType([256, 256], "float32", data=buffer_data))
func = (
tirx.Function(
params=[A, B],
@@ -258,8 +258,8 @@ with Ts.sblock("block", no_realize=True):
def test_match_tensor_region():
- src = tirx.decl_tensor((128, 128), "float32", name="src")
- tgt = tirx.decl_tensor((64, 64), "float32", name="tgt")
+ src = tirx.Var("src", tirx.TensorType((128, 128), "float32"))
+ tgt = tirx.Var("tgt", tirx.TensorType((64, 64), "float32"))
obj = s_tir.MatchTensorRegion(
tgt,
TensorRegion(
@@ -395,8 +395,8 @@ def main():
def test_private_function():
- A = tirx.decl_tensor(shape=[128, 128], dtype="float32", name="A")
- B = tirx.decl_tensor(shape=[256, 256], dtype="float32", name="B")
+ A = tirx.Var("A", tirx.TensorType([128, 128], "float32"))
+ B = tirx.Var("B", tirx.TensorType([256, 256], "float32"))
func = tirx.Function(
params=[A, B],
ret_type=None,
@@ -543,8 +543,8 @@ def test_predicated_tensor_load_store():
a = tirx.Var("a", "handle")
b = tirx.Var("b", "handle")
tensors = {
- a: tirx.decl_tensor(shape=[128, 128], dtype="float32", name="A"),
- b: tirx.decl_tensor(shape=[256, 256], dtype="float32", name="B"),
+ a: tirx.Var("A", tirx.TensorType([128, 128], "float32")),
+ b: tirx.Var("B", tirx.TensorType([256, 256], "float32")),
}
tensor_load = tirx.call_intrin(
"tirx.masked_load",
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 de1e00b415..65f1de775f 100644
--- a/tests/python/s_tir/test_s_tir_renew_defs.py
+++ b/tests/python/s_tir/test_s_tir_renew_defs.py
@@ -194,7 +194,9 @@ def test_tensor_params():
def test_compound_tensor_param_shape_var():
n = tvm_tirx.tirx.Var("n", "int32")
- A = tvm_tirx.tirx.decl_tensor((tvm_tirx.tirx.max(n, 1),), layout=None)
+ A = tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType((tvm_tirx.tirx.max(n, 1),),
"float32", layout=None)
+ )
f1 = tvm_tirx.tirx.Function([A], tvm_tirx.ir.Evaluate(n))
f2 = tvm_tirx.tirx.renew_def(f1)
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
index a81351920e..a37b502acd 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
@@ -28,7 +28,7 @@ from tvm_tirx.script import tirx as T
@pytest.mark.parametrize("scope,shape", [("mma.matrixA", (32, 8)),
("mma.matrixB", (8, 32))])
@pytest.mark.parametrize("access_kind", ["load", "store"])
def test_explicit_matrix_ab_access_is_rejected(scope, shape, access_kind):
- tensor = tirx.decl_tensor(shape, "float32", scope=scope)
+ tensor = tirx.Var("tensor", tirx.TensorType(shape, "float32", scope=scope))
if access_kind == "load":
body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(tensor, [0, 0]))
else:
diff --git a/tests/python/te/test_te_tensor.py
b/tests/python/te/test_te_tensor.py
index 21071eca89..1f401258ad 100644
--- a/tests/python/te/test_te_tensor.py
+++ b/tests/python/te/test_te_tensor.py
@@ -184,6 +184,12 @@ def test_extern():
assert tuple(B.shape) == (m,)
+def test_extern_none_output_dtype_defaults_to_float32():
+ source = te.placeholder((8,), dtype="int32")
+ result = te.extern((8,), [source], lambda ins, outs:
tvm_tirx.ir.Evaluate(0), dtype=[None])
+ assert result.dtype == "float32"
+
+
def test_extern_multi_out():
m = te.var("m")
A = te.placeholder((m,), name="A")
diff --git a/tests/python/tirx/base/test_tirx_constructor.py
b/tests/python/tirx/base/test_tirx_constructor.py
index 5aaa6059f5..32d231571c 100644
--- a/tests/python/tirx/base/test_tirx_constructor.py
+++ b/tests/python/tirx/base/test_tirx_constructor.py
@@ -109,7 +109,7 @@ def test_expr_constructor():
assert x.condition == a
tensor_var = tvm_tirx.tirx.Var("buf",
tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32")))
- tensor = tvm_tirx.tirx.decl_tensor([16], "float32", data=tensor_var)
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([16],
"float32", data=tensor_var))
x = tvm_tirx.ir.TensorLoad(tensor, [1])
assert isinstance(x, tvm_tirx.ir.TensorLoad)
assert x.ty == tvm_tirx.ir.PrimType("float32")
@@ -235,7 +235,7 @@ def test_expr_constructor():
def test_tensor_region_call_wrappers_reject():
- tensor = tvm_tirx.tirx.decl_tensor([4], "int32")
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([4],
"int32"))
region = tensor[0:4]
calls = [
lambda: tvm_tirx.tirx.call_intrin("tirx.reinterpret", region,
ty="int32"),
@@ -255,22 +255,22 @@ def test_tensor_region_call_wrappers_reject():
def test_tensor_region_type_is_singleton():
- lhs = tvm_tirx.tirx.decl_tensor([1], "int32")[0:1]
- rhs = tvm_tirx.tirx.decl_tensor([2], "float32")[0:2]
+ lhs = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([1],
"int32"))[0:1]
+ rhs = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([2],
"float32"))[0:2]
assert isinstance(lhs, TensorRegion)
assert isinstance(rhs, TensorRegion)
assert lhs.ty.same_as(rhs.ty)
def test_tensor_region_is_not_arithmetic_operand():
- int_region = tvm_tirx.tirx.decl_tensor([4], "int32")[0:4]
+ int_region = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([4],
"int32"))[0:4]
with pytest.raises(TypeError, match="construct a TensorLoad explicitly"):
tvm.s_tir.IterVar((0, 4), "i", tvm.s_tir.IterVar.DataPar) + int_region
def test_operator_base_categories_have_primitive_type():
var = tvm_tirx.tirx.Var("x", "int32")
- tensor = tvm_tirx.tirx.decl_tensor([4], "float32")
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([4],
"float32"))
lhs, rhs = tvm_tirx.tirx.Var("lhs", "int32"), tvm_tirx.tirx.Var("rhs",
"int32")
combiner = tvm.te.CommReducer([lhs], [rhs], [lhs + rhs],
[tvm_tirx.tirx.const(0, "int32")])
expressions = [
@@ -324,7 +324,7 @@ def test_stmt_constructor():
assert len(x.body) == 1 and x.body[0] == nop
tensor_var = tvm_tirx.tirx.Var("buf",
tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("bool")))
- tensor = tvm_tirx.tirx.decl_tensor([16], "bool", data=tensor_var)
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([16],
"bool", data=tensor_var))
x = tvm_tirx.ir.TensorStore(tensor, [10], tvm_tirx.ir.IntImm("bool", 1))
assert isinstance(x, tvm_tirx.ir.TensorStore)
assert x.dest == tensor
@@ -333,7 +333,7 @@ def test_stmt_constructor():
assert list(x.indices) == [10]
assert x.value.value == 1
- buf = tvm_tirx.tirx.decl_tensor([10], "float32")
+ buf = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([10],
"float32"))
x = tvm_tirx.ir.Bind(
buf,
tvm_tirx.ir.Call(
diff --git a/tests/python/tirx/script/test_tirx_script_ir_builder.py
b/tests/python/tirx/script/test_tirx_script_ir_builder.py
index 3ad62e1fe9..053f698712 100644
--- a/tests/python/tirx/script/test_tirx_script_ir_builder.py
+++ b/tests/python/tirx/script/test_tirx_script_ir_builder.py
@@ -268,13 +268,13 @@ def test_ir_builder_tir_inline():
],
)
def test_concrete_tensor_indices(shape, index, expected):
- tensor = tirx.decl_tensor(shape, "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType(shape, "float32"))
assert T.tensor_indices(tensor, index) == expected
def test_symbolic_tensor_indices():
m, n, k = [tirx.Var(name, "int32") for name in ("m", "n", "k")]
- tensor = tirx.decl_tensor((m, n), "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType((m, n), "float32"))
actual = T.tensor_indices(tensor, k)
for index, expected in zip(actual, [k // n, k % n]):
tvm_tirx.ir.assert_structural_equal(index, expected)
@@ -282,7 +282,9 @@ def test_symbolic_tensor_indices():
@pytest.mark.parametrize("layout", [None, TileLayout(S[(2, 3) : (1, 2)])])
def test_flat_tensor_store_preserves_identity_and_emits_once(layout):
- tensor = tirx.decl_tensor((2, 3), "float32", strides=(5, 1),
elem_offset=2, layout=layout)
+ tensor = tirx.Var(
+ "tensor", tirx.TensorType((2, 3), "float32", strides=(5, 1),
elem_offset=2, layout=layout)
+ )
indices = T.tensor_indices(tensor, 4)
load = tensor[indices]
assert load.source.same_as(tensor)
@@ -411,7 +413,10 @@ def test_gpu_imperative_tensors(kind):
elif kind == "scan":
from tvm.topi.gpu.scan import exclusive_scan_ir
- body = exclusive_scan_ir(tirx.decl_tensor((2, 4)),
tirx.decl_tensor((2, 4)))
+ body = exclusive_scan_ir(
+ tirx.Var("tensor", tirx.TensorType((2, 4), "float32")),
+ tirx.Var("tensor", tirx.TensorType((2, 4), "float32")),
+ )
assert isinstance(body, tvm_tirx.ir.Stmt)
return
elif kind == "scatter_nd":
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py
b/tests/python/tirx/script/test_tirx_script_printer.py
index c631a944da..aac6946f34 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -152,7 +152,7 @@ def test_reflected_type_recursive_values(monkeypatch):
def test_tensor():
- a = tirx.decl_tensor((128, 128), "float16", name="A")
+ a = tirx.Var("A", tirx.TensorType((128, 128), "float16"))
_assert_print(
a,
"""A = T.Var("A", T.Tensor((128, 128), T.float16))
@@ -165,7 +165,7 @@ def _assert_print(obj, expected):
def test_tensor_region():
- src = tirx.decl_tensor((128, 128), "float32", name="src")
+ src = tirx.Var("src", tirx.TensorType((128, 128), "float32"))
obj = TensorRegion(
src,
[
@@ -184,7 +184,7 @@ src[64:128, 64:128]
def test_tensor_load():
- a = tirx.decl_tensor((128, 128), "float16", name="A")
+ a = tirx.Var("A", tirx.TensorType((128, 128), "float16"))
obj = tvm_tirx.ir.TensorLoad(a, [128, 128])
_assert_print(
obj,
@@ -194,7 +194,7 @@ A[128, 128]""",
def test_tensor_store():
- a = tirx.decl_tensor((128, 128), "float16", name="A")
+ a = tirx.Var("A", tirx.TensorType((128, 128), "float16"))
with IRBuilder() as ib:
TB.tensor_store(a, [128, 128], a[128, 128] + 1)
obj = ib.get()
@@ -1016,8 +1016,8 @@ T.nvshmem.putmem_signal_nbi_block(p, p, 16, p, 1, "set",
0)
def test_printer_nki_namespace():
- A = tir.decl_tensor([1], dtype="float16", name="A")
- B = tir.decl_tensor([1], dtype="float16", name="B")
+ A = tir.Var("A", tir.TensorType([1], "float16"))
+ B = tir.Var("B", tir.TensorType([1], "float16"))
a0 = A[0]
b0 = B[0]
_assert_namespace_print(
diff --git a/tirx/python/tvm_tirx/backend/cuda/op.py
b/tirx/python/tvm_tirx/backend/cuda/op.py
index d21a8952e8..b88aeb16da 100644
--- a/tirx/python/tvm_tirx/backend/cuda/op.py
+++ b/tirx/python/tvm_tirx/backend/cuda/op.py
@@ -24,12 +24,10 @@ from functools import cache
import tvm_ffi
-import tvm_tirx
from tvm_tirx import tirx
from tvm_tirx.ir import Attrs, Call, Op, StringImm, const
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm_tirx.ir.op import _init_op_api, _make_op_api
-from tvm_tirx.ir.type import PrimType, PtrType
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.tirx.op import bitwise_and, call_intrin, ptr_byte_offset
@@ -208,7 +206,6 @@ _WAIT_UNTIL_SCOPE = ("cta", "cluster", "gpu", "sys")
# allocation address out of a mailbox -- are ordinary reads after a barrier,
# with no spin and no predicate, so they are not protocols at all.
_WAIT_UNTIL_SPACE = ("global",)
-_WAIT_UNTIL_DTYPE = ("int32", "uint32", "int64", "uint64")
# Widths, so a requested PTX type can be checked against the word it accesses.
# Which spellings an instruction actually takes is the instruction table's
# business: `add` has no `.b32` form, so a bit-typed arrival is rejected there
@@ -347,24 +344,6 @@ def wait_until(
)
-_MBARRIER_ARRIVE_SEM = ("", "release", "relaxed")
-_MBARRIER_ARRIVE_SCOPE = ("", "cta", "cluster")
-_MBARRIER_ARRIVE_SPACE = ("shared", "shared::cta", "shared::cluster")
-
-
-def _validate_mbarrier_arrive_attrs(sem, scope, space, remote):
- if (sem == "") != (scope == ""):
- raise ValueError("mbarrier.arrive sem and scope must be specified
together")
- if sem not in _MBARRIER_ARRIVE_SEM:
- raise ValueError(f"invalid sem={sem!r}; expected one of
{_MBARRIER_ARRIVE_SEM}")
- if scope not in _MBARRIER_ARRIVE_SCOPE:
- raise ValueError(f"invalid scope={scope!r}; expected one of
{_MBARRIER_ARRIVE_SCOPE}")
- if space not in _MBARRIER_ARRIVE_SPACE:
- raise ValueError(f"invalid space={space!r}; expected one of
{_MBARRIER_ARRIVE_SPACE}")
- if remote is not None and space != "shared::cluster":
- raise ValueError("remote mbarrier.arrive requires
space='shared::cluster'")
-
-
_cp_async_raw = _make_op_api(Op.get("tirx.s_tir.cp_async_raw"), __name__)
@@ -390,15 +369,6 @@ def ptx_cp_async_legacy(
return _cp_async_raw(dst_ptr, 0, src_ptr, 0, cp_size, ty=ty, loc=loc)
-def _is_static_unicast_cta_mask(cta_mask):
- if isinstance(cta_mask, int):
- return cta_mask == 0 or cta_mask & (cta_mask - 1) == 0
- if isinstance(cta_mask, tvm_tirx.ir.IntImm):
- value = int(cta_mask)
- return value == 0 or value & (value - 1) == 0
- return False
-
-
def mov_sreg(bits, reg_name, *, ty=None, loc: Location = UNKNOWN_LOC):
"""TVM intrinsic to tvm instrinsics to fetch PTX pre-defined registers
@@ -539,32 +509,6 @@ def mma_fill(local_size, local_ptr, offset, *, ty):
return call_intrin("tirx.cuda.mma_fill", local_size, local_ptr, offset,
ty=ty)
-_PTX_TO_NUMPY_DTYPE = {
- "fp16": "float16",
- "fp32": "float32",
- "fp64": "float64",
- "bf16": "bfloat16",
- "tf32": "float32",
- "s8": "int8",
- "u8": "uint8",
- "s32": "int32",
- "s4": "int4",
- "u4": "uint4",
- "b1": "int1",
- "b16": "uint16",
- "e4m3": "float8_e4m3fn",
- "e5m2": "float8_e5m2",
-}
-
-
-def _ptx_to_numpy_dtype(dtype_str):
- """Map a PTX-abbreviation or numpy dtype string to a numpy dtype string
- used by the legacy intrinsic wrappers. Unknown strings pass through
unchanged
- so a caller may also pass an already-numpy dtype."""
- s = dtype_str if isinstance(dtype_str, str) else str(dtype_str)
- return _PTX_TO_NUMPY_DTYPE.get(s, s)
-
-
_legacy_ldmatrix = _make_op_api(Op.get("tirx.ptx_legacy.ldmatrix"), __name__)
@@ -1052,126 +996,6 @@ def sm100_2sm_leader_smem_addr(ptr):
return bitwise_and(globals()["cvta_generic_to_shared"](ptr),
const(0xFEFFFFFF, dtype="uint32"))
-_PTX_CVT_TYPES = {
- "u8",
- "u16",
- "u32",
- "u64",
- "s8",
- "s16",
- "s32",
- "s64",
- "bf16",
- "f16",
- "f32",
- "f64",
- "f16x2",
- "bf16x2",
- "tf32",
- "e4m3x2",
- "e5m2x2",
- "e2m1x2",
- "e2m3x2",
- "e3m2x2",
- "e4m3x4",
- "e5m2x4",
- "e2m1x4",
- "e2m3x4",
- "e3m2x4",
- "ue8m0x2",
- "s2f6x2",
-}
-_PTX_CVT_ROUNDING = {"", "rni", "rzi", "rmi", "rpi", "rn", "rz", "rm", "rp",
"rna", "rs"}
-_PTX_CVT_SCALED = {"", "n2::ue8m0"}
-_PTX_CVT_RETURN_TYPE = {
- "u8": "uint8",
- "s8": "int8",
- "u16": "uint16",
- "s16": "int16",
- "u32": "uint32",
- "s32": "int32",
- "u64": "uint64",
- "s64": "int64",
- "f32": "float32",
- "f64": "float64",
- "f16": "uint16",
- "bf16": "uint16",
- "e4m3x2": "uint16",
- "e5m2x2": "uint16",
- "e2m1x2": "uint8",
- "e2m3x2": "uint16",
- "e3m2x2": "uint16",
- "ue8m0x2": "uint16",
- "s2f6x2": "uint16",
- "tf32": "uint32",
- "f16x2": "uint32",
- "bf16x2": "uint32",
- "e2m1x4": "uint16",
- "e4m3x4": "uint32",
- "e5m2x4": "uint32",
- "e2m3x4": "uint32",
- "e3m2x4": "uint32",
-}
-
-
-_CP_ASYNC_BULK_CACHE_HINT = ("", "evict_last", "evict_first", "evict_normal",
"evict_last_use")
-
-_PTX_CACHE_POLICY = {
- "evict_normal": 0x1000000000000000,
- "evict_first": 0x12F0000000000000,
- "evict_last": 0x14F0000000000000,
-}
-
-
-def _resolve_cache_policy(cache_hint, cache_policy,
choices=_CP_ASYNC_BULK_CACHE_HINT):
- _choice("cache_hint", cache_hint, choices)
- if cache_policy is not None:
- return cache_policy, True
- if cache_hint:
- if cache_hint not in _PTX_CACHE_POLICY:
- raise ValueError(
- f"Unsupported built-in cache policy {cache_hint!r}; pass
cache_policy explicitly"
- )
- return const(_PTX_CACHE_POLICY[cache_hint], dtype="uint64"), True
- return const(0, dtype="uint64"), False
-
-
-def _ptx_vec_len(vec):
- return int(vec[1:]) if vec else 1
-
-
-def _normalize_ptx_ld_dst(dst, vec, op_name):
- if dst is None:
- if vec:
- raise ValueError(f"vec {op_name} requires dst")
- return [], 0
- if isinstance(dst, list | tuple):
- if not vec:
- raise ValueError(f"{op_name} scatter dst requires vec")
- vec_len = _ptx_vec_len(vec)
- if len(dst) != vec_len:
- raise ValueError(f"{op_name} scatter dst length must match {vec}:
got {len(dst)}")
- return list(dst), vec_len
- return [dst], 1
-
-
-def _validate_ptx_address(addr, space, op_name):
- """Validate pointer and raw shared-memory address forms."""
- addr_ty = getattr(addr, "ty", None)
- if isinstance(addr_ty, PtrType):
- return
- if isinstance(addr_ty, PrimType):
- if addr_ty.dtype == "uint32":
- if not str(space).startswith("shared"):
- raise ValueError(f"{op_name} uint32 address requires shared
state space")
- return
- if addr_ty.dtype.startswith(("int", "uint")):
- raise ValueError(
- f"{op_name} integer address must be uint32 in shared state
space, "
- f"got {addr_ty.dtype}"
- )
-
-
def atomic_cas(ptr, old_val, new_val, *, ty=None, loc: Location = UNKNOWN_LOC):
"""TVM intrinsic to call cuda atomic cas instruction
diff --git a/tirx/python/tvm_tirx/script/ir_builder/base.py
b/tirx/python/tvm_tirx/script/ir_builder/base.py
index ddc32545fd..4b9ee703d1 100644
--- a/tirx/python/tvm_tirx/script/ir_builder/base.py
+++ b/tirx/python/tvm_tirx/script/ir_builder/base.py
@@ -474,21 +474,3 @@ def _return_annotation(annotation):
if callable(annotation) and not isinstance(annotation, ir.Expr | ir.Type):
return annotation()
return annotation
-
-
-def annotation_constructor(constructor):
- """Expose a constructor as a real annotation class supporting Python
unions.
-
- Calls construct ordinary native values directly. The class preserves the
- constructor's signature and documentation without adapting its arguments.
- """
- return type(
- constructor.__name__,
- (),
- {
- "__new__": lambda cls, *args, **kwargs: constructor(*args,
**kwargs),
- "__signature__": signature(constructor),
- "__doc__": constructor.__doc__,
- "__module__": constructor.__module__,
- },
- )
diff --git a/tirx/python/tvm_tirx/tirx/__init__.py
b/tirx/python/tvm_tirx/tirx/__init__.py
index ca2fc9604c..6ef8dfcc8f 100644
--- a/tirx/python/tvm_tirx/tirx/__init__.py
+++ b/tirx/python/tvm_tirx/tirx/__init__.py
@@ -31,7 +31,6 @@ from tvm_tirx.ir import const
from .tensor import (
TensorType,
tensor_data_ptr,
- decl_tensor,
is_tensor_var,
)
from .type import TensorMapType
diff --git a/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
index a474e6af5a..bc50dcc2c6 100644
--- a/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
+++ b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
@@ -158,12 +158,6 @@ def finalize_device_passes(): # pylint:
disable=unused-argument
return tvm_tirx.ir.transform.Sequential(device_pass_list)
-def finalize_device_passes_tirx(): # pylint: disable=unused-argument
- """The TIRx finalization passes for TIR backend."""
- device_pass_list = [tirx.transform.LowerIntrin()]
- return tvm_tirx.ir.transform.Sequential(device_pass_list)
-
-
# global map of pre-built pipelines
PIPELINE_MAP = {"default": default_tir_pipeline, "tirx": tirx_pipeline}
diff --git a/tirx/python/tvm_tirx/tirx/function.py
b/tirx/python/tvm_tirx/tirx/function.py
index 724d5c9957..9ce3bd1487 100644
--- a/tirx/python/tvm_tirx/tirx/function.py
+++ b/tirx/python/tvm_tirx/tirx/function.py
@@ -140,7 +140,8 @@ class Function(BaseFunc, Scriptable):
.. code-block:: python
a, _, m, n = mem_copy.params
- func = mem_copy.specialize({a: tirx.decl_tensor((16, 16))})
+ tensor = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
+ func = mem_copy.specialize({a: tensor})
# or
func = mem_copy.specialize({n: 16, m: 16})
diff --git a/tirx/python/tvm_tirx/tirx/tensor.py
b/tirx/python/tvm_tirx/tirx/tensor.py
index 43c476ec30..aeae4b3b70 100644
--- a/tirx/python/tvm_tirx/tirx/tensor.py
+++ b/tirx/python/tvm_tirx/tirx/tensor.py
@@ -108,6 +108,8 @@ class TensorType(Type):
if data is not None:
if not isinstance(data, tvm_tirx.ir.Expr) or not
isinstance(data.ty, PtrType):
raise TypeError("Tensor data must be an Expr with PtrType")
+ if not isinstance(data.ty.element_type, PrimType):
+ raise TypeError("Tensor data must point to a primitive type")
scope = data.ty.scope
if layout is MISSING:
layout = None
@@ -125,7 +127,7 @@ class TensorType(Type):
data_alignment,
offset_factor,
layout,
- UNKNOWN_LOC,
+ loc if isinstance(loc, Location) else UNKNOWN_LOC,
)
at(loc, self)
@@ -186,7 +188,7 @@ class TensorType(Type):
The storage scope associated with this tensor.
"""
_check_tensor_receiver(self, expr)
- return _ffi_api.TensorStorageScope(expr) # type: ignore
+ return self.scope
def is_scalar(self, expr, alloc_or_decl=True):
"""Check if the tensor is a scalar.
@@ -314,49 +316,6 @@ def _check_tensor_property_receiver(ty, expr):
_check_tensor_receiver(ty, expr)
-def decl_tensor(
- shape,
- dtype=None,
- name="tensor",
- data=None,
- strides=None,
- elem_offset=None,
- scope=_SCOPE_GLOBAL,
- data_alignment=-1,
- offset_factor=0,
- loc: Location = UNKNOWN_LOC,
- layout=None,
-):
- # pylint: disable=import-outside-toplevel
- from .layout import S, TileLayout
-
- shape = (shape,) if tvm_tirx.ir.is_prim_expr(shape) or isinstance(shape,
Integral) else shape
- dtype = "float32" if dtype is None else dtype
- strides = () if strides is None else strides
-
- if layout == "default":
- layout = TileLayout(S[tuple(shape)]) if shape else None
-
- if data is not None:
- if not isinstance(data, tvm_tirx.ir.Expr) or not isinstance(data.ty,
PtrType):
- raise TypeError("Tensor data must be an Expr with PtrType")
- if not isinstance(data.ty.element_type, PrimType):
- raise TypeError("Tensor data must point to a primitive type")
- scope = data.ty.scope
- tensor_type = _ffi_api.TensorType( # type: ignore
- shape,
- dtype,
- scope,
- strides,
- elem_offset,
- data_alignment,
- offset_factor,
- layout,
- loc,
- )
- return Var(name, tensor_type, loc=loc)
-
-
def tensor_data_ptr(tensor, *, ty=None, loc: Location = UNKNOWN_LOC):
"""Project a tensor variable's physical pointer.
diff --git a/tirx/src/tirx/ir/specialize.cc b/tirx/src/tirx/ir/specialize.cc
index 8e94586a3f..169395c8c6 100644
--- a/tirx/src/tirx/ir/specialize.cc
+++ b/tirx/src/tirx/ir/specialize.cc
@@ -352,7 +352,7 @@ class FunctionSpecializer : public StmtExprMutator {
* For example, we define a tensor in Function:
* A: T.Tensor([m, n])
*
- * Then we match it with a tensor B = tirx.decl_tensor((8, 16))
+ * Then we match it with a tensor B = Var("tensor", TensorType((8, 16),
"float32"))
*
* It means we have two var mappings here: m = 8 and n = 16
*
diff --git a/tirx/src/tirx/ir/tensor.cc b/tirx/src/tirx/ir/tensor.cc
index c1a8fe13f4..28a7b95e1a 100644
--- a/tirx/src/tirx/ir/tensor.cc
+++ b/tirx/src/tirx/ir/tensor.cc
@@ -443,9 +443,7 @@ bool TensorVar::IsScalar(bool alloc_or_decl) const { return
ty()->IsScalar(alloc
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def_method("tirx.TensorStorageScope", &TensorVar::scope)
- .def_method("tirx.TensorIsScalar", &TensorVar::IsScalar);
+ refl::GlobalDef().def_method("tirx.TensorIsScalar", &TensorVar::IsScalar);
}
} // namespace tirx
diff --git a/tirx/tests/python/codegen/c/test_basic.py
b/tirx/tests/python/codegen/c/test_basic.py
index 9dd738f428..52d5bde2df 100644
--- a/tirx/tests/python/codegen/c/test_basic.py
+++ b/tirx/tests/python/codegen/c/test_basic.py
@@ -257,7 +257,7 @@ def test_local_alloc_buffer_uses_plain_c_pointer():
def test_vector_pointer_uses_byte_offset():
- buffer = tvm_tirx.tirx.decl_tensor((8,), "float32x2", name="A")
+ buffer = tvm_tirx.tirx.Var("A", tvm_tirx.tirx.TensorType((8,),
"float32x2"))
pointer = tvm_tirx.tirx.ptr_byte_offset(buffer.data, 16, ty=buffer.data.ty)
body = tvm_tirx.ir.Evaluate(tvm_tirx.tirx.call_extern("consume", pointer,
ty="void"))
func = tvm_tirx.tirx.Function([buffer], body).with_attr("global_symbol",
"main")
diff --git a/tirx/tests/python/sym/test_deduce_bound.py
b/tirx/tests/python/sym/test_deduce_bound.py
index ec41000dc9..574c21a46b 100644
--- a/tirx/tests/python/sym/test_deduce_bound.py
+++ b/tirx/tests/python/sym/test_deduce_bound.py
@@ -19,7 +19,7 @@ import pytest
import tvm_tirx
import tvm_tirx.testing
-from tvm_tirx.tirx.tensor import decl_tensor
+from tvm_tirx.tirx import TensorType, Var
def test_deduce():
@@ -236,7 +236,7 @@ def test_deduce_non_support():
test_non_support(tvm_tirx.tirx.EQ(a, 16))
test_non_support(tvm_tirx.tirx.NE(a, 16))
test_non_support(tvm_tirx.tirx.log(a))
- test_non_support(tvm_tirx.ir.TensorLoad(decl_tensor([16], "int32"), [a]))
+ test_non_support(tvm_tirx.ir.TensorLoad(Var("tensor", TensorType([16],
"int32")), [a]))
def test_deduce_floordiv():
diff --git a/tirx/tests/python/sym/test_iter_affine_map.py
b/tirx/tests/python/sym/test_iter_affine_map.py
index 03b675bd44..84ded92be0 100644
--- a/tirx/tests/python/sym/test_iter_affine_map.py
+++ b/tirx/tests/python/sym/test_iter_affine_map.py
@@ -1496,7 +1496,7 @@ def test_detect_iter_map_with_bufferload_recursion():
i = tvm_tirx.tirx.Var("i", "int32")
j = tvm_tirx.tirx.Var("j", "int32")
- buffer = tvm_tirx.tirx.decl_tensor((n,), "int32", name="seqlen")
+ buffer = tvm_tirx.tirx.Var("seqlen", tvm_tirx.tirx.TensorType((n,),
"int32"))
indices = [(buffer[i] + j) // divisor]
iter_vars = {
diff --git a/tirx/tests/python/sym/test_rewrite_simplify.py
b/tirx/tests/python/sym/test_rewrite_simplify.py
index 4a6e1b80d4..326f960f0b 100644
--- a/tirx/tests/python/sym/test_rewrite_simplify.py
+++ b/tirx/tests/python/sym/test_rewrite_simplify.py
@@ -1596,7 +1596,7 @@ class TestDivZero(BaseCompare):
class TestSubBufferload(BaseCompare):
- buf = tvm_tirx.tirx.decl_tensor([1], dtype="float32")
+ buf = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([1], "float32"))
load = tvm_tirx.ir.TensorLoad(buf, [0])
test_case = tvm_tirx.testing.parameter(
diff --git a/tirx/tests/python/sym/test_z3.py b/tirx/tests/python/sym/test_z3.py
index 755d0469a2..70b7e7121f 100644
--- a/tirx/tests/python/sym/test_z3.py
+++ b/tirx/tests/python/sym/test_z3.py
@@ -691,7 +691,7 @@ def test_z3_memo_pool_reuse_survives_clone():
# TensorLoad is read-state, so these expressions are memoized only for the
# duration of the query and then erased. This leaves holes in the Z3 memo
# pool that subsequent pure expressions must be able to reuse safely.
- buffer = tirx.decl_tensor((16,), "int32")
+ buffer = tirx.Var("tensor", tirx.TensorType((16,), "int32"))
read_expr = tirx.all(*(buffer[i] >= 0 for i in range(16)))
assert not analyzer.can_prove(read_expr, SB)
diff --git a/tirx/tests/python/tirx/analysis/test_check_well_defined.py
b/tirx/tests/python/tirx/analysis/test_check_well_defined.py
index 016fcc4bce..ad39f4cee5 100644
--- a/tirx/tests/python/tirx/analysis/test_check_well_defined.py
+++ b/tirx/tests/python/tirx/analysis/test_check_well_defined.py
@@ -352,7 +352,7 @@ def
test_tensor_load_asserted_type_matches_source_and_indices():
def test_tensor_load_malformed_indices_return_false_without_asserting():
- tensor = tvm_tirx.tirx.decl_tensor((4, 4), "float32")
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((4, 4),
"float32"))
vector_index = tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4)
load = tvm_tirx.ir.TensorLoad(tensor, [0, vector_index])
func = tvm_tirx.tirx.Function([tensor], tvm_tirx.ir.Evaluate(load))
diff --git a/tirx/tests/python/tirx/analysis/test_undefined_vars.py
b/tirx/tests/python/tirx/analysis/test_undefined_vars.py
index d3658116f4..4f3b68acd4 100644
--- a/tirx/tests/python/tirx/analysis/test_undefined_vars.py
+++ b/tirx/tests/python/tirx/analysis/test_undefined_vars.py
@@ -32,7 +32,7 @@ def test_decl_tensor_data_is_use():
from tvm_tirx.ir import PrimType, PtrType
data_ptr = tirx.Var("buf_data", PtrType(PrimType("float32")))
- buf = tirx.decl_tensor((n,), "float32", "buf", data=data_ptr)
+ buf = tirx.Var("buf", tirx.TensorType((n,), "float32", data=data_ptr))
body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(buf, [0]))
decl = tvm_tirx.ir.Bind(
@@ -62,7 +62,7 @@ def test_decl_tensor_elem_offset_is_use():
n = tirx.Var("n", "int32")
data_ptr = tirx.Var("buf_data", PtrType(PrimType("float32")))
elem_off = tirx.Var("buf_elem_offset", "int32")
- buf = tirx.decl_tensor((n,), "float32", "buf", data=data_ptr,
elem_offset=elem_off)
+ buf = tirx.Var("buf", tirx.TensorType((n,), "float32", data=data_ptr,
elem_offset=elem_off))
body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(buf, [0]))
decl = tvm_tirx.ir.Bind(
@@ -90,7 +90,7 @@ def test_alloc_tensor_data_is_def():
not a reference to an external variable.
"""
n = tirx.Var("n", "int32")
- buf = tirx.decl_tensor((n,), "float32", "buf")
+ buf = tirx.Var("buf", tirx.TensorType((n,), "float32"))
body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(buf, [0]))
alloc = tvm_tirx.ir.Bind(
@@ -114,7 +114,7 @@ def test_alloc_tensor_data_is_def():
def test_tensor_data_projection_is_tensor_use():
"""An opaque data projection must retain the TensorVar identity."""
- buf = tirx.decl_tensor((16,), "float32", "buf")
+ buf = tirx.Var("buf", tirx.TensorType((16,), "float32"))
stmt = tvm_tirx.ir.Evaluate(buf.data)
undef = tvm_tirx.tirx.analysis.undefined_vars(stmt, [])
diff --git a/tirx/tests/python/tirx/base/test_nodes.py
b/tirx/tests/python/tirx/base/test_nodes.py
index b8ec2b40e9..fa7f00b042 100644
--- a/tirx/tests/python/tirx/base/test_nodes.py
+++ b/tirx/tests/python/tirx/base/test_nodes.py
@@ -83,7 +83,7 @@ def test_ir2():
storage_type = ir.PrimType("int32")
handle_type = ir.PtrType(storage_type)
array = tvm_tirx.tirx.Var("array", handle_type)
- buf = tvm_tirx.tirx.decl_tensor([buf_size], "int32", data=array)
+ buf = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([buf_size],
"int32", data=array))
st = tvm_tirx.ir.TensorStore(buf, [1], x + 1)
assert isinstance(st, tvm_tirx.ir.TensorStore)
@@ -306,7 +306,7 @@ def test_equality_string_imm():
def test_function():
x = tvm_tirx.tirx.Var("x", "int32")
y = tvm_tirx.tirx.Var("y", "int32")
- b = tvm_tirx.tirx.decl_tensor((x,), "float32")
+ b = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((x,), "float32"))
stmt = tvm_tirx.ir.SeqStmt([tvm_tirx.ir.Bind(x, 10),
tvm_tirx.ir.Evaluate(x + 1)])
func = tvm_tirx.tirx.Function([x, y, b], stmt)
@@ -341,7 +341,7 @@ def test_scoped_storage_vars():
def test_tensor_load_store():
- b = tvm_tirx.tirx.decl_tensor((10,), "float32")
+ b = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((10,), "float32"))
x = tvm_tirx.ir.TensorLoad(b, [0])
assert isinstance(x, tvm_tirx.ir.TensorLoad)
assert x.ty.dtype == "float32"
@@ -356,7 +356,9 @@ def test_tensor_load_store():
s = tvm_tirx.ir.TensorStore(b, [0], 0.1)
assert isinstance(s, tvm_tirx.ir.TensorStore)
- tensor = tvm_tirx.tirx.decl_tensor((4, 8), "float32", elem_offset=100)
+ tensor = tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType((4, 8), "float32", elem_offset=100)
+ )
tvm_tirx.ir.assert_structural_equal(
tensor[2, 3].indices,
[tvm_tirx.ir.IntImm("int32", 2), tvm_tirx.ir.IntImm("int32", 3)],
@@ -364,7 +366,7 @@ def test_tensor_load_store():
def test_tensor_load_validation():
- tensor = tvm_tirx.tirx.decl_tensor((8, 8), "float32")
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((8, 8),
"float32"))
with pytest.raises(RuntimeError, match="dimensional"):
tvm_tirx.ir.TensorLoad(tensor, [0])
with pytest.raises(RuntimeError, match="Only the last index"):
@@ -382,14 +384,14 @@ def test_tensor_load_validation():
],
)
def test_tensor_load_vector_elements(element_dtype, index_lanes, result_dtype):
- tensor = tvm_tirx.tirx.decl_tensor((32,), element_dtype)
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((32,),
element_dtype))
lanes = 4 * tvm_tirx.tirx.vscale() if index_lanes == "scalable" else
index_lanes
load = tvm_tirx.ir.TensorLoad(tensor,
[tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, lanes)])
assert load.ty.dtype == result_dtype
def test_tensor_load_rejects_two_scalable_types():
- tensor = tvm_tirx.tirx.decl_tensor((32,), "float32xvscalex2")
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((32,),
"float32xvscalex2"))
with pytest.raises(RuntimeError, match="can't both be scalable"):
tvm_tirx.ir.TensorLoad(
tensor, [tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4 *
tvm_tirx.tirx.vscale())]
@@ -458,7 +460,7 @@ def test_broadcast_to_scalable_vec():
def test_tensor_load_scalable_vec():
- buf = tvm_tirx.tirx.decl_tensor((24,), "float32")
+ buf = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((24,),
"float32"))
index = tvm_tirx.ir.prim.vector_expr.Ramp(1, 1, 8 * tvm_tirx.tirx.vscale())
load = tvm_tirx.ir.TensorLoad(buf, [index])
@@ -467,7 +469,7 @@ def test_tensor_load_scalable_vec():
def test_tensor_store_scalable_vec():
- b = tvm_tirx.tirx.decl_tensor((24,), "int32")
+ b = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((24,), "int32"))
value = tvm_tirx.ir.prim.vector_expr.Broadcast(1, 4 *
tvm_tirx.tirx.vscale())
index = tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4 * tvm_tirx.tirx.vscale())
store = tvm_tirx.ir.TensorStore(b, [index], value)
@@ -477,7 +479,7 @@ def test_tensor_store_scalable_vec():
def test_scalable_vec_cast():
- b = tvm_tirx.tirx.decl_tensor((24,), "float32")
+ b = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((24,), "float32"))
value = tvm_tirx.ir.prim.vector_expr.Broadcast(1, 12 *
tvm_tirx.tirx.vscale()).astype(
"float32xvscalex12"
)
diff --git a/tirx/tests/python/tirx/base/test_op_types.py
b/tirx/tests/python/tirx/base/test_op_types.py
index 324f33af3a..8c0caf23a5 100644
--- a/tirx/tests/python/tirx/base/test_op_types.py
+++ b/tirx/tests/python/tirx/base/test_op_types.py
@@ -47,11 +47,13 @@ def test_tir_op_tvm_struct_set():
def test_tir_op_address_of():
- tensor = tirx.decl_tensor((128), "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType(128, "float32"))
expr = tirx.address_of(tensor[0])
assert expr.op.name == "tirx.address_of"
storage = tirx.Var("storage",
tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("uint8"), "shared.dyn"))
- pooled_tensor = tirx.decl_tensor((128), "float32", data=storage,
scope="shared.dyn")
+ pooled_tensor = tirx.Var(
+ "tensor", tirx.TensorType(128, "float32", data=storage,
scope="shared.dyn")
+ )
expected_ty = tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32"),
"shared.dyn")
assert tirx.address_of(pooled_tensor).ty == expected_ty
assert tirx.address_of(pooled_tensor[0]).ty == expected_ty
@@ -102,7 +104,7 @@ def test_primitive_result_type_keyword(op_name):
def test_masked_memory_wrapper_parity_and_inference():
from tvm_tirx.script import tirx as T
- tensor = tirx.decl_tensor((16,), "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType((16,), "float32"))
index = tvm_tirx.ir.prim.Ramp(0, 1, 4)
mask = tvm_tirx.ir.prim.Broadcast(True, 4)
loc = tvm_tirx.ir.SourceLoc(tvm_tirx.ir.SourceName("op.py"), 2, 0, 2, 8)
@@ -168,8 +170,8 @@ def test_tir_op_call_likely():
def test_tir_op_tvm_thread_allreduce():
- tensor = tirx.decl_tensor((128,), "float32")
- result = tirx.decl_tensor((1,), "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType((128,), "float32"))
+ result = tirx.Var("tensor", tirx.TensorType((1,), "float32"))
axis = tirx.Var("thread", "int32")
combine = tvm_tirx.ir.LambdaExpr(["float32", "float32"], lambda lhs, rhs:
(lhs + rhs,))
expr = tirx.gpu_thread_allreduce(
@@ -179,7 +181,7 @@ def test_tir_op_tvm_thread_allreduce():
def test_tir_op_ptr_byte_offset():
- tensor = tirx.decl_tensor((128,), "float32", scope="shared")
+ tensor = tirx.Var("tensor", tirx.TensorType((128,), "float32",
scope="shared"))
ty = tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("uint8"), "shared")
offset_expr = tirx.ptr_byte_offset(tensor.data, 16, ty=ty)
assert offset_expr.op.name == "tirx.ptr_byte_offset"
@@ -193,24 +195,24 @@ def test_tir_op_tvm_throw_last_error():
def test_tir_op_tvm_load_matrix_sync():
- tensor = tirx.decl_tensor((16, 16), "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
x = tirx.Var("x", "handle")
expr = tvm_tirx.tirx.gpu_load_matrix_sync(tensor.data, 16, 16, 16, 0, x,
128, "row_major")
assert expr.op.name == "tirx.gpu_load_matrix_sync"
def test_tir_op_tvm_store_matrix_sync():
- tensor = tirx.decl_tensor((16, 16), "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
x = tirx.Var("x", "handle")
expr = tvm_tirx.tirx.gpu_store_matrix_sync(tensor.data, 16, 16, 16, 0, x,
128, "row_major")
assert expr.op.name == "tirx.gpu_store_matrix_sync"
def test_tir_op_tvm_mma_sync():
- tensor_0 = tirx.decl_tensor((16, 16), "float32")
- tensor_1 = tirx.decl_tensor((16, 16), "float32")
- tensor_2 = tirx.decl_tensor((16, 16), "float32")
- tensor_3 = tirx.decl_tensor((16, 16), "float32")
+ tensor_0 = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
+ tensor_1 = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
+ tensor_2 = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
+ tensor_3 = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
expr = tvm_tirx.tirx.gpu_mma_sync(
tensor_0.data, 0, tensor_1.data, 0, tensor_2.data, 0, tensor_3.data, 0
)
@@ -218,10 +220,10 @@ def test_tir_op_tvm_mma_sync():
def test_tir_op_tvm_bmma_sync():
- tensor_0 = tirx.decl_tensor((16, 16), "float32")
- tensor_1 = tirx.decl_tensor((16, 16), "float32")
- tensor_2 = tirx.decl_tensor((16, 16), "float32")
- tensor_3 = tirx.decl_tensor((16, 16), "float32")
+ tensor_0 = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
+ tensor_1 = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
+ tensor_2 = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
+ tensor_3 = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
expr = tvm_tirx.backend.cuda.op.bmma_sync(
tensor_0.data, 0, tensor_1.data, 0, tensor_2.data, 0, tensor_3.data, 0
)
@@ -229,15 +231,15 @@ def test_tir_op_tvm_bmma_sync():
def test_tir_op_tvm_fill_fragment():
- tensor = tirx.decl_tensor((16, 16), "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
expr = tvm_tirx.tirx.gpu_fill_fragment(tensor.data, 16, 16, 16, 0, 0)
assert expr.op.name == "tirx.gpu_fill_fragment"
def test_tir_op_ptx_mma():
- tensor_a = tirx.decl_tensor([32], "int4", scope="local")
- tensor_b = tirx.decl_tensor([16], "uint4", scope="local")
- tensor_c = tirx.decl_tensor([4], "int32", scope="local")
+ tensor_a = tirx.Var("tensor", tirx.TensorType([32], "int4", scope="local"))
+ tensor_b = tirx.Var("tensor", tirx.TensorType([16], "uint4",
scope="local"))
+ tensor_c = tirx.Var("tensor", tirx.TensorType([4], "int32", scope="local"))
expr = _cuda_op.ptx_legacy_mma(
"m8n8k32",
"row",
@@ -259,9 +261,10 @@ def test_tir_op_ptx_mma():
def test_tir_op_mma_store():
x = tirx.Var("x", ty="int32")
y = tirx.Var("y", ty="int32")
- tensor_w = tirx.decl_tensor([16, 8], dtype="int32", scope="warp",
offset_factor=1)
- tensor = tirx.decl_tensor(
- [16, 16], dtype="int32", scope="global", offset_factor=1, strides=[x,
y]
+ tensor_w = tirx.Var("tensor", tirx.TensorType([16, 8], "int32",
scope="warp", offset_factor=1))
+ tensor = tirx.Var(
+ "tensor",
+ tirx.TensorType([16, 16], "int32", scope="global", offset_factor=1,
strides=[x, y]),
)
expr = _cuda_op.mma_store(
16, 16, tensor.ptr_to([0, 0]), tensor_w.data, tensor_w.ty.elem_offset,
x, ty="int32"
@@ -270,14 +273,14 @@ def test_tir_op_mma_store():
def test_tir_op_mma_fill():
- tensor_w = tirx.decl_tensor([16, 8], dtype="int32", scope="warp",
offset_factor=1)
+ tensor_w = tirx.Var("tensor", tirx.TensorType([16, 8], "int32",
scope="warp", offset_factor=1))
expr = _cuda_op.mma_fill(8, tensor_w.data, tensor_w.ty.elem_offset,
ty="int32")
assert expr.op.name == "tirx.cuda.mma_fill"
def test_op_ptx_cp_async():
- tensor_shared = tirx.decl_tensor([16, 16], "float16", scope="shared")
- tensor_local = tirx.decl_tensor([8], "float16", scope="local")
+ tensor_shared = tirx.Var("tensor", tirx.TensorType([16, 16], "float16",
scope="shared"))
+ tensor_local = tirx.Var("tensor", tirx.TensorType([8], "float16",
scope="local"))
expr = _cuda_op.ptx_cp_async_legacy(tensor_shared.data, 0,
tensor_local.data, 0, 16)
assert expr.op.name == "tirx.s_tir.cp_async_raw"
@@ -295,14 +298,14 @@ def test_op_ptx_cp_async():
def test_tir_op_vector_low():
- tensor = tirx.decl_tensor((4, 4), "int8", offset_factor=1)
+ tensor = tirx.Var("tensor", tirx.TensorType((4, 4), "int8",
offset_factor=1))
vec = tensor[0, tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 16)]
expr = tirx.vector_low(vec, ty="int8x8")
assert expr.op.name == "tirx.vector_low"
def test_tir_op_vector_high():
- tensor = tirx.decl_tensor((4, 4), "int8", offset_factor=1)
+ tensor = tirx.Var("tensor", tirx.TensorType((4, 4), "int8",
offset_factor=1))
vec = tensor[0, tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 16)]
expr = tirx.vector_high(vec, ty="int8x8")
assert expr.op.name == "tirx.vector_high"
@@ -317,7 +320,7 @@ def test_tir_op_dp4a():
def test_tir_op_vector_combine():
- tensor = tirx.decl_tensor((4, 4), "int8", offset_factor=1)
+ tensor = tirx.Var("tensor", tirx.TensorType((4, 4), "int8",
offset_factor=1))
vec = tensor[0, tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 16)]
expr = tirx.vector_combine(vec, vec, ty="int8x8")
assert expr.op.name == "tirx.vector_combine"
@@ -356,7 +359,7 @@ def test_tir_op_TVMBackendAllocWorkspace():
def test_tir_op_TVMBackendFreeWorkspace():
- tensor = tirx.decl_tensor((128), "float32")
+ tensor = tirx.Var("tensor", tirx.TensorType(128, "float32"))
expr = tirx.free_workspace(0, 1, tensor.data)
assert expr.op.name == "tirx.free_workspace"
diff --git a/tirx/tests/python/tirx/base/test_specialize.py
b/tirx/tests/python/tirx/base/test_specialize.py
index f2c13332fe..d15506b9eb 100644
--- a/tirx/tests/python/tirx/base/test_specialize.py
+++ b/tirx/tests/python/tirx/base/test_specialize.py
@@ -195,7 +195,9 @@ def test_specialize_nothing():
def test_specialize_matmul():
a, _, _, n = matmul.params
# fully specialized
- func = matmul.specialize({a: tvm_tirx.tirx.decl_tensor((128, 128))})
+ func = matmul.specialize(
+ {a: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((128, 128),
"float32"))}
+ )
assert_structural_equal_ignore_global_symbol(func, matmul_128)
# partially specialized
func = matmul.specialize({n: 128})
@@ -209,10 +211,14 @@ def test_specialize_elemwise():
a, c = element_wise.params
C = c
# fully specialized
- func = element_wise.specialize({a: tvm_tirx.tirx.decl_tensor((128, 64))})
+ func = element_wise.specialize(
+ {a: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((128, 64),
"float32"))}
+ )
assert_structural_equal_ignore_global_symbol(func, element_wise_128_64)
# partially specialized
- func = element_wise.specialize({c: tvm_tirx.tirx.decl_tensor((128,
C.ty.shape[1]))})
+ func = element_wise.specialize(
+ {c: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((128,
C.ty.shape[1]), "float32"))}
+ )
assert_structural_equal_ignore_global_symbol(func, element_wise_128_n)
@@ -220,7 +226,14 @@ def test_specialize_mem_copy():
a, _, m, n, p, q = mem_copy.params
# fully specialized
func = mem_copy.specialize(
- {a: tvm_tirx.tirx.decl_tensor((16, 16), strides=[8, 1],
elem_offset=T.int32(4))}
+ {
+ a: tvm_tirx.tirx.Var(
+ "tensor",
+ tvm_tirx.tirx.TensorType(
+ (16, 16), "float32", strides=[8, 1], elem_offset=T.int32(4)
+ ),
+ )
+ }
)
assert_structural_equal_ignore_global_symbol(func, mem_copy_16_16_8_4)
func = mem_copy.specialize({n: 16, m: 16, p: 8, q: 4})
@@ -249,7 +262,9 @@ def test_specialize_with_const_folding():
B[i] = A[i // 8, i % 8] + 714
b = before.params[1]
- after = before.specialize({b: tvm_tirx.tirx.decl_tensor([16],
dtype="int32")})
+ after = before.specialize(
+ {b: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([16],
"int32"))}
+ )
assert_structural_equal_ignore_global_symbol(expected, after)
diff --git a/tirx/tests/python/tirx/base/test_tensor_type.py
b/tirx/tests/python/tirx/base/test_tensor_type.py
index 4973f77593..41b999474d 100644
--- a/tirx/tests/python/tirx/base/test_tensor_type.py
+++ b/tirx/tests/python/tirx/base/test_tensor_type.py
@@ -28,7 +28,7 @@ from tvm_tirx.script import tirx as T
def test_tensor_pointer_type_derived_from_dtype_and_scope():
data = tvm_tirx.ir.Var("storage",
tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("uint8"), "local"))
- tensor = tvm_tirx.tirx.decl_tensor((8,), "float16", data=data)
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((8,),
"float16", data=data))
assert tensor.ty.dtype == tvm_tirx.ir.PrimType("float16")
assert tensor.ty.scope == tvm_tirx.runtime.MemoryScope("local")
@@ -39,14 +39,17 @@ def test_tensor_pointer_type_derived_from_dtype_and_scope():
def test_tensor_type_constructor_scope_order(scope):
direct = tvm_tirx.tirx.TensorType((4, 8), "float32", scope, (16, 1), 4,
64, 2)
scripted = T.Tensor((4, 8), "float32", scope, (16, 1), 4, 64, 2)
- declared = tvm_tirx.tirx.decl_tensor(
- (4, 8),
- "float32",
- scope=scope,
- strides=(16, 1),
- elem_offset=4,
- data_alignment=64,
- offset_factor=2,
+ declared = tvm_tirx.tirx.Var(
+ "tensor",
+ tvm_tirx.tirx.TensorType(
+ (4, 8),
+ "float32",
+ scope=scope,
+ strides=(16, 1),
+ elem_offset=4,
+ data_alignment=64,
+ offset_factor=2,
+ ),
)
tvm_tirx.ir.assert_structural_equal(direct, scripted)
tvm_tirx.ir.assert_structural_equal(direct, declared.ty)
@@ -60,11 +63,11 @@ def test_tensor_type_rejects_unregistered_scope(scope):
with pytest.raises(ValueError):
T.Tensor((4,), "float32", scope)
with pytest.raises(ValueError):
- tvm_tirx.tirx.decl_tensor((4,), "float32", scope=scope)
+ tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((4,), "float32",
scope=scope))
def test_decl_tensor_physical_data_binding():
- tensor = tvm_tirx.tirx.decl_tensor((8,), "float32")
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((8,),
"float32"))
data = tvm_tirx.tirx.Var("data", tensor.data.ty)
decl = tvm_tirx.ir.Bind(
@@ -82,13 +85,17 @@ def test_decl_tensor_physical_data_binding():
def test_tensor_ptr_to():
m = tvm_tirx.tirx.Var("m", "int32")
n = tvm_tirx.tirx.Var("n", "int32")
- tensor = tvm_tirx.tirx.decl_tensor((m, n), "float32", strides=[n + 1, 1],
elem_offset=7)
+ tensor = tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType((m, n), "float32", strides=[n + 1,
1], elem_offset=7)
+ )
pointer = tensor.ptr_to([2, 3])
assert pointer.op.name == "tirx.address_of"
assert pointer.ty == tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32"))
assert pointer.args[0].source.same_as(tensor)
tvm_tirx.ir.assert_structural_equal(pointer.args[0].indices, [T.int32(2),
T.int32(3)])
- shared = tvm_tirx.tirx.decl_tensor((m, n), "float32", scope="shared")
+ shared = tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType((m, n), "float32", scope="shared")
+ )
assert shared.ptr_to([0, 0]).ty == tvm_tirx.ir.PtrType(
tvm_tirx.ir.PrimType("float32"), "shared"
)
@@ -100,7 +107,7 @@ def test_tensor_type_and_projection(dtype, scope):
tvm_tirx.ir.Var("m", tvm_tirx.ir.PrimType("int32")),
tvm_tirx.ir.Var("n", tvm_tirx.ir.PrimType("int32")),
)
- tensor = tvm_tirx.tirx.decl_tensor((m, n), dtype, scope=scope)
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((m, n),
dtype, scope=scope))
assert type(tensor) is tvm_tirx.ir.Var
assert tvm_tirx.tirx.is_tensor_var(tensor)
assert not tvm_tirx.tirx.is_tensor_var(m)
@@ -124,7 +131,9 @@ def test_tensor_type_and_projection(dtype, scope):
def test_tensor_type_default_offset():
- tensor = tvm_tirx.tirx.decl_tensor((8,), "float32", elem_offset=None)
+ tensor = tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType((8,), "float32", elem_offset=None)
+ )
assert int(tensor.ty.elem_offset) == 0
@@ -145,7 +154,7 @@ def test_tensor_type_default_offset():
def test_tensor_type_low_precision_dtype(dtype):
m = tvm_tirx.ir.Var("m", tvm_tirx.ir.PrimType("int32"))
n = tvm_tirx.ir.Var("n", tvm_tirx.ir.PrimType("int32"))
- tensor = tvm_tirx.tirx.decl_tensor((m, n), dtype)
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((m, n),
dtype))
assert tensor.dtype == dtype
@@ -153,7 +162,7 @@ def test_memory_scope_typed_conversion():
custom = MemoryScope.register("test.tensor.scope." + uuid.uuid4().hex)
for scope in (custom, custom.name):
ptr_ty = tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32"), scope)
- tensor = tvm_tirx.tirx.decl_tensor((4,), "float32", scope=scope)
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((4,),
"float32", scope=scope))
tensor_ty = tvm_tirx.tirx.TensorType((4,), "float32", scope)
assert tensor_ty.scope.same_as(custom)
assert ptr_ty.scope.same_as(custom)
@@ -164,7 +173,9 @@ def test_memory_scope_typed_conversion():
unknown = "test.memory_scope.unknown." + uuid.uuid4().hex
for construct in (
lambda: tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32"), unknown),
- lambda: tvm_tirx.tirx.decl_tensor((4,), "float32", scope=unknown),
+ lambda: tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType((4,), "float32", scope=unknown)
+ ),
lambda: T.Tensor((4,), "float32", scope=unknown),
):
with pytest.raises(ValueError):
@@ -176,17 +187,25 @@ def test_memory_scope_typed_conversion():
def test_memory_scope_global_defaults():
global_scope = MemoryScope("global")
assert
tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32")).scope.same_as(global_scope)
- assert tvm_tirx.tirx.decl_tensor((4,),
"float32").scope().same_as(global_scope)
+ assert (
+ tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((4,), "float32"))
+ .scope()
+ .same_as(global_scope)
+ )
for scope in ("global", global_scope):
assert tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32"),
scope).scope.same_as(
global_scope
)
- assert tvm_tirx.tirx.decl_tensor((4,), "float32",
scope=scope).scope().same_as(global_scope)
+ assert (
+ tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((4,),
"float32", scope=scope))
+ .scope()
+ .same_as(global_scope)
+ )
assert T.Tensor((4,), "float32",
scope=scope).scope.same_as(global_scope)
for construct in (
lambda: MemoryScope(""),
lambda: tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("float32"), ""),
- lambda: tvm_tirx.tirx.decl_tensor((4,), "float32", scope=""),
+ lambda: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((4,),
"float32", scope="")),
lambda: T.Tensor((4,), "float32", scope=""),
):
with pytest.raises(ValueError):
@@ -202,5 +221,54 @@ def test_tensor_type_memory_scope_json_roundtrip():
tvm_tirx.ir.assert_structural_equal(ty, restored)
+def test_tensor_type_pointer_validation_and_scope():
+ custom = MemoryScope.register("test.tensor.pointer." + uuid.uuid4().hex)
+ data = tvm_tirx.ir.Var("data",
tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("uint8"), custom))
+ tensor = tvm_tirx.ir.Var(
+ "tensor", tvm_tirx.tirx.TensorType(8, "float32", scope="global",
data=data)
+ )
+ assert tensor.scope().same_as(custom)
+ assert tensor.data.ty.scope.same_as(custom)
+ assert tensor.data.ty.element_type == tvm_tirx.ir.PrimType("float32")
+ for invalid in (0, tvm_tirx.ir.Var("scalar", "int32")):
+ with pytest.raises(TypeError, match="Tensor data must be an Expr with
PtrType"):
+ tvm_tirx.tirx.TensorType(8, "float32", data=invalid)
+ pointer = tvm_tirx.ir.Var("pointer",
tvm_tirx.ir.PtrType(tvm_tirx.ir.TupleType([])))
+ with pytest.raises(TypeError, match="Tensor data must point to a primitive
type"):
+ tvm_tirx.tirx.TensorType(8, "float32", data=pointer)
+
+
+def test_tensor_type_explicit_location():
+ loc = tvm_tirx.ir.SourceLoc(tvm_tirx.ir.SourceName("tensor.py"), 2, 3, 2,
12)
+ ty = tvm_tirx.tirx.TensorType(8, "float32", loc=loc)
+ tensor = tvm_tirx.ir.Var("tensor", ty, loc=loc)
+ assert ty.loc.same_as(loc)
+ assert tensor.loc.same_as(loc)
+
+
+def test_tensor_scope_receiver_validation():
+ ty = tvm_tirx.tirx.TensorType(8, "float32")
+ tensor = tvm_tirx.ir.Var("tensor", ty)
+ assert ty._scope(tensor).same_as(ty.scope)
+ with pytest.raises(TypeError, match="Tensor methods expect a Var with
TensorType"):
+ ty._scope(tvm_tirx.ir.Var("scalar", "int32"))
+ with pytest.raises(TypeError, match="Tensor method type must be the
operand's type"):
+ tvm_tirx.tirx.TensorType(8, "float32")._scope(tensor)
+
+
[email protected]("alignment", [-1, 0, 64])
+def test_tensor_type_construction_offsets_and_alignment(alignment):
+ ty = tvm_tirx.tirx.TensorType(
+ (), "float32", data_alignment=alignment, byte_offset=16, strides=()
+ )
+ tensor = tvm_tirx.ir.Var("scalar", ty)
+ assert list(tensor.shape) == []
+ assert tensor.layout is None
+ assert int(tensor.elem_offset) == 4
+ assert int(tensor.byte_offset) == 16
+ assert tensor.data_alignment == 64
+ assert tensor.offset_factor == 1
+
+
if __name__ == "__main__":
tvm_tirx.testing.main()
diff --git a/tirx/tests/python/tirx/codegen/test_cuda.py
b/tirx/tests/python/tirx/codegen/test_cuda.py
index d34b68520c..299ec0f355 100644
--- a/tirx/tests/python/tirx/codegen/test_cuda.py
+++ b/tirx/tests/python/tirx/codegen/test_cuda.py
@@ -254,7 +254,7 @@ def test_typed_buffer_bindings(monkeypatch, alignment):
def test_vector_pointer_preserves_packed_offset(monkeypatch):
- tensor = tvm_tirx.tirx.decl_tensor((8,), "int4x4", name="A")
+ tensor = tvm_tirx.tirx.Var("A", tvm_tirx.tirx.TensorType((8,), "int4x4"))
data = tvm_tirx.tirx.Var("A_data", tensor.data.ty)
pointer = tvm_tirx.tirx.ptr_byte_offset(tensor.data, 4, ty=tensor.data.ty)
body = tvm_tirx.ir.SeqStmt(
diff --git a/tirx/tests/python/tirx/codegen/test_ptx_dialect.py
b/tirx/tests/python/tirx/codegen/test_ptx_dialect.py
index a30a0e516a..0ab3dff696 100644
--- a/tirx/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tirx/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -3155,8 +3155,8 @@ def test_ptx_tcgen05_mapa_address_coercion():
raw_u64 = tvm_tirx.tirx.Var("a64", "uint64")
ncols = tvm_tirx.tirx.Var("n", "uint32")
mask = tvm_tirx.tirx.Var("mask", "uint16")
- out32 = tvm_tirx.tirx.decl_tensor((1,), "uint32", name="out32",
scope="local")
- out64 = tvm_tirx.tirx.decl_tensor((1,), "uint64", name="out64",
scope="local")
+ out32 = tvm_tirx.tirx.Var("out32", tvm_tirx.tirx.TensorType((1,),
"uint32", scope="local"))
+ out64 = tvm_tirx.tirx.Var("out64", tvm_tirx.tirx.TensorType((1,),
"uint64", scope="local"))
bare_alloc =
T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.b32(generic_ptr, ncols)
assert bare_alloc.args[0].same_as(generic_ptr)
@@ -3813,7 +3813,7 @@ def test_ptx_coercion_ir_forms():
call = T.ptx.st.release.gpu.global_.b32(global_ptr, val, pred=flag)
assert call.args[2].same_as(flag)
assert len(call.args) == 2 + 1 + 8 + 1 # operands + pred + slot tokens +
marker
- out = tvm_tirx.tirx.decl_tensor((1,), "uint32", name="out", scope="local")
+ out = tvm_tirx.tirx.Var("out", tvm_tirx.tirx.TensorType((1,), "uint32",
scope="local"))
call = T.ptx.ld.global_.b32(out[0], global_ptr, pred=flag)
assert str(call.args[-1]).strip('"') == "pred"
call = T.ptx.ld.global_.b32(out[0], global_ptr, pred=flag,
preserve_dst=True)
diff --git a/tirx/tests/python/tirx/script/test_tensors.py
b/tirx/tests/python/tirx/script/test_tensors.py
index e65f451833..850028d7f5 100644
--- a/tirx/tests/python/tirx/script/test_tensors.py
+++ b/tirx/tests/python/tirx/script/test_tensors.py
@@ -348,7 +348,7 @@ def test_tensor():
@pytest.mark.skipif(sys.version_info < (3, 12), reason="PEP 695 requires
Python 3.12")
def test_tensor_shape_repeated_var_prints_out_of_line():
n = tvm_tirx.tirx.Var("n", "int32")
- tensor = tvm_tirx.tirx.decl_tensor((n + n,), name="A")
+ tensor = tvm_tirx.tirx.Var("A", tvm_tirx.tirx.TensorType((n + n,),
"float32"))
func = tvm_tirx.tirx.Function([tensor], tvm_tirx.ir.Evaluate(0))
code = func.script()
@@ -987,8 +987,13 @@ def test_tensor_sub_multi_iter_dim_ir():
def test_tensor_sub_multi_iter_misaligned_rejected():
- buf = tvm_tirx.tirx.decl_tensor(
- (8, 16), "float16", layout=tvm_tirx.tirx.layout.TileLayout(T.S[(2, 4,
16) : (1024, 64, 1)])
+ buf = tvm_tirx.tirx.Var(
+ "tensor",
+ tvm_tirx.tirx.TensorType(
+ (8, 16),
+ "float16",
+ layout=tvm_tirx.tirx.layout.TileLayout(T.S[(2, 4, 16) : (1024, 64,
1)]),
+ ),
)
# sub[2:6] narrows the multi-iter dim 0 at a misaligned offset.
with pytest.raises(ValueError, match="multiples of the inner iter block"):
@@ -1039,11 +1044,17 @@ def test_tensor_sub_ir():
def test_tensor_view_surgery_static_bounds_rejected():
"""Statically-known out-of-range sub arguments must be rejected loudly
(review finding: OOB offsets were silent)."""
- buf = tvm_tirx.tirx.decl_tensor(
- (10,), "float16", layout=tvm_tirx.tirx.layout.TileLayout(T.S[(10,) :
(1,)])
+ buf = tvm_tirx.tirx.Var(
+ "tensor",
+ tvm_tirx.tirx.TensorType(
+ (10,), "float16", layout=tvm_tirx.tirx.layout.TileLayout(T.S[(10,)
: (1,)])
+ ),
)
- grid = tvm_tirx.tirx.decl_tensor(
- (4, 8), "float16", layout=tvm_tirx.tirx.layout.TileLayout(T.S[(4, 8) :
(8, 1)])
+ grid = tvm_tirx.tirx.Var(
+ "tensor",
+ tvm_tirx.tirx.TensorType(
+ (4, 8), "float16", layout=tvm_tirx.tirx.layout.TileLayout(T.S[(4,
8) : (8, 1)])
+ ),
)
# int index: static bounds
with pytest.raises(ValueError, match="out of range"):
@@ -1230,9 +1241,14 @@ def test_tensor_factored_views():
def test_tensor_contiguous_regions():
"""Explicit contiguous regions preserve rank, offsets and dynamic
starts."""
- A = tvm_tirx.tirx.decl_tensor((4, 8, 16), "float16", layout="default")
- C = tvm_tirx.tirx.decl_tensor(
- (4, 512), "bfloat16", layout=T.ComposeLayout(3, 3, 3,
T.TileLayout(T.S[(4, 512)]))
+ A = tvm_tirx.tirx.Var(
+ "tensor", tvm_tirx.tirx.TensorType((4, 8, 16), "float16",
layout="default")
+ )
+ C = tvm_tirx.tirx.Var(
+ "tensor",
+ tvm_tirx.tirx.TensorType(
+ (4, 512), "bfloat16", layout=T.ComposeLayout(3, 3, 3,
T.TileLayout(T.S[(4, 512)]))
+ ),
)
reg = A[:, :, 8:16]
assert isinstance(reg, TensorRegion)
@@ -1277,7 +1293,7 @@ def test_tensor_view_dtype_ir():
def test_tensor_slice_region():
"""Verify A[slice] returns TensorRegion (not DeclTensor)."""
- buf = tvm_tirx.tirx.decl_tensor((128, 64), "float16")
+ buf = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((128, 64),
"float16"))
br = buf[32:64, 0:32]
assert isinstance(br, TensorRegion)
assert br.source.same_as(buf)
diff --git a/tirx/tests/python/tirx/test_op.py
b/tirx/tests/python/tirx/test_op.py
index c0be7f08fe..9897c91879 100644
--- a/tirx/tests/python/tirx/test_op.py
+++ b/tirx/tests/python/tirx/test_op.py
@@ -23,13 +23,13 @@ import tvm_tirx
from tvm_tirx.ir import Call, Evaluate, assert_structural_equal, load_json,
save_json
from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Var, decl_tensor
+from tvm_tirx.tirx import TensorType, Var
from tvm_tirx.tirx.analysis import undefined_vars
def test_instruction_is_a_void_call():
- src = decl_tensor((16, 4), "float32", scope="global")
- dst = decl_tensor((16, 4), "float32", scope="local")
+ src = Var("tensor", TensorType((16, 4), "float32", scope="global"))
+ dst = Var("tensor", TensorType((16, 4), "float32", scope="local"))
call = T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], dst, src, scope="warp")
assert isinstance(call, Call)
assert call.ty == tvm_tirx.ir.PrimType("void")
@@ -67,8 +67,8 @@ def test_roundtrip_and_exactly_one_evaluate(composite):
def test_runtime_operands_are_seen_and_replaced():
- src = decl_tensor((16, 16), "float32", scope="global")
- dst = decl_tensor((4, 16), "float32", scope="shared")
+ src = Var("tensor", TensorType((16, 16), "float32", scope="global"))
+ dst = Var("tensor", TensorType((4, 16), "float32", scope="shared"))
mbar = Var("mbar", "handle")
row = Var("row", "int32")
mask = Var("mask", "uint16")
@@ -101,8 +101,8 @@ def test_runtime_operands_are_seen_and_replaced():
def test_workspace_is_an_optional_call_operand():
tvm_tirx.backend.load("trn")
- a = decl_tensor((16, 16), "float32", scope="shared")
- acc = decl_tensor((16, 16), "float32", scope="trn.psum")
+ a = Var("tensor", TensorType((16, 16), "float32", scope="shared"))
+ acc = Var("tensor", TensorType((16, 16), "float32", scope="trn.psum"))
call = T.trn.tile.matmul(a, a, a, acc_psum=acc)
assert call.args[-1].same_as(acc)
assert any(v.same_as(acc) for v in undefined_vars(Evaluate(call)))
@@ -110,7 +110,7 @@ def test_workspace_is_an_optional_call_operand():
def test_direct_call_validation_cannot_bypass_signature():
- a = decl_tensor((16, 1), "float32", scope="local")
+ a = Var("tensor", TensorType((16, 1), "float32", scope="local"))
valid = T.cuda.tile.apply(T.ptx["mov.b32"], a, T.float32(0))
for call in [
Call(
diff --git a/tirx/tests/python/tirx/test_op_namespace_cleanup.py
b/tirx/tests/python/tirx/test_op_namespace_cleanup.py
index 2f3c6b9635..6e1617c435 100644
--- a/tirx/tests/python/tirx/test_op_namespace_cleanup.py
+++ b/tirx/tests/python/tirx/test_op_namespace_cleanup.py
@@ -161,7 +161,7 @@ def
test_device_intrinsic_namespaces_are_canonical_and_classified():
assert T.metal is builder_op.metal
assert T.nki is builder_op.nki
- tensor = tvm_tirx.tirx.decl_tensor((1,), "float32")
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((1,),
"float32"))
calls = [
T.cuda.elect_sync(),
T.cuda.thread_fence(),
diff --git a/tirx/tests/python/tirx/test_tensor_data_reinfer_type.py
b/tirx/tests/python/tirx/test_tensor_data_reinfer_type.py
index 309911d95f..174bbe03e1 100644
--- a/tirx/tests/python/tirx/test_tensor_data_reinfer_type.py
+++ b/tirx/tests/python/tirx/test_tensor_data_reinfer_type.py
@@ -22,8 +22,8 @@ from tvm_tirx import ir, tirx
def test_tensor_data_reinfer_type_from_rewritten_argument():
- global_tensor = tirx.decl_tensor((8,), "float32", name="global_tensor",
scope="global")
- local_tensor = tirx.decl_tensor((8,), "float16", name="local_tensor",
scope="local")
+ global_tensor = tirx.Var("global_tensor", tirx.TensorType((8,), "float32",
scope="global"))
+ local_tensor = tirx.Var("local_tensor", tirx.TensorType((8,), "float16",
scope="local"))
stale_type = global_tensor.data.ty
expected_type = local_tensor.data.ty
diff --git a/tirx/tests/python/tirx/test_thread_scope.py
b/tirx/tests/python/tirx/test_thread_scope.py
index 536afd71ed..a205e84b00 100644
--- a/tirx/tests/python/tirx/test_thread_scope.py
+++ b/tirx/tests/python/tirx/test_thread_scope.py
@@ -26,7 +26,7 @@ from tvm_ffi.dataclasses import Enum
import tvm_tirx
from tvm_tirx.ir import Call, Evaluate, make_node
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function, decl_tensor
+from tvm_tirx.tirx import Function, TensorType, Var
from tvm_tirx.tirx.thread_scope import ThreadScope
from tvm_tirx.tirx.tile import DispatchContext
from tvm_tirx.tirx.tile.instruction import TensorCall
@@ -79,7 +79,7 @@ def test_scope_roundtrips_keep_singleton(name):
@pytest.mark.parametrize("scope", ["warp", "cuda.warp",
T.cuda.ThreadScope("warp")])
def test_instruction_normalizes_scope_before_constructing_call(scope):
- tensor = decl_tensor((16,), "float32", scope="local", layout="default")
+ tensor = Var("tensor", TensorType((16,), "float32", scope="local",
layout="default"))
call = T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], tensor, tensor, scope=scope)
tag = ThreadScope("cuda.warp")
assert call.attrs.scope.same_as(tag)
@@ -106,7 +106,7 @@ def test_raw_scope_attrs_are_required_and_typed():
def test_cross_backend_scopes_are_rejected_at_each_boundary():
- tensor = decl_tensor((16,), "float32", scope="local", layout="default")
+ tensor = Var("tensor", TensorType((16,), "float32", scope="local",
layout="default"))
trn_scope = ThreadScope("trn.thread")
for scope in (trn_scope, "trn.thread"):
with pytest.raises(ValueError, match="belongs to trn"):
@@ -123,7 +123,7 @@ def
test_cross_backend_scopes_are_rejected_at_each_boundary():
raw.validate()
with pytest.raises(ValueError, match="does not match target"):
DispatchContext(tvm_tirx.target.Target("cuda"), trn_scope, {}, {})
- trn_tensor = decl_tensor((16,), "float32", scope="trn.sbuf",
layout="default")
+ trn_tensor = Var("tensor", TensorType((16,), "float32", scope="trn.sbuf",
layout="default"))
trn_call = T.trn.tile.tensortensor(trn_tensor, trn_tensor, trn_tensor)
func = Function([], Evaluate(trn_call)).with_attr("target",
tvm_tirx.target.Target("cuda"))
with pytest.raises(ValueError, match="does not match target"):
diff --git a/tirx/tests/python/tirx/test_tile_namespace.py
b/tirx/tests/python/tirx/test_tile_namespace.py
index eed33aad83..5526fb9346 100644
--- a/tirx/tests/python/tirx/test_tile_namespace.py
+++ b/tirx/tests/python/tirx/test_tile_namespace.py
@@ -27,7 +27,7 @@ from tvm_tirx.ir import assert_structural_equal, load_json,
save_json
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.testing import env
-from tvm_tirx.tirx import decl_tensor
+from tvm_tirx.tirx import TensorType, Var
from tvm_tirx.tirx.layout import Axis, S, TileLayout
from tvm_tirx.tirx.script import ir_builder as I
from tvm_tirx.tirx.thread_scope import ThreadScope
@@ -42,12 +42,12 @@ def test_generic_layout_without_backend_autoload():
import sys
import tvm_tirx
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import decl_tensor
+from tvm_tirx.tirx import Var, TensorType
from tvm_tirx.tirx.layout import Axis, S, TileLayout
assert Axis.m.backend is None
assert T.Axis.m.same_as(Axis.m)
assert TileLayout(S[8]).apply(3, shape=[8])["m"] == 3
-assert decl_tensor((8,), "float32").layout is None
+assert Var("tensor", TensorType((8,), "float32")).layout is None
tensor_type = T.Tensor((8,), "float32")
assert tensor_type.layout is None
tvm_tirx.ir.assert_structural_equal(tensor_type.effective_layout(),
TileLayout(S[8]))
@@ -98,7 +98,7 @@ def test_backend_axis_and_scope_identity():
def test_omitted_and_explicit_none_layout(scope):
assert T.Tensor((4, 8), "float32", scope=scope).layout is None
assert T.Tensor((4, 8), "float32", scope=scope, layout=None).layout is None
- assert decl_tensor((4, 8), "float32", scope=scope).layout is None
+ assert Var("tensor", TensorType((4, 8), "float32", scope=scope)).layout is
None
explicit = T.Tensor((4, 8), "float32", scope=scope, layout="default")
assert_structural_equal(explicit.layout, TileLayout(S[4, 8]))
with IRBuilder():
@@ -121,7 +121,7 @@ def test_effective_layout_on_tensor_type(shape, strides):
def test_effective_global_layout_preserves_source_and_region():
- tensor = decl_tensor((8, 16), "float32", elem_offset=13)
+ tensor = Var("tensor", TensorType((8, 16), "float32", elem_offset=13))
before = save_json(tensor)
region = tensor[2:4, 3:7]
layout = region.source.ty.effective_layout()
@@ -130,14 +130,14 @@ def
test_effective_global_layout_preserves_source_and_region():
assert int(layout.apply(*starts, shape=tensor.shape)["m"]) == 35
assert int(layout.apply(starts[0] + 1, starts[1] + 3,
shape=tensor.shape)["m"]) == 54
assert save_json(tensor) == before
- compact = decl_tensor((8, 16), "float32", strides=(16, 1))
+ compact = Var("tensor", TensorType((8, 16), "float32", strides=(16, 1)))
assert_structural_equal(compact.ty.effective_layout(), layout)
- strided = decl_tensor((8, 16), "float32", strides=(32, 1))
+ strided = Var("tensor", TensorType((8, 16), "float32", strides=(32, 1)))
with pytest.raises(ValueError, match="noncompact"):
strided.ty.effective_layout()
with pytest.raises(ValueError, match="explicit layout"):
- decl_tensor((8, 16), "float32", scope="shared").ty.effective_layout()
- explicit = decl_tensor((8, 16), "float32", scope="shared", layout=layout)
+ Var("tensor", TensorType((8, 16), "float32",
scope="shared")).ty.effective_layout()
+ explicit = Var("tensor", TensorType((8, 16), "float32", scope="shared",
layout=layout))
assert explicit.ty.effective_layout().same_as(layout)
@@ -209,7 +209,7 @@ def test_none_layout_dtype_view_accepts_symbolic_extent():
def test_tile_call_rejects_cross_backend_scope():
- tensor = decl_tensor((8,), "float32", scope="local", layout="default")
+ tensor = Var("tensor", TensorType((8,), "float32", scope="local",
layout="default"))
with pytest.raises(ValueError, match="scope belongs to trn"):
T.cuda.tile.apply(
T.ptx["add.f32"], tensor[:], tensor[:], tensor[:],
scope=T.trn.ThreadScope("thread")
diff --git a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
index 571faf8e53..caeea7e4f9 100644
--- a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
+++ b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
@@ -47,8 +47,12 @@ def _make_dsmem_dispatch_call(shape, dtype, src_layout,
dst_layout, chunks=1):
# original physical maps; flattening only changes the logical shape.
canonical = (math.prod(shape),) if chunks == 1 else shape
chunk_bytes = math.prod(shape) * tvm_ffi.dtype(dtype).bits // 8 // chunks
- src = tvm_tirx.tirx.decl_tensor(canonical, dtype, "A", scope="shared.dyn",
layout=src_layout)
- dst = tvm_tirx.tirx.decl_tensor(canonical, dtype, "B", scope="shared.dyn",
layout=dst_layout)
+ src = tvm_tirx.tirx.Var(
+ "A", tvm_tirx.tirx.TensorType(canonical, dtype, scope="shared.dyn",
layout=src_layout)
+ )
+ dst = tvm_tirx.tirx.Var(
+ "B", tvm_tirx.tirx.TensorType(canonical, dtype, scope="shared.dyn",
layout=dst_layout)
+ )
op_call = T.cuda.tile.apply(
T.ptx[BULK], dst, src, chunk_bytes=chunk_bytes, mbar=Var("mbar",
"handle"), remote_cta_id=1
)
@@ -309,7 +313,9 @@ def test_dsmem_dispatch_missing_config():
from tvm_tirx.ir import Range
layout = TileLayout(S[64])
- buf = tvm_tirx.tirx.decl_tensor((64,), "float16", "A", scope="shared.dyn",
layout=layout)
+ buf = tvm_tirx.tirx.Var(
+ "A", tvm_tirx.tirx.TensorType((64,), "float16", scope="shared.dyn",
layout=layout)
+ )
br = TensorRegion(buf, [Range.from_min_extent(0, 64)],
ty=TensorRegionType())
target = tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_90a"})
sctx = DispatchContext(target, ThreadScope("cuda.thread"), {}, {})
diff --git a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
index acd4180769..45ba41d18c 100644
--- a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
+++ b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
@@ -848,35 +848,39 @@ def _recipe_views(kwargs, recipe):
shape = tuple(reversed(recipe.dims))
strides = tuple(reversed((1, *(s * 8 // bits for s in recipe.strides))))
box = tuple(reversed(recipe.boxes))
- g = tvm_tirx.tirx.decl_tensor(
- kwargs["g_shape"], dtype, "global_storage",
layout=kwargs.get("g_layout")
+ g = tvm_tirx.tirx.Var(
+ "global_storage",
+ tvm_tirx.tirx.TensorType(kwargs["g_shape"], dtype,
layout=kwargs.get("g_layout")),
)
sshape = kwargs.get("s_shape", kwargs["g_shape"])
- s = tvm_tirx.tirx.decl_tensor(
- sshape,
- dtype,
+ s = tvm_tirx.tirx.Var(
"shared_storage",
- scope="shared.dyn",
- layout=kwargs.get("s_layout"),
- data_alignment=1024,
+ tvm_tirx.tirx.TensorType(
+ sshape, dtype, scope="shared.dyn", layout=kwargs.get("s_layout"),
data_alignment=1024
+ ),
)
- gv = tvm_tirx.tirx.decl_tensor(
- shape, recipe.dtype, "global_view", data=g.data,
layout=_plain_layout(shape, strides)
+ gv = tvm_tirx.tirx.Var(
+ "global_view",
+ tvm_tirx.tirx.TensorType(
+ shape, recipe.dtype, data=g.data, layout=_plain_layout(shape,
strides)
+ ),
)
assert int(sa.min()) * original_bits % bits == 0
base = int(sa.min()) * original_bits // bits
layout = _plain_layout(box)
if recipe.enums[1]:
layout = ComposeLayout(7 - int(math.log2(bits)), recipe.enums[1], 3,
layout)
- sv = tvm_tirx.tirx.decl_tensor(
- box,
- recipe.dtype,
+ sv = tvm_tirx.tirx.Var(
"shared_view",
- data=s.data,
- scope="shared.dyn",
- layout=layout,
- elem_offset=base,
- data_alignment=1024,
+ tvm_tirx.tirx.TensorType(
+ box,
+ recipe.dtype,
+ data=s.data,
+ scope="shared.dyn",
+ layout=layout,
+ elem_offset=base,
+ data_alignment=1024,
+ ),
)
# Compare address PAIRS, not separate address sets: a transpose can
preserve
# both sets while copying every element to the wrong destination.
@@ -946,17 +950,22 @@ def _direct(
**fields,
):
s_shape = s_shape or g_shape
- g = tvm_tirx.tirx.decl_tensor(
- g_shape, dtype, "A", layout=g_layout or _plain_layout(g_shape),
elem_offset=g_elem_offset
+ g = tvm_tirx.tirx.Var(
+ "A",
+ tvm_tirx.tirx.TensorType(
+ g_shape, dtype, layout=g_layout or _plain_layout(g_shape),
elem_offset=g_elem_offset
+ ),
)
- s = tvm_tirx.tirx.decl_tensor(
- s_shape,
- dtype,
+ s = tvm_tirx.tirx.Var(
"A_smem",
- scope="shared.dyn",
- layout=s_layout or _plain_layout(s_shape),
- elem_offset=s_elem_offset,
- data_alignment=1024,
+ tvm_tirx.tirx.TensorType(
+ s_shape,
+ dtype,
+ scope="shared.dyn",
+ layout=s_layout or _plain_layout(s_shape),
+ elem_offset=s_elem_offset,
+ data_alignment=1024,
+ ),
)
box = box or tuple(e for _, e in s_region) if s_region else box or s_shape
coords = coords or (0,) * len(box)
@@ -1189,9 +1198,12 @@ def
test_caller_descriptor_view_preserves_original_addresses(kwargs, recipe):
def test_caller_explicit_issue_coordinates_and_pointer():
# Original shape/layout retained, now explicitly ask for 512 one-row boxes.
shape = (512, 64)
- g = tvm_tirx.tirx.decl_tensor(shape, "float16", "A", layout="default")
- s = tvm_tirx.tirx.decl_tensor(
- (512, 1, 64), "float16", "S", scope="shared", layout="default",
data_alignment=128
+ g = tvm_tirx.tirx.Var("A", tvm_tirx.tirx.TensorType(shape, "float16",
layout="default"))
+ s = tvm_tirx.tirx.Var(
+ "S",
+ tvm_tirx.tirx.TensorType(
+ (512, 1, 64), "float16", scope="shared", layout="default",
data_alignment=128
+ ),
)
impl, host, _ = _emit(_load(2), g, s, box=(1, 64), coords=(0, 0),
coord_steps=((1, 0),))
assert _count_tma(impl).total == 512
@@ -1493,11 +1505,17 @@ def test_caller_dynamic_stage_coordinates_and_pointer():
stage: Range.from_min_extent(0, 3),
}
)
- gv = tvm_tirx.tirx.decl_tensor(
- (32, 4096, 64), "float16", "G", layout=_plain_layout((32, 4096, 64),
(64, 2048, 1))
+ gv = tvm_tirx.tirx.Var(
+ "G",
+ tvm_tirx.tirx.TensorType(
+ (32, 4096, 64), "float16", layout=_plain_layout((32, 4096, 64),
(64, 2048, 1))
+ ),
)
- sv = tvm_tirx.tirx.decl_tensor(
- (3, 2, 128, 64), "float16", "S", scope="shared", layout="default",
data_alignment=128
+ sv = tvm_tirx.tirx.Var(
+ "S",
+ tvm_tirx.tirx.TensorType(
+ (3, 2, 128, 64), "float16", scope="shared", layout="default",
data_alignment=128
+ ),
)
impl, host, _ = _emit(
_load(3),
@@ -1910,14 +1928,18 @@ def test_caller_tiles_nested_shared_layout(case_id):
def test_caller_explicit_strided_outer_issues():
kwargs = next(p.values[1] for p in TMA_CASES if p.values[0] ==
"g2s-3d-stride-gap-outer")
- g = tvm_tirx.tirx.decl_tensor(kwargs["g_shape"], "float16", "A",
layout=kwargs["g_layout"])
- s = tvm_tirx.tirx.decl_tensor(
- (8, 1, 32, 64),
- "float16",
+ g = tvm_tirx.tirx.Var(
+ "A", tvm_tirx.tirx.TensorType(kwargs["g_shape"], "float16",
layout=kwargs["g_layout"])
+ )
+ s = tvm_tirx.tirx.Var(
"S",
- scope="shared",
- layout=_plain_layout((8, 1, 32, 64), (4096, 0, 64, 1)),
- data_alignment=128,
+ tvm_tirx.tirx.TensorType(
+ (8, 1, 32, 64),
+ "float16",
+ scope="shared",
+ layout=_plain_layout((8, 1, 32, 64), (4096, 0, 64, 1)),
+ data_alignment=128,
+ ),
)
impl, _, _ = _emit(_load(3), g, s, box=(1, 32, 64), coords=(0, 0, 0),
coord_steps=((1, 0, 0),))
assert _count_tma(impl).total == 8
@@ -1952,9 +1974,14 @@ def
test_descriptor_cache_distinguishes_explicit_dtype_views():
data = Var("A", PtrType(PrimType("uint8"), "global"))
sctx = _context()
for dtype in ("uint8", "uint16"):
- g = tvm_tirx.tirx.decl_tensor((256,), dtype, "A_view", data=data,
layout="default")
- s = tvm_tirx.tirx.decl_tensor(
- (256,), dtype, "S", scope="shared", layout="default",
data_alignment=128
+ g = tvm_tirx.tirx.Var(
+ "A_view", tvm_tirx.tirx.TensorType((256,), dtype, data=data,
layout="default")
+ )
+ s = tvm_tirx.tirx.Var(
+ "S",
+ tvm_tirx.tirx.TensorType(
+ (256,), dtype, scope="shared", layout="default",
data_alignment=128
+ ),
)
_emit(_load(1), g, s, box=(256,), coords=(0,), sctx=sctx)
_emit(_load(1), g, s, box=(256,), coords=(0,), sctx=sctx)
diff --git a/tirx/tests/python/tirx/tile/test_dispatcher.py
b/tirx/tests/python/tirx/tile/test_dispatcher.py
index 6c02f195fc..0163d5df9b 100644
--- a/tirx/tests/python/tirx/tile/test_dispatcher.py
+++ b/tirx/tests/python/tirx/tile/test_dispatcher.py
@@ -32,14 +32,14 @@ from tvm_tirx.ir import (
register_op_attr,
)
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx import Function, decl_tensor
+from tvm_tirx.tirx import Function, TensorType, Var
from tvm_tirx.tirx.tile import DispatchContext
from tvm_tirx.tirx.tile.dispatch import run_dispatch
from tvm_tirx.tirx.tile.instruction import SPECS
def test_selected_lowerer_reports_context_and_cause(monkeypatch):
- a = decl_tensor((16,), "float32", scope="local", layout="default")
+ a = Var("tensor", TensorType((16,), "float32", scope="local",
layout="default"))
call = T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], a, a)
def failing(view, context):
@@ -56,7 +56,7 @@ def
test_selected_lowerer_reports_context_and_cause(monkeypatch):
def test_backend_mismatch_is_rejected():
- a = decl_tensor((16,), "float32", scope="local", layout="default")
+ a = Var("tensor", TensorType((16,), "float32", scope="local",
layout="default"))
context = SimpleNamespace(
target=SimpleNamespace(kind=SimpleNamespace(name="trn")),
scope_kind="thread"
)
@@ -67,7 +67,7 @@ def test_backend_mismatch_is_rejected():
@pytest.mark.parametrize("backend", ["cuda", "trn"])
@pytest.mark.parametrize("scope", ["kernel", "unknown"])
def test_instruction_scope_contract(backend, scope):
- tensor = decl_tensor((16,), "float32", scope="local", layout="default")
+ tensor = Var("tensor", TensorType((16,), "float32", scope="local",
layout="default"))
instruction = (
(lambda *args, **kw: T.cuda.tile.apply(T.ptx["add.f32"], *args, **kw))
if backend == "cuda"
@@ -186,7 +186,7 @@ def test_dispatch_requires_stmt_result():
@pytest.mark.parametrize("result", [Function([], Evaluate(0)), const(0), None])
def test_selected_lowerer_requires_stmt_result(monkeypatch, result):
- tensor = decl_tensor((16,), "float32", scope="local", layout="default")
+ tensor = Var("tensor", TensorType((16,), "float32", scope="local",
layout="default"))
call = T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], tensor, tensor)
monkeypatch.setattr(SPECS[call.op.name], "lower", lambda view, context:
result)
context = SimpleNamespace(target=tvm_tirx.target.Target("cuda"),
scope_kind="thread")
diff --git a/tirx/tests/python/tirx/tile/trn/test_copy.py
b/tirx/tests/python/tirx/tile/trn/test_copy.py
index 8862ea9241..4e82e44e0f 100644
--- a/tirx/tests/python/tirx/tile/trn/test_copy.py
+++ b/tirx/tests/python/tirx/tile/trn/test_copy.py
@@ -25,7 +25,7 @@ from tvm_tirx.backend.trn.tile.instruction_generator import
InstructionGenerator
from tvm_tirx.ir import assert_structural_equal as _assert_structural_equal
from tvm_tirx.script import tirx as T
from tvm_tirx.sym import Analyzer
-from tvm_tirx.tirx import decl_tensor
+from tvm_tirx.tirx import TensorType, Var
from tvm_tirx.tirx.layout import S, TileLayout
target = tvm_tirx.target.Target("aws/trn1/trn1.2xlarge")
@@ -53,7 +53,7 @@ def assert_structural_equal(lhs, rhs, *args, **kwargs):
def test_instruction_generator_compact_global_layout():
- tensor = decl_tensor((8, 16), "float32")
+ tensor = Var("tensor", TensorType((8, 16), "float32"))
region = tensor[0:8, 0:16]
generator = InstructionGenerator((region,), Analyzer())
layout = generator.split_layout_views[region]
