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
The following commit(s) were added to refs/heads/tqchen/printer-refinement by
this push:
new 42170ee91c [FIX][Script] Preserve expression and parameter references
42170ee91c is described below
commit 42170ee91cbb10cdf56d3346afc2d0ce806681ee
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 1 16:29:09 2026 +0000
[FIX][Script] Preserve expression and parameter references
Keep compact expressions free of deferred declarations, emit pending
variables in encounter order, retain shared parameter identities and map
parameter diagnostics to their visible components.
---
src/script/printer/doc_translator.cc | 8 ++++----
src/script/printer/printer.cc | 30 ++++++++++++++++++------------
src/tirx/script/printer/function.cc | 22 +++++++++++++++++++---
3 files changed, 41 insertions(+), 19 deletions(-)
diff --git a/src/script/printer/doc_translator.cc
b/src/script/printer/doc_translator.cc
index a2fe6ce197..cdf9e0fa73 100644
--- a/src/script/printer/doc_translator.cc
+++ b/src/script/printer/doc_translator.cc
@@ -169,11 +169,12 @@ details::TranslationResult TranslateComplete(ffi::AnyView
obj, TranslationEngine
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()) {
- if (!engine->GetExtraConfig<bool>("script.verbose_expr", true))
- engine->CurrentScopeDocs().clear();
+ compact_expression =
!engine->GetExtraConfig<bool>("script.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()) {
@@ -187,10 +188,9 @@ details::TranslationResult TranslateComplete(ffi::AnyView
obj, TranslationEngine
}
});
auto root_declarations = engine->WithDocScope([&]() {
+ if (compact_expression) return;
auto pending = engine->GetImplicitDefs();
std::vector<std::pair<Var, IdDoc>> remaining(pending.begin(),
pending.end());
- std::sort(remaining.begin(), remaining.end(),
- [](const auto& a, const auto& b) { return a.second->name <
b.second->name; });
for (const auto& [var, _] : remaining) {
if (!pending.count(var)) continue;
if (auto rhs = engine->Translate(var, var)) {
diff --git a/src/script/printer/printer.cc b/src/script/printer/printer.cc
index dafae45eda..1c9b184df2 100644
--- a/src/script/printer/printer.cc
+++ b/src/script/printer/printer.cc
@@ -72,7 +72,7 @@ void RecoverDocPaths(
ffi::AnyView value, const std::vector<AccessPath>& enclosing_ir_paths,
const std::unordered_map<const ffi::Object*, ffi::ObjectRef>& origins,
const std::unordered_map<const ffi::Object*, std::vector<AccessPath>>&
ir_paths,
- std::unordered_set<const ffi::Object*>* active) {
+ std::unordered_set<const ffi::Object*>* active, bool function_parameter =
false) {
const ffi::Object* object = value.as<ffi::Object>();
if (!object || !active->insert(object).second) return;
std::vector<AccessPath> context = enclosing_ir_paths;
@@ -90,10 +90,15 @@ void RecoverDocPaths(
}
}
if (context.empty()) context = candidates->second;
- for (const AccessPath& path : context) {
- if (std::none_of(doc->source_paths.begin(), doc->source_paths.end(),
- [&](const AccessPath& existing) { return
SamePath(existing, path); })) {
- doc->source_paths.push_back(path);
+ // Parameter wrappers provide occurrence context; the name and type
+ // carry their own visible ranges.
+ if (!(function_parameter && value.as<AssignDocNode>())) {
+ for (const AccessPath& path : context) {
+ if (std::none_of(
+ doc->source_paths.begin(), doc->source_paths.end(),
+ [&](const AccessPath& existing) { return
SamePath(existing, path); })) {
+ doc->source_paths.push_back(path);
+ }
}
}
}
@@ -101,19 +106,20 @@ void RecoverDocPaths(
}
if (auto array = value.as<ffi::Array<ffi::Any>>()) {
for (const auto& child : array.value()) {
- RecoverDocPaths(child, context, origins, ir_paths, active);
+ RecoverDocPaths(child, context, origins, ir_paths, active,
function_parameter);
}
} else if (auto map = value.as<ffi::Map<ffi::Any, ffi::Any>>()) {
for (const auto& entry : map.value()) {
RecoverDocPaths(entry.second, context, origins, ir_paths, active);
}
} else {
- ffi::reflection::ForEachFieldInfo(TVMFFIGetTypeInfo(object->type_index()),
- [&](const TVMFFIFieldInfo* field) {
- if (ffi::String(field->name) ==
"source_paths") return;
-
RecoverDocPaths(ffi::reflection::FieldGetter(field)(object),
- context, origins,
ir_paths, active);
- });
+ ffi::reflection::ForEachFieldInfo(
+ TVMFFIGetTypeInfo(object->type_index()), [&](const TVMFFIFieldInfo*
field) {
+ if (ffi::String(field->name) == "source_paths") return;
+ RecoverDocPaths(ffi::reflection::FieldGetter(field)(object),
context, origins, ir_paths,
+ active,
+ value.as<FunctionDocNode>() &&
ffi::String(field->name) == "args");
+ });
}
active->erase(object);
}
diff --git a/src/tirx/script/printer/function.cc
b/src/tirx/script/printer/function.cc
index 59bbd3c214..23fef1f784 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -27,6 +27,7 @@
#include <algorithm>
#include <optional>
#include <unordered_map>
+#include <unordered_set>
#include <utility>
#include <vector>
@@ -90,19 +91,34 @@ ffi::Optional<ExprDoc>
PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView in
FunctionDoc doc(ffi::UnsafeInit{});
{
ffi::Array<AssignDoc> args;
+ std::unordered_set<const VarNode*> signature_symbols;
+ for (const Var& var : func->params) {
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+ var->ty, [&](const Var& symbol) -> ffi::Expected<ffi::WalkResult> {
+ signature_symbols.insert(symbol.get());
+ return ffi::WalkResult::Skip();
+ });
+ }
ffi::Array<IdDoc> param_ids;
- for (const Var& var : func->params) param_ids.push_back(VarDoc(d, var));
+ for (const Var& var : func->params) {
+ param_ids.push_back(VarDoc(d, var, !signature_symbols.count(var.get())));
+ }
size_t param_index = 0;
for (const Var& var : func->params) {
IdDoc lhs = param_ids[param_index++];
- ExprDoc annotation = d->Translate(var->ty).value();
- d->RecordOrigin(annotation, var->ty);
+ // Reuse the captured scalar in dependent annotations instead of creating
+ // a new parameter identity after an earlier annotation has read it.
+ bool shared_symbol = signature_symbols.count(var.get());
+ ExprDoc annotation =
+ shared_symbol ? d->Translate(var).value() :
d->Translate(var->ty).value();
+ d->RecordOrigin(annotation, shared_symbol ? ffi::ObjectRef(var) :
ffi::ObjectRef(var->ty));
AssignDoc argument(lhs, std::nullopt, annotation);
d->RecordOrigin(argument, var);
args.push_back(argument);
}
auto signature_candidates = CopyImplicitDefs(d);
+ for (const Var& var : func->params) signature_candidates.erase(var);
std::unordered_map<const ffi::Object*, size_t> thread_counts;
std::vector<tirx::IterVar> thread_vars;
ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
