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