This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tqchen/printer-refinement in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 28c119e031b5e6136bdad6da8579d19234dc5cb3 Author: Tianqi Chen <[email protected]> AuthorDate: Thu Oct 1 19:53:33 2026 +0000 [REFACTOR][Script] Expose namespace aliases through the printer API --- include/tvm/script/printer/printer.h | 11 ++++++++++ src/relax/script/printer/function.cc | 4 ++-- src/s_tir/script/printer/config.cc | 5 ++--- src/script/printer/config.cc | 24 +++++++++++---------- src/script/printer/dialect_prefix.h | 41 ------------------------------------ src/script/printer/doc_printer.cc | 6 ++---- src/script/printer/doc_translator.cc | 9 +++++--- src/script/printer/ir/ir.cc | 4 ++-- src/script/printer/script_printer.cc | 2 -- src/tirx/script/printer/function.cc | 4 ++-- 10 files changed, 40 insertions(+), 70 deletions(-) diff --git a/include/tvm/script/printer/printer.h b/include/tvm/script/printer/printer.h index 83f313013b..22d4c52e56 100644 --- a/include/tvm/script/printer/printer.h +++ b/include/tvm/script/printer/printer.h @@ -23,6 +23,7 @@ #ifndef TVM_SCRIPT_PRINTER_PRINTER_H_ #define TVM_SCRIPT_PRINTER_PRINTER_H_ +#include <tvm/ffi/container/map.h> #include <tvm/ffi/optional.h> #include <tvm/script/printer/config.h> @@ -43,6 +44,16 @@ TVM_DLL std::string Script(const ffi::ObjectRef& node, namespace script { namespace printer { +/*! + * \brief Register a namespace alias during dialect static initialization. + * \param key The existing prefix configuration key, such as "tirx.prefix". + * \param default_alias The alias reserved before translation assigns variable names. + */ +TVM_DLL void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& default_alias); + +/*! \brief Read the registered namespace aliases. */ +TVM_DLL const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases(); + /*! * \brief Translate IR, recover diagnostic paths, and render Python text. * \param obj The input IR object. diff --git a/src/relax/script/printer/function.cc b/src/relax/script/printer/function.cc index 8913af7eca..42cba66e0a 100644 --- a/src/relax/script/printer/function.cc +++ b/src/relax/script/printer/function.cc @@ -18,6 +18,7 @@ */ #include <tvm/ffi/extra/structural_equal.h> #include <tvm/ir/function.h> +#include <tvm/script/printer/printer.h> #include <tvm/tirx/type.h> #include <algorithm> @@ -25,7 +26,6 @@ #include <utility> #include <vector> -#include "../../../script/printer/dialect_prefix.h" #include "../../../script/printer/ir/utils.h" #include "utils.h" @@ -116,7 +116,7 @@ ffi::Optional<ExprDoc> FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in } TVM_FFI_STATIC_INIT_BLOCK() { - RegisterDialectPrefix("relax.prefix", "R"); + RegisterNamespaceAlias("relax.prefix", "R"); ffi::reflection::TypeAttrDef<relax::FunctionNode>().attr( kDocTranslate, FDocTranslate::FromNative<&FunctionDocTranslate>()); } diff --git a/src/s_tir/script/printer/config.cc b/src/s_tir/script/printer/config.cc index 461858edad..c0e135d7de 100644 --- a/src/s_tir/script/printer/config.cc +++ b/src/s_tir/script/printer/config.cc @@ -17,14 +17,13 @@ * under the License. */ #include <tvm/ffi/reflection/registry.h> - -#include "../../../script/printer/dialect_prefix.h" +#include <tvm/script/printer/printer.h> namespace tvm { namespace script { namespace printer { -TVM_FFI_STATIC_INIT_BLOCK() { RegisterDialectPrefix("s_tir.prefix", "Ts"); } +TVM_FFI_STATIC_INIT_BLOCK() { RegisterNamespaceAlias("s_tir.prefix", "Ts"); } } // namespace printer } // namespace script diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc index 08462543d8..15a1d8bf1a 100644 --- a/src/script/printer/config.cc +++ b/src/script/printer/config.cc @@ -17,22 +17,21 @@ * under the License. */ #include <tvm/script/printer/config.h> +#include <tvm/script/printer/printer.h> #include <algorithm> #include <cctype> -#include <map> #include <string> #include <unordered_set> #include <utility> - -#include "./dialect_prefix.h" +#include <vector> namespace tvm { namespace { -std::map<ffi::String, ffi::String>& DialectPrefixes() { - static std::map<ffi::String, ffi::String> prefixes; - return prefixes; +ffi::Map<ffi::String, ffi::String>& NamespaceAliases() { + static ffi::Map<ffi::String, ffi::String> aliases; + return aliases; } bool IsIdentifier(const std::string& name) { @@ -54,12 +53,12 @@ bool IsIdentifier(const std::string& name) { namespace script { namespace printer { -void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix) { - TVM_FFI_ICHECK(DialectPrefixes().emplace(key, default_prefix).second) - << "Duplicate printer dialect prefix: " << key; +void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& default_alias) { + TVM_FFI_ICHECK(!NamespaceAliases().count(key)) << "Duplicate printer dialect prefix: " << key; + NamespaceAliases().Set(key, default_alias); } -const std::map<ffi::String, ffi::String>& GetDialectPrefixes() { return DialectPrefixes(); } +const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases() { return NamespaceAliases(); } } // namespace printer } // namespace script @@ -172,7 +171,10 @@ ffi::Array<ffi::String> PrinterConfigNode::GetBuiltinKeywords() { TVM_FFI_ICHECK(module_alias.empty() || IsIdentifier(std::string(module_alias))) << "Invalid `module_alias`: " << module_alias; ffi::Array<ffi::String> result{ir_prefix}; - for (const auto& [key, default_prefix] : DialectPrefixes()) { + const auto& aliases = NamespaceAliases(); + std::vector<std::pair<ffi::String, ffi::String>> ordered(aliases.begin(), aliases.end()); + std::sort(ordered.begin(), ordered.end()); + for (const auto& [key, default_prefix] : ordered) { ffi::String prefix = GetExtraConfig<ffi::String>(key, default_prefix); TVM_FFI_ICHECK(IsIdentifier(std::string(prefix))) << "Invalid `" << key << "`: " << prefix; result.push_back(prefix); diff --git a/src/script/printer/dialect_prefix.h b/src/script/printer/dialect_prefix.h deleted file mode 100644 index 0eac91b5ea..0000000000 --- a/src/script/printer/dialect_prefix.h +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -#ifndef TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_ -#define TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_ - -#include <tvm/ffi/string.h> - -#include <map> - -namespace tvm { -namespace script { -namespace printer { - -// Register during dialect static initialization so configuration can validate -// and reserve the prefix before translation assigns variable names. -void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix); - -// Read the existing registration table without introducing a second namespace registry. -const std::map<ffi::String, ffi::String>& GetDialectPrefixes(); - -} // namespace printer -} // namespace script -} // namespace tvm - -#endif // TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_ diff --git a/src/script/printer/doc_printer.cc b/src/script/printer/doc_printer.cc index 302ba6d069..37080b6c4d 100644 --- a/src/script/printer/doc_printer.cc +++ b/src/script/printer/doc_printer.cc @@ -24,6 +24,7 @@ #include <tvm/ffi/reflection/registry.h> #include <tvm/ir/object_functor.h> #include <tvm/ir/prim/expr.h> +#include <tvm/script/printer/printer.h> #include <algorithm> #include <cmath> @@ -42,7 +43,6 @@ #include "../../support/str_escape.h" #include "../../support/utils.h" -#include "dialect_prefix.h" namespace tvm { namespace script { @@ -1158,9 +1158,7 @@ void PythonDocPrinter::PrintTypedDoc(const IdDoc& doc) { output_ << doc->name; } void PythonDocPrinter::PrintTypedDoc(const NamespaceDoc& doc) { const ffi::String& name = doc->canonical_name; ffi::String key = std::string(name) + ".prefix"; - const auto& prefixes = GetDialectPrefixes(); - auto registered = prefixes.find(key); - ffi::String fallback = registered == prefixes.end() ? name : registered->second; + ffi::String fallback = GetNamespaceAliases().Get(key).value_or(name); output_ << (name == "ir" ? config()->ir_prefix : config()->GetExtraConfig<ffi::String>(key, fallback)); } diff --git a/src/script/printer/doc_translator.cc b/src/script/printer/doc_translator.cc index 6ca872d9a4..96bc70a7ee 100644 --- a/src/script/printer/doc_translator.cc +++ b/src/script/printer/doc_translator.cc @@ -24,6 +24,7 @@ #include <tvm/runtime/base.h> #include <tvm/script/printer/doc.h> #include <tvm/script/printer/doc_translator.h> +#include <tvm/script/printer/printer.h> #include <algorithm> #include <cctype> @@ -32,8 +33,6 @@ #include <utility> #include <vector> -#include "dialect_prefix.h" - namespace tvm { namespace script { namespace printer { @@ -149,7 +148,11 @@ details::TranslationResult TranslateComplete(ffi::AnyView obj, TranslationEngine ffi::Map<ffi::String, ffi::Any> extra_config) { auto binding_names = engine->GetExtraConfig<ffi::Array<ffi::String>>("ir.binding_names", {}); // Reserve the same registered or explicitly configured aliases used by rendering. - for (const auto& [key, fallback] : GetDialectPrefixes()) { + // Keep key order so overlapping aliases reserve the same generated names. + const auto& aliases = GetNamespaceAliases(); + std::vector<std::pair<ffi::String, ffi::String>> ordered(aliases.begin(), aliases.end()); + std::sort(ordered.begin(), ordered.end()); + for (const auto& [key, fallback] : ordered) { if (!extra_config.count(key)) engine->AllocId(fallback); } for (const auto& [key, value] : extra_config) { diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc index 11982d4e1d..5e561b421b 100644 --- a/src/script/printer/ir/ir.cc +++ b/src/script/printer/ir/ir.cc @@ -20,6 +20,7 @@ #include <tvm/ir/prim/op.h> #include <tvm/relax/expr.h> #include <tvm/relax/type.h> +#include <tvm/script/printer/printer.h> #include <tvm/tirx/function.h> #include <algorithm> @@ -27,7 +28,6 @@ #include <utility> #include <vector> -#include "../dialect_prefix.h" #include "utils.h" namespace tvm { @@ -161,7 +161,7 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView in } TVM_FFI_STATIC_INIT_BLOCK() { - RegisterDialectPrefix("ir.prefix", "I"); + RegisterNamespaceAlias("ir.prefix", "I"); ffi::reflection::TypeAttrDef<IRModuleNode>().attr( kDocTranslate, FDocTranslate::FromNative<&IRModuleDocTranslate>()); } diff --git a/src/script/printer/script_printer.cc b/src/script/printer/script_printer.cc index 878b6f2440..56d831ad55 100644 --- a/src/script/printer/script_printer.cc +++ b/src/script/printer/script_printer.cc @@ -42,8 +42,6 @@ #include <sstream> #include <utility> -#include "./dialect_prefix.h" - namespace tvm { namespace { diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc index 7a8fdad386..4a043bd1da 100644 --- a/src/tirx/script/printer/function.cc +++ b/src/tirx/script/printer/function.cc @@ -21,6 +21,7 @@ #include <tvm/ir/global_info.h> #include <tvm/ir/prim/op.h> #include <tvm/s_tir/stmt.h> +#include <tvm/script/printer/printer.h> #include <tvm/tirx/function.h> #include <tvm/tirx/stmt_functor.h> @@ -31,7 +32,6 @@ #include <utility> #include <vector> -#include "../../../script/printer/dialect_prefix.h" #include "../../../script/printer/ir/utils.h" #include "utils.h" @@ -239,7 +239,7 @@ ffi::Optional<ExprDoc> PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView in } TVM_FFI_STATIC_INIT_BLOCK() { - RegisterDialectPrefix("tirx.prefix", "T"); + RegisterNamespaceAlias("tirx.prefix", "T"); ffi::reflection::TypeAttrDef<tirx::PrimFuncNode>().attr( kDocTranslate, FDocTranslate::FromNative<&PrimFuncDocTranslate>()); }
