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 22:57:56 UTC
(tvm) branch tqchen/printer-refinement updated (dcd027c9df -> 6927b91cf9)
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 dcd027c9df Keep float32 attribute constants concise
new e5799a4b59 Separate complete Doc translation from script display policy
new 10b4a14e9a [FIX][TIRx] Preserve storage synchronization calls in named syntax
new 8d9ee4f90d [REFACTOR][Script] Preserve existing CUDA load emission
new ed13ec061b [CI] Use Python 3.12 for TVMScript round trips
new dd740ff684 Keep metadata serialization and identifier display in rendering
new 58867a45e5 Read global information directly from the current module
new 3ddd9c249e Align printer expectations with explicit blocks and generic display policy
new 6927b91cf9 Keep complete output expectations aligned with renderer headers
The 8 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:
.github/actions/setup/action.yml | 10 +-
.github/workflows/main.yml | 6 +
docs/contribute/testing.rst | 5 +
docs/reference/api/python/script/printer.rst | 10 +
include/tvm/script/printer/config.h | 3 -
python/tvm/ir/base.py | 4 +-
python/tvm/relax/base_py_module.py | 2 -
python/tvm/runtime/script_printer.py | 13 -
python/tvm/tirx/op.py | 21 +-
src/backend/cuda/codegen/codegen_cuda.cc | 4 +-
src/relax/script/printer/distributed.cc | 1 -
src/relax/script/printer/global_info.cc | 9 +-
src/s_tir/schedule/error.cc | 1 -
src/script/printer/config.cc | 5 -
src/script/printer/doc_translator.cc | 79 +----
src/script/printer/doc_translator.h | 57 ----
src/script/printer/ir/ir.cc | 46 +--
src/script/printer/ir/utils.cc | 44 +--
src/script/printer/ir/utils.h | 4 -
src/script/printer/printer.cc | 330 ++++++++++++++-------
src/tirx/script/printer/expr.cc | 37 ++-
src/tirx/script/printer/function.cc | 59 +---
src/tirx/script/printer/stmt.cc | 1 -
.../relax/script/test_relax_script_printer.py | 52 +++-
tests/python/relax/test_analysis_well_formed.py | 2 +-
.../s_tir/script/test_s_tir_script_printer.py | 105 +++----
.../transform/test_s_tir_transform_thread_sync.py | 9 +-
.../python/tirx/script/test_tirx_script_buffers.py | 28 +-
.../python/tirx/script/test_tirx_script_printer.py | 31 +-
.../test_tirx_script_printer_structural_equal.py | 4 +-
.../script/test_tirx_script_printer_underlining.py | 7 +-
31 files changed, 456 insertions(+), 533 deletions(-)
delete mode 100644 src/script/printer/doc_translator.h
(tvm) 01/08: Separate complete Doc translation from script display policy
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 e5799a4b592d9e60a7910096418331dec34aff4a
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:11:44 2026 +0000
Separate complete Doc translation from script display policy
---
include/tvm/script/printer/config.h | 3 -
python/tvm/ir/base.py | 4 +-
python/tvm/relax/base_py_module.py | 2 -
python/tvm/runtime/script_printer.py | 13 -
src/relax/script/printer/distributed.cc | 1 -
src/script/printer/config.cc | 3 -
src/script/printer/doc_translator.cc | 63 +----
src/script/printer/doc_translator.h | 57 -----
src/script/printer/ir/ir.cc | 13 +-
src/script/printer/ir/utils.cc | 37 +--
src/script/printer/ir/utils.h | 4 -
src/script/printer/printer.cc | 281 ++++++++++++---------
src/tirx/script/printer/expr.cc | 2 +-
src/tirx/script/printer/function.cc | 59 +----
src/tirx/script/printer/stmt.cc | 1 -
.../s_tir/script/test_s_tir_script_printer.py | 2 +-
.../python/tirx/script/test_tirx_script_printer.py | 8 +-
.../test_tirx_script_printer_structural_equal.py | 4 +-
18 files changed, 202 insertions(+), 355 deletions(-)
diff --git a/include/tvm/script/printer/config.h b/include/tvm/script/printer/config.h
index e770de90f1..9d0a562864 100644
--- a/include/tvm/script/printer/config.h
+++ b/include/tvm/script/printer/config.h
@@ -62,8 +62,6 @@ class PrinterConfigNode : public ffi::Object {
bool print_line_numbers = false;
/*! \brief Number of context lines to print around the underlined text */
int num_context_lines = -1;
- /*! \brief Whether to output with syntax sugar, set false for complete printing. */
- bool syntax_sugar = true;
/*! \brief Whether to render access-path context for invisible underlined paths. Defaults true. */
bool render_invisible_path_info = true;
@@ -114,7 +112,6 @@ class PrinterConfigNode : public ffi::Object {
.def_ro("indent_spaces", &PrinterConfigNode::indent_spaces)
.def_ro("print_line_numbers", &PrinterConfigNode::print_line_numbers)
.def_ro("num_context_lines", &PrinterConfigNode::num_context_lines)
- .def_ro("syntax_sugar", &PrinterConfigNode::syntax_sugar)
.def_ro("render_invisible_path_info", &PrinterConfigNode::render_invisible_path_info)
.def_ro("path_to_underline", &PrinterConfigNode::path_to_underline)
.def_ro("path_to_annotate", &PrinterConfigNode::path_to_annotate)
diff --git a/python/tvm/ir/base.py b/python/tvm/ir/base.py
index 267dffee74..f572b02e37 100644
--- a/python/tvm/ir/base.py
+++ b/python/tvm/ir/base.py
@@ -198,8 +198,8 @@ def assert_structural_equal(lhs, rhs, map_free_vars=False):
# Diagnostics use the same display policy as Object.script(), including
# dialect selection and commented imports. The unbound method also
# accepts IR objects that do not inherit the convenience mixin.
- lhs_script = Scriptable.script(lhs, syntax_sugar=False, path_to_underline=[lhs_path])
- rhs_script = Scriptable.script(rhs, syntax_sugar=False, path_to_underline=[rhs_path])
+ lhs_script = Scriptable.script(lhs, path_to_underline=[lhs_path])
+ rhs_script = Scriptable.script(rhs, path_to_underline=[rhs_path])
raise ValueError(
f"StructuralEqual check failed, caused by lhs at {lhs_path}:\n"
f"{lhs_script}\n"
diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py
index 0db6a57659..6495a77e19 100644
--- a/python/tvm/relax/base_py_module.py
+++ b/python/tvm/relax/base_py_module.py
@@ -548,7 +548,6 @@ class BasePyModule:
indent_spaces: int = 4,
print_line_numbers: bool = False,
num_context_lines: int = -1,
- syntax_sugar: bool = True,
show_all_ty: bool = True,
extra_config: dict | None = None,
) -> str:
@@ -569,7 +568,6 @@ class BasePyModule:
indent_spaces=indent_spaces,
print_line_numbers=print_line_numbers,
num_context_lines=num_context_lines,
- syntax_sugar=syntax_sugar,
show_all_ty=show_all_ty,
extra_config=extra_config,
)
diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py
index 2fe649927e..483e02e3b3 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -49,7 +49,6 @@ class PrinterConfig(Object):
indent_spaces: int
print_line_numbers: bool
num_context_lines: int
- syntax_sugar: bool
extra_config: dict
path_to_underline: list[AccessPath] | None
path_to_annotate: dict[AccessPath, str] | None
@@ -70,7 +69,6 @@ class PrinterConfig(Object):
indent_spaces: int = 4,
print_line_numbers: bool = False,
num_context_lines: int | None = None,
- syntax_sugar: bool = True,
show_all_ty: bool = True,
extra_config: dict | None = None,
path_to_underline: list[AccessPath] | None = None,
@@ -91,7 +89,6 @@ class PrinterConfig(Object):
"indent_spaces": indent_spaces,
"print_line_numbers": print_line_numbers,
"num_context_lines": num_context_lines,
- "syntax_sugar": syntax_sugar,
"path_to_underline": path_to_underline,
"path_to_annotate": path_to_annotate,
"obj_to_underline": obj_to_underline,
@@ -135,7 +132,6 @@ class Scriptable:
indent_spaces: int = 4,
print_line_numbers: bool = False,
num_context_lines: int = -1,
- syntax_sugar: bool = True,
show_all_ty: bool = True,
extra_config: dict | None = None,
path_to_underline: list[AccessPath] | None = None,
@@ -167,8 +163,6 @@ class Scriptable:
Whether to print line numbers
num_context_lines : int = -1
The number of lines of context to print before and after the line to underline.
- syntax_sugar: bool = True
- Whether to output with syntax sugar, set false for complete printing.
show_all_ty: bool = True
If True (default), annotate all variable bindings with the struct
info of that variable. If False, only add annotations where
@@ -242,7 +236,6 @@ class Scriptable:
indent_spaces=indent_spaces,
print_line_numbers=print_line_numbers,
num_context_lines=num_context_lines,
- syntax_sugar=syntax_sugar,
show_all_ty=show_all_ty,
extra_config=merged_extra if merged_extra else None,
path_to_underline=path_to_underline,
@@ -265,7 +258,6 @@ class Scriptable:
indent_spaces: int = 4,
print_line_numbers: bool = False,
num_context_lines: int = -1,
- syntax_sugar: bool = True,
extra_config: dict | None = None,
path_to_underline: list[AccessPath] | None = None,
path_to_annotate: dict[AccessPath, str] | None = None,
@@ -286,7 +278,6 @@ class Scriptable:
indent_spaces=indent_spaces,
print_line_numbers=print_line_numbers,
num_context_lines=num_context_lines,
- syntax_sugar=syntax_sugar,
extra_config=merged_extra,
path_to_underline=path_to_underline,
path_to_annotate=path_to_annotate,
@@ -310,7 +301,6 @@ class Scriptable:
indent_spaces: int = 4,
print_line_numbers: bool = False,
num_context_lines: int = -1,
- syntax_sugar: bool = True,
show_all_ty: bool = True,
extra_config: dict | None = None,
path_to_underline: list[AccessPath] | None = None,
@@ -365,8 +355,6 @@ class Scriptable:
Whether to print line numbers
num_context_lines : int = -1
The number of lines of context to print before and after the line to underline.
- syntax_sugar: bool = True
- Whether to output with syntax sugar, set false for complete printing.
show_all_ty: bool = True
If True (default), annotate all variable bindings with the struct
info of that variable. If False, only add annotations where
@@ -401,7 +389,6 @@ class Scriptable:
indent_spaces=indent_spaces,
print_line_numbers=print_line_numbers,
num_context_lines=num_context_lines,
- syntax_sugar=syntax_sugar,
show_all_ty=show_all_ty,
extra_config=extra_config,
path_to_underline=path_to_underline,
diff --git a/src/relax/script/printer/distributed.cc b/src/relax/script/printer/distributed.cc
index f76a3c2626..e0386ef8f7 100644
--- a/src/relax/script/printer/distributed.cc
+++ b/src/relax/script/printer/distributed.cc
@@ -74,7 +74,6 @@ ffi::Optional<ExprDoc> DTensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
ExprDoc mesh = d->Translate(ty->device_mesh).value();
if (auto selector = GlobalInfoSelector(d, ty->device_mesh)) {
mesh = LiteralDoc::Str(selector.value(), std::nullopt);
- d->ExchangeExtraState("script.future_annotations", true);
}
ExprDoc placement = d->Translate(ty->placement).value();
if (require_keywords) {
diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc
index 15a1d8bf1a..4c1e9a087b 100644
--- a/src/script/printer/config.cc
+++ b/src/script/printer/config.cc
@@ -128,9 +128,6 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
v.value().as_or_throw<ffi::Optional<ffi::Map<ffi::ObjectRef, ffi::String>>>().value_or(
ffi::Map<ffi::ObjectRef, ffi::String>());
}
- if (auto v = get("syntax_sugar")) {
- n->syntax_sugar = v.value().cast<bool>();
- }
if (auto v = get("render_invisible_path_info")) {
n->render_invisible_path_info = v.value().cast<bool>();
}
diff --git a/src/script/printer/doc_translator.cc b/src/script/printer/doc_translator.cc
index 96bc70a7ee..bbb2f3938f 100644
--- a/src/script/printer/doc_translator.cc
+++ b/src/script/printer/doc_translator.cc
@@ -16,8 +16,6 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include "doc_translator.h"
-
#include <tvm/ffi/extra/json.h>
#include <tvm/ffi/extra/serialization.h>
#include <tvm/ffi/reflection/registry.h>
@@ -93,6 +91,7 @@ class TranslationEngine : public DocTranslatorObj {
}
IdDoc id = Allocate(d, var->name);
auto& frame = self->var_scopes_.back();
+ self->RecordOrigin(id, var);
frame.mapped.Set(var, id);
if (!explicit_def) frame.implicit_defs.Set(var, id);
return id;
@@ -144,9 +143,8 @@ class TranslationEngine : public DocTranslatorObj {
std::vector<VariableEnvironment> var_scopes_{1};
};
-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>>("ir.binding_names", {});
+Doc TranslateComplete(ffi::AnyView obj, TranslationEngine* engine,
+ const ffi::Map<ffi::String, ffi::Any>& extra_config) {
// Reserve the same registered or explicitly configured aliases used by rendering.
// Keep key order so overlapping aliases reserve the same generated names.
const auto& aliases = GetNamespaceAliases();
@@ -160,38 +158,18 @@ details::TranslationResult TranslateComplete(ffi::AnyView obj, TranslationEngine
if (option.size() >= 7 && option.compare(option.size() - 7, 7, ".prefix") == 0)
engine->AllocId(value.cast<ffi::String>());
}
- for (const ffi::String& name : binding_names) engine->AllocId(name);
- if (binding_names.empty()) {
- if (auto name =
- engine->GetExtraConfig<ffi::Optional<ffi::String>>("ir.module_name", std::nullopt)) {
- engine->AllocId(name.value());
- }
- }
// Allocate generated names after all configured names. A dialect or module
// named tvm/metadata must not shadow the loader or metadata registry.
IdDoc metadata_id = engine->AllocId("metadata");
IdDoc metadata_loader = engine->AllocId("tvm");
engine->ExchangeExtraState("ir.metadata_id", metadata_id);
- bool compact_expression = false;
auto docs = engine->WithDocScope([&]() {
ffi::Optional<ExprDoc> expression = engine->Translate(obj);
if (expression.has_value()) {
- compact_expression = !engine->GetExtraConfig<bool>("ir.verbose_expr", true);
- if (compact_expression) engine->CurrentScopeDocs().clear();
engine->Emit(ExprStmtDoc(expression.value()), obj.as<ffi::ObjectRef>());
}
- if (!binding_names.empty() && !engine->CurrentScopeDocs().empty()) {
- Doc root = engine->CurrentScopeDocs().back();
- ffi::String name = binding_names.back();
- if (auto function = root.as<FunctionDoc>()) {
- function.value()->name = IdDoc(name);
- } else if (auto module = root.as<ClassDoc>()) {
- module.value()->name = IdDoc(name);
- }
- }
});
auto root_declarations = engine->WithDocScope([&]() {
- if (compact_expression) return;
auto pending = engine->GetImplicitDefs();
std::vector<std::pair<Var, IdDoc>> remaining(pending.begin(), pending.end());
for (const auto& [var, _] : remaining) {
@@ -208,7 +186,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>("ir.show_meta", true)) {
+ if (!metadata.empty()) {
// 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;
@@ -221,38 +199,23 @@ details::TranslationResult TranslateComplete(ffi::AnyView obj, TranslationEngine
}
ffi::String json = ffi::json::Stringify(
ffi::ToJSONGraph(serialized, ffi::json::Object{{"tvm_version", TVM_VERSION}}), 2);
- body.push_back(AssignDoc(IdDoc(metadata_id->name),
- IdDoc(metadata_loader->name)
- ->Attr("ir")
- ->Attr("load_json")
- ->Call({LiteralDoc::Str(json, std::nullopt)}),
- std::nullopt));
+ body.push_back(AssignDoc(
+ metadata_id,
+ metadata_loader->Attr("ir")->Attr("load_json")->Call({LiteralDoc::Str(json, std::nullopt)}),
+ std::nullopt));
}
for (const auto& stmt : ToStmtDocArray(docs)) body.push_back(stmt);
- 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>("ir.show_meta", true)};
+ return StmtBlockDoc(body);
}
} // namespace
-namespace details {
-TranslationResult TranslateWithOptions(ffi::AnyView ir,
- ffi::Map<ffi::String, ffi::Any> extra_config) {
- auto engine = ffi::make_object<TranslationEngine>(extra_config);
- return TranslateComplete(ir, engine.get(), std::move(extra_config));
-}
-} // namespace details
-
Doc DocTranslate(ffi::AnyView ir, ffi::Dict<Doc, ffi::ObjectRef>* doc_origins,
ffi::Map<ffi::String, ffi::Any> extra_config) {
- auto result = details::TranslateWithOptions(ir, std::move(extra_config));
- if (doc_origins) *doc_origins = std::move(result.origins);
- return result.doc;
+ auto engine = ffi::make_object<TranslationEngine>(extra_config);
+ Doc doc = TranslateComplete(ir, engine.get(), extra_config);
+ if (doc_origins) *doc_origins = engine->GetDocOrigins();
+ return doc;
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/script/printer/doc_translator.h b/src/script/printer/doc_translator.h
deleted file mode 100644
index 54a85f431a..0000000000
--- a/src/script/printer/doc_translator.h
+++ /dev/null
@@ -1,57 +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_DOC_TRANSLATOR_PRIVATE_H_
-#define TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_PRIVATE_H_
-
-#include <tvm/script/printer/doc_translator.h>
-
-namespace tvm {
-namespace script {
-namespace printer {
-namespace details {
-
-/*! \brief Complete translated body, origins and explicit header facts. */
-struct TranslationResult {
- /*! \brief Complete body, including declarations and metadata. */
- Doc doc;
- /*! \brief Original IR objects associated with generated Docs. */
- ffi::Dict<Doc, ffi::ObjectRef> origins;
- /*! \brief Whether the script needs postponed annotation evaluation. */
- bool future_annotations;
- /*! \brief Allocated import name for the metadata loader. */
- ffi::String metadata_loader;
- /*! \brief Whether the body uses the metadata loader. */
- bool imports_metadata;
-};
-
-/*!
- * \brief Translate an input into its complete Doc representation.
- * \param ir The input to translate.
- * \param extra_config Read-only options for this translation.
- * \return Complete body, origins and header facts.
- */
-TranslationResult TranslateWithOptions(ffi::AnyView ir,
- ffi::Map<ffi::String, ffi::Any> extra_config);
-
-} // namespace details
-} // namespace printer
-} // namespace script
-} // namespace tvm
-
-#endif // TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_PRIVATE_H_
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index 5e561b421b..fa0145f339 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -72,10 +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>>("ir.binding_names", {});
- ffi::String module_name =
- names.empty() ? d->GetExtraConfig<ffi::String>("ir.module_name", "Module") : names.back();
- d->AllocId(module_name);
+ ffi::String module_name = d->GetExtraConfig<ffi::String>("ir.module_name", "Module");
+ IdDoc module_id = d->AllocId(module_name);
+ auto saved_module = d->ExchangeExtraState("ir.module_id", module_id);
GlobalInfoScope global_infos(d, mod);
std::vector<std::pair<GlobalVar, BaseFunc>> functions(mod->functions.begin(),
mod->functions.end());
@@ -154,9 +153,9 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView in
}
});
global_infos.Close();
- d->Emit(
- ClassDoc(IdDoc(module_name), {NamespaceDoc("ir")->Attr("ir_module")}, ToStmtDocArray(body)),
- ffi::GetRef<ffi::ObjectRef>(mod));
+ d->ExchangeExtraState("ir.module_id", std::move(saved_module));
+ d->Emit(ClassDoc(module_id, {NamespaceDoc("ir")->Attr("ir_module")}, ToStmtDocArray(body)),
+ ffi::GetRef<ffi::ObjectRef>(mod));
return std::nullopt;
}
diff --git a/src/script/printer/ir/utils.cc b/src/script/printer/ir/utils.cc
index 6c597fd0a7..082ecc2d2e 100644
--- a/src/script/printer/ir/utils.cc
+++ b/src/script/printer/ir/utils.cc
@@ -62,32 +62,6 @@ IdDoc VarDoc(DocTranslatorObj* d, const Var& var, bool explicit_def) {
return doc;
}
-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>>("ir.annotated_objects", {})) {
- if (item.same_as(object)) return true;
- }
- return false;
-}
-
-bool IsUnderlined(DocTranslatorObj* d, const ffi::ObjectRef& object) {
- for (const auto& item :
- d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("ir.underlined_objects", {})) {
- if (item.same_as(object)) return true;
- }
- return false;
-}
-
-bool HasAnnotatedDescendant(DocTranslatorObj* d, const ffi::ObjectRef& object) {
- for (const auto& item :
- d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("ir.annotation_ancestors", {})) {
- if (item.same_as(object)) return true;
- }
- return false;
-}
-
ExprDoc NamedCallCallee(const ffi::String& canonical_name) {
std::string name = canonical_name;
size_t end = name.find('.');
@@ -101,10 +75,12 @@ ExprDoc NamedCallCallee(const ffi::String& canonical_name) {
}
ExprDoc GlobalReference(DocTranslatorObj* d, const ffi::String& name) {
- 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);
+ auto module = d->GetOrCreateExtraState<ffi::Optional<IdDoc>>("ir.module_id");
+ if (!module.has_value()) {
+ module = IdDoc(d->GetExtraConfig<ffi::String>("ir.module_name", "Module"));
+ d->ExchangeExtraState("ir.module_id", module);
+ }
+ return module.value()->Attr(name);
}
namespace {
@@ -240,7 +216,6 @@ void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var, IdDoc
}
std::sort(parameters.begin(), parameters.end(),
[&](const Var& a, const Var& b) { return signature[a]->name < signature[b]->name; });
- if (!parameters.empty()) d->ExchangeExtraState("script.future_annotations", true);
for (const Var& var : parameters) {
IdDoc id = VarDoc(d, var);
if (var->ty.as_or_throw<PrimType>()->dtype == (DLDataType{kDLInt, 64, 1})) {
diff --git a/src/script/printer/ir/utils.h b/src/script/printer/ir/utils.h
index 272aa1ad44..2dd96e46ec 100644
--- a/src/script/printer/ir/utils.h
+++ b/src/script/printer/ir/utils.h
@@ -35,10 +35,6 @@ namespace details {
ExprDoc AddMetadata(DocTranslatorObj* d, ffi::Any value);
IdDoc VarDoc(DocTranslatorObj* d, const Var& var, bool explicit_def = true);
-bool SyntaxSugar(DocTranslatorObj* d);
-bool IsAnnotated(DocTranslatorObj* d, const ffi::ObjectRef& object);
-bool IsUnderlined(DocTranslatorObj* d, const ffi::ObjectRef& object);
-bool HasAnnotatedDescendant(DocTranslatorObj* d, const ffi::ObjectRef& object);
ExprDoc GlobalReference(DocTranslatorObj* d, const ffi::String& name);
ExprDoc NamedCallCallee(const ffi::String& canonical_name);
diff --git a/src/script/printer/printer.cc b/src/script/printer/printer.cc
index 7469dafd44..dcb016d93f 100644
--- a/src/script/printer/printer.cc
+++ b/src/script/printer/printer.cc
@@ -18,15 +18,12 @@
*/
#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/module.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/relax/expr.h>
#include <tvm/script/printer/doc.h>
#include <tvm/script/printer/doc_translator.h>
#include <tvm/script/printer/printer.h>
-#include <tvm/tirx/function.h>
#include <algorithm>
+#include <functional>
#include <string>
#include <unordered_map>
#include <unordered_set>
@@ -34,33 +31,119 @@
#include <vector>
#include "doc_printer.h"
-#include "doc_translator.h"
namespace tvm {
namespace script {
namespace printer {
namespace {
-void CollectNamespaces(ffi::AnyView value, std::unordered_set<std::string>* names,
- std::unordered_set<const ffi::Object*>* visited) {
+// Visit only the completed Doc graph, never the IR carried by literal values.
+void VisitDocs(ffi::AnyView value, const std::function<void(const Doc&)>& visit,
+ std::unordered_set<const ffi::Object*>* visited) {
const ffi::Object* object = value.as<ffi::Object>();
if (!object || !visited->insert(object).second) return;
- if (auto name = value.as<NamespaceDocNode>()) names->insert(name->canonical_name);
if (auto array = value.as<ffi::Array<ffi::Any>>()) {
- for (const auto& item : array.value()) CollectNamespaces(item, names, visited);
- } else if (auto map = value.as<ffi::Map<ffi::Any, ffi::Any>>()) {
- for (const auto& [key, item] : map.value()) {
- CollectNamespaces(key, names, visited);
- CollectNamespaces(item, names, visited);
- }
- } else {
+ for (const auto& item : array.value()) VisitDocs(item, visit, visited);
+ } else if (auto doc = value.as<Doc>()) {
+ visit(doc.value());
ffi::reflection::ForEachFieldInfo(
TVMFFIGetTypeInfo(object->type_index()), [&](const TVMFFIFieldInfo* field) {
- CollectNamespaces(ffi::reflection::FieldGetter(field)(object), names, visited);
+ if (ffi::String(field->name) != "source_paths") {
+ VisitDocs(ffi::reflection::FieldGetter(field)(object), visit, visited);
+ }
});
}
}
+ffi::Any ReflectedField(ffi::AnyView value, const char* name) {
+ ffi::Any result;
+ if (const auto* object = value.as<ffi::Object>()) {
+ ffi::reflection::ForEachFieldInfo(TVMFFIGetTypeInfo(object->type_index()),
+ [&](const TVMFFIFieldInfo* field) {
+ if (ffi::String(field->name) == name)
+ result = ffi::reflection::FieldGetter(field)(object);
+ });
+ }
+ return result;
+}
+
+// The translator materializes metadata through the ordinary load_json call.
+// Recover its allocated loader symbol from that expression for the import.
+ffi::Optional<IdDoc> MetadataLoader(const StmtDoc& statement) {
+ auto assignment = statement.as<AssignDoc>();
+ if (!assignment || !assignment.value()->rhs) return std::nullopt;
+ auto call = assignment.value()->rhs.value().as<CallDoc>();
+ if (!call) return std::nullopt;
+ auto load = call.value()->callee.as<AttrAccessDoc>();
+ if (!load || load.value()->name != "load_json") return std::nullopt;
+ auto module = load.value()->value.as<AttrAccessDoc>();
+ if (!module || module.value()->name != "ir") return std::nullopt;
+ return module.value()->value.as<IdDoc>();
+}
+
+void RenameRoot(const Doc& root, const ffi::String& name,
+ const ffi::Dict<Doc, ffi::ObjectRef>& origins) {
+ auto block = root.as<StmtBlockDoc>();
+ if (!block || block.value()->stmts.empty()) return;
+ Doc definition = block.value()->stmts.back();
+ ffi::Optional<IdDoc> root_id;
+ if (auto function = definition.as<FunctionDoc>()) {
+ root_id = function.value()->name;
+ // A public function's definition must preserve its declared symbol. This
+ // invariant applies through the common reflected attributes, without
+ // dispatching on the function's dialect.
+ if (auto original = origins.Get(definition)) {
+ auto attrs = ReflectedField(ReflectedField(original.value(), "attrs"), "dict");
+ if (auto dict = attrs.as<ffi::Map<ffi::String, ffi::Any>>()) {
+ if (auto symbol = dict->Get("global_symbol")) {
+ TVM_FFI_CHECK(symbol.value().cast<ffi::String>() == name, TypeError)
+ << "printer function global_symbol must match its definition name";
+ }
+ }
+ }
+ } else if (auto module = definition.as<ClassDoc>()) {
+ root_id = module.value()->name;
+ }
+ if (!root_id) return;
+ std::vector<IdDoc> identifiers;
+ std::unordered_set<std::string> used;
+ std::unordered_set<const ffi::Object*> definitions;
+ std::unordered_set<const ffi::Object*> visited;
+ VisitDocs(
+ root,
+ [&](const Doc& doc) {
+ if (auto id = doc.as<IdDoc>()) {
+ identifiers.push_back(id.value());
+ used.insert(id.value()->name);
+ } else if (auto function = doc.as<FunctionDoc>()) {
+ definitions.insert(function.value()->name.get());
+ } else if (auto module = doc.as<ClassDoc>()) {
+ definitions.insert(module.value()->name.get());
+ }
+ },
+ &visited);
+ std::string replacement;
+ for (size_t suffix = 1;; ++suffix) {
+ replacement = std::string(name) + "_" + std::to_string(suffix);
+ if (!used.count(replacement)) break;
+ }
+ std::unordered_set<const ffi::Object*> generated;
+ if (!block.value()->stmts.empty()) {
+ if (auto loader = MetadataLoader(block.value()->stmts.front()))
+ generated.insert(loader.value().get());
+ auto assignment = block.value()->stmts.front().as<AssignDoc>();
+ if (assignment && assignment.value()->lhs.as<IdDoc>()) {
+ generated.insert(assignment.value()->lhs.get());
+ }
+ }
+ for (const IdDoc& id : identifiers) {
+ if (id->name == name && !definitions.count(id.get()) &&
+ (origins.count(id) || generated.count(id.get())))
+ id->name = replacement;
+ }
+ root_id.value()->name = name;
+}
+
bool SamePath(const AccessPath& a, const AccessPath& b) {
return a->depth == b->depth && a->IsPrefixOf(b);
}
@@ -131,66 +214,6 @@ class MapDocPaths {
Index(root, AccessPath::Root(), &ancestors);
}
- ffi::Array<ffi::ObjectRef> AnnotatedObjects() const {
- ffi::Array<ffi::ObjectRef> objects;
- ffi::Array<AccessPath> paths;
- for (const auto& [object, message] : config_->obj_to_annotate) objects.push_back(object);
- for (const auto& [path, message] : config_->path_to_annotate) paths.push_back(path);
- return RequestedObjects(objects, paths);
- }
-
- ffi::Array<ffi::ObjectRef> UnderlinedObjects() const {
- return RequestedObjects(config_->obj_to_underline, config_->path_to_underline);
- }
-
- ffi::Array<ffi::ObjectRef> AnnotationAncestors(
- const ffi::Array<ffi::ObjectRef>& annotated_objects) const {
- ffi::Array<ffi::ObjectRef> objects = annotated_objects;
- std::unordered_set<const ffi::Object*> annotated;
- for (const auto& object : annotated_objects) annotated.insert(object.get());
- std::vector<AccessPath> targets;
- for (const auto& [object, path] : occurrences_) {
- if (annotated.count(object.get())) targets.push_back(path);
- }
- // Shared identities can occur in several subtrees. Include each original
- // ancestor conservatively, without handing diagnostic traversal to hooks.
- std::unordered_set<const ffi::Object*> seen = annotated;
- for (const auto& [object, path] : occurrences_) {
- if (seen.count(object.get())) continue;
- for (const auto& target : targets) {
- if (path->IsPrefixOf(target)) {
- seen.insert(object.get());
- objects.push_back(object);
- break;
- }
- }
- }
- return objects;
- }
-
- ffi::Array<ffi::ObjectRef> RequestedObjects(const ffi::Array<ffi::ObjectRef>& requested_objects,
- const ffi::Array<AccessPath>& requested_paths) const {
- ffi::Array<ffi::ObjectRef> objects;
- std::unordered_set<const ffi::Object*> seen;
- auto append = [&](const ffi::ObjectRef& object) {
- if (seen.insert(object.get()).second) objects.push_back(object);
- };
- for (const auto& object : requested_objects) append(object);
- // Hooks only need original identities to avoid hiding diagnostics behind
- // compact syntax. Resolve scalar fields to their deepest object owner here;
- // paths, messages, and final Doc-path correspondence remain entry concerns.
- for (const auto& requested : requested_paths) {
- int deepest = -1;
- for (const auto& [object, path] : occurrences_) {
- if (path->IsPrefixOf(requested)) deepest = std::max(deepest, path->depth);
- }
- for (const auto& [object, path] : occurrences_) {
- if (path->depth == deepest && path->IsPrefixOf(requested)) append(object);
- }
- }
- return objects;
- }
-
std::pair<ffi::Array<AccessPath>, ffi::Map<AccessPath, ffi::String>> Map(
const Doc& root, const ffi::Dict<Doc, ffi::ObjectRef>& doc_origins) {
auto underline_paths = config_->path_to_underline;
@@ -249,56 +272,84 @@ class MapDocPaths {
} // namespace
ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config) {
- MapDocPaths paths(obj, config);
- auto extra_config = config->extra_config;
- 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("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::Dict<Doc, ffi::ObjectRef> origins;
+ Doc doc = DocTranslate(obj, &origins, config->extra_config);
+ if (!config->binding_names.empty()) RenameRoot(doc, config->binding_names.back(), origins);
+
+ auto block = doc.as<StmtBlockDoc>();
+ if (block && !config->verbose_expr && !block.value()->stmts.empty()) {
+ if (auto expression = block.value()->stmts.back().as<ExprStmtDoc>()) {
+ auto original = origins.Get(expression.value()->expr);
+ if (original && original.value().same_as(obj)) {
+ // A returned root expression carries the input origin. Statement
+ // emissions retain their own input and child-expression origins.
+ ffi::Array<StmtDoc> compact;
+ if (MetadataLoader(block.value()->stmts.front()))
+ compact.push_back(block.value()->stmts.front());
+ compact.push_back(block.value()->stmts.back());
+ block.value()->stmts = compact;
+ }
+ }
+ }
ffi::Array<ffi::Any> header;
- if (result.future_annotations) {
- header.push_back(ffi::String("from __future__ import annotations\n\n"));
+ if (block && !block.value()->stmts.empty()) {
+ if (auto loader = MetadataLoader(block.value()->stmts.front())) {
+ if (config->show_meta) {
+ header.push_back(
+ ffi::String(loader.value()->name == "tvm"
+ ? "import tvm\n"
+ : "import tvm as " + std::string(loader.value()->name) + "\n"));
+ } else {
+ block.value()->stmts.erase(block.value()->stmts.begin());
+ block.value()->stmts.push_back(
+ CommentDoc("Metadata omitted. Use show_meta=True in script() method to show it."));
+ }
+ }
}
- if (result.imports_metadata) {
- header.push_back(
- ffi::String(result.metadata_loader == "tvm"
- ? "import tvm\n"
- : "import tvm as " + std::string(result.metadata_loader) + "\n"));
+ bool definition = false;
+ bool future_annotations = false;
+ std::unordered_set<std::string> namespaces;
+ std::unordered_set<const ffi::Object*> visited;
+ VisitDocs(
+ doc,
+ [&](const Doc& child) {
+ if (auto name = child.as<NamespaceDoc>()) namespaces.insert(name.value()->canonical_name);
+ if (auto function = child.as<FunctionDoc>()) {
+ definition = true;
+ // Defer annotations uniformly, including forward references and type
+ // parameters. Header policy depends on syntax, not dialect semantics.
+ future_annotations |=
+ std::any_of(function.value()->args.begin(), function.value()->args.end(),
+ [](const AssignDoc& arg) { return arg->annotation.has_value(); }) ||
+ function.value()->return_type.has_value() || !function.value()->type_params.empty();
+ }
+ definition |= child.as<ClassDoc>().has_value();
+ },
+ &visited);
+ if (future_annotations) {
+ header.insert(header.begin(), ffi::String("from __future__ import annotations\n\n"));
}
- if (obj.as<IRModuleNode>() || obj.as<tirx::PrimFuncNode>() || obj.as<relax::FunctionNode>()) {
- std::unordered_set<std::string> namespaces;
- std::unordered_set<const ffi::Object*> visited;
- CollectNamespaces(result.doc, &namespaces, &visited);
- for (const char* canonical : {"ir", "tirx", "s_tir", "relax"}) {
+ if (definition) {
+ std::vector<std::pair<ffi::String, ffi::String>> aliases(GetNamespaceAliases().begin(),
+ GetNamespaceAliases().end());
+ std::sort(aliases.begin(), aliases.end());
+ for (const auto& [key, fallback] : aliases) {
+ std::string canonical = key;
+ canonical.resize(canonical.size() - std::string(".prefix").size());
if (!namespaces.count(canonical)) continue;
- ffi::String name = canonical;
- ffi::String alias = name == "ir" ? config->ir_prefix
- : name == "tirx" ? config->GetExtraConfig<ffi::String>("tirx.prefix", "T")
- : name == "s_tir"
- ? config->GetExtraConfig<ffi::String>("s_tir.prefix", "Ts")
- : config->GetExtraConfig<ffi::String>("relax.prefix", "R");
- // 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>("ir.comment_imports", false) &&
- !(config->show_meta && name == "tirx")) {
+ ffi::String alias = config->GetExtraConfig<ffi::String>(key, fallback);
+ ffi::String import = "from tvm.script import " + canonical + " as " + std::string(alias);
+ if (config->GetExtraConfig<bool>("ir.comment_imports", false)) {
header.push_back(CommentDoc(import));
} else {
header.push_back(ffi::String(std::string(import) + "\n"));
}
- if (name == "tirx" && config->GetExtraConfig<bool>("ir.comment_imports", false)) {
- header.push_back(CommentDoc("from tvm.tirx.layout import Axis"));
- }
}
if (!header.empty()) header.push_back(ffi::String("\n"));
}
- return details::RenderPythonScript(result.doc, config, header, underline_paths, annotations);
+ MapDocPaths paths(obj, config);
+ auto [underline_paths, annotations] = paths.Map(doc, origins);
+ return details::RenderPythonScript(doc, config, header, underline_paths, annotations);
}
TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::GlobalDef().def("script.printer.Script", Script); }
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index d27dbe46c8..45c4437ccf 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -75,7 +75,7 @@ ffi::Optional<ExprDoc> VarDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
}
return ffi::WalkResult::Skip();
});
- d->Emit(AssignDoc(IdDoc(id->name), rhs, annotation), var);
+ d->Emit(AssignDoc(VarDoc(d, var), rhs, annotation), var);
return std::nullopt;
}
return IdDoc(id->name);
diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc
index 4a043bd1da..61c1e7a601 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -72,16 +72,6 @@ ffi::Optional<ExprDoc> PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView in
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const tirx::PrimFuncNode>(input);
TVM_FFI_CHECK(destination == nullptr, TypeError)
<< "printer statement-only node cannot fulfill a destination";
- 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>>("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)
- << "printer PrimFunc global_symbol must match its definition name";
- }
- }
VarScope vars(d);
PrimFuncScope function_scope(d, ffi::GetRef<tirx::PrimFunc>(func));
@@ -157,54 +147,7 @@ ffi::Optional<ExprDoc> PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView in
}
if (!decorator_keys.empty()) decorator = decorator->Call({}, decorator_keys, decorator_values);
ffi::Array<StmtDoc> body;
- if (auto root = func->body.as<s_tir::SBlockRealizeNode>();
- root && root->iter_values.empty() && tvm::prim::is_one(root->predicate) &&
- root->block->annotations.empty() && root->block->match_buffers.empty() &&
- root->block->reads.empty() && root->block->writes.empty() &&
- !root->block->init.has_value() && root->block->iter_vars.empty() &&
- (!root->block->alloc_buffers.empty() ||
- (!root->block->body.as<s_tir::SBlockRealizeNode>() &&
- tirx::ContainsNode<s_tir::SBlockRealizeNode>(root->block->body)))) {
- auto docs = d->WithDocScope([&]() {
- for (const tirx::BufferVar& buffer : root->block->alloc_buffers) {
- CallDoc rhs = d->Translate(buffer.var()->ty).value().as_or_throw<CallDoc>();
- TVM_FFI_CHECK(rhs->callee.as_or_throw<AttrAccessDoc>()->name == "Buffer", TypeError)
- << "Ts.sblock_alloc_buffer cannot reconstruct this nonrepresentable BufferType";
- const auto* buffer_type = buffer.var()->ty.as<tirx::BufferTypeNode>();
- TVM_FFI_CHECK(
- buffer_type->allocated_addr.empty() || (buffer_type->storage_scope != "global" &&
- buffer_type->storage_scope != "shared" &&
- buffer_type->storage_scope != "shared.dyn" &&
- buffer_type->storage_scope != "local"),
- TypeError)
- << "Ts.sblock_alloc_buffer does not accept allocated_addr in "
- << buffer_type->storage_scope;
- rhs->callee = NamespaceDoc("s_tir")->Attr("sblock_alloc_buffer");
- IdDoc lhs = VarDoc(d, buffer);
- d->Emit(AssignDoc(lhs, rhs, std::nullopt), buffer);
- }
- d->Translate(root->block->body);
- });
- ffi::Array<StmtDoc> root_body = ToStmtDocArray(docs);
- ScopeDoc scope(std::nullopt,
- NamespaceDoc("s_tir")->Attr("sblock")->Call(
- {LiteralDoc::Str(root->block->name_hint, std::nullopt)}),
- root_body);
- d->RecordOrigin(scope, root->block);
- ffi::Array<StmtDoc> elided = {CommentDoc("with Ts.sblock(\"root\"):")};
- elided.insert(elided.end(), root_body.begin(), root_body.end());
- if (SyntaxSugar(d) && !IsAnnotated(d, root->block) &&
- !IsAnnotated(d, ffi::GetRef<s_tir::SBlockRealize>(root)) &&
- !IsUnderlined(d, root->block) &&
- !IsUnderlined(d, ffi::GetRef<s_tir::SBlockRealize>(root))) {
- body = elided;
- } else {
- d->RecordOrigin(scope->rhs, ffi::GetRef<s_tir::SBlockRealize>(root));
- body = {scope};
- }
- } else {
- body = Body(func->body, d);
- }
+ body = Body(func->body, d);
body.insert(body.begin(), thread_declarations.begin(), thread_declarations.end());
std::vector<std::pair<ffi::String, ffi::Any>> attrs;
for (const auto& [key, value] : func->attrs->dict) {
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 323c5900a0..4b0c70f3fe 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -317,7 +317,6 @@ ffi::Optional<ExprDoc> SeqStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView inp
continue;
auto scalar = docs.back().as<AssignDoc>();
if (!IsScalarBuffer(d, alloc->var) || !scalar.has_value() ||
- HasAnnotatedDescendant(d, ffi::GetRef<tirx::BufferStore>(store)) ||
!std::all_of(store->indices.begin(), store->indices.end(), tvm::prim::is_zero))
continue;
bool reads_allocation = false;
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 89c5a9dcda..a6a2d00bca 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
@@ -434,7 +434,7 @@ def func():
return 5
"""
_assert_print(func, expected_output)
- assert func.script(verbose_expr=True, syntax_sugar=False).strip() == expected_output.strip()
+ assert func.script(verbose_expr=True).strip() == expected_output.strip()
CUSTOM_FLOAT_DTYPES = [
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py b/tests/python/tirx/script/test_tirx_script_printer.py
index fe38a5c118..09d1e03ed7 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -54,7 +54,7 @@ def test_config_extension_passthrough():
{
"extension.option": 7,
"custom_key": "value",
- "syntax_sugar": False,
+ "verbose_expr": False,
"render_invisible_path_info": True,
"tirx.prefix": "invalid-prefix",
"extra_config": {
@@ -66,11 +66,11 @@ def test_config_extension_passthrough():
)
assert cfg.extra_config["extension.option"] == 9
assert cfg.extra_config["custom_key"] == "value"
- assert "syntax_sugar" not in cfg.extra_config
+ assert "verbose_expr" not in cfg.extra_config
assert "extra_config" not in cfg.extra_config
- assert not cfg.syntax_sugar
+ assert not cfg.verbose_expr
assert not cfg.render_invisible_path_info
- assert make_config({}).syntax_sugar
+ assert not make_config({}).verbose_expr
assert _script(tirx.Var("Custom", "int32"), cfg) == "Custom_1"
diff --git a/tests/python/tirx/script/test_tirx_script_printer_structural_equal.py b/tests/python/tirx/script/test_tirx_script_printer_structural_equal.py
index 7f9e95db83..9a178a5a6a 100644
--- a/tests/python/tirx/script/test_tirx_script_printer_structural_equal.py
+++ b/tests/python/tirx/script/test_tirx_script_printer_structural_equal.py
@@ -80,9 +80,9 @@ def test_prim_func_buffer_param():
def _expected_result(func1, func2, objpath1, objpath2):
return f"""StructuralEqual check failed, caused by lhs at {objpath1}:
-{func1.script(path_to_underline=[objpath1], syntax_sugar=False)}
+{func1.script(path_to_underline=[objpath1])}
and rhs at {objpath2}:
-{func2.script(path_to_underline=[objpath2], syntax_sugar=False)}"""
+{func2.script(path_to_underline=[objpath2])}"""
def _error_message(exception):
(tvm) 03/08: [REFACTOR][Script] Preserve existing CUDA load emission
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 8d9ee4f90ddf26557487d4ebc40f163539a939af
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:06:15 2026 +0000
[REFACTOR][Script] Preserve existing CUDA load emission
---
src/backend/cuda/codegen/codegen_cuda.cc | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index 471cba6edc..8b365cdfa5 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1295,14 +1295,14 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
std::string global_buffer =
this->PrintExpr(addr_buffer->source.as_or_throw<tvm::tirx::BufferVar>().data());
std::string local_addr = this->PrintExpr(op->args[3]);
- std::string destination = op->args[0].as<TensorLoadNode>() ? reg : reg + "[" + local_addr + "]";
this->stream << "asm volatile (\n";
this->stream << "\"{.reg .pred p;\\n\"\n";
this->stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n";
this->stream << "\" @!p mov.b32 %0, 0;\\n\"\n";
this->stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n";
// stream << "\" @p ld.global.nc.L2::128B.f32 %0, [%1];}\\n\"\n" ;
- stream << ": \"=f\"(" << destination << ")\n";
+ stream << ": \"=f\"(" << reg << "[" << local_addr << "]"
+ << ")\n";
stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr << ")), \"r\"((int)"
<< guard << ")\n";
stream << ");\n";
(tvm) 06/08: Read global information directly from the current module
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 58867a45e539e815474d5b8a0967777277c05065
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:48:50 2026 +0000
Read global information directly from the current module
---
src/relax/script/printer/global_info.cc | 9 ++++----
src/script/printer/ir/ir.cc | 37 ++++-----------------------------
src/script/printer/ir/utils.cc | 7 ++-----
3 files changed, 10 insertions(+), 43 deletions(-)
diff --git a/src/relax/script/printer/global_info.cc b/src/relax/script/printer/global_info.cc
index 9c06dcac71..d6ad01bf25 100644
--- a/src/relax/script/printer/global_info.cc
+++ b/src/relax/script/printer/global_info.cc
@@ -16,6 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ir/module.h>
#include <tvm/target/target.h>
#include <optional>
@@ -29,14 +30,12 @@ namespace script {
namespace printer {
namespace details {
-using GlobalInfoMap = ffi::Dict<ffi::String, ffi::List<GlobalInfo>>;
-
// Selectors are derived from the module's forward registry, never cached as
// object-to-string state. Definitions and standalone objects keep constructors.
ffi::Optional<ffi::String> GlobalInfoSelector(DocTranslatorObj* d, const GlobalInfo& info) {
- auto scope = d->GetOrCreateExtraState<ffi::Optional<GlobalInfoMap>>("ir.global_info_map");
- if (!scope.has_value()) return std::nullopt;
- auto infos = scope.value();
+ auto module = d->GetOrCreateExtraState<ffi::Optional<IRModule>>("ir.module");
+ if (!module.has_value()) return std::nullopt;
+ const auto& infos = module.value()->global_infos;
if (auto query = info.as<relax::VDevice>()) {
if (auto devices = infos.Get("vdevice")) {
ffi::String kind = query.value()->target->kind->name;
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index fa0145f339..15bc776cf6 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -35,37 +35,8 @@ namespace script {
namespace printer {
namespace details {
-using GlobalInfoMap = ffi::Dict<ffi::String, ffi::List<GlobalInfo>>;
-
namespace {
-class GlobalInfoScope {
- public:
- GlobalInfoScope(DocTranslatorObj* d, const IRModuleNode* mod) : d_(d) {
- GlobalInfoMap infos;
- for (const auto& [key, entries] : mod->global_infos) {
- infos.Set(key, ffi::List<GlobalInfo>(entries.begin(), entries.end()));
- }
- saved_ = d_->ExchangeExtraState("ir.global_info_map", infos);
- }
- ~GlobalInfoScope() noexcept {
- if (d_) {
- try {
- Close();
- } catch (...) { /* Preserve translation failure. */
- }
- }
- }
- void Close() {
- DocTranslatorObj* d = std::exchange(d_, nullptr);
- d->ExchangeExtraState("ir.global_info_map", std::move(saved_));
- }
-
- private:
- DocTranslatorObj* d_;
- std::optional<ffi::Any> saved_;
-};
-
ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
const ffi::Object* destination) {
const auto* mod =
@@ -74,8 +45,8 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView in
<< "printer statement-only node cannot fulfill a destination";
ffi::String module_name = d->GetExtraConfig<ffi::String>("ir.module_name", "Module");
IdDoc module_id = d->AllocId(module_name);
- auto saved_module = d->ExchangeExtraState("ir.module_id", module_id);
- GlobalInfoScope global_infos(d, mod);
+ d->ExchangeExtraState("ir.module_id", module_id);
+ d->ExchangeExtraState("ir.module", ffi::GetRef<IRModule>(mod));
std::vector<std::pair<GlobalVar, BaseFunc>> functions(mod->functions.begin(),
mod->functions.end());
auto rank = [](const BaseFunc& func) {
@@ -152,8 +123,8 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView in
}
}
});
- global_infos.Close();
- d->ExchangeExtraState("ir.module_id", std::move(saved_module));
+ d->ExchangeExtraState("ir.module", std::nullopt);
+ d->ExchangeExtraState("ir.module_id", std::nullopt);
d->Emit(ClassDoc(module_id, {NamespaceDoc("ir")->Attr("ir_module")}, ToStmtDocArray(body)),
ffi::GetRef<ffi::ObjectRef>(mod));
return std::nullopt;
diff --git a/src/script/printer/ir/utils.cc b/src/script/printer/ir/utils.cc
index 082ecc2d2e..fec3bbf20f 100644
--- a/src/script/printer/ir/utils.cc
+++ b/src/script/printer/ir/utils.cc
@@ -19,7 +19,7 @@
#include "utils.h"
#include <tvm/ffi/reflection/accessor.h>
-#include <tvm/ir/global_info.h>
+#include <tvm/ir/module.h>
#include <tvm/ir/op.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
@@ -204,10 +204,7 @@ ffi::Dict<Var, IdDoc> CopyImplicitDefs(DocTranslatorObj* d) {
void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var, IdDoc>& signature,
const FunctionDoc& function) {
- // A module's existing scoped registry also identifies shared external captures.
- auto module =
- d->GetOrCreateExtraState<ffi::Optional<ffi::Dict<ffi::String, ffi::List<GlobalInfo>>>>(
- "ir.global_info_map");
+ auto module = d->GetOrCreateExtraState<ffi::Optional<IRModule>>("ir.module");
auto pending = d->GetImplicitDefs();
std::vector<Var> parameters;
if (!module.has_value()) {
(tvm) 05/08: Keep metadata serialization and identifier display in rendering
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 dd740ff684874b181f42163e667fcef4c6245ec8
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:48:49 2026 +0000
Keep metadata serialization and identifier display in rendering
---
src/s_tir/schedule/error.cc | 1 -
src/script/printer/config.cc | 2 -
src/script/printer/doc_translator.cc | 26 +----
src/script/printer/printer.cc | 189 ++++++++++++++++++++++-------------
4 files changed, 120 insertions(+), 98 deletions(-)
diff --git a/src/s_tir/schedule/error.cc b/src/s_tir/schedule/error.cc
index b6aed81da9..b233e91061 100644
--- a/src/s_tir/schedule/error.cc
+++ b/src/s_tir/schedule/error.cc
@@ -51,7 +51,6 @@ ffi::String ScheduleErrorContextObj::RenderReport(const ffi::String& primitive)
int n_locs = locs.size();
std::string msg = DetailRenderTemplate();
PrinterConfig cfg;
- cfg->syntax_sugar = false;
if (n_locs > 0) {
for (int i = 0; i < n_locs; ++i) {
std::string name = locs[i]->GetTypeKey() + '#' + std::to_string(i);
diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc
index 4c1e9a087b..3e384a2628 100644
--- a/src/script/printer/config.cc
+++ b/src/script/printer/config.cc
@@ -147,8 +147,6 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
}
}
- n->extra_config.Set("ir.prefix", n->ir_prefix);
-
// 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] :
diff --git a/src/script/printer/doc_translator.cc b/src/script/printer/doc_translator.cc
index bbb2f3938f..d1b043a4da 100644
--- a/src/script/printer/doc_translator.cc
+++ b/src/script/printer/doc_translator.cc
@@ -16,8 +16,6 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/extra/json.h>
-#include <tvm/ffi/extra/serialization.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/base.h>
#include <tvm/script/printer/doc.h>
@@ -158,10 +156,8 @@ Doc TranslateComplete(ffi::AnyView obj, TranslationEngine* engine,
if (option.size() >= 7 && option.compare(option.size() - 7, 7, ".prefix") == 0)
engine->AllocId(value.cast<ffi::String>());
}
- // Allocate generated names after all configured names. A dialect or module
- // named tvm/metadata must not shadow the loader or metadata registry.
+ // Metadata references share one allocated identifier across the completed Doc.
IdDoc metadata_id = engine->AllocId("metadata");
- IdDoc metadata_loader = engine->AllocId("tvm");
engine->ExchangeExtraState("ir.metadata_id", metadata_id);
auto docs = engine->WithDocScope([&]() {
ffi::Optional<ExprDoc> expression = engine->Translate(obj);
@@ -183,27 +179,7 @@ Doc TranslateComplete(ffi::AnyView obj, TranslationEngine* engine,
for (const Doc& doc : docs) root_declarations.push_back(doc);
docs = root_declarations;
}
- 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()) {
- // 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;
- std::vector<ffi::String> keys;
- for (const auto& [key, entries] : metadata) keys.push_back(key);
- std::sort(keys.begin(), keys.end());
- for (const auto& key : keys) {
- auto entries = metadata[key];
- serialized.Set(key, ffi::Array<ffi::Any>(entries.begin(), entries.end()));
- }
- ffi::String json = ffi::json::Stringify(
- ffi::ToJSONGraph(serialized, ffi::json::Object{{"tvm_version", TVM_VERSION}}), 2);
- body.push_back(AssignDoc(
- metadata_id,
- metadata_loader->Attr("ir")->Attr("load_json")->Call({LiteralDoc::Str(json, std::nullopt)}),
- std::nullopt));
- }
for (const auto& stmt : ToStmtDocArray(docs)) body.push_back(stmt);
return StmtBlockDoc(body);
}
diff --git a/src/script/printer/printer.cc b/src/script/printer/printer.cc
index dcb016d93f..43821facd6 100644
--- a/src/script/printer/printer.cc
+++ b/src/script/printer/printer.cc
@@ -16,6 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/serialization.h>
#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/script/printer/doc.h>
@@ -67,47 +69,88 @@ ffi::Any ReflectedField(ffi::AnyView value, const char* name) {
return result;
}
-// The translator materializes metadata through the ordinary load_json call.
-// Recover its allocated loader symbol from that expression for the import.
-ffi::Optional<IdDoc> MetadataLoader(const StmtDoc& statement) {
- auto assignment = statement.as<AssignDoc>();
- if (!assignment || !assignment.value()->rhs) return std::nullopt;
- auto call = assignment.value()->rhs.value().as<CallDoc>();
- if (!call) return std::nullopt;
- auto load = call.value()->callee.as<AttrAccessDoc>();
- if (!load || load.value()->name != "load_json") return std::nullopt;
- auto module = load.value()->value.as<AttrAccessDoc>();
- if (!module || module.value()->name != "ir") return std::nullopt;
- return module.value()->value.as<IdDoc>();
+// Metadata references retain their original values in the ordinary origin map.
+// Recognize the existing registry[type_key][index] expression rather than a
+// reserved spelling, since the generated identifier may have been renamed.
+ffi::Map<ffi::String, ffi::Array<ffi::Any>> CollectMetadata(
+ const Doc& root, const ffi::Dict<Doc, ffi::ObjectRef>& origins,
+ ffi::Optional<IdDoc>* metadata_id) {
+ ffi::Map<ffi::String, ffi::Array<ffi::Any>> metadata;
+ std::unordered_set<const ffi::Object*> visited;
+ VisitDocs(
+ root,
+ [&](const Doc& doc) {
+ auto reference = doc.as<IndexDoc>();
+ auto original = origins.Get(doc);
+ if (!reference || !original || reference.value()->indices.size() != 1) return;
+ auto entries = reference.value()->value.as<IndexDoc>();
+ if (!entries || entries.value()->indices.size() != 1) return;
+ auto id = entries.value()->value.as<IdDoc>();
+ auto key_doc = entries.value()->indices[0].as<LiteralDoc>();
+ auto index_doc = reference.value()->indices[0].as<LiteralDoc>();
+ if (!id || origins.count(id.value()) || !key_doc || !index_doc) return;
+ auto key = key_doc.value()->value.as<ffi::String>();
+ const auto* index = index_doc.value()->value.as<IntImmNode>();
+ if (!key || key.value() != original.value()->GetTypeKey() || !index || index->value < 0)
+ return;
+ if (metadata_id->has_value() && !metadata_id->value().same_as(id.value())) return;
+ *metadata_id = id;
+ auto values = metadata.Get(key.value()).value_or(ffi::Array<ffi::Any>{});
+ size_t slot = static_cast<size_t>(index->value);
+ while (values.size() <= slot) values.push_back(nullptr);
+ values.Set(slot, original.value());
+ metadata.Set(key.value(), values);
+ },
+ &visited);
+ return metadata;
}
-void RenameRoot(const Doc& root, const ffi::String& name,
- const ffi::Dict<Doc, ffi::ObjectRef>& origins) {
+ffi::Array<ffi::String> DisplayAliases(const PrinterConfig& config) {
+ ffi::Array<ffi::String> aliases;
+ for (const auto& [key, fallback] : GetNamespaceAliases()) {
+ aliases.push_back(key == "ir.prefix" ? config->ir_prefix
+ : config->GetExtraConfig<ffi::String>(key, fallback));
+ }
+ return aliases;
+}
+
+void ApplyDisplayNames(const Doc& root, const PrinterConfig& config,
+ const ffi::Dict<Doc, ffi::ObjectRef>& origins,
+ const ffi::Optional<IdDoc>& metadata_id) {
auto block = root.as<StmtBlockDoc>();
if (!block || block.value()->stmts.empty()) return;
Doc definition = block.value()->stmts.back();
ffi::Optional<IdDoc> root_id;
+ ffi::Optional<ffi::String> name;
+ if (!config->binding_names.empty()) name = config->binding_names.back();
if (auto function = definition.as<FunctionDoc>()) {
root_id = function.value()->name;
- // A public function's definition must preserve its declared symbol. This
- // invariant applies through the common reflected attributes, without
- // dispatching on the function's dialect.
- if (auto original = origins.Get(definition)) {
- auto attrs = ReflectedField(ReflectedField(original.value(), "attrs"), "dict");
- if (auto dict = attrs.as<ffi::Map<ffi::String, ffi::Any>>()) {
- if (auto symbol = dict->Get("global_symbol")) {
- TVM_FFI_CHECK(symbol.value().cast<ffi::String>() == name, TypeError)
- << "printer function global_symbol must match its definition name";
+ if (name) {
+ // Preserve a public function's declared symbol through common reflected
+ // attributes, without dispatching on the function's dialect.
+ if (auto original = origins.Get(definition)) {
+ auto attrs = ReflectedField(ReflectedField(original.value(), "attrs"), "dict");
+ if (auto dict = attrs.as<ffi::Map<ffi::String, ffi::Any>>()) {
+ if (auto symbol = dict->Get("global_symbol")) {
+ TVM_FFI_CHECK(symbol.value().cast<ffi::String>() == name.value(), TypeError)
+ << "printer function global_symbol must match its definition name";
+ }
}
}
}
} else if (auto module = definition.as<ClassDoc>()) {
root_id = module.value()->name;
+ if (!name)
+ name = config->GetExtraConfig<ffi::Optional<ffi::String>>("ir.module_name", std::nullopt);
}
- if (!root_id) return;
+ auto reserved = DisplayAliases(config);
+ if (auto module_name =
+ config->GetExtraConfig<ffi::Optional<ffi::String>>("ir.module_name", std::nullopt)) {
+ reserved.push_back(module_name.value());
+ }
+ if (root_id && name) reserved.push_back(name.value());
std::vector<IdDoc> identifiers;
- std::unordered_set<std::string> used;
- std::unordered_set<const ffi::Object*> definitions;
+ std::unordered_set<std::string> used(reserved.begin(), reserved.end());
std::unordered_set<const ffi::Object*> visited;
VisitDocs(
root,
@@ -115,33 +158,26 @@ void RenameRoot(const Doc& root, const ffi::String& name,
if (auto id = doc.as<IdDoc>()) {
identifiers.push_back(id.value());
used.insert(id.value()->name);
- } else if (auto function = doc.as<FunctionDoc>()) {
- definitions.insert(function.value()->name.get());
- } else if (auto module = doc.as<ClassDoc>()) {
- definitions.insert(module.value()->name.get());
}
},
&visited);
- std::string replacement;
- for (size_t suffix = 1;; ++suffix) {
- replacement = std::string(name) + "_" + std::to_string(suffix);
- if (!used.count(replacement)) break;
- }
- std::unordered_set<const ffi::Object*> generated;
- if (!block.value()->stmts.empty()) {
- if (auto loader = MetadataLoader(block.value()->stmts.front()))
- generated.insert(loader.value().get());
- auto assignment = block.value()->stmts.front().as<AssignDoc>();
- if (assignment && assignment.value()->lhs.as<IdDoc>()) {
- generated.insert(assignment.value()->lhs.get());
+ for (const ffi::String& spelling : reserved) {
+ std::string replacement;
+ for (size_t suffix = 1;; ++suffix) {
+ replacement = std::string(spelling) + "_" + std::to_string(suffix);
+ if (!used.count(replacement)) break;
}
+ bool renamed = false;
+ for (const IdDoc& id : identifiers) {
+ if (id->name == spelling && !(root_id && root_id.value().same_as(id)) &&
+ (origins.count(id) || (metadata_id && metadata_id.value().same_as(id)))) {
+ id->name = replacement;
+ renamed = true;
+ }
+ }
+ if (renamed) used.insert(replacement);
}
- for (const IdDoc& id : identifiers) {
- if (id->name == name && !definitions.count(id.get()) &&
- (origins.count(id) || generated.count(id.get())))
- id->name = replacement;
- }
- root_id.value()->name = name;
+ if (root_id && name) root_id.value()->name = name.value();
}
bool SamePath(const AccessPath& a, const AccessPath& b) {
@@ -274,36 +310,48 @@ class MapDocPaths {
ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config) {
ffi::Dict<Doc, ffi::ObjectRef> origins;
Doc doc = DocTranslate(obj, &origins, config->extra_config);
- if (!config->binding_names.empty()) RenameRoot(doc, config->binding_names.back(), origins);
-
auto block = doc.as<StmtBlockDoc>();
if (block && !config->verbose_expr && !block.value()->stmts.empty()) {
if (auto expression = block.value()->stmts.back().as<ExprStmtDoc>()) {
auto original = origins.Get(expression.value()->expr);
if (original && original.value().same_as(obj)) {
- // A returned root expression carries the input origin. Statement
- // emissions retain their own input and child-expression origins.
- ffi::Array<StmtDoc> compact;
- if (MetadataLoader(block.value()->stmts.front()))
- compact.push_back(block.value()->stmts.front());
- compact.push_back(block.value()->stmts.back());
- block.value()->stmts = compact;
+ // Statement emissions retain their own input and child-expression origins.
+ block.value()->stmts = {block.value()->stmts.back()};
}
}
}
+ ffi::Optional<IdDoc> metadata_id;
+ auto metadata = CollectMetadata(doc, origins, &metadata_id);
+ ApplyDisplayNames(doc, config, origins, metadata_id);
ffi::Array<ffi::Any> header;
- if (block && !block.value()->stmts.empty()) {
- if (auto loader = MetadataLoader(block.value()->stmts.front())) {
- if (config->show_meta) {
- header.push_back(
- ffi::String(loader.value()->name == "tvm"
- ? "import tvm\n"
- : "import tvm as " + std::string(loader.value()->name) + "\n"));
- } else {
- block.value()->stmts.erase(block.value()->stmts.begin());
- block.value()->stmts.push_back(
- CommentDoc("Metadata omitted. Use show_meta=True in script() method to show it."));
- }
+ if (!metadata.empty()) {
+ if (config->show_meta) {
+ auto aliases = DisplayAliases(config);
+ std::unordered_set<std::string> used(aliases.begin(), aliases.end());
+ std::unordered_set<const ffi::Object*> visited;
+ VisitDocs(
+ doc,
+ [&](const Doc& child) {
+ if (auto id = child.as<IdDoc>()) used.insert(id.value()->name);
+ },
+ &visited);
+ std::string loader = "tvm";
+ for (size_t suffix = 1; used.count(loader); ++suffix)
+ loader = "tvm_" + std::to_string(suffix);
+ ffi::String json = ffi::json::Stringify(
+ ffi::ToJSONGraph(metadata, ffi::json::Object{{"tvm_version", TVM_VERSION}}), 2);
+ block.value()->stmts.insert(block.value()->stmts.begin(),
+ AssignDoc(metadata_id.value(),
+ IdDoc(loader)
+ ->Attr("ir")
+ ->Attr("load_json")
+ ->Call({LiteralDoc::Str(json, std::nullopt)}),
+ std::nullopt));
+ header.push_back(
+ ffi::String(loader == "tvm" ? "import tvm\n" : "import tvm as " + loader + "\n"));
+ } else {
+ block.value()->stmts.push_back(
+ CommentDoc("Metadata omitted. Use show_meta=True in script() method to show it."));
}
}
bool definition = false;
@@ -337,7 +385,8 @@ ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config) {
std::string canonical = key;
canonical.resize(canonical.size() - std::string(".prefix").size());
if (!namespaces.count(canonical)) continue;
- ffi::String alias = config->GetExtraConfig<ffi::String>(key, fallback);
+ ffi::String alias = canonical == "ir" ? config->ir_prefix
+ : config->GetExtraConfig<ffi::String>(key, fallback);
ffi::String import = "from tvm.script import " + canonical + " as " + std::string(alias);
if (config->GetExtraConfig<bool>("ir.comment_imports", false)) {
header.push_back(CommentDoc(import));
(tvm) 07/08: Align printer expectations with explicit blocks and generic display policy
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 3ddd9c249e186a69bbe7c60b5291b6a75ddfc18b
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:48:50 2026 +0000
Align printer expectations with explicit blocks and generic display policy
---
.../relax/script/test_relax_script_printer.py | 50 +++++++++++++----
tests/python/relax/test_analysis_well_formed.py | 2 +-
.../s_tir/script/test_s_tir_script_printer.py | 65 +++++++++++-----------
.../python/tirx/script/test_tirx_script_buffers.py | 28 +++++-----
.../python/tirx/script/test_tirx_script_printer.py | 22 ++++----
5 files changed, 98 insertions(+), 69 deletions(-)
diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py
index 173fbe23ba..a21098dc2b 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -116,6 +116,8 @@ def test_func():
_assert_print(
TestModule["foo"],
"""
+from __future__ import annotations
+
# from tvm.script import relax as R
@R.function
@@ -133,10 +135,9 @@ def test_module():
from __future__ import annotations
# from tvm.script import ir as I
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
-# from tvm.script import s_tir as Ts
# from tvm.script import relax as R
+# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@I.ir_module
class Module:
@@ -181,6 +182,8 @@ def test_function():
_assert_print_lines(
func,
"""
+from __future__ import annotations
+
# from tvm.script import relax as R
@R.function
@@ -230,7 +233,7 @@ def test_function_dependent_shape_source_spans():
definition, underline = render(cast_path.attr("dtype"))
dtype_start = definition.index(dtype_literal)
assert underline[dtype_start : dtype_start + len(dtype_literal)] == "^" * len(dtype_literal)
- assert underline.strip() == "^" * len(dtype_literal)
+ assert underline.strip() == "^" * len(expression)
definition, underline = render(cast_path.attr("value"))
variable_start = definition.index(expression) + expression.rindex("n")
@@ -248,6 +251,8 @@ def test_lone_private_function():
_assert_print_lines(
func,
"""
+from __future__ import annotations
+
# from tvm.script import relax as R
@R.function(private=True)
@@ -271,6 +276,8 @@ def test_extern_func():
_assert_print_lines(
obj,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
# from tvm.script import relax as R
@@ -342,6 +349,8 @@ def test_nested_function():
_assert_print_lines(
NestedFunction,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
# from tvm.script import relax as R
@@ -498,9 +507,10 @@ def test_prim_value():
_assert_print_lines(
func,
"""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import relax as R
+# from tvm.script import tirx as T
@R.function
def func() -> T.int64:
@@ -871,11 +881,12 @@ def test_module_cross_func_call():
_assert_print_lines(
TestModule,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
-# from tvm.script import s_tir as Ts
# from tvm.script import relax as R
+# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@I.ir_module
class Module:
@@ -895,11 +906,12 @@ class Module:
_assert_print_lines(
module_str,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
-# from tvm.script import s_tir as Ts
# from tvm.script import relax as R
+# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@I.ir_module
class Module:
@@ -926,6 +938,8 @@ def test_assert_op():
_assert_print_lines(
AssertOpMod,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
# from tvm.script import relax as R
@@ -950,6 +964,8 @@ def test_print():
_assert_print_lines(
PrintMod,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
# from tvm.script import relax as R
@@ -974,6 +990,8 @@ def test_private_function():
_assert_print_lines(
AddMod,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
# from tvm.script import relax as R
@@ -1027,6 +1045,8 @@ def test_directly_construct_private_funcs():
_assert_print_lines(
obj,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
# from tvm.script import relax as R
@@ -1072,6 +1092,8 @@ def test_reused_extern_func():
_assert_print_lines(
func,
"""
+from __future__ import annotations
+
# from tvm.script import relax as R
@R.function
@@ -1100,6 +1122,8 @@ def test_inline_extern_func():
_assert_print_lines(
func,
"""
+from __future__ import annotations
+
# from tvm.script import relax as R
@R.function
@@ -1148,6 +1172,8 @@ def test_hide_inferable_ty():
_assert_print_lines(
func.script(show_all_ty=False),
"""
+from __future__ import annotations
+
# from tvm.script import relax as R
@R.function
diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py
index daf834999a..a6fad5df7d 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -669,7 +669,7 @@ def test_impure_in_dataflow_block():
# The throwing form surfaces the offending impure call in its message.
with pytest.raises(Exception) as excinfo:
rx.analysis.well_formed(mod)
- assert "R.print" in str(excinfo.value)
+ assert 'I.Call("relax.print", ["{}", x], ty=R.Tuple())' in str(excinfo.value)
def test_well_formed_function():
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 a6a2d00bca..bea1c3b4fb 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
@@ -279,9 +279,8 @@ def test_bind():
_assert_print(
obj,
"""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func(private=True)
def main():
@@ -315,9 +314,8 @@ def test_remap():
v4, v5 = Ts.axis.remap("RS", [i4, i5])
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main():
@@ -361,20 +359,21 @@ def test_root_block():
TB.evaluate(0)
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main():
- # with Ts.sblock("root"):
- buffer = Ts.sblock_alloc_buffer((128, 128), "float32")
- for i in range(128):
- for j in range(128):
- with Ts.sblock(""):
- Ts.reads()
- Ts.writes()
- T.evaluate(0)"""
+ with Ts.sblock("root"):
+ Ts.reads()
+ Ts.writes()
+ v = Ts.sblock_alloc_buffer((128, 128), "float32")
+ for i in range(128):
+ for j in range(128):
+ with Ts.sblock(""):
+ Ts.reads()
+ Ts.writes()
+ T.evaluate(0)"""
_assert_print(root_block_implicitly.with_attr("global_symbol", "main"), expected_output)
_assert_print(root_block_explicitly.with_attr("global_symbol", "main"), expected_output)
@@ -409,9 +408,10 @@ def test_prim_func_different_symbol():
TB.evaluate(0)
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
@@ -504,9 +504,10 @@ def test_predicated_load_store():
)
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
@@ -548,9 +549,10 @@ def test_predicated_buffer_load_store():
).with_attr("s_tir", True)
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func(private=True)
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -577,9 +579,10 @@ def test_predicated_scalable_load_store():
)
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
@@ -617,9 +620,10 @@ def test_vload_with_explicit_scalable_data_type():
B[0 : TB.vscale() * 4] = A.vload([TB.Ramp(0, 1, TB.vscale() * 4)], dtype="float32xvscalex4")
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
@@ -637,9 +641,10 @@ def test_vectorize_llvm_pure_intrin():
)
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
@@ -660,9 +665,10 @@ def test_func_with_loop_jumps():
break
expected_output = """
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
@@ -2077,7 +2083,6 @@ def test_annotation_multi_access_paths():
assert (
result
== """# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
@Ts.prim_func
@@ -2105,7 +2110,6 @@ def test_annotate_from_multi_obj():
assert (
result
== """# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
@Ts.prim_func
@@ -2138,7 +2142,6 @@ def test_disable_concise_scoping_when_scope_annotated():
assert (
result
== """# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
@Ts.prim_func
diff --git a/tests/python/tirx/script/test_tirx_script_buffers.py b/tests/python/tirx/script/test_tirx_script_buffers.py
index 152c94ddd6..625c083c47 100644
--- a/tests/python/tirx/script/test_tirx_script_buffers.py
+++ b/tests/python/tirx/script/test_tirx_script_buffers.py
@@ -442,7 +442,7 @@ def test_buffer_local_ir():
# Round-trip
code = func.script()
assert (
- 'buffer_1 = T.decl_buffer((2,), "float16", data=buffer.data, '
+ 'v_1 = T.decl_buffer((2,), "float16", data=v.data, '
'scope="local", layout="default")'
) in code
assert from_source(code).script() == code
@@ -518,11 +518,11 @@ def test_buffer_local_physical_order():
code = func.script()
assert (
- 'buffer_1 = T.decl_buffer((32,), "float32", data=buffer.data, '
+ 'v_1 = T.decl_buffer((32,), "float32", data=v.data, '
'scope="local", layout="default")'
) in code
assert (
- 'buffer_2 = T.decl_buffer((4, 8), "float32", data=buffer.data, '
+ 'v_2 = T.decl_buffer((4, 8), "float32", data=v.data, '
'scope="local", layout="default")'
) in code
assert from_source(code).script() == code
@@ -553,13 +553,13 @@ def test_buffer_local_layout_overrides_roundtrip():
assert not b_custom.ty.layout.is_trivial()
code = func.script()
- storage_line = next(line for line in code.splitlines() if "buffer_1 =" in line)
- custom_line = next(line for line in code.splitlines() if "buffer_2 =" in line)
+ storage_line = next(line for line in code.splitlines() if "v_1 =" in line)
+ custom_line = next(line for line in code.splitlines() if "v_2 =" in line)
assert (
- 'T.decl_buffer((32,), "float32", data=buffer.data, scope="local", layout='
+ 'T.decl_buffer((32,), "float32", data=v.data, scope="local", layout='
) in storage_line
assert (
- 'T.decl_buffer((2, 4), "float32", data=buffer.data, scope="local", '
+ 'T.decl_buffer((2, 4), "float32", data=v.data, scope="local", '
'layout='
) in custom_line
assert_structural_equal(func, from_source(code))
@@ -608,7 +608,7 @@ def test_buffer_local_compose_layout_printer_roundtrip():
assert [int(dim) for dim in b_buf.ty.shape] == [64]
assert b_buf.ty.layout.is_trivial()
code = func.script()
- local_line = next(line for line in code.splitlines() if "buffer =" in line)
+ local_line = next(line for line in code.splitlines() if "v =" in line)
assert 'T.decl_buffer((64,), "float32", data=A.data, scope="local", layout=' in local_line
parsed = from_source(code)
assert_structural_equal(func, parsed)
@@ -661,8 +661,8 @@ def test_buffer_local_physical_span_includes_gaps_and_offset():
assert int(b_storage.ty.layout.apply(1, shape=list(b_storage.ty.shape))["m"]) == 5
code = func.script()
- storage_line = next(line for line in code.splitlines() if "buffer_3 =" in line)
- assert 'T.decl_buffer((2,), "float32", data=buffer.data, scope="local", layout=' in storage_line
+ storage_line = next(line for line in code.splitlines() if "v_3 =" in line)
+ assert 'T.decl_buffer((2,), "float32", data=v.data, scope="local", layout=' in storage_line
assert_structural_equal(func, from_source(code))
assert from_source(code).script() == code
@@ -685,16 +685,16 @@ def test_buffer_local_printer_is_stable_with_multiple_aliases():
expected = func.script()
assert (
- 'buffer_1 = T.decl_buffer((32,), "float32", data=buffer.data, '
+ 'v_1 = T.decl_buffer((32,), "float32", data=v.data, '
'scope="local", layout="default")'
) in expected
assert (
- 'buffer_2 = T.decl_buffer((4, 8), "float32", data=buffer.data, '
+ 'v_2 = T.decl_buffer((4, 8), "float32", data=v.data, '
'scope="local", layout="default")'
) in expected
- storage_line = next(line for line in expected.splitlines() if "buffer_3 =" in line)
+ storage_line = next(line for line in expected.splitlines() if "v_3 =" in line)
assert (
- 'T.decl_buffer((32,), "float32", data=buffer.data, scope="local", layout='
+ 'T.decl_buffer((32,), "float32", data=v.data, scope="local", layout='
) in storage_line
for _ in range(20):
parsed = from_source(expected)
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py b/tests/python/tirx/script/test_tirx_script_printer.py
index 09d1e03ed7..6e455f934e 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -228,7 +228,7 @@ def test_allocate():
_assert_print(
obj.body,
"""
-buffer = T.alloc_buffer((128, 128), "float32", layout="default")
+v = T.alloc_buffer((128, 128), "float32", layout="default")
T.evaluate(1)
""",
)
@@ -246,8 +246,8 @@ def test_allocate_with_decl_buffer_sugar():
_assert_print(
obj.body,
"""
-buffer = T.alloc_buffer((128, 128), "float32", layout="default")
-buffer_1 = T.decl_buffer((128, 128), "float32", data=buffer.data, layout="default")
+v = T.alloc_buffer((128, 128), "float32", layout="default")
+v_1 = T.decl_buffer((128, 128), "float32", data=v.data, layout="default")
T.evaluate(1)
""",
)
@@ -265,9 +265,9 @@ def test_allocate_with_decl_buffer_sugar_multi_usage():
_assert_print(
obj.body,
"""
-buffer = T.alloc_buffer((128, 128), "float32", layout="default")
-buffer_1 = T.decl_buffer((128, 128), "float32", data=buffer.data, layout="default")
-T.evaluate(buffer.data)
+v = T.alloc_buffer((128, 128), "float32", layout="default")
+v_1 = T.decl_buffer((128, 128), "float32", data=v.data, layout="default")
+T.evaluate(v.data)
""",
)
@@ -283,9 +283,9 @@ def test_allocate_with_decl_buffer_no_sugar_mismatch():
_assert_print(
obj.body,
"""
-buffer = T.alloc_buffer((128, 128), "float32", layout="default")
-buffer_1 = T.decl_buffer((256, 256), "float32", data=buffer.data, layout="default")
-T.evaluate(buffer.data)
+v = T.alloc_buffer((128, 128), "float32", layout="default")
+v_1 = T.decl_buffer((256, 256), "float32", data=v.data, layout="default")
+T.evaluate(v.data)
""",
)
@@ -302,8 +302,8 @@ def test_decl_buffer():
_assert_print(
obj.body,
"""
-v: T.handle("float32", "global") = T.handle("float32", "global")
-buffer = T.decl_buffer((10, 10), "float32", data=v, layout="default")
+v_1: T.handle("float32", "global") = T.handle("float32", "global")
+v = T.decl_buffer((10, 10), "float32", data=v_1, layout="default")
T.evaluate(1)
""",
)
(tvm) 02/08: [FIX][TIRx] Preserve storage synchronization calls in named syntax
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 10b4a14e9a000ac17a31dd200b413c103a1b87b2
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:09:03 2026 +0000
[FIX][TIRx] Preserve storage synchronization calls in named syntax
---
python/tvm/tirx/op.py | 21 ++++++++++---
src/tirx/script/printer/expr.cc | 35 ++++++++++++++++++----
.../transform/test_s_tir_transform_thread_sync.py | 9 ++----
3 files changed, 49 insertions(+), 16 deletions(-)
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index d583d8930f..4eff21210a 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -833,7 +833,7 @@ def tvm_thread_invariant(cond):
return call_intrin(_primexpr_ty(cond), "tirx.tvm_thread_invariant", cond)
-def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1):
+def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1, *, dtype="void"):
"""Perform synchronization in specified scope.
Parameters
@@ -841,18 +841,31 @@ def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1):
storage_scope : str
The storage scope to perform synchronization.
- is_load : bool
+ is_load : bool or Expr or None
Whether to perform load synchronization. (for global sync only)
+ Set both ``is_load`` and ``num_blocks`` to None to omit these operands.
- num_blocks : int
+ num_blocks : int or Expr or None
The number of blocks to synchronize. (for global sync only)
+ Set to None to omit this operand.
+
+ dtype : str or tvm.ir.Type
+ The stored result type. Defaults to void.
Returns
-------
call : Expr
The call expression.
"""
- return call_intrin("void", "tirx.tvm_storage_sync", storage_scope, is_load, num_blocks)
+ args = [storage_scope]
+ if is_load is None:
+ if num_blocks is not None:
+ raise ValueError("num_blocks must be None when is_load is omitted")
+ else:
+ args.append(is_load)
+ if num_blocks is not None:
+ args.append(num_blocks)
+ return call_intrin(dtype, "tirx.tvm_storage_sync", *args)
def tvm_kernel_replace_point():
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index 45c4437ccf..87f07f2fd5 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -181,6 +181,35 @@ bool CanTranslateExplicitResultCall(const CallNode* call) {
});
}
+ffi::Optional<ExprDoc> StorageSyncDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* call =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(input);
+ if (!CanTranslateExplicitResultCall(call) || call->args.empty() || call->args.size() > 3) {
+ return RawCall(d, call, false);
+ }
+ ffi::Array<ExprDoc> args;
+ for (const Expr& arg : call->args) {
+ args.push_back(MaterializeCallArgument(d, arg, d->Translate(arg).value()));
+ }
+ // Explicit None operands suppress the helper's defaults without changing
+ // the stored argument list of native one- and two-operand calls.
+ while (args.size() < 3) args.push_back(LiteralDoc::None(std::nullopt));
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ if (!ffi::StructuralEqual()(call->ty, PrimType::Void())) {
+ keys.push_back("dtype");
+ values.push_back(TypeValue(d, call->ty));
+ }
+ return NamespaceDoc("tirx")->Attr("tvm_storage_sync")->Call(args, keys, values);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ OpDef("tirx.tvm_storage_sync")
+ .set_attr<FDocTranslate>(kOpCallDocTranslate,
+ FDocTranslate::FromNative<&StorageSyncDocTranslate>());
+}
+
ffi::Optional<ExprDoc> CallExternDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
const ffi::Object*) {
const auto* call =
@@ -446,12 +475,6 @@ ffi::Optional<ExprDoc> TIRCallPrefixDocTranslate(DocTranslatorObj* d, const Call
return RawCall(d, call, false);
}
}
- // The named storage-sync helper supplies two default operands. A shorter
- // native Call must keep its actual operand list and explicit stored type.
- if (call->op.same_as(tirx::builtin::tvm_storage_sync()) &&
- (call->args.size() != 3 || !ffi::StructuralEqual()(call->ty, PrimType::Void()))) {
- return RawCall(d, call, false);
- }
if (auto op = call->op.as<Op>(); op && op.value()->name == "tirx.isnan") {
auto input_type = call->args.size() == 1 ? call->args[0]->ty.as<PrimType>() : std::nullopt;
// The named helper folds constants and widens float16 operands, and its
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
index 93587b8022..5d5a52e5f6 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
@@ -59,7 +59,7 @@ def test_sync_read_thread_id_independent_location():
result_local[0] = result_local[0] + temp_shared[0] * p1[1]
mod = run_passes(func)
- assert 'I.Call("tirx.tvm_storage_sync", ["shared"], ty="int32")' in str(mod)
+ assert 'T.tvm_storage_sync("shared", None, None, dtype="int32")' in str(mod)
def test_sync_inside_condition():
@@ -101,7 +101,7 @@ def test_sync_inside_condition():
for func in (func2, func3):
mod = s_tir.transform.ThreadSync("shared")(tvm.IRModule.from_expr(func))
- assert 'I.Call("tirx.tvm_storage_sync", ["shared"], ty="int32")' in mod.script()
+ assert 'T.tvm_storage_sync("shared", None, None, dtype="int32")' in mod.script()
def test_sync_shared_dyn():
@@ -165,10 +165,7 @@ def test_sync_shared_aliasing_buffer_views():
# In addition to the explicit write-to-read barrier, the shared physical
# storage needs a read-to-next-write barrier across loop iterations.
- assert (
- str(mod["main"]).count("T.tvm_storage_sync")
- + str(mod["main"]).count('I.Call("tirx.tvm_storage_sync"')
- ) == 2
+ assert str(mod["main"]).count("T.tvm_storage_sync") == 2
@pytest.mark.gpu
(tvm) 08/08: Keep complete output expectations aligned with renderer headers
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 6927b91cf913ce9d89637abc7d7fbc7d8b46dd94
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:54:39 2026 +0000
Keep complete output expectations aligned with renderer headers
---
.../relax/script/test_relax_script_printer.py | 2 +-
.../s_tir/script/test_s_tir_script_printer.py | 38 ++++++++++++----------
.../python/tirx/script/test_tirx_script_printer.py | 1 -
.../script/test_tirx_script_printer_underlining.py | 7 ++--
4 files changed, 23 insertions(+), 25 deletions(-)
diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py
index a21098dc2b..08b12bc2e8 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -1179,7 +1179,7 @@ from __future__ import annotations
@R.function
def func(A: R.Tensor((10, 20), dtype="float32"), B: R.Tensor(dtype="float32", ndim=2)) -> R.Tensor((10, 20), dtype="float32"):
B2 = R.match_cast(B, R.Tensor((10, 20), dtype="float32"))
- C = R.add(A, B2)
+ C: R.Tensor((10, 20), dtype="float32") = R.add(A, B2)
D = C
E: R.Tensor((10, 20), dtype="float32") = R.add(D, B)
return E""",
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 bea1c3b4fb..5d9c9b07a8 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
@@ -129,9 +129,10 @@ def test_prim_func():
_assert_print(
func,
expected="""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -154,9 +155,10 @@ def test_prim_func_buffer_data_use():
_assert_print(
func,
expected="""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -181,9 +183,10 @@ def test_prim_func_buffer_data_argument_is_scope_hint():
_assert_print(
func,
expected="""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -389,9 +392,10 @@ def test_private_primfunc():
_assert_print(
func,
expected="""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func(private=True)
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -464,9 +468,8 @@ def test_custom_float_types(dtype):
TB.evaluate(getattr(TB, dtype)(0.0))
expected_output = f"""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def func():
@@ -2082,8 +2085,8 @@ def test_annotation_multi_access_paths():
)
assert (
result
- == """# from tvm.script import tirx as T
-# from tvm.script import s_tir as Ts
+ == """# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main():
@@ -2109,8 +2112,8 @@ def test_annotate_from_multi_obj():
)
assert (
result
- == """# from tvm.script import tirx as T
-# from tvm.script import s_tir as Ts
+ == """# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main():
@@ -2141,8 +2144,8 @@ def test_disable_concise_scoping_when_scope_annotated():
)
assert (
result
- == """# from tvm.script import tirx as T
-# from tvm.script import s_tir as Ts
+ == """# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main():
@@ -2169,9 +2172,8 @@ def test_ir_module():
mod,
"""
# from tvm.script import ir as I
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@I.ir_module
class Module:
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py b/tests/python/tirx/script/test_tirx_script_printer.py
index 6e455f934e..b55e962323 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -674,7 +674,6 @@ def test_print_kwargs_schedule_op_full_code():
expected = (
"# from tvm.script import tirx as T\n"
- "# from tvm.tirx.layout import Axis\n"
"\n"
"@T.prim_func\n"
"def test():\n"
diff --git a/tests/python/tirx/script/test_tirx_script_printer_underlining.py b/tests/python/tirx/script/test_tirx_script_printer_underlining.py
index 854ac67158..6d8706d763 100644
--- a/tests/python/tirx/script/test_tirx_script_printer_underlining.py
+++ b/tests/python/tirx/script/test_tirx_script_printer_underlining.py
@@ -39,8 +39,9 @@ def test_underline_from_obj():
)
assert result == format_script(
"""
+ from __future__ import annotations
+
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@T.prim_func
def main(a: T.int32, b: T.int32):
@@ -102,7 +103,6 @@ def test_underline_from_multi_obj():
assert result == format_script(
"""
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@T.prim_func
def main():
@@ -136,7 +136,6 @@ def test_underline_func():
assert result == format_script(
"""
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@T.prim_func
^^^^^^^^^^^^
@@ -165,7 +164,6 @@ def test_underline_func_in_irmodule():
"""
# from tvm.script import ir as I
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@I.ir_module
class Module:
@@ -196,7 +194,6 @@ def test_underline_irmodule():
"""
# from tvm.script import ir as I
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@I.ir_module
^^^^^^^^^^^^
(tvm) 04/08: [CI] Use Python 3.12 for TVMScript round trips
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 ed13ec061b0e2b7fd459bf704e7f1aebe7af383f
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:06:15 2026 +0000
[CI] Use Python 3.12 for TVMScript round trips
---
.github/actions/setup/action.yml | 10 ++++++++--
.github/workflows/main.yml | 6 ++++++
docs/contribute/testing.rst | 5 +++++
docs/reference/api/python/script/printer.rst | 10 ++++++++++
4 files changed, 29 insertions(+), 2 deletions(-)
diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml
index 980e5d087f..0ce25f3e4f 100644
--- a/.github/actions/setup/action.yml
+++ b/.github/actions/setup/action.yml
@@ -1,6 +1,12 @@
name: Setup TVM CI environment
description: Set up the cached conda build environment and install tvm-ffi.
+inputs:
+ python-version:
+ description: Python version for the build and test environment.
+ required: false
+ default: "3.10"
+
runs:
using: "composite"
steps:
@@ -19,7 +25,7 @@ runs:
environment-file: ci/scripts/package/build-environment.yaml
auto-activate-base: false
miniforge-version: latest
- python-version: "3.10"
+ python-version: ${{ inputs.python-version }}
condarc-file: tests/conda/condarc
conda-remove-defaults: true
- uses: conda-incubator/setup-miniconda@v4.0.1
@@ -31,7 +37,7 @@ runs:
auto-activate-base: false
miniforge-version: latest
use-only-tar-bz2: true
- python-version: "3.10"
+ python-version: ${{ inputs.python-version }}
condarc-file: tests/conda/condarc
conda-remove-defaults: true
- name: Conda info
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 7a524732cf..a177aa4054 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -47,6 +47,9 @@ jobs:
submodules: 'recursive'
- name: Set up environment
uses: ./.github/actions/setup
+ with:
+ # Generated TVMScript type parameter headers require Python 3.12+.
+ python-version: "3.12"
- name: Detect environment variables
uses: ./.github/actions/detect-env-vars
id: env_vars
@@ -103,6 +106,9 @@ jobs:
submodules: 'recursive'
- name: Set up environment
uses: ./.github/actions/setup
+ with:
+ # Generated TVMScript type parameter headers require Python 3.12+.
+ python-version: "3.12"
- name: Detect environment variables
uses: ./.github/actions/detect-env-vars
id: env_vars
diff --git a/docs/contribute/testing.rst b/docs/contribute/testing.rst
index 8dd71ca4ca..af50ea714f 100644
--- a/docs/contribute/testing.rst
+++ b/docs/contribute/testing.rst
@@ -212,6 +212,11 @@ Running Locally
To run the Python unit tests locally, use the command ``pytest`` in
the ``${TVM_HOME}`` directory.
+Use Python 3.12 or newer for TVMScript printer and round-trip tests, including
+when running the full Python test suite. Eligible printed functions use PEP 695
+type parameter headers, which older Python versions cannot parse. This generated
+syntax requirement is separate from the package's minimum Python version.
+
- Environment variables
- ``TVM_TEST_TARGETS`` should be a semicolon-separated list of
targets to run. If unset, will default to the targets defined in
diff --git a/docs/reference/api/python/script/printer.rst b/docs/reference/api/python/script/printer.rst
index cc16c55c3b..adb8bcc34f 100644
--- a/docs/reference/api/python/script/printer.rst
+++ b/docs/reference/api/python/script/printer.rst
@@ -18,6 +18,16 @@
tvm.script.printer
-------------------------
+Generated syntax
+****************
+
+Printed functions use PEP 695 type parameter headers, such as
+``def main[n](...)``, when their symbolic parameters can be declared locally.
+Parsing or executing this generated syntax requires Python 3.12 or newer.
+The printer uses the same syntax on older Python versions; it does not select
+a compatibility spelling based on the host interpreter. Symbols shared across
+functions retain their shared bindings instead of becoming local type parameters.
+
tvm.script.printer
*************************
.. automodule:: tvm.script.printer