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
The following commit(s) were added to refs/heads/tqchen/printer-refinement by
this push:
new 5905967cdf [REFACTOR][Script] Use ir namespace for printer options
5905967cdf is described below
commit 5905967cdfc10a9db3db6c61a8b26cd9a8444a12
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 1 16:43:33 2026 +0000
[REFACTOR][Script] Use ir namespace for printer options
Name printer configuration keys consistently under ir while preserving
option behavior and existing constructor arguments.
---
python/tvm/ir/base.py | 2 +-
python/tvm/runtime/script_printer.py | 14 +++++++-------
src/ir/transform.cc | 2 +-
src/script/printer/config.cc | 4 ++--
src/script/printer/doc_translator.cc | 14 +++++++-------
src/script/printer/ir/ir.cc | 4 ++--
src/script/printer/ir/prim_type.cc | 2 +-
src/script/printer/ir/utils.cc | 16 +++++++---------
src/script/printer/printer.cc | 18 +++++++++---------
src/script/printer/script_printer.cc | 2 +-
src/tirx/script/printer/function.cc | 2 +-
.../relax/script/test_relax_script_dynamic_shape.py | 2 +-
tests/python/s_tir/script/test_s_tir_script_printer.py | 4 ++--
tests/python/tirx/script/test_tirx_script_buffers.py | 2 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 4 ++--
15 files changed, 45 insertions(+), 47 deletions(-)
diff --git a/python/tvm/ir/base.py b/python/tvm/ir/base.py
index fa7bba369b..267dffee74 100644
--- a/python/tvm/ir/base.py
+++ b/python/tvm/ir/base.py
@@ -35,7 +35,7 @@ class Node(Object):
try:
return _script(
self,
- PrinterConfig(extra_config={"script.comment_imports": True}),
+ PrinterConfig(extra_config={"ir.comment_imports": True}),
)
except Exception:
return super().__repr__()
diff --git a/python/tvm/runtime/script_printer.py
b/python/tvm/runtime/script_printer.py
index 45874a8e49..9a3e25d0d7 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -33,7 +33,7 @@ class PrinterConfig(Object):
"""Configuration of TVMScript printer.
The ``int_dtype`` and ``float_dtype`` arguments populate the string-valued
- ``ir.int_dtype`` and ``script.float_dtype`` entries in ``extra_config``.
+ ``ir.int_dtype`` and ``ir.float_dtype`` entries in ``extra_config``.
Explicit entries in ``extra_config`` override those arguments.
"""
@@ -83,7 +83,7 @@ class PrinterConfig(Object):
"module_alias": module_alias,
"buffer_dtype": buffer_dtype,
"ir.int_dtype": int_dtype,
- "script.float_dtype": float_dtype,
+ "ir.float_dtype": float_dtype,
"verbose_expr": verbose_expr,
"indent_spaces": indent_spaces,
"print_line_numbers": print_line_numbers,
@@ -198,9 +198,9 @@ class Scriptable:
merged_extra.update(extra_config)
# Keep the historical interactive display header. Direct printer calls
# retain executable imports for standalone parser round-trips.
- merged_extra.setdefault("script.comment_imports", True)
- if "script.use_pep695" not in merged_extra:
- merged_extra["script.use_pep695"] = merged_extra.get(
+ 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)
)
@@ -274,8 +274,8 @@ class Scriptable:
obj_to_annotate: dict[Object, str] | None = None,
) -> str:
merged_extra = dict(extra_config or {})
- if "script.use_pep695" not in merged_extra:
- merged_extra["script.use_pep695"] = merged_extra.get(
+ 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(
diff --git a/src/ir/transform.cc b/src/ir/transform.cc
index baf088a5a9..d1fe44f357 100644
--- a/src/ir/transform.cc
+++ b/src/ir/transform.cc
@@ -315,7 +315,7 @@ std::optional<std::string> RenderScriptWithUnderline(const
ffi::ObjectRef& node,
ffi::Map<ffi::String, ffi::Any> config_dict;
// Error snippets use the same display conventions as script() and repr().
config_dict.Set("extra_config", ffi::Map<ffi::String, ffi::Any>{
- {"script.comment_imports", true},
+ {"ir.comment_imports", true},
});
config_dict.Set("path_to_underline",
ffi::Array<ffi::reflection::AccessPath>{path});
// Show enough context that a small function renders in full (no
diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc
index 518f93bb91..08462543d8 100644
--- a/src/script/printer/config.cc
+++ b/src/script/printer/config.cc
@@ -93,7 +93,7 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any>
config_dict) {
}
if (auto v = get("float_dtype")) {
n->extra_config.Set(
- "script.float_dtype",
+ "ir.float_dtype",
ffi::DLDataTypeToString(ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>())));
}
if (auto v = get("verbose_expr")) {
@@ -156,7 +156,7 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any>
config_dict) {
// Keep legacy constructor arguments as adapters; explicit extra_config
values
// take precedence. Store validated dtype strings with the existing defaults.
for (const auto& [key, fallback] :
- {std::pair{"ir.int_dtype", "int32"}, std::pair{"script.float_dtype",
"void"}}) {
+ {std::pair{"ir.int_dtype", "int32"}, std::pair{"ir.float_dtype",
"void"}}) {
ffi::String value = n->GetExtraConfig<ffi::String>(key, fallback);
n->extra_config.Set(key,
ffi::DLDataTypeToString(ffi::StringToDLDataType(value)));
}
diff --git a/src/script/printer/doc_translator.cc
b/src/script/printer/doc_translator.cc
index cdf9e0fa73..6ca872d9a4 100644
--- a/src/script/printer/doc_translator.cc
+++ b/src/script/printer/doc_translator.cc
@@ -147,7 +147,7 @@ class TranslationEngine : public DocTranslatorObj {
details::TranslationResult TranslateComplete(ffi::AnyView obj,
TranslationEngine* engine,
ffi::Map<ffi::String, ffi::Any>
extra_config) {
- auto binding_names =
engine->GetExtraConfig<ffi::Array<ffi::String>>("script.binding_names", {});
+ 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()) {
if (!extra_config.count(key)) engine->AllocId(fallback);
@@ -159,8 +159,8 @@ details::TranslationResult TranslateComplete(ffi::AnyView
obj, TranslationEngine
}
for (const ffi::String& name : binding_names) engine->AllocId(name);
if (binding_names.empty()) {
- if (auto name =
engine->GetExtraConfig<ffi::Optional<ffi::String>>("script.module_name",
-
std::nullopt)) {
+ if (auto name =
+
engine->GetExtraConfig<ffi::Optional<ffi::String>>("ir.module_name",
std::nullopt)) {
engine->AllocId(name.value());
}
}
@@ -173,7 +173,7 @@ details::TranslationResult TranslateComplete(ffi::AnyView
obj, TranslationEngine
auto docs = engine->WithDocScope([&]() {
ffi::Optional<ExprDoc> expression = engine->Translate(obj);
if (expression.has_value()) {
- compact_expression =
!engine->GetExtraConfig<bool>("script.verbose_expr", true);
+ compact_expression = !engine->GetExtraConfig<bool>("ir.verbose_expr",
true);
if (compact_expression) engine->CurrentScopeDocs().clear();
engine->Emit(ExprStmtDoc(expression.value()), obj.as<ffi::ObjectRef>());
}
@@ -205,7 +205,7 @@ details::TranslationResult TranslateComplete(ffi::AnyView
obj, TranslationEngine
using MetadataMap = ffi::Dict<ffi::String, ffi::List<ffi::Any>>;
auto metadata =
engine->GetOrCreateExtraState<MetadataMap>("ir.metadata_map");
ffi::Array<StmtDoc> body;
- if (!metadata.empty() && engine->GetExtraConfig<bool>("script.show_meta",
true)) {
+ if (!metadata.empty() && engine->GetExtraConfig<bool>("ir.show_meta", true))
{
// Serialize the whole registry as one graph so aliases shared across
// entries survive. Immutable maps/arrays preserve the loader's contract.
ffi::Map<ffi::String, ffi::Any> serialized;
@@ -226,13 +226,13 @@ details::TranslationResult TranslateComplete(ffi::AnyView
obj, TranslationEngine
std::nullopt));
}
for (const auto& stmt : ToStmtDocArray(docs)) body.push_back(stmt);
- if (!metadata.empty() && !engine->GetExtraConfig<bool>("script.show_meta",
true)) {
+ if (!metadata.empty() && !engine->GetExtraConfig<bool>("ir.show_meta",
true)) {
body.push_back(
CommentDoc("Metadata omitted. Use show_meta=True in script() method to
show it."));
}
return {StmtBlockDoc(body), engine->GetDocOrigins(),
engine->GetOrCreateExtraState<bool>("script.future_annotations"),
metadata_loader->name,
- !metadata.empty() &&
engine->GetExtraConfig<bool>("script.show_meta", true)};
+ !metadata.empty() && engine->GetExtraConfig<bool>("ir.show_meta",
true)};
}
} // namespace
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index 24c758e3a2..11982d4e1d 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -72,9 +72,9 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj*
d, ffi::AnyView in
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
IRModuleNode>(input);
TVM_FFI_CHECK(destination == nullptr, TypeError)
<< "printer statement-only node cannot fulfill a destination";
- auto names =
d->GetExtraConfig<ffi::Array<ffi::String>>("script.binding_names", {});
+ auto names = d->GetExtraConfig<ffi::Array<ffi::String>>("ir.binding_names",
{});
ffi::String module_name =
- names.empty() ? d->GetExtraConfig<ffi::String>("script.module_name",
"Module") : names.back();
+ names.empty() ? d->GetExtraConfig<ffi::String>("ir.module_name",
"Module") : names.back();
d->AllocId(module_name);
GlobalInfoScope global_infos(d, mod);
std::vector<std::pair<GlobalVar, BaseFunc>> functions(mod->functions.begin(),
diff --git a/src/script/printer/ir/prim_type.cc
b/src/script/printer/ir/prim_type.cc
index 22debd5500..335307be51 100644
--- a/src/script/printer/ir/prim_type.cc
+++ b/src/script/printer/ir/prim_type.cc
@@ -55,7 +55,7 @@ ffi::Optional<ExprDoc> FloatImmDocTranslate(DocTranslatorObj*
d, ffi::AnyView in
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FloatImmNode>(input);
DLDataType dtype = imm->ty.as_or_throw<PrimType>()->dtype;
DLDataType default_dtype =
-
ffi::StringToDLDataType(d->GetExtraConfig<ffi::String>("script.float_dtype",
"void"));
+ ffi::StringToDLDataType(d->GetExtraConfig<ffi::String>("ir.float_dtype",
"void"));
ExprDoc value = LiteralDoc::Float(imm->value, std::nullopt);
if (dtype == default_dtype) return value;
return
NamespaceDoc("tirx")->Attr(ffi::DLDataTypeToString(dtype))->Call({value});
diff --git a/src/script/printer/ir/utils.cc b/src/script/printer/ir/utils.cc
index b3f35c3d53..a2f01b05b0 100644
--- a/src/script/printer/ir/utils.cc
+++ b/src/script/printer/ir/utils.cc
@@ -62,13 +62,11 @@ IdDoc VarDoc(DocTranslatorObj* d, const Var& var, bool
explicit_def) {
return doc;
}
-bool SyntaxSugar(DocTranslatorObj* d) {
- return d->GetExtraConfig<bool>("script.syntax_sugar", true);
-}
+bool SyntaxSugar(DocTranslatorObj* d) { return
d->GetExtraConfig<bool>("ir.syntax_sugar", true); }
bool IsAnnotated(DocTranslatorObj* d, const ffi::ObjectRef& object) {
for (const auto& item :
-
d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("script.annotated_objects", {})) {
+ d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("ir.annotated_objects",
{})) {
if (item.same_as(object)) return true;
}
return false;
@@ -76,7 +74,7 @@ bool IsAnnotated(DocTranslatorObj* d, const ffi::ObjectRef&
object) {
bool IsUnderlined(DocTranslatorObj* d, const ffi::ObjectRef& object) {
for (const auto& item :
-
d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("script.underlined_objects", {}))
{
+ d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("ir.underlined_objects",
{})) {
if (item.same_as(object)) return true;
}
return false;
@@ -84,7 +82,7 @@ bool IsUnderlined(DocTranslatorObj* d, const ffi::ObjectRef&
object) {
bool HasAnnotatedDescendant(DocTranslatorObj* d, const ffi::ObjectRef& object)
{
for (const auto& item :
-
d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("script.annotation_ancestors",
{})) {
+
d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("ir.annotation_ancestors", {})) {
if (item.same_as(object)) return true;
}
return false;
@@ -103,8 +101,8 @@ ExprDoc NamedCallCallee(const ffi::String& canonical_name) {
}
ExprDoc GlobalReference(DocTranslatorObj* d, const ffi::String& name) {
- auto names =
d->GetExtraConfig<ffi::Array<ffi::String>>("script.binding_names", {});
- return IdDoc(names.empty() ?
d->GetExtraConfig<ffi::String>("script.module_name", "Module")
+ auto names = d->GetExtraConfig<ffi::Array<ffi::String>>("ir.binding_names",
{});
+ return IdDoc(names.empty() ?
d->GetExtraConfig<ffi::String>("ir.module_name", "Module")
: names.back())
->Attr(name);
}
@@ -239,7 +237,7 @@ 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>("script.use_pep695", false);
+ bool use_pep695 = !module.has_value() &&
d->GetExtraConfig<bool>("ir.use_pep695", false);
auto pending = d->GetImplicitDefs();
std::vector<Var> parameters;
if (use_pep695) {
diff --git a/src/script/printer/printer.cc b/src/script/printer/printer.cc
index 1c9b184df2..7469dafd44 100644
--- a/src/script/printer/printer.cc
+++ b/src/script/printer/printer.cc
@@ -251,14 +251,14 @@ class MapDocPaths {
ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config) {
MapDocPaths paths(obj, config);
auto extra_config = config->extra_config;
- extra_config.Set("script.binding_names", config->binding_names);
- extra_config.Set("script.show_meta", config->show_meta);
- extra_config.Set("script.verbose_expr", config->verbose_expr);
- extra_config.Set("script.syntax_sugar", config->syntax_sugar);
+ extra_config.Set("ir.binding_names", config->binding_names);
+ extra_config.Set("ir.show_meta", config->show_meta);
+ extra_config.Set("ir.verbose_expr", config->verbose_expr);
+ extra_config.Set("ir.syntax_sugar", config->syntax_sugar);
auto annotated_objects = paths.AnnotatedObjects();
- extra_config.Set("script.annotated_objects", annotated_objects);
- extra_config.Set("script.annotation_ancestors",
paths.AnnotationAncestors(annotated_objects));
- extra_config.Set("script.underlined_objects", paths.UnderlinedObjects());
+ extra_config.Set("ir.annotated_objects", annotated_objects);
+ extra_config.Set("ir.annotation_ancestors",
paths.AnnotationAncestors(annotated_objects));
+ extra_config.Set("ir.underlined_objects", paths.UnderlinedObjects());
auto result = details::TranslateWithOptions(obj, std::move(extra_config));
auto [underline_paths, annotations] = paths.Map(result.doc, result.origins);
ffi::Array<ffi::Any> header;
@@ -286,13 +286,13 @@ ffi::String Script(const ffi::ObjectRef& obj, const
PrinterConfig& config) {
// Header policy is fixed here; comments remain ordinary CommentDocs.
ffi::String import =
"from tvm.script import " + std::string(name) + " as " +
std::string(alias);
- if (config->GetExtraConfig<bool>("script.comment_imports", false) &&
+ if (config->GetExtraConfig<bool>("ir.comment_imports", false) &&
!(config->show_meta && name == "tirx")) {
header.push_back(CommentDoc(import));
} else {
header.push_back(ffi::String(std::string(import) + "\n"));
}
- if (name == "tirx" &&
config->GetExtraConfig<bool>("script.comment_imports", false)) {
+ if (name == "tirx" && config->GetExtraConfig<bool>("ir.comment_imports",
false)) {
header.push_back(CommentDoc("from tvm.tirx.layout import Axis"));
}
}
diff --git a/src/script/printer/script_printer.cc
b/src/script/printer/script_printer.cc
index 547381a0fb..878b6f2440 100644
--- a/src/script/printer/script_printer.cc
+++ b/src/script/printer/script_printer.cc
@@ -88,7 +88,7 @@ std::string Script(const ffi::ObjectRef& node, const
ffi::Optional<PrinterConfig
std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) {
try {
PrinterConfig config;
- config->extra_config.Set("script.comment_imports", true);
+ config->extra_config.Set("ir.comment_imports", true);
// Call translation directly so an unsupported type cannot recurse through
ffi repr.
return std::string(script::printer::Script(obj, config));
} catch (const tvm::ffi::Error& e) {
diff --git a/src/tirx/script/printer/function.cc
b/src/tirx/script/printer/function.cc
index 23fef1f784..7a8fdad386 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -75,7 +75,7 @@ ffi::Optional<ExprDoc> PrimFuncDocTranslate(DocTranslatorObj*
d, ffi::AnyView in
auto module =
d->GetOrCreateExtraState<ffi::Optional<ffi::Dict<ffi::String,
ffi::List<GlobalInfo>>>>(
"ir.global_info_map");
- auto names =
d->GetExtraConfig<ffi::Array<ffi::String>>("script.binding_names", {});
+ auto names = d->GetExtraConfig<ffi::Array<ffi::String>>("ir.binding_names",
{});
if (!module.has_value() && !names.empty()) {
if (auto symbol = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
TVM_FFI_CHECK(symbol.value() == names.back(), TypeError)
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 2ee370a9aa..095b64d50b 100644
--- a/tests/python/relax/script/test_relax_script_dynamic_shape.py
+++ b/tests/python/relax/script/test_relax_script_dynamic_shape.py
@@ -229,7 +229,7 @@ def test_call_tir_with_tir_var():
Y[vi] = X[vi]
_check(Module)
- portable = Module.script(show_meta=True,
extra_config={"script.use_pep695": False})
+ portable = Module.script(show_meta=True, extra_config={"ir.use_pep695":
False})
tvm.ir.assert_structural_equal(
Module,
tvm.script.from_source(
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 d721837cad..2640bed191 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,7 +46,7 @@ def test_prim_func_symbolic_buffer_param_roundtrip():
.with_attr("s_tir", True)
)
- source = func.script(extra_config={"script.use_pep695": False})
+ source = func.script(extra_config={"ir.use_pep695": False})
assert "T.Buffer((n + 1, n)" in source
assert source.index('n = I.dynamic("n", dtype="int32")') <
source.index("T.evaluate(n)")
tvm.ir.assert_structural_equal(
@@ -66,7 +66,7 @@ def
test_prim_func_compound_buffer_shape_first_use_roundtrip():
.with_attr("s_tir", True)
)
- source = func.script(extra_config={"script.use_pep695": False})
+ source = func.script(extra_config={"ir.use_pep695": False})
assert "T.Buffer((T.max(n, 1),)" in source
assert source.index('n = I.dynamic("n", dtype="int32")') <
source.index("T.evaluate(n)")
tvm.ir.assert_structural_equal(
diff --git a/tests/python/tirx/script/test_tirx_script_buffers.py
b/tests/python/tirx/script/test_tirx_script_buffers.py
index 9c7e1239e8..df5d2a557c 100644
--- a/tests/python/tirx/script/test_tirx_script_buffers.py
+++ b/tests/python/tirx/script/test_tirx_script_buffers.py
@@ -325,7 +325,7 @@ 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={"script.use_pep695": False})
+ code = func.script(extra_config={"ir.use_pep695": False})
assert 'n = I.dynamic("n", dtype="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 ddbde17128..5fa9a11aef 100644
--- a/tests/python/tirx/script/test_tirx_script_dynamic_shape.py
+++ b/tests/python/tirx/script/test_tirx_script_dynamic_shape.py
@@ -152,7 +152,7 @@ def func[M: int](A: T.Buffer((M, M * 2), "float32")):
assert 'M = I.dynamic("M", dtype="int64")' in script
assert "M = T.int64()" not in script
- portable = func.script(extra_config={"script.use_pep695": False})
+ 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
@@ -182,7 +182,7 @@ def test_dynamic_int32_roundtrip():
assert "def main[n: T.int32](" in source
else:
assert 'n = I.dynamic("n", dtype="int32")' in source
- portable = func.script(extra_config={"script.use_pep695": False})
+ portable = func.script(extra_config={"ir.use_pep695": False})
assert 'n = I.dynamic("n", dtype="int32")' in portable
tvm.ir.assert_structural_equal(
func, tvm.script.from_source(source, extra_vars={"I": tvm.script.ir,
"T": tvm.script.tirx})
