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