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 32cde7d391 Migrate TIRx Axis to native tvm-ffi enum (#20470)
32cde7d391 is described below
commit 32cde7d3916d3fcc3b4458b6a68efd5b3d132038
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 26 13:20:02 2026 -0400
Migrate TIRx Axis to native tvm-ffi enum (#20470)
Move TIRx Axis to native tvm-ffi Enum registration and shared attribute
storage, removing its separate AttrRegistry machinery.
- Preserve axis identity, thread and memory classification, typed scope
and callback behavior, name-based creation, and serialization.
- Keep Python access lazy through a small Enum subclass, including
runtime-only imports.
---
include/tvm/tirx/layout.h | 131 +--------
python/tvm/tirx/layout.py | 33 +--
src/tirx/analysis/exec_context.cc | 10 +-
src/tirx/ir/layout/axis_registry.cc | 542 +++++++++++++++++------------------
src/tirx/ir/layout/compose_layout.cc | 2 +-
src/tirx/ir/layout/tile_core.cc | 18 +-
src/tirx/ir/layout/tile_tile_ops.cc | 4 +-
src/tirx/ir/layout/utils.cc | 6 +-
src/tirx/script/printer/buffer.cc | 12 +-
tests/python/test_tirx_axis_enum.py | 65 +++++
10 files changed, 382 insertions(+), 441 deletions(-)
diff --git a/include/tvm/tirx/layout.h b/include/tvm/tirx/layout.h
index 529ea35823..31729ef718 100644
--- a/include/tvm/tirx/layout.h
+++ b/include/tvm/tirx/layout.h
@@ -25,9 +25,9 @@
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/enum.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/object.h>
-#include <tvm/ir/attr_registry_map.h>
#include <tvm/ir/module.h>
#include <tvm/target/target.h>
#include <tvm/tirx/exec_scope.h>
@@ -35,13 +35,7 @@
namespace tvm {
-// Forward declaration
-template <typename, typename>
-class AttrRegistry;
-
namespace tirx {
-template <typename>
-class AxisAttrMap;
class Layout;
class TileLayout;
@@ -146,15 +140,8 @@ using FAxisFuser =
ffi::TypedFunction<ffi::Optional<Iter>(Target, ffi::String, f
using FAxisSplitter = ffi::TypedFunction<ffi::Array<Iter, void>(Target,
ffi::String, Iter)>;
// Axis
-class AxisNode : public ffi::Object {
+class AxisNode : public ffi::EnumObj {
public:
- ffi::String name;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AxisNode>().def_ro("name", &AxisNode::name);
- }
-
/*! \brief Check if the axis is a thread axis. */
bool IsThreadAxis() const;
@@ -173,118 +160,26 @@ class AxisNode : public ffi::Object {
/*! \brief Get the splitter of the (thread) axis. */
ffi::Optional<FAxisSplitter> GetSplitter() const;
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Axis", AxisNode, ffi::Object);
-
- private:
- // Iternals necessary for AttrRegistry
- template <typename>
- friend class tvm::AttrRegistryMapContainerMap;
- template <typename, typename>
- friend class tvm::AttrRegistry;
- friend class AxisRegEntry;
- /*! \brief Program internal unique index of operator. */
- uint32_t index_{0};
- /*! \brief Return the index stored in attr registry */
- uint32_t AttrRegistryIndex() const { return index_; }
- /*! \brief Return the name stored in attr registry */
- ffi::String AttrRegistryName() const { return name; }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Axis", AxisNode, ffi::EnumObj);
};
-class Axis : public ffi::ObjectRef {
+class Axis : public ffi::Enum {
public:
Axis() = default;
- /*! \brief Get the axis object by name. */
- TVM_DLL static Axis Get(const ffi::String& name);
-
- /*! \brief Get the attribute map for the axis. */
- template <typename ValueType>
- inline static AxisAttrMap<ValueType> GetAttrMap(const ffi::String&
attr_name);
-
- explicit Axis(ffi::ObjectPtr<AxisNode> data) : ObjectRef(ffi::UnsafeInit{}) {
- TVM_FFI_ICHECK(data != nullptr);
- data_ = std::move(data);
- }
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Axis, ffi::ObjectRef,
AxisNode);
-
- private:
- // Internals necessary for AttrRegistry
- template <typename, typename>
- friend class tvm::AttrRegistry;
- friend class AxisRegEntry;
-};
-
-// AxisRegistry
-class AxisRegEntry {
- public:
- /*! \brief List all axis names. */
- TVM_DLL static ffi::Array<ffi::String> ListAxisNames();
+ /*! \brief Return the canonical axis name. */
+ const ffi::String& name() const { return get()->_str_index; }
- /*! \brief Register or get the axis entry by name. */
- TVM_DLL static AxisRegEntry& RegisterOrGet(const ffi::String& name);
-
- /*! \brief Set the attribute for the axis. */
- template <typename ValueType>
- inline AxisRegEntry& set_attr(const ffi::String& attr_name, const ValueType&
value,
- int plevel = 10);
-
- /*! \brief Set the scope of the axis. */
- inline AxisRegEntry& set_scope(const ffi::String& scope_name, int plevel =
10);
-
- /*! \brief Set the subscope of the axis. */
- inline AxisRegEntry& set_subscope(const ffi::String& subscope_name, int
plevel = 10);
-
- /*! \brief Set the fuser of the axis. */
- inline AxisRegEntry& set_fuser(const FAxisFuser& fuser);
-
- /*! \brief Set the splitter of the axis. */
- inline AxisRegEntry& set_splitter(const FAxisSplitter& splitter);
-
- private:
- // return internal pointer to op.
- inline AxisNode* get();
- TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
-
- // Internals necessary for AttrRegistry
- Axis axis_;
- ffi::String name;
- explicit AxisRegEntry(uint32_t index);
- template <typename, typename>
- friend class tvm::AttrRegistry;
- friend class Axis;
-};
-
-using AxisRegistry = AttrRegistry<AxisRegEntry, Axis>;
+ /*!
+ * \brief Get or create the axis with the given name.
+ * \param name Name of the axis to look up or register.
+ * \return The registered axis singleton.
+ */
+ TVM_DLL static Axis Get(const ffi::String& name);
-// AxisAttrffi::Map
-template <typename ValueType>
-class AxisAttrMap : public AttrRegistryMap<Axis, ValueType> {
- public:
- using TParent = AttrRegistryMap<Axis, ValueType>;
- using TParent::count;
- using TParent::get;
- using TParent::operator[];
-
- private:
- friend class Axis;
- explicit AxisAttrMap(const AttrRegistryMapContainerMap<Axis>& map) :
TParent(map) {}
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Axis, ffi::Enum, AxisNode);
};
-// Helper macro for token concatenation
-#ifndef TVM_STR_CONCAT
-#define TVM_STR_CONCAT_(__x, __y) __x##__y
-#define TVM_STR_CONCAT(__x, __y) TVM_STR_CONCAT_(__x, __y)
-#endif
-
-// Define a macro to register the axis entry.
-#define TVM_AXIS_REGISTER_VAR_DEF [[maybe_unused]] static
::tvm::tirx::AxisRegEntry& __make_##Axis
-
-#define TVM_REGISTER_AXIS(AxisName) \
- TVM_STR_CONCAT(TVM_AXIS_REGISTER_VAR_DEF, __COUNTER__) = \
- ::tvm::tirx::AxisRegEntry::RegisterOrGet(AxisName)
-
class IterNode : public ffi::Object {
public:
PrimExpr extent;
diff --git a/python/tvm/tirx/layout.py b/python/tvm/tirx/layout.py
index 02878df314..f1015092c0 100644
--- a/python/tvm/tirx/layout.py
+++ b/python/tvm/tirx/layout.py
@@ -21,9 +21,10 @@ import functools
import operator
import re
from collections.abc import Sequence
-from typing import ClassVar, Optional, Union
+from typing import Optional, Union
import tvm_ffi
+from tvm_ffi.dataclasses import Enum
import tvm
from tvm.runtime import Object
@@ -502,7 +503,7 @@ _AXIS_NAMES = (
)
-class _AxisMeta(type(Object)):
+class _AxisMeta(type(Enum)):
"""Metaclass: lazy resolve `Axis.<name>` for registered axes."""
def __getattr__(cls, name):
@@ -511,27 +512,21 @@ class _AxisMeta(type(Object)):
raise AttributeError(f"type object 'Axis' has no attribute {name!r}")
-@tvm_ffi.register_object("tirx.Axis")
-class Axis(Object, metaclass=_AxisMeta):
+class Axis(Enum, metaclass=_AxisMeta, type_key="tirx.Axis", init=False):
"""Layout axis wrapper."""
- # ---- forbid direct construction ----
- def __init__(self, *args, **kwargs):
- raise RuntimeError("Cannot create Axis directly; use Axis.get()")
+ @property
+ def name(self) -> str:
+ """Return the canonical axis name."""
+ return self._str_index
- @staticmethod
- def _register_axis(name: str) -> "Axis":
- return _ffi_api.AxisGet(name) # pylint: disable=no-member
-
- # Singleton cache, populated lazily as names are accessed.
- reg_dict: ClassVar[dict[str, "Axis"]] = {}
+ @classmethod
+ def get(cls, name: str) -> "Axis":
+ """Get or create the axis singleton named `name`.
- @staticmethod
- def get(name: str) -> "Axis":
- """Get or create an axis by name. Unknown names are auto-registered."""
- if name not in Axis.reg_dict:
- Axis.reg_dict[name] = Axis._register_axis(name)
- return Axis.reg_dict[name]
+ Unknown names are registered without thread or memory attributes.
+ """
+ return _ffi_api.AxisGet(name) # pylint: disable=no-member
def is_thread(self) -> bool:
"""Check if the axis is a thread axis."""
diff --git a/src/tirx/analysis/exec_context.cc
b/src/tirx/analysis/exec_context.cc
index d47114bd6c..fcfffa387b 100644
--- a/src/tirx/analysis/exec_context.cc
+++ b/src/tirx/analysis/exec_context.cc
@@ -77,8 +77,8 @@ std::vector<std::pair<std::string, AxisRange>>
AxisRanges(const ActiveSet& A) {
std::vector<std::pair<std::string, AxisRange>> axes;
for (const auto& iter : A.layout->shard) {
AxisRange range;
- TVM_FFI_ICHECK(A.GetAxis(iter->axis->name.operator std::string(), &range));
- axes.push_back({iter->axis->name.operator std::string(), range});
+ TVM_FFI_ICHECK(A.GetAxis(iter->axis.name().operator std::string(),
&range));
+ axes.push_back({iter->axis.name().operator std::string(), range});
}
return axes;
}
@@ -381,10 +381,10 @@ bool AxisRange::Modulo(int64_t modulus, int64_t residue,
AxisRange* out) const {
bool ActiveSet::GetAxis(const std::string& axis, AxisRange* out) const {
if (!layout.defined()) return false;
for (const auto& iter : layout->shard) {
- if (iter->axis->name != axis) continue;
+ if (iter->axis.name() != axis) continue;
PrimExpr off = I64(0);
for (const auto& kv : layout->offset) {
- if (kv.first->name == axis) {
+ if (kv.first.name() == axis) {
off = kv.second;
break;
}
@@ -418,7 +418,7 @@ std::vector<std::string> ActiveSet::AxisNames() const {
std::vector<std::string> names;
if (!layout.defined()) return names;
for (const auto& iter : layout->shard) {
- names.push_back(iter->axis->name.operator std::string());
+ names.push_back(iter->axis.name().operator std::string());
}
return names;
}
diff --git a/src/tirx/ir/layout/axis_registry.cc
b/src/tirx/ir/layout/axis_registry.cc
index 2474390789..ce8d8aaa5c 100644
--- a/src/tirx/ir/layout/axis_registry.cc
+++ b/src/tirx/ir/layout/axis_registry.cc
@@ -20,144 +20,84 @@
/*
* Axis definitions, attributes, fusers/splitters, and registrations.
*/
+#include <tvm/ffi/reflection/enum_def.h>
+
#include "utils.h"
namespace tvm {
namespace tirx {
+namespace refl = tvm::ffi::reflection;
/**************** Axis ****************/
-// AxisNode
ffi::ObjectPtr<ffi::Object> CreateAxis(const std::string& name) {
- // Hack use ffi::Any as exchange
auto axis = Axis::Get(name);
- TVM_FFI_ICHECK(axis.defined()) << "Cannot find axis '" << name << '\'';
return ffi::details::ObjectUnsafe::ObjectPtrFromObjectRef<ffi::Object>(axis);
}
+const ffi::EnumState& AxisEnumState() {
+ static refl::TypeAttrColumn state_column(refl::type_attr::kEnumState);
+ static ffi::EnumState state =
+
ffi::Any(state_column[AxisNode::_GetOrAllocRuntimeTypeIndex()]).as_or_throw<ffi::EnumState>();
+ return state;
+}
+
+template <typename T>
+ffi::Optional<T> AxisAttr(const AxisNode* axis, const char* key) {
+ const ffi::EnumState& state = AxisEnumState();
+ auto column = state->attrs.Get(ffi::String(key));
+ if (!column) return std::nullopt;
+ auto value = column.value().Get(ffi::GetRef<Axis>(axis));
+ if (!value) return std::nullopt;
+ return value.value().cast<T>();
+}
+
bool AxisNode::IsThreadAxis() const {
- static const auto& thread_attr_map = Axis::GetAttrMap<bool>("thread");
- return thread_attr_map[ffi::GetRef<Axis>(this)];
+ auto thread = AxisAttr<bool>(this, "thread");
+ TVM_FFI_ICHECK(thread.has_value()) << "Axis '" << _str_index << "' has no
thread classification";
+ return thread.value();
}
bool AxisNode::IsMemoryAxis() const {
- static const auto& thread_attr_map = Axis::GetAttrMap<bool>("thread");
- return !thread_attr_map[ffi::GetRef<Axis>(this)];
+ auto thread = AxisAttr<bool>(this, "thread");
+ TVM_FFI_ICHECK(thread.has_value()) << "Axis '" << _str_index << "' has no
thread classification";
+ return !thread.value();
}
-ffi::Optional<ExecScope> AxisNode::GetScope() const {
- static const auto& scope_attr_map =
Axis::GetAttrMap<ffi::Optional<ExecScope>>("scope");
- return scope_attr_map.get(ffi::GetRef<Axis>(this), std::nullopt);
-}
+ffi::Optional<ExecScope> AxisNode::GetScope() const { return
AxisAttr<ExecScope>(this, "scope"); }
ffi::Optional<ExecScope> AxisNode::GetSubscope() const {
- static const auto& subscope_attr_map =
Axis::GetAttrMap<ffi::Optional<ExecScope>>("subscope");
- return subscope_attr_map.get(ffi::GetRef<Axis>(this), std::nullopt);
+ return AxisAttr<ExecScope>(this, "subscope");
}
-ffi::Optional<FAxisFuser> AxisNode::GetFuser() const {
- static const auto& fuser_attr_map =
Axis::GetAttrMap<ffi::Optional<FAxisFuser>>("fuser");
- return fuser_attr_map.get(ffi::GetRef<Axis>(this), std::nullopt);
-}
+ffi::Optional<FAxisFuser> AxisNode::GetFuser() const { return
AxisAttr<FAxisFuser>(this, "fuser"); }
ffi::Optional<FAxisSplitter> AxisNode::GetSplitter() const {
- static const auto& splitter_attr_map =
Axis::GetAttrMap<ffi::Optional<FAxisSplitter>>("splitter");
- return splitter_attr_map.get(ffi::GetRef<Axis>(this), std::nullopt);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.AxisIsThreadAxis", [](Axis axis) { return
axis->IsThreadAxis(); });
+ return AxisAttr<FAxisSplitter>(this, "splitter");
}
TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.AxisIsMemoryAxis", [](Axis axis) { return
axis->IsMemoryAxis(); });
-}
+ refl::ObjectDef<AxisNode>(refl::init(false));
+ refl::TypeAttrDef<AxisNode>()
+ .def("__data_to_json__", [](const AxisNode* node) -> ffi::String {
return node->_str_index; })
+ .def("__data_from_json__", [](const ffi::String& name) -> Axis { return
Axis::Get(name); });
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.AxisGetScope", [](Axis axis) { return
axis->GetScope(); });
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.AxisGetSubscope", [](Axis axis) { return
axis->GetSubscope(); });
+ refl::GlobalDef()
+ .def("tirx.AxisIsThreadAxis", [](Axis axis) { return
axis->IsThreadAxis(); })
+ .def("tirx.AxisIsMemoryAxis", [](Axis axis) { return
axis->IsMemoryAxis(); })
+ .def("tirx.AxisGetScope", [](Axis axis) { return axis->GetScope(); })
+ .def("tirx.AxisGetSubscope", [](Axis axis) { return axis->GetSubscope();
})
+ .def("tirx.AxisGet", [](ffi::String name) -> Axis { return
Axis::Get(name); });
}
// Axis
Axis Axis::Get(const ffi::String& name) {
- const AxisRegEntry* reg = AxisRegistry::Global()->Get(name);
- if (reg != nullptr) {
- return reg->axis_;
+ const ffi::EnumState& state = AxisEnumState();
+ auto existing = state->indexes.Get(ffi::Any(name));
+ if (!existing) {
+ refl::EnumDef<AxisNode>(name.c_str());
+ existing = state->indexes.Get(ffi::Any(name));
}
- // Auto-register unknown axes on the fly
- return AxisRegEntry::RegisterOrGet(name).axis_;
-}
-
-template <typename ValueType>
-inline AxisAttrMap<ValueType> Axis::GetAttrMap(const ffi::String& attr_name) {
- return AxisAttrMap<ValueType>(AxisRegistry::Global()->GetAttrMap(attr_name));
-}
-
-// AxisRegEntry
-inline AxisNode* AxisRegEntry::get() { return
const_cast<AxisNode*>(axis_.operator->()); }
-
-AxisRegEntry::AxisRegEntry(uint32_t index) {
- ffi::ObjectPtr<AxisNode> n = ffi::make_object<AxisNode>();
- n->index_ = index;
- axis_ = Axis(n);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- AxisNode::RegisterReflection();
- refl::TypeAttrDef<AxisNode>()
- .def("__data_to_json__", [](const AxisNode* node) -> ffi::String {
return node->name; })
- .def("__data_from_json__", [](const ffi::String& name) -> Axis { return
Axis::Get(name); });
-}
-
-AxisRegEntry& AxisRegEntry::RegisterOrGet(const ffi::String& name) {
- auto& entry = AxisRegistry::Global()->RegisterOrGet(name);
- entry.get()->name = name;
- return entry;
-}
-
-ffi::Array<ffi::String> AxisRegEntry::ListAxisNames() {
- return AxisRegistry::Global()->ListAllNames();
-}
-
-template <typename ValueType>
-inline AxisRegEntry& AxisRegEntry::set_attr(const ffi::String& key, const
ValueType& value,
- int plevel) {
- TVM_FFI_ICHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
- ffi::Any rv;
- rv = value;
- UpdateAttr(key, rv, plevel);
- return *this;
-}
-
-AxisRegEntry& AxisRegEntry::set_scope(const ffi::String& scope_name, int
plevel) {
- set_attr<ffi::Optional<ExecScope>>("scope", ExecScope(scope_name), plevel);
- return *this;
-}
-
-AxisRegEntry& AxisRegEntry::set_subscope(const ffi::String& subscope_name, int
plevel) {
- set_attr<ffi::Optional<ExecScope>>("subscope", ExecScope(subscope_name),
plevel);
- return *this;
-}
-
-AxisRegEntry& AxisRegEntry::set_fuser(const FAxisFuser& fuser) {
- set_attr<ffi::Optional<FAxisFuser>>("fuser", fuser);
- return *this;
-}
-
-AxisRegEntry& AxisRegEntry::set_splitter(const FAxisSplitter& splitter) {
- set_attr<ffi::Optional<FAxisSplitter>>("splitter", splitter);
- return *this;
-}
-
-void AxisRegEntry::UpdateAttr(const ffi::String& key, ffi::Any value, int
plevel) {
- AxisRegistry::Global()->UpdateAttr(key, axis_, value, plevel);
+ return existing.value().as_or_throw<Axis>();
}
// register thread axis split/fuse helpers
@@ -178,179 +118,225 @@ ffi::Array<Iter> SplitterGen(const Iter& iter, const
Axis& axis_outer, const Axi
return {};
}
-// register thread axes
-TVM_REGISTER_AXIS("bx").set_attr<bool>("thread",
true).set_scope("thread").set_subscope("cta");
-TVM_REGISTER_AXIS("by").set_attr<bool>("thread",
true).set_scope("thread").set_subscope("cta");
-TVM_REGISTER_AXIS("bz").set_attr<bool>("thread",
true).set_scope("thread").set_subscope("cta");
-TVM_REGISTER_AXIS("cbx").set_attr<bool>("thread",
true).set_scope("cluster").set_subscope("cta");
-TVM_REGISTER_AXIS("cby").set_attr<bool>("thread",
true).set_scope("cluster").set_subscope("cta");
-TVM_REGISTER_AXIS("cbz").set_attr<bool>("thread",
true).set_scope("cluster").set_subscope("cta");
-TVM_REGISTER_AXIS("tx")
- .set_attr<bool>("thread", true)
- .set_scope("cta")
- .set_subscope("thread")
- .set_fuser([](Target target, ffi::String subscope, ffi::String scope,
- Iter iter) -> ffi::Optional<Iter> {
- if (target->kind->default_device_type == kDLCUDA) {
- return std::nullopt;
- }
- return std::nullopt;
- })
- .set_splitter([](Target target, ffi::String scope, Iter iter) ->
ffi::Array<Iter> {
- sym::Analyzer analyzer;
- if (target->kind->default_device_type == kDLCUDA) {
- if (scope == "warp") {
- // tx -> warpid, laneid
- return SplitterGen(iter, Axis::Get("warpid"), Axis::Get("laneid"),
32);
- } else if (scope == "warpgroup") {
- // tx -> wgid, tid_in_wg
- return SplitterGen(iter, Axis::Get("wgid"), Axis::Get("tid_in_wg"),
128);
- }
- LOG(FATAL) << "Cannot split cta->thread axis into cta->" << scope <<
"->thread";
- }
- return {};
- });
-TVM_REGISTER_AXIS("warpid")
- .set_attr<bool>("thread", true)
- .set_scope("cta")
- .set_subscope("warp")
- .set_fuser([](Target target, ffi::String subscope, ffi::String scope,
- Iter iter) -> ffi::Optional<Iter> {
- if (target->kind->default_device_type == kDLCUDA) {
- // cta->warp ===> cta->thread (tx)
- if (subscope == "thread" && scope == "cta") {
- return Iter(iter->extent, 32 * iter->stride, Axis::Get("tx"));
- }
- return std::nullopt;
- }
- return std::nullopt;
- })
- .set_splitter([](Target target, ffi::String scope, Iter iter) ->
ffi::Array<Iter> {
- sym::Analyzer analyzer;
- if (target->kind->default_device_type == kDLCUDA) {
- if (scope == "warp") {
- // warpid -> wgid, wid_in_wg
- return SplitterGen(iter, Axis::Get("wgid"), Axis::Get("wid_in_wg"),
4);
- }
- LOG(FATAL) << "Cannot split cta->warp axis into cta->" << scope <<
"->warp";
- }
- return {};
- });
-TVM_REGISTER_AXIS("laneid")
- .set_attr<bool>("thread", true)
- .set_scope("warp")
- .set_subscope("thread")
- .set_fuser([](Target target, ffi::String subscope, ffi::String scope,
- Iter iter) -> ffi::Optional<Iter> {
- if (target->kind->default_device_type == kDLCUDA) {
- if (subscope == "thread" && scope == "warpgroup") {
- // warp->thread ===> warpgroup->thread (tid_in_wg)
- return Iter(iter->extent, iter->stride, Axis::Get("tid_in_wg"));
- } else if (subscope == "thread" && scope == "cta") {
- // warp->thread ===> cta->thread (tx)
- return Iter(iter->extent, iter->stride, Axis::Get("tx"));
- }
- return std::nullopt;
- }
- return std::nullopt;
- })
- .set_splitter([](Target target, ffi::String scope, Iter iter) ->
ffi::Array<Iter> {
- sym::Analyzer analyzer;
- if (target->kind->default_device_type == kDLCUDA) {
- LOG(FATAL) << "laneid can not be split any more";
- }
- return {};
- });
-TVM_REGISTER_AXIS("wgid")
- .set_attr<bool>("thread", true)
- .set_scope("cta")
- .set_subscope("warpgroup")
- .set_fuser([](Target target, ffi::String subscope, ffi::String scope,
- Iter iter) -> ffi::Optional<Iter> {
- if (target->kind->default_device_type == kDLCUDA) {
- if (subscope == "thread" && scope == "cta") {
- // cta->warpgroup ===> cta->thread (tx)
- return Iter(iter->extent, iter->stride * 128, Axis::Get("tx"));
- } else if (subscope == "warp" && scope == "cta") {
- // cta->warpgroup ===> cta->warp (warpid)
- return Iter(iter->extent, iter->stride * 4, Axis::Get("wgid"));
- }
- }
- return std::nullopt;
- })
- .set_splitter([](Target target, ffi::String scope, Iter iter) ->
ffi::Array<Iter> {
- sym::Analyzer analyzer;
- if (target->kind->default_device_type == kDLCUDA) {
- LOG(FATAL) << "wgid can not be split any more";
- }
- return {};
- });
-TVM_REGISTER_AXIS("tid_in_wg")
- .set_attr<bool>("thread", true)
- .set_scope("warpgroup")
- .set_subscope("thread")
- .set_fuser([](Target target, ffi::String subscope, ffi::String scope,
- Iter iter) -> ffi::Optional<Iter> {
- if (target->kind->default_device_type == kDLCUDA) {
- if (subscope == "thread" && scope == "cta") {
- // warpgroup->thread ===> cta->thread (tx)
- return Iter(iter->extent, iter->stride, Axis::Get("tx"));
- }
- return std::nullopt;
- }
- return std::nullopt;
- })
- .set_splitter([](Target target, ffi::String scope, Iter iter) ->
ffi::Array<Iter> {
- sym::Analyzer analyzer;
- if (target->kind->default_device_type == kDLCUDA) {
- if (scope == "warp") {
- // tid_in_wg -> wid_in_wg, laneid
- return SplitterGen(iter, Axis::Get("wid_in_wg"),
Axis::Get("laneid"), 32);
- }
- LOG(FATAL) << "Cannot split warpgroup->thread axis into warpgroup->"
<< scope << "->thread";
- }
- return {};
- });
-TVM_REGISTER_AXIS("wid_in_wg")
- .set_attr<bool>("thread", true)
- .set_scope("warpgroup")
- .set_subscope("warp")
- .set_fuser([](Target target, ffi::String subscope, ffi::String scope,
- Iter iter) -> ffi::Optional<Iter> {
- if (target->kind->default_device_type == kDLCUDA) {
- if (subscope == "thread" && scope == "warpgroup") {
- // warpgroup->warp ===> warpgroup->thread (tid_in_wg)
- return Iter(iter->extent, iter->stride * 32, Axis::Get("tid_in_wg"));
- } else if (subscope == "thread" && scope == "cta") {
- // warpgroup->warp ===> cta->thread (tx)
- return Iter(iter->extent, iter->stride * 32, Axis::Get("tx"));
- } else if (subscope == "warp" && scope == "cta") {
- // warpgroup->warp ===> cta->warp (warpid)
- return Iter(iter->extent, iter->stride, Axis::Get("warpid"));
- }
- return std::nullopt;
- }
- return std::nullopt;
- })
- .set_splitter([](Target target, ffi::String scope, Iter iter) ->
ffi::Array<Iter> {
- sym::Analyzer analyzer;
- if (target->kind->default_device_type == kDLCUDA) {
- LOG(FATAL) << "wid_in_wg can not be split any more";
- }
- return {};
- });
-
-// register memory axis
-TVM_REGISTER_AXIS("m").set_attr<bool>("thread", false);
-TVM_REGISTER_AXIS("P").set_attr<bool>("thread", false);
-TVM_REGISTER_AXIS("F").set_attr<bool>("thread", false);
-TVM_REGISTER_AXIS("Bank").set_attr<bool>("thread", false);
-TVM_REGISTER_AXIS("TCol").set_attr<bool>("thread", false);
-TVM_REGISTER_AXIS("TLane").set_attr<bool>("thread", false);
-
TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.AxisGet", [](ffi::String name) -> Axis { return
Axis::Get(name); });
+ // Register thread axes in declaration order.
+ refl::EnumDef<AxisNode>("bx")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("thread"))
+ .set_attr("subscope", ExecScope("cta"));
+
+ refl::EnumDef<AxisNode>("by")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("thread"))
+ .set_attr("subscope", ExecScope("cta"));
+
+ refl::EnumDef<AxisNode>("bz")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("thread"))
+ .set_attr("subscope", ExecScope("cta"));
+
+ refl::EnumDef<AxisNode>("cbx")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("cluster"))
+ .set_attr("subscope", ExecScope("cta"));
+
+ refl::EnumDef<AxisNode>("cby")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("cluster"))
+ .set_attr("subscope", ExecScope("cta"));
+
+ refl::EnumDef<AxisNode>("cbz")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("cluster"))
+ .set_attr("subscope", ExecScope("cta"));
+
+ refl::EnumDef<AxisNode>("tx")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("cta"))
+ .set_attr("subscope", ExecScope("thread"))
+ .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ Iter iter) -> ffi::Optional<Iter> {
+ if (target->kind->default_device_type == kDLCUDA) {
+ return std::nullopt;
+ }
+ return std::nullopt;
+ }))
+ .set_attr("splitter",
+ FAxisSplitter([](Target target, ffi::String scope, Iter iter)
-> ffi::Array<Iter> {
+ sym::Analyzer analyzer;
+ if (target->kind->default_device_type == kDLCUDA) {
+ if (scope == "warp") {
+ // tx -> warpid, laneid
+ return SplitterGen(iter, Axis::Get("warpid"),
Axis::Get("laneid"), 32);
+ } else if (scope == "warpgroup") {
+ // tx -> wgid, tid_in_wg
+ return SplitterGen(iter, Axis::Get("wgid"),
Axis::Get("tid_in_wg"), 128);
+ }
+ LOG(FATAL) << "Cannot split cta->thread axis into cta->"
<< scope << "->thread";
+ }
+ return {};
+ }));
+
+ refl::EnumDef<AxisNode>("warpid")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("cta"))
+ .set_attr("subscope", ExecScope("warp"))
+ .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ Iter iter) -> ffi::Optional<Iter> {
+ if (target->kind->default_device_type == kDLCUDA) {
+ // cta->warp ===> cta->thread (tx)
+ if (subscope == "thread" && scope == "cta") {
+ return Iter(iter->extent, 32 * iter->stride,
Axis::Get("tx"));
+ }
+ return std::nullopt;
+ }
+ return std::nullopt;
+ }))
+ .set_attr("splitter",
+ FAxisSplitter([](Target target, ffi::String scope, Iter iter)
-> ffi::Array<Iter> {
+ sym::Analyzer analyzer;
+ if (target->kind->default_device_type == kDLCUDA) {
+ if (scope == "warp") {
+ // warpid -> wgid, wid_in_wg
+ return SplitterGen(iter, Axis::Get("wgid"),
Axis::Get("wid_in_wg"), 4);
+ }
+ LOG(FATAL) << "Cannot split cta->warp axis into cta->" <<
scope << "->warp";
+ }
+ return {};
+ }));
+
+ refl::EnumDef<AxisNode>("laneid")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("warp"))
+ .set_attr("subscope", ExecScope("thread"))
+ .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ Iter iter) -> ffi::Optional<Iter> {
+ if (target->kind->default_device_type == kDLCUDA) {
+ if (subscope == "thread" && scope == "warpgroup") {
+ // warp->thread ===> warpgroup->thread (tid_in_wg)
+ return Iter(iter->extent, iter->stride,
Axis::Get("tid_in_wg"));
+ } else if (subscope == "thread" && scope == "cta") {
+ // warp->thread ===> cta->thread (tx)
+ return Iter(iter->extent, iter->stride, Axis::Get("tx"));
+ }
+ return std::nullopt;
+ }
+ return std::nullopt;
+ }))
+ .set_attr("splitter",
+ FAxisSplitter([](Target target, ffi::String scope, Iter iter)
-> ffi::Array<Iter> {
+ sym::Analyzer analyzer;
+ if (target->kind->default_device_type == kDLCUDA) {
+ LOG(FATAL) << "laneid can not be split any more";
+ }
+ return {};
+ }));
+
+ refl::EnumDef<AxisNode>("wgid")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("cta"))
+ .set_attr("subscope", ExecScope("warpgroup"))
+ .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ Iter iter) -> ffi::Optional<Iter> {
+ if (target->kind->default_device_type == kDLCUDA) {
+ if (subscope == "thread" && scope == "cta") {
+ // cta->warpgroup ===> cta->thread (tx)
+ return Iter(iter->extent, iter->stride * 128,
Axis::Get("tx"));
+ } else if (subscope == "warp" && scope == "cta") {
+ // cta->warpgroup ===> cta->warp (warpid)
+ return Iter(iter->extent, iter->stride * 4,
Axis::Get("wgid"));
+ }
+ }
+ return std::nullopt;
+ }))
+ .set_attr("splitter",
+ FAxisSplitter([](Target target, ffi::String scope, Iter iter)
-> ffi::Array<Iter> {
+ sym::Analyzer analyzer;
+ if (target->kind->default_device_type == kDLCUDA) {
+ LOG(FATAL) << "wgid can not be split any more";
+ }
+ return {};
+ }));
+
+ refl::EnumDef<AxisNode>("tid_in_wg")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("warpgroup"))
+ .set_attr("subscope", ExecScope("thread"))
+ .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ Iter iter) -> ffi::Optional<Iter> {
+ if (target->kind->default_device_type == kDLCUDA) {
+ if (subscope == "thread" && scope == "cta") {
+ // warpgroup->thread ===> cta->thread (tx)
+ return Iter(iter->extent, iter->stride, Axis::Get("tx"));
+ }
+ return std::nullopt;
+ }
+ return std::nullopt;
+ }))
+ .set_attr("splitter",
+ FAxisSplitter([](Target target, ffi::String scope, Iter iter)
-> ffi::Array<Iter> {
+ sym::Analyzer analyzer;
+ if (target->kind->default_device_type == kDLCUDA) {
+ if (scope == "warp") {
+ // tid_in_wg -> wid_in_wg, laneid
+ return SplitterGen(iter, Axis::Get("wid_in_wg"),
Axis::Get("laneid"), 32);
+ }
+ LOG(FATAL) << "Cannot split warpgroup->thread axis into
warpgroup->" << scope
+ << "->thread";
+ }
+ return {};
+ }));
+
+ refl::EnumDef<AxisNode>("wid_in_wg")
+ .set_attr("thread", true)
+ .set_attr("scope", ExecScope("warpgroup"))
+ .set_attr("subscope", ExecScope("warp"))
+ .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ Iter iter) -> ffi::Optional<Iter> {
+ if (target->kind->default_device_type == kDLCUDA) {
+ if (subscope == "thread" && scope == "warpgroup") {
+ // warpgroup->warp ===> warpgroup->thread (tid_in_wg)
+ return Iter(iter->extent, iter->stride * 32,
Axis::Get("tid_in_wg"));
+ } else if (subscope == "thread" && scope == "cta") {
+ // warpgroup->warp ===> cta->thread (tx)
+ return Iter(iter->extent, iter->stride * 32,
Axis::Get("tx"));
+ } else if (subscope == "warp" && scope == "cta") {
+ // warpgroup->warp ===> cta->warp (warpid)
+ return Iter(iter->extent, iter->stride,
Axis::Get("warpid"));
+ }
+ return std::nullopt;
+ }
+ return std::nullopt;
+ }))
+ .set_attr("splitter",
+ FAxisSplitter([](Target target, ffi::String scope, Iter iter)
-> ffi::Array<Iter> {
+ sym::Analyzer analyzer;
+ if (target->kind->default_device_type == kDLCUDA) {
+ LOG(FATAL) << "wid_in_wg can not be split any more";
+ }
+ return {};
+ }));
+
+ // Register memory axes.
+ // clang-format off
+ refl::EnumDef<AxisNode>("m")
+ .set_attr("thread", false);
+
+ refl::EnumDef<AxisNode>("P")
+ .set_attr("thread", false);
+
+ refl::EnumDef<AxisNode>("F")
+ .set_attr("thread", false);
+
+ refl::EnumDef<AxisNode>("Bank")
+ .set_attr("thread", false);
+
+ refl::EnumDef<AxisNode>("TCol")
+ .set_attr("thread", false);
+
+ refl::EnumDef<AxisNode>("TLane")
+ .set_attr("thread", false);
+ // clang-format on
}
} // namespace tirx
diff --git a/src/tirx/ir/layout/compose_layout.cc
b/src/tirx/ir/layout/compose_layout.cc
index 1f79b13e41..8a19e3fdc5 100644
--- a/src/tirx/ir/layout/compose_layout.cc
+++ b/src/tirx/ir/layout/compose_layout.cc
@@ -246,7 +246,7 @@ ffi::Map<ffi::String, PrimExpr> ApplyStructured(const
ComposeLayoutNode* layout,
}
for (const auto& [axis, offset] : tile->offset) {
- if (axis->name != "m") return fallback();
+ if (axis.name() != "m") return fallback();
std::vector<PrimExpr> dynamic_terms;
bool valid = true;
CollectOffsetTerms(offset, 1, &dynamic_terms, &constant, &valid, analyzer);
diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc
index e6681adee1..e3f4dedff4 100644
--- a/src/tirx/ir/layout/tile_core.cc
+++ b/src/tirx/ir/layout/tile_core.cc
@@ -212,7 +212,7 @@ bool TileLayoutNode::VerifyWellFormed() const {
// std::unordered_map<String, std::vector<Iter>> thread_axes;
// auto collect_thread_axis = [&thread_axes](const Iter& iter) {
// if (iter->axis->IsThreadAxis()) {
- // thread_axes[iter->axis->name].push_back(iter);
+ // thread_axes[iter->axis.name()].push_back(iter);
// }
// };
// for (const auto& iter : shard) {
@@ -238,7 +238,7 @@ bool TileLayoutNode::VerifyWellFormed() const {
PrimExpr TileLayoutNode::GetSize(ffi::Optional<ffi::String> axis_name) const {
auto filter = [&](const Iter& iter, PrimExpr acc) {
- if (!axis_name.has_value() || iter->axis->name == axis_name.value()) {
+ if (!axis_name.has_value() || iter->axis.name() == axis_name.value()) {
return acc * iter->extent;
}
return acc;
@@ -314,20 +314,20 @@ ffi::Map<ffi::String, PrimExpr>
TileLayoutNode::Apply(Array<PrimExpr> coord) con
<< "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);
+ auto it = result.find(shard[i]->axis.name());
if (it == result.end()) {
- result[shard[i]->axis->name] = analyzer->Simplify(coord[i] *
shard[i]->stride);
+ result[shard[i]->axis.name()] = analyzer->Simplify(coord[i] *
shard[i]->stride);
} else {
- result[shard[i]->axis->name] = analyzer->Simplify(it->second + coord[i]
* shard[i]->stride);
+ result[shard[i]->axis.name()] = analyzer->Simplify(it->second + coord[i]
* shard[i]->stride);
}
}
// Add offset to the result
for (const auto& [axis, off] : offset) {
- auto it = result.find(axis->name);
+ auto it = result.find(axis.name());
if (it == result.end()) {
- result[axis->name] = analyzer->Simplify(off);
+ result[axis.name()] = analyzer->Simplify(off);
} else {
- result[axis->name] = analyzer->Simplify(it->second + off);
+ result[axis.name()] = analyzer->Simplify(it->second + off);
}
}
return result;
@@ -387,7 +387,7 @@ ffi::Optional<ffi::Tuple<ExecScope, ExecScope>>
TileLayoutNode::GetScope() const
auto subtile_primitivet = axis->GetSubscope();
auto tile_primitivet = axis->GetScope();
TVM_FFI_ICHECK(subtile_primitivet.has_value() &&
tile_primitivet.has_value())
- << "Thread axis " << axis->name << " has no subscope or scope";
+ << "Thread axis " << axis.name() << " has no subscope or scope";
ffi::String subscope = subtile_primitivet.value()->name();
ffi::String scope = tile_primitivet.value()->name();
diff --git a/src/tirx/ir/layout/tile_tile_ops.cc
b/src/tirx/ir/layout/tile_tile_ops.cc
index 22450684f1..69c9f54afa 100644
--- a/src/tirx/ir/layout/tile_tile_ops.cc
+++ b/src/tirx/ir/layout/tile_tile_ops.cc
@@ -343,7 +343,7 @@ Layout TileLayoutNode::Tile(const TileLayout& outer_in,
const Array<PrimExpr>& o
auto inner_span_map = BuildSpanMap(inner);
std::vector<Iter> new_shard;
for (size_t i = 0; i < outer->shard.size(); ++i) {
- auto it = inner_span_map.find(outer->shard[i]->axis->name);
+ auto it = inner_span_map.find(outer->shard[i]->axis.name());
if (it != inner_span_map.end()) {
new_shard.push_back(Iter(outer->shard[i]->extent,
outer->shard[i]->stride * (*it).second,
outer->shard[i]->axis));
@@ -493,7 +493,7 @@ ffi::Optional<TileLayout> TileLayoutNode::IsTileInner(
// Get the span map of the inner layout of each axis
auto inner_span_map = BuildSpanMap(layout);
auto rescale_by_inner_span = [&](const Iter& iter) -> ffi::Optional<Iter> {
- auto it = inner_span_map.find(iter->axis->name);
+ auto it = inner_span_map.find(iter->axis.name());
if (it != inner_span_map.end() && !is_one(iter->extent)) {
if (!analyzer->CanProveEqual(floormod(iter->stride, (*it).second), 0)) {
return std::nullopt;
diff --git a/src/tirx/ir/layout/utils.cc b/src/tirx/ir/layout/utils.cc
index fe41905028..988c168549 100644
--- a/src/tirx/ir/layout/utils.cc
+++ b/src/tirx/ir/layout/utils.cc
@@ -44,8 +44,8 @@ PrimExpr FlattenCoord(const Array<PrimExpr>& coord, const
Array<PrimExpr>& shape
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));
+ if (span_map.find(iter->axis.name()) == span_map.end()) {
+ span_map.Set(iter->axis.name(), layout->GetSpan(iter->axis.name()));
}
}
return span_map;
@@ -74,7 +74,7 @@ std::vector<PrimExpr> GetDefaultStrides(const
ffi::Array<PrimExpr>& data, PrimEx
bool AxisMatchesFilter(const Axis& axis, const ffi::Optional<ffi::String>&
axis_name) {
return (!axis_name.has_value() && axis->IsMemoryAxis()) ||
- (axis_name.has_value() && axis->name == axis_name.value());
+ (axis_name.has_value() && axis.name() == axis_name.value());
}
TileLayout IdentityTileLayout(const ffi::Array<PrimExpr>& shape) {
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index 96fe33a918..76f1a93611 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -419,7 +419,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tirx::Axis>(
"", [](tirx::Axis axis, AccessPath p, IRDocsifier d) -> Doc {
- return LiteralDoc::Str(axis->name, p->Attr("name"));
+ return LiteralDoc::Str(axis.name(), p->Attr("name"));
});
}
@@ -428,7 +428,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"", [](tirx::Iter iter, AccessPath p, IRDocsifier d) -> Doc {
return TIR(d, "Iter")->Call({d->AsDoc<ExprDoc>(iter->extent,
p->Attr("extent")),
d->AsDoc<ExprDoc>(iter->stride,
p->Attr("stride")),
- d->AsDoc<ExprDoc>(iter->axis->name,
p->Attr("axis"))},
+ d->AsDoc<ExprDoc>(iter->axis.name(),
p->Attr("axis"))},
{}, {});
});
}
@@ -438,8 +438,8 @@ Doc PrintTileLayout(tirx::TileLayout layout, IRDocsifier d,
AccessPath p) {
// `value @ Axis.<name>`, but elide `@m` (the default memory axis).
auto bind_axis = [&](ExprDoc value, const tirx::Axis& axis) -> ExprDoc {
- if (axis->name == "m") return value;
- return OperationDoc(OpKind::kMatMul, {value,
IdDoc("Axis")->Attr(axis->name)});
+ if (axis.name() == "m") return value;
+ return OperationDoc(OpKind::kMatMul, {value,
IdDoc("Axis")->Attr(axis.name())});
};
// Build `head[(e0, e1, ...) : (s0@a0, s1@a1, ...)]` (or 1D shorthand
@@ -465,7 +465,7 @@ Doc PrintTileLayout(tirx::TileLayout layout, IRDocsifier d,
AccessPath p) {
if (layout->offset.size() > 0) {
ffi::Array<ExprDoc> offset_keys, offset_values;
for (const auto& [axis, off] : layout->offset) {
- offset_keys.push_back(LiteralDoc::Str(axis->name, p->Attr("axis")));
+ offset_keys.push_back(LiteralDoc::Str(axis.name(), p->Attr("axis")));
offset_values.push_back(d->AsDoc<ExprDoc>(off, p->Attr("offset")));
}
keys.push_back("offset");
@@ -496,7 +496,7 @@ Doc PrintTileLayout(tirx::TileLayout layout, IRDocsifier d,
AccessPath p) {
std::vector<std::pair<tirx::Axis, PrimExpr>>
sorted_offset(layout->offset.begin(),
layout->offset.end());
std::sort(sorted_offset.begin(), sorted_offset.end(),
- [](const auto& a, const auto& b) { return a.first->name <
b.first->name; });
+ [](const auto& a, const auto& b) { return a.first.name() <
b.first.name(); });
// Build the offset as a single arithmetic expression first, then add it
// to the spec in one `+`. Chaining `spec + term1 + term2` would re-enter
diff --git a/tests/python/test_tirx_axis_enum.py
b/tests/python/test_tirx_axis_enum.py
new file mode 100644
index 0000000000..aa0329ffc3
--- /dev/null
+++ b/tests/python/test_tirx_axis_enum.py
@@ -0,0 +1,65 @@
+# 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.
+"""Axis enum identity and import behavior across compiler and runtime
builds."""
+
+import os
+import subprocess
+import sys
+
+import pytest
+from tvm_ffi.dataclasses import Enum
+
+import tvm
+from tvm.tirx.layout import Axis, S, TileLayout, laneid
+
+
+def test_axis_enum_identity_and_attributes():
+ axis = Axis.laneid
+ assert isinstance(axis, Enum)
+ assert axis.same_as(Axis.get("laneid"))
+ assert axis.same_as(laneid)
+ assert axis.name == "laneid"
+ assert axis in Axis.all_entries()
+ assert hash(axis) == hash(Axis.get("laneid"))
+ assert Axis.def_attr("thread")[axis] is True
+ assert Axis.m.is_memory()
+ assert Axis.bx.get_scope().name == "thread"
+ assert Axis.bx.get_subscope().name == "cta"
+ assert TileLayout(S[4 : 1 @ axis]).shard[0].axis.same_as(axis)
+
+
+def test_axis_unknown_name_and_json_roundtrip():
+ axis = Axis.get("axis_enum_unknown_test")
+ assert axis.same_as(Axis.get("axis_enum_unknown_test"))
+ assert axis.get_scope() is None
+ assert axis.get_subscope() is None
+ with pytest.raises(Exception, match="has no thread classification"):
+ axis.is_thread()
+ with pytest.raises(Exception, match="has no thread classification"):
+ axis.is_memory()
+ assert tvm.ir.load_json(tvm.ir.save_json(axis)).same_as(axis)
+ assert tvm.ir.load_json(tvm.ir.save_json(Axis.laneid)).same_as(Axis.laneid)
+
+
+def test_axis_runtime_only_import():
+ env = os.environ.copy()
+ env["TVM_USE_RUNTIME_LIB"] = "1"
+ subprocess.run(
+ [sys.executable, "-c", "import tvm; import tvm.tirx.layout; assert
tvm.base._RUNTIME_ONLY"],
+ env=env,
+ check=True,
+ )
