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-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 27404a5eb [FIX] Remap logical tensor geometry before offset
linearization (#47)
27404a5eb is described below
commit 27404a5eba659db7a0d0bd10c1d27f1360860613
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Oct 24 18:30:12 2026 -0400
[FIX] Remap logical tensor geometry before offset linearization (#47)
LowerTensorLayouts can retain references to an original tensor when
another tensor uses its loads in shape-derived offsets. Remap the
logical tensor type before applying its layout or computing element
offsets so shape and layout expressions use the rebuilt bindings.
Preserve the original logical rank and already-lowered access indices.
---
src/tirx/transform/lower_tensor_layouts.cc | 21 ++++++++--------
.../tirx/transform/test_lower_tensor_layouts.py | 28 ++++++++++++++++++++++
2 files changed, 39 insertions(+), 10 deletions(-)
diff --git a/src/tirx/transform/lower_tensor_layouts.cc
b/src/tirx/transform/lower_tensor_layouts.cc
index d06164a14..c65638175 100644
--- a/src/tirx/transform/lower_tensor_layouts.cc
+++ b/src/tirx/transform/lower_tensor_layouts.cc
@@ -301,17 +301,18 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
ffi::Array<PrimExpr> GetSimplifiedElemOffset(const TensorVar& tensor,
const ffi::Array<PrimExpr>&
indices) {
- if (tensor.ty()->layout.has_value()) {
- if (auto tile = tensor.ty()->layout.value().as<TileLayoutObj>();
+ // Remap the ty first since the type may depend on previously remapped
value in mutator.
+ auto logical_type = StmtExprMutator::Mutate(ffi::AnyView(tensor.ty()),
InplaceMode::kDisallow)
+ .ValueOrUnchanged(tensor.ty())
+ .as_or_throw<TensorType>();
+ if (logical_type->layout.has_value()) {
+ if (auto tile = logical_type->layout.value().as<TileLayoutObj>();
tile && ffi::GetRef<TileLayout>(tile).HasThreadAxis()) {
LOG(FATAL) << "Cannot lower direct TensorLoad/TensorStore on a tensor
with thread-axis "
<< "layout: unable to verify that the coordinate matches
the current thread. "
<< "Use .view() + .local() to decompose thread and memory
axes.";
}
- auto layout = tensor.ty()->layout.value();
- layout = StmtExprMutator::Mutate(ffi::AnyView(layout),
InplaceMode::kDisallow)
- .ValueOrUnchanged(layout)
- .as_or_throw<Layout>();
+ auto layout = logical_type->layout.value();
auto flat = GetFlattenedTensor(tensor);
auto pattern = MatchShapeStridePattern(flat.ty()).value();
if (!prim::IsZero(pattern.get<2>())) {
@@ -319,14 +320,14 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
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)};
+ auto res = layout.Canonicalize().Apply(indices, logical_type->shape);
+ if (res.empty() && logical_type->shape.empty()) {
+ return {IntImm(logical_type->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);
+ auto flattened_indices = logical_type->ElemOffset(indices);
TVM_FFI_ICHECK_EQ(flattened_indices.size(), 1) << "Expected a single
element offset";
return {analyzer_->Simplify(flattened_indices[0])};
}
diff --git a/tests/python/tirx/transform/test_lower_tensor_layouts.py
b/tests/python/tirx/transform/test_lower_tensor_layouts.py
index 91b470844..08ac5a11f 100644
--- a/tests/python/tirx/transform/test_lower_tensor_layouts.py
+++ b/tests/python/tirx/transform/test_lower_tensor_layouts.py
@@ -1536,3 +1536,31 @@ def test_alloc_tensor_swizzle_storage_span(scope):
with tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_100a"}):
After = LowerTensorLayouts()(Before)
tvm_tirx.ir.assert_structural_equal(After, Expected)
+
+
+def test_shape_tensor_load_is_remapped_in_flattened_offset():
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(out: T.Tensor((1,), "float32")):
+ d = T.alloc_local((1,), "int32", layout="default")
+ d[0] = 1
+ b = T.alloc_local((d[0] + 1, d[0] + 1), "float32", layout=None)
+ b[d[0], d[0]] = T.float32(1)
+ out[0] = b[d[0], d[0]]
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(out: T.Tensor((1,), "float32")):
+ d = T.alloc_local((1,), "int32", layout=None)
+ d[0] = 1
+ b = T.alloc_local(((d[0] + 1) * (d[0] + 1),), "float32",
layout=None)
+ b[(d[0] + 2) * d[0]] = T.float32(1)
+ out[0] = b[(d[0] + 2) * d[0]]
+
+ with tvm_tirx.target.Target("llvm"):
+ after = LowerTensorLayouts()(Before)
+ assert not tvm_tirx.tirx.analysis.undefined_vars(after["main"].body,
after["main"].params)
+ tvm_tirx.ir.assert_structural_equal(after, Expected)
+ tvm_tirx.tirx.transform.FlattenTensor()(after)
