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 <[email protected]> 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):
