This is an automated email from the ASF dual-hosted git repository.
tlopex 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 cd2c4f3ef2 [TIRx] Reuse pass-through input names for inverse index map
vars (#19906)
cd2c4f3ef2 is described below
commit cd2c4f3ef271536b6980fff76bb8b3daee5b6154
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Fri Jul 10 07:20:19 2026 +0800
[TIRx] Reuse pass-through input names for inverse index map vars (#19906)
## Why
Inverse index map input variables were always named axis0, axis1, ...,
making generated IR harder to read.
## How
- Name a pass-through inverse input after its source Var instead of
axis{i}.
- Keep the axis{i} fallback for computed (non-pass-through) indices.
- Add test_inverse_preserves_passthrough_var_names covering a transpose
map.
Signed-off-by: Guan-Ming (Wesley) Chiu
<[email protected]>
---
src/tirx/ir/index_map.cc | 18 ++++++++++--------
tests/python/tirx-base/test_tir_index_map.py | 6 ++++++
2 files changed, 16 insertions(+), 8 deletions(-)
diff --git a/src/tirx/ir/index_map.cc b/src/tirx/ir/index_map.cc
index 5ae7af4b4b..02e156c50d 100644
--- a/src/tirx/ir/index_map.cc
+++ b/src/tirx/ir/index_map.cc
@@ -82,14 +82,16 @@ std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const
IndexMap& self,
ffi::Array<PrimVar> output_vars;
for (size_t i = 0; i < self->final_indices.size(); i++) {
PrimExpr index = self->final_indices[i];
- // TODO(Lunderberg): Better names for these variables. A variable
- // that is passed through unmodified (`index` is an element of
- // `initial_indices`) should use that input index's name. A pair
- // of output indices variables split from a single input index
- // should be named (X.outer,X.inner).
- std::stringstream ss;
- ss << "axis" << i;
- PrimVar var_index(ss.str(), index.ty());
+ // A pass-through index keeps the input variable's name for readability.
+ // TODO(Lunderberg): Name a pair of output indices split from a single
+ // input index as (X.outer, X.inner).
+ std::string name;
+ if (auto* var = index.as<VarNode>()) {
+ name = var->name_hint;
+ } else {
+ name = "axis" + std::to_string(i);
+ }
+ PrimVar var_index(name, index.ty());
output_vars.push_back(var_index);
}
diff --git a/tests/python/tirx-base/test_tir_index_map.py
b/tests/python/tirx-base/test_tir_index_map.py
index 2e2d6fa223..5443bb4851 100644
--- a/tests/python/tirx-base/test_tir_index_map.py
+++ b/tests/python/tirx-base/test_tir_index_map.py
@@ -102,6 +102,12 @@ def test_inverse():
assert index_map.inverse([16]).is_equivalent_to(expected_inverse)
+def test_inverse_preserves_passthrough_var_names():
+ index_map = IndexMap.from_func(lambda i, j: [j, i], index_dtype="int32")
+ inverse = index_map.inverse([8, 16])
+ assert [v.name for v in inverse.initial_indices] == ["j", "i"]
+
+
def test_inverse_accepts_external_analyzer():
tile = tvm.tirx.Var("tile", "int32")
index_map = IndexMap.from_func(lambda i: [i // tile, i % tile],
index_dtype="int32")
