This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new dd242c3295 [TIRx] Preserve parameter strides and physical view extents
(#20788)
dd242c3295 is described below
commit dd242c3295279433f232a559ade6f0f428d6966e
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 18 13:20:00 2026 -0400
[TIRx] Preserve parameter strides and physical view extents (#20788)
---
docs/tirx/arch/lowering_pipeline.rst | 5 +-
docs/tirx/native_basics/cuda/tensors.rst | 39 ++++---
.../backend/trn/transform/lower_trainium_layout.cc | 14 ++-
tirx/src/tirx/transform/flatten_tensor.cc | 5 +-
tirx/src/tirx/transform/flatten_tensor_utils.h | 119 +++++++++++++++++++++
tirx/src/tirx/transform/flattened_tensor.h | 77 -------------
tirx/src/tirx/transform/lower_tensor_layouts.cc | 2 +-
tirx/tests/python/tirx/tile/trn/test_copy.py | 34 +++++-
.../python/tirx/transform/test_flatten_tensor.py | 21 +++-
.../tirx/transform/test_lower_tensor_layouts.py | 20 ++--
10 files changed, 218 insertions(+), 118 deletions(-)
diff --git a/docs/tirx/arch/lowering_pipeline.rst
b/docs/tirx/arch/lowering_pipeline.rst
index b449425fa6..44b0369f81 100644
--- a/docs/tirx/arch/lowering_pipeline.rst
+++ b/docs/tirx/arch/lowering_pipeline.rst
@@ -140,8 +140,9 @@ The generic pipeline explicitly composes two independently
usable passes:
finite CUDA index calls inside ordinary ``Bind`` statements.
- **``LowerTensorLayouts``** then runs the ``LayoutApplier``: it resolves every
``TileLayout``-typed tensor access into concrete physical address arithmetic
- (``addr = data + elem_offset + layout.apply(coord)``), flattens the tensors,
- and removes tensor offsets that have been folded into the resulting views.
+ (``addr = data + layout.apply(coord, shape=shape)["m"]``, in elements),
+ flattens the tensors, and accounts for the layout's memory origin exactly
once
+ in the resulting view or its access indices.
A backend can pair ``TileDispatch`` with its own layout lowering, such as
Trainium's ``LowerTrainiumLayout``. The generic pipeline inserts ``PrintIR``
diff --git a/docs/tirx/native_basics/cuda/tensors.rst
b/docs/tirx/native_basics/cuda/tensors.rst
index 8c66d78484..aaba07af26 100644
--- a/docs/tirx/native_basics/cuda/tensors.rst
+++ b/docs/tirx/native_basics/cuda/tensors.rst
@@ -57,11 +57,8 @@ The tensor descriptors share these parameters:
* - ``shape``
- logical shape (a tuple of extents)
* - ``layout``
- - physical mapping (:doc:`TileLayout <../../layout>`); ``"default"`` =
dense
- row-major
- * - ``elem_offset``
- - ``elem_offset`` (or ``byte_offset``) places a *view* at an offset into
- ``data``
+ - physical mapping, including strides and element origin
+ (:doc:`TileLayout <../../layout>`); ``"default"`` = dense row-major
* - ``align``
- alignment of the data pointer, in bytes
@@ -103,21 +100,23 @@ The ``scope`` argument selects the memory space:
**A ptr-based tensor is just metadata over a pointer.** For any non-tmem
tensor,
the declaration is a pointer plus a layout, and indexing resolves to an
address::
- addr(tensor[coord]) = tensor.data + elem_offset + layout.apply(coord,
shape=shape)["m"]
+ addr(tensor[coord]) = tensor.data + layout.apply(coord, shape=shape)["m"]
(``layout.apply`` returns the per-axis mapping; its ``"m"`` component is the
-element offset.) So the *same* logical access compiles to different address
-arithmetic depending purely on the tensor's metadata. Writing
+element offset, including the layout's origin.) Pointer arithmetic here is in
+elements. Strides and origins are represented only by the layout, not by
separate
+tensor arguments or properties. So the *same* logical access compiles to
different
+address arithmetic depending purely on the tensor's metadata. Writing
``B[i, j] = A[i, j] + 1`` over a 4×8 region, with ``B`` annotated four ways in
the function signature:
.. code-block:: python
- from tvm.tirx.layout import TileLayout, S
+ from tvm_tirx.tirx.layout import TileLayout, S
B: Tx.Tensor((4, 8), "float32") # row-major
B: Tx.Tensor((4, 8), "float32", layout=TileLayout(S[(4, 8) : (1, 4)])) #
column-major
- B: Tx.Tensor((4, 8), "float32", elem_offset=64) # shifted view
+ B: Tx.Tensor((4, 8), "float32", layout=TileLayout(S[(4, 8)] + 64)) #
shifted view
B: Tx.Tensor((4, 8), "float32", layout=TileLayout(S[(4, 8) : (16, 1)])) #
row stride 16
each makes ``B[i, j]`` lower to a different index in the generated CUDA (the
@@ -127,9 +126,14 @@ each makes ``B[i, j]`` lower to a different index in the
generated CUDA (the
B_ptr[((i * 8) + j)] = ...; // row-major: i*8 + j
B_ptr[((j * 4) + i)] = ...; // column-major: j*4 + i
- B_ptr[(((i * 8) + j) + 64)] = ...; // elem_offset=64: i*8 + j + 64
+ B_ptr[(((i * 8) + j) + 64)] = ...; // layout origin 64: i*8 + j + 64
B_ptr[((i * 16) + j)] = ...; // row stride 16: i*16 + j
+Layouts must describe the actual tensor domain; lowering does not resize or
+extend mismatched tile extents. For nonzero shapes and non-negative strides,
+allocation footprints cover the maximum mapped memory address. A declaration
+over a borrowed pointer remains responsible for the storage that pointer
addresses.
+
Shared memory
-------------
@@ -176,7 +180,7 @@ Dynamic
``sharedMemBytes`` launch parameter), not at compile time. A kernel may have
**only
one** dynamic-shared allocation — the *arena*. So you allocate it once and
``decl``
each tensor as a view into it: ``Tx.decl_tensor`` with the arena pointer as its
-first argument and an ``elem_offset``:
+first argument and a memory-axis origin in its layout:
.. code-block:: python
@@ -187,7 +191,10 @@ first argument and an ``elem_offset``:
arena.data, ty_args=[Tx.Tensor((64,), "float32", scope="shared.dyn")]
) # offset 0
Bs = Tx.decl_tensor(
- arena.data, ty_args=[Tx.Tensor((64,), "float32", elem_offset=64,
scope="shared.dyn")]
+ arena.data,
+ ty_args=[Tx.Tensor(
+ (64,), "float32", scope="shared.dyn",
layout=Tx.TileLayout(Tx.S[(64,)] + 64)
+ )],
) # offset 64
As[tx] = A[tx]
Bs[tx] = B[tx]
@@ -404,9 +411,9 @@ or hand you a pointer — they emit no runtime op of their
own.
through an expression: ``B.view(...)`` binds the operand exactly as
``B.ty.view(B, ...)``. ``TensorType.__expr_properties__`` maps property names
to getters taking the type and original expression. It preserves conveniences
-such as ``B.shape``, ``B.strides``, ``B.dtype``, ``B.byte_offset``, ``B.sub``
and
-``B.data`` without modifying the shared ``Var`` class. ``B.dtype`` returns a
-runtime ``DataType``; ``B.ty.dtype`` stores a ``PrimType``.
+such as ``B.shape``, ``B.dtype``, ``B.layout``, ``B.sub`` and ``B.data``
without
+modifying the shared ``Var`` class. ``B.dtype`` returns a runtime ``DataType``;
+``B.ty.dtype`` stores a ``PrimType``.
Existing expression attributes and reflected fields take precedence, followed
by declared properties and then declared methods. Declared properties are
diff --git a/tirx/src/backend/trn/transform/lower_trainium_layout.cc
b/tirx/src/backend/trn/transform/lower_trainium_layout.cc
index 896b86ad5d..ef2739e67d 100644
--- a/tirx/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/tirx/src/backend/trn/transform/lower_trainium_layout.cc
@@ -39,7 +39,7 @@
#include <utility>
#include <vector>
-#include "../../../tirx/transform/flattened_tensor.h"
+#include "../../../tirx/transform/flatten_tensor_utils.h"
namespace tvm {
namespace tirx {
@@ -73,9 +73,13 @@ class TrainiumLayoutApplier : public
tirx::StmtExprMutatorWithAnalyzer {
TensorVar flattened =
storage_lower->GetFlattenedTensor(tensor.value());
auto type =
ffi::make_object<TensorTypeNode>(*tensor.value().ty().get());
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();
+ // Trainium's physical axes are embedded in the lowered indices.
Ordinary
+ // memory strides still define the parameter ABI, even when they are
static.
+ // Retain symbolic physical geometry that may define parameter
variables.
+ bool static_geometry =
+ IsTrainiumLayout(tile) && tile->replica.empty() &&
+ std::any_of(tile->shard.begin(), tile->shard.end(),
+ [](const Iter& iter) { return
iter->axis->IsMemoryAxis(); });
for (const auto& iter : tile->shard) {
static_geometry &= iter->extent.as<IntImmNode>() &&
iter->stride.as<IntImmNode>();
}
@@ -211,7 +215,7 @@ class TrainiumLayoutApplier : public
tirx::StmtExprMutatorWithAnalyzer {
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
}
} else {
- flattened = FlattenedTensor(buf, ffi::GetRef<sym::Analyzer>(analyzer_));
+ flattened = FlattenedTensor(buf, ffi::GetRef<sym::Analyzer>(analyzer_),
/*fold_origin=*/true);
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
}
if (flattened.dtype()->dtype == DLDataType{kDLBool, 8, 1}) {
diff --git a/tirx/src/tirx/transform/flatten_tensor.cc
b/tirx/src/tirx/transform/flatten_tensor.cc
index bf944975c5..65db1f2ded 100644
--- a/tirx/src/tirx/transform/flatten_tensor.cc
+++ b/tirx/src/tirx/transform/flatten_tensor.cc
@@ -34,7 +34,7 @@
#include <unordered_set>
-#include "flattened_tensor.h"
+#include "flatten_tensor_utils.h"
#include "ir_utils.h"
namespace tvm {
@@ -134,7 +134,8 @@ class TensorFlattener : public StmtExprMutatorWithAnalyzer {
// buf': the storage husk. The linearized indices carry layout and
// elem_offset, so the husk keeps neither.
- auto flat = FlattenedTensor(fold_view,
ffi::GetRef<sym::Analyzer>(analyzer_));
+ auto flat = FlattenedTensor(fold_view,
ffi::GetRef<sym::Analyzer>(analyzer_),
+ /*fold_origin=*/true);
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]));
diff --git a/tirx/src/tirx/transform/flatten_tensor_utils.h
b/tirx/src/tirx/transform/flatten_tensor_utils.h
new file mode 100644
index 0000000000..5e001b100d
--- /dev/null
+++ b/tirx/src/tirx/transform/flatten_tensor_utils.h
@@ -0,0 +1,119 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file flatten_tensor_utils.h
+ * \brief Physical tensor view construction for layout lowering.
+ */
+#ifndef TVM_TIRX_TRANSFORM_FLATTEN_TENSOR_UTILS_H_
+#define TVM_TIRX_TRANSFORM_FLATTEN_TENSOR_UTILS_H_
+
+#include <tvm/ir/prim/op.h>
+#include <tvm/tirx/expr.h>
+#include <tvm/tirx/layout.h>
+
+namespace tvm {
+namespace tirx {
+
+/*!
+ * \brief Construct a one-dimensional physical tensor variable without
allocating storage.
+ *
+ * The result preserves name, dtype, scope, alignment, and source location,
and has
+ * no layout. Its shape covers the physical footprint, including stride gaps
and
+ * conservative padding where needed. The extent is a required backing-storage
+ * bound in elements, not bytes or necessarily the minimum storage size.
General
+ * layouts include memory-axis offsets and replicas, and exclude
hardware-thread
+ * axes to give per-thread storage. Swizzles round up to blocks covering the
bits
+ * they may change; a bare swizzle uses the tensor shape as its domain.
Without a
+ * layout, the shape is the logical element count.
+ *
+ * Callers bind a changed variable to a decl_tensor over the source pointer
for a
+ * borrowed view, and preserve any relative origin in that binding's geometry.
+ * Allocation callers apply their own per-thread/backend storage contract; this
+ * helper neither allocates nor rejects borrowed strided or offset views.
+ *
+ * \param tensor Tensor to flatten, assuming valid nonnegative strides and
nonzero
+ * logical/tile extents, with any layout describing the actual tensor domain.
+ * Mismatched tile domains are not resized or extended.
+ * \param analyzer Optional analyzer; a fresh analyzer is used when absent.
+ * \param fold_origin For shape-aligned regular strides, include the memory
origin
+ * in the extent for absolute flat indexing. Otherwise the caller retains that
+ * origin separately. General layouts already use an absolute bound regardless
+ * of this flag.
+ * \return Fresh physical tensor variable, or tensor itself if it is already
+ * one-dimensional with no layout.
+ */
+TVM_FFI_INLINE TensorVar FlattenedTensor(const TensorVar& tensor,
+ ffi::Optional<sym::Analyzer> analyzer
= std::nullopt,
+ bool fold_origin = false) {
+ if (!analyzer.has_value()) analyzer = sym::Analyzer();
+ const auto& ana = analyzer.value();
+ ffi::Array<PrimExpr> output_shape{1};
+ const auto* tile = tensor.ty()->layout.as<TileLayoutNode>();
+ bool explicit_strides = tile &&
MatchShapeStridePattern(tensor.ty()).has_value();
+ if (explicit_strides) {
+ // Preserve the outer stride footprint and cover every addressed element.
+ PrimExpr span = 1;
+ for (const auto& iter : tile->shard) {
+ span = span + (iter->extent - 1) * iter->stride;
+ }
+ if (tile->IsContiguous(ana)) {
+ 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));
+ }
+ if (fold_origin) {
+ // Absolute flat indices include the origin.
+ static const Axis memory_axis = Axis::Get("m");
+ if (auto origin = tile->offset.Get(memory_axis)) {
+ output_shape.Set(0, output_shape[0] + origin.value());
+ }
+ }
+ } else if (tile) {
+ output_shape.Set(0, ana->Simplify(tile->GetSpan(ffi::String("m"))));
+ } else if (const auto* compose =
tensor.ty()->layout.as<ComposeLayoutNode>()) {
+ PrimExpr span = compose->tile_layout->GetSpan(ffi::String("m"));
+ if (compose->tile_layout->IsTrivial()) {
+ span = IntImm(tensor.ty()->DefaultIndexType(), 1);
+ for (const auto& extent : tensor.shape()) span = span * extent;
+ }
+ if (compose->swizzle_len != 0) {
+ int changed_bits = compose->per_element + compose->swizzle_len;
+ if (!compose->swizzle_inner) changed_bits += compose->atom_len;
+ PrimExpr block = IntImm(span.ty(), int64_t{1} << changed_bits);
+ span = floordiv(span + block - 1, block) * block;
+ }
+ output_shape.Set(0, ana->Simplify(span));
+ } else {
+ for (const auto& extent : tensor.shape()) output_shape.Set(0,
output_shape[0] * extent);
+ }
+
+ if (output_shape.size() == tensor.shape().size() &&
!tensor.ty()->layout.has_value()) {
+ return tensor;
+ }
+ return TensorVar(tensor->name,
+ TensorType(output_shape, tensor.dtype(), tensor.scope(),
+ tensor.ty()->data_alignment, std::nullopt),
+ tensor->loc);
+}
+
+} // namespace tirx
+} // namespace tvm
+
+#endif // TVM_TIRX_TRANSFORM_FLATTEN_TENSOR_UTILS_H_
diff --git a/tirx/src/tirx/transform/flattened_tensor.h
b/tirx/src/tirx/transform/flattened_tensor.h
deleted file mode 100644
index e71669a279..0000000000
--- a/tirx/src/tirx/transform/flattened_tensor.h
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file flattened_tensor.h
- * \brief Physical tensor view construction for layout lowering.
- */
-#ifndef TVM_TIRX_TRANSFORM_FLATTENED_TENSOR_H_
-#define TVM_TIRX_TRANSFORM_FLATTENED_TENSOR_H_
-
-#include <tvm/ir/prim/op.h>
-#include <tvm/tirx/expr.h>
-#include <tvm/tirx/layout.h>
-
-namespace tvm {
-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,
- ffi::Optional<sym::Analyzer> analyzer
= std::nullopt) {
- ffi::Array<PrimExpr> output_shape{1};
- 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 (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));
- }
- } 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].
- for (size_t i = 0; i < tensor.shape().size(); i++) {
- output_shape.Set(0, output_shape[0] * tensor.shape()[i]);
- }
- }
-
- 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.ty()->data_alignment, std::nullopt),
- tensor->loc);
- }
-}
-
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_TIRX_TRANSFORM_FLATTENED_TENSOR_H_
diff --git a/tirx/src/tirx/transform/lower_tensor_layouts.cc
b/tirx/src/tirx/transform/lower_tensor_layouts.cc
index a3a861ec2a..8f55b07f57 100644
--- a/tirx/src/tirx/transform/lower_tensor_layouts.cc
+++ b/tirx/src/tirx/transform/lower_tensor_layouts.cc
@@ -41,7 +41,7 @@
#include <utility>
#include <vector>
-#include "flattened_tensor.h"
+#include "flatten_tensor_utils.h"
#include "ir_utils.h"
namespace tvm {
diff --git a/tirx/tests/python/tirx/tile/trn/test_copy.py
b/tirx/tests/python/tirx/tile/trn/test_copy.py
index 4e82e44e0f..996f45f04f 100644
--- a/tirx/tests/python/tirx/tile/trn/test_copy.py
+++ b/tirx/tests/python/tirx/tile/trn/test_copy.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
+import numpy as np
import pytest
import tvm_ffi
@@ -25,6 +26,7 @@ from tvm_tirx.backend.trn.tile.instruction_generator import
InstructionGenerator
from tvm_tirx.ir import assert_structural_equal as _assert_structural_equal
from tvm_tirx.script import tirx as T
from tvm_tirx.sym import Analyzer
+from tvm_tirx.testing import env
from tvm_tirx.tirx import TensorType, Var
from tvm_tirx.tirx.layout import S, TileLayout
@@ -61,6 +63,32 @@ def test_instruction_generator_compact_global_layout():
assert tensor.layout is None
[email protected](not env.has_llvm(), reason="need llvm")
+def test_parameter_stride_contract():
+ """Physical index lowering must preserve the host's DLTensor stride
checks."""
+ layout = TileLayout(S[(2, 2) : (3, 1)])
+
+ @T.function
+ def copy(
+ A: T.Tensor((2, 2), "float32", layout=layout, data_alignment=1),
+ B: T.Tensor((2, 2), "float32"),
+ ):
+ for i, j in T.grid(2, 2):
+ B[i, j] = A[i, j]
+
+ mod = tvm_tirx.IRModule({"main": copy})
+ lowered = tvm_tirx.backend.trn.transform.LowerTrainiumLayout()(mod)
+ # Execute the lowered scalar accesses and packed wrapper on the host,
without
+ # requiring a Trainium device or compiling Trainium-specific instructions.
+ lib = tvm_tirx.tirx.compile(lowered, target="llvm")
+ data = np.arange(6, dtype="float32").reshape(2, 3)[:, :2]
+ result = tvm_tirx.runtime.empty((2, 2), "float32")
+ lib(data, result)
+ np.testing.assert_array_equal(result.numpy(), data)
+ with pytest.raises(ValueError, match="strides"):
+ lib(np.ascontiguousarray(data), result)
+
+
def test_simple_copy():
src_shape = [128, 512]
src_layout = T.TileLayout(T.S[(128, 512) : (512, 1)])
@@ -348,7 +376,7 @@ def test_copy_different_shape_dim():
T.trn.tile.load(A_sbuf, A[i, :, :])
@T.function
- def expected(A: T.Tensor((32, 128, 512), "float32", layout=None)):
+ def expected(A: T.Tensor((32, 128, 512), "float32", layout=src_layout)):
T.func_attr({"global_symbol": "copy"})
A_1 = T.decl_tensor(
@@ -500,7 +528,7 @@ def test_copy_with_complex_index():
T.trn.tile.load(A_sbuf[1, 0:2048, 0:1024], A[2048: 4096, 3072:4096])
@T.function
- def expected(A: T.Tensor((4096, 4096), "float32", layout=None)):
+ def expected(A: T.Tensor((4096, 4096), "float32", layout=A_layout)):
T.func_attr({"global_symbol": "copy"})
A_1 = T.decl_tensor(
@@ -536,7 +564,7 @@ def test_copy_with_complex_index_2():
T.trn.tile.load(A_sbuf[2048: 4096, 3072:4096], A[1, 0:2048, 0:1024])
@T.function
- def expected(A: T.Tensor((2, 2048, 1024), "float32", layout=None)):
+ def expected(A: T.Tensor((2, 2048, 1024), "float32", layout=A_layout)):
T.func_attr({"global_symbol": "copy"})
A_1 = T.decl_tensor(
diff --git a/tirx/tests/python/tirx/transform/test_flatten_tensor.py
b/tirx/tests/python/tirx/transform/test_flatten_tensor.py
index 266f59c759..71aad5dbe2 100644
--- a/tirx/tests/python/tirx/transform/test_flatten_tensor.py
+++ b/tirx/tests/python/tirx/transform/test_flatten_tensor.py
@@ -516,7 +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", scope="shared",
layout=None)],
+ ty_args=[T.Tensor((n[0] + 1,), "uint64", scope="shared",
layout=None)],
)
mbar[n[0]] = T.uint64(1)
@@ -563,7 +563,7 @@ def test_tensor_flatten_offset():
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)])
+ flat = T.decl_tensor(A.data, ty_args=[T.Tensor((m * n + 100,),
"float32", layout=None)])
T.evaluate(flat[n * 2 + 103])
After = FlattenTensor()(Before)
@@ -702,5 +702,22 @@ def test_tensor_index_merge_rank_four(factored):
tvm_tirx.ir.assert_structural_equal(After, Expected)
+def test_split_layout_physical_storage_extent():
+ layout = T.TileLayout(T.S[(3, 4) : (16, 2)])
+
+ @T.function(private=True)
+ def main(data: T.handle):
+ B = T.decl_tensor(data, ty_args=[T.Tensor((12,), "float32",
layout=layout)])
+ B[11] = T.float32(7)
+
+ # A borrowed view's flattened extent must cover the physical gaps.
+ after = FlattenTensor()(tvm_tirx.IRModule({"main": main}))
+ body = after["main"].body
+ flat = body[0].var
+ assert int(flat.ty.shape[0]) == 39
+ assert int(body[-1].indices[0]) == 38
+ assert flat.ty.layout is None
+
+
if __name__ == "__main__":
tvm_tirx.testing.main()
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 538f90c894..adc2526ea9 100644
--- a/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
+++ b/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
@@ -88,7 +88,7 @@ def test_lower_view_get():
lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
A = T.alloc_local((2,), "float16", layout=None)
B = T.decl_tensor(
- A.data, ty_args=[T.Tensor((64,), "float16", scope="local",
layout=None)]
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
)
A_local = T.decl_tensor(
A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
@@ -96,7 +96,7 @@ def test_lower_view_get():
for i in T.vectorized(2):
A_local[i] = T.Cast("float16",
in_buf_1[T.cuda.thread_idx("x") * 2 + i])
B_1 = T.decl_tensor(
- A.data, ty_args=[T.Tensor((64,), "float16", scope="local",
layout=None)]
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
)
A_local_1 = T.decl_tensor(
A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
@@ -260,7 +260,7 @@ def test_lower_view_get():
lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
A = T.alloc_local((8,), layout=None)
B = T.decl_tensor(
- A.data, ty_args=[T.Tensor((256,), "float32",
scope="local", layout=None)]
+ A.data, ty_args=[T.Tensor((8,), "float32", scope="local",
layout=None)]
)
A_local = T.decl_tensor(
A.data, ty_args=[T.Tensor((8,), "float32", scope="local",
layout=None)]
@@ -275,7 +275,7 @@ def test_lower_view_get():
+ T.cuda.thread_idx("x") % 4
]
B_1 = T.decl_tensor(
- A.data, ty_args=[T.Tensor((256,), "float32",
scope="local", layout=None)]
+ A.data, ty_args=[T.Tensor((8,), "float32", scope="local",
layout=None)]
)
A_local_1 = T.decl_tensor(
A.data, ty_args=[T.Tensor((8,), "float32", scope="local",
layout=None)]
@@ -477,7 +477,7 @@ def test_lower_view_get():
lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
acc = T.alloc_local((64,), layout=None)
B = T.decl_tensor(
- acc.data, ty_args=[T.Tensor((16384,), "float32",
scope="local", layout=None)]
+ acc.data, ty_args=[T.Tensor((64,), "float32",
scope="local", layout=None)]
)
acc_local = T.decl_tensor(
acc.data, ty_args=[T.Tensor((64,), "float32",
scope="local", layout=None)]
@@ -494,7 +494,7 @@ def test_lower_view_get():
+ vec
]
B_1 = T.decl_tensor(
- acc.data, ty_args=[T.Tensor((16384,), "float32",
scope="local", layout=None)]
+ acc.data, ty_args=[T.Tensor((64,), "float32",
scope="local", layout=None)]
)
acc_local_1 = T.decl_tensor(
acc.data, ty_args=[T.Tensor((64,), "float32",
scope="local", layout=None)]
@@ -700,10 +700,10 @@ def test_lower_view_get():
lane_id: T.let[T.int32] = T.cuda.thread_idx("x") % 32
A = T.alloc_local((2,), "float16", layout=None)
B = T.decl_tensor(
- A.data, ty_args=[T.Tensor((64,), "float16", scope="local",
layout=None)]
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
)
B_1 = T.decl_tensor(
- A.data, ty_args=[T.Tensor((64,), "float16", scope="local",
layout=None)]
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
)
A_local = T.decl_tensor(
A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
@@ -714,10 +714,10 @@ def test_lower_view_get():
)
A_local_1[1] = T.Cast("float16",
in_buf_1[T.cuda.thread_idx("x") * 2 + 1])
B_2 = T.decl_tensor(
- A.data, ty_args=[T.Tensor((64,), "float16", scope="local",
layout=None)]
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
)
B_3 = T.decl_tensor(
- A.data, ty_args=[T.Tensor((64,), "float16", scope="local",
layout=None)]
+ A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
)
A_local_2 = T.decl_tensor(
A.data, ty_args=[T.Tensor((2,), "float16", scope="local",
layout=None)]
