You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/17 20:13:11 UTC
(tvm) 04/10: [REFACTOR][TIRx] Remove redundant helpers and construct tensor variables directly (#20776)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx-standalone-root
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 98be1c59bdaec59046369efa646668668cf89896
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Sat Oct 17 15:35:59 2026 -0400
[REFACTOR][TIRx] Remove redundant helpers and construct tensor variables directly (#20776)
---
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 +-
31 files changed, 316 insertions(+), 422 deletions(-)
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)
+
+
+@pytest.mark.parametrize("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]