You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/18 03:58:15 UTC
(tvm) 04/10: [TIRx] Store tensor address geometry in layouts (#20786)
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 7d631e4ea822bcb430a83cbf41e265d916fe3630
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Sat Oct 17 22:44:53 2026 -0400
[TIRx] Store tensor address geometry in layouts (#20786)
Make layout the single source of tensor strides and element offsets.
Tensor and S-TIR declarations accept explicit layouts, while views
forward layout geometry directly. Shape-aligned memory layouts support
dynamic strides and linear offset expressions through the existing
parameter binder.
Provide structural shape/stride matching and distinguish TileLayout
contiguity from TensorType default shape coverage and zero origin. Both
predicates reuse an optional analyzer after deep-equality checks. Python
exposes the optional named matching result only in S-TIR. Shared-memory
pools select their physical allocation base with pointer offsets, and
Relax preserves contiguous native tensor origins at call boundaries.
---
tirx/include/tvm/tirx/expr.h | 6 -
tirx/include/tvm/tirx/layout.h | 10 +
tirx/include/tvm/tirx/type.h | 87 ++++++---
.../tvm_tirx/backend/cuda/lang/alloc_pool.py | 8 +-
.../tvm_tirx/backend/cuda/tile/patterns/common.py | 4 +-
.../tvm_tirx/backend/cuda/tile/patterns/tma.py | 8 +-
tirx/python/tvm_tirx/tirx/_tensor_view.py | 214 ++++++++++-----------
tirx/python/tvm_tirx/tirx/layout.py | 8 +
tirx/python/tvm_tirx/tirx/op/__init__.py | 17 +-
tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py | 26 +--
tirx/python/tvm_tirx/tirx/tensor.py | 53 ++---
tirx/python/tvm_tirx/tirx/transform/transform.py | 3 +-
tirx/src/backend/cuda/op/target_builtin.cc | 2 -
.../backend/trn/transform/lower_trainium_layout.cc | 38 ++--
tirx/src/tirx/analysis/check_well_formed.cc | 22 ++-
tirx/src/tirx/ir/layout/layout.cc | 1 +
tirx/src/tirx/ir/layout/tile_core.cc | 28 +++
tirx/src/tirx/ir/specialize.cc | 67 ++-----
tirx/src/tirx/ir/tensor.cc | 35 ++--
tirx/src/tirx/ir/type.cc | 112 +++++------
tirx/src/tirx/script/ir_builder/ir.cc | 11 +-
tirx/src/tirx/script/printer/tensor.cc | 23 +--
tirx/src/tirx/transform/flatten_tensor.cc | 32 +--
tirx/src/tirx/transform/flattened_tensor.h | 35 ++--
tirx/src/tirx/transform/lower_tensor_layouts.cc | 54 ++++--
tirx/src/tirx/transform/remove_no_op.cc | 4 -
tirx/src/tirx/transform/stmt_simplify.cc | 6 +-
tirx/src/tirx/transform/storage_rewrite.cc | 22 +--
tirx/src/tirx/transform/tvm_ffi_binder.cc | 149 +++-----------
tirx/src/tirx/transform/tvm_ffi_binder.h | 24 ---
tirx/src/tirx/transform/vectorize_loop.cc | 38 ++--
tirx/tests/cpp/tirx/tensor_var_test.cc | 7 +-
.../python/backend/vulkan/codegen/test_basic.py | 14 +-
tirx/tests/python/codegen/llvm/test_basic.py | 31 +++
tirx/tests/python/codegen/llvm/test_riscv.py | 6 +-
.../python/tirx/analysis/test_undefined_vars.py | 7 +-
tirx/tests/python/tirx/base/test_intrin.py | 12 +-
tirx/tests/python/tirx/base/test_nodes.py | 5 +-
tirx/tests/python/tirx/base/test_op_types.py | 27 ++-
tirx/tests/python/tirx/base/test_specialize.py | 42 ++--
tirx/tests/python/tirx/base/test_tensor_type.py | 37 ++--
tirx/tests/python/tirx/codegen/test_blackwell.py | 13 +-
tirx/tests/python/tirx/script/test_tensors.py | 96 +++++----
tirx/tests/python/tirx/test_tile_namespace.py | 71 ++++---
.../python/tirx/tile/cuda/test_cp_async_bulk.py | 8 +-
.../tirx/tile/cuda/test_cp_async_bulk_tensor.py | 49 +++--
.../tirx/tile/cuda/test_explicit_patterns.py | 11 +-
.../python/tirx/tile/cuda/test_ld_st_registers.py | 1 -
.../python/tirx/transform/test_convert_ssa.py | 26 ++-
.../python/tirx/transform/test_flatten_tensor.py | 63 ++++--
.../transform/test_force_narrow_index_to_i32.py | 54 +++---
.../tirx/transform/test_lower_tensor_layouts.py | 31 ++-
.../python/tirx/transform/test_make_packed_api.py | 26 +++
.../python/tirx/transform/test_narrow_datatype.py | 8 +-
tirx/tests/python/tirx/transform/test_simplify.py | 24 ++-
.../python/tirx/transform/test_storage_rewrite.py | 24 ++-
56 files changed, 952 insertions(+), 888 deletions(-)
diff --git a/tirx/include/tvm/tirx/expr.h b/tirx/include/tvm/tirx/expr.h
index afca08f7ec..3a86467390 100644
--- a/tirx/include/tvm/tirx/expr.h
+++ b/tirx/include/tvm/tirx/expr.h
@@ -65,12 +65,6 @@ class TensorVar : public Var {
/*! \brief Return the primitive element type. */
PrimType dtype() const { return ty()->dtype; }
- /*! \brief Return explicit strides, or an empty array for a compact tensor. */
- const ffi::Array<PrimExpr>& strides() const { return ty()->strides; }
-
- /*! \brief Return the offset in tensor elements. */
- const PrimExpr& elem_offset() const { return ty()->elem_offset; }
-
/*! \brief Project the physical pointer established by the definition site. */
Expr data() const { return Call(ty()->DataPtrType(), tensor_data_ptr_op(), {*this}); }
diff --git a/tirx/include/tvm/tirx/layout.h b/tirx/include/tvm/tirx/layout.h
index 933e0e2e0e..18297a1a68 100644
--- a/tirx/include/tvm/tirx/layout.h
+++ b/tirx/include/tvm/tirx/layout.h
@@ -30,6 +30,7 @@
#include <tvm/ffi/object.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/module.h>
+#include <tvm/sym/analyzer.h>
#include <tvm/target/target.h>
#include <tvm/tirx/tile/thread_scope.h>
@@ -191,6 +192,15 @@ class TileLayoutNode : public LayoutNode {
/*! \brief Get the shape of the shard */
ffi::Array<PrimExpr> GetShardShape() const;
+ /*!
+ * \brief Whether the tile describes contiguous row-major memory strides.
+ *
+ * Memory-axis origins do not affect this predicate. Replicas and non-memory
+ * axes are rejected; unit extents may have arbitrary strides. Deep equality
+ * is tried first, reusing the supplied analyzer or lazily creating one for proof.
+ */
+ TVM_TIRX_DLL bool IsContiguous(ffi::Optional<sym::Analyzer> analyzer = std::nullopt) const;
+
/*! \brief Is the layout trivial (pure memory, identical mapping) */
bool IsTrivial() const;
diff --git a/tirx/include/tvm/tirx/type.h b/tirx/include/tvm/tirx/type.h
index 57e6fc8515..1cd7a1e361 100644
--- a/tirx/include/tvm/tirx/type.h
+++ b/tirx/include/tvm/tirx/type.h
@@ -25,8 +25,10 @@
#define TVM_TIRX_TYPE_H_
#include <tvm/ir/expr.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/ir/type.h>
#include <tvm/runtime/device_api.h>
+#include <tvm/sym/analyzer.h>
#include <tvm/tirx/layout.h>
#include <optional>
@@ -56,9 +58,8 @@ inline PrimType DefaultIndexType() {
*/
class TensorTypeNode : public TypeNode {
public:
- explicit TensorTypeNode(ffi::UnsafeInit tag) : elem_offset(tag) {}
-
- explicit TensorTypeNode(PrimExpr elem_offset) : elem_offset(std::move(elem_offset)) {}
+ explicit TensorTypeNode(ffi::UnsafeInit) {}
+ TensorTypeNode() = default;
/*! \brief dtype in the content of the tensor */
PrimType dtype = PrimType::Void();
@@ -71,20 +72,8 @@ class TensorTypeNode : public TypeNode {
* generators.
*/
ffi::Array<PrimExpr> shape;
- /*!
- * \brief The strides of each dimension
- * This can be an empty array, indicating array is contiguous
- */
- ffi::Array<PrimExpr> strides;
- /*! \brief The offset in terms of number of dtype elements (including lanes) */
- PrimExpr elem_offset;
/*! \brief Alignment requirement of data pointer in bytes. */
int data_alignment;
- /*!
- * \brief Factor of elem_offset field,
- * elem_offset is guaranteed to be multiple of offset_factor.
- */
- int offset_factor;
/*! \brief The layout of the tensor */
ffi::Optional<Layout> layout;
@@ -94,17 +83,22 @@ class TensorTypeNode : public TypeNode {
.def_ro("shape", &TensorTypeNode::shape, refl::AttachFieldFlag::SEqHashDefSimple())
.def_ro("dtype", &TensorTypeNode::dtype)
.def_ro("scope", &TensorTypeNode::scope, refl::DefaultValue(runtime::MemoryScope::Global()))
- .def_ro("strides", &TensorTypeNode::strides, refl::AttachFieldFlag::SEqHashDefSimple(),
- refl::DefaultValue(ffi::Array<PrimExpr>()))
- .def_ro("elem_offset", &TensorTypeNode::elem_offset,
- refl::AttachFieldFlag::SEqHashDefSimple(),
- refl::DefaultValue(IntImm(tvm::tirx::DefaultIndexType(), 0)))
.def_ro("data_alignment", &TensorTypeNode::data_alignment,
refl::DefaultValue(static_cast<int>(runtime::kAllocAlignment)))
- .def_ro("offset_factor", &TensorTypeNode::offset_factor, refl::DefaultValue(1))
.def_ro("layout", &TensorTypeNode::layout, refl::DefaultValue(nullptr));
}
+ /*!
+ * \brief Whether the layout is provably equivalent to implicit row-major storage.
+ *
+ * An absent layout returns true. Explicit layouts must have zero origin and
+ * contiguous memory-only tiles covering exactly the tensor shape, without replicas.
+ * Equivalent split tiles and arbitrary strides on unit extents are accepted.
+ * Deep expression equality is tried before symbolic proof. A supplied analyzer
+ * retains its constraints; otherwise one is created only on the first proof fallback.
+ */
+ TVM_TIRX_DLL bool HasDefaultLayout(ffi::Optional<sym::Analyzer> analyzer = std::nullopt) const;
+
/*! \return preferred index type for this tensor node */
PrimType DefaultIndexType() const {
return !shape.empty() ? shape[0].ty() : tvm::tirx::DefaultIndexType();
@@ -126,10 +120,8 @@ class TensorTypeNode : public TypeNode {
* float16x4 elements in a tensor of type float16x4.)
*
* \param index The index to be accessed.
- * \param inner Ignore the elem_offset, return inner offset only
*/
- TVM_TIRX_DLL ffi::Array<PrimExpr> ElemOffset(ffi::Array<PrimExpr> index,
- bool inner = false) const;
+ TVM_TIRX_DLL ffi::Array<PrimExpr> ElemOffset(ffi::Array<PrimExpr> index) const;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorType", TensorTypeNode, TypeNode);
};
@@ -140,9 +132,7 @@ class TensorTypeNode : public TypeNode {
class TensorType : public Type {
public:
TVM_TIRX_DLL TensorType(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
- runtime::MemoryScope scope = {}, ffi::Array<PrimExpr> strides = {},
- ffi::Optional<PrimExpr> elem_offset = std::nullopt,
- int data_alignment = 0, int offset_factor = 0,
+ runtime::MemoryScope scope = {}, int data_alignment = 0,
ffi::Optional<Layout> layout = std::nullopt, Location loc = UnknownLoc());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type, TensorTypeNode);
@@ -153,6 +143,49 @@ class TensorType : public Type {
}
};
+/*!
+ * \brief Match a tensor's shape-aligned strided memory representation.
+ *
+ * Returns (shape, strides, element offset). An absent layout yields contiguous
+ * row-major strides and zero offset. A TileLayout must have one memory iterator
+ * per shape dimension, no replicas, and only a memory-axis offset. Other layouts
+ * or mismatched extents return nullopt. Strides are always actual expressions;
+ * an empty stride array means rank zero, not an implicit compact contract.
+ *
+ * ExprDeepEqual preserves variable identity without remapping or symbolic proofs,
+ * avoiding the overhead and broader equivalences of general structural matching.
+ */
+inline ffi::Optional<ffi::Tuple<ffi::Array<PrimExpr>, ffi::Array<PrimExpr>, PrimExpr>>
+MatchShapeStridePattern(const TensorType& tensor) {
+ ffi::Array<PrimExpr> strides;
+ PrimExpr offset = IntImm(tensor->DefaultIndexType(), 0);
+ if (!tensor->layout.has_value()) {
+ PrimExpr stride = IntImm(tensor->DefaultIndexType(), 1);
+ strides = ffi::Array<PrimExpr>(tensor->shape.size(), stride);
+ for (size_t i = tensor->shape.size(); i > 0; --i) {
+ strides.Set(i - 1, stride);
+ stride = stride * tensor->shape[i - 1];
+ }
+ } else {
+ const auto* tile = tensor->layout.as<TileLayoutNode>();
+ if (!tile || !tile->replica.empty() || tile->shard.size() != tensor->shape.size()) {
+ return std::nullopt;
+ }
+ prim::ExprDeepEqual equal;
+ static const Axis memory_axis = Axis::Get("m");
+ for (size_t i = 0; i < tile->shard.size(); ++i) {
+ if (tile->shard[i]->axis != memory_axis || !equal(tile->shard[i]->extent, tensor->shape[i]))
+ return std::nullopt;
+ strides.push_back(tile->shard[i]->stride);
+ }
+ for (const auto& [axis, value] : tile->offset) {
+ if (axis != memory_axis) return std::nullopt;
+ offset = value;
+ }
+ }
+ return ffi::Tuple(tensor->shape, std::move(strides), std::move(offset));
+}
+
/*!
* \brief The type of tensor map.
* \sa TensorMapType
diff --git a/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py b/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
index d1ff8bad2e..3f790f6951 100644
--- a/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
+++ b/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
@@ -201,23 +201,23 @@ class SMEMPool:
self,
shape,
dtype="float32",
- strides=None,
scope="shared.dyn",
align=0,
layout="default",
):
+ from tvm_tirx.ir import PtrType
+ from tvm_tirx.tirx import ptr_byte_offset
+
ir = _get_ir()
if align > 0:
self.offset = (self.offset + align - 1) // align * align
res = ir._bind_tensor(
ir.decl_tensor(
- self.ptr,
+ ptr_byte_offset(self.ptr, self.offset, ty=PtrType(dtype, scope)),
ty_args=[
ir.Tensor(
shape,
dtype,
- strides=strides,
- byte_offset=self.offset,
scope=scope,
align=align,
layout=layout,
diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
index 97797b7cda..c89948448a 100644
--- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
+++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
@@ -276,8 +276,6 @@ class Region:
def physical(self, atom, repeat=None):
result = dict(self.layout.apply(*self.indices(atom, repeat), shape=self.ty.shape))
- if self.ty.scope != MemoryScope("tmem"):
- result["m"] = result.get("m", 0) + self.ty.elem_offset
return {key: self.ctx.analyzer.simplify(value) for key, value in result.items()}
def memory(self, atom):
@@ -357,7 +355,7 @@ class Region:
if self.ty.scope != MemoryScope("local"):
raise ValueError(f"{self.name}: scalar values must be in registers")
if self.raw is None:
- span = self.layout.span("m") + self.ty.elem_offset
+ span = self.layout.span("m")
self.raw = Var(
"tile_" + self.name,
T.Tensor((span,), str(self.dtype), scope="local", layout="default"),
diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
index 08a73682e0..1601f2cb9f 100644
--- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
+++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
@@ -125,11 +125,7 @@ def _tensor_map(ctx, region, rank, box, encoding):
# checked at the same coordinates, without folding or permuting its axes.
layout = ty.layout
if layout is None:
- strides = (
- list(ty.strides)
- if ty.strides
- else [prod(ty.shape[i + 1 :]) for i in range(tensor_rank)]
- )
+ strides = [prod(ty.shape[i + 1 :]) for i in range(tensor_rank)]
base = {"m": sum(i * s for i, s in zip(starts, strides))}
else:
base = layout.apply(*starts, shape=ty.shape)
@@ -164,7 +160,7 @@ def _tensor_map(ctx, region, rank, box, encoding):
ctx.require(strides[-1] == channels, "TMA W stride must equal one interleave slice")
alignment = 32 if packed or interleave == 2 else 16
ctx.require(
- (base["m"] + ty.elem_offset) * bits % (8 * alignment) == 0,
+ base["m"] * bits % (8 * alignment) == 0,
f"TMA global origin requires {alignment}-byte alignment",
)
for stride in strides[:-1]:
diff --git a/tirx/python/tvm_tirx/tirx/_tensor_view.py b/tirx/python/tvm_tirx/tirx/_tensor_view.py
index 01c2ff4245..1e96411d83 100644
--- a/tirx/python/tvm_tirx/tirx/_tensor_view.py
+++ b/tirx/python/tvm_tirx/tirx/_tensor_view.py
@@ -46,14 +46,11 @@ if TYPE_CHECKING:
from tvm_tirx.ir import Var
-def _redecl(
- buf: Var, shape, layout, *, dtype=None, elem_offset=None, addr_offset=None, strides=None
-):
+def _redecl(buf: Var, shape, layout, *, dtype=None, addr_offset=None):
"""Re-declare a derived view over ``buf`` storage.
Routes the tensor identity by storage kind: a ``tmem`` tensor aliases by
- column address, while every other scope carries
- ``data``/``strides``/``elem_offset`` through.
+ column address, while every other scope carries its physical pointer and layout.
"""
if _is_tmem_scope(buf.scope()):
from tvm_tirx.backend.cuda.op import (
@@ -81,33 +78,50 @@ def _redecl(
tensor_type = tvm_tirx.tirx.script.ir_builder.Tensor(
shape,
buf.dtype if dtype is None else dtype,
- strides=buf.strides if strides is None else strides,
- elem_offset=buf.elem_offset if elem_offset is None else elem_offset,
scope=buf.scope(),
align=buf.data_alignment,
- offset_factor=buf.offset_factor,
layout=layout,
)
return bind_(tvm_tirx.tirx.script.ir_builder.decl_tensor(buf.data, ty_args=[tensor_type]))
-def _strides(buf):
- if buf.strides:
- return list(buf.strides)
- return tvm_tirx.tirx.layout.Layout._get_default_strides(list(buf.shape), 1)
+def _view_layout(buf):
+ """Materialize implicit compact storage for layout transformations."""
+ from .layout import S, TileLayout
+ if buf.layout is not None:
+ return buf.layout
+ return TileLayout(S[tuple(buf.shape)])
+
+
+def _is_memory_tile(layout, shape):
+ """Whether view surgery can address each memory dimension directly."""
+ from .layout import Axis, TileLayout
+
+ return (
+ isinstance(layout, TileLayout)
+ and not layout.replica
+ and len(layout.shard) == len(shape)
+ and all(it.axis == Axis.m for it in layout.shard)
+ and all(axis == Axis.m for axis in layout.offset)
+ and all(
+ tvm_tirx.ir.prim.expr_deep_equal(it.extent, extent)
+ for it, extent in zip(layout.shard, shape)
+ )
+ )
+
+
+def _reshape_layout(tile, shape):
+ """Reshape contiguous iterator groups without changing their storage addresses."""
+ from .layout import Iter, TileLayout
-def _reshape_strides(buf, shape):
- """Reshape contiguous stride groups without changing their storage addresses."""
analyzer = tvm_tirx.sym.Analyzer()
- old_shape, old_strides = list(buf.shape), _strides(buf)
- if not analyzer.can_prove_equal(
- functools.reduce(lambda a, b: a * b, old_shape, 1),
- functools.reduce(lambda a, b: a * b, shape, 1),
+ old_shape = [it.extent for it in tile.shard]
+ old_strides = [it.stride for it in tile.shard]
+ if len(shape) == len(old_shape) and all(
+ analyzer.can_prove_equal(a, b) for a, b in zip(shape, old_shape)
):
- raise ValueError("view: source and result must have equal element counts")
- if not buf.strides:
- return []
+ return tile
chunks = []
for extent, stride in reversed(list(zip(old_shape, old_strides))):
if analyzer.can_prove_equal(extent, 1):
@@ -127,13 +141,16 @@ def _reshape_strides(buf, shape):
if analyzer.can_prove_equal(product, extent):
break
if not analyzer.can_prove_equal(product, extent):
- raise ValueError("view: reshape crosses non-contiguous stride groups")
+ return None
while i >= 0 and analyzer.can_prove_equal(shape[i], 1):
result[i] = analyzer.simplify(result[i + 1] * shape[i + 1]) if i + 1 < len(shape) else 1
i -= 1
if i >= 0:
- raise ValueError("view: reshape crosses non-contiguous stride groups")
- return result
+ return None
+ return TileLayout.from_iters(
+ [Iter(extent, stride, "m") for extent, stride in zip(shape, result)],
+ offset=dict(tile.offset.items()),
+ )
def view(buf: Var, *args, **kwargs) -> Var:
@@ -166,44 +183,45 @@ def view(buf: Var, *args, **kwargs) -> Var:
)
return shape
+ from .layout import Iter, TileLayout
+
+ source = _view_layout(buf)
+ regular = _is_memory_tile(_view_layout(buf), buf.shape)
if len(args) == 1 and isinstance(args[0], str | tvm_ffi.dtype) and not kwargs:
cast_dtype = tvm_ffi.dtype(args[0])
cur_dtype = tvm_ffi.dtype(buf.dtype)
- storage_axis = "TCol" if _is_tmem_scope(buf.scope()) else "m"
- if cast_dtype.bits > cur_dtype.bits:
- assert cast_dtype.bits % cur_dtype.bits == 0
- ratio = cast_dtype.bits // cur_dtype.bits
- layout = None if buf.layout is None else buf.layout.pack(ratio, storage_axis)
- shape = [s for s in buf.shape[:-1]] + [buf.shape[-1] // ratio]
- new_elem_offset = buf.elem_offset // ratio
- else:
- assert cur_dtype.bits % cast_dtype.bits == 0
- ratio = cur_dtype.bits // cast_dtype.bits
- layout = None if buf.layout is None else buf.layout.unpack(ratio, storage_axis)
- shape = [s for s in buf.shape[:-1]] + [buf.shape[-1] * ratio]
- new_elem_offset = buf.elem_offset * ratio
- strides = None
- if buf.layout is None:
+ pack = cast_dtype.bits > cur_dtype.bits
+ ratio = (cast_dtype.bits // cur_dtype.bits) if pack else (cur_dtype.bits // cast_dtype.bits)
+ assert max(cast_dtype.bits, cur_dtype.bits) % min(cast_dtype.bits, cur_dtype.bits) == 0
+ shape = [*buf.shape[:-1], buf.shape[-1] // ratio if pack else buf.shape[-1] * ratio]
+ if regular:
analyzer = tvm_tirx.sym.Analyzer()
- source_strides = _strides(buf)
- if not analyzer.can_prove_equal(source_strides[-1], 1):
+ if not analyzer.can_prove_equal(source.shard[-1].stride, 1):
raise ValueError("view(dtype): the last dimension must have unit stride")
- if cast_dtype.bits > cur_dtype.bits:
- if any(
- analyzer.can_prove(value % ratio != 0)
- for value in [buf.shape[-1], buf.elem_offset, *source_strides[:-1]]
- ):
- raise ValueError(
- "view(dtype): shape, offset and strides must align to the new dtype"
- )
- strides = [analyzer.simplify(value // ratio) for value in source_strides[:-1]] + [1]
- else:
- strides = [analyzer.simplify(value * ratio) for value in source_strides[:-1]] + [1]
- if not buf.strides:
- strides = []
- return _redecl(
- buf, shape, layout, dtype=cast_dtype, elem_offset=new_elem_offset, strides=strides
- )
+ values = [
+ buf.shape[-1],
+ *source.offset.values(),
+ *[it.stride for it in source.shard[:-1]],
+ ]
+ if pack and any(analyzer.can_prove(value % ratio != 0) for value in values):
+ raise ValueError(
+ "view(dtype): shape, offset and strides must align to the new dtype"
+ )
+
+ def scale(value):
+ return analyzer.simplify(value // ratio if pack else value * ratio)
+
+ layout = TileLayout.from_iters(
+ [
+ Iter(extent, scale(it.stride) if i + 1 < len(shape) else 1, "m")
+ for i, (extent, it) in enumerate(zip(shape, source.shard))
+ ],
+ offset={axis: scale(value) for axis, value in source.offset.items()},
+ )
+ else:
+ axis = "TCol" if _is_tmem_scope(buf.scope()) else "m"
+ layout = source.pack(ratio, axis) if pack else source.unpack(ratio, axis)
+ return _redecl(buf, shape, None if buf.layout is None else layout, dtype=cast_dtype)
shape = args
assert all(
@@ -217,10 +235,21 @@ def view(buf: Var, *args, **kwargs) -> Var:
)
if layout is None:
shape = _infer_shape(shape)
- strides = (
- [] if layout is not None else _reshape_strides(buf, shape) if buf.layout is None else None
- )
- return _redecl(buf, shape, buf.layout if layout is None else layout, strides=strides)
+ if layout is not None:
+ return _redecl(buf, shape, layout)
+ analyzer = tvm_tirx.sym.Analyzer()
+ if not analyzer.can_prove_equal(
+ functools.reduce(lambda a, b: a * b, buf.shape, 1),
+ functools.reduce(lambda a, b: a * b, shape, 1),
+ ):
+ raise ValueError("view: source and result must have equal element counts")
+ if buf.layout is None:
+ return _redecl(buf, shape, None)
+ if regular:
+ reshaped = _reshape_layout(source, shape)
+ if reshaped is not None:
+ return _redecl(buf, shape, reshaped)
+ return _redecl(buf, shape, source)
def permute(buf: Var, *dims) -> Var:
@@ -228,11 +257,11 @@ def permute(buf: Var, *dims) -> Var:
if sorted(dims) != list(range(len(buf.shape))):
raise ValueError(f"permute: invalid dimensions {dims}")
new_shape = [buf.shape[d] for d in dims]
- if buf.layout is None:
- strides = _strides(buf)
- return _redecl(buf, new_shape, None, strides=[strides[d] for d in dims])
layout, swizzle = _surgery_parts(buf)
- grouped, seps = layout.group(list(buf.shape))
+ if swizzle is None and _is_memory_tile(_view_layout(buf), buf.shape):
+ grouped, seps = layout, list(range(len(buf.shape) + 1))
+ else:
+ grouped, seps = layout.group(list(buf.shape))
new_layout = grouped.permute_by_groups(seps, list(dims))
if swizzle is not None:
new_layout = _rewrap_swizzle(new_layout, swizzle)
@@ -325,7 +354,7 @@ def _concrete_int(value):
def _surgery_parts(buf: Var):
"""Split a layout into its inner tile layout and optional swizzle params."""
- layout = buf.layout
+ layout = _view_layout(buf)
swizzle = None
if isinstance(layout, tvm_tirx.tirx.layout.ComposeLayout):
swizzle = (
@@ -380,31 +409,11 @@ def _dim_group_offset(group, index):
return offset
-def _swizzle_offset_commutes(swizzle, extra_offset):
- """Return whether an offset may be folded outside a swizzle permutation."""
- if swizzle is None or extra_offset is None:
- return True
- per_element, swizzle_len, atom_len, _ = swizzle
- if swizzle_len == 0:
- return True
- period = 1 << (int(per_element) + int(atom_len) + swizzle_len)
- offset_c = _concrete_int(extra_offset)
- if offset_c is not None:
- return offset_c % period == 0
- from ..sym import Analyzer # pylint: disable=import-outside-toplevel
-
- return Analyzer().can_prove_equal(tvm_tirx.tirx.floormod(extra_offset, period), 0)
-
-
def _rebuild_view(buf: Var, new_shape, new_shard, grouped, swizzle, extra_offset):
"""Rebuild a derived view while preserving its physical addresses."""
offset_map = dict(grouped.offset.items())
is_tmem = _is_tmem_scope(buf.scope())
- if (
- not is_tmem
- and extra_offset is not None
- and not _swizzle_offset_commutes(swizzle, extra_offset)
- ):
+ if not is_tmem and extra_offset is not None:
m_axis = tvm_tirx.tirx.layout.Axis.get("m")
previous = offset_map.get(m_axis)
offset_map[m_axis] = extra_offset if previous is None else previous + extra_offset
@@ -416,10 +425,7 @@ def _rebuild_view(buf: Var, new_shape, new_shard, grouped, swizzle, extra_offset
if is_tmem:
addr_offset = _tmem_element_offset_to_column_offset(buf, extra_offset)
return _redecl(buf, new_shape, new_layout, addr_offset=addr_offset)
- elem_offset = buf.elem_offset
- if extra_offset is not None:
- elem_offset = elem_offset + extra_offset
- return _redecl(buf, new_shape, new_layout, elem_offset=elem_offset)
+ return _redecl(buf, new_shape, new_layout)
def _tmem_element_offset_to_column_offset(buf: Var, element_offset):
@@ -475,17 +481,11 @@ def _view_drop(buf: Var, dim, index):
f"index {index_c} out of range for dim {dim} "
f"of extent {extent_c if extent_c is not None else buf.shape[dim]}"
)
- if buf.layout is None:
- strides = _strides(buf)
- return _redecl(
- buf,
- list(buf.shape[:dim]) + list(buf.shape[dim + 1 :]),
- None,
- elem_offset=buf.elem_offset + index * strides[dim],
- strides=strides[:dim] + strides[dim + 1 :],
- )
layout, swizzle = _surgery_parts(buf)
- grouped, seps = layout.group(list(buf.shape))
+ if swizzle is None and _is_memory_tile(_view_layout(buf), buf.shape):
+ grouped, seps = layout, list(range(len(buf.shape) + 1))
+ else:
+ grouped, seps = layout.group(list(buf.shape))
iters = list(grouped.shard)
lo, hi = seps[dim], seps[dim + 1]
_tmem_offset_axis_check(buf, iters[lo:hi], index_c, "index")
@@ -518,15 +518,11 @@ def _view_narrow(buf: Var, dim, start, length):
raise ValueError(
f"sub: range [{start_c}, {start_c + length_c}) exceeds dim {dim} extent {extent_c}"
)
- if buf.layout is None:
- strides = _strides(buf)
- shape = list(buf.shape)
- shape[dim] = length
- return _redecl(
- buf, shape, None, elem_offset=buf.elem_offset + start * strides[dim], strides=strides
- )
layout, swizzle = _surgery_parts(buf)
- grouped, seps = layout.group(list(buf.shape))
+ if swizzle is None and _is_memory_tile(_view_layout(buf), buf.shape):
+ grouped, seps = layout, list(range(len(buf.shape) + 1))
+ else:
+ grouped, seps = layout.group(list(buf.shape))
iters = list(grouped.shard)
lo, hi = seps[dim], seps[dim + 1]
group = iters[lo:hi]
diff --git a/tirx/python/tvm_tirx/tirx/layout.py b/tirx/python/tvm_tirx/tirx/layout.py
index 057ead1e31..15f0e59424 100644
--- a/tirx/python/tvm_tirx/tirx/layout.py
+++ b/tirx/python/tvm_tirx/tirx/layout.py
@@ -574,6 +574,14 @@ class TileLayout(Layout):
offset = {Axis.get(k) if isinstance(k, str) else k: v for k, v in offset.items()}
return _ffi_api.TileLayout(shard, replica, offset or {}) # pylint: disable=no-member
+ def is_contiguous(self, analyzer=None) -> bool:
+ """Whether memory-only tile strides are contiguous, regardless of origin.
+
+ Replicas and non-memory axes are rejected. An optional analyzer supplies
+ existing constraints; symbolic proof is otherwise initialized lazily.
+ """
+ return _ffi_api.TileLayoutIsContiguous(self, analyzer)
+
def is_trivial(self) -> bool:
"""Check if the layout is trivial."""
return _ffi_api.TileLayoutIsTrivial(self) # pylint: disable=no-member
diff --git a/tirx/python/tvm_tirx/tirx/op/__init__.py b/tirx/python/tvm_tirx/tirx/op/__init__.py
index 1ffc80c7a0..17e9573772 100644
--- a/tirx/python/tvm_tirx/tirx/op/__init__.py
+++ b/tirx/python/tvm_tirx/tirx/op/__init__.py
@@ -192,14 +192,25 @@ def _pack_tensor(buf, loc: Location = UNKNOWN_LOC):
loc=loc,
ty=PtrType(tvm_tirx.ir.PrimType("int64")),
)
+ from .._tensor_view import _is_memory_tile, _view_layout
+ from ..layout import Axis
+
+ layout = _view_layout(buf)
+ if not _is_memory_tile(layout, buf.ty.shape):
+ raise ValueError("DLTensor packing requires a shape-aligned strided memory layout")
+ stride_exprs = [it.stride for it in layout.shard]
+ index_type = buf.ty.shape[0].ty if buf.ty.shape else "int64"
+ origin = layout.offset.get(Axis.m, const(0, index_type))
+ # The DLTensor carries its origin separately from the compact-stride sentinel.
+ contiguous = layout.is_contiguous()
strides = (
Call(
"tirx.stack_make_shape",
- buf.ty.strides,
+ stride_exprs,
loc=loc,
ty=PtrType(tvm_tirx.ir.PrimType("int64")),
)
- if buf.ty.strides
+ if not contiguous
else 0
)
pack_args = [
@@ -208,7 +219,7 @@ def _pack_tensor(buf, loc: Location = UNKNOWN_LOC):
strides,
len(buf.ty.shape),
const(0, dtype=buf.ty.dtype),
- buf.ty.elem_offset,
+ origin,
]
return Call(
"tirx.stack_make_dltensor",
diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
index 7bce0c7e00..a13cfd9b2f 100644
--- a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
+++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
@@ -82,17 +82,6 @@ def _normalize_prim_type(dtype) -> ir.PrimType:
return ir.PrimType(dtype)
-def _get_elem_offset(elem_offset, byte_offset, dtype: str):
- assert elem_offset is None or byte_offset is None, (
- "elem_offset and byte_offset cannot be set at the same time"
- )
- if elem_offset is not None:
- return elem_offset
- if byte_offset is None:
- return None
- return byte_offset * 8 // (tvm_ffi.dtype(_normalize_prim_type(dtype).dtype).bits)
-
-
_meta_construction_state = threading.local()
_THIS_FILE = __file__
@@ -269,12 +258,8 @@ def _alloc_tensor(
shape,
dtype="float32",
data=None,
- strides=None,
- elem_offset=None,
- byte_offset=None,
scope=_SCOPE_GLOBAL,
align=-1,
- offset_factor=0,
layout=MISSING,
allocated_addr=None,
annotations=None,
@@ -284,7 +269,7 @@ def _alloc_tensor(
if allocated_addr is not None:
if scope in (_SCOPE_GLOBAL, _SCOPE_SHARED, _SCOPE_SHARED_DYN, _SCOPE_LOCAL):
raise ValueError(f"{scope} does not support allocation placement")
- if data is not None or elem_offset is not None or byte_offset is not None or offset_factor:
+ if data is not None:
raise ValueError(
"Allocation placement cannot be combined with pointer or offset metadata"
)
@@ -292,13 +277,9 @@ def _alloc_tensor(
shape,
dtype,
scope,
- strides,
- elem_offset,
align,
- offset_factor,
layout,
data=data,
- byte_offset=byte_offset,
)
return _bind_tensor(
alloc_tensor(allocated_addr, ty_args=[tensor_type], annotations=annotations)
@@ -391,7 +372,7 @@ def alloc_scalar(
@_register_mutable_decl("tirx.decl_scalar")
-def decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) -> _ir.TensorLoad:
+def decl_scalar(dtype, data, scope, layout=None) -> _ir.TensorLoad:
"""Declare a zero-dimensional tensor (scalar) from a pointer."""
buf = _bind_tensor(
decl_tensor(
@@ -401,9 +382,8 @@ def decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) -> _ir.T
(1,),
dtype,
scope=scope,
- elem_offset=_get_elem_offset(elem_offset, byte_offset, dtype),
align=-1,
- layout=None,
+ layout=layout,
)
],
)
diff --git a/tirx/python/tvm_tirx/tirx/tensor.py b/tirx/python/tvm_tirx/tirx/tensor.py
index aeae4b3b70..b67b3019b8 100644
--- a/tirx/python/tvm_tirx/tirx/tensor.py
+++ b/tirx/python/tvm_tirx/tirx/tensor.py
@@ -55,10 +55,8 @@ class TensorType(Type):
dtype: PrimType
scope: MemoryScope
shape: list
- strides: list
- elem_offset: tvm_tirx.ir.Expr
data_alignment: int
- offset_factor: int
+
layout: object | None
def __init__(
@@ -66,24 +64,17 @@ class TensorType(Type):
shape,
dtype,
scope=_SCOPE_GLOBAL,
- strides=None,
- elem_offset=None,
align=0,
- offset_factor=0,
layout=None,
*,
data=None,
- byte_offset=None,
loc: Location = UNKNOWN_LOC,
data_alignment=None,
):
"""Construct a tensor access contract.
- An omitted element offset is zero in the configured index dtype,
- independent of shape and offset_factor. Integer offsets use that dtype;
- explicitly typed expressions retain theirs. Pass a symbolic offset explicitly
- when the contract permits a variable starting position. ``align`` remains a
- compatibility spelling for ``data_alignment``.
+ Strides and element offsets are described by ``layout``.
+ ``align`` is an alias for ``data_alignment``.
"""
from tvm_tirx.ir.type import _normalize_type
from tvm_tirx.script.ir_builder.base import MISSING, at
@@ -97,10 +88,6 @@ class TensorType(Type):
dtype = _normalize_type(dtype)
if not isinstance(dtype, PrimType):
raise TypeError("Tensor dtype must be a primitive type")
- if elem_offset is not None and byte_offset is not None:
- raise ValueError("elem_offset and byte_offset cannot be set at the same time")
- if byte_offset is not None:
- elem_offset = byte_offset * 8 // tvm_ffi.dtype(dtype.dtype).bits
if data_alignment is None:
data_alignment = align
elif align not in (0, data_alignment):
@@ -114,7 +101,8 @@ class TensorType(Type):
if layout is MISSING:
layout = None
if layout == "default":
- layout = TileLayout(S[shape])
+ default = TileLayout(S[shape])
+ layout = default
elif layout is not None and not isinstance(layout, Layout):
raise TypeError("layout must be a Layout, None, or 'default'")
self.__init_handle_by_constructor__(
@@ -122,10 +110,7 @@ class TensorType(Type):
shape,
dtype,
scope,
- () if strides is None else strides,
- elem_offset,
data_alignment,
- offset_factor,
layout,
loc if isinstance(loc, Location) else UNKNOWN_LOC,
)
@@ -142,17 +127,22 @@ class TensorType(Type):
__expr_properties__: ClassVar[dict[str, Callable]] = {
"scope": lambda ty, expr: lambda: ty._scope(expr),
"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"),
"data": lambda ty, expr: ty._data(expr),
"dtype": lambda ty, expr: ty._dtype(expr),
- "byte_offset": lambda ty, expr: ty._byte_offset(expr),
"sub": lambda ty, expr: ty._sub(expr),
}
+ def has_default_layout(self, analyzer=None) -> bool:
+ """Whether the layout matches implicit row-major storage with zero origin.
+
+ Equivalent split tiles are accepted when their extent covers this shape.
+ An optional analyzer supplies existing bindings and constraints; otherwise
+ symbolic analysis is initialized only when structural equality is insufficient.
+ """
+ return _ffi_api.TensorTypeHasDefaultLayout(self, analyzer)
+
def effective_layout(self):
"""Return the explicit layout or infer a compact global layout.
@@ -165,19 +155,8 @@ class TensorType(Type):
if self.scope != _SCOPE_GLOBAL:
raise ValueError(f"An explicit layout is required for {self.scope} tensors")
- from tvm_tirx.sym import Analyzer
-
from .layout import S, TileLayout
- analyzer = Analyzer()
- expected = 1
- if self.strides:
- for extent, stride in reversed(list(zip(self.shape, self.strides))):
- if not analyzer.can_prove_equal(extent, 1) and not analyzer.can_prove_equal(
- stride, expected
- ):
- raise ValueError("An explicit layout is required for noncompact global tensors")
- expected = analyzer.simplify(expected * extent)
return TileLayout(S[tuple(self.shape)])
def _scope(self, expr):
@@ -284,10 +263,6 @@ class TensorType(Type):
_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_ffi.dtype(self.dtype).bits // 8
-
def _sub(self, expr):
_check_tensor_property_receiver(self, expr)
return _tensor_view.sub(expr)
diff --git a/tirx/python/tvm_tirx/tirx/transform/transform.py b/tirx/python/tvm_tirx/tirx/transform/transform.py
index 67a9954518..49213a9482 100644
--- a/tirx/python/tvm_tirx/tirx/transform/transform.py
+++ b/tirx/python/tvm_tirx/tirx/transform/transform.py
@@ -255,8 +255,7 @@ def MakePackedAPI():
`TensorType`. This pass consumes those annotations to generate
arguments that implement the packed based TVM FFI API.
- For static shapes, the `TensorType::shape`, `TensorType::strides`,
- and `TensorType::elem_offset` fields are used to
+ For static shapes, the tensor shape and layout are used to
generate runtime checks on the corresponding member variables in
the user-provided `DLTensor*` or `tvm_tirx.runtime.tensor` argument. (e.g. A
Function that accepts a tensor of shape `[16,32]` validates that
diff --git a/tirx/src/backend/cuda/op/target_builtin.cc b/tirx/src/backend/cuda/op/target_builtin.cc
index 0dbaf36fca..edbc2e0e54 100644
--- a/tirx/src/backend/cuda/op/target_builtin.cc
+++ b/tirx/src/backend/cuda/op/target_builtin.cc
@@ -62,8 +62,6 @@ Type InferTypeDeclTmem(const CallNode* call) {
addr_type.bits() == 32,
TypeError)
<< "decl_tmem address must be a scalar 32-bit integer";
- TVM_FFI_CHECK(prim::IsZero(type->elem_offset) && type->strides.empty(), ValueError)
- << "decl_tmem addresses use physical columns; express tensor offsets through the layout";
return type;
}
diff --git a/tirx/src/backend/trn/transform/lower_trainium_layout.cc b/tirx/src/backend/trn/transform/lower_trainium_layout.cc
index c90b4fc7e2..896b86ad5d 100644
--- a/tirx/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/tirx/src/backend/trn/transform/lower_trainium_layout.cc
@@ -72,7 +72,19 @@ class TrainiumLayoutApplier : public tirx::StmtExprMutatorWithAnalyzer {
tensor.value().scope() == runtime::MemoryScope::Global()) {
TensorVar flattened = storage_lower->GetFlattenedTensor(tensor.value());
auto type = ffi::make_object<TensorTypeNode>(*tensor.value().ty().get());
- type->layout = std::nullopt;
+ if (const auto* tile = type->layout.as<TileLayoutNode>()) {
+ // Static physical geometry is already embedded in the lowered indices.
+ // Retain symbolic geometry that may still define parameter variables.
+ bool static_geometry = tile->replica.empty();
+ for (const auto& iter : tile->shard) {
+ static_geometry &= iter->extent.as<IntImmNode>() && iter->stride.as<IntImmNode>();
+ }
+ bool lower_layout = static_geometry || type->HasDefaultLayout(ana);
+ for (const auto& [axis, offset] : tile->offset) {
+ lower_layout &= ana->CanProveEqual(offset, 0);
+ }
+ if (lower_layout) type->layout = std::nullopt;
+ }
TensorVar source =
TensorVar(tensor.value()->name, TensorType(std::move(type)), tensor.value()->loc);
param_flattened_tensors.emplace_back(flattened, source);
@@ -169,7 +181,7 @@ class TrainiumLayoutApplier : public tirx::StmtExprMutatorWithAnalyzer {
flattened = buf;
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
type->shape = new_shape;
- type->strides = {};
+ type->layout = std::nullopt;
} else if (is_alloc) {
if (auto tile_layout = buf.ty()->layout.as<TileLayoutNode>();
tile_layout && tile_layout->HasThreadAxis()) {
@@ -193,13 +205,13 @@ class TrainiumLayoutApplier : public tirx::StmtExprMutatorWithAnalyzer {
flattened = buf;
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
type->shape = {ana->Simplify(mem_span)};
- type->strides = {};
+ type->layout = std::nullopt;
} else {
- flattened = FlattenedTensor(buf);
+ flattened = FlattenedTensor(buf, ffi::GetRef<sym::Analyzer>(analyzer_));
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
}
} else {
- flattened = FlattenedTensor(buf);
+ flattened = FlattenedTensor(buf, ffi::GetRef<sym::Analyzer>(analyzer_));
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
}
if (flattened.dtype()->dtype == DLDataType{kDLBool, 8, 1}) {
@@ -208,9 +220,7 @@ class TrainiumLayoutApplier : public tirx::StmtExprMutatorWithAnalyzer {
for (size_t i = 0; i < type->shape.size(); ++i) {
type->shape.Set(i, analyzer_->canonical_simplify(type->shape[i]));
}
- type->layout = std::nullopt;
- type->elem_offset = StmtExprMutator::Mutate(buf.elem_offset(), InplaceMode::kDisallow)
- .ValueOrUnchanged(buf.elem_offset());
+ if (!IsTrainiumLayout(trn_layout)) type->layout = std::nullopt;
flattened = TensorVar(flattened->name, TensorType(std::move(type)), flattened->loc);
VarRemapSet(buf, flattened);
@@ -261,9 +271,13 @@ class TrainiumLayoutApplier : public tirx::StmtExprMutatorWithAnalyzer {
const ffi::Array<PrimExpr>& indices) {
static const auto scope_trn_psum = runtime::MemoryScope::Get("trn.psum");
if (tensor.ty()->layout.has_value()) {
- auto tile_layout = tensor.ty()->layout.value().as<TileLayoutNode>();
+ auto layout = tensor.ty()->layout.value();
+ layout = StmtExprMutator::Mutate(ffi::AnyView(layout), InplaceMode::kDisallow)
+ .ValueOrUnchanged(layout)
+ .as_or_throw<Layout>();
+ auto tile_layout = layout.as<TileLayoutNode>();
if (IsTrainiumLayout(tile_layout)) {
- auto coord = tensor.ty()->layout.value()->Apply(indices, tensor.shape());
+ auto coord = layout->Apply(indices, tensor.shape());
std::vector<PrimExpr> res;
for (const auto& axis : tensor.scope() == scope_trn_psum
? ffi::Array<ffi::String>{"Bank", "P", "F"}
@@ -282,11 +296,11 @@ class TrainiumLayoutApplier : public tirx::StmtExprMutatorWithAnalyzer {
<< "layout: unable to verify that the coordinate matches the current thread. "
<< "Use .view() + .local() to decompose thread and memory axes.";
}
- auto res = tensor.ty()->layout.value()->Canonicalize()->Apply(indices, tensor.shape());
+ auto res = layout->Canonicalize()->Apply(indices, tensor.shape());
TVM_FFI_ICHECK_EQ(res.size(), 1) << "Expected a single element offset";
return {analyzer_->Simplify((*res.begin()).second)};
}
- auto flattened_indices = tensor.ty()->ElemOffset(indices, true);
+ auto flattened_indices = tensor.ty()->ElemOffset(indices);
TVM_FFI_ICHECK_EQ(flattened_indices.size(), 1) << "Expected a single element offset";
return {analyzer_->Simplify(flattened_indices[0])};
}
diff --git a/tirx/src/tirx/analysis/check_well_formed.cc b/tirx/src/tirx/analysis/check_well_formed.cc
index 2dce2f9fec..8b9261c368 100644
--- a/tirx/src/tirx/analysis/check_well_formed.cc
+++ b/tirx/src/tirx/analysis/check_well_formed.cc
@@ -24,9 +24,11 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/function.h>
#include <tvm/ir/module.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/layout.h>
#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
@@ -255,13 +257,21 @@ class WellFormednessVerifier : public StmtExprVisitor {
};
ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(tensor.shape()[i], walk);
}
- for (size_t i = 0; i < tensor.strides().size(); ++i) {
- if (auto var = tensor.strides()[i].as<Var>()) {
- define_if_needed(var.value(), field->Attr("strides")->ArrayItem(i));
+ if (const auto* tile = tensor.ty()->layout.as<TileLayoutNode>();
+ tile && MatchShapeStridePattern(tensor.ty()).has_value()) {
+ auto site = field->Attr("layout");
+ for (size_t i = 0; i < tile->shard.size(); ++i) {
+ if (auto var = tile->shard[i]->stride.as<Var>()) {
+ define_if_needed(var.value(), site->Attr("shard")->ArrayItem(i)->Attr("stride"));
+ }
+ }
+ for (const auto& [axis, offset] : tile->offset) {
+ auto walk = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ define_if_needed(var, site->Attr("offset"));
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(offset, walk);
}
- }
- if (auto var = tensor.elem_offset().as<Var>()) {
- define_if_needed(var.value(), field->Attr("elem_offset"));
}
}
diff --git a/tirx/src/tirx/ir/layout/layout.cc b/tirx/src/tirx/ir/layout/layout.cc
index b56058c026..e7c1e733f9 100644
--- a/tirx/src/tirx/ir/layout/layout.cc
+++ b/tirx/src/tirx/ir/layout/layout.cc
@@ -16,6 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
+
#include "utils.h"
namespace tvm {
diff --git a/tirx/src/tirx/ir/layout/tile_core.cc b/tirx/src/tirx/ir/layout/tile_core.cc
index 744c399718..e9ee7cd3b9 100644
--- a/tirx/src/tirx/ir/layout/tile_core.cc
+++ b/tirx/src/tirx/ir/layout/tile_core.cc
@@ -307,6 +307,30 @@ ffi::Array<PrimExpr> TileLayoutNode::GetShardShape() const {
return shard.Map([](const Iter& iter) { return iter->extent; });
}
+bool TileLayoutNode::IsContiguous(ffi::Optional<sym::Analyzer> analyzer) const {
+ if (!replica.empty()) return false;
+ static const Axis memory_axis = Axis::Get("m");
+ for (const auto& [axis, origin] : offset) {
+ if (axis != memory_axis) return false;
+ }
+ for (const auto& iter : shard) {
+ if (iter->axis != memory_axis) return false;
+ }
+ prim::ExprDeepEqual deep_equal;
+ auto equal = [&](const PrimExpr& lhs, const PrimExpr& rhs) {
+ if (deep_equal(lhs, rhs)) return true;
+ if (!analyzer.has_value()) analyzer = sym::Analyzer();
+ return analyzer.value()->CanProveEqual(lhs, rhs);
+ };
+ PrimExpr stride = IntImm(shard.empty() ? PrimType::Int(64) : shard.back()->stride.ty(), 1);
+ for (size_t i = shard.size(); i > 0; --i) {
+ const auto& iter = shard[i - 1];
+ if (!prim::IsOne(iter->extent) && !equal(iter->stride, stride)) return false;
+ stride = stride * iter->extent;
+ }
+ return true;
+}
+
bool TileLayoutNode::IsTrivial() const {
if (shard.size() > 1) return false;
if (shard.size() == 1) {
@@ -317,6 +341,10 @@ bool TileLayoutNode::IsTrivial() const {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef().def("tirx.TileLayoutIsContiguous",
+ [](const TileLayout& layout, ffi::Optional<sym::Analyzer> analyzer) {
+ return layout->IsContiguous(std::move(analyzer));
+ });
refl::GlobalDef().def("tirx.TileLayoutIsTrivial", [](const TileLayout& layout) {
return layout->Canonicalize().as<TileLayout>().value()->IsTrivial();
});
diff --git a/tirx/src/tirx/ir/specialize.cc b/tirx/src/tirx/ir/specialize.cc
index 169395c8c6..7abf56b4e3 100644
--- a/tirx/src/tirx/ir/specialize.cc
+++ b/tirx/src/tirx/ir/specialize.cc
@@ -260,46 +260,18 @@ class FunctionSpecializer : public StmtExprMutator {
ffi::Array<PrimExpr> shape =
tensor.shape().Map([this](const PrimExpr& e) { return Mutate(e).ValueOrUnchanged(e); });
- ffi::Array<PrimExpr> strides =
- tensor.strides().Map([this](const PrimExpr& e) { return Mutate(e).ValueOrUnchanged(e); });
-
- PrimExpr elem_offset =
- Mutate(tensor.elem_offset(), InplaceMode::kDisallow).ValueOrUnchanged(tensor.elem_offset());
-
- // Layout iter extents/strides may reference the same shape vars; remap
- // them in lock-step with shape (otherwise the specialized tensor keeps
- // stale layout extents from before specialization).
ffi::Optional<Layout> layout = tensor.ty()->layout;
- bool layout_changed = false;
- if (tensor.ty()->layout.has_value()) {
- if (auto opt_tile = tensor.ty()->layout.value().as<TileLayoutNode>()) {
- auto remap_iter = [this](const Iter& it) -> Iter {
- PrimExpr new_extent = Mutate(it->extent).ValueOrUnchanged(it->extent);
- PrimExpr new_stride = Mutate(it->stride).ValueOrUnchanged(it->stride);
- if (new_extent.same_as(it->extent) && new_stride.same_as(it->stride)) {
- return it;
- }
- return Iter(new_extent, new_stride, it->axis);
- };
- auto new_shard = opt_tile->shard.Map(remap_iter);
- auto new_replica = opt_tile->replica.Map(remap_iter);
- if (!new_shard.same_as(opt_tile->shard) || !new_replica.same_as(opt_tile->replica)) {
- layout = TileLayout(new_shard, new_replica, opt_tile->offset);
- layout_changed = true;
- }
- }
- }
+ auto mapped_layout = Mutate(ffi::AnyView(layout), InplaceMode::kDisallow);
+ bool layout_changed = !mapped_layout.UnchangedOrSameAs(layout);
+ layout = std::move(mapped_layout).ValueOrUnchanged(layout).cast<ffi::Optional<Layout>>();
bool storage_scope_changed = specialized_storage_scope.has_value() &&
specialized_storage_scope.value() != tensor.scope();
- if (tensor.elem_offset().same_as(elem_offset) && tensor.shape().same_as(shape) &&
- tensor.strides().same_as(strides) && !layout_changed && !storage_scope_changed) {
+ if (tensor.shape().same_as(shape) && !layout_changed && !storage_scope_changed) {
return tensor;
} else {
auto n = ffi::make_object<TensorTypeNode>(*tensor.ty().get());
- n->elem_offset = std::move(elem_offset);
n->shape = std::move(shape);
- n->strides = std::move(strides);
if (layout_changed) {
n->layout = std::move(layout);
}
@@ -378,7 +350,12 @@ void UpdateSpecializeVarMap(const Function& func, const Var& param, const Tensor
}
return lhs.same_as(rhs);
};
- if (!expr_equal(new_expr, old_expr)) {
+ Expr resolved = old_expr;
+ if (auto prim_expr = old_expr.as<PrimExpr>()) {
+ auto mutator = ffi::make_object<FunctionSpecializer>(*var_map);
+ resolved = mutator->Mutate(prim_expr.value()).ValueOrUnchanged(prim_expr.value());
+ }
+ if (!expr_equal(new_expr, resolved)) {
auto maybe_var = old_expr.as<Var>();
TVM_FFI_CHECK(maybe_var, TypeError)
<< "The signature of target tensor exprected an independent Var, but got " << old_expr
@@ -400,33 +377,29 @@ void UpdateSpecializeVarMap(const Function& func, const Var& param, const Tensor
<< "The tensor dimensions mismatched" << buf_to_specialize.shape().size() << " vs. "
<< specific_buf.shape().size() << ".";
- TVM_FFI_CHECK(specific_buf.strides().size() == buf_to_specialize.strides().size(), ValueError)
- << "The tensor strides dimensions mismatched" << buf_to_specialize.strides().size() << " vs. "
- << specific_buf.strides().size() << ".";
-
- // Updating var mapping using specific_expr
for (size_t i = 0; i < specific_buf.shape().size(); ++i) {
build_var_mapping(specific_buf.shape()[i], buf_to_specialize.shape()[i]);
}
- for (size_t i = 0; i < specific_buf.strides().size(); ++i) {
- build_var_mapping(specific_buf.strides()[i], buf_to_specialize.strides()[i]);
+ auto specific = MatchShapeStridePattern(specific_buf.ty());
+ auto original = MatchShapeStridePattern(buf_to_specialize.ty());
+ TVM_FFI_CHECK(specific.has_value() && original.has_value(), ValueError)
+ << "Tensor specialization requires shape-aligned strided memory layouts";
+ auto new_strides = specific.value().get<1>();
+ auto old_strides = original.value().get<1>();
+ for (size_t i = 0; i < old_strides.size(); ++i) {
+ build_var_mapping(new_strides[i], old_strides[i]);
}
- build_var_mapping(specific_buf.elem_offset(), buf_to_specialize.elem_offset());
+ build_var_mapping(specific.value().get<2>(), original.value().get<2>());
// The specializer distinguishes a concrete tensor constraint (which keeps
// the parameter) from another function parameter (which aliases it).
build_var_mapping(specific_buf, buf_to_specialize);
- // Check data_alignment and offset_factor.
+ // Check data_alignment.
// These two signatures are int, so we do not need map them.
TVM_FFI_CHECK_EQ(specific_buf.ty()->data_alignment, buf_to_specialize.ty()->data_alignment,
ValueError)
<< "The tensor data_alignment mismatched" << buf_to_specialize.ty()->data_alignment << " vs. "
<< specific_buf.ty()->data_alignment << ".";
-
- TVM_FFI_CHECK_EQ(specific_buf.ty()->offset_factor, buf_to_specialize.ty()->offset_factor,
- ValueError)
- << "The tensor offset_factor mismatched" << buf_to_specialize.ty()->offset_factor << " vs. "
- << specific_buf.ty()->offset_factor << ".";
}
/*!
diff --git a/tirx/src/tirx/ir/tensor.cc b/tirx/src/tirx/ir/tensor.cc
index 28a7b95e1a..9cb8d2f734 100644
--- a/tirx/src/tirx/ir/tensor.cc
+++ b/tirx/src/tirx/ir/tensor.cc
@@ -402,35 +402,26 @@ inline PrimExpr MergeMulMod(sym::AnalyzerObj* analyzer, const PrimExpr& base) {
// The tensor offset in convention of number of elements of
// original data ignoring number of lanes.
// We also perform optimization to simplify the indexing expression.
-ffi::Array<PrimExpr> TensorTypeNode::ElemOffset(ffi::Array<PrimExpr> input_indices,
- bool inner) const {
+ffi::Array<PrimExpr> TensorTypeNode::ElemOffset(ffi::Array<PrimExpr> input_indices) const {
TVM_FFI_ICHECK_EQ(shape.size(), input_indices.size())
<< "TensorType is " << shape.size() << "-dimensional, cannot be indexed with the "
<< input_indices.size() << "-dimensional indices provided.";
- if (strides.size()) {
- TVM_FFI_ICHECK_EQ(this->strides.size(), input_indices.size())
- << "If strides are defined, "
- << "the index's dimensionality must match the dimensionality of the index given.";
- }
-
- PrimExpr output_index = 0;
sym::Analyzer ana;
-
- for (size_t i = 0; i < input_indices.size(); i++) {
- if (strides.size()) {
- output_index = output_index + input_indices[i] * strides[i];
- } else {
- output_index = output_index * this->shape[i] + input_indices[i];
- }
-
- if (i > 0) {
- output_index = MergeMulMod(ana.get(), output_index);
+ if (layout.has_value()) {
+ auto mapped = layout.value()->Apply(input_indices, shape);
+ if (mapped.empty() && shape.empty()) {
+ return {IntImm(DefaultIndexType(), 0)};
}
+ TVM_FFI_CHECK(mapped.size() == 1 && mapped.find("m") != mapped.end(), ValueError)
+ << "Tensor indexing requires a memory-only layout";
+ PrimExpr offset = mapped["m"];
+ return SimplifyArray(ana.get(), {offset});
}
-
- if (elem_offset.defined() && !IsZero(elem_offset) && !inner) {
- output_index = output_index + elem_offset;
+ PrimExpr output_index = 0;
+ for (size_t i = 0; i < input_indices.size(); ++i) {
+ output_index = output_index * shape[i] + input_indices[i];
+ if (i > 0) output_index = MergeMulMod(ana.get(), output_index);
}
return SimplifyArray(ana.get(), {output_index});
diff --git a/tirx/src/tirx/ir/type.cc b/tirx/src/tirx/ir/type.cc
index fc0623acd1..13a62ab23e 100644
--- a/tirx/src/tirx/ir/type.cc
+++ b/tirx/src/tirx/ir/type.cc
@@ -21,12 +21,12 @@
* \file tirx/ir/type.cc
* \brief Types specific to TIRX.
*/
-#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
#include <tvm/runtime/device_api.h>
+#include <tvm/sym/analyzer.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/type.h>
@@ -34,13 +34,38 @@
namespace tvm::tirx {
+bool TensorTypeNode::HasDefaultLayout(ffi::Optional<sym::Analyzer> analyzer) const {
+ if (!layout.has_value()) return true;
+ const auto* tile = layout.as<TileLayoutNode>();
+ if (!tile || !tile->replica.empty()) return false;
+ static const Axis memory_axis = Axis::Get("m");
+ for (const auto& iter : tile->shard) {
+ if (iter->axis != memory_axis) return false;
+ }
+ prim::ExprDeepEqual deep_equal;
+ auto equal = [&](const PrimExpr& lhs, const PrimExpr& rhs) {
+ if (deep_equal(lhs, rhs)) return true;
+ if (!analyzer.has_value()) analyzer = sym::Analyzer();
+ return analyzer.value()->CanProveEqual(lhs, rhs);
+ };
+ for (const auto& [axis, offset] : tile->offset) {
+ if (axis != memory_axis || !equal(offset, IntImm(offset.ty(), 0))) return false;
+ }
+ PrimExpr extent = IntImm(DefaultIndexType(), 1);
+ for (size_t i = tile->shard.size(); i > 0; --i) extent = extent * tile->shard[i - 1]->extent;
+ PrimExpr size = IntImm(DefaultIndexType(), 1);
+ for (size_t i = shape.size(); i > 0; --i) size = size * shape[i - 1];
+ if (!equal(extent, size)) return false;
+ // Delegate last so any analyzer created above is reused by the stride proof.
+ return tile->IsContiguous(std::move(analyzer));
+}
+
bool TensorTypeNode::IsScalar(bool alloc_or_decl) const {
// TODO(@bohan): logical scope is not considered
- return shape.size() == 1 && tvm::prim::IsOne(shape[0]) && strides.empty() &&
- (!alloc_or_decl || tvm::prim::IsZero(elem_offset)) && data_alignment == 64 &&
- offset_factor == 1 &&
+ return shape.size() == 1 && tvm::prim::IsOne(shape[0]) && data_alignment == 64 &&
(!layout.has_value() ||
- ffi::StructuralEqual()(layout.value(), TileLayoutNode::DefaultLayout({1})));
+ (layout.as<TileLayoutNode>() &&
+ layout.value()->Canonicalize().as_or_throw<TileLayout>()->IsTrivial()));
}
std::optional<int64_t> TensorTypeNode::ConstantAllocationSize() const {
@@ -74,92 +99,55 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorMapTypeMaybeInpla
TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TensorTypeVisit(
ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
- // skips: scope, data_alignment, offset_factor
+ // skips: scope, data_alignment
const TensorTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->dtype));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->shape));
- // Empty strides denote the common compact layout. Broad callbacks do not see the empty
- // container; explicit strides retain normal container descent and callback behavior.
- if (!self->strides.empty()) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->strides));
- }
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->elem_offset));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->layout));
return std::nullopt;
}
TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorTypeMutate(
ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: scope, data_alignment, offset_factor
+ // skips: scope, data_alignment
const TensorTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorTypeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimType>, mapped_dtype,
mutator->MutateExpected(self->dtype));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<PrimExpr>>, mapped_shape,
mutator->MutateExpected(self->shape));
- // Empty strides denote the common compact layout. Broad callbacks do not see the empty
- // container; explicit strides retain normal container descent and callback behavior.
- ffi::UnchangedOr<ffi::Array<PrimExpr>> mapped_strides = ffi::Unchanged();
- if (!self->strides.empty()) {
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<PrimExpr>>, descended_strides,
- mutator->MutateExpected(self->strides));
- mapped_strides = std::move(descended_strides);
- }
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_elem_offset,
- mutator->MutateExpected(self->elem_offset));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<Layout>>, mapped_layout,
mutator->MutateExpected(self->layout));
if (mapped_dtype.UnchangedOrSameAs(self->dtype) && mapped_shape.UnchangedOrSameAs(self->shape) &&
- mapped_strides.UnchangedOrSameAs(self->strides) &&
- mapped_elem_offset.UnchangedOrSameAs(self->elem_offset) &&
mapped_layout.UnchangedOrSameAs(self->layout)) {
return ffi::Unchanged();
}
ffi::ObjectPtr<TensorTypeNode> copy = ffi::make_object<TensorTypeNode>(*self);
copy->dtype = std::move(mapped_dtype).ValueOrUnchanged(std::move(copy->dtype));
copy->shape = std::move(mapped_shape).ValueOrUnchanged(std::move(copy->shape));
- copy->strides = std::move(mapped_strides).ValueOrUnchanged(std::move(copy->strides));
- copy->elem_offset = std::move(mapped_elem_offset).ValueOrUnchanged(std::move(copy->elem_offset));
copy->layout = std::move(mapped_layout).ValueOrUnchanged(std::move(copy->layout));
return ffi::Any(std::move(copy));
}
TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorTypeMaybeInplaceMutate(
ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: scope, data_alignment, offset_factor
+ // skips: scope, data_alignment
TensorTypeNode* self = const_cast<TensorTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorTypeNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimType>, mapped_dtype,
mutator->MutateExpected(self->dtype, ffi::InplaceMode::kAllow));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<PrimExpr>>, mapped_shape,
mutator->MutateExpected(self->shape, ffi::InplaceMode::kAllow));
- // Empty strides denote the common compact layout. Broad callbacks do not see the empty
- // container; explicit strides retain normal container descent and callback behavior.
- ffi::UnchangedOr<ffi::Array<PrimExpr>> mapped_strides = ffi::Unchanged();
- if (!self->strides.empty()) {
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
- ffi::UnchangedOr<ffi::Array<PrimExpr>>, descended_strides,
- mutator->MutateExpected(self->strides, ffi::InplaceMode::kAllow));
- mapped_strides = std::move(descended_strides);
- }
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
- ffi::UnchangedOr<PrimExpr>, mapped_elem_offset,
- mutator->MutateExpected(self->elem_offset, ffi::InplaceMode::kAllow));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
ffi::UnchangedOr<ffi::Optional<Layout>>, mapped_layout,
mutator->MutateExpected(self->layout, ffi::InplaceMode::kAllow));
if (mapped_dtype.UnchangedOrSameAs(self->dtype) && mapped_shape.UnchangedOrSameAs(self->shape) &&
- mapped_strides.UnchangedOrSameAs(self->strides) &&
- mapped_elem_offset.UnchangedOrSameAs(self->elem_offset) &&
mapped_layout.UnchangedOrSameAs(self->layout)) {
return ffi::Unchanged();
}
if (!mapped_dtype.IsUnchanged()) self->dtype = std::move(mapped_dtype).ValueUnchecked();
if (!mapped_shape.IsUnchanged()) self->shape = std::move(mapped_shape).ValueUnchecked();
- if (!mapped_strides.IsUnchanged()) self->strides = std::move(mapped_strides).ValueUnchecked();
- if (!mapped_elem_offset.IsUnchanged())
- self->elem_offset = std::move(mapped_elem_offset).ValueUnchecked();
if (!mapped_layout.IsUnchanged()) self->layout = std::move(mapped_layout).ValueUnchecked();
return ffi::Unchanged();
}
@@ -167,19 +155,14 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorTypeMaybeInplaceM
} // namespace
TensorType::TensorType(ffi::Array<PrimExpr> shape, PrimType dtype, runtime::MemoryScope scope,
- ffi::Array<PrimExpr> strides, ffi::Optional<PrimExpr> elem_offset,
- int data_alignment, int offset_factor, ffi::Optional<Layout> layout,
- Location loc)
+ int data_alignment, ffi::Optional<Layout> layout, Location loc)
: Type(ffi::UnsafeInit{}) {
- PrimExpr offset = elem_offset.value_or(IntImm(DefaultIndexType(), 0));
- auto n = ffi::make_object<TensorTypeNode>(std::move(offset));
+ auto n = ffi::make_object<TensorTypeNode>();
n->dtype = std::move(dtype);
n->scope = std::move(scope);
n->shape = std::move(shape);
- n->strides = std::move(strides);
n->data_alignment =
data_alignment <= 0 ? static_cast<int>(runtime::kAllocAlignment) : data_alignment;
- n->offset_factor = offset_factor == 0 ? 1 : offset_factor;
n->layout = std::move(layout);
n->loc = loc;
data_ = std::move(n);
@@ -188,6 +171,10 @@ TensorType::TensorType(ffi::Array<PrimExpr> shape, PrimType dtype, runtime::Memo
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TensorTypeNode::RegisterReflection();
+ refl::GlobalDef().def("tirx.TensorTypeHasDefaultLayout",
+ [](const TensorType& tensor, ffi::Optional<sym::Analyzer> analyzer) {
+ return tensor->HasDefaultLayout(std::move(analyzer));
+ });
refl::TypeAttrDef<TensorTypeNode>()
.attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TensorTypeVisit>())
.attr(ffi::type_attr::kStructuralMutate,
@@ -195,21 +182,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(ffi::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&TensorTypeMaybeInplaceMutate>());
- refl::GlobalDef().def(
- "tirx.TensorType",
- [](ffi::Array<PrimExpr> shape, PrimType dtype, runtime::MemoryScope scope,
- ffi::Array<PrimExpr> strides, ffi::AnyView elem_offset, int data_alignment,
- int offset_factor, ffi::Optional<Layout> layout, Location loc) {
- ffi::Optional<PrimExpr> offset;
- if (auto integer = elem_offset.as<int64_t>()) {
- offset = IntImm(DefaultIndexType(), integer.value());
- } else {
- offset = elem_offset.cast<ffi::Optional<PrimExpr>>();
- }
- return TensorType(std::move(shape), std::move(dtype), std::move(scope), std::move(strides),
- std::move(offset), data_alignment, offset_factor, std::move(layout),
- std::move(loc));
- });
+ refl::GlobalDef().def("tirx.TensorType",
+ [](ffi::Array<PrimExpr> shape, PrimType dtype, runtime::MemoryScope scope,
+ int data_alignment, ffi::Optional<Layout> layout, Location loc) {
+ return TensorType(std::move(shape), std::move(dtype), std::move(scope),
+ data_alignment, std::move(layout), std::move(loc));
+ });
}
TensorMapType::TensorMapType(Location loc) : Type(ffi::UnsafeInit{}) {
diff --git a/tirx/src/tirx/script/ir_builder/ir.cc b/tirx/src/tirx/script/ir_builder/ir.cc
index 6c6f690869..34d802cc86 100644
--- a/tirx/src/tirx/script/ir_builder/ir.cc
+++ b/tirx/src/tirx/script/ir_builder/ir.cc
@@ -45,17 +45,12 @@ using tvm::tirx::Layout;
namespace {
tvm::tirx::TensorType TensorTypeDecl(ffi::Array<PrimExpr> shape, PrimType dtype,
- runtime::MemoryScope storage_scope,
- ffi::Optional<ffi::Array<PrimExpr>> strides,
- ffi::Optional<PrimExpr> elem_offset, int align,
- int offset_factor, ffi::Optional<Layout> layout,
- ffi::Optional<Expr> data) {
+ runtime::MemoryScope storage_scope, int align,
+ ffi::Optional<Layout> layout, ffi::Optional<Expr> data) {
if (data.has_value()) {
storage_scope = data.value()->ty.as_or_throw<PtrType>()->scope;
}
- return tvm::tirx::TensorType(shape, dtype, storage_scope,
- strides.value_or(ffi::Array<PrimExpr>()), elem_offset, align,
- offset_factor, layout);
+ return tvm::tirx::TensorType(shape, dtype, storage_scope, align, layout);
}
} // namespace
diff --git a/tirx/src/tirx/script/printer/tensor.cc b/tirx/src/tirx/script/printer/tensor.cc
index 2a7abcb91a..2a417c90e0 100644
--- a/tirx/src/tirx/script/printer/tensor.cc
+++ b/tirx/src/tirx/script/printer/tensor.cc
@@ -44,19 +44,14 @@ ffi::Optional<ExprDoc> TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const tirx::TensorTypeNode>(input);
// The tensor type constructor normalizes these fields.
// Keep the explicit type constructor when that normalization would change the IR.
- if (tensor->data_alignment <= 0 || tensor->offset_factor == 0) {
+ if (tensor->data_alignment <= 0) {
return NamespaceDoc("ir")
->Attr("make_node")
->Call({LiteralDoc::Str("tirx.TensorType", std::nullopt)},
- {"dtype", "scope", "shape", "strides", "elem_offset", "data_alignment",
- "offset_factor", "layout"},
+ {"dtype", "scope", "shape", "data_alignment", "layout"},
{d->Translate(tensor->dtype).value(), d->Translate(tensor->scope).value(),
- d->Translate(tensor->shape).value(), d->Translate(tensor->strides).value(),
- d->WithExtraStateScope<bool>(
- "ir.explicit_value_type", true,
- [&] { return d->Translate(tensor->elem_offset).value(); }),
+ d->Translate(tensor->shape).value(),
LiteralDoc::Int(tensor->data_alignment, std::nullopt),
- LiteralDoc::Int(tensor->offset_factor, std::nullopt),
tensor->layout.has_value() ? d->Translate(tensor->layout.value()).value()
: LiteralDoc::None(std::nullopt)});
}
@@ -86,18 +81,10 @@ ffi::Optional<ExprDoc> TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
if (tensor->scope != runtime::MemoryScope::Global()) {
optional("scope", LiteralDoc::Str(tensor->scope.name(), std::nullopt));
}
- if (!tensor->strides.empty()) optional("strides", d->Translate(tensor->strides).value());
- if (!ffi::StructuralEqual()(tensor->elem_offset, IntImm(tirx::DefaultIndexType(), 0))) {
- optional("elem_offset", d->WithExtraStateScope<bool>("ir.explicit_value_type", true, [&] {
- return d->Translate(tensor->elem_offset).value();
- }));
- }
if (tensor->data_alignment != static_cast<int>(runtime::kAllocAlignment)) {
optional("data_alignment", LiteralDoc::Int(tensor->data_alignment, std::nullopt));
}
- if (tensor->offset_factor != 1) {
- optional("offset_factor", LiteralDoc::Int(tensor->offset_factor, std::nullopt));
- }
+
if (tensor->layout.has_value()) optional("layout", d->Translate(tensor->layout.value()).value());
return NamespaceDoc("tirx")->Attr("Tensor")->Call(args, keys, values);
}
@@ -174,6 +161,7 @@ ffi::Optional<ExprDoc> TileLayoutDocTranslate(DocTranslatorObj* d, ffi::AnyView
keys.push_back("offset");
values.push_back(DictDoc(offset_keys, offset_values));
}
+
return NamespaceDoc("tirx")->Attr("TileLayout")->Attr("from_iters")->Call({}, keys, values);
}
@@ -194,6 +182,7 @@ ffi::Optional<ExprDoc> ComposeLayoutDocTranslate(DocTranslatorObj* d, ffi::AnyVi
keys.push_back("swizzle_inner");
values.push_back(LiteralDoc::Boolean(false, std::nullopt));
}
+
return NamespaceDoc("tirx")
->Attr("ComposeLayout")
->Call({LiteralDoc::Int(layout->per_element, std::nullopt),
diff --git a/tirx/src/tirx/transform/flatten_tensor.cc b/tirx/src/tirx/transform/flatten_tensor.cc
index cb3a31b6f3..bf944975c5 100644
--- a/tirx/src/tirx/transform/flatten_tensor.cc
+++ b/tirx/src/tirx/transform/flatten_tensor.cc
@@ -45,7 +45,7 @@ using namespace tvm::prim;
* \brief Flatten each n-d tensor ``buf`` into a 1-d storage view ``buf'``,
* rewriting every access ``buf[x]`` into ``buf'[f(x)]``.
*
- * The invariant: ``f(x) = layout.apply(x, shape) + elem_offset`` is fully
+ * The invariant: ``f(x) = layout.apply(x, shape)`` is fully
* determined by ``buf``'s geometry, and ``buf'`` is only a storage husk —
* same data origin, dtype, alignment and scope; no layout, no elem_offset.
*
@@ -125,41 +125,21 @@ class TensorFlattener : public StmtExprMutatorWithAnalyzer {
auto view_type = ffi::make_object<TensorTypeNode>(*buf.ty().get());
auto mutate_expr = [this](const PrimExpr& expr) { return Mutate(expr).ValueOrUnchanged(expr); };
view_type->shape = view_type->shape.Map(mutate_expr);
- view_type->strides = view_type->strides.Map(mutate_expr);
- if (view_type->elem_offset.defined()) {
- view_type->elem_offset = this->Mutate(view_type->elem_offset, InplaceMode::kDisallow)
- .ValueOrUnchanged(view_type->elem_offset);
- }
- if (auto tile = view_type->layout.as<TileLayoutNode>()) {
- auto remap_iter = [this](const Iter& iter) {
- PrimExpr extent =
- this->Mutate(iter->extent, InplaceMode::kDisallow).ValueOrUnchanged(iter->extent);
- PrimExpr stride =
- this->Mutate(iter->stride, InplaceMode::kDisallow).ValueOrUnchanged(iter->stride);
- if (extent.same_as(iter->extent) && stride.same_as(iter->stride)) {
- return iter;
- }
- return Iter(extent, stride, iter->axis);
- };
- auto shard = tile->shard.Map(remap_iter);
- auto replica = tile->replica.Map(remap_iter);
- if (!shard.same_as(tile->shard) || !replica.same_as(tile->replica)) {
- view_type->layout = TileLayout(shard, replica, tile->offset);
- }
+ if (view_type->layout.has_value()) {
+ view_type->layout = this->Mutate(ffi::AnyView(view_type->layout), InplaceMode::kDisallow)
+ .ValueOrUnchanged(view_type->layout)
+ .cast<ffi::Optional<Layout>>();
}
TensorVar fold_view = TensorVar(buf->name, TensorType(std::move(view_type)), buf->loc);
// buf': the storage husk. The linearized indices carry layout and
// elem_offset, so the husk keeps neither.
- auto flat = FlattenedTensor(fold_view);
+ auto flat = FlattenedTensor(fold_view, ffi::GetRef<sym::Analyzer>(analyzer_));
auto type = ffi::make_object<TensorTypeNode>(*flat.ty().get());
for (size_t i = 0; i < type->shape.size(); ++i) {
type->shape.Set(i, analyzer_->canonical_simplify(type->shape[i]));
}
type->layout = std::nullopt;
- if (type->elem_offset.defined() && !IsZero(type->elem_offset)) {
- type->elem_offset = IntImm(type->elem_offset.ty().as_or_throw<PrimType>(), 0);
- }
// Body-local tensors keep their identity when flattening changes nothing.
// Function-parameter tensors always rebuild: the epilogue aliases the
// rebuilt view onto the argument tensor with an explicit DeclTensor, and
diff --git a/tirx/src/tirx/transform/flattened_tensor.h b/tirx/src/tirx/transform/flattened_tensor.h
index 02f014e86c..e71669a279 100644
--- a/tirx/src/tirx/transform/flattened_tensor.h
+++ b/tirx/src/tirx/transform/flattened_tensor.h
@@ -33,17 +33,25 @@ namespace tirx {
// A changed type requires a fresh variable. Callers bind that variable to a
// decl_tensor over the source's physical pointer when using it as a view.
-TVM_FFI_INLINE TensorVar FlattenedTensor(const TensorVar& tensor) {
+TVM_FFI_INLINE TensorVar FlattenedTensor(const TensorVar& tensor,
+ ffi::Optional<sym::Analyzer> analyzer = std::nullopt) {
ffi::Array<PrimExpr> output_shape{1};
- if (tensor.strides().size()) {
- // Preserve the outer allocation footprint and cover every addressed element.
- // Permuted views can have their largest physical stride on an inner axis.
- TVM_FFI_ICHECK_EQ(tensor.shape().size(), tensor.strides().size());
+ const auto* tile = tensor.ty()->layout.as<TileLayoutNode>();
+ bool explicit_strides =
+ tile && !tile->shard.empty() && MatchShapeStridePattern(tensor.ty()).has_value();
+ if (explicit_strides) {
+ // Keep the outer allocation footprint and cover every addressed element.
PrimExpr span = 1;
- for (size_t i = 0; i < tensor.shape().size(); ++i) {
- span = span + (tensor.shape()[i] - 1) * tensor.strides()[i];
+ for (const auto& iter : tile->shard) {
+ span = span + (iter->extent - 1) * iter->stride;
+ }
+ // The flattened footprint excludes the origin, which the caller applies separately.
+ if (tile->IsContiguous(analyzer)) {
+ output_shape.Set(0, PrimExpr(1));
+ for (const auto& extent : tensor.shape()) output_shape.Set(0, output_shape[0] * extent);
+ } else {
+ output_shape.Set(0, tvm::max(tile->shard[0]->stride * tensor.shape()[0], span));
}
- output_shape.Set(0, prim::Max(tensor.strides()[0] * tensor.shape()[0], span));
} else {
// Otherwise, the flattened extent is the product of the input extents.
// This also flattens rank-0 tensors to a rank-1 tensor of shape [1].
@@ -52,15 +60,14 @@ TVM_FFI_INLINE TensorVar FlattenedTensor(const TensorVar& tensor) {
}
}
- if (output_shape.size() == tensor.shape().size() && tensor.strides().empty()) {
+ if (output_shape.size() == tensor.shape().size() && !tensor.ty()->layout.has_value()) {
return tensor;
} else {
// The original layout describes the old rank; reset it for the flattened shape.
- return TensorVar(
- tensor->name,
- TensorType(output_shape, tensor.dtype(), tensor.scope(), {}, tensor.elem_offset(),
- tensor.ty()->data_alignment, tensor.ty()->offset_factor, std::nullopt),
- tensor->loc);
+ return TensorVar(tensor->name,
+ TensorType(output_shape, tensor.dtype(), tensor.scope(),
+ tensor.ty()->data_alignment, std::nullopt),
+ tensor->loc);
}
}
diff --git a/tirx/src/tirx/transform/lower_tensor_layouts.cc b/tirx/src/tirx/transform/lower_tensor_layouts.cc
index 8a02e8ae37..a3a861ec2a 100644
--- a/tirx/src/tirx/transform/lower_tensor_layouts.cc
+++ b/tirx/src/tirx/transform/lower_tensor_layouts.cc
@@ -67,7 +67,7 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
storage_lower->tensor_aliases_.Set(tensor.value(), tensor.value());
TensorVar flattened = storage_lower->GetFlattenedTensor(tensor.value());
TensorVar source = tensor.value();
- if (source.ty()->layout.has_value()) {
+ if (source.ty()->layout.has_value() && source.ty()->HasDefaultLayout(ana)) {
auto type = ffi::make_object<TensorTypeNode>(*source.ty().get());
type->layout = std::nullopt;
source = TensorVar(source->name, TensorType(std::move(type)), source->loc);
@@ -212,7 +212,7 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
flattened = buf;
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
type->shape = {ana->Simplify(mem_span)};
- type->strides = {};
+ type->layout = std::nullopt;
} else if (auto compose = buf.ty()->layout.as<ComposeLayoutNode>()) {
// A swizzle permutes addresses, so the logical element count is not
// the allocation footprint of a strided view. Round the underlying
@@ -231,13 +231,13 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
}
type = ffi::make_object<TensorTypeNode>(*buf.ty().get());
type->shape = {analyzer_->Simplify(mem_span)};
- type->strides = {};
+ type->layout = std::nullopt;
} else {
- flattened = FlattenedTensor(buf);
+ flattened = FlattenedTensor(buf, ffi::GetRef<sym::Analyzer>(analyzer_));
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
}
} else {
- flattened = FlattenedTensor(buf);
+ flattened = FlattenedTensor(buf, ffi::GetRef<sym::Analyzer>(analyzer_));
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
}
// Remap variables the pass has already rebuilt (a shape may load from
@@ -249,17 +249,23 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
.ValueOrUnchanged(type->shape[i])
.as_or_throw<PrimExpr>()));
}
- for (size_t i = 0; i < type->strides.size(); ++i) {
- type->strides.Set(
- i, StmtExprMutator::Mutate(ffi::AnyView(type->strides[i]), InplaceMode::kDisallow)
- .ValueOrUnchanged(type->strides[i])
- .as_or_throw<PrimExpr>());
- }
type->layout = std::nullopt;
- type->elem_offset =
- StmtExprMutator::Mutate(ffi::AnyView(buf.elem_offset()), InplaceMode::kDisallow)
- .ValueOrUnchanged(buf.elem_offset())
- .as_or_throw<PrimExpr>();
+ // Keep a regular view's origin in its flat memory layout. Allocations have
+ // already included that origin in their span and use absolute indices.
+ if (!is_alloc) {
+ if (auto pattern = MatchShapeStridePattern(buf.ty()); pattern.has_value()) {
+ PrimExpr origin = pattern.value().get<2>();
+ origin = analyzer_->canonical_simplify(
+ StmtExprMutator::Mutate(ffi::AnyView(origin), InplaceMode::kDisallow)
+ .ValueOrUnchanged(origin)
+ .as_or_throw<PrimExpr>());
+ if (!prim::IsZero(origin)) {
+ static const auto memory_axis = Axis::Get("m");
+ type->layout =
+ TileLayout({Iter(type->shape[0], 1, memory_axis)}, {}, {{memory_axis, origin}});
+ }
+ }
+ }
if (ffi::StructuralEqual()(buf.ty(), TensorType(type))) {
return buf;
}
@@ -302,11 +308,25 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
<< "layout: unable to verify that the coordinate matches the current thread. "
<< "Use .view() + .local() to decompose thread and memory axes.";
}
- auto res = tensor.ty()->layout.value()->Canonicalize()->Apply(indices, tensor.shape());
+ auto layout = tensor.ty()->layout.value();
+ layout = StmtExprMutator::Mutate(ffi::AnyView(layout), InplaceMode::kDisallow)
+ .ValueOrUnchanged(layout)
+ .as_or_throw<Layout>();
+ auto flat = GetFlattenedTensor(tensor);
+ auto pattern = MatchShapeStridePattern(flat.ty()).value();
+ if (!prim::IsZero(pattern.get<2>())) {
+ // The flat view retains this origin, so only lower its relative indices.
+ auto tile = layout.as_or_throw<TileLayout>();
+ layout = TileLayout(tile->shard, tile->replica, {});
+ }
+ auto res = layout->Canonicalize()->Apply(indices, tensor.shape());
+ if (res.empty() && tensor.shape().empty()) {
+ return {IntImm(tensor.ty()->DefaultIndexType(), 0)};
+ }
TVM_FFI_ICHECK_EQ(res.size(), 1) << "Expected a single element offset";
return {analyzer_->Simplify((*res.begin()).second)};
}
- auto flattened_indices = tensor.ty()->ElemOffset(indices, true);
+ auto flattened_indices = tensor.ty()->ElemOffset(indices);
TVM_FFI_ICHECK_EQ(flattened_indices.size(), 1) << "Expected a single element offset";
return {analyzer_->Simplify(flattened_indices[0])};
}
diff --git a/tirx/src/tirx/transform/remove_no_op.cc b/tirx/src/tirx/transform/remove_no_op.cc
index e693b3bfd4..ed40b90b68 100644
--- a/tirx/src/tirx/transform/remove_no_op.cc
+++ b/tirx/src/tirx/transform/remove_no_op.cc
@@ -179,10 +179,6 @@ class NoOpRemover : public StmtExprMutatorWithAnalyzer {
if (const TensorLoadNode* load = store->value.as<TensorLoadNode>()) {
TensorVar tensor = load->source.as_or_throw<tvm::tirx::TensorVar>();
if (tensor.same_as(store->dest.as_or_throw<TensorVar>()) &&
- analyzer_->CanProveEqual(tensor.elem_offset(),
- store->dest.as_or_throw<TensorVar>().elem_offset()) &&
- ArrayValueEqual(tensor.shape(), store->dest.as_or_throw<TensorVar>().shape()) &&
- ArrayValueEqual(tensor.strides(), store->dest.as_or_throw<TensorVar>().strides()) &&
ArrayValueEqual(load->indices, store->indices)) {
return only_side_effects();
}
diff --git a/tirx/src/tirx/transform/stmt_simplify.cc b/tirx/src/tirx/transform/stmt_simplify.cc
index f80354f537..8de2d373ad 100644
--- a/tirx/src/tirx/transform/stmt_simplify.cc
+++ b/tirx/src/tirx/transform/stmt_simplify.cc
@@ -200,11 +200,7 @@ UnchangedOr<Stmt> StmtSimplifier::Mutate_(const TensorStoreNode* op, InplaceMode
if (const TensorLoadNode* load = store->value.as<TensorLoadNode>()) {
TensorVar tensor = load->source.as_or_throw<tvm::tirx::TensorVar>();
if (tensor.same_as(store->dest.as_or_throw<TensorVar>()) &&
- ArrayDeepEqual(load->indices, store->indices) &&
- prim::ExprDeepEqual()(tensor.elem_offset(),
- store->dest.as_or_throw<TensorVar>().elem_offset()) &&
- ArrayDeepEqual(tensor.shape(), store->dest.as_or_throw<TensorVar>().shape()) &&
- ArrayDeepEqual(tensor.strides(), store->dest.as_or_throw<TensorVar>().strides())) {
+ ArrayDeepEqual(load->indices, store->indices)) {
return Evaluate(0);
}
}
diff --git a/tirx/src/tirx/transform/storage_rewrite.cc b/tirx/src/tirx/transform/storage_rewrite.cc
index 6e0a2dc17c..9f128401b1 100644
--- a/tirx/src/tirx/transform/storage_rewrite.cc
+++ b/tirx/src/tirx/transform/storage_rewrite.cc
@@ -822,9 +822,7 @@ class StoragePlanRewriter : public StmtExprMutator {
}
// Flat packing and lifetime reuse cannot reconstruct arbitrary physical views.
// Preserve their original allocation before either planning step can combine it.
- if (type->shape.size() != 1 || type->layout.has_value() || !IsConstInt(type->elem_offset, 0) ||
- (!type->strides.empty() &&
- (type->strides.size() != 1 || !IsConstInt(type->strides[0], 1)))) {
+ if (type->shape.size() != 1 || type->layout.has_value()) {
return true;
}
if (!call->args.empty()) return true;
@@ -842,8 +840,7 @@ class StoragePlanRewriter : public StmtExprMutator {
static bool HasCompatibleAlignment(const StorageEntry* entry, const CallNode* call) {
auto first = entry->allocs[0]->value.as<CallNode>()->ty_args[0].as_or_throw<TensorType>();
auto type = call->ty_args[0].as_or_throw<TensorType>();
- return first->data_alignment == type->data_alignment &&
- first->offset_factor == type->offset_factor;
+ return first->data_alignment == type->data_alignment;
}
// The owning lowering handles single-allocation layouts that cannot be packed
@@ -856,10 +853,7 @@ class StoragePlanRewriter : public StmtExprMutator {
for (const auto* allocation : entry->allocs) {
const auto* call = allocation->value.as<CallNode>();
auto type = call->ty_args[0].as_or_throw<TensorType>();
- if (type->shape.size() != 1 || type->dtype.IsScalableVector() || type->layout.has_value() ||
- !IsConstInt(type->elem_offset, 0) ||
- (!type->strides.empty() &&
- (type->strides.size() != 1 || !IsConstInt(type->strides[0], 1))))
+ if (type->shape.size() != 1 || type->dtype.IsScalableVector() || type->layout.has_value())
return false;
}
return true;
@@ -883,7 +877,6 @@ class StoragePlanRewriter : public StmtExprMutator {
auto type = allocation->value.as<CallNode>()->ty_args[0].as_or_throw<TensorType>();
uint64_t element_bits = type->dtype.bits() * type->dtype.lanes();
if (!include_alignment(element_bits) ||
- !include_alignment(element_bits * std::max(type->offset_factor, 1)) ||
!include_alignment(static_cast<uint64_t>(std::max(type->data_alignment, 1)) * 8)) {
return std::nullopt;
}
@@ -1071,9 +1064,8 @@ class StoragePlanRewriter : public StmtExprMutator {
combo_size = analyzer_->Simplify(combo_size);
TensorVar buf(
e->alloc_var->name,
- TensorType({combo_size}, alloc_type, e->scope, {}, std::nullopt,
- first_call->ty_args[0].as_or_throw<TensorType>()->data_alignment,
- first_call->ty_args[0].as_or_throw<TensorType>()->offset_factor));
+ TensorType({combo_size}, alloc_type, e->scope,
+ first_call->ty_args[0].as_or_throw<TensorType>()->data_alignment));
e->alloc_var = buf;
ffi::Map<ffi::String, ffi::Any> annotations;
if (e->is_volatile) {
@@ -1116,8 +1108,8 @@ class StoragePlanRewriter : public StmtExprMutator {
shape_type = PrimType::Int(64);
}
PrimExpr alloc_size = MakeConst(shape_type, alloc_elements);
- TensorVar buf(e->alloc_var->name, TensorType({alloc_size}, e->elem_type, e->scope, {},
- std::nullopt, static_cast<int>(align / 8), 0));
+ TensorVar buf(e->alloc_var->name,
+ TensorType({alloc_size}, e->elem_type, e->scope, static_cast<int>(align / 8)));
e->alloc_var = buf;
for (StorageEntry* child : e->merged_children) {
child->alloc_var = e->alloc_var;
diff --git a/tirx/src/tirx/transform/tvm_ffi_binder.cc b/tirx/src/tirx/transform/tvm_ffi_binder.cc
index bad702479a..b86ffa5e46 100644
--- a/tirx/src/tirx/transform/tvm_ffi_binder.cc
+++ b/tirx/src/tirx/transform/tvm_ffi_binder.cc
@@ -445,99 +445,6 @@ TVMFFIABIBuilder::Result TVMFFIABIBuilder::Finalize() {
return result;
}
-// ============================================================
-// BindArray (array bind with ffi::reflection::AccessPath)
-// ============================================================
-
-void TVMFFIABIBuilder::BindArray(const ffi::Array<PrimExpr>& arg, const ffi::Array<PrimExpr>& value,
- const ffi::reflection::AccessPath& base_path) {
- TVM_FFI_ICHECK_EQ(arg.size(), value.size())
- << "Array size mismatch at " << RenderAccessPath(base_path);
- for (size_t i = 0; i < arg.size(); ++i) {
- ffi::reflection::AccessPath elem_path = base_path->ArrayItem(i);
- BindScalar(arg[i], value[i], elem_path, false);
- }
-}
-
-// ============================================================
-// BindTensor (tensor-to-tensor bind with ffi::reflection::AccessPath)
-// ============================================================
-
-void TVMFFIABIBuilder::BindTensor(const TensorVar& arg, const TensorVar& value,
- ffi::reflection::AccessPath base_path, bool fuzzy_match) {
- TVM_FFI_ICHECK_EQ(arg.scope(), value.scope())
- << "Argument " << arg->name << " TensorVar bind scope mismatch";
- TVM_FFI_ICHECK_EQ(arg.dtype(), value.dtype())
- << "Argument " << arg->name << " TensorVar bind data type mismatch";
- if (value.ty()->data_alignment % arg.ty()->data_alignment != 0) {
- LOG(WARNING) << "Trying to bind tensor to another one with lower alignment requirement "
- << " required alignment=" << arg.ty()->data_alignment
- << ", provided alignment=" << value.ty()->data_alignment;
- }
-
- if (value.elem_offset().defined()) {
- if (IsZero(arg.elem_offset())) {
- TVM_FFI_ICHECK(IsZero(value.elem_offset()))
- << "Trying to bind a TensorVar with offset into one without offset "
- << " required elem_offset=" << arg.elem_offset()
- << ", provided elem_offset=" << value.elem_offset();
- }
- ffi::reflection::AccessPath data_path = base_path->Attr(ffi::String("data"));
- ffi::reflection::AccessPath offset_path = base_path->Attr(ffi::String("elem_offset"));
- if (BindScalar(arg.elem_offset(), value.elem_offset(), offset_path, false)) {
- if (arg.ty()->offset_factor > 1) {
- PrimExpr offset = value.elem_offset();
- PrimExpr factor = IntImm(offset.ty(), arg.ty()->offset_factor);
- PrimExpr zero = IntImm(offset.ty(), 0);
- PrimExpr acond = analyzer_->Simplify(truncmod(offset, factor) == zero);
- if (IsZero(acond)) {
- TVM_FFI_THROW(InternalError)
- << "Bind have an unmet assertion at " << RenderAccessPath(offset_path);
- }
- if (!IsOne(acond)) {
- int param_index = GetParamIndex(base_path);
- int data_bytes =
- ((((arg.dtype()->dtype).bits * static_cast<int16_t>((arg.dtype()->dtype).lanes)) +
- 7) /
- 8);
- EmitAssert(acond, "ValueError", //
- "Misaligned tensor data on argument #", std::to_string(param_index),
- when_calling_imm_, sig_imm_, "`,\n expected data alignment=",
- std::to_string(arg.ty()->offset_factor * data_bytes), " bytes");
- }
- }
- }
- }
-
- ffi::reflection::AccessPath shape_path = base_path->Attr(ffi::String("shape"));
- ffi::reflection::AccessPath strides_path = base_path->Attr(ffi::String("strides"));
-
- if (arg.shape().size() < value.shape().size()) {
- TVM_FFI_ICHECK(fuzzy_match) << "TensorVar size mismatch at " << RenderAccessPath(base_path);
- size_t diff = value.shape().size() - arg.shape().size();
- for (size_t i = 0; i < diff; ++i) {
- TVM_FFI_ICHECK(IsOne(analyzer_->Simplify(value.shape()[i])))
- << "TensorVar shape mismatch at " << RenderAccessPath(base_path) << ": " << arg.shape()
- << " vs " << value.shape();
- }
- for (size_t i = 0; i < arg.shape().size(); ++i) {
- ffi::reflection::AccessPath shape_k_path = shape_path->ArrayItem(i);
- BindScalar(arg.shape()[i], value.shape()[i + diff], shape_k_path, false);
- }
- if (value.strides().size() != 0) {
- TVM_FFI_ICHECK_EQ(arg.strides().size(), arg.shape().size());
- TVM_FFI_ICHECK_EQ(value.strides().size(), value.shape().size());
- for (size_t i = 0; i < arg.strides().size(); ++i) {
- ffi::reflection::AccessPath strides_k_path = strides_path->ArrayItem(i);
- BindScalar(arg.strides()[i], value.strides()[i + diff], strides_k_path, false);
- }
- }
- } else {
- BindArray(arg.shape(), value.shape(), shape_path);
- BindArray(arg.strides(), value.strides(), strides_path);
- }
-}
-
// ============================================================
// Per-dtype type-check + value-load methods
// ============================================================
@@ -708,11 +615,11 @@ void TVMFFIABIBuilder::DecodeAllParams() {
}
}
}
- // Shape pass 2: Infer only variables not defined by a bare dimension.
+ // Integer-expression pass 2: Infer only variables not defined by a bare dimension.
for (auto& shape : pending_int_exprs_) {
shape.needs_check = !TryBindLinearIntExpr(shape.arg, shape.value, shape.path);
}
- // Shape pass 3: Check the remaining expressions after all bindings are known.
+ // Integer-expression pass 3: Check the remaining expressions after all bindings are known.
for (const auto& shape : pending_int_exprs_) {
if (shape.needs_check) {
BindScalar(shape.arg, shape.value, shape.path, true);
@@ -779,13 +686,16 @@ void TVMFFIABIBuilder::BindRegularStrides(const TensorVar& tensor, const Var& st
const Var& shape_ptr, const PrimExpr& v_strides_is_null,
const ffi::reflection::AccessPath& param_path) {
PrimExpr stride_from_shape = 1;
- for (int k = tensor.strides().size() - 1; k >= 0; k--) {
- PrimExpr explicit_stride = cast(LoadInt64ArrayElem(strides_ptr, k), tensor.shape()[k].ty());
+ auto strides = MatchShapeStridePattern(tensor.ty()).value().get<1>();
+ for (int k = strides.size() - 1; k >= 0; k--) {
+ PrimType stride_type = strides[k].ty();
+ PrimExpr explicit_stride = cast(LoadInt64ArrayElem(strides_ptr, k), stride_type);
ffi::reflection::AccessPath strides_k_path =
param_path->Attr(ffi::String("strides"))->ArrayItem(k);
- BindScalar(tensor.strides()[k],
- tvm::if_then_else(v_strides_is_null, stride_from_shape, explicit_stride),
- strides_k_path, true);
+ BindScalar(
+ strides[k],
+ tvm::if_then_else(v_strides_is_null, cast(stride_from_shape, stride_type), explicit_stride),
+ strides_k_path, true);
stride_from_shape *= cast(LoadInt64ArrayElem(shape_ptr, k), tensor.shape()[k].ty());
}
}
@@ -798,6 +708,10 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const TensorVar& tensor, const PrimEx
const PrimExpr& device_id, const Var& handle,
const std::string& arg_name,
ffi::reflection::AccessPath base_path) {
+ auto pattern = MatchShapeStridePattern(tensor.ty());
+ TVM_FFI_CHECK(pattern.has_value(), ValueError)
+ << "Tensor parameters require a shape-aligned strided memory layout";
+ PrimExpr elem_offset = pattern.value().get<2>();
const PrimType tvm_ndim_type = PrimType::Int(32);
std::string buf_name = tensor->name;
@@ -864,7 +778,7 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const TensorVar& tensor, const PrimEx
Var strides_ptr = DLTensorGetFieldPtr(handle, tirx::kDLTensorStrides, arg_name + "_strides");
PrimExpr v_strides_is_null =
Call(PrimType::Bool(), tirx::isnullptr_op(), {strides_ptr}).as_or_throw<PrimExpr>();
- if (tensor.strides().size() == 0) {
+ if (!tensor.ty()->layout.has_value()) {
BindCompactStrides(tensor, strides_ptr, v_strides_is_null, param_path);
} else {
BindRegularStrides(tensor, strides_ptr, shape_ptr, v_strides_is_null, param_path);
@@ -873,7 +787,7 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const TensorVar& tensor, const PrimEx
// ── Section: byte_offset ─────────────────────────────────────
int data_bytes = static_cast<int>(tensor.dtype().StorageBytes());
ffi::reflection::AccessPath byte_offset_path = param_path->Attr(ffi::String("byte_offset"));
- if (const auto* const_offset = tensor.elem_offset().as<IntImmNode>()) {
+ if (const auto* const_offset = elem_offset.as<IntImmNode>()) {
BindScalar(IntImm(PrimType::UInt(64), const_offset->value * data_bytes),
Call(PrimType::UInt(64), tirx::abi_field_get_op(),
{handle, IntImm::Int32(0), IntImm::Int32(tirx::kDLTensorByteOffset)})
@@ -887,29 +801,14 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const TensorVar& tensor, const PrimEx
EmitAssert(truncmod(byte_offset, MakeConst(PrimType::UInt(64), data_bytes)) ==
MakeConst(PrimType::UInt(64), 0),
"ValueError", "Tensor byte offset must be divisible by its element size");
- if (BindScalar(tensor.elem_offset(),
- cast((Call(PrimType::UInt(64), tirx::abi_field_get_op(),
- {handle, IntImm::Int32(0), IntImm::Int32(tirx::kDLTensorByteOffset)})
- .as_or_throw<PrimExpr>() /
- MakeConst(PrimType::UInt(64), data_bytes)),
- tensor.elem_offset().ty()),
- byte_offset_path, true)) {
- if (tensor.ty()->offset_factor > 1) {
- PrimExpr offset = tensor.elem_offset();
- PrimExpr factor = IntImm(offset.ty(), tensor.ty()->offset_factor);
- PrimExpr zero = IntImm(offset.ty(), 0);
- PrimExpr acond = analyzer_->Simplify(truncmod(offset, factor) == zero);
- if (IsZero(acond)) {
- TVM_FFI_THROW(InternalError)
- << "Bind have an unmet assertion at " << RenderAccessPath(byte_offset_path);
- }
- if (!IsOne(acond)) {
- EmitAssert(acond, "ValueError", //
- "Misaligned Tensor data on argument #", std::to_string(param_index),
- when_calling_imm_, sig_imm_, "`,\n expected data alignment=",
- std::to_string(tensor.ty()->offset_factor * data_bytes), " bytes");
- }
- }
+ PrimExpr runtime_origin =
+ cast(byte_offset / MakeConst(PrimType::UInt(64), data_bytes), elem_offset.ty());
+ if (elem_offset.as<VarNode>()) {
+ BindScalar(elem_offset, runtime_origin, byte_offset_path, true);
+ } else {
+ // Linear offsets bind in element units, after byte-offset conversion.
+ pending_int_exprs_.push_back(
+ {elem_offset, runtime_origin, param_path->Attr(ffi::String("elem_offset"))});
}
}
diff --git a/tirx/src/tirx/transform/tvm_ffi_binder.h b/tirx/src/tirx/transform/tvm_ffi_binder.h
index 832a2c3450..96846ec2a6 100644
--- a/tirx/src/tirx/transform/tvm_ffi_binder.h
+++ b/tirx/src/tirx/transform/tvm_ffi_binder.h
@@ -258,30 +258,6 @@ class TVMFFIABIBuilder {
bool BindPointer(const Var& arg, const Expr& value, const ffi::reflection::AccessPath& path,
bool with_lets);
- /*!
- * \brief Array bind: binds element-wise with ffi::reflection::AccessPath[k] for each element.
- *
- * \param arg The expected array of expressions.
- * \param value The actual array of expressions to bind against.
- * \param base_path Base ffi::reflection::AccessPath; each element appends ArrayItem(k).
- */
- void BindArray(const ffi::Array<PrimExpr>& arg, const ffi::Array<PrimExpr>& value,
- const ffi::reflection::AccessPath& base_path);
-
- /*!
- * \brief TensorVar-to-tensor bind with ffi::reflection::AccessPath.
- *
- * Binds data, elem_offset, shape, and strides of \p arg against \p value,
- * emitting assertions for any mismatches.
- *
- * \param arg The expected tensor definition.
- * \param value The actual tensor to bind against.
- * \param base_path Base ffi::reflection::AccessPath for the tensor parameter.
- * \param fuzzy_match If true, allow value to have more dimensions than arg.
- */
- void BindTensor(const TensorVar& arg, const TensorVar& value,
- ffi::reflection::AccessPath base_path, bool fuzzy_match);
-
/*!
* \brief DLTensor bind: ndim/dtype/shape/strides/data/device assertions.
*
diff --git a/tirx/src/tirx/transform/vectorize_loop.cc b/tirx/src/tirx/transform/vectorize_loop.cc
index 4fb1012155..34991aff73 100644
--- a/tirx/src/tirx/transform/vectorize_loop.cc
+++ b/tirx/src/tirx/transform/vectorize_loop.cc
@@ -395,21 +395,10 @@ class VecAllocAccess : public StmtExprMutator {
// flattening, the strides are not technically necessary, but
// are updated for consistency.
- // Update strides if defined.
- ffi::Array<PrimExpr> strides = node->dest.template as_or_throw<TensorVar>().strides();
- for (size_t i = 0; i < strides.size(); i++) {
- PrimExpr stride = strides[i];
- if (i != strides.size() - 1) {
- stride *= var_lanes_;
- }
- strides.Set(i, analyzer_->Simplify(stride));
- }
-
- // Copy everything into the new tensor.
auto type = ffi::make_object<TensorTypeNode>(
*node->dest.template as_or_throw<TensorVar>().ty().get());
+ WidenLayout(type.get(), shape);
type->shape = shape;
- type->strides = strides;
buf =
TensorVar(node->dest.template as_or_throw<TensorVar>()->name, TensorType(std::move(type)),
node->dest.template as_or_throw<TensorVar>()->loc);
@@ -438,15 +427,9 @@ class VecAllocAccess : public StmtExprMutator {
} else {
ffi::Array<PrimExpr> shape = tensor.shape();
shape.Set(shape.size() - 1, analyzer_->Simplify(shape.back() * var_lanes_));
- ffi::Array<PrimExpr> strides = tensor.strides();
- for (size_t i = 0; i < strides.size(); ++i) {
- PrimExpr stride = strides[i];
- if (i + 1 != strides.size()) stride *= var_lanes_;
- strides.Set(i, analyzer_->Simplify(stride));
- }
auto type = ffi::make_object<TensorTypeNode>(*tensor.ty().get());
+ WidenLayout(type.get(), shape);
type->shape = shape;
- type->strides = strides;
buf = TensorVar(tensor->name, TensorType(std::move(type)), tensor->loc);
VarRemapSet(tensor, buf);
}
@@ -456,6 +439,23 @@ class VecAllocAccess : public StmtExprMutator {
return TensorLoad(buf, indices, node->loc);
}
+ void WidenLayout(TensorTypeNode* type, const ffi::Array<PrimExpr>& shape) {
+ if (!type->layout.has_value()) return;
+ auto tile = type->layout.value().as_or_throw<TileLayout>();
+ TVM_FFI_CHECK(
+ MatchShapeStridePattern(TensorType(ffi::make_object<TensorTypeNode>(*type))).has_value(),
+ ValueError)
+ << "Vectorized allocation requires a regular memory layout";
+ ffi::Array<Iter> shard;
+ for (size_t i = 0; i < tile->shard.size(); ++i) {
+ PrimExpr stride = tile->shard[i]->stride;
+ if (i + 1 != tile->shard.size()) stride *= var_lanes_;
+ shard.push_back(Iter(shape[i], analyzer_->Simplify(stride), Axis::Get("m")));
+ }
+ tile.CopyOnWrite()->shard = std::move(shard);
+ type->layout = tile;
+ }
+
// tensor var
const VarNode* buf_;
// variable to be replaced
diff --git a/tirx/tests/cpp/tirx/tensor_var_test.cc b/tirx/tests/cpp/tirx/tensor_var_test.cc
index 83b3ad5072..add39a59d8 100644
--- a/tirx/tests/cpp/tirx/tensor_var_test.cc
+++ b/tirx/tests/cpp/tirx/tensor_var_test.cc
@@ -45,8 +45,8 @@ TEST(TensorVar, CheckedViewPreservesOrdinaryVariableIdentity) {
EXPECT_TRUE(tensor.ty().same_as(type));
EXPECT_TRUE(tensor.shape().same_as(type->shape));
EXPECT_TRUE(tensor.dtype().same_as(type->dtype));
- EXPECT_TRUE(tensor.strides().empty());
- EXPECT_TRUE(prim::IsZero(tensor.elem_offset()));
+ EXPECT_EQ(MatchShapeStridePattern(tensor.ty()).value().get<1>().size(), 1);
+ EXPECT_TRUE(prim::IsZero(MatchShapeStridePattern(tensor.ty()).value().get<2>()));
EXPECT_EQ(tensor.scope(), type->scope);
EXPECT_TRUE(tensor == TensorVar(var));
Var scalar("scalar", PrimType::Int(32));
@@ -69,7 +69,6 @@ TEST(TensorVar, RebuiltTypeRetainsMetadataWithFreshIdentity) {
EXPECT_EQ(rebuilt->name, original->name);
EXPECT_TRUE(rebuilt->loc.same_as(original->loc));
EXPECT_EQ(rebuilt.ty()->data_alignment, runtime::kAllocAlignment);
- EXPECT_EQ(rebuilt.ty()->offset_factor, 1);
EXPECT_FALSE(rebuilt.ty()->layout.has_value());
}
@@ -81,7 +80,7 @@ TEST(TensorVar, PreferredIndexTypeUsesShapeOrConfiguredDefault) {
EXPECT_EQ(DefaultIndexType().bits(), 32);
#endif
EXPECT_TRUE(scalar->DefaultIndexType().same_as(DefaultIndexType()));
- EXPECT_TRUE(scalar->elem_offset.ty().same_as(DefaultIndexType()));
+ EXPECT_TRUE(MatchShapeStridePattern(scalar).value().get<2>().ty().same_as(DefaultIndexType()));
for (int bits : {32, 64}) {
PrimType index_type = PrimType::Int(bits);
TensorType tensor({IntImm(index_type, 16)});
diff --git a/tirx/tests/python/backend/vulkan/codegen/test_basic.py b/tirx/tests/python/backend/vulkan/codegen/test_basic.py
index be20d5e9ff..a8f01b124e 100644
--- a/tirx/tests/python/backend/vulkan/codegen/test_basic.py
+++ b/tirx/tests/python/backend/vulkan/codegen/test_basic.py
@@ -346,12 +346,15 @@ def test_vectorized_index_ramp():
n = 4
ramp_index = tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4)
+ module_a_offset = T.dynamic("module_a_offset", "int32")
+ module_b_offset = T.dynamic("module_b_offset", "int32")
+
@I.ir_module
class Module:
@T.function
def main(
- A: T.Tensor((n,), "int32", offset_factor=1, elem_offset=T.int32()),
- B: T.Tensor((n,), "int32", offset_factor=1, elem_offset=T.int32()),
+ A: T.Tensor((n,), "int32", layout=T.TileLayout(T.S[(n,)] + (module_a_offset))),
+ B: T.Tensor((n,), "int32", layout=T.TileLayout(T.S[(n,)] + (module_b_offset))),
):
T.func_attr({"tirx.noalias": True})
@@ -385,12 +388,15 @@ def test_vectorized_index_broadcast():
broadcast_index = tvm_tirx.ir.prim.vector_expr.Broadcast(0, 4)
ramp_index = tvm_tirx.ir.prim.vector_expr.Ramp(0, 1, 4)
+ module_a_offset_1 = T.dynamic("module_a_offset_1", "int32")
+ module_b_offset_1 = T.dynamic("module_b_offset_1", "int32")
+
@I.ir_module
class Module:
@T.function
def main(
- A: T.Tensor((n,), "int32", offset_factor=1, elem_offset=T.int32()),
- B: T.Tensor((n,), "int32", offset_factor=1, elem_offset=T.int32()),
+ A: T.Tensor((n,), "int32", layout=T.TileLayout(T.S[(n,)] + (module_a_offset_1))),
+ B: T.Tensor((n,), "int32", layout=T.TileLayout(T.S[(n,)] + (module_b_offset_1))),
):
T.func_attr({"tirx.noalias": True})
diff --git a/tirx/tests/python/codegen/llvm/test_basic.py b/tirx/tests/python/codegen/llvm/test_basic.py
index 7802879ff1..af1d5ef70b 100644
--- a/tirx/tests/python/codegen/llvm/test_basic.py
+++ b/tirx/tests/python/codegen/llvm/test_basic.py
@@ -257,6 +257,37 @@ def test_llvm_add_linear_shape():
np.testing.assert_equal(c.numpy(), a.numpy() + b.numpy())
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_llvm_strides_and_linear_layout_offset():
+ s = T.dynamic("s", "int64")
+ m = T.dynamic("m", "int64")
+
+ @T.function
+ def copy(
+ A: T.Tensor((3, 4), "float32", layout=T.TileLayout(T.S[(3, 4) : (s, 1)] + 4 * m), align=4),
+ B: T.Tensor((3, 4), "float32"),
+ ):
+ for i, j in T.grid(3, 4):
+ B[i, j] = A[i, j] + T.Cast("float32", m + s)
+
+ lib = tvm_tirx.tirx.compile(copy, target="llvm")
+ data = np.arange(64, dtype="float32")
+ allocation = tvm_tirx.runtime.tensor(data)
+ result = tvm_tirx.runtime.tensor(np.zeros((3, 4), dtype="float32"))
+ for offset in (4, 8):
+ # CreateView keeps the allocation pointer and sets DLTensor.byte_offset.
+ view = allocation._create_view((3, 4), relative_byte_offset=offset * 4)
+ lib(view, result)
+ expected = data[offset : offset + 12].reshape(3, 4) + offset // 4 + 4
+ np.testing.assert_equal(result.numpy(), expected)
+ strided = data.reshape(8, 8)[:3, :4]
+ lib(strided, result)
+ np.testing.assert_equal(result.numpy(), strided + 8)
+ invalid = allocation._create_view((3, 4), relative_byte_offset=4)
+ with pytest.raises(ValueError, match="divisible"):
+ lib(invalid, result)
+
+
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_llvm_madd_pipeline():
def check_llvm(nn, base, stride):
diff --git a/tirx/tests/python/codegen/llvm/test_riscv.py b/tirx/tests/python/codegen/llvm/test_riscv.py
index 4182e84925..16450fb009 100644
--- a/tirx/tests/python/codegen/llvm/test_riscv.py
+++ b/tirx/tests/python/codegen/llvm/test_riscv.py
@@ -108,8 +108,12 @@ def test_rvv_vscale_llvm_dbginfo(target):
pytest.skip(f"{target} not enabled")
# fmt: off
+ rvv_with_vscale_a_offset = T.dynamic("rvv_with_vscale_a_offset", "int32")
+ rvv_with_vscale_b_offset = T.dynamic("rvv_with_vscale_b_offset", "int32")
+ rvv_with_vscale_c_offset = T.dynamic("rvv_with_vscale_c_offset", "int32")
+
@T.function
- def rvv_with_vscale(A: T.Tensor((8,), dtype='float32', align=4, offset_factor=1, elem_offset=T.int32()), B: T.Tensor((4, 8), dtype='float32', align=4, offset_factor=1, elem_offset=T.int32(), strides=[8, 1]), C: T.Tensor((4,), dtype='float32', align=4, offset_factor=1, elem_offset=T.int32())):
+ def rvv_with_vscale(A: T.Tensor((8,), dtype='float32', align=4, layout=T.TileLayout(T.S[(8,)] + (rvv_with_vscale_a_offset))), B: T.Tensor((4, 8), dtype='float32', align=4, layout=T.TileLayout(T.S[(4, 8) : [8, 1]] + (rvv_with_vscale_b_offset))), C: T.Tensor((4,), dtype='float32', align=4, layout=T.TileLayout(T.S[(4,)] + (rvv_with_vscale_c_offset)))):
zero = T.call_llvm_intrin( 'llvm.riscv.vfmv.v.f', T.Broadcast(T.float32(0.0), T.vscale() * 2), C[0], T.uint64(1), ty='float32xvscalex2')
vec_A = T.call_llvm_intrin( 'llvm.riscv.vle', T.Broadcast(T.float32(0.0), T.vscale() * 4), A.ptr_to([0] * len(A.shape)), T.int64(8), ty='float32xvscalex4')
diff --git a/tirx/tests/python/tirx/analysis/test_undefined_vars.py b/tirx/tests/python/tirx/analysis/test_undefined_vars.py
index 4f3b68acd4..30aea75807 100644
--- a/tirx/tests/python/tirx/analysis/test_undefined_vars.py
+++ b/tirx/tests/python/tirx/analysis/test_undefined_vars.py
@@ -62,7 +62,12 @@ 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.Var("buf", tirx.TensorType((n,), "float32", data=data_ptr, elem_offset=elem_off))
+ buf = tirx.Var(
+ "buf",
+ tirx.TensorType(
+ (n,), "float32", data=data_ptr, layout=tirx.TileLayout(tirx.layout.S[(n,)] + (elem_off))
+ ),
+ )
body = tvm_tirx.ir.Evaluate(tvm_tirx.ir.TensorLoad(buf, [0]))
decl = tvm_tirx.ir.Bind(
diff --git a/tirx/tests/python/tirx/base/test_intrin.py b/tirx/tests/python/tirx/base/test_intrin.py
index edc3bf1f8d..faef661b39 100644
--- a/tirx/tests/python/tirx/base/test_intrin.py
+++ b/tirx/tests/python/tirx/base/test_intrin.py
@@ -32,7 +32,7 @@ from tvm_tirx.support import clang, utils
def _unary_kernel(op, dtype="float32", out_dtype=None, gpu=False):
out_dtype = out_dtype or dtype
- n = T.int32()
+ n = T.dynamic("n", "int32")
@T.function
def kernel(A: T.Tensor((n,), dtype), B: T.Tensor((n,), out_dtype)):
@@ -49,7 +49,7 @@ def _unary_kernel(op, dtype="float32", out_dtype=None, gpu=False):
def _binary_kernel(op, rhs_dtype="float32"):
- n = T.int32()
+ n = T.dynamic("n", "int32")
@T.function
def kernel(
@@ -329,10 +329,10 @@ stride_3 = T.dynamic("stride_3", "int32")
class Module:
@T.function
def test_tir_fma(
- A_1: T.Tensor([n], "float32", strides=[stride], elem_offset=0, align=64, offset_factor=1),
- B_1: T.Tensor([n], "float32", strides=[stride_1], elem_offset=0, align=64, offset_factor=1),
- C_1: T.Tensor([n], "float32", strides=[stride_2], elem_offset=0, align=64, offset_factor=1),
- d_1: T.Tensor([n], "float32", strides=[stride_3], elem_offset=0, align=64, offset_factor=1),
+ A_1: T.Tensor([n], "float32", layout=T.TileLayout(T.S[[n] : [stride]]), align=64),
+ B_1: T.Tensor([n], "float32", layout=T.TileLayout(T.S[[n] : [stride_1]]), align=64),
+ C_1: T.Tensor([n], "float32", layout=T.TileLayout(T.S[[n] : [stride_2]]), align=64),
+ d_1: T.Tensor([n], "float32", layout=T.TileLayout(T.S[[n] : [stride_3]]), align=64),
) -> None:
# function attr dict
T.func_attr({"global_symbol": "test_fma", "tirx.noalias": True})
diff --git a/tirx/tests/python/tirx/base/test_nodes.py b/tirx/tests/python/tirx/base/test_nodes.py
index fa7f00b042..44b683f586 100644
--- a/tirx/tests/python/tirx/base/test_nodes.py
+++ b/tirx/tests/python/tirx/base/test_nodes.py
@@ -357,7 +357,10 @@ def test_tensor_load_store():
assert isinstance(s, tvm_tirx.ir.TensorStore)
tensor = tvm_tirx.tirx.Var(
- "tensor", tvm_tirx.tirx.TensorType((4, 8), "float32", elem_offset=100)
+ "tensor",
+ tvm_tirx.tirx.TensorType(
+ (4, 8), "float32", layout=tvm_tirx.tirx.TileLayout(tvm_tirx.tirx.layout.S[(4, 8)] + 100)
+ ),
)
tvm_tirx.ir.assert_structural_equal(
tensor[2, 3].indices,
diff --git a/tirx/tests/python/tirx/base/test_op_types.py b/tirx/tests/python/tirx/base/test_op_types.py
index 8c0caf23a5..30873293e2 100644
--- a/tirx/tests/python/tirx/base/test_op_types.py
+++ b/tirx/tests/python/tirx/base/test_op_types.py
@@ -261,20 +261,31 @@ 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.Var("tensor", tirx.TensorType([16, 8], "int32", scope="warp", offset_factor=1))
+ tensor_w = tirx.Var("tensor", tirx.TensorType([16, 8], "int32", scope="warp"))
tensor = tirx.Var(
"tensor",
- tirx.TensorType([16, 16], "int32", scope="global", offset_factor=1, strides=[x, y]),
+ tirx.TensorType(
+ [16, 16],
+ "int32",
+ scope="global",
+ layout=tirx.layout.TileLayout(tirx.layout.S[[16, 16] : [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"
+ 16,
+ 16,
+ tensor.ptr_to([0, 0]),
+ tensor_w.data,
+ 0,
+ x,
+ ty="int32",
)
assert expr.op.name == "tirx.cuda.mma_store"
def test_tir_op_mma_fill():
- 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")
+ tensor_w = tirx.Var("tensor", tirx.TensorType([16, 8], "int32", scope="warp"))
+ expr = _cuda_op.mma_fill(8, tensor_w.data, 0, ty="int32")
assert expr.op.name == "tirx.cuda.mma_fill"
@@ -298,14 +309,14 @@ def test_op_ptx_cp_async():
def test_tir_op_vector_low():
- tensor = tirx.Var("tensor", tirx.TensorType((4, 4), "int8", offset_factor=1))
+ tensor = tirx.Var("tensor", tirx.TensorType((4, 4), "int8"))
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.Var("tensor", tirx.TensorType((4, 4), "int8", offset_factor=1))
+ tensor = tirx.Var("tensor", tirx.TensorType((4, 4), "int8"))
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"
@@ -320,7 +331,7 @@ def test_tir_op_dp4a():
def test_tir_op_vector_combine():
- tensor = tirx.Var("tensor", tirx.TensorType((4, 4), "int8", offset_factor=1))
+ tensor = tirx.Var("tensor", tirx.TensorType((4, 4), "int8"))
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"
diff --git a/tirx/tests/python/tirx/base/test_specialize.py b/tirx/tests/python/tirx/base/test_specialize.py
index d15506b9eb..8bdc8685b2 100644
--- a/tirx/tests/python/tirx/base/test_specialize.py
+++ b/tirx/tests/python/tirx/base/test_specialize.py
@@ -133,15 +133,23 @@ def element_wise_128_n(A: T.Tensor((128, n), "float32"), C: T.Tensor((128, n), "
C[i, j] = B[i, j] + 1.0
-mem_copy_m = T.int32()
+mem_copy_m = T.dynamic("mem_copy_m", "int32")
-mem_copy_n = T.int32()
+mem_copy_n = T.dynamic("mem_copy_n", "int32")
@T.function
def mem_copy(
- A: T.Tensor((mem_copy_m, mem_copy_n), "float32", strides=[p, 1], elem_offset=q), # noqa: F821
- B: T.Tensor((mem_copy_m, mem_copy_n), "float32", strides=[p, 1], elem_offset=q), # noqa: F821
+ A: T.Tensor(
+ (mem_copy_m, mem_copy_n),
+ "float32",
+ layout=T.TileLayout(T.S[(mem_copy_m, mem_copy_n) : [p, 1]] + (q)), # noqa: F821
+ ),
+ B: T.Tensor(
+ (mem_copy_m, mem_copy_n),
+ "float32",
+ layout=T.TileLayout(T.S[(mem_copy_m, mem_copy_n) : [p, 1]] + (q)), # noqa: F821
+ ),
m: mem_copy_m,
n: mem_copy_n,
p: T.int32,
@@ -153,16 +161,16 @@ def mem_copy(
@T.function
def mem_copy_16_16_8_4(
- A: T.Tensor((16, 16), "float32", strides=[8, 1], elem_offset=T.int32(4)),
- B: T.Tensor((16, 16), "float32", strides=[8, 1], elem_offset=T.int32(4)),
+ A: T.Tensor((16, 16), "float32", layout=T.TileLayout(T.S[(16, 16) : [8, 1]] + (T.int32(4)))),
+ B: T.Tensor((16, 16), "float32", layout=T.TileLayout(T.S[(16, 16) : [8, 1]] + (T.int32(4)))),
) -> None:
for i, j in T.grid(16, 16):
B[i, j] = A[i, j]
-mem_copy_m_n_p_n_m = T.int32()
+mem_copy_m_n_p_n_m = T.dynamic("mem_copy_m_n_p_n_m", "int32")
-mem_copy_m_n_p_n_n = T.int32()
+mem_copy_m_n_p_n_n = T.dynamic("mem_copy_m_n_p_n_n", "int32")
@T.function
@@ -170,14 +178,16 @@ def mem_copy_m_n_p_n(
A: T.Tensor(
(mem_copy_m_n_p_n_m, mem_copy_m_n_p_n_n),
"float32",
- strides=[p, 1], # noqa: F821
- elem_offset=mem_copy_m_n_p_n_n,
+ layout=T.TileLayout(
+ T.S[(mem_copy_m_n_p_n_m, mem_copy_m_n_p_n_n) : [p, 1]] + (mem_copy_m_n_p_n_n) # noqa: F821
+ ),
),
B: T.Tensor(
(mem_copy_m_n_p_n_m, mem_copy_m_n_p_n_n),
"float32",
- strides=[p, 1], # noqa: F821
- elem_offset=mem_copy_m_n_p_n_n,
+ layout=T.TileLayout(
+ T.S[(mem_copy_m_n_p_n_m, mem_copy_m_n_p_n_n) : [p, 1]] + (mem_copy_m_n_p_n_n) # noqa: F821
+ ),
),
m: mem_copy_m_n_p_n_m,
n: mem_copy_m_n_p_n_n,
@@ -230,7 +240,11 @@ def test_specialize_mem_copy():
a: tvm_tirx.tirx.Var(
"tensor",
tvm_tirx.tirx.TensorType(
- (16, 16), "float32", strides=[8, 1], elem_offset=T.int32(4)
+ (16, 16),
+ "float32",
+ layout=tvm_tirx.tirx.layout.TileLayout(
+ tvm_tirx.tirx.layout.S[(16, 16) : [8, 1]] + (T.int32(4))
+ ),
),
)
}
@@ -290,7 +304,7 @@ def test_specialize_decl_tensor():
def test_specialize_preserves_decl_tensor_alias():
- before_n = T.int32()
+ before_n = T.dynamic("before_n", "int32")
@T.function(private=True)
def before(A: T.Tensor((before_n,), "int32"), n: before_n):
diff --git a/tirx/tests/python/tirx/base/test_tensor_type.py b/tirx/tests/python/tirx/base/test_tensor_type.py
index 41b999474d..75f2562dec 100644
--- a/tirx/tests/python/tirx/base/test_tensor_type.py
+++ b/tirx/tests/python/tirx/base/test_tensor_type.py
@@ -37,18 +37,18 @@ def test_tensor_pointer_type_derived_from_dtype_and_scope():
@pytest.mark.parametrize("scope", ["global", "shared", tvm_tirx.runtime.MemoryScope("local")])
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)
+ direct = tvm_tirx.tirx.TensorType(
+ (4, 8), "float32", scope, 64, T.TileLayout(T.S[(4, 8) : (16, 1)] + 4)
+ )
+ scripted = T.Tensor((4, 8), "float32", scope, 64, T.TileLayout(T.S[(4, 8) : (16, 1)] + 4))
declared = tvm_tirx.tirx.Var(
"tensor",
tvm_tirx.tirx.TensorType(
(4, 8),
"float32",
scope=scope,
- strides=(16, 1),
- elem_offset=4,
+ layout=tvm_tirx.tirx.layout.TileLayout(tvm_tirx.tirx.layout.S[(4, 8) : (16, 1)] + (4)),
data_alignment=64,
- offset_factor=2,
),
)
tvm_tirx.ir.assert_structural_equal(direct, scripted)
@@ -86,7 +86,14 @@ def test_tensor_ptr_to():
m = tvm_tirx.tirx.Var("m", "int32")
n = tvm_tirx.tirx.Var("n", "int32")
tensor = tvm_tirx.tirx.Var(
- "tensor", tvm_tirx.tirx.TensorType((m, n), "float32", strides=[n + 1, 1], elem_offset=7)
+ "tensor",
+ tvm_tirx.tirx.TensorType(
+ (m, n),
+ "float32",
+ layout=tvm_tirx.tirx.layout.TileLayout(
+ tvm_tirx.tirx.layout.S[(m, n) : [n + 1, 1]] + (7)
+ ),
+ ),
)
pointer = tensor.ptr_to([2, 3])
assert pointer.op.name == "tirx.address_of"
@@ -131,10 +138,8 @@ def test_tensor_type_and_projection(dtype, scope):
def test_tensor_type_default_offset():
- tensor = tvm_tirx.tirx.Var(
- "tensor", tvm_tirx.tirx.TensorType((8,), "float32", elem_offset=None)
- )
- assert int(tensor.ty.elem_offset) == 0
+ tensor = tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((8,), "float32"))
+ assert tensor.ty.layout is None
@pytest.mark.parametrize(
@@ -259,15 +264,17 @@ def test_tensor_scope_receiver_validation():
@pytest.mark.parametrize("alignment", [-1, 0, 64])
def test_tensor_type_construction_offsets_and_alignment(alignment):
ty = tvm_tirx.tirx.TensorType(
- (), "float32", data_alignment=alignment, byte_offset=16, strides=()
+ (), "float32", data_alignment=alignment, layout=T.TileLayout(T.S[()] + 4)
)
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 int(tensor.layout.apply(*([0] * len(tensor.shape)), shape=tensor.shape).get("m", 0)) == 4
+ assert int(tensor.layout.apply(*([0] * len(tensor.shape)), shape=tensor.shape).get("m", 0)) == 4
+ assert (
+ int(tensor.layout.apply(*([0] * len(tensor.shape)), shape=tensor.shape).get("m", 0) * 4)
+ == 16
+ )
assert tensor.data_alignment == 64
- assert tensor.offset_factor == 1
if __name__ == "__main__":
diff --git a/tirx/tests/python/tirx/codegen/test_blackwell.py b/tirx/tests/python/tirx/codegen/test_blackwell.py
index dbf8cb69af..46f2f3d3f4 100644
--- a/tirx/tests/python/tirx/codegen/test_blackwell.py
+++ b/tirx/tests/python/tirx/codegen/test_blackwell.py
@@ -569,19 +569,20 @@ def test_tcgen05_mma_ss_no_tma(swizzle):
lane_id = T.cuda.lane_id()
tx = T.cuda.thread_idx('x')
dyn = T.alloc_tensor(ty_args=[T.Tensor((dyn_smem_bytes,), "uint8", scope="shared")])
- tmem_addr = T.decl_scalar("uint32", dyn.data, scope="shared", elem_offset=0)
+ tmem_addr = T.decl_scalar("uint32", dyn.data, scope="shared")
A_smem = T.decl_tensor(
- dyn.data,
- ty_args=[T.Tensor((M, K), a_type, elem_offset=256, layout=A_layout, data=dyn.data)],
+ T.ptr_byte_offset(dyn.data, 512, ty=T.Ptr(a_type, "shared")),
+ ty_args=[T.Tensor((M, K), a_type, layout=A_layout, data=dyn.data)],
)
B_smem = T.decl_tensor(
- dyn.data,
+ T.ptr_byte_offset(dyn.data, (256 + M * K) * 2, ty=T.Ptr(b_type, "shared")),
ty_args=[
- T.Tensor((N, K), b_type, elem_offset=256 + M * K, layout=B_layout, data=dyn.data)
+ T.Tensor((N, K), b_type, layout=B_layout, data=dyn.data)
],
)
bar = T.decl_tensor(
- dyn.data, ty_args=[T.Tensor((1,), "uint64", scope="shared", elem_offset=8)]
+ dyn.data, ty_args=[T.Tensor((1,), "uint64", scope="shared",
+ layout=T.TileLayout(T.S[(1,)] + (8)))]
)
reg = T.alloc_tensor(ty_args=[T.Tensor((N,), d_type, scope="local")])
diff --git a/tirx/tests/python/tirx/script/test_tensors.py b/tirx/tests/python/tirx/script/test_tensors.py
index 850028d7f5..1816895209 100644
--- a/tirx/tests/python/tirx/script/test_tensors.py
+++ b/tirx/tests/python/tirx/script/test_tensors.py
@@ -245,7 +245,10 @@ def test_alloc_apis():
self.inner_pool = inner_pool
self.Tb = T.shared_scalar("float16")
self.idx = T.local_scalar("int32")
- self.inner_pool2 = T.decl_scalar("float16", self.inner_pool.data, "shared.dyn", 5)
+ self.inner_pool2 = T.decl_scalar(
+ "float16", self.inner_pool.data, "shared.dyn",
+ layout=T.TileLayout(T.S[(1,)] + 5),
+ )
@T.inline
def init(self):
@@ -271,7 +274,7 @@ def test_alloc_apis():
D: T.float16
pool = T.alloc_tensor(ty_args=[T.Tensor([10], "uint8", scope="shared.dyn")])
# scalar tensor (decl)
- E = T.decl_scalar("float16", pool.data, "shared.dyn", 0)
+ E = T.decl_scalar("float16", pool.data, "shared.dyn")
# normal 1-dim tensor with shape (1,)
F = T.alloc_local((1,), "float16")
Ta: T.float16
@@ -754,18 +757,20 @@ def test_tensor_physical_view_printer_preserves_inherited_metadata():
T.Tensor(
[32, 2],
"float32",
- elem_offset=8,
+
scope="local",
- layout=T.TileLayout(T.S[(32, 2) : (1 @ laneid, 2)]),
+ layout=T.TileLayout(T.S[(32, 2) : (1 @ laneid, 2)] + (8)),
)
]
)
B_align = T.decl_tensor(
- A.data, ty_args=[T.Tensor((2,), "float32", elem_offset=8, scope="local", align=128)]
+ A.data, ty_args=[T.Tensor((2,), "float32",
+ layout=T.TileLayout(T.S[(2,)] + (8)), scope="local", align=128)]
)
B_factor = T.decl_tensor(
A.data,
- ty_args=[T.Tensor((2,), "float32", elem_offset=8, scope="local", offset_factor=8)],
+ ty_args=[T.Tensor((2,), "float32",
+ layout=T.TileLayout(T.S[(2,)] + (8)), scope="local")],
)
B_align[0] = B_factor[0]
# fmt: on
@@ -774,7 +779,7 @@ def test_tensor_physical_view_printer_preserves_inherited_metadata():
align_line = next(line for line in code.splitlines() if "B_align =" in line)
factor_line = next(line for line in code.splitlines() if "B_factor =" in line)
assert "T.decl_tensor" in align_line and 'data_alignment=128' in align_line
- assert "T.decl_tensor" in factor_line and "offset_factor=8" in factor_line
+ assert "T.decl_tensor" in factor_line and "offset_factor" not in factor_line
parsed = from_source(code)
assert_structural_equal(func, parsed)
assert parsed.script() == code
@@ -978,9 +983,14 @@ def test_tensor_sub_multi_iter_dim_ir():
a_buf, b_buf = _collect_tensors(func)
# 5 -> (5 // 4, 5 % 4) = (1, 1) -> 1 * 1024 + 1 * 64
- assert int(tvm_tirx.sym.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 1088
+ assert int(tvm_tirx.sym.Analyzer().simplify(
+ b_buf.layout.apply(0).get("m", 0) - a_buf.layout.apply(0).get("m", 0)
+ )) == 1088
assert [int(s) for s in b_buf.shape] == [16]
- assert_structural_equal(b_buf.layout, tvm_tirx.tirx.layout.TileLayout(T.S[(16,) : (1,)]))
+ assert_structural_equal(
+ b_buf.layout,
+ tvm_tirx.tirx.layout.TileLayout(T.S[(16,) : (1,)] + 1088),
+ )
code = func.script()
assert from_source(code).script() == code
@@ -1028,13 +1038,25 @@ def test_tensor_sub_ir():
a_buf, _, b_buf, _, c_buf = _collect_tensors(func)
# sub[1, 2:6]: drop dim 0 at 1 (1 * 256) then narrow dim 1 to [2, 6) (2 * 16)
assert [int(s) for s in b_buf.shape] == [4, 16]
- assert int(tvm_tirx.sym.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 288
- assert_structural_equal(b_buf.layout, tvm_tirx.tirx.layout.TileLayout(T.S[(4, 16) : (16, 1)]))
+ assert int(tvm_tirx.sym.Analyzer().simplify(
+ b_buf.layout.apply(0).get("m", 0) - a_buf.layout.apply(0).get("m", 0)
+ )) == 288
+ assert_structural_equal(
+ b_buf.layout,
+ tvm_tirx.tirx.layout.TileLayout(
+ T.S[(4, 16) : (16, 1)] + b_buf.layout.apply(0, 0, shape=b_buf.shape)["m"]
+ ),
+ )
# sub[:, 1::2]: keep dim 0, split dim 1 into (4, 2) and fix the remainder at 1
assert [int(s) for s in c_buf.shape] == [4, 4, 16]
- assert int(tvm_tirx.sym.Analyzer().simplify(c_buf.elem_offset - a_buf.elem_offset)) == 16
+ assert int(tvm_tirx.sym.Analyzer().simplify(
+ c_buf.layout.apply(0).get("m", 0) - a_buf.layout.apply(0).get("m", 0)
+ )) == 16
assert_structural_equal(
- c_buf.layout, tvm_tirx.tirx.layout.TileLayout(T.S[(4, 4, 16) : (256, 32, 1)])
+ c_buf.layout,
+ tvm_tirx.tirx.layout.TileLayout(
+ T.S[(4, 4, 16) : (256, 32, 1)] + c_buf.layout.apply(0, 0, 0, shape=c_buf.shape)["m"]
+ ),
)
code = func.script()
@@ -1078,12 +1100,7 @@ def test_tensor_view_surgery_static_bounds_rejected():
def test_tensor_sub_swizzle_commutation():
- """A folded view offset moves into elem_offset only when it commutes
- with the swizzle, i.e. is a multiple of the swizzle period
- 2^(per_element + atom_len + swizzle_len). Sub-period offsets stay inside
- the derived tile layout's offset so the swizzle keeps applying to them
- (review finding: folding them outside produced wrong addresses). Both
- placements must be address-equivalent to the parent layout."""
+ """Subview offsets stay inside the tile layout and preserve parent addresses."""
def addr(buf, base, *coords):
analyzer = tvm_tirx.sym.Analyzer()
@@ -1091,7 +1108,7 @@ def test_tensor_sub_swizzle_commutation():
rel = buf.layout.apply(coords[0])["m"]
else:
rel = buf.layout.apply(*coords, shape=[int(s) for s in buf.shape])["m"]
- return int(analyzer.simplify((buf.elem_offset - base) + rel))
+ return int(analyzer.simplify(rel - base))
analyzer = tvm_tirx.sym.Analyzer()
compose = T.ComposeLayout(
@@ -1105,15 +1122,15 @@ def test_tensor_sub_swizzle_commutation():
A = T.alloc_tensor(
ty_args=[T.Tensor([4, 1024], "bfloat16", scope="shared.dyn", layout=compose)]
)
- B = A.sub[1] # offset 1024 = 2 * period: folds into elem_offset
+ B = A.sub[1] # offset 1024 = 2 * period
B[0] = T.bfloat16(0)
- C = A.sub[:, 512:1024] # offset 512 = period: folds into elem_offset
+ C = A.sub[:, 512:1024] # offset 512 = period
C[0, 0] = T.bfloat16(0)
# fmt: on
a_buf, b_buf, c_buf = _collect_tensors(func)
- base = a_buf.elem_offset
- assert int(analyzer.simplify(b_buf.elem_offset - base)) == 1024
+ base = a_buf.layout.apply(0).get("m", 0)
+ assert int(analyzer.simplify(b_buf.layout.apply(0).get("m", 0) - base)) == 1024
for j in (0, 1, 63, 511, 1023):
assert addr(a_buf, base, 1024 + j) == addr(b_buf, base, j)
for j in (0, 1, 255, 511):
@@ -1123,7 +1140,7 @@ def test_tensor_sub_swizzle_commutation():
assert from_source(code).script() == code
# Sub-period offsets do not commute: they stay inside the tile layout's
- # offset (elem_offset unchanged) and every address matches the parent.
+ # offset and every address matches the parent.
compose2 = T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(2, 16, 8) : (128, 8, 1)]))
# fmt: off
@@ -1147,7 +1164,7 @@ def test_tensor_sub_swizzle_commutation():
# fmt: on
a2, b2, c2, d2, e2, _ = _collect_tensors(func2)
- base2 = a2.elem_offset
+ base2 = a2.layout.apply(0).get("m", 0)
shape2 = [2, 16, 8]
for name, child, to_parent in [
("B", b2, lambda c: (c[0], 1, c[1])),
@@ -1155,7 +1172,6 @@ def test_tensor_sub_swizzle_commutation():
("D", d2, lambda c: (c[0], 1 + c[1], c[2])),
("E", e2, lambda c: (c[0], 1 + c[1], c[2])),
]:
- assert int(analyzer.simplify(child.elem_offset - base2)) == 0
child_shape = [int(s) for s in child.shape]
for flat in range(math.prod(child_shape)):
coords, rem = [], flat
@@ -1188,7 +1204,8 @@ def test_tensor_sub_swizzle_commutation():
a3, b3 = _collect_tensors(func3)
for j in range(8):
- assert addr(a3, a3.elem_offset, 8 + j) == addr(b3, a3.elem_offset, j) == 8 + j
+ assert (addr(a3, a3.layout.apply(0).get("m", 0), 8 + j)
+ == addr(b3, a3.layout.apply(0).get("m", 0), j) == 8 + j)
def test_tensor_factored_views():
@@ -1207,16 +1224,15 @@ def test_tensor_factored_views():
assert list(map(int, B.shape)) == [3, 16, 512]
analyzer = tvm_tirx.sym.Analyzer()
for row in range(16):
- address = B.layout.apply(1, row, 7, shape=B.shape)["m"] + B.elem_offset
+ address = B.layout.apply(1, row, 7, shape=B.shape)["m"]
expected = (64 + row // 4 * 16 + w * 4 + row % 4) * 512 + 7
assert analyzer.can_prove_equal(address, expected)
assert (
- int(D.layout.apply(1, row, 7, shape=D.shape)["m"] + D.elem_offset)
+ int(D.layout.apply(1, row, 7, shape=D.shape)["m"])
== (64 + row * 4 + 2) * 512 + 7
)
assert (
- int(F.layout.apply(1, row, 7, shape=F.shape)["m"] + F.elem_offset)
- == (64 + 32 + row) * 512 + 7
+ int(F.layout.apply(1, row, 7, shape=F.shape)["m"]) == (64 + 32 + row) * 512 + 7
)
M = A.view(192, 512, layout="default")
multi = (
@@ -1229,13 +1245,13 @@ def test_tensor_factored_views():
)
assert list(map(int, multi.shape)) == [48, 256]
for row, col in [(0, 0), (13, 19), (47, 255)]:
- actual = multi.layout.apply(row, col, shape=multi.shape)["m"] + multi.elem_offset
+ actual = multi.layout.apply(row, col, shape=multi.shape)["m"]
expected = (row // 2 * 8 + w * 2 + row % 2) * 512 + col // 8 * 16 + 8 + col % 8
assert analyzer.can_prove_equal(actual, expected)
a = tvm_tirx.tirx.Var("a", "int32")
b = tvm_tirx.tirx.Var("b", "int32")
selected = M.view(2, 4, 24, 512).sub[a, b]
- actual = selected.layout.apply(7, 11, shape=selected.shape)["m"] + selected.elem_offset
+ actual = selected.layout.apply(7, 11, shape=selected.shape)["m"]
assert analyzer.can_prove_equal(actual, (a * 96 + b * 24 + 7) * 512 + 11)
@@ -1481,7 +1497,9 @@ def test_memory_scope_canonical_constructor_aliases_and_reflected_scopes():
@T.function
def func(
pointer: T.Ptr(T.float32, scope),
- tensor: T.Tensor((4,), T.float32, scope, byte_offset=8, align=128, layout="default"),
+ tensor: T.Tensor(
+ (4,), T.float32, scope, align=128, layout=T.TileLayout(T.S[(4,)] + T.int64(2))
+ ),
):
tensor[0] = T.float32(1)
@@ -1492,17 +1510,17 @@ def test_memory_scope_canonical_constructor_aliases_and_reflected_scopes():
tensor_type = func.params[1].ty
assert pointer_type.scope.same_as(expected_scope)
assert tensor_type.scope.same_as(expected_scope)
- assert tensor_type.elem_offset.value == 2
+ assert tensor_type.layout.apply(0).get("m", 0).value == 2
assert tensor_type.data_alignment == 128
script = func.script()
if expected_scope == MemoryScope("global"):
assert "pointer: T.Ptr(T.float32)" in script
- assert "tensor: T.Tensor((4,), T.float32, elem_offset=" in script
+ assert "tensor: T.Tensor((4,), T.float32, data_alignment=" in script
else:
assert f'pointer: T.Ptr(T.float32, "{expected_scope.name}")' in script
assert f'tensor: T.Tensor((4,), T.float32, scope="{expected_scope.name}", ' in script
- assert "elem_offset=T.int64(2), data_alignment=128" in script
- assert 'layout=T.TileLayout.from_iters(shard=[T.Iter(4, 1, "m")], replica=[])' in script
+ assert "data_alignment=128" in script
+ assert "T.int64(2)" in script
restored = tvm_tirx.script.from_source(script, extra_vars={"T": T})
tvm_tirx.ir.assert_structural_equal(func, restored)
assert restored.params[0].ty.scope.same_as(expected_scope)
diff --git a/tirx/tests/python/tirx/test_tile_namespace.py b/tirx/tests/python/tirx/test_tile_namespace.py
index 5526fb9346..1f8a564f3f 100644
--- a/tirx/tests/python/tirx/test_tile_namespace.py
+++ b/tirx/tests/python/tirx/test_tile_namespace.py
@@ -100,7 +100,7 @@ def test_omitted_and_explicit_none_layout(scope):
assert T.Tensor((4, 8), "float32", scope=scope, layout=None).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]))
+ assert_structural_equal(explicit.layout, TileLayout.from_iters(TileLayout(S[4, 8]).shard))
with IRBuilder():
with I.function():
tensor_type = I.Tensor((4, 8), "float32", scope=scope)
@@ -114,27 +114,31 @@ def test_omitted_and_explicit_none_layout(scope):
[((8, 16), None), ((8, 16), (16, 1)), ((8, 1, 16), (16, 99, 1))],
)
def test_effective_layout_on_tensor_type(shape, strides):
- tensor_type = T.Tensor(shape, "float32", strides=strides)
+ tensor_type = T.Tensor(
+ shape, "float32", layout=T.TileLayout(T.S[shape:strides]) if strides else None
+ )
before = save_json(tensor_type)
- assert_structural_equal(tensor_type.effective_layout(), TileLayout(S[shape]))
+ effective = tensor_type.effective_layout()
+ assert_structural_equal(
+ effective, TileLayout(S[shape] if strides is None else S[shape:strides])
+ )
assert save_json(tensor_type) == before
def test_effective_global_layout_preserves_source_and_region():
- tensor = Var("tensor", TensorType((8, 16), "float32", elem_offset=13))
+ tensor = Var("tensor", TensorType((8, 16), "float32", layout=T.TileLayout(T.S[(8, 16)] + (13))))
before = save_json(tensor)
region = tensor[2:4, 3:7]
layout = region.source.ty.effective_layout()
- assert int(layout.apply(2, 3, shape=tensor.shape)["m"]) == 35
+ assert int(layout.apply(2, 3, shape=tensor.shape)["m"]) == 48
starts = [r.min for r in region.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 int(layout.apply(*starts, shape=tensor.shape)["m"]) == 48
+ assert int(layout.apply(starts[0] + 1, starts[1] + 3, shape=tensor.shape)["m"]) == 67
assert save_json(tensor) == before
- compact = Var("tensor", TensorType((8, 16), "float32", strides=(16, 1)))
- assert_structural_equal(compact.ty.effective_layout(), layout)
- strided = Var("tensor", TensorType((8, 16), "float32", strides=(32, 1)))
- with pytest.raises(ValueError, match="noncompact"):
- strided.ty.effective_layout()
+ compact = Var("tensor", TensorType((8, 16), "float32", layout=TileLayout(S[(8, 16) : (16, 1)])))
+ assert int(compact.ty.effective_layout().apply(2, 3, shape=compact.shape)["m"]) == 35
+ strided = Var("tensor", TensorType((8, 16), "float32", layout=TileLayout(S[(8, 16) : (32, 1)])))
+ assert int(strided.ty.effective_layout().apply(2, 3, shape=strided.shape)["m"]) == 67
with pytest.raises(ValueError, match="explicit layout"):
Var("tensor", TensorType((8, 16), "float32", scope="shared")).ty.effective_layout()
explicit = Var("tensor", TensorType((8, 16), "float32", scope="shared", layout=layout))
@@ -147,7 +151,8 @@ def test_none_layout_views_keep_offsets_and_strides():
with I.function():
source = I.bind(
I.decl_tensor(
- I.ptr("float32"), ty_args=[I.Tensor((4, 8), "float32", elem_offset=5)]
+ I.ptr("float32"),
+ ty_args=[I.Tensor((4, 8), "float32", layout=T.TileLayout(T.S[(4, 8)] + (5)))],
)
)
reshaped = source.view(2, 2, 8)
@@ -156,16 +161,29 @@ def test_none_layout_views_keep_offsets_and_strides():
row = source.sub[2, :]
stepped = source.sub[:, 1::2]
for view in (reshaped, permuted, sliced, row, stepped):
- assert view.layout is None
+ assert view.layout is not None
assert list(map(int, reshaped.shape)) == [2, 2, 8]
- assert list(map(int, permuted.strides)) == [1, 8]
- assert list(map(int, sliced.strides)) == [8, 1]
- assert int(sliced.elem_offset) == 15
- assert int(row.elem_offset) == 21
- assert list(map(int, stepped.strides)) == [8, 2]
- assert int(stepped.elem_offset) == 6
- with pytest.raises(ValueError, match="non-contiguous"):
- permuted.view(32)
+ assert list(map(int, [it.stride for it in permuted.layout.shard])) == [1, 8]
+ assert list(map(int, [it.stride for it in sliced.layout.shard])) == [8, 1]
+ assert (
+ int(sliced.layout.apply(*([0] * len(sliced.shape)), shape=sliced.shape).get("m", 0))
+ == 15
+ )
+ assert int(row.layout.apply(*([0] * len(row.shape)), shape=row.shape).get("m", 0)) == 21
+ assert list(map(int, [it.stride for it in stepped.layout.shard])) == [8, 2]
+ assert (
+ int(
+ stepped.layout.apply(*([0] * len(stepped.shape)), shape=stepped.shape).get(
+ "m", 0
+ )
+ )
+ == 6
+ )
+ linear = permuted.view(32)
+ for i in range(32):
+ assert int(linear.layout.apply(i, shape=linear.shape)["m"]) == (
+ int(permuted.layout.apply(i // 4, i % 4, shape=permuted.shape)["m"])
+ )
parents = {bind.var: bind.value.args[0].args[0] for bind in builder.get().body.seq[1:]}
for view in (reshaped, permuted, sliced, row, stepped):
root = view
@@ -202,7 +220,10 @@ def test_none_layout_dtype_view_accepts_symbolic_extent():
assert viewed.layout is None
assert_structural_equal(viewed.shape[0], size // 2)
misaligned = I.bind(
- I.decl_tensor(I.ptr("float32"), ty_args=[I.Tensor((8,), "float32", elem_offset=1)])
+ I.decl_tensor(
+ I.ptr("float32"),
+ ty_args=[I.Tensor((8,), "float32", layout=T.TileLayout(T.S[(8,)] + (1)))],
+ )
)
with pytest.raises(ValueError, match="align"):
misaligned.view("uint64")
@@ -268,7 +289,9 @@ def test_tile_load_infers_full_global_layout_with_offset():
@T.function
def kernel(A: T.Tensor((144,), "float32"), D: T.Tensor((2, 4), "float32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
- source = T.decl_tensor(A.data, ty_args=[T.Tensor((8, 16), "float32", elem_offset=7)])
+ source = T.decl_tensor(
+ A.data, ty_args=[T.Tensor((8, 16), "float32", layout=T.TileLayout(T.S[(8, 16)] + (7)))]
+ )
registers = T.alloc_local((2, 4), "float32", layout="default")
source_words = source.sub[2:4, 3:7].view(
2, 2, 2, layout=TileLayout(S[(2, 2, 2) : (16, 2, 1)])
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 caeea7e4f9..19e86c38e2 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
@@ -176,17 +176,13 @@ def test_dsmem(shape, dtype, src_spec, dst_spec, expected):
src_raw = pool.alloc([src_phys], dtype, align=128)
src_smem = T.decl_tensor(
src_raw.data,
- ty_args=[
- T.Tensor(list(shape), dtype, elem_offset=0, scope="shared.dyn", layout=src_layout)
- ],
+ ty_args=[T.Tensor(list(shape), dtype, scope="shared.dyn", layout=src_layout)],
)
# dst_smem: dispatch writes here (on remote CTA), CTA 1 reads
dst_raw = pool.alloc([dst_phys], dtype, align=128)
dst_smem = T.decl_tensor(
dst_raw.data,
- ty_args=[
- T.Tensor(list(shape), dtype, elem_offset=0, scope="shared.dyn", layout=dst_layout)
- ],
+ ty_args=[T.Tensor(list(shape), dtype, scope="shared.dyn", layout=dst_layout)],
)
mbar = MBarrier(pool, 1)
pool.commit()
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 45ba41d18c..a5f483754d 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
@@ -867,7 +867,7 @@ def _recipe_views(kwargs, recipe):
)
assert int(sa.min()) * original_bits % bits == 0
base = int(sa.min()) * original_bits // bits
- layout = _plain_layout(box)
+ layout = TileLayout(S[box] + base)
if recipe.enums[1]:
layout = ComposeLayout(7 - int(math.log2(bits)), recipe.enums[1], 3, layout)
sv = tvm_tirx.tirx.Var(
@@ -878,7 +878,6 @@ def _recipe_views(kwargs, recipe):
data=s.data,
scope="shared.dyn",
layout=layout,
- elem_offset=base,
data_alignment=1024,
),
)
@@ -887,7 +886,7 @@ def _recipe_views(kwargs, recipe):
coords = np.indices(box, dtype=np.int64).reshape(len(box), -1)
coords += np.asarray(tuple(reversed(recipe.coordinates)))[:, None]
new_ga = np.sum(coords * np.asarray(strides)[:, None], axis=0)
- new_sa = _address_array(layout, box, [(0, x) for x in box], base)
+ new_sa = _address_array(layout, box, [(0, x) for x in box])
def bytes_at(address, width):
return (address[:, None] * width + np.arange(width)[None, :]).reshape(-1)
@@ -944,17 +943,13 @@ def _direct(
box=None,
coords=None,
s_region=None,
- g_elem_offset=0,
- s_elem_offset=0,
sctx=None,
**fields,
):
s_shape = s_shape or g_shape
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
- ),
+ tvm_tirx.tirx.TensorType(g_shape, dtype, layout=g_layout or _plain_layout(g_shape)),
)
s = tvm_tirx.tirx.Var(
"A_smem",
@@ -963,7 +958,6 @@ def _direct(
dtype,
scope="shared.dyn",
layout=s_layout or _plain_layout(s_shape),
- elem_offset=s_elem_offset,
data_alignment=1024,
),
)
@@ -1001,12 +995,20 @@ def test_explicit_rejects_unknown_or_nonunit_implicit_inner_stride(stride):
def test_explicit_view_base_offset_is_encoded_without_coordinate_rewrite():
- _, host, _ = _direct(_load(2), g_shape=(64, 64), g_elem_offset=8)
+ _, host, _ = _direct(_load(2), g_shape=(64, 64), g_layout=TileLayout(S[(64, 64)] + 8))
signature = _encode_signature(_collect_encodes(host)[0])
assert isinstance(signature["base"], tvm_tirx.ir.Call)
assert signature["base"].op.name == "tirx.address_of"
base = signature["base"].args[0]
- assert int(base.source.ty.elem_offset) * 2 == 16
+ assert (
+ int(
+ base.source.ty.layout.apply(
+ *([0] * len(base.source.ty.shape)), shape=base.source.ty.shape
+ ).get("m", 0)
+ )
+ * 2
+ == 16
+ )
assert _ints(base.indices) == (0, 0)
assert _ints(signature["dims"]) == (64, 64)
@@ -1046,7 +1048,10 @@ def test_explicit_allows_different_operand_ranks_with_equal_payload_bytes():
],
)
mbar = T.decl_tensor(
- dyn.data, ty_args=[T.Tensor((1,), "uint64", elem_offset=16, data=dyn.data)]
+ dyn.data,
+ ty_args=[
+ T.Tensor((1,), "uint64", layout=T.TileLayout(T.S[(1,)] + (16)), data=dyn.data)
+ ],
)
T.cuda.tile.apply(
T.ptx[
@@ -1656,27 +1661,27 @@ def _build_sparse_decode_qo_tma_regression():
],
)
ts = T.decl_tensor(
- dyn.data,
+ T.ptr_byte_offset(dyn.data, q_elements * 2, ty=T.Ptr("bfloat16", "shared.dyn")),
ty_args=[
T.Tensor(
(64, 64),
"bfloat16",
scope="shared.dyn",
layout=tail_layout,
- elem_offset=q_elements,
align=1024,
)
],
)
os = T.decl_tensor(
- dyn.data,
+ T.ptr_byte_offset(
+ dyn.data, (q_elements + tail_elements) * 2, ty=T.Ptr("bfloat16", "shared.dyn")
+ ),
ty_args=[
T.Tensor(
(64, 512),
"bfloat16",
scope="shared.dyn",
layout=q_layout,
- elem_offset=q_elements + tail_elements,
align=1024,
)
],
@@ -1688,8 +1693,7 @@ def _build_sparse_decode_qo_tma_regression():
(1,),
"uint64",
scope="shared.dyn",
- layout="default",
- elem_offset=shared_bytes // 8,
+ layout=T.TileLayout(T.S[(1,)] + (shared_bytes // 8)),
)
],
)
@@ -2093,7 +2097,7 @@ def test_caller_stages_unaligned_shared_region(warp_offset):
# Original offset-count regressions wrote at 32 + warp_offset*64 + 2*64
# halfwords, which is 64 bytes off the native TMA alignment. Preserve the
# destination through explicit aligned staging, then caller scalar stores.
- layout = TileLayout(S[(4, 64)] + warp_offset * 64)
+ layout = TileLayout(S[(4, 64)] + warp_offset * 64 + 32)
with pytest.raises(RuntimeError, match="128-byte aligned"):
_direct(
_load(2),
@@ -2101,7 +2105,6 @@ def test_caller_stages_unaligned_shared_region(warp_offset):
s_shape=(4, 64),
s_region=((2, 2), (0, 64)),
s_layout=layout,
- s_elem_offset=32,
)
@T.function
@@ -2111,11 +2114,7 @@ def test_caller_stages_unaligned_shared_region(warp_offset):
backing = T.alloc_shared((512,), "float16", layout="default", align=128)
target = T.decl_tensor(
backing.data,
- ty_args=[
- T.Tensor(
- (4, 64), "float16", scope="shared", layout=layout, elem_offset=32, align=128
- )
- ],
+ ty_args=[T.Tensor((4, 64), "float16", scope="shared", layout=layout, align=128)],
)
stage = T.alloc_shared((2, 64), "float16", layout="default", align=128)
mbar = T.alloc_shared((1,), "uint64", layout="default")
diff --git a/tirx/tests/python/tirx/tile/cuda/test_explicit_patterns.py b/tirx/tests/python/tirx/tile/cuda/test_explicit_patterns.py
index e3894a1301..a243682396 100644
--- a/tirx/tests/python/tirx/tile/cuda/test_explicit_patterns.py
+++ b/tirx/tests/python/tirx/tile/cuda/test_explicit_patterns.py
@@ -122,7 +122,6 @@ def _dispatch(
opcode=LD,
scope="thread",
align=16,
- elem_offset=0,
active=32,
):
matrix = opcode.startswith("ldmatrix")
@@ -134,7 +133,6 @@ def _dispatch(
scope="shared" if matrix else "global",
layout=src_layout or TileLayout(S[src_shape]),
align=align,
- elem_offset=elem_offset,
),
)
dst_shape = dst_shape or src_shape
@@ -156,7 +154,7 @@ def test_contract_rejections():
with pytest.raises(RuntimeError, match="consecutive elements"):
_dispatch(src_layout=TileLayout(S[(2, 4) : (8, 2)]))
with pytest.raises(RuntimeError, match="16-byte aligned"):
- _dispatch(elem_offset=1)
+ _dispatch(src_layout=TileLayout(S[(2, 4)] + 1))
with pytest.raises(RuntimeError, match="known 16-byte-aligned base"):
_dispatch(align=4)
with pytest.raises(RuntimeError, match="Repeat shapes must agree"):
@@ -236,11 +234,14 @@ def test_region_starts_and_element_offsets_gpu():
for i in T.unroll(32):
raw[i] = T.uint32(777)
src = T.decl_tensor(
- A.data, ty_args=[T.Tensor((3, 8), "uint32", elem_offset=4, align=16, layout="default")]
+ A.data,
+ ty_args=[T.Tensor((3, 8), "uint32", align=16, layout=T.TileLayout(T.S[(3, 8)] + (4)))],
)
dst = T.decl_tensor(
raw.data,
- ty_args=[T.Tensor((3, 8), "uint32", elem_offset=4, scope="local", layout="default")],
+ ty_args=[
+ T.Tensor((3, 8), "uint32", scope="local", layout=T.TileLayout(T.S[(3, 8)] + (4)))
+ ],
)
T.cuda.tile.apply(T.ptx[LD], dst[1:3, 2:6], src[0:2, 4:8])
for i in T.unroll(32):
diff --git a/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py b/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py
index dc06902ae0..05aed14aa8 100644
--- a/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py
+++ b/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py
@@ -756,7 +756,6 @@ def test_reg_view_preserves_shared_and_register_addresses(case):
for layout, view in ((r_layout, r_view), (s_layout, s_view)):
before = layout.apply(row, col, shape=shape)
after = dict(view.ty.layout.apply(*coords, shape=view.ty.shape))
- after["m"] = after.get("m", 0) + view.ty.elem_offset
assert {k: int(analyzer.simplify(v)) for k, v in before.items()} == {
k: int(analyzer.simplify(v)) for k, v in after.items()
}
diff --git a/tirx/tests/python/tirx/transform/test_convert_ssa.py b/tirx/tests/python/tirx/transform/test_convert_ssa.py
index ebdaf92c23..e550a6bb56 100644
--- a/tirx/tests/python/tirx/transform/test_convert_ssa.py
+++ b/tirx/tests/python/tirx/transform/test_convert_ssa.py
@@ -421,7 +421,13 @@ def test_reused_loop_var_in_decl_tensor_elem_offset():
tensor = T.decl_tensor(
data,
ty_args=[
- T.Tensor((128,), "float32", data=data, elem_offset=i * 128, scope="shared.dyn")
+ T.Tensor(
+ (128,),
+ "float32",
+ data=data,
+ layout=T.TileLayout(T.S[(128,)] + (i * 128)),
+ scope="shared.dyn",
+ )
],
)
T.evaluate(tensor[0])
@@ -446,7 +452,11 @@ def test_reused_loop_var_in_decl_tensor_elem_offset():
data,
ty_args=[
T.Tensor(
- (128,), "float32", data=data, elem_offset=i * 128, scope="shared.dyn"
+ (128,),
+ "float32",
+ data=data,
+ layout=T.TileLayout(T.S[(128,)] + (i * 128)),
+ scope="shared.dyn",
)
],
)
@@ -456,7 +466,11 @@ def test_reused_loop_var_in_decl_tensor_elem_offset():
data,
ty_args=[
T.Tensor(
- (128,), "float32", data=data, elem_offset=j * 128, scope="shared.dyn"
+ (128,),
+ "float32",
+ data=data,
+ layout=T.TileLayout(T.S[(128,)] + (j * 128)),
+ scope="shared.dyn",
)
],
)
@@ -466,7 +480,11 @@ def test_reused_loop_var_in_decl_tensor_elem_offset():
data,
ty_args=[
T.Tensor(
- (128,), "float32", data=data, elem_offset=k * 128, scope="shared.dyn"
+ (128,),
+ "float32",
+ data=data,
+ layout=T.TileLayout(T.S[(128,)] + (k * 128)),
+ scope="shared.dyn",
)
],
)
diff --git a/tirx/tests/python/tirx/transform/test_flatten_tensor.py b/tirx/tests/python/tirx/transform/test_flatten_tensor.py
index 5914e43015..266f59c759 100644
--- a/tirx/tests/python/tirx/transform/test_flatten_tensor.py
+++ b/tirx/tests/python/tirx/transform/test_flatten_tensor.py
@@ -338,7 +338,15 @@ def test_strided():
for i0 in T.serial(4):
B = T.alloc_tensor(ty_args=[T.Tensor([4, 17], "float32")])
B_1 = T.decl_tensor(
- B.data, ty_args=[T.Tensor([4, 16], "float32", strides=[17, 1], data=B.data)]
+ B.data,
+ ty_args=[
+ T.Tensor(
+ [4, 16],
+ "float32",
+ layout=T.TileLayout(T.S[[4, 16] : [17, 1]]),
+ data=B.data,
+ )
+ ],
)
for i1, j in T.grid(4, 16):
B_1[i1, j] = A[i0 * 4 + i1, j] + 1.0
@@ -490,7 +498,12 @@ def test_flatten_remaps_loads_in_folded_elem_offset():
n[0] = 4
base = T.alloc_shared([128], "uint64")
mbar = T.decl_tensor(
- base.data, ty_args=[T.Tensor((1,), "uint64", elem_offset=n[0], scope="shared")]
+ base.data,
+ ty_args=[
+ T.Tensor(
+ (1,), "uint64", layout=T.TileLayout(T.S[(1,)] + (n[0])), scope="shared"
+ )
+ ],
)
mbar[0] = T.uint64(1)
@@ -503,9 +516,7 @@ def test_flatten_remaps_loads_in_folded_elem_offset():
base = T.alloc_shared([128], "uint64", layout=None)
mbar = T.decl_tensor(
base.data,
- ty_args=[
- T.Tensor((1,), "uint64", elem_offset=T.int32(0), scope="shared", layout=None)
- ],
+ ty_args=[T.Tensor((1,), "uint64", scope="shared", layout=None)],
)
mbar[n[0]] = T.uint64(1)
@@ -537,16 +548,22 @@ def test_tensor_flatten_offset():
@I.ir_module
class Before:
@T.function(private=True)
- def main(m: T.int32, n: T.int32, A: T.Tensor((m, n), "float32", elem_offset=T.int32(100))):
+ def main(
+ m: T.int32,
+ n: T.int32,
+ A: T.Tensor((m, n), "float32", layout=T.TileLayout(T.S[(m, n)] + (T.int32(100)))),
+ ):
T.evaluate(A[2, 3])
@I.ir_module
class Expected:
@T.function(private=True)
- def main(m: T.int32, n: T.int32, A: T.Tensor((m, n), "float32", elem_offset=T.int32(100))):
- flat = T.decl_tensor(
- A.data, ty_args=[T.Tensor((m * n,), "float32", elem_offset=T.int32(0), layout=None)]
- )
+ def main(
+ m: T.int32,
+ n: T.int32,
+ A: T.Tensor((m, n), "float32", layout=T.TileLayout(T.S[(m, n)] + (T.int32(100)))),
+ ):
+ flat = T.decl_tensor(A.data, ty_args=[T.Tensor((m * n,), "float32", layout=None)])
T.evaluate(flat[n * 2 + 103])
After = FlattenTensor()(Before)
@@ -559,16 +576,22 @@ def test_tensor_flatten_offset():
ids=["matrix", "vector", "scalar", "strided"],
)
def test_tensor_flatten(shape, strides, span):
+ layout = T.TileLayout(T.S[shape:strides]) if strides else None
+
@I.ir_module
class Before:
@T.function(private=True)
- def main(A: T.Tensor(shape, "float32", strides=strides)):
+ def main(
+ A: T.Tensor(shape, "float32", layout=layout),
+ ):
T.evaluate(A[tuple(0 for _ in shape)])
@I.ir_module
class Expected:
@T.function(private=True)
- def main(A: T.Tensor(shape, "float32", strides=strides)):
+ def main(
+ A: T.Tensor(shape, "float32", layout=layout),
+ ):
flat = T.decl_tensor(A.data, ty_args=[T.Tensor((span,), "float32", layout=None)])
T.evaluate(flat[0])
@@ -589,8 +612,8 @@ def test_tensor_index_merge_mult_mod(pattern):
n = T.dynamic("n", "int32")
s = T.dynamic("s", "int32")
- def tensor_strides(s):
- return (s, 1) if pattern == 1 else None
+ def tensor_layout(m, n, s):
+ return T.TileLayout(T.S[(m, n) : (s, 1)]) if pattern == 1 else None
def storage_span(m, n, s):
return T.max(s * m, (m - 1) * s + n) if pattern == 1 else m * n
@@ -621,7 +644,11 @@ def test_tensor_index_merge_mult_mod(pattern):
s: T.int32,
k0: T.int32,
k1: T.int32,
- A: T.Tensor((m, n), "float32", strides=tensor_strides(s)),
+ A: T.Tensor(
+ (m, n),
+ "float32",
+ layout=tensor_layout(m, n, s),
+ ),
):
T.evaluate(A[indices(k0, k1, s, n)])
@@ -634,7 +661,11 @@ def test_tensor_index_merge_mult_mod(pattern):
s: T.int32,
k0: T.int32,
k1: T.int32,
- A: T.Tensor((m, n), "float32", strides=tensor_strides(s)),
+ A: T.Tensor(
+ (m, n),
+ "float32",
+ layout=tensor_layout(m, n, s),
+ ),
):
flat = T.decl_tensor(
A.data, ty_args=[T.Tensor((storage_span(m, n, s),), "float32", layout=None)]
diff --git a/tirx/tests/python/tirx/transform/test_force_narrow_index_to_i32.py b/tirx/tests/python/tirx/transform/test_force_narrow_index_to_i32.py
index 495ac205ad..bb503dcbb9 100644
--- a/tirx/tests/python/tirx/transform/test_force_narrow_index_to_i32.py
+++ b/tirx/tests/python/tirx/transform/test_force_narrow_index_to_i32.py
@@ -33,8 +33,8 @@ def test_thread_axis1():
@T.function(private=True)
def expected(
- A: T.Tensor((64,), "float32", elem_offset=T.int32(0)),
- B: T.Tensor((64,), "float32", elem_offset=T.int32(0)),
+ A: T.Tensor((64,), "float32"),
+ B: T.Tensor((64,), "float32"),
):
blockIdx_x = T.launch_thread("blockIdx.x", 2)
threadIdx_x = T.launch_thread("threadIdx.x", 32)
@@ -152,9 +152,9 @@ def test_thread_axis2():
@T.function
def expected(
- T_reshape: T.Tensor((1, 12, 384, 384), "float32", elem_offset=T.int32(0)),
- placeholder_1: T.Tensor((1, 12, 384, 384), "bool", elem_offset=T.int32(0)),
- T_where: T.Tensor((1, 12, 384, 384), "float32", elem_offset=T.int32(0)),
+ T_reshape: T.Tensor((1, 12, 384, 384), "float32"),
+ placeholder_1: T.Tensor((1, 12, 384, 384), "bool"),
+ T_where: T.Tensor((1, 12, 384, 384), "float32"),
):
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
with T.launch_thread("blockIdx.x", 256) as i0_i1_i2_i3_fused_1:
@@ -247,8 +247,8 @@ def test_block():
@T.function(private=True)
def expected(
- A: T.Tensor((128,), "float32", elem_offset=T.int32(0)),
- B: T.Tensor((128,), "float32", elem_offset=T.int32(0)),
+ A: T.Tensor((128,), "float32"),
+ B: T.Tensor((128,), "float32"),
):
for i in T.serial(0, T.int32(16)):
for j in T.serial(0, T.int32(8)):
@@ -271,8 +271,8 @@ def test_i16_tensor():
@T.function(private=True)
def expected(
- A: T.Tensor((128,), "int16", elem_offset=T.int32(0)),
- B: T.Tensor((128,), "int16", elem_offset=T.int32(0)),
+ A: T.Tensor((128,), "int16"),
+ B: T.Tensor((128,), "int16"),
):
for i in T.serial(0, 16):
for j in T.serial(0, 16):
@@ -334,8 +334,8 @@ def test_pod_params_and_select():
class Expected:
@T.function
def main(
- A: T.Tensor((4,), "float32", elem_offset=T.int32(0)),
- B: T.Tensor((4,), "float32", elem_offset=T.int32(0)),
+ A: T.Tensor((4,), "float32"),
+ B: T.Tensor((4,), "float32"),
n: T.int32,
):
for i in range(4):
@@ -356,8 +356,8 @@ def test_if_then_else_index():
class Expected:
@T.function
def main(
- A: T.Tensor((4,), "float32", elem_offset=T.int32(0)),
- B: T.Tensor((1,), "float32", elem_offset=T.int32(0)),
+ A: T.Tensor((4,), "float32"),
+ B: T.Tensor((1,), "float32"),
n: T.int32,
):
B[0] = A[T.if_then_else(n < 0, n + 1, n)]
@@ -381,13 +381,11 @@ def test_conditional_index_mixed_width_branches():
class Expected:
@T.function
def main(
- A: T.Tensor((4,), "float32", elem_offset=T.int32(0)),
- B: T.Tensor((4,), "float32", elem_offset=T.int32(0)),
+ A: T.Tensor((4,), "float32"),
+ B: T.Tensor((4,), "float32"),
n: T.int32,
):
- opaque_index = T.alloc_tensor(
- ty_args=[T.Tensor((1,), "int64", scope="local", elem_offset=T.int32(0))]
- )
+ opaque_index = T.alloc_tensor(ty_args=[T.Tensor((1,), "int64", scope="local")])
opaque_index[0] = T.call_extern("opaque_index", n, ty="int64")
B[0] = A[T.if_then_else(n < 0, opaque_index[0], T.Cast("int64", n))]
B[1] = A[T.if_then_else(n < 0, T.Cast("int64", n), opaque_index[0])]
@@ -409,7 +407,7 @@ def test_clz():
@tvm_tirx.script.ir_module
class Expected:
@T.function
- def main(B: T.Tensor((4,), "int32", elem_offset=T.int32(0))):
+ def main(B: T.Tensor((4,), "int32")):
for i in range(4):
B[i] = T.clz(i) - 32 + 64
@@ -428,8 +426,8 @@ def test_right_shift_preserves_sign_extension_after_narrowing():
class Expected:
@T.function
def main(
- A: T.Tensor((6,), "float32", elem_offset=T.int32(0)),
- B: T.Tensor((1,), "float32", elem_offset=T.int32(0)),
+ A: T.Tensor((6,), "float32"),
+ B: T.Tensor((1,), "float32"),
n: T.int32,
):
B[0] = A[T.shift_right(T.truncmod(n - 8, 6), 31)]
@@ -463,8 +461,8 @@ def test_right_shift_dynamic_and_vector_amounts():
class Expected:
@T.function
def main(
- A: T.Tensor((6,), "float32", elem_offset=T.int32(0)),
- B: T.Tensor((5,), "float32", elem_offset=T.int32(0)),
+ A: T.Tensor((6,), "float32"),
+ B: T.Tensor((5,), "float32"),
n: T.int32,
shift: T.int32,
):
@@ -502,8 +500,8 @@ def test_left_shift_dynamic_and_vector_amounts_remain_valid():
class Expected:
@T.function
def main(
- A: T.Tensor((1,), "float32", elem_offset=T.int32(0)),
- B: T.Tensor((5,), "float32", elem_offset=T.int32(0)),
+ A: T.Tensor((1,), "float32"),
+ B: T.Tensor((5,), "float32"),
n: T.int32,
shift: T.int32,
):
@@ -535,13 +533,11 @@ def test_let_binding():
@tvm_tirx.script.ir_module
class Expected:
@T.function
- def main(Buf: T.Tensor([n], "int32", elem_offset=T.int32(0))):
+ def main(Buf: T.Tensor([n], "int32")):
# The pass narrows indexing variables (n, the For extent) but leaves
# an explicitly-typed `T.Cast(..., "int64")` storage alone; a Cast to
# int32 is inserted at the use site (the For iter) instead.
- ceil_log2 = T.alloc_tensor(
- ty_args=[T.Tensor((1,), "int64", scope="local", elem_offset=T.int32(0))]
- )
+ ceil_log2 = T.alloc_tensor(ty_args=[T.Tensor((1,), "int64", scope="local")])
ceil_log2[0] = T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n))))
for i in range(T.Cast("int32", ceil_log2[0])):
T.evaluate(0)
diff --git a/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py b/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
index 1172f15062..538f90c894 100644
--- a/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
+++ b/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
@@ -1011,7 +1011,12 @@ def test_lower_decl_tensor_pointer():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
buf = T.alloc_tensor(ty_args=[T.Tensor([1024], "uint8", scope="shared.dyn")])
A = T.decl_tensor(
- buf.data, ty_args=[T.Tensor([128], "float16", elem_offset=32, data=buf.data)]
+ buf.data,
+ ty_args=[
+ T.Tensor(
+ [128], "float16", layout=T.TileLayout(T.S[[128]] + (32)), data=buf.data
+ )
+ ],
)
T.evaluate(A.ptr_to([64]))
@@ -1030,7 +1035,12 @@ def test_lower_decl_tensor_pointer():
A = T.decl_tensor(
buf.data,
ty_args=[
- T.Tensor((128,), "float16", elem_offset=32, scope="shared.dyn", layout=None)
+ T.Tensor(
+ (128,),
+ "float16",
+ scope="shared.dyn",
+ layout=T.TileLayout(T.S[(128,)] + 32),
+ )
],
)
T.address_of(A[64])
@@ -1052,7 +1062,12 @@ def test_lower_decl_tensor_pointer():
A = T.decl_tensor(
buf.data,
ty_args=[
- T.Tensor((128,), T.float16, scope="shared.dyn", elem_offset=T.int64(32))
+ T.Tensor(
+ (128,),
+ T.float16,
+ scope="shared.dyn",
+ layout=T.TileLayout(T.S[(128,)] + 32),
+ )
],
)
T.address_of(A[64])
@@ -1210,7 +1225,7 @@ def test_lower_tensor_offset():
A = T.alloc_tensor(ty_args=[T.Tensor([64, 64], "float16", scope="local")])
A0 = T.decl_tensor(
A.ptr_to([32, 32]),
- ty_args=[T.Tensor([64], "float16", elem_offset=0, data=A.ptr_to([32, 32]))],
+ ty_args=[T.Tensor([64], "float16", data=A.ptr_to([32, 32]))],
)
T.evaluate(T.address_of(A0[32]))
@@ -1226,7 +1241,7 @@ def test_lower_tensor_offset():
A = T.alloc_local((4096,), "float16", layout=None)
A0 = T.decl_tensor(
T.address_of(A[2080]),
- ty_args=[T.Tensor((64,), "float16", elem_offset=0, scope="local", layout=None)],
+ ty_args=[T.Tensor((64,), "float16", scope="local", layout=None)],
)
T.address_of(A0[32])
@@ -1263,9 +1278,7 @@ def test_lower_alloc_decl_tensor_outside_of_parser():
self.A = T.alloc_local([1], "float16")
self.B = T.alloc_local([1], "float16")
self.C = T.bind(
- T.decl_tensor(
- smem, ty_args=[T.Tensor([1], "float16", elem_offset=0, scope="shared.dyn")]
- )
+ T.decl_tensor(smem, ty_args=[T.Tensor([1], "float16", scope="shared.dyn")])
)
@register_mutable_decl("TestMutableCells.int_var1")
@@ -1314,7 +1327,7 @@ def test_lower_alloc_decl_tensor_outside_of_parser():
B = T.alloc_local((1,), "float16", layout=None)
C = T.decl_tensor(
smem.data,
- ty_args=[T.Tensor((1,), "float16", elem_offset=0, scope="shared.dyn", layout=None)],
+ ty_args=[T.Tensor((1,), "float16", scope="shared.dyn", layout=None)],
)
A[0] = T.float16(1)
B[0] = T.float16(2)
diff --git a/tirx/tests/python/tirx/transform/test_make_packed_api.py b/tirx/tests/python/tirx/transform/test_make_packed_api.py
index 98d96a5fdb..e4df85ee7b 100644
--- a/tirx/tests/python/tirx/transform/test_make_packed_api.py
+++ b/tirx/tests/python/tirx/transform/test_make_packed_api.py
@@ -1112,6 +1112,32 @@ def test_linear_shape_binding():
tvm_tirx.ir.assert_structural_equal(After, Expected)
+def test_strides_and_linear_layout_offset():
+ s = T.dynamic("s", "int64")
+ m = T.dynamic("m", "int64")
+
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(
+ A: T.Tensor((3, 4), "float32", layout=T.TileLayout(T.S[(3, 4) : (s, 1)] + 4 * m)),
+ ):
+ T.func_attr({"target": T.target("llvm", host="llvm")})
+ T.evaluate(s + m)
+
+ after = tvm_tirx.tirx.transform.MakePackedAPI()(Before)
+ script = after.script()
+ assert "A.strides" in script
+ offset_check = next(
+ line for line in script.splitlines() if '"A.elem_offset"' in line and "assert" in line
+ )
+ assert "T.int64(4)) == T.int64(0)" in offset_check
+ assert "divisible" in offset_check
+ assert 'm: T.let[T.int64] = T.Div(T.Cast("int64", T.Div(' in script
+ assert "T.uint64(4))), T.int64(4))" in script
+ assert 'offset={"m": T.int64(4) * m}' in script
+
+
def test_forward_reference_symbolic_variable():
"""MakePackedAPI succeeds when a symbolic variable is used before it is defined.
diff --git a/tirx/tests/python/tirx/transform/test_narrow_datatype.py b/tirx/tests/python/tirx/transform/test_narrow_datatype.py
index 65d472d14f..ba0be0f4a2 100644
--- a/tirx/tests/python/tirx/transform/test_narrow_datatype.py
+++ b/tirx/tests/python/tirx/transform/test_narrow_datatype.py
@@ -89,12 +89,8 @@ def test_symbolic_extent(dtype):
class Expected:
@T.function
def main(A: T.handle("float32"), B: T.handle("float32"), m: scalar, n: scalar):
- A_buf = T.decl_tensor(
- A, ty_args=[T.Tensor((m * n,), "float32", data=A, elem_offset=T.int32(0))]
- )
- B_buf = T.decl_tensor(
- B, ty_args=[T.Tensor((m * n,), "float32", data=B, elem_offset=T.int32(0))]
- )
+ A_buf = T.decl_tensor(A, ty_args=[T.Tensor((m * n,), "float32", data=A)])
+ B_buf = T.decl_tensor(B, ty_args=[T.Tensor((m * n,), "float32", data=B)])
for i in T.serial(m):
for j in T.serial(n):
B_buf[i * n + j] = A_buf[i * n + j] + T.float32(1)
diff --git a/tirx/tests/python/tirx/transform/test_simplify.py b/tirx/tests/python/tirx/transform/test_simplify.py
index ffc878ab4b..747293cfa1 100644
--- a/tirx/tests/python/tirx/transform/test_simplify.py
+++ b/tirx/tests/python/tirx/transform/test_simplify.py
@@ -738,7 +738,10 @@ def test_simplify_trivial_let_elem_offset():
def before(A_ptr: T.handle("float32"), A_offset: T.int32):
A_offset_redef = A_offset
A = T.decl_tensor(
- A_ptr, ty_args=[T.Tensor(1, "float32", elem_offset=A_offset_redef, data=A_ptr)]
+ A_ptr,
+ ty_args=[
+ T.Tensor(1, "float32", layout=T.TileLayout(T.S[1] + (A_offset_redef)), data=A_ptr)
+ ],
)
A[0] = 42.0
@@ -746,7 +749,10 @@ def test_simplify_trivial_let_elem_offset():
def expected(A_ptr: T.handle("float32"), A_offset: T.int32):
A_offset_redef = A_offset
A = T.decl_tensor(
- A_ptr, ty_args=[T.Tensor(1, "float32", elem_offset=A_offset_redef, data=A_ptr)]
+ A_ptr,
+ ty_args=[
+ T.Tensor(1, "float32", layout=T.TileLayout(T.S[1] + (A_offset_redef)), data=A_ptr)
+ ],
)
A[0] = 42.0
@@ -780,7 +786,12 @@ def test_simplify_trivial_let_stride():
def before(A_ptr: T.handle("float32"), A_stride: T.int32):
A_stride_redef = A_stride
A = T.decl_tensor(
- A_ptr, ty_args=[T.Tensor(1, "float32", strides=[A_stride_redef], data=A_ptr)]
+ A_ptr,
+ ty_args=[
+ T.Tensor(
+ 1, "float32", layout=T.TileLayout(T.S[(1,) : [A_stride_redef]]), data=A_ptr
+ )
+ ],
)
A[0] = 42.0
@@ -788,7 +799,12 @@ def test_simplify_trivial_let_stride():
def expected(A_ptr: T.handle("float32"), A_stride: T.int32):
A_stride_redef = A_stride
A = T.decl_tensor(
- A_ptr, ty_args=[T.Tensor(1, "float32", strides=[A_stride_redef], data=A_ptr)]
+ A_ptr,
+ ty_args=[
+ T.Tensor(
+ 1, "float32", layout=T.TileLayout(T.S[(1,) : [A_stride_redef]]), data=A_ptr
+ )
+ ],
)
A[0] = 42.0
diff --git a/tirx/tests/python/tirx/transform/test_storage_rewrite.py b/tirx/tests/python/tirx/transform/test_storage_rewrite.py
index 0797d6efc2..3f85d07473 100644
--- a/tirx/tests/python/tirx/transform/test_storage_rewrite.py
+++ b/tirx/tests/python/tirx/transform/test_storage_rewrite.py
@@ -167,7 +167,14 @@ def test_allocation_constraints_prevent_lifetime_reuse(property_value, stride, o
@T.function
def func(output: T.Tensor((2,), "float32")):
A = T.alloc_tensor(
- ty_args=[T.Tensor((128,), "float32", scope, strides=[stride], elem_offset=offset)]
+ ty_args=[
+ T.Tensor(
+ (128,),
+ "float32",
+ scope,
+ layout=T.TileLayout(T.S[(128,) : [stride]] + (offset)),
+ )
+ ]
)
for i in T.serial(128):
A[i] = T.float32(i)
@@ -184,7 +191,12 @@ def test_allocation_constraints_prevent_lifetime_reuse(property_value, stride, o
B = T.alloc_tensor(ty_args=[T.Tensor((128,), T.float32, scope=scope)])
A = T.alloc_tensor(
ty_args=[
- T.Tensor((128,), T.float32, scope=scope, strides=[stride], elem_offset=offset)
+ T.Tensor(
+ (128,),
+ T.float32,
+ scope=scope,
+ layout=T.TileLayout(T.S[(128,) : [stride]] + (offset)),
+ )
]
)
for i in range(128):
@@ -245,9 +257,7 @@ def test_packing_preserves_different_alignment_contracts():
for i in T.serial(17):
A[i] = T.float32(i)
output[0] = A[16]
- B = T.alloc_tensor(
- ty_args=[T.Tensor((17,), "float32", scope, elem_offset=0, align=256, offset_factor=128)]
- )
+ B = T.alloc_tensor(ty_args=[T.Tensor((17,), "float32", scope, align=512)])
for i in T.serial(17):
B[i] = T.float32(i + 1)
output[1] = B[16]
@@ -265,9 +275,7 @@ def test_packing_preserves_different_alignment_contracts():
output[0] = A_1[16]
A_2 = T.decl_tensor(
A.data,
- ty_args=[
- T.Tensor((17,), T.float32, scope=scope, data_alignment=256, offset_factor=128)
- ],
+ ty_args=[T.Tensor((17,), T.float32, scope=scope, data_alignment=512)],
)
for i_1 in range(17):
A_2[128 + i_1] = T.Cast("float32", i_1 + 1)