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 <[email protected]> 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) + + [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]
