This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new aa95685731 [TIRx] Own expression methods and properties on TensorType
(#20622)
aa95685731 is described below
commit aa956857314861a8859eba102b02aa660ebf0a0c
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 9 08:29:49 2026 -0400
[TIRx] Own expression methods and properties on TensorType (#20622)
Declare Python expression methods and read-only properties on their type
with `Type.__expr_methods__` and `Type.__expr_properties__`. Generic
expression lookup preserves ordinary and reflected attributes, then
resolves declared getters or binds the original operand to an exposed
type method.
Move the complete tensor method and property surface into `TensorType`
in the tensor module, preserving `A.data`, metadata conveniences, and
canonical pointer semantics. This removes the global method and property
patches on `Var` while retaining compatibility imports.
---
docs/tirx/native_basics/cuda/buffers.rst | 21 +-
python/tvm/ir/expr.py | 53 ++
python/tvm/ir/type.py | 17 +-
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/_buffer_view.py | 20 +-
python/tvm/tirx/buffer.py | 544 +--------------------
python/tvm/tirx/op/__init__.py | 2 +-
python/tvm/tirx/{buffer.py => tensor.py} | 293 +++++------
tests/python/sym/test_sym_deduce_bound.py | 2 +-
tests/python/tirx-base/test_tir_buffer.py | 2 +-
.../operator/tile_primitive/test_dispatcher.py | 2 +-
.../script/test_tirx_script_source_locations.py | 8 +-
tests/python/tirx/test_hint.py | 2 +-
tests/python/tirx/test_op.py | 2 +-
14 files changed, 277 insertions(+), 693 deletions(-)
diff --git a/docs/tirx/native_basics/cuda/buffers.rst
b/docs/tirx/native_basics/cuda/buffers.rst
index 314b577b96..e39b527462 100644
--- a/docs/tirx/native_basics/cuda/buffers.rst
+++ b/docs/tirx/native_basics/cuda/buffers.rst
@@ -414,7 +414,24 @@ Tensor variable APIs
A tensor variable is an ``ir.Var`` carrying ``tirx.TensorType`` metadata
(see *Declaring buffers* above), so most of
its methods are *compile-time* reshapes/reinterprets that change index
arithmetic
-or hand you a pointer — they emit no runtime op of their own. The common ones:
+or hand you a pointer — they emit no runtime op of their own.
+
+``TensorType.__expr_methods__`` explicitly names the Python methods available
+through an expression: ``B.view(...)`` binds the operand exactly as
+``B.ty.view(B, ...)``. ``TensorType.__expr_properties__`` maps property names
+to getters taking the type and original expression. It preserves conveniences
+such as ``B.shape``, ``B.strides``, ``B.dtype``, ``B.byte_offset``, ``B.sub``
and
+``B.data`` without modifying the shared ``Var`` class. ``B.dtype`` returns a
+runtime ``DataType``; ``B.ty.dtype`` stores a ``PrimType``.
+
+Existing expression attributes and reflected fields take precedence, followed
+by declared properties and then declared methods. Declared properties are
+read-only. Tensor methods and properties require an ordinary tensor variable;
+merely retrieving a method does not construct IR. A property getter executes
+when the property is read: for example, ``B.data`` constructs the canonical
+physical-pointer projection.
+
+The common methods and properties:
.. list-table::
:header-rows: 1
@@ -423,7 +440,7 @@ or hand you a pointer — they emit no runtime op of their
own. The common ones:
* - Method
- What it is
* - ``B.data``
- - the raw data pointer (a ``Var``); prints as ``B_ptr``
+ - the physical-pointer projection (a ``Call``); lowers to ``B_ptr``
* - ``B.ptr_to([i, j])``
- a typed pointer to an element (``address_of``); prints as ``&B_ptr[…]``
* - ``B.vload([i], dtype="float32x4")`` / ``B.vstore([i], v)``
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index e2da21fb75..a8cd15afba 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -17,6 +17,8 @@
"""Common expressions data structures in the IR."""
from collections.abc import Callable
+from functools import partial
+from inspect import getattr_static
from numbers import Number
import tvm_ffi
@@ -28,6 +30,8 @@ from . import _ffi_api, _tensor_expr_overload
from ._constant import const
from .base import Node, Scriptable, Span
+_ATTRIBUTE_MISSING = object()
+
def _convert_subscript_index(index):
"""Convert Python indexing syntax into an FFI subscript descriptor."""
@@ -51,6 +55,55 @@ class Expr(Node):
span: Span | None
ty: "tvm.ir.Type"
+ def __getattr__(self, name):
+ # Reflected fields are ordinary Python properties and have already had
+ # their chance to resolve before this type-directed fallback.
+ try:
+ ty = object.__getattribute__(self, "ty")
+ except AttributeError:
+ ty = None
+ if ty is not None and not name.startswith("_"):
+ if name in ty.__expr_properties__:
+ getter = ty.__expr_properties__[name]
+ if not callable(getter):
+ raise TypeError(
+ f"Declared expression property {name!r} must have a
callable getter"
+ )
+ return getter(ty, self)
+ if name in ty.__expr_methods__:
+ method = getattr(ty, name)
+ if not callable(method):
+ raise TypeError(f"Declared expression method {name!r} must
be callable")
+ bound = partial(method, self)
+ bound.__doc__ = method.__doc__
+ return bound
+ raise AttributeError(f"{type(self).__name__!r} object has no attribute
{name!r}")
+
+ def __setattr__(self, name, value):
+ # Do not invoke getters to distinguish real attributes from the
fallback.
+ # Existing descriptors and instance attributes keep their usual
behavior.
+ if (
+ not name.startswith("_")
+ and getattr_static(self, name, _ATTRIBUTE_MISSING) is
_ATTRIBUTE_MISSING
+ ):
+ try:
+ ty = object.__getattribute__(self, "ty")
+ except AttributeError:
+ ty = None
+ if ty is not None and name in ty.__expr_properties__:
+ raise AttributeError(f"Expression property {name!r} is
read-only")
+ super().__setattr__(name, value)
+
+ def __dir__(self):
+ names = set(super().__dir__())
+ try:
+ ty = object.__getattribute__(self, "ty")
+ except AttributeError:
+ return sorted(names)
+ names.update(name for name in ty.__expr_methods__ if not
name.startswith("_"))
+ names.update(name for name in ty.__expr_properties__ if not
name.startswith("_"))
+ return sorted(names)
+
def __getitem__(self, index):
if isinstance(self.ty, tvm.ir.MissingType):
# Preserve Relax's pre-normalization tuple access: operator calls
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 2d67c0ade2..694f4dbd41 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -16,6 +16,9 @@
# under the License.
"""Unified type system in the project."""
+from collections.abc import Callable
+from typing import ClassVar
+
import tvm_ffi
from . import _ffi_api
@@ -24,7 +27,19 @@ from .base import Node, Scriptable
@tvm_ffi.register_object("ir.Type")
class Type(Node, Scriptable):
- """The base class of all types."""
+ """The base class of all types.
+
+ ``__expr_methods__`` names methods exposed on expressions of this type.
+ Each method receives the type instance and then the expression operand:
+ ``expr.method(...)`` is equivalent to ``expr.ty.method(expr, ...)``.
+ ``__expr_properties__`` maps read-only expression property names to
+ getters taking ``(type_instance, expression)``. Ordinary expression
+ attributes take precedence, then declared properties, then methods.
+ Subclasses inherit these declarations or replace them explicitly.
+ """
+
+ __expr_methods__: tuple[str, ...] = ()
+ __expr_properties__: ClassVar[dict[str, Callable]] = {}
@staticmethod
def missing():
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 7306911876..a7d9c2a4c8 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -26,7 +26,7 @@ tvm.script.register_dialect("tirx", "tvm.tirx.script",
builder_path="tvm.tirx.sc
from tvm.ir import Expr
from tvm.ir import const
-from .buffer import (
+from .tensor import (
BufferAccessKind,
TensorType,
tensor_data_ptr,
diff --git a/python/tvm/tirx/_buffer_view.py b/python/tvm/tirx/_buffer_view.py
index 8bef33dded..84628330b7 100644
--- a/python/tvm/tirx/_buffer_view.py
+++ b/python/tvm/tirx/_buffer_view.py
@@ -68,7 +68,7 @@ def _redecl(buf: Var, shape, layout, *, dtype=None,
elem_offset=None, addr_offse
def view(buf: Var, *args, **kwargs) -> Var:
- """Implement :meth:`Var.view`."""
+ """Implement :meth:`TensorType.view`."""
def _infer_shape(shape):
shape = list(shape)
@@ -124,7 +124,7 @@ def view(buf: Var, *args, **kwargs) -> Var:
def local(buf: Var, *shape, layout=None) -> Var:
- """Implement :meth:`Var.local`."""
+ """Implement :meth:`TensorType.local`."""
if not shape:
if buf.layout is None:
raise ValueError(
@@ -151,7 +151,7 @@ def local(buf: Var, *shape, layout=None) -> Var:
def permute(buf: Var, *dims) -> Var:
- """Implement :meth:`Var.permute`."""
+ """Implement :meth:`TensorType.permute`."""
new_shape = [buf.shape[d] for d in dims]
layout, swizzle = _surgery_parts(buf)
grouped, seps = layout.group(list(buf.shape))
@@ -162,7 +162,7 @@ def permute(buf: Var, *dims) -> Var:
def rearrange(buf: Var, pattern: str, /, **sizes) -> Var:
- """Implement :meth:`Var.rearrange`."""
+ """Implement :meth:`TensorType.rearrange`."""
def _groups(side):
out = []
@@ -220,12 +220,12 @@ def rearrange(buf: Var, pattern: str, /, **sizes) -> Var:
def sub(buf: Var) -> SubIndexer:
- """Return the indexer used by :attr:`Var.sub`."""
+ """Return the indexer used by ``expr.sub``."""
return SubIndexer(buf)
def tile(buf: Var, *specs) -> TileIndexer:
- """Implement :meth:`Var.tile`."""
+ """Implement :meth:`TensorType.tile`."""
if specs and isinstance(specs[0], int | Integral):
if len(specs) != 2:
raise ValueError("tile(dim, factors) takes a dim and a factors
tuple")
@@ -247,7 +247,7 @@ def tile(buf: Var, *specs) -> TileIndexer:
def chunk(buf: Var, spec) -> ChunkIndexer:
- """Implement :meth:`Var.chunk`."""
+ """Implement :meth:`TensorType.chunk`."""
if not isinstance(spec, tuple | list):
raise ValueError(f"chunk: spec must be a per-dim tuple, got {spec!r}")
if len(spec) != len(buf.shape):
@@ -497,7 +497,7 @@ def _view_narrow(buf: Var, dim, start, length):
class SubIndexer:
- """Indexer returned by :attr:`Var.sub`."""
+ """Indexer returned by ``expr.sub``."""
def __init__(self, buffer: Var):
self._buffer = buffer
@@ -548,7 +548,7 @@ class SubIndexer:
class ChunkIndexer:
- """Indexer returned by :meth:`Var.chunk`."""
+ """Indexer returned by :meth:`TensorType.chunk`."""
def __init__(self, buffer: Var, spec):
self._buffer = buffer
@@ -576,7 +576,7 @@ class ChunkIndexer:
class TileIndexer:
- """Indexer returned by :meth:`Var.tile`."""
+ """Indexer returned by :meth:`TensorType.tile`."""
def __init__(self, buffer: Var, specs):
self._buffer = buffer
diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py
index 2d744397db..9e8d21fd9c 100644
--- a/python/tvm/tirx/buffer.py
+++ b/python/tvm/tirx/buffer.py
@@ -14,528 +14,22 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Abstraction for array data structures."""
-
-from enum import IntEnum
-from numbers import Integral
-
-import tvm_ffi
-
-import tvm
-from tvm.ir import Call, PointerType, PrimType, Type, Var
-from tvm.runtime import convert
-
-from . import _buffer_view, _ffi_api
-
-_REARRANGE_PATTERN_UNSET = object()
-
-
-@tvm_ffi.register_object("tirx.TensorType")
-class TensorType(Type):
- """The structural type carried by an ordinary TIRx tensor variable."""
-
- dtype: PrimType
- storage_scope: str
- shape: list
- strides: list
- elem_offset: tvm.ir.Expr
- data_alignment: int
- offset_factor: int
- layout: object | None
- allocated_addr: list
-
-
-def is_tensor_var(value) -> bool:
- """Return whether ``value`` is an ordinary Var carrying TensorType.
-
- Unlike ``isinstance(value, Var)``, this predicate distinguishes tensor
- variables from scalar or pointer variables.
- """
-
- return isinstance(value, tvm.ir.Var) and isinstance(value.ty, TensorType)
-
-
-class BufferAccessKind(IntEnum):
- """Buffer access modes accepted by :func:`buffer_access_ptr`."""
-
- READ = 1
- WRITE = 2
-
-
-class _TensorMethods:
- """Operations on tensor variables carrying layout and storage metadata.
-
- Do not construct directly, use :py:func:`~decl_tensor` instead.
- See the documentation of :py:func:`decl_tensor` for more details.
-
- See Also
- --------
- decl_tensor : Declare a buffer
- """
-
- def access_ptr(self, access_mask, ptr_type="handle", content_lanes=1,
offset=0, extent=None):
- """Get an access pointer to the head of buffer.
-
- This is the recommended method to get buffer data
- ptress when interacting with external functions.
-
- Parameters
- ----------
- access_mask : int
- The access pattern MASK. Indicate whether the
- access will read or write to the data content.
-
- ptr_type : str or tvm.ir.Type, optional
- The data type of the result pointer. Do not specify
- unless we want to cast pointer to specific type.
-
- content_lanes: int, optional
- The number of lanes for the data type. This value
- is greater than one for vector types.
-
- offset: Expr, optional
- The offset of pointer. We can use it to offset by
- the number of elements from the address of ptr.
-
- extent: Expr, optional
- The extent of pointer.
-
- Examples
- --------
- .. code-block:: python
-
- # Get access ptr for read
- buffer.access_ptr("r")
- # Get access ptr for read/write with bitmask
- buffer.access_ptr(BufferAccessKind.READ | BufferAccessKind.WRITE)
- # Get access ptr for read/write with str flag
- buffer.access_ptr("rw")
- # Get access ptr for read with offset
- buffer.access_ptr("r", offset = 100)
- # Get access ptr for read with extent
- buffer.access_ptr("r", extent = 100)
- """
- if isinstance(access_mask, str):
- mask = 0
- for value in access_mask:
- if value == "r":
- mask = mask | BufferAccessKind.READ
- elif value == "w":
- mask = mask | BufferAccessKind.WRITE
- else:
- raise ValueError(f"Unknown access_mask {access_mask}")
- access_mask = mask
- if isinstance(ptr_type, str):
- ptr_type = (
- PointerType(PrimType("void"))
- if ptr_type == "handle"
- else PointerType(PrimType(ptr_type))
- )
- elif isinstance(ptr_type, PrimType):
- ptr_type = PointerType(ptr_type)
- offset = convert(offset)
- extent = convert(extent)
- return _ffi_api.TensorAccessPtr(
- self,
- access_mask,
- ptr_type,
- content_lanes,
- offset,
- extent, # type: ignore
- )
-
- def vload(self, begin, dtype=None):
- """Generate an Expr that loads dtype from begin index.
-
- Parameters
- ----------
- begin : Array of Expr
- The beginning index in unit of Var.dtype
-
- dtype : str
- The data type to be loaded,
- can be vector type which have lanes that is multiple of Var.dtype
-
- Returns
- -------
- load : Expr
- The corresponding load expression.
- """
- begin = (begin,) if isinstance(begin, int) or
tvm.ir.is_prim_expr(begin) else begin
- dtype = dtype if dtype else self.ty.dtype
- return _ffi_api.TensorVLoad(self, begin, dtype) # type: ignore
-
- def vstore(self, begin, value):
- """Generate a Stmt that store value into begin index.
-
- Parameters
- ----------
- begin : Array of Expr
- The beginning index in unit of Var.dtype
-
- value : Expr
- The value to be stored.
-
- Returns
- -------
- store : Stmt
- The corresponding store stmt.
- """
- begin = (begin,) if isinstance(begin, int) or
tvm.ir.is_prim_expr(begin) else begin
- return _ffi_api.TensorVStore(self, begin, value) # type: ignore
-
- def scope(self):
- """Return the storage scope associated with this buffer.
- Returns
- -------
- scope : str
- The storage scope associated with this buffer.
- """
- return _ffi_api.TensorStorageScope(self) # type: ignore
-
- def get_flattened_buffer(self):
- """Generate a Var that is a flattened version of this buffer.
-
- Returns
- -------
- flattened : Var
- The corresponding flat buffer.
- """
- return _ffi_api.TensorGetFlattenedTensor(self) # type: ignore
-
- def with_allocated_addr(self, allocated_addr):
- """Return a new buffer with the allocated address."""
- return _ffi_api.TensorWithAllocatedAddr(self, allocated_addr) # type:
ignore
-
- def with_dtype(self, dtype):
- """Return a new buffer with the dtype."""
- return _ffi_api.TensorWithDtype(self, dtype) # type: ignore
-
- def offset_of(self, indices):
- """Determine the offset of the provided indices in the flattened
buffer.
-
- Parameters
- ----------
- indices : Union[Expr, List[Expr]]
-
- The indices of the element in the original buffer.
-
- Returns
- -------
- flattened_indices: List[Expr]
-
- The offset indices of the element in the flattened buffer.
- """
- return _ffi_api.TensorOffsetOf(self, indices) # type: ignore
-
- @property
- def byte_offset(self):
- """Get the byte offset of the buffer."""
- return self.ty.elem_offset * tvm.DataType(self.ty.dtype).bits // 8
-
- def is_scalar(self, alloc_or_decl=True):
- """Check if the buffer is a scalar.
-
- Parameters
- ----------
- alloc_or_decl : bool, optional
- Whether to consider alloc_scalar and decl_scalar as scalar. True
for alloc_scalar,
- False for decl_scalar.
-
- Returns
- -------
- bool: True if the buffer is a scalar, False otherwise.
- """
- return _ffi_api.TensorIsScalar(self, alloc_or_decl)
-
- def ptr_to(self, indices):
- """Get the pointer to the buffer at the given indices (logical
indices).
-
- Note that the bufferload inside requires LowerTIPp pass to apply the
layout to get the physical indices.
- """ # noqa: E501
- assert len(indices) == len(self.ty.shape), (
- f"The number of indices {indices} does not match "
- f"the shape of the buffer {self.ty.shape}"
- )
- return tvm.tirx.address_of(self[tuple(indices)])
-
- def view(self, *args, **kwargs) -> "Var":
- """Creates a new view of the buffer. (used by parser)
-
- Supported signatures are ``view(*shape, layout=None)``, where shape
can contain
- ``-1`` to indicate that the dimension size is auto-inferred, and
- ``view(dtype: Union[str, tvm.DataType])``.
-
- Returns
- -------
- view : Var
- The corresponding view buffer.
- """
-
- return _buffer_view.view(self, *args, **kwargs)
-
- def local(self, *shape, layout=None) -> "Var":
- """Create a thread-local view of this buffer.
-
- By default, both the inferred and explicit-shape forms address the
- raw physical storage span. ``local()[k]`` is the k-th physical
- storage element, including any gaps or layout offset, while
- ``local(d0, d1, ...)`` is a row-major reshape of that same span.
- Pass ``layout=`` to request a mediated view explicitly. This is an
- escape hatch whose shape is interpreted by the supplied layout. When
- that shape is explicit, the parent buffer does not need a layout.
-
- When called with no shape arguments, auto-infers a 1D shape from
- the span of the parent layout's non-thread component (i.e.
- ``self.layout.storage().span()``). The explicit-``layout=`` form
- instead infers the parent layout's ``storage().size()`` for
- compatibility. Either inference requires the parent buffer to have a
- layout.
-
- Parameters
- ----------
- shape : tuple of Expr
- The shape of the local view for indexing. Without ``layout=``,
- its product must equal the per-thread physical storage span.
- With an explicit layout, the shape is not constrained by the raw
- span. If omitted, a matching 1D shape is computed automatically.
-
- layout : optional
- Override layout. If None, the default (identity) layout is used.
-
- Returns
- -------
- local : Var
- The corresponding local buffer.
- """
- return _buffer_view.local(self, *shape, layout=layout)
-
- def permute(self, *dims) -> "Var":
- """Permute the dimensions of the buffer.
-
- Parameters
- ----------
- dims : tuple of int
- The permutation of dimensions.
-
- Returns
- -------
- permuted : Var
- The buffer with permuted dimensions.
- """
- return _buffer_view.permute(self, *dims)
-
- def rearrange(self, pattern: str = _REARRANGE_PATTERN_UNSET, /, **sizes)
-> "Var":
- """einops-style relayout in one line: ``buf.rearrange("b (2 r) -> 2 b
r")``.
-
- A pure reshape+permute+reshape over the SAME physical bytes, spelled as
- an einops pattern. Lowers to ``view`` (split lhs groups) → ``permute``
- (reorder to rhs atom order) → ``view`` (merge rhs groups), so it
inherits
- whatever the underlying axis machinery does: a plain (unswizzled)
buffer
- collapses to a flat layout, a swizzled buffer keeps its swizzle,
- and a tmem buffer carries ``allocated_addr`` through. It therefore does
- NOT flatten a swizzle atom — the same pattern on a swizzled SMEM buffer
- vs an unswizzled TMEM buffer legitimately yields different physical
- layouts (that is the point: rearrange acts on the operand, not a
string).
-
- ``pattern`` is ``"lhs -> rhs"``; each side is space-separated axis
names,
- with ``(a b)`` grouping a product axis. Every lhs group's product must
- equal that input dim; at most one axis per group may be unknown
(inferred
- from the dim), the rest supplied via ``**sizes``. Cannot express a
- replica (``R[...]``), a stride-fiction/padded view, or a reshape
crossing
- a swizzle-atom boundary — keep those as explicit ``view(layout=...)``.
- """
- if pattern is _REARRANGE_PATTERN_UNSET:
- if "pattern" not in sizes:
- raise TypeError("Var.rearrange() missing required argument:
'pattern'")
- pattern = sizes.pop("pattern")
- return _buffer_view.rearrange(self, pattern, **sizes)
-
- @property
- def sub(self) -> "_buffer_view.SubIndexer":
- """Numpy-style view indexer: ``buf.sub[2, 4:8, ::4]``.
-
- Unlike plain ``buf[...]`` (TensorLoad for scalar indices, extent-1
- TensorRegion dims for tile-primitive operands), ``sub`` follows numpy
- basic-indexing semantics as a *view constructor*: an integer index
- removes the dim (``select``), ``a:b`` narrows it, and ``a::s`` takes
- every s-th element (requires the extent divisible by ``s`` and
- ``a < s``). Trailing dims are kept whole.
- """
- return _buffer_view.sub(self)
-
- def tile(self, *specs) -> "_buffer_view.TileIndexer":
- """Chunk a dim: split it into factors, pick a chunk, keep the rest.
-
- Rank-preserving — the picked dim's remaining factors merge back into
- that one dim, and every other dim is untouched, so N dims in gives N
- dims out. Chunk multiple dims by chaining (dims never shift):
- ``buf.tile(0, (nx, -1))[cx, :].tile(1, (-1, ny))[:, cy]``.
-
- Call as ``tile(dim, factors)`` for one dim, or pass several
- ``(dim, factors)`` specs as sugar for a chain. ``factors`` is the
- tuple the dim splits into (row-major, like :meth:`unflatten`; one
- ``-1`` inferred). The indexer takes one entry per factor: an ``int`` /
- ``Expr`` **picks** it (fixing the chunk, dropping the axis, folding
- its offset) and ``:`` **keeps** it. At least one factor per dim must be
- picked — a pure keep-everything split is :meth:`unflatten`, not a
- chunk::
-
- # 64 rows split into (stripe, warp, row) = (-1, WARPS, 4); this
- # warp's 16 interleaved rows (stripe x row merged):
- buf.tile(1, (-1, WARPS, 4))[:, warp, :]
-
- tile(d, (n, -1))[c, :] # contiguous block c
- tile(d, (-1, n))[:, c] # round-robin chunk c
-
- A picked index may be a dynamic Expr (e.g. a warp id); picking
- several factors of one dim is allowed.
- """
- return _buffer_view.tile(self, *specs)
-
- def chunk(self, spec) -> "_buffer_view.ChunkIndexer":
- """Split dims into equal contiguous chunks and pick a chunk per dim —
- **rank-preserving**. Index the result with ``[picks]``.
-
- ``spec`` is a per-dim tuple (length = rank). Each entry is ``None``
- (leave the dim) or a positive int ``n`` (split that dim, extent ``E``
- with ``E % n == 0``, into ``n`` equal chunks of ``E // n``). Then
- ``chunk(spec)[picks]`` takes one entry per dim: a chunked dim's pick is
- the chunk index (int / Expr) and **narrows that dim** to the chunk's
- ``[c*E//n : (c+1)*E//n)`` range — the dim is kept at ``E // n``, no
- dimension is added; an unchunked dim's pick is a normal index (``:`` /
- int / slice). The result is the *same TensorRegion* as the hand-written
- slice — one line instead of the ``c*k : (c+1)*k`` arithmetic::
-
- X[.., c * k : (c + 1) * k, ..] # before (k = E // n)
- X.chunk((None, .., n, ..))[.., c, ..] # after (k inferred)
- """
- return _buffer_view.chunk(self, spec)
-
-
-# These results expose the complete buffer namespace defined above. Indexer
-# definitions precede this class to avoid a circular buffer/view import.
-
-
-def decl_tensor(
- shape,
- dtype=None,
- name="buffer",
- data=None,
- strides=None,
- elem_offset=None,
- scope="",
- data_alignment=-1,
- offset_factor=0,
- span=None,
- layout="default",
-):
- # pylint: disable=import-outside-toplevel
- from .expr import Var
- from .layout import S, TileLayout
-
- shape = (shape,) if tvm.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 offset_factor != 0 and elem_offset is None:
- shape_ty = shape[0].ty if shape and tvm.ir.is_prim_expr(shape[0]) else
"int32"
- elem_offset = Var(f"{name}_elem_offset", shape_ty)
- storage_scope = scope
- if data is not None:
- if not isinstance(data, tvm.ir.Expr) or not isinstance(data.ty,
PointerType):
- raise TypeError("Tensor data must be an Expr with PointerType")
- if not isinstance(data.ty.element_type, PrimType):
- raise TypeError("Tensor data must point to a primitive type")
- storage_scope = data.ty.storage_scope
- buffer_type = _ffi_api.TensorType( # type: ignore
- storage_scope,
- dtype,
- shape,
- strides,
- elem_offset,
- data_alignment,
- offset_factor,
- layout,
- (),
- span,
- )
- return _ffi_api.TensorVar(name, buffer_type, span) # type: ignore
-
-
-def tensor_data_ptr(tensor, *, ty=None, span=None):
- """Project a tensor variable's physical pointer.
-
- The result type is inferred from its element type and storage scope.
- ``ty`` may supply an explicit result type; ``span`` records the source
location.
- """
- if not is_tensor_var(tensor):
- raise TypeError("tensor_data_ptr expects a Var with TensorType")
- return Call("tirx.tensor_data_ptr", [tensor], ty=ty, span=span)
-
-
-def buffer_data_pointer_type(buffer):
- """Return the pointer type produced by :func:`tensor_data_ptr`."""
-
- if not is_tensor_var(buffer):
- raise TypeError("buffer_data_pointer_type expects a Var with
TensorType")
- return _ffi_api.TensorDataPointerType(buffer)
-
-
-# Tensor variables intentionally retain runtime type key ``ir.Var``. Importing
-# ``tvm.tirx`` therefore augments ``tvm.ir.Var`` process-wide with the legacy
-# buffer operation and metadata surface. Non-buffer Vars reject the metadata
-# properties with AttributeError, preserving correct ``hasattr`` behavior.
-for _name, _value in _TensorMethods.__dict__.items():
- if _name.startswith("__"):
- continue
- if callable(_value) or isinstance(_value, property):
- setattr(tvm.ir.Var, _name, _value)
-
-
-def _tensor_type_field(name):
- def getter(value):
- if not is_tensor_var(value):
- raise AttributeError(f"{name} is only available on a Var with
TensorType")
- return getattr(value.ty, name)
-
- return property(getter)
-
-
-# Preserve the tensor variable metadata surface while keeping TensorType as the
-# single source of truth.
-for _name in (
- "shape",
- "strides",
- "elem_offset",
- "data_alignment",
- "offset_factor",
- "layout",
- "allocated_addr",
-):
- setattr(tvm.ir.Var, _name, _tensor_type_field(_name))
-
-
-def _tensor_dtype_property(value):
- if not is_tensor_var(value):
- raise AttributeError("dtype is only available on a Var with
TensorType")
- # Preserve the tensor variable Python dtype surface. TensorType stores a
- # PrimType, while Python callers historically receive its runtime DataType.
- return value.ty.dtype.dtype
-
-
-tvm.ir.Var.dtype = property(_tensor_dtype_property)
-
-
-# Keep the established ``A.data`` TVMScript surface as syntax sugar. Compiler
-# and builder code calls ``tensor_data_ptr(A)`` directly.
-def _tensor_data_property(value):
- if not is_tensor_var(value):
- raise AttributeError("data is only available on a Var with TensorType")
- return tensor_data_ptr(value)
-
-
-tvm.ir.Var.data = property(_tensor_data_property)
+"""Compatibility imports for the tensor module."""
+
+from .tensor import (
+ BufferAccessKind,
+ TensorType,
+ buffer_data_pointer_type,
+ decl_tensor,
+ is_tensor_var,
+ tensor_data_ptr,
+)
+
+__all__ = [
+ "BufferAccessKind",
+ "TensorType",
+ "buffer_data_pointer_type",
+ "decl_tensor",
+ "is_tensor_var",
+ "tensor_data_ptr",
+]
diff --git a/python/tvm/tirx/op/__init__.py b/python/tvm/tirx/op/__init__.py
index 9c709d4733..241dc97b04 100644
--- a/python/tvm/tirx/op/__init__.py
+++ b/python/tvm/tirx/op/__init__.py
@@ -110,9 +110,9 @@ from tvm.ir.prim.op import truncmod as truncmod
from tvm.ir.prim.op import vscale as vscale
from .. import _ffi_api
-from ..buffer import is_tensor_var, tensor_data_ptr
from ..expr import ExprOp, IntImm
from ..expr import TensorLoad as _make_tensor_load
+from ..tensor import is_tensor_var, tensor_data_ptr
from ..type import TensorMapType
tir = tirx # alias for backward compat with upstream tir.convert() calls
diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/tensor.py
similarity index 74%
copy from python/tvm/tirx/buffer.py
copy to python/tvm/tirx/tensor.py
index 2d744397db..c54395d1cc 100644
--- a/python/tvm/tirx/buffer.py
+++ b/python/tvm/tirx/tensor.py
@@ -14,10 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Abstraction for array data structures."""
+"""Tensor types, declarations, and type-owned expression methods."""
+from collections.abc import Callable
from enum import IntEnum
from numbers import Integral
+from typing import ClassVar
import tvm_ffi
@@ -30,9 +32,23 @@ from . import _buffer_view, _ffi_api
_REARRANGE_PATTERN_UNSET = object()
+def _tensor_type_field(name):
+ def getter(ty, expr):
+ _check_tensor_property_receiver(ty, expr)
+ return getattr(ty, name)
+
+ return getter
+
+
@tvm_ffi.register_object("tirx.TensorType")
class TensorType(Type):
- """The structural type carried by an ordinary TIRx tensor variable."""
+ """The structural type carried by an ordinary TIRx tensor variable.
+
+ Declared expression methods receive the tensor variable as ``expr`` after
+ the type receiver. Access through ``expr.method`` binds both receivers.
+ Read-only expression properties receive the same type and operand through
+ the getters in ``__expr_properties__``.
+ """
dtype: PrimType
storage_scope: str
@@ -44,40 +60,32 @@ class TensorType(Type):
layout: object | None
allocated_addr: list
+ __expr_methods__ = (
+ "access_ptr",
+ "vload",
+ "vstore",
+ "scope",
+ "get_flattened_buffer",
+ "with_allocated_addr",
+ "with_dtype",
+ "offset_of",
+ "is_scalar",
+ "ptr_to",
+ "view",
+ "local",
+ "permute",
+ "rearrange",
+ "tile",
+ "chunk",
+ )
-def is_tensor_var(value) -> bool:
- """Return whether ``value`` is an ordinary Var carrying TensorType.
-
- Unlike ``isinstance(value, Var)``, this predicate distinguishes tensor
- variables from scalar or pointer variables.
- """
-
- return isinstance(value, tvm.ir.Var) and isinstance(value.ty, TensorType)
-
-
-class BufferAccessKind(IntEnum):
- """Buffer access modes accepted by :func:`buffer_access_ptr`."""
-
- READ = 1
- WRITE = 2
-
-
-class _TensorMethods:
- """Operations on tensor variables carrying layout and storage metadata.
-
- Do not construct directly, use :py:func:`~decl_tensor` instead.
- See the documentation of :py:func:`decl_tensor` for more details.
-
- See Also
- --------
- decl_tensor : Declare a buffer
- """
-
- def access_ptr(self, access_mask, ptr_type="handle", content_lanes=1,
offset=0, extent=None):
+ def access_ptr(
+ self, expr, access_mask, ptr_type="handle", content_lanes=1, offset=0,
extent=None
+ ):
"""Get an access pointer to the head of buffer.
This is the recommended method to get buffer data
- ptress when interacting with external functions.
+ pointers when interacting with external functions.
Parameters
----------
@@ -115,6 +123,7 @@ class _TensorMethods:
# Get access ptr for read with extent
buffer.access_ptr("r", extent = 100)
"""
+ _check_tensor_receiver(self, expr)
if isinstance(access_mask, str):
mask = 0
for value in access_mask:
@@ -136,7 +145,7 @@ class _TensorMethods:
offset = convert(offset)
extent = convert(extent)
return _ffi_api.TensorAccessPtr(
- self,
+ expr,
access_mask,
ptr_type,
content_lanes,
@@ -144,7 +153,7 @@ class _TensorMethods:
extent, # type: ignore
)
- def vload(self, begin, dtype=None):
+ def vload(self, expr, begin, dtype=None):
"""Generate an Expr that loads dtype from begin index.
Parameters
@@ -161,11 +170,12 @@ class _TensorMethods:
load : Expr
The corresponding load expression.
"""
+ _check_tensor_receiver(self, expr)
begin = (begin,) if isinstance(begin, int) or
tvm.ir.is_prim_expr(begin) else begin
- dtype = dtype if dtype else self.ty.dtype
- return _ffi_api.TensorVLoad(self, begin, dtype) # type: ignore
+ dtype = dtype if dtype else self.dtype
+ return _ffi_api.TensorVLoad(expr, begin, dtype) # type: ignore
- def vstore(self, begin, value):
+ def vstore(self, expr, begin, value):
"""Generate a Stmt that store value into begin index.
Parameters
@@ -181,19 +191,21 @@ class _TensorMethods:
store : Stmt
The corresponding store stmt.
"""
+ _check_tensor_receiver(self, expr)
begin = (begin,) if isinstance(begin, int) or
tvm.ir.is_prim_expr(begin) else begin
- return _ffi_api.TensorVStore(self, begin, value) # type: ignore
+ return _ffi_api.TensorVStore(expr, begin, value) # type: ignore
- def scope(self):
+ def scope(self, expr):
"""Return the storage scope associated with this buffer.
Returns
-------
scope : str
The storage scope associated with this buffer.
"""
- return _ffi_api.TensorStorageScope(self) # type: ignore
+ _check_tensor_receiver(self, expr)
+ return _ffi_api.TensorStorageScope(expr) # type: ignore
- def get_flattened_buffer(self):
+ def get_flattened_buffer(self, expr):
"""Generate a Var that is a flattened version of this buffer.
Returns
@@ -201,17 +213,20 @@ class _TensorMethods:
flattened : Var
The corresponding flat buffer.
"""
- return _ffi_api.TensorGetFlattenedTensor(self) # type: ignore
+ _check_tensor_receiver(self, expr)
+ return _ffi_api.TensorGetFlattenedTensor(expr) # type: ignore
- def with_allocated_addr(self, allocated_addr):
+ def with_allocated_addr(self, expr, allocated_addr):
"""Return a new buffer with the allocated address."""
- return _ffi_api.TensorWithAllocatedAddr(self, allocated_addr) # type:
ignore
+ _check_tensor_receiver(self, expr)
+ return _ffi_api.TensorWithAllocatedAddr(expr, allocated_addr) # type:
ignore
- def with_dtype(self, dtype):
+ def with_dtype(self, expr, dtype):
"""Return a new buffer with the dtype."""
- return _ffi_api.TensorWithDtype(self, dtype) # type: ignore
+ _check_tensor_receiver(self, expr)
+ return _ffi_api.TensorWithDtype(expr, dtype) # type: ignore
- def offset_of(self, indices):
+ def offset_of(self, expr, indices):
"""Determine the offset of the provided indices in the flattened
buffer.
Parameters
@@ -226,14 +241,10 @@ class _TensorMethods:
The offset indices of the element in the flattened buffer.
"""
- return _ffi_api.TensorOffsetOf(self, indices) # type: ignore
+ _check_tensor_receiver(self, expr)
+ return _ffi_api.TensorOffsetOf(expr, indices) # type: ignore
- @property
- def byte_offset(self):
- """Get the byte offset of the buffer."""
- return self.ty.elem_offset * tvm.DataType(self.ty.dtype).bits // 8
-
- def is_scalar(self, alloc_or_decl=True):
+ def is_scalar(self, expr, alloc_or_decl=True):
"""Check if the buffer is a scalar.
Parameters
@@ -246,20 +257,21 @@ class _TensorMethods:
-------
bool: True if the buffer is a scalar, False otherwise.
"""
- return _ffi_api.TensorIsScalar(self, alloc_or_decl)
+ _check_tensor_receiver(self, expr)
+ return _ffi_api.TensorIsScalar(expr, alloc_or_decl)
- def ptr_to(self, indices):
+ def ptr_to(self, expr, indices):
"""Get the pointer to the buffer at the given indices (logical
indices).
Note that the bufferload inside requires LowerTIPp pass to apply the
layout to get the physical indices.
""" # noqa: E501
- assert len(indices) == len(self.ty.shape), (
- f"The number of indices {indices} does not match "
- f"the shape of the buffer {self.ty.shape}"
+ _check_tensor_receiver(self, expr)
+ assert len(indices) == len(self.shape), (
+ f"The number of indices {indices} does not match the shape of the
buffer {self.shape}"
)
- return tvm.tirx.address_of(self[tuple(indices)])
+ return tvm.tirx.address_of(expr[tuple(indices)])
- def view(self, *args, **kwargs) -> "Var":
+ def view(self, expr, *args, **kwargs) -> "Var":
"""Creates a new view of the buffer. (used by parser)
Supported signatures are ``view(*shape, layout=None)``, where shape
can contain
@@ -271,10 +283,11 @@ class _TensorMethods:
view : Var
The corresponding view buffer.
"""
+ _check_tensor_receiver(self, expr)
- return _buffer_view.view(self, *args, **kwargs)
+ return _buffer_view.view(expr, *args, **kwargs)
- def local(self, *shape, layout=None) -> "Var":
+ def local(self, expr, *shape, layout=None) -> "Var":
"""Create a thread-local view of this buffer.
By default, both the inferred and explicit-shape forms address the
@@ -287,7 +300,7 @@ class _TensorMethods:
When called with no shape arguments, auto-infers a 1D shape from
the span of the parent layout's non-thread component (i.e.
- ``self.layout.storage().span()``). The explicit-``layout=`` form
+ ``expr.layout.storage().span()``). The explicit-``layout=`` form
instead infers the parent layout's ``storage().size()`` for
compatibility. Either inference requires the parent buffer to have a
layout.
@@ -308,9 +321,10 @@ class _TensorMethods:
local : Var
The corresponding local buffer.
"""
- return _buffer_view.local(self, *shape, layout=layout)
+ _check_tensor_receiver(self, expr)
+ return _buffer_view.local(expr, *shape, layout=layout)
- def permute(self, *dims) -> "Var":
+ def permute(self, expr, *dims) -> "Var":
"""Permute the dimensions of the buffer.
Parameters
@@ -323,9 +337,10 @@ class _TensorMethods:
permuted : Var
The buffer with permuted dimensions.
"""
- return _buffer_view.permute(self, *dims)
+ _check_tensor_receiver(self, expr)
+ return _buffer_view.permute(expr, *dims)
- def rearrange(self, pattern: str = _REARRANGE_PATTERN_UNSET, /, **sizes)
-> "Var":
+ def rearrange(self, expr, pattern: str = _REARRANGE_PATTERN_UNSET, /,
**sizes) -> "Var":
"""einops-style relayout in one line: ``buf.rearrange("b (2 r) -> 2 b
r")``.
A pure reshape+permute+reshape over the SAME physical bytes, spelled as
@@ -345,26 +360,14 @@ class _TensorMethods:
replica (``R[...]``), a stride-fiction/padded view, or a reshape
crossing
a swizzle-atom boundary — keep those as explicit ``view(layout=...)``.
"""
+ _check_tensor_receiver(self, expr)
if pattern is _REARRANGE_PATTERN_UNSET:
if "pattern" not in sizes:
raise TypeError("Var.rearrange() missing required argument:
'pattern'")
pattern = sizes.pop("pattern")
- return _buffer_view.rearrange(self, pattern, **sizes)
-
- @property
- def sub(self) -> "_buffer_view.SubIndexer":
- """Numpy-style view indexer: ``buf.sub[2, 4:8, ::4]``.
-
- Unlike plain ``buf[...]`` (TensorLoad for scalar indices, extent-1
- TensorRegion dims for tile-primitive operands), ``sub`` follows numpy
- basic-indexing semantics as a *view constructor*: an integer index
- removes the dim (``select``), ``a:b`` narrows it, and ``a::s`` takes
- every s-th element (requires the extent divisible by ``s`` and
- ``a < s``). Trailing dims are kept whole.
- """
- return _buffer_view.sub(self)
+ return _buffer_view.rearrange(expr, pattern, **sizes)
- def tile(self, *specs) -> "_buffer_view.TileIndexer":
+ def tile(self, expr, *specs) -> "_buffer_view.TileIndexer":
"""Chunk a dim: split it into factors, pick a chunk, keep the rest.
Rank-preserving — the picked dim's remaining factors merge back into
@@ -391,9 +394,10 @@ class _TensorMethods:
A picked index may be a dynamic Expr (e.g. a warp id); picking
several factors of one dim is allowed.
"""
- return _buffer_view.tile(self, *specs)
+ _check_tensor_receiver(self, expr)
+ return _buffer_view.tile(expr, *specs)
- def chunk(self, spec) -> "_buffer_view.ChunkIndexer":
+ def chunk(self, expr, spec) -> "_buffer_view.ChunkIndexer":
"""Split dims into equal contiguous chunks and pick a chunk per dim —
**rank-preserving**. Index the result with ``[picks]``.
@@ -410,11 +414,68 @@ class _TensorMethods:
X[.., c * k : (c + 1) * k, ..] # before (k = E // n)
X.chunk((None, .., n, ..))[.., c, ..] # after (k inferred)
"""
- return _buffer_view.chunk(self, spec)
+ _check_tensor_receiver(self, expr)
+ return _buffer_view.chunk(expr, spec)
+
+ def _data(self, expr):
+ _check_tensor_property_receiver(self, expr)
+ return tensor_data_ptr(expr)
+
+ def _dtype(self, expr):
+ _check_tensor_property_receiver(self, expr)
+ return self.dtype.dtype
+
+ def _byte_offset(self, expr):
+ _check_tensor_property_receiver(self, expr)
+ return self.elem_offset * tvm.DataType(self.dtype).bits // 8
+
+ def _sub(self, expr):
+ _check_tensor_property_receiver(self, expr)
+ return _buffer_view.sub(expr)
+
+ __expr_properties__: ClassVar[dict[str, Callable]] = {
+ "shape": _tensor_type_field("shape"),
+ "strides": _tensor_type_field("strides"),
+ "elem_offset": _tensor_type_field("elem_offset"),
+ "data_alignment": _tensor_type_field("data_alignment"),
+ "offset_factor": _tensor_type_field("offset_factor"),
+ "layout": _tensor_type_field("layout"),
+ "allocated_addr": _tensor_type_field("allocated_addr"),
+ "data": _data,
+ "dtype": _dtype,
+ "byte_offset": _byte_offset,
+ "sub": _sub,
+ }
-# These results expose the complete buffer namespace defined above. Indexer
-# definitions precede this class to avoid a circular buffer/view import.
+def is_tensor_var(value) -> bool:
+ """Return whether ``value`` is an ordinary Var carrying TensorType.
+
+ Unlike ``isinstance(value, Var)``, this predicate distinguishes tensor
+ variables from scalar or pointer variables.
+ """
+
+ return isinstance(value, tvm.ir.Var) and isinstance(value.ty, TensorType)
+
+
+class BufferAccessKind(IntEnum):
+ """Buffer access modes accepted by :func:`buffer_access_ptr`."""
+
+ READ = 1
+ WRITE = 2
+
+
+def _check_tensor_receiver(ty, expr):
+ if not is_tensor_var(expr):
+ raise TypeError("Tensor methods expect a Var with TensorType")
+ if not ty.same_as(expr.ty):
+ raise TypeError("Tensor method type must be the operand's type")
+
+
+def _check_tensor_property_receiver(ty, expr):
+ if not is_tensor_var(expr):
+ raise AttributeError("Tensor properties are only available on a Var
with TensorType")
+ _check_tensor_receiver(ty, expr)
def decl_tensor(
@@ -483,59 +544,3 @@ def buffer_data_pointer_type(buffer):
if not is_tensor_var(buffer):
raise TypeError("buffer_data_pointer_type expects a Var with
TensorType")
return _ffi_api.TensorDataPointerType(buffer)
-
-
-# Tensor variables intentionally retain runtime type key ``ir.Var``. Importing
-# ``tvm.tirx`` therefore augments ``tvm.ir.Var`` process-wide with the legacy
-# buffer operation and metadata surface. Non-buffer Vars reject the metadata
-# properties with AttributeError, preserving correct ``hasattr`` behavior.
-for _name, _value in _TensorMethods.__dict__.items():
- if _name.startswith("__"):
- continue
- if callable(_value) or isinstance(_value, property):
- setattr(tvm.ir.Var, _name, _value)
-
-
-def _tensor_type_field(name):
- def getter(value):
- if not is_tensor_var(value):
- raise AttributeError(f"{name} is only available on a Var with
TensorType")
- return getattr(value.ty, name)
-
- return property(getter)
-
-
-# Preserve the tensor variable metadata surface while keeping TensorType as the
-# single source of truth.
-for _name in (
- "shape",
- "strides",
- "elem_offset",
- "data_alignment",
- "offset_factor",
- "layout",
- "allocated_addr",
-):
- setattr(tvm.ir.Var, _name, _tensor_type_field(_name))
-
-
-def _tensor_dtype_property(value):
- if not is_tensor_var(value):
- raise AttributeError("dtype is only available on a Var with
TensorType")
- # Preserve the tensor variable Python dtype surface. TensorType stores a
- # PrimType, while Python callers historically receive its runtime DataType.
- return value.ty.dtype.dtype
-
-
-tvm.ir.Var.dtype = property(_tensor_dtype_property)
-
-
-# Keep the established ``A.data`` TVMScript surface as syntax sugar. Compiler
-# and builder code calls ``tensor_data_ptr(A)`` directly.
-def _tensor_data_property(value):
- if not is_tensor_var(value):
- raise AttributeError("data is only available on a Var with TensorType")
- return tensor_data_ptr(value)
-
-
-tvm.ir.Var.data = property(_tensor_data_property)
diff --git a/tests/python/sym/test_sym_deduce_bound.py
b/tests/python/sym/test_sym_deduce_bound.py
index 7a58b72f42..3b840543d4 100644
--- a/tests/python/sym/test_sym_deduce_bound.py
+++ b/tests/python/sym/test_sym_deduce_bound.py
@@ -19,7 +19,7 @@ import pytest
import tvm
import tvm.testing
-from tvm.tirx.buffer import decl_tensor
+from tvm.tirx.tensor import decl_tensor
def test_deduce():
diff --git a/tests/python/tirx-base/test_tir_buffer.py
b/tests/python/tirx-base/test_tir_buffer.py
index 7abfd86e2c..c9aebb04c3 100644
--- a/tests/python/tirx-base/test_tir_buffer.py
+++ b/tests/python/tirx-base/test_tir_buffer.py
@@ -61,7 +61,7 @@ def test_tensor_var_identity_and_global_var_properties():
for name in ("shape", "dtype", "data"):
assert not hasattr(scalar, name)
- with pytest.raises(AttributeError, match="only available on a Var with
TensorType"):
+ with pytest.raises(AttributeError, match="has no attribute"):
getattr(scalar, name)
diff --git a/tests/python/tirx/operator/tile_primitive/test_dispatcher.py
b/tests/python/tirx/operator/tile_primitive/test_dispatcher.py
index 225909579d..1d00e7c858 100644
--- a/tests/python/tirx/operator/tile_primitive/test_dispatcher.py
+++ b/tests/python/tirx/operator/tile_primitive/test_dispatcher.py
@@ -140,9 +140,9 @@ def test_dispatch_prints_real_opcall_ir():
"""Create a real TileOpCall via BufferRegions and ensure its IR is in the
table."""
_import_and_register()
from tvm.ir import Op
- from tvm.tirx.buffer import decl_tensor
from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
from tvm.tirx.stmt import TileOpCall
+ from tvm.tirx.tensor import decl_tensor
# Build a real TIRx TileOpCall: tirx.tile.copy(A[0:64], B[0:64])
A = decl_tensor((64,), "float32", scope="global")
diff --git a/tests/python/tirx/script/test_tirx_script_source_locations.py
b/tests/python/tirx/script/test_tirx_script_source_locations.py
index 683a8c591c..7adb114899 100644
--- a/tests/python/tirx/script/test_tirx_script_source_locations.py
+++ b/tests/python/tirx/script/test_tirx_script_source_locations.py
@@ -246,14 +246,14 @@ def
test_native_view_keeps_producer_identity_name_and_span(monkeypatch):
from tvm import ir, tirx
from tvm.script.ir_builder import base
- original = ir.Var.view
+ original = tirx.TensorType.view
seen, produced, observed = [], [], []
span = ir.Span(ir.SourceName("producer.py"), 7, 7, 2, 19)
@wraps(original)
- def view(buffer, *args):
+ def view(ty, buffer, *args):
seen.append("view")
- value = base.at_(span, original(buffer, *args))
+ value = base.at_(span, original(ty, buffer, *args))
produced.append((value, value.name, value.span))
return value
@@ -264,7 +264,7 @@ def
test_native_view_keeps_producer_identity_name_and_span(monkeypatch):
def observe(value):
observed.append((value, value.span))
- monkeypatch.setattr(ir.Var, "view", view)
+ monkeypatch.setattr(tirx.TensorType, "view", view)
@T.function
def main(A: T.Tensor((4, 4), "float32")):
diff --git a/tests/python/tirx/test_hint.py b/tests/python/tirx/test_hint.py
index df0f15eb1e..a8f380da07 100644
--- a/tests/python/tirx/test_hint.py
+++ b/tests/python/tirx/test_hint.py
@@ -29,8 +29,8 @@ def from_source(code):
def test_hint_keyword_arg_on_tx_op():
"""Tx.op(..., hint="msg") stores hint in TileOpCall.config."""
- from tvm.tirx.buffer import decl_tensor
from tvm.tirx.stmt import TileOpCall
+ from tvm.tirx.tensor import decl_tensor
A = decl_tensor((64, 64), "float32", scope="global")
A_sm = decl_tensor((64, 64), "float32", scope="shared")
diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py
index e892e8707a..44657a6144 100644
--- a/tests/python/tirx/test_op.py
+++ b/tests/python/tirx/test_op.py
@@ -19,9 +19,9 @@ import pickle
import pytest
from tvm.ir import Op, assert_structural_equal
-from tvm.tirx.buffer import decl_tensor
from tvm.tirx.exec_scope import ExecScope
from tvm.tirx.stmt import TileOpCall
+from tvm.tirx.tensor import decl_tensor
def _test(op: str, *args):
