[REFACTOR][IR] Inline ReplaceGlobalVars into AttachGlobalSymbol#19625
Conversation
`ReplaceGlobalVars` was declared in `include/tvm/ir/replace_global_vars.h` and dispatched per-function-type via a NodeFunctor vtable populated by relax and tirx static-init blocks. The only in-tree C++ caller was `AttachGlobalSymbol` in `src/relax/transform/attach_global_symbol.cc`, and the only python user was an `IRModule.replace_global_vars` method exercised only by its own tests. Move the dispatch into `attach_global_symbol.cc` as a file-local helper. The helper branches directly on `tirx::PrimFuncNode` / `relax::FunctionNode` / `relax::ExternFuncNode` instead of going through a NodeFunctor vtable, since `attach_global_symbol.cc` already includes the relax + tirx headers. Delete the IR-layer driver, the per-dialect dispatch registration files, the public header, the python wrapper, and its tests. The remaining behavior is still covered by `tests/python/relax/test_transform_attach_global_symbol.py` and by every relax pipeline that runs `AttachGlobalSymbol`. Removes the runtime registration-based coupling from the IR layer to relax/tirx dialects for this specific transform.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request removes the generic ReplaceGlobalVars utility and its associated source, header, and test files, refactoring the logic into a file-local utility ReplaceGlobalVarsInModule inside src/relax/transform/attach_global_symbol.cc. The review feedback identifies opportunities to optimize the kGlobalSymbol update logic within this new utility by directly comparing new_gvar and old_gvar instead of performing a linear scan over the replacements map.
| // Update kGlobalSymbol if the function is externally exposed and being renamed. | ||
| if (auto opt = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) { | ||
| auto name = opt.value(); | ||
| for (const auto& [before, after] : replacements) { | ||
| if (before->name_hint == name) { | ||
| if (after->name_hint != name) { | ||
| func = WithAttr(func, tvm::attr::kGlobalSymbol, after->name_hint); | ||
| } | ||
| break; | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Since we are already iterating over mod->functions and have direct access to old_gvar and new_gvar, we do not need to perform a linear scan over replacements to find the matching GlobalVar. We can directly check if new_gvar has a different name than old_gvar and update the attribute accordingly. This simplifies the code and improves the lookup complexity from O(N) to O(1).
// Update kGlobalSymbol if the function is externally exposed and being renamed.
if (func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
if (new_gvar->name_hint != old_gvar->name_hint) {
func = WithAttr(func, tvm::attr::kGlobalSymbol, new_gvar->name_hint);
}
}| // Update kGlobalSymbol if the function is externally exposed and being renamed. | ||
| if (auto opt = new_relax_func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) { | ||
| auto name = opt.value(); | ||
| for (const auto& [before, after] : replacements) { | ||
| if (before->name_hint == name) { | ||
| if (after->name_hint != name) { | ||
| new_relax_func = WithAttr(new_relax_func, tvm::attr::kGlobalSymbol, after->name_hint); | ||
| } | ||
| break; | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Similarly, since we have direct access to old_gvar and new_gvar, we can avoid the linear scan over replacements here as well. We can directly check if new_gvar has a different name than old_gvar and update the attribute in O(1) time.
// Update kGlobalSymbol if the function is externally exposed and being renamed.
if (new_relax_func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
if (new_gvar->name_hint != old_gvar->name_hint) {
new_relax_func = WithAttr(new_relax_func, tvm::attr::kGlobalSymbol, new_gvar->name_hint);
}
}Since the loop top already computes new_gvar as replacements.Get(old_gvar).value_or(old_gvar), the kGlobalSymbol rename step can directly compare new_gvar->name_hint vs old_gvar->name_hint in O(1) instead of scanning replacements in O(N) on each iteration. Applied identically to the tirx PrimFunc branch and the relax Function branch — they were duplicated O(N) scans.
2 participants
Summary
ReplaceGlobalVarswas a public IR-layer API with only one in-tree C++caller (
relax::AttachGlobalSymbol). The mechanism used a NodeFunctorvtable populated at static-init time by per-dialect
.ccfiles inrelax and tirx, which made the IR layer logically depend on its
dialects even though the include graph did not show it.
Move the dispatch logic into the consumer as file-local mutators and
a private helper. Delete the public header, the IR-layer driver, both
per-dialect dispatch registrations, the
IRModule.replace_global_varspython method, and its dedicated test file. The behavior is still
covered by
tests/python/relax/test_transform_attach_global_symbol.pyand by the pipelines that include the
AttachGlobalSymbolpass.