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 e9c49853 [REFACTOR] Keep layout and document data nodes nonvirtual (#8)
e9c49853 is described below
commit e9c49853ca72a06759ca7dc64db923de12a21c14
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 20 20:52:15 2026 -0400
[REFACTOR] Keep layout and document data nodes nonvirtual (#8)
Keep layout and document data nodes nonvirtual. Expose layout operations
as exported nonvirtual reference-wrapper methods with explicit FFI
runtime-type dispatch and source-private concrete helpers, preserving
mapping behavior and Python interfaces.
Keep LayoutNode, TileLayoutNode, and ComposeLayoutNode data-only, and
use concrete FFI destruction for document nodes.
---
include/tvm/script/printer/doc.h | 3 -
include/tvm/tirx/layout.h | 100 ++++++--------------
src/backend/trn/transform/lower_trainium_layout.cc | 20 ++--
src/tirx/ir/layout/compose_layout.cc | 77 ++++++++-------
src/tirx/ir/layout/layout.cc | 78 +++++++++++++---
src/tirx/ir/layout/tile_canonicalize.cc | 8 +-
src/tirx/ir/layout/tile_core.cc | 104 +++++++++++----------
src/tirx/ir/layout/tile_internal.h | 28 ++++++
src/tirx/ir/layout/tile_tile_ops.cc | 11 ++-
src/tirx/ir/layout/utils.cc | 2 +-
src/tirx/ir/tensor.cc | 2 +-
src/tirx/ir/type.cc | 4 +-
src/tirx/transform/flatten_tensor_utils.h | 8 +-
src/tirx/transform/lower_tensor_layouts.cc | 8 +-
14 files changed, 249 insertions(+), 204 deletions(-)
diff --git a/include/tvm/script/printer/doc.h b/include/tvm/script/printer/doc.h
index 6441b40b..16aff11d 100644
--- a/include/tvm/script/printer/doc.h
+++ b/include/tvm/script/printer/doc.h
@@ -68,9 +68,6 @@ class DocNode : public ffi::Object {
static constexpr bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("script.printer.Doc", DocNode, ffi::Object);
-
- public:
- virtual ~DocNode() = default;
};
/*!
diff --git a/include/tvm/tirx/layout.h b/include/tvm/tirx/layout.h
index 7c8baab9..7c2527ce 100644
--- a/include/tvm/tirx/layout.h
+++ b/include/tvm/tirx/layout.h
@@ -46,35 +46,35 @@ using ffi::Tuple;
// Base class for layout
class LayoutNode : public ffi::Object {
+ public:
+ static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
+ TVM_FFI_DECLARE_OBJECT_INFO("tirx.Layout", LayoutNode, ffi::Object);
+};
+
+class Layout : public ffi::ObjectRef {
public:
/*! \brief Verify if the layout is well-formed */
- virtual void CheckWellFormed() const = 0;
+ TVM_TIRX_DLL void CheckWellFormed() const;
/*! \brief Get the size of the layout (of some axis) */
- virtual PrimExpr GetSize(ffi::Optional<ffi::String> axis_name =
std::nullopt) const = 0;
+ TVM_TIRX_DLL PrimExpr GetSize(ffi::Optional<ffi::String> axis_name =
std::nullopt) const;
/*! \brief Get the span of the layout (of some axis) */
- virtual PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name =
std::nullopt) const = 0;
+ TVM_TIRX_DLL PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name =
std::nullopt) const;
/*! \brief Apply layout on the input coordinate and get the mapped output */
- virtual ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord)
const = 0;
- virtual ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const = 0;
- virtual ffi::Map<ffi::String, PrimExpr> Apply(const ffi::Array<PrimExpr>&
coord,
- const ffi::Array<PrimExpr>&
shape) const;
+ TVM_TIRX_DLL ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr>
coord) const;
+ TVM_TIRX_DLL ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const;
+ TVM_TIRX_DLL ffi::Map<ffi::String, PrimExpr> Apply(const
ffi::Array<PrimExpr>& coord,
+ const
ffi::Array<PrimExpr>& shape) const;
/*! \brief Turn the layout to canonical form */
- virtual Layout Canonicalize() const = 0;
+ TVM_TIRX_DLL Layout Canonicalize() const;
/*! \brief Tile the current layout with a given layout */
- virtual Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>&
outer_shape,
- const ffi::Array<PrimExpr>& inner_shape) const = 0;
+ TVM_TIRX_DLL Layout Tile(const TileLayout& outer, const
ffi::Array<PrimExpr>& outer_shape,
+ const ffi::Array<PrimExpr>& inner_shape) const;
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO("tirx.Layout", LayoutNode, ffi::Object);
-};
-
-class Layout : public ffi::ObjectRef {
- public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Layout, ffi::ObjectRef,
LayoutNode);
};
@@ -163,34 +163,16 @@ class TileLayoutNode : public LayoutNode {
.def_ro("offset", &TileLayoutNode::offset);
}
- /*! \brief Verify if the layout is well-formed */
- void CheckWellFormed() const final;
-
- /*! \brief Get the size of the layout (of some axis) */
- PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const
final;
-
- /*! \brief Get the span of the layout (of some axis) */
- PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const
final;
-
- /*! \brief Apply the input coordinate and get the mapped output */
- ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const
final;
- ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const final;
- /*! \brief Group-first override: if this layout can be regrouped by
``shape``,
- * split each ``coord[d]`` against its group's local extents (cleaner
- * symbolic form than flatten+split-against-shard-shape). Otherwise fall
- * back to flatten+split. */
- ffi::Map<ffi::String, PrimExpr> Apply(const ffi::Array<PrimExpr>& coord,
- const ffi::Array<PrimExpr>& shape)
const final;
-
- /*! \brief Turn the layout to canonical form */
- Layout Canonicalize() const final;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TileLayout", TileLayoutNode,
LayoutNode);
+};
- /*! \brief Tile the layout with an outer layout */
- Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
- const ffi::Array<PrimExpr>& inner_shape) const final;
+class TileLayout : public Layout {
+ public:
+ TVM_TIRX_DLL explicit TileLayout(ffi::Array<Iter> shard, ffi::Array<Iter>
replica,
+ ffi::Map<Axis, PrimExpr> offset);
/*! \brief Get the shape of the shard */
- ffi::Array<PrimExpr> GetShardShape() const;
+ TVM_TIRX_DLL ffi::Array<PrimExpr> GetShardShape() const;
/*!
* \brief Whether the tile describes contiguous row-major memory strides.
@@ -202,25 +184,17 @@ class TileLayoutNode : public LayoutNode {
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;
+ TVM_TIRX_DLL bool IsTrivial() const;
/*! \brief Has Thread Axis */
- bool HasThreadAxis() const;
+ TVM_TIRX_DLL bool HasThreadAxis() const;
/*! \brief Get the scope pair of the layout */
- ffi::Optional<Tuple<ThreadScope, ThreadScope>> GetScope() const;
+ TVM_TIRX_DLL ffi::Optional<Tuple<ThreadScope, ThreadScope>> GetScope() const;
/*! \brief Get the default layout for the shape */
TVM_TIRX_DLL static TileLayout DefaultLayout(ffi::Array<PrimExpr> shape);
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TileLayout", TileLayoutNode,
LayoutNode);
-};
-
-class TileLayout : public Layout {
- public:
- TVM_TIRX_DLL explicit TileLayout(ffi::Array<Iter> shard, ffi::Array<Iter>
replica,
- ffi::Map<Axis, PrimExpr> offset);
-
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TileLayout, Layout,
TileLayoutNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(TileLayoutNode);
};
@@ -253,28 +227,6 @@ class ComposeLayoutNode : public LayoutNode {
.def_ro("tile_layout", &ComposeLayoutNode::tile_layout);
}
- /*! \brief Verify if the layout is well-formed */
- void CheckWellFormed() const final;
-
- /*! \brief Get the size (of some axis) of the layout */
- PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const
final;
-
- /*! \brief Get the span (of some axis) of the layout */
- PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const
final;
-
- /*! \brief Apply the input coordinate and get the mapped output */
- ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const
final;
- ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const final;
- ffi::Map<ffi::String, PrimExpr> Apply(const ffi::Array<PrimExpr>& coord,
- const ffi::Array<PrimExpr>& shape)
const final;
-
- /*! \brief Turn the layout to canonical form */
- Layout Canonicalize() const final;
-
- /*! \brief Tile the layout with an outer layout */
- Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
- const ffi::Array<PrimExpr>& inner_shape) const final;
-
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ComposeLayout", ComposeLayoutNode,
LayoutNode);
};
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc
b/src/backend/trn/transform/lower_trainium_layout.cc
index ef2739e6..9166d4a5 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -175,20 +175,20 @@ class TrainiumLayoutApplier : public
tirx::StmtExprMutatorWithAnalyzer {
TensorVar flattened = buf;
ffi::ObjectPtr<TensorTypeNode> type;
if (IsTrainiumLayout(trn_layout)) {
+ const Layout& layout = buf.ty()->layout.value();
ffi::Array<PrimExpr> new_shape =
- buf.scope() == scope_trn_psum
- ? ffi::Array<PrimExpr>{trn_layout->GetSpan(ffi::String("Bank")),
- trn_layout->GetSize(ffi::String("P")),
- trn_layout->GetSpan(ffi::String("F"))}
- : ffi::Array<PrimExpr>{trn_layout->GetSize(ffi::String("P")),
- trn_layout->GetSpan(ffi::String("F"))};
+ buf.scope() == scope_trn_psum ?
ffi::Array<PrimExpr>{layout.GetSpan(ffi::String("Bank")),
+
layout.GetSize(ffi::String("P")),
+
layout.GetSpan(ffi::String("F"))}
+ :
ffi::Array<PrimExpr>{layout.GetSize(ffi::String("P")),
+
layout.GetSpan(ffi::String("F"))};
flattened = buf;
type = ffi::make_object<TensorTypeNode>(*flattened.ty().get());
type->shape = new_shape;
type->layout = std::nullopt;
} else if (is_alloc) {
if (auto tile_layout = buf.ty()->layout.as<TileLayoutNode>();
- tile_layout && tile_layout->HasThreadAxis()) {
+ tile_layout && ffi::GetRef<TileLayout>(tile_layout).HasThreadAxis())
{
sym::Analyzer ana;
PrimExpr mem_span = IntImm::Int32(1);
for (const auto& iter : tile_layout->shard) {
@@ -281,7 +281,7 @@ class TrainiumLayoutApplier : public
tirx::StmtExprMutatorWithAnalyzer {
.as_or_throw<Layout>();
auto tile_layout = layout.as<TileLayoutNode>();
if (IsTrainiumLayout(tile_layout)) {
- auto coord = layout->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"}
@@ -295,12 +295,12 @@ class TrainiumLayoutApplier : public
tirx::StmtExprMutatorWithAnalyzer {
}
return res;
}
- if (tile_layout && tile_layout->HasThreadAxis()) {
+ if (tile_layout && ffi::GetRef<TileLayout>(tile_layout).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 res = layout->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)};
}
diff --git a/src/tirx/ir/layout/compose_layout.cc
b/src/tirx/ir/layout/compose_layout.cc
index b09095aa..1f91214f 100644
--- a/src/tirx/ir/layout/compose_layout.cc
+++ b/src/tirx/ir/layout/compose_layout.cc
@@ -153,7 +153,7 @@ ffi::Map<ffi::String, PrimExpr> ApplyStructured(const
ComposeLayoutNode* layout,
}
}
- auto mapped = tile->Apply(coord);
+ auto mapped = tile.Apply(coord);
TVM_FFI_ICHECK(mapped.size() == 1 && mapped.find("m") != mapped.end());
PrimExpr m = mapped["m"];
auto fallback = [&]() {
@@ -308,7 +308,7 @@ ComposeLayout::ComposeLayout(int per_element, int
swizzle_len, int atom_len, Til
n->swizzle_len = swizzle_len;
n->atom_len = atom_len;
n->swizzle_inner = swizzle_inner;
- n->CheckWellFormed();
+ layout_details::CheckWellFormed(n.get());
int swizzle_mask = (1 << swizzle_len) - 1;
n->inner_mask = swizzle_mask;
n->outer_mask = swizzle_mask << atom_len;
@@ -323,56 +323,64 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-void ComposeLayoutNode::CheckWellFormed() const {
- TVM_FFI_CHECK(per_element >= 0 && swizzle_len >= 0 && atom_len >=
swizzle_len, ValueError)
+void layout_details::CheckWellFormed(const ComposeLayoutNode* layout) {
+ TVM_FFI_CHECK(layout->per_element >= 0 && layout->swizzle_len >= 0 &&
+ layout->atom_len >= layout->swizzle_len,
+ ValueError)
<< "A compose layout requires nonnegative per_element and swizzle_len,
and "
<< "atom_len >= swizzle_len";
- tile_layout->CheckWellFormed();
+ layout->tile_layout.CheckWellFormed();
}
-PrimExpr ComposeLayoutNode::GetSize(ffi::Optional<ffi::String> axis_name)
const {
+PrimExpr layout_details::GetSize(const ComposeLayoutNode* layout,
+ ffi::Optional<ffi::String> axis_name) {
TVM_FFI_ICHECK(!axis_name.has_value())
<< "ValueError: axis_name is not supported for compose layout";
- return tile_layout->GetSize(axis_name);
+ return layout->tile_layout.GetSize(axis_name);
}
-PrimExpr ComposeLayoutNode::GetSpan(ffi::Optional<ffi::String> axis_name)
const {
+PrimExpr layout_details::GetSpan(const ComposeLayoutNode* layout,
+ ffi::Optional<ffi::String> axis_name) {
TVM_FFI_ICHECK(!axis_name.has_value())
<< "ValueError: axis_name is not supported for compose layout";
- return tile_layout->GetSpan(axis_name);
+ return layout->tile_layout.GetSpan(axis_name);
}
-ffi::Map<ffi::String, PrimExpr> ComposeLayoutNode::Apply(ffi::Array<PrimExpr>
coord) const {
- return ApplyStructured(this, tile_layout, std::move(coord));
+ffi::Map<ffi::String, PrimExpr> layout_details::Apply(const ComposeLayoutNode*
layout,
+ ffi::Array<PrimExpr>
coord) {
+ return ApplyStructured(layout, layout->tile_layout, std::move(coord));
}
-ffi::Map<ffi::String, PrimExpr> ComposeLayoutNode::Apply(PrimExpr coord) const
{
- auto res = tile_layout->Apply(coord);
+ffi::Map<ffi::String, PrimExpr> layout_details::Apply(const ComposeLayoutNode*
layout,
+ PrimExpr coord) {
+ auto res = layout->tile_layout.Apply(coord);
TVM_FFI_ICHECK(res.size() == 1 && res.find("m") != res.end());
auto m = res["m"];
// Inline the swizzle XOR (formerly SwizzleLayoutNode::Apply): the swizzle
// operates on the tile-mapped coordinate ``m``.
auto f = [&](const PrimExpr& x) -> PrimExpr {
- if (swizzle_inner) {
- return x ^ ((x & outer_mask) >> atom_len);
+ if (layout->swizzle_inner) {
+ return x ^ ((x & layout->outer_mask) >> layout->atom_len);
} else {
- return x ^ ((x & inner_mask) << atom_len);
+ return x ^ ((x & layout->inner_mask) << layout->atom_len);
}
};
- auto base = 1 << per_element;
+ auto base = 1 << layout->per_element;
sym::Analyzer analyzer;
- return {{"m", analyzer->Simplify((f(floordiv(m, base)) << per_element) +
floormod(m, base))}};
+ return {
+ {"m", analyzer->Simplify((f(floordiv(m, base)) << layout->per_element) +
floormod(m, base))}};
}
-ffi::Map<ffi::String, PrimExpr> ComposeLayoutNode::Apply(const
ffi::Array<PrimExpr>& coord,
- const
ffi::Array<PrimExpr>& shape) const {
+ffi::Map<ffi::String, PrimExpr> layout_details::Apply(const ComposeLayoutNode*
layout,
+ const
ffi::Array<PrimExpr>& coord,
+ const
ffi::Array<PrimExpr>& shape) {
TVM_FFI_ICHECK_EQ(coord.size(), shape.size())
<< "ValueError: The size of coord and shape should be equal";
- auto grouped = TryGroup(tile_layout, shape);
+ auto grouped = TryGroup(layout->tile_layout, shape);
if (!grouped.has_value()) {
- auto mapped = tile_layout->Apply(coord, shape);
+ auto mapped = layout->tile_layout.Apply(coord, shape);
TVM_FFI_ICHECK(mapped.size() == 1 && mapped.find("m") != mapped.end());
- return {{"m", ApplyFullSwizzle(this, mapped["m"])}};
+ return {{"m", ApplyFullSwizzle(layout, mapped["m"])}};
}
const auto& [grouped_tile, separators] = grouped.value();
@@ -394,25 +402,28 @@ ffi::Map<ffi::String, PrimExpr>
ComposeLayoutNode::Apply(const ffi::Array<PrimEx
shard_coord.Set(i, split[j]);
}
}
- return ApplyStructured(this, grouped_tile, std::move(shard_coord));
+ return ApplyStructured(layout, grouped_tile, std::move(shard_coord));
}
-Layout ComposeLayoutNode::Canonicalize() const {
- auto tile_normalized = tile_layout->Canonicalize().as<TileLayout>().value();
- return ComposeLayout(per_element, swizzle_len, atom_len, tile_normalized,
swizzle_inner);
+Layout layout_details::Canonicalize(const ComposeLayoutNode* layout) {
+ auto tile_normalized =
layout->tile_layout.Canonicalize().as<TileLayout>().value();
+ return ComposeLayout(layout->per_element, layout->swizzle_len,
layout->atom_len, tile_normalized,
+ layout->swizzle_inner);
}
-Layout ComposeLayoutNode::Tile(const TileLayout& outer, const
ffi::Array<PrimExpr>& outer_shape,
- const ffi::Array<PrimExpr>& inner_shape) const {
+Layout layout_details::Tile(const ComposeLayoutNode* layout, const TileLayout&
outer,
+ const ffi::Array<PrimExpr>& outer_shape,
+ const ffi::Array<PrimExpr>& inner_shape) {
// A bare swizzle (ComposeLayout with a trivial tile) carries only the
swizzle
// period, not a tile matching `inner_shape`; substitute an identity tile
over
// the inner product first, matching the former SwizzleLayoutNode::Tile.
- TileLayout base = this->tile_layout;
- if (base->IsTrivial()) {
+ TileLayout base = layout->tile_layout;
+ if (base.IsTrivial()) {
base = IdentityTileLayout(inner_shape);
}
- auto tiled_B = base->Tile(outer, outer_shape,
inner_shape).as<TileLayout>().value();
- return ComposeLayout(per_element, swizzle_len, atom_len, tiled_B,
swizzle_inner);
+ auto tiled_B = base.Tile(outer, outer_shape,
inner_shape).as<TileLayout>().value();
+ return ComposeLayout(layout->per_element, layout->swizzle_len,
layout->atom_len, tiled_B,
+ layout->swizzle_inner);
}
} // namespace tirx
diff --git a/src/tirx/ir/layout/layout.cc b/src/tirx/ir/layout/layout.cc
index e1135fb4..6dc91ef1 100644
--- a/src/tirx/ir/layout/layout.cc
+++ b/src/tirx/ir/layout/layout.cc
@@ -17,41 +17,89 @@
* under the License.
*/
-#include "utils.h"
+#include "tile_internal.h"
namespace tvm {
namespace tirx {
-/**************** Layout ****************/
-ffi::Map<ffi::String, PrimExpr> LayoutNode::Apply(const ffi::Array<PrimExpr>&
coord,
- const ffi::Array<PrimExpr>&
shape) const {
- TVM_FFI_ICHECK_EQ(coord.size(), shape.size())
- << "ValueError: The size of coord and shape should be equal";
- return Apply(FlattenCoord(coord, shape));
+namespace {
+
+template <typename F>
+decltype(auto) DispatchLayout(const Layout& layout, F&& f) {
+ if (const auto* tile = layout.as<TileLayoutNode>()) {
+ return f(tile);
+ }
+ if (const auto* compose = layout.as<ComposeLayoutNode>()) {
+ return f(compose);
+ }
+ TVM_FFI_THROW(TypeError) << "Unsupported layout type: " <<
layout->GetTypeKey();
+}
+
+} // namespace
+
+void Layout::CheckWellFormed() const {
+ DispatchLayout(*this, [](const auto* node) {
layout_details::CheckWellFormed(node); });
+}
+
+PrimExpr Layout::GetSize(ffi::Optional<ffi::String> axis_name) const {
+ return DispatchLayout(*this,
+ [&](const auto* node) { return
layout_details::GetSize(node, axis_name); });
+}
+
+PrimExpr Layout::GetSpan(ffi::Optional<ffi::String> axis_name) const {
+ return DispatchLayout(*this,
+ [&](const auto* node) { return
layout_details::GetSpan(node, axis_name); });
+}
+
+ffi::Map<ffi::String, PrimExpr> Layout::Apply(ffi::Array<PrimExpr> coord)
const {
+ return DispatchLayout(
+ *this, [&](const auto* node) { return layout_details::Apply(node,
std::move(coord)); });
+}
+
+ffi::Map<ffi::String, PrimExpr> Layout::Apply(PrimExpr coord) const {
+ return DispatchLayout(
+ *this, [&](const auto* node) { return layout_details::Apply(node,
std::move(coord)); });
+}
+
+ffi::Map<ffi::String, PrimExpr> Layout::Apply(const ffi::Array<PrimExpr>&
coord,
+ const ffi::Array<PrimExpr>&
shape) const {
+ return DispatchLayout(
+ *this, [&](const auto* node) { return layout_details::Apply(node, coord,
shape); });
+}
+
+Layout Layout::Canonicalize() const {
+ return DispatchLayout(*this, [](const auto* node) { return
layout_details::Canonicalize(node); });
+}
+
+Layout Layout::Tile(const TileLayout& outer, const ffi::Array<PrimExpr>&
outer_shape,
+ const ffi::Array<PrimExpr>& inner_shape) const {
+ return DispatchLayout(*this, [&](const auto* node) {
+ return layout_details::Tile(node, outer, outer_shape, inner_shape);
+ });
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
auto def = refl::GlobalDef();
- def.def("tirx.LayoutCheckWellFormed", [](Layout layout) { return
layout->CheckWellFormed(); });
+ def.def("tirx.LayoutCheckWellFormed", [](Layout layout) { return
layout.CheckWellFormed(); });
def.def("tirx.LayoutGetSize", [](Layout layout, ffi::Optional<ffi::String>
axis_name) {
- return layout->GetSize(axis_name);
+ return layout.GetSize(axis_name);
});
def.def("tirx.LayoutGetSpan", [](Layout layout, ffi::Optional<ffi::String>
axis_name) {
- return layout->GetSpan(axis_name);
+ return layout.GetSpan(axis_name);
});
def.def("tirx.LayoutApplyWithShape",
[](Layout layout, ffi::Array<PrimExpr> coord, ffi::Array<PrimExpr>
shape) {
- return layout->Apply(coord, shape);
+ return layout.Apply(coord, shape);
});
def.def("tirx.LayoutApply",
- [](Layout layout, ffi::Array<PrimExpr> coord) { return
layout->Apply(coord); });
+ [](Layout layout, ffi::Array<PrimExpr> coord) { return
layout.Apply(coord); });
def.def("tirx.LayoutApplyLinear",
- [](Layout layout, PrimExpr coord) { return layout->Apply(coord); });
- def.def("tirx.LayoutCanonicalize", [](Layout layout) { return
layout->Canonicalize(); });
+ [](Layout layout, PrimExpr coord) { return layout.Apply(coord); });
+ def.def("tirx.LayoutCanonicalize", [](Layout layout) { return
layout.Canonicalize(); });
def.def("tirx.LayoutTile", [](Layout layout, TileLayout outer,
ffi::Array<PrimExpr> outer_shape,
ffi::Array<PrimExpr> inner_shape) {
- return layout->Tile(outer, outer_shape, inner_shape);
+ return layout.Tile(outer, outer_shape, inner_shape);
});
}
diff --git a/src/tirx/ir/layout/tile_canonicalize.cc
b/src/tirx/ir/layout/tile_canonicalize.cc
index 37687556..f0be540d 100644
--- a/src/tirx/ir/layout/tile_canonicalize.cc
+++ b/src/tirx/ir/layout/tile_canonicalize.cc
@@ -20,7 +20,7 @@
/*
* Canonicalization routines for TileLayout.
*/
-#include "utils.h"
+#include "tile_internal.h"
namespace tvm {
namespace tirx {
@@ -81,7 +81,7 @@ TileLayout FuseContiguousShardIters(TileLayout layout) {
TileLayout FuseAxesByScope(TileLayout layout) {
// Step 1: Get the target and scope information
- auto scope_pair_opt = layout->GetScope();
+ auto scope_pair_opt = layout.GetScope();
ffi::Optional<Target> target = Target::Current();
if (!scope_pair_opt.has_value() || !target.has_value()) {
return layout;
@@ -119,9 +119,9 @@ TileLayout FuseAxesByScope(TileLayout layout) {
return result;
}
-Layout TileLayoutNode::Canonicalize() const {
+Layout layout_details::Canonicalize(const TileLayoutNode* layout) {
// 0. Remove unit iters in shard
- TileLayout res = RemoveUnitIters(ffi::GetRef<TileLayout>(this));
+ TileLayout res = RemoveUnitIters(ffi::GetRef<TileLayout>(layout));
// 1. Remove zero offset
res = RemoveZeroOffsets(res);
// 2. Try fuse axes
diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc
index 12c224a7..38b0b7b4 100644
--- a/src/tirx/ir/layout/tile_core.cc
+++ b/src/tirx/ir/layout/tile_core.cc
@@ -196,12 +196,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-void TileLayoutNode::CheckWellFormed() const {
- TVM_FFI_CHECK(GetScope().has_value() || !HasThreadAxis(), ValueError)
+void layout_details::CheckWellFormed(const TileLayoutNode* layout) {
+ const auto tile = ffi::GetRef<TileLayout>(layout);
+ TVM_FFI_CHECK(tile.GetScope().has_value() || !tile.HasThreadAxis(),
ValueError)
<< "A tile layout with thread axes must have a connected scope";
}
-PrimExpr TileLayoutNode::GetSize(ffi::Optional<ffi::String> axis_name) const {
+PrimExpr layout_details::GetSize(const TileLayoutNode* layout,
+ ffi::Optional<ffi::String> axis_name) {
auto filter = [&](const Iter& iter, PrimExpr acc) {
if (!axis_name.has_value() || iter->axis.name() == axis_name.value()) {
return acc * iter->extent;
@@ -209,35 +211,38 @@ PrimExpr
TileLayoutNode::GetSize(ffi::Optional<ffi::String> axis_name) const {
return acc;
};
PrimExpr res = 1;
- for (const auto& iter : shard) {
+ for (const auto& iter : layout->shard) {
res = filter(iter, res);
}
return res;
}
-PrimExpr TileLayoutNode::GetSpan(ffi::Optional<ffi::String> axis_name) const {
+PrimExpr layout_details::GetSpan(const TileLayoutNode* layout,
+ ffi::Optional<ffi::String> axis_name) {
sym::Analyzer analyzer;
PrimExpr result = 1;
auto filter = [&](const Axis& axis) { return AxisMatchesFilter(axis,
axis_name); };
- for (const auto& iter : shard) {
+ for (const auto& iter : layout->shard) {
if (filter(iter->axis)) result += (iter->extent - 1) * iter->stride;
}
- for (const auto& iter : replica) {
+ for (const auto& iter : layout->replica) {
if (filter(iter->axis)) result += (iter->extent - 1) * iter->stride;
}
- for (const auto& [axis, off] : offset) {
+ for (const auto& [axis, off] : layout->offset) {
if (filter(axis)) result += off;
}
return analyzer->Simplify(result);
}
-ffi::Map<ffi::String, PrimExpr> TileLayoutNode::Apply(PrimExpr coord) const {
- return Apply(SplitCoord(coord, GetShardShape()));
+ffi::Map<ffi::String, PrimExpr> layout_details::Apply(const TileLayoutNode*
layout,
+ PrimExpr coord) {
+ return Apply(layout, SplitCoord(coord,
ffi::GetRef<TileLayout>(layout).GetShardShape()));
}
-ffi::Map<ffi::String, PrimExpr> TileLayoutNode::Apply(const
ffi::Array<PrimExpr>& coord,
- const
ffi::Array<PrimExpr>& shape) const {
+ffi::Map<ffi::String, PrimExpr> layout_details::Apply(const TileLayoutNode*
layout,
+ const
ffi::Array<PrimExpr>& coord,
+ const
ffi::Array<PrimExpr>& shape) {
TVM_FFI_ICHECK_EQ(coord.size(), shape.size())
<< "ValueError: The size of coord and shape should be equal";
// Group-first path: if this layout can be regrouped by ``shape`` (the input
@@ -247,7 +252,8 @@ ffi::Map<ffi::String, PrimExpr> TileLayoutNode::Apply(const
ffi::Array<PrimExpr>
// symbolic form small (local mod/divs) and avoids the cross-dim noise of the
// flatten+split-against-shard-shape round-trip. Equivalent numerical output,
// much friendlier for sym.Analyzer downstream.
- if (auto grouped_opt = TryGroup(ffi::GetRef<TileLayout>(this), shape);
grouped_opt.has_value()) {
+ if (auto grouped_opt = TryGroup(ffi::GetRef<TileLayout>(layout), shape);
+ grouped_opt.has_value()) {
auto& [grouped, seps] = *grouped_opt;
ffi::Array<PrimExpr> per_shard_coords;
per_shard_coords.reserve(grouped->shard.size());
@@ -267,29 +273,30 @@ ffi::Map<ffi::String, PrimExpr>
TileLayoutNode::Apply(const ffi::Array<PrimExpr>
per_shard_coords.Set(i, split[j]);
}
}
- return grouped->Apply(per_shard_coords);
+ return grouped.Apply(per_shard_coords);
}
// Fallback: flatten across input shape then split against shard shape.
- return LayoutNode::Apply(coord, shape);
+ return Apply(layout, FlattenCoord(coord, shape));
}
-ffi::Map<ffi::String, PrimExpr> TileLayoutNode::Apply(Array<PrimExpr> coord)
const {
+ffi::Map<ffi::String, PrimExpr> layout_details::Apply(const TileLayoutNode*
layout,
+ Array<PrimExpr> coord) {
sym::Analyzer analyzer;
- TVM_FFI_ICHECK_EQ(coord.size(), shard.size())
+ TVM_FFI_ICHECK_EQ(coord.size(), layout->shard.size())
<< "Coordinate size must match the number of shard axes";
std::unordered_map<ffi::String, PrimExpr> result;
- for (size_t i = 0; i < shard.size(); ++i) {
- auto it = result.find(shard[i]->axis.name());
+ for (size_t i = 0; i < layout->shard.size(); ++i) {
+ auto it = result.find(layout->shard[i]->axis.name());
if (it == result.end()) {
- result.insert_or_assign(shard[i]->axis.name(),
- analyzer->Simplify(coord[i] * shard[i]->stride));
+ result.insert_or_assign(layout->shard[i]->axis.name(),
+ analyzer->Simplify(coord[i] *
layout->shard[i]->stride));
} else {
- result.insert_or_assign(shard[i]->axis.name(),
- analyzer->Simplify(it->second + coord[i] *
shard[i]->stride));
+ result.insert_or_assign(layout->shard[i]->axis.name(),
+ analyzer->Simplify(it->second + coord[i] *
layout->shard[i]->stride));
}
}
// Add offset to the result
- for (const auto& [axis, off] : offset) {
+ for (const auto& [axis, off] : layout->offset) {
auto it = result.find(axis.name());
if (it == result.end()) {
result.insert_or_assign(axis.name(), analyzer->Simplify(off));
@@ -300,17 +307,17 @@ ffi::Map<ffi::String, PrimExpr>
TileLayoutNode::Apply(Array<PrimExpr> coord) con
return result;
}
-ffi::Array<PrimExpr> TileLayoutNode::GetShardShape() const {
- return shard.Map([](const Iter& iter) { return iter->extent; });
+ffi::Array<PrimExpr> TileLayout::GetShardShape() const {
+ return get()->shard.Map([](const Iter& iter) { return iter->extent; });
}
-bool TileLayoutNode::IsContiguous(ffi::Optional<sym::Analyzer> analyzer) const
{
- if (!replica.empty()) return false;
+bool TileLayout::IsContiguous(ffi::Optional<sym::Analyzer> analyzer) const {
+ if (!get()->replica.empty()) return false;
static const Axis memory_axis = Axis::Get("m");
- for (const auto& [axis, origin] : offset) {
+ for (const auto& [axis, origin] : get()->offset) {
if (axis != memory_axis) return false;
}
- for (const auto& iter : shard) {
+ for (const auto& iter : get()->shard) {
if (iter->axis != memory_axis) return false;
}
prim::ExprDeepEqual deep_equal;
@@ -319,40 +326,41 @@ bool
TileLayoutNode::IsContiguous(ffi::Optional<sym::Analyzer> analyzer) const {
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];
+ PrimExpr stride =
+ IntImm(get()->shard.empty() ? PrimType::Int(64) :
get()->shard.back()->stride.ty(), 1);
+ for (size_t i = get()->shard.size(); i > 0; --i) {
+ const auto& iter = get()->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) {
- if (!shard[0]->axis->IsMemoryAxis() || !IsOne(shard[0]->stride)) return
false;
+bool TileLayout::IsTrivial() const {
+ if (get()->shard.size() > 1) return false;
+ if (get()->shard.size() == 1) {
+ if (!get()->shard[0]->axis->IsMemoryAxis() ||
!IsOne(get()->shard[0]->stride)) return false;
}
- return replica.size() == 0 && offset.size() == 0;
+ return get()->replica.size() == 0 && get()->offset.size() == 0;
}
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));
+ return layout.IsContiguous(std::move(analyzer));
});
refl::GlobalDef().def("tirx.TileLayoutIsTrivial", [](const TileLayout&
layout) {
- return layout->Canonicalize().as<TileLayout>().value()->IsTrivial();
+ return layout.Canonicalize().as<TileLayout>().value().IsTrivial();
});
}
-bool TileLayoutNode::HasThreadAxis() const {
- return std::any_of(shard.begin(), shard.end(),
+bool TileLayout::HasThreadAxis() const {
+ return std::any_of(get()->shard.begin(), get()->shard.end(),
[](const Iter& iter) { return iter->axis->IsThreadAxis();
});
}
-ffi::Optional<ffi::Tuple<ThreadScope, ThreadScope>> TileLayoutNode::GetScope()
const {
+ffi::Optional<ffi::Tuple<ThreadScope, ThreadScope>> TileLayout::GetScope()
const {
if (!HasThreadAxis()) return std::nullopt;
std::unordered_map<ThreadScope, ThreadScope, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual> scope_map;
@@ -382,9 +390,9 @@ ffi::Optional<ffi::Tuple<ThreadScope, ThreadScope>>
TileLayoutNode::GetScope() c
TVM_FFI_ICHECK(inserted || it->second.same_as(scope))
<< "Ill-formed tile layout: conflicting scopes for " << subscope;
};
- for (const auto& iter : shard) check_axis(iter->axis);
- for (const auto& iter : replica) check_axis(iter->axis);
- for (const auto& [axis, off] : offset) check_axis(axis);
+ for (const auto& iter : get()->shard) check_axis(iter->axis);
+ for (const auto& iter : get()->replica) check_axis(iter->axis);
+ for (const auto& [axis, off] : get()->offset) check_axis(axis);
ThreadScope outer_most = inner_most.value();
size_t count = 0;
@@ -398,7 +406,7 @@ ffi::Optional<ffi::Tuple<ThreadScope, ThreadScope>>
TileLayoutNode::GetScope() c
return Tuple<ThreadScope, ThreadScope>{inner_most.value(), outer_most};
}
-TileLayout TileLayoutNode::DefaultLayout(ffi::Array<PrimExpr> shape) {
+TileLayout TileLayout::DefaultLayout(ffi::Array<PrimExpr> shape) {
Array<Iter> shard;
auto strides = GetDefaultStrides(shape);
for (size_t i = 0; i < shape.size(); ++i) {
@@ -412,7 +420,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def(
"tirx.TileLayoutGetScope",
[](const TileLayout& layout) -> ffi::Optional<ffi::Tuple<ThreadScope,
ThreadScope>> {
- return layout->GetScope();
+ return layout.GetScope();
});
}
diff --git a/src/tirx/ir/layout/tile_internal.h
b/src/tirx/ir/layout/tile_internal.h
index 1c3fbcb7..c0233840 100644
--- a/src/tirx/ir/layout/tile_internal.h
+++ b/src/tirx/ir/layout/tile_internal.h
@@ -30,6 +30,34 @@
namespace tvm {
namespace tirx {
+namespace layout_details {
+
+void CheckWellFormed(const TileLayoutNode* layout);
+PrimExpr GetSize(const TileLayoutNode* layout, ffi::Optional<ffi::String>
axis_name);
+PrimExpr GetSpan(const TileLayoutNode* layout, ffi::Optional<ffi::String>
axis_name);
+ffi::Map<ffi::String, PrimExpr> Apply(const TileLayoutNode* layout,
ffi::Array<PrimExpr> coord);
+ffi::Map<ffi::String, PrimExpr> Apply(const TileLayoutNode* layout, PrimExpr
coord);
+ffi::Map<ffi::String, PrimExpr> Apply(const TileLayoutNode* layout,
+ const ffi::Array<PrimExpr>& coord,
+ const ffi::Array<PrimExpr>& shape);
+Layout Canonicalize(const TileLayoutNode* layout);
+Layout Tile(const TileLayoutNode* layout, const TileLayout& outer,
+ const ffi::Array<PrimExpr>& outer_shape, const
ffi::Array<PrimExpr>& inner_shape);
+
+void CheckWellFormed(const ComposeLayoutNode* layout);
+PrimExpr GetSize(const ComposeLayoutNode* layout, ffi::Optional<ffi::String>
axis_name);
+PrimExpr GetSpan(const ComposeLayoutNode* layout, ffi::Optional<ffi::String>
axis_name);
+ffi::Map<ffi::String, PrimExpr> Apply(const ComposeLayoutNode* layout,
ffi::Array<PrimExpr> coord);
+ffi::Map<ffi::String, PrimExpr> Apply(const ComposeLayoutNode* layout,
PrimExpr coord);
+ffi::Map<ffi::String, PrimExpr> Apply(const ComposeLayoutNode* layout,
+ const ffi::Array<PrimExpr>& coord,
+ const ffi::Array<PrimExpr>& shape);
+Layout Canonicalize(const ComposeLayoutNode* layout);
+Layout Tile(const ComposeLayoutNode* layout, const TileLayout& outer,
+ const ffi::Array<PrimExpr>& outer_shape, const
ffi::Array<PrimExpr>& inner_shape);
+
+} // namespace layout_details
+
// Group a tile layout's shard by a logical shape, returning the grouped
layout and separators.
std::pair<TileLayout, std::vector<int64_t>> Group(TileLayout layout,
const ffi::Array<PrimExpr>&
shape);
diff --git a/src/tirx/ir/layout/tile_tile_ops.cc
b/src/tirx/ir/layout/tile_tile_ops.cc
index 961365b7..b430527b 100644
--- a/src/tirx/ir/layout/tile_tile_ops.cc
+++ b/src/tirx/ir/layout/tile_tile_ops.cc
@@ -122,10 +122,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-Layout TileLayoutNode::Tile(const TileLayout& outer_in, const Array<PrimExpr>&
outer_shape,
- const Array<PrimExpr>& inner_shape) const {
- auto outer = outer_in->Canonicalize().as<TileLayout>().value();
- auto inner =
ffi::GetRef<TileLayout>(this)->Canonicalize().as<TileLayout>().value();
+Layout layout_details::Tile(const TileLayoutNode* layout, const TileLayout&
outer_in,
+ const Array<PrimExpr>& outer_shape,
+ const Array<PrimExpr>& inner_shape) {
+ auto outer = outer_in.Canonicalize().as<TileLayout>().value();
+ auto inner =
ffi::GetRef<TileLayout>(layout).Canonicalize().as<TileLayout>().value();
TVM_FFI_ICHECK_EQ(outer_shape.size(), inner_shape.size())
<< "Outer and inner shape size must match";
@@ -184,7 +185,7 @@ Layout TileLayoutNode::Tile(const TileLayout& outer_in,
const Array<PrimExpr>& o
}
}
- return TileLayout(tile_shard, tile_rep, tile_offset)->Canonicalize();
+ return TileLayout(tile_shard, tile_rep, tile_offset).Canonicalize();
}
} // namespace tirx
diff --git a/src/tirx/ir/layout/utils.cc b/src/tirx/ir/layout/utils.cc
index 56b4dd0f..785a1b47 100644
--- a/src/tirx/ir/layout/utils.cc
+++ b/src/tirx/ir/layout/utils.cc
@@ -45,7 +45,7 @@ ffi::Map<ffi::String, PrimExpr> BuildSpanMap(const
TileLayout& layout) {
ffi::Map<ffi::String, PrimExpr> span_map;
for (const auto& iter : layout->shard) {
if (span_map.find(iter->axis.name()) == span_map.end()) {
- span_map.Set(iter->axis.name(), layout->GetSpan(iter->axis.name()));
+ span_map.Set(iter->axis.name(), layout.GetSpan(iter->axis.name()));
}
}
return span_map;
diff --git a/src/tirx/ir/tensor.cc b/src/tirx/ir/tensor.cc
index 00d08bfe..2537e867 100644
--- a/src/tirx/ir/tensor.cc
+++ b/src/tirx/ir/tensor.cc
@@ -400,7 +400,7 @@ ffi::Array<PrimExpr>
TensorTypeNode::ElemOffset(ffi::Array<PrimExpr> input_indic
sym::Analyzer ana;
if (layout.has_value()) {
- auto mapped = layout.value()->Apply(input_indices, shape);
+ auto mapped = layout.value().Apply(input_indices, shape);
if (mapped.empty() && shape.empty()) {
return {IntImm(DefaultIndexType(), 0)};
}
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
index 13a62ab2..e8754ad4 100644
--- a/src/tirx/ir/type.cc
+++ b/src/tirx/ir/type.cc
@@ -57,7 +57,7 @@ bool
TensorTypeNode::HasDefaultLayout(ffi::Optional<sym::Analyzer> analyzer) con
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));
+ return ffi::GetRef<TileLayout>(tile).IsContiguous(std::move(analyzer));
}
bool TensorTypeNode::IsScalar(bool alloc_or_decl) const {
@@ -65,7 +65,7 @@ bool TensorTypeNode::IsScalar(bool alloc_or_decl) const {
return shape.size() == 1 && tvm::prim::IsOne(shape[0]) && data_alignment ==
64 &&
(!layout.has_value() ||
(layout.as<TileLayoutNode>() &&
-
layout.value()->Canonicalize().as_or_throw<TileLayout>()->IsTrivial()));
+
layout.value().Canonicalize().as_or_throw<TileLayout>().IsTrivial()));
}
std::optional<int64_t> TensorTypeNode::ConstantAllocationSize() const {
diff --git a/src/tirx/transform/flatten_tensor_utils.h
b/src/tirx/transform/flatten_tensor_utils.h
index 5e001b10..ffc406e9 100644
--- a/src/tirx/transform/flatten_tensor_utils.h
+++ b/src/tirx/transform/flatten_tensor_utils.h
@@ -73,7 +73,7 @@ TVM_FFI_INLINE TensorVar FlattenedTensor(const TensorVar&
tensor,
for (const auto& iter : tile->shard) {
span = span + (iter->extent - 1) * iter->stride;
}
- if (tile->IsContiguous(ana)) {
+ if (ffi::GetRef<TileLayout>(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));
@@ -86,10 +86,10 @@ TVM_FFI_INLINE TensorVar FlattenedTensor(const TensorVar&
tensor,
}
}
} else if (tile) {
- output_shape.Set(0, ana->Simplify(tile->GetSpan(ffi::String("m"))));
+ output_shape.Set(0,
ana->Simplify(tensor.ty()->layout.value().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()) {
+ 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;
}
diff --git a/src/tirx/transform/lower_tensor_layouts.cc
b/src/tirx/transform/lower_tensor_layouts.cc
index 32fdb59c..85766f9c 100644
--- a/src/tirx/transform/lower_tensor_layouts.cc
+++ b/src/tirx/transform/lower_tensor_layouts.cc
@@ -218,8 +218,8 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
// the allocation footprint of a strided view. Round the underlying
// memory span to a block whose upper bits the swizzle cannot change.
// A bare swizzle applies to the full tensor, not just its atom period.
- PrimExpr mem_span = compose->tile_layout->GetSpan(ffi::String("m"));
- if (compose->tile_layout->IsTrivial()) {
+ PrimExpr mem_span = compose->tile_layout.GetSpan(ffi::String("m"));
+ if (compose->tile_layout.IsTrivial()) {
mem_span = IntImm::Int32(1);
for (const auto& extent : buf.shape()) mem_span = mem_span * extent;
}
@@ -303,7 +303,7 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer {
const ffi::Array<PrimExpr>&
indices) {
if (tensor.ty()->layout.has_value()) {
if (auto tile = tensor.ty()->layout.value().as<TileLayoutNode>();
- tile && tile->HasThreadAxis()) {
+ 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.";
@@ -319,7 +319,7 @@ 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());
+ auto res = layout.Canonicalize().Apply(indices, tensor.shape());
if (res.empty() && tensor.shape().empty()) {
return {IntImm(tensor.ty()->DefaultIndexType(), 0)};
}
