You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/01 20:01:03 UTC
(tvm) branch tqchen/printer-refinement updated (0dbebccb24 -> 3b94b1a2cf)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch tqchen/printer-refinement
in repository https://gitbox.apache.org/repos/asf/tvm.git
from 0dbebccb24 [REFACTOR][TIRx] Configure scalar buffer variable syntax in its owning printer
new c2f712c67e [REFACTOR][Script] Always emit eligible function type parameters
new 28c119e031 [REFACTOR][Script] Expose namespace aliases through the printer API
new 3b94b1a2cf [TEST][TIRx] Construct source-location resources from typed parameters
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
include/tvm/script/printer/printer.h | 11 +++++
python/tvm/runtime/script_printer.py | 9 ----
src/relax/script/printer/function.cc | 4 +-
src/s_tir/script/printer/config.cc | 5 +--
src/script/printer/config.cc | 24 +++++-----
src/script/printer/dialect_prefix.h | 41 -----------------
src/script/printer/doc_printer.cc | 6 +--
src/script/printer/doc_translator.cc | 9 ++--
src/script/printer/ir/ir.cc | 4 +-
src/script/printer/ir/utils.cc | 3 +-
src/script/printer/script_printer.cc | 2 -
src/tirx/script/printer/function.cc | 4 +-
.../script/test_relax_script_dynamic_shape.py | 45 ++++++++-----------
.../s_tir/script/test_s_tir_script_printer.py | 8 ++--
.../python/tirx/script/test_tirx_script_buffers.py | 4 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 52 +++++++++-------------
.../script/test_tirx_script_source_locations.py | 30 +++++++------
17 files changed, 105 insertions(+), 156 deletions(-)
delete mode 100644 src/script/printer/dialect_prefix.h
(tvm) 01/03: [REFACTOR][Script] Always emit eligible function type parameters
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tqchen/printer-refinement
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit c2f712c67eb8116223cbbc02a10a07d120058d92
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 19:49:22 2026 +0000
[REFACTOR][Script] Always emit eligible function type parameters
---
python/tvm/runtime/script_printer.py | 9 ----
src/script/printer/ir/utils.cc | 3 +-
.../script/test_relax_script_dynamic_shape.py | 45 ++++++++-----------
.../s_tir/script/test_s_tir_script_printer.py | 8 ++--
.../python/tirx/script/test_tirx_script_buffers.py | 4 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 52 +++++++++-------------
6 files changed, 48 insertions(+), 73 deletions(-)
diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py
index 4fe8739dd8..2fe649927e 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -17,7 +17,6 @@
"""Configuration of TVMScript printer"""
import os
-import sys
from collections.abc import Sequence
from tvm_ffi import get_global_func, register_object
@@ -203,10 +202,6 @@ class Scriptable:
# Keep the historical interactive display header. Direct printer calls
# retain executable imports for standalone parser round-trips.
merged_extra.setdefault("ir.comment_imports", True)
- if "ir.use_pep695" not in merged_extra:
- merged_extra["ir.use_pep695"] = merged_extra.get(
- "relax.use_pep695", sys.version_info >= (3, 12)
- )
# Only auto-switch if the caller has not already set a tirx.prefix override.
if "tirx.prefix" not in merged_extra:
@@ -278,10 +273,6 @@ class Scriptable:
obj_to_annotate: dict[Object, str] | None = None,
) -> str:
merged_extra = dict(extra_config or {})
- if "ir.use_pep695" not in merged_extra:
- merged_extra["ir.use_pep695"] = merged_extra.get(
- "relax.use_pep695", sys.version_info >= (3, 12)
- )
return _relax_script(
self,
PrinterConfig(
diff --git a/src/script/printer/ir/utils.cc b/src/script/printer/ir/utils.cc
index a2f01b05b0..4674395bcb 100644
--- a/src/script/printer/ir/utils.cc
+++ b/src/script/printer/ir/utils.cc
@@ -237,10 +237,9 @@ void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var, IdDoc
auto module =
d->GetOrCreateExtraState<ffi::Optional<ffi::Dict<ffi::String, ffi::List<GlobalInfo>>>>(
"ir.global_info_map");
- bool use_pep695 = !module.has_value() && d->GetExtraConfig<bool>("ir.use_pep695", false);
auto pending = d->GetImplicitDefs();
std::vector<Var> parameters;
- if (use_pep695) {
+ if (!module.has_value()) {
for (const auto& [var, id] : signature) {
if (pending.count(var) && var->ty.as<PrimTypeNode>()) parameters.push_back(var);
}
diff --git a/tests/python/relax/script/test_relax_script_dynamic_shape.py b/tests/python/relax/script/test_relax_script_dynamic_shape.py
index 095b64d50b..4051941cfb 100644
--- a/tests/python/relax/script/test_relax_script_dynamic_shape.py
+++ b/tests/python/relax/script/test_relax_script_dynamic_shape.py
@@ -19,8 +19,6 @@
from __future__ import annotations
-import sys
-
import pytest
import tvm
@@ -46,31 +44,26 @@ def test_type_vars_roundtrip():
return x
script = func.script()
- if sys.version_info >= (3, 12):
- assert script.startswith("from __future__ import annotations\n\n")
- assert "def main[M](" in script
- assert 'R.Tensor((M, M * 2), dtype="float32")' in script
- assert "M = T.int64()" not in script
- typed = tvm.script.from_source(
- """
+ assert script.startswith("from __future__ import annotations\n\n")
+ assert "def main[M](" in script
+ assert 'R.Tensor((M, M * 2), dtype="float32")' in script
+ assert "M = T.int64()" not in script
+ typed = tvm.script.from_source(
+ """
@R.function(private=True)
def func[M: int](x: R.Tensor((M, M * 2), "float32")):
return x
""",
- extra_vars={"I": tvm.script.ir, "R": tvm.script.relax},
- )
- tvm.ir.assert_structural_equal(func, typed)
- else:
- assert "from __future__ import annotations" not in script
- assert 'M = I.dynamic("M", dtype="int64")' in script
- assert "M = T.int64()" not in script
- assert 'R.Tensor((M, M * 2), dtype="float32")' in script
-
- portable = func.script(extra_config={"relax.use_pep695": False})
- assert "from __future__ import annotations" not in portable
- assert 'M = I.dynamic("M", dtype="int64")' in portable
- assert 'R.Tensor((M, M * 2), dtype="float32")' in portable
- assert "M = T.int64()" not in portable
+ extra_vars={"I": tvm.script.ir, "R": tvm.script.relax},
+ )
+ tvm.ir.assert_structural_equal(func, typed)
+ assert 'M = I.dynamic("M", dtype="int64")' not in script
+
+ repeated = func.script()
+ assert "from __future__ import annotations" in repeated
+ assert "def main[M](" in repeated
+ assert 'R.Tensor((M, M * 2), dtype="float32")' in repeated
+ assert "M = T.int64()" not in repeated
assert "UNUSED_GENERIC" not in script
assert [param.name for param in func.params] == ["x"]
assert not hasattr(func, "type_params")
@@ -80,7 +73,7 @@ def func[M: int](x: R.Tensor((M, M * 2), "float32")):
)
tvm.ir.assert_structural_equal(
func,
- tvm.script.from_source(portable, extra_vars={"I": tvm.script.ir, "R": tvm.script.relax}),
+ tvm.script.from_source(repeated, extra_vars={"I": tvm.script.ir, "R": tvm.script.relax}),
)
@@ -229,11 +222,11 @@ def test_call_tir_with_tir_var():
Y[vi] = X[vi]
_check(Module)
- portable = Module.script(show_meta=True, extra_config={"ir.use_pep695": False})
+ repeated = Module.script(show_meta=True)
tvm.ir.assert_structural_equal(
Module,
tvm.script.from_source(
- portable,
+ repeated,
extra_vars={
"I": tvm.script.ir,
"R": tvm.script.relax,
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py b/tests/python/s_tir/script/test_s_tir_script_printer.py
index 2640bed191..89c5a9dcda 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -46,9 +46,9 @@ def test_prim_func_symbolic_buffer_param_roundtrip():
.with_attr("s_tir", True)
)
- source = func.script(extra_config={"ir.use_pep695": False})
+ source = func.script()
assert "T.Buffer((n + 1, n)" in source
- assert source.index('n = I.dynamic("n", dtype="int32")') < source.index("T.evaluate(n)")
+ assert source.index("def main[n: T.int32](") < source.index("T.evaluate(n)")
tvm.ir.assert_structural_equal(
tvm.script.from_source(
source, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx, "Ts": tvm.script.s_tir}
@@ -66,9 +66,9 @@ def test_prim_func_compound_buffer_shape_first_use_roundtrip():
.with_attr("s_tir", True)
)
- source = func.script(extra_config={"ir.use_pep695": False})
+ source = func.script()
assert "T.Buffer((T.max(n, 1),)" in source
- assert source.index('n = I.dynamic("n", dtype="int32")') < source.index("T.evaluate(n)")
+ assert source.index("def main[n: T.int32](") < source.index("T.evaluate(n)")
tvm.ir.assert_structural_equal(
tvm.script.from_source(
source, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx, "Ts": tvm.script.s_tir}
diff --git a/tests/python/tirx/script/test_tirx_script_buffers.py b/tests/python/tirx/script/test_tirx_script_buffers.py
index df5d2a557c..152c94ddd6 100644
--- a/tests/python/tirx/script/test_tirx_script_buffers.py
+++ b/tests/python/tirx/script/test_tirx_script_buffers.py
@@ -325,8 +325,8 @@ def test_buffer_shape_repeated_var_prints_out_of_line():
buffer = tvm.tirx.decl_buffer((n + n,), name="A")
func = tvm.tirx.PrimFunc([buffer], tvm.tirx.Evaluate(0))
- code = func.script(extra_config={"ir.use_pep695": False})
- assert 'n = I.dynamic("n", dtype="int32")' in code
+ code = func.script()
+ assert "def main[n: T.int32](" in code
assert_structural_equal(func, from_source(code))
diff --git a/tests/python/tirx/script/test_tirx_script_dynamic_shape.py b/tests/python/tirx/script/test_tirx_script_dynamic_shape.py
index 5fa9a11aef..3278d2758e 100644
--- a/tests/python/tirx/script/test_tirx_script_dynamic_shape.py
+++ b/tests/python/tirx/script/test_tirx_script_dynamic_shape.py
@@ -19,8 +19,6 @@
from __future__ import annotations
-import sys
-
import pytest
import tvm
@@ -133,31 +131,27 @@ def test_type_vars_roundtrip():
A[0, 0] = T.float32(1)
script = func.script()
- if sys.version_info >= (3, 12):
- assert script.startswith("from __future__ import annotations\n\n")
- assert "def main[M](" in script
- assert 'T.Buffer((M, M * T.int64(2)), "float32", layout="default")' in script
- assert "M = T.int64()" not in script
- typed = tvm.script.from_source(
- """
+ assert script.startswith("from __future__ import annotations\n\n")
+ assert "def main[M](" in script
+ assert 'T.Buffer((M, M * T.int64(2)), "float32", layout="default")' in script
+ assert "M = T.int64()" not in script
+ typed = tvm.script.from_source(
+ """
@T.prim_func(private=True)
def func[M: int](A: T.Buffer((M, M * 2), "float32")):
A[0, 0] = T.float32(1)
""",
- extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx},
- )
- tvm.ir.assert_structural_equal(func, typed)
- else:
- assert "from __future__ import annotations" not in script
- assert 'M = I.dynamic("M", dtype="int64")' in script
- assert "M = T.int64()" not in script
-
- portable = func.script(extra_config={"ir.use_pep695": False})
- assert "from __future__ import annotations" not in portable
- assert 'M = I.dynamic("M", dtype="int64")' in portable
- assert 'T.Buffer((M, M * T.int64(2)), "float32", layout="default")' in portable
+ extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx},
+ )
+ tvm.ir.assert_structural_equal(func, typed)
+ assert 'M = I.dynamic("M", dtype="int64")' not in script
+
+ repeated = func.script()
+ assert "from __future__ import annotations" in repeated
+ assert "def main[M](" in repeated
+ assert 'T.Buffer((M, M * T.int64(2)), "float32", layout="default")' in repeated
assert "UNUSED" not in script
- assert "M = T.int64()" not in portable
+ assert "M = T.int64()" not in repeated
assert len(func.params) == 1
assert not hasattr(func, "type_params")
assert func.attrs.get("tirx.type_vars") is None
@@ -166,7 +160,7 @@ def func[M: int](A: T.Buffer((M, M * 2), "float32")):
)
tvm.ir.assert_structural_equal(
func,
- tvm.script.from_source(portable, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}),
+ tvm.script.from_source(repeated, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}),
)
@@ -178,18 +172,16 @@ def test_dynamic_int32_roundtrip():
A[0] = T.float32(1)
source = func.script()
- if sys.version_info >= (3, 12):
- assert "def main[n: T.int32](" in source
- else:
- assert 'n = I.dynamic("n", dtype="int32")' in source
- portable = func.script(extra_config={"ir.use_pep695": False})
- assert 'n = I.dynamic("n", dtype="int32")' in portable
+ assert "def main[n: T.int32](" in source
+ assert 'n = I.dynamic("n", dtype="int32")' not in source
+ repeated = func.script()
+ assert "def main[n: T.int32](" in repeated
tvm.ir.assert_structural_equal(
func, tvm.script.from_source(source, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx})
)
tvm.ir.assert_structural_equal(
func,
- tvm.script.from_source(portable, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}),
+ tvm.script.from_source(repeated, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}),
)
(tvm) 03/03: [TEST][TIRx] Construct source-location resources from typed parameters
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tqchen/printer-refinement
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 3b94b1a2cf17ca5505272fe0d0c95fc7b0a0bc7f
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 19:54:16 2026 +0000
[TEST][TIRx] Construct source-location resources from typed parameters
---
.../script/test_tirx_script_source_locations.py | 30 ++++++++++++----------
1 file changed, 17 insertions(+), 13 deletions(-)
diff --git a/tests/python/tirx/script/test_tirx_script_source_locations.py b/tests/python/tirx/script/test_tirx_script_source_locations.py
index 87cec0809c..642676e4ab 100644
--- a/tests/python/tirx/script/test_tirx_script_source_locations.py
+++ b/tests/python/tirx/script/test_tirx_script_source_locations.py
@@ -246,8 +246,7 @@ def test_native_view_keeps_producer_identity_name_and_span(monkeypatch):
from tvm import ir, tirx
from tvm.script.ir_builder import base
- captured = T.buffer((4, 4), "float32")
- original = type(captured).view
+ original = ir.Var.view
seen, produced, observed = [], [], []
span = ir.Span(ir.SourceName("producer.py"), 7, 7, 2, 19)
@@ -265,17 +264,18 @@ def test_native_view_keeps_producer_identity_name_and_span(monkeypatch):
def observe(value):
observed.append((value, value.span))
- monkeypatch.setattr(type(captured), "view", view)
+ monkeypatch.setattr(ir.Var, "view", view)
@T.prim_func
- def main(A: captured):
- renamed = captured.view(mark())
+ def main(A: T.Buffer((4, 4), "float32")):
+ renamed = A.view(mark())
alias = renamed
observe(renamed)
observe(alias)
alias[0] = 0
- captured.view(mark())
+ A.view(mark())
+ captured = main.params[0]
assert seen == ["argument", "view", "argument", "view"]
value, name, produced_span = produced[0]
assert len(produced) == len(observed) == 2
@@ -299,8 +299,6 @@ def test_native_binding_preserves_metadata_but_binds_buffer_expressions():
from tvm.script.ir_builder import base
producer_span = ir.Span(ir.SourceName("producer.py"), 7, 7, 2, 19)
- buffer = T.buffer((4,), "float32")
- base.at_(producer_span, buffer)
layout = T.TileLayout(T.S[4])
@T.meta_class
@@ -308,13 +306,16 @@ def test_native_binding_preserves_metadata_but_binds_buffer_expressions():
def __init__(self, resource):
self.resource = resource
- holder = Holder(buffer)
- projection = ir.TupleGetItem(ir.Tuple([buffer]), 0)
- values = [layout, holder, projection]
+ values = []
+
+ def initialize(buffer):
+ base.at_(producer_span, buffer)
+ values.extend([layout, Holder(buffer), ir.TupleGetItem(ir.Tuple([buffer]), 0)])
+
calls, observed, resource_spans = [], [], []
def make(index):
- resource_spans.append(buffer.span)
+ resource_spans.append(values[1].resource.span)
calls.append(index)
return values[index]
@@ -322,7 +323,8 @@ def test_native_binding_preserves_metadata_but_binds_buffer_expressions():
observed.extend(items)
@T.prim_func
- def main(A: buffer):
+ def main(A: T.Buffer((4,), "float32")):
+ initialize(A)
renamed_layout = make(0)
renamed_holder = make(1)
bound = make(2)
@@ -330,6 +332,8 @@ def test_native_binding_preserves_metadata_but_binds_buffer_expressions():
make(0)
make(1)
+ buffer = main.params[0]
+ holder, projection = values[1:]
assert calls == [0, 1, 2, 0, 1]
assert observed[0].same_as(layout)
assert observed[1] is holder and holder.resource.same_as(buffer)
(tvm) 02/03: [REFACTOR][Script] Expose namespace aliases through the printer API
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tqchen/printer-refinement
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 28c119e031b5e6136bdad6da8579d19234dc5cb3
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 19:53:33 2026 +0000
[REFACTOR][Script] Expose namespace aliases through the printer API
---
include/tvm/script/printer/printer.h | 11 ++++++++++
src/relax/script/printer/function.cc | 4 ++--
src/s_tir/script/printer/config.cc | 5 ++---
src/script/printer/config.cc | 24 +++++++++++----------
src/script/printer/dialect_prefix.h | 41 ------------------------------------
src/script/printer/doc_printer.cc | 6 ++----
src/script/printer/doc_translator.cc | 9 +++++---
src/script/printer/ir/ir.cc | 4 ++--
src/script/printer/script_printer.cc | 2 --
src/tirx/script/printer/function.cc | 4 ++--
10 files changed, 40 insertions(+), 70 deletions(-)
diff --git a/include/tvm/script/printer/printer.h b/include/tvm/script/printer/printer.h
index 83f313013b..22d4c52e56 100644
--- a/include/tvm/script/printer/printer.h
+++ b/include/tvm/script/printer/printer.h
@@ -23,6 +23,7 @@
#ifndef TVM_SCRIPT_PRINTER_PRINTER_H_
#define TVM_SCRIPT_PRINTER_PRINTER_H_
+#include <tvm/ffi/container/map.h>
#include <tvm/ffi/optional.h>
#include <tvm/script/printer/config.h>
@@ -43,6 +44,16 @@ TVM_DLL std::string Script(const ffi::ObjectRef& node,
namespace script {
namespace printer {
+/*!
+ * \brief Register a namespace alias during dialect static initialization.
+ * \param key The existing prefix configuration key, such as "tirx.prefix".
+ * \param default_alias The alias reserved before translation assigns variable names.
+ */
+TVM_DLL void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& default_alias);
+
+/*! \brief Read the registered namespace aliases. */
+TVM_DLL const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases();
+
/*!
* \brief Translate IR, recover diagnostic paths, and render Python text.
* \param obj The input IR object.
diff --git a/src/relax/script/printer/function.cc b/src/relax/script/printer/function.cc
index 8913af7eca..42cba66e0a 100644
--- a/src/relax/script/printer/function.cc
+++ b/src/relax/script/printer/function.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ir/function.h>
+#include <tvm/script/printer/printer.h>
#include <tvm/tirx/type.h>
#include <algorithm>
@@ -25,7 +26,6 @@
#include <utility>
#include <vector>
-#include "../../../script/printer/dialect_prefix.h"
#include "../../../script/printer/ir/utils.h"
#include "utils.h"
@@ -116,7 +116,7 @@ ffi::Optional<ExprDoc> FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in
}
TVM_FFI_STATIC_INIT_BLOCK() {
- RegisterDialectPrefix("relax.prefix", "R");
+ RegisterNamespaceAlias("relax.prefix", "R");
ffi::reflection::TypeAttrDef<relax::FunctionNode>().attr(
kDocTranslate, FDocTranslate::FromNative<&FunctionDocTranslate>());
}
diff --git a/src/s_tir/script/printer/config.cc b/src/s_tir/script/printer/config.cc
index 461858edad..c0e135d7de 100644
--- a/src/s_tir/script/printer/config.cc
+++ b/src/s_tir/script/printer/config.cc
@@ -17,14 +17,13 @@
* under the License.
*/
#include <tvm/ffi/reflection/registry.h>
-
-#include "../../../script/printer/dialect_prefix.h"
+#include <tvm/script/printer/printer.h>
namespace tvm {
namespace script {
namespace printer {
-TVM_FFI_STATIC_INIT_BLOCK() { RegisterDialectPrefix("s_tir.prefix", "Ts"); }
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterNamespaceAlias("s_tir.prefix", "Ts"); }
} // namespace printer
} // namespace script
diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc
index 08462543d8..15a1d8bf1a 100644
--- a/src/script/printer/config.cc
+++ b/src/script/printer/config.cc
@@ -17,22 +17,21 @@
* under the License.
*/
#include <tvm/script/printer/config.h>
+#include <tvm/script/printer/printer.h>
#include <algorithm>
#include <cctype>
-#include <map>
#include <string>
#include <unordered_set>
#include <utility>
-
-#include "./dialect_prefix.h"
+#include <vector>
namespace tvm {
namespace {
-std::map<ffi::String, ffi::String>& DialectPrefixes() {
- static std::map<ffi::String, ffi::String> prefixes;
- return prefixes;
+ffi::Map<ffi::String, ffi::String>& NamespaceAliases() {
+ static ffi::Map<ffi::String, ffi::String> aliases;
+ return aliases;
}
bool IsIdentifier(const std::string& name) {
@@ -54,12 +53,12 @@ bool IsIdentifier(const std::string& name) {
namespace script {
namespace printer {
-void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix) {
- TVM_FFI_ICHECK(DialectPrefixes().emplace(key, default_prefix).second)
- << "Duplicate printer dialect prefix: " << key;
+void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& default_alias) {
+ TVM_FFI_ICHECK(!NamespaceAliases().count(key)) << "Duplicate printer dialect prefix: " << key;
+ NamespaceAliases().Set(key, default_alias);
}
-const std::map<ffi::String, ffi::String>& GetDialectPrefixes() { return DialectPrefixes(); }
+const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases() { return NamespaceAliases(); }
} // namespace printer
} // namespace script
@@ -172,7 +171,10 @@ ffi::Array<ffi::String> PrinterConfigNode::GetBuiltinKeywords() {
TVM_FFI_ICHECK(module_alias.empty() || IsIdentifier(std::string(module_alias)))
<< "Invalid `module_alias`: " << module_alias;
ffi::Array<ffi::String> result{ir_prefix};
- for (const auto& [key, default_prefix] : DialectPrefixes()) {
+ const auto& aliases = NamespaceAliases();
+ std::vector<std::pair<ffi::String, ffi::String>> ordered(aliases.begin(), aliases.end());
+ std::sort(ordered.begin(), ordered.end());
+ for (const auto& [key, default_prefix] : ordered) {
ffi::String prefix = GetExtraConfig<ffi::String>(key, default_prefix);
TVM_FFI_ICHECK(IsIdentifier(std::string(prefix))) << "Invalid `" << key << "`: " << prefix;
result.push_back(prefix);
diff --git a/src/script/printer/dialect_prefix.h b/src/script/printer/dialect_prefix.h
deleted file mode 100644
index 0eac91b5ea..0000000000
--- a/src/script/printer/dialect_prefix.h
+++ /dev/null
@@ -1,41 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-#ifndef TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
-#define TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
-
-#include <tvm/ffi/string.h>
-
-#include <map>
-
-namespace tvm {
-namespace script {
-namespace printer {
-
-// Register during dialect static initialization so configuration can validate
-// and reserve the prefix before translation assigns variable names.
-void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix);
-
-// Read the existing registration table without introducing a second namespace registry.
-const std::map<ffi::String, ffi::String>& GetDialectPrefixes();
-
-} // namespace printer
-} // namespace script
-} // namespace tvm
-
-#endif // TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
diff --git a/src/script/printer/doc_printer.cc b/src/script/printer/doc_printer.cc
index 302ba6d069..37080b6c4d 100644
--- a/src/script/printer/doc_printer.cc
+++ b/src/script/printer/doc_printer.cc
@@ -24,6 +24,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/object_functor.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/script/printer/printer.h>
#include <algorithm>
#include <cmath>
@@ -42,7 +43,6 @@
#include "../../support/str_escape.h"
#include "../../support/utils.h"
-#include "dialect_prefix.h"
namespace tvm {
namespace script {
@@ -1158,9 +1158,7 @@ void PythonDocPrinter::PrintTypedDoc(const IdDoc& doc) { output_ << doc->name; }
void PythonDocPrinter::PrintTypedDoc(const NamespaceDoc& doc) {
const ffi::String& name = doc->canonical_name;
ffi::String key = std::string(name) + ".prefix";
- const auto& prefixes = GetDialectPrefixes();
- auto registered = prefixes.find(key);
- ffi::String fallback = registered == prefixes.end() ? name : registered->second;
+ ffi::String fallback = GetNamespaceAliases().Get(key).value_or(name);
output_ << (name == "ir" ? config()->ir_prefix
: config()->GetExtraConfig<ffi::String>(key, fallback));
}
diff --git a/src/script/printer/doc_translator.cc b/src/script/printer/doc_translator.cc
index 6ca872d9a4..96bc70a7ee 100644
--- a/src/script/printer/doc_translator.cc
+++ b/src/script/printer/doc_translator.cc
@@ -24,6 +24,7 @@
#include <tvm/runtime/base.h>
#include <tvm/script/printer/doc.h>
#include <tvm/script/printer/doc_translator.h>
+#include <tvm/script/printer/printer.h>
#include <algorithm>
#include <cctype>
@@ -32,8 +33,6 @@
#include <utility>
#include <vector>
-#include "dialect_prefix.h"
-
namespace tvm {
namespace script {
namespace printer {
@@ -149,7 +148,11 @@ details::TranslationResult TranslateComplete(ffi::AnyView obj, TranslationEngine
ffi::Map<ffi::String, ffi::Any> extra_config) {
auto binding_names = engine->GetExtraConfig<ffi::Array<ffi::String>>("ir.binding_names", {});
// Reserve the same registered or explicitly configured aliases used by rendering.
- for (const auto& [key, fallback] : GetDialectPrefixes()) {
+ // Keep key order so overlapping aliases reserve the same generated names.
+ const auto& aliases = GetNamespaceAliases();
+ std::vector<std::pair<ffi::String, ffi::String>> ordered(aliases.begin(), aliases.end());
+ std::sort(ordered.begin(), ordered.end());
+ for (const auto& [key, fallback] : ordered) {
if (!extra_config.count(key)) engine->AllocId(fallback);
}
for (const auto& [key, value] : extra_config) {
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index 11982d4e1d..5e561b421b 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -20,6 +20,7 @@
#include <tvm/ir/prim/op.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
+#include <tvm/script/printer/printer.h>
#include <tvm/tirx/function.h>
#include <algorithm>
@@ -27,7 +28,6 @@
#include <utility>
#include <vector>
-#include "../dialect_prefix.h"
#include "utils.h"
namespace tvm {
@@ -161,7 +161,7 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView in
}
TVM_FFI_STATIC_INIT_BLOCK() {
- RegisterDialectPrefix("ir.prefix", "I");
+ RegisterNamespaceAlias("ir.prefix", "I");
ffi::reflection::TypeAttrDef<IRModuleNode>().attr(
kDocTranslate, FDocTranslate::FromNative<&IRModuleDocTranslate>());
}
diff --git a/src/script/printer/script_printer.cc b/src/script/printer/script_printer.cc
index 878b6f2440..56d831ad55 100644
--- a/src/script/printer/script_printer.cc
+++ b/src/script/printer/script_printer.cc
@@ -42,8 +42,6 @@
#include <sstream>
#include <utility>
-#include "./dialect_prefix.h"
-
namespace tvm {
namespace {
diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc
index 7a8fdad386..4a043bd1da 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -21,6 +21,7 @@
#include <tvm/ir/global_info.h>
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/script/printer/printer.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt_functor.h>
@@ -31,7 +32,6 @@
#include <utility>
#include <vector>
-#include "../../../script/printer/dialect_prefix.h"
#include "../../../script/printer/ir/utils.h"
#include "utils.h"
@@ -239,7 +239,7 @@ ffi::Optional<ExprDoc> PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView in
}
TVM_FFI_STATIC_INIT_BLOCK() {
- RegisterDialectPrefix("tirx.prefix", "T");
+ RegisterNamespaceAlias("tirx.prefix", "T");
ffi::reflection::TypeAttrDef<tirx::PrimFuncNode>().attr(
kDocTranslate, FDocTranslate::FromNative<&PrimFuncDocTranslate>());
}