gemini-code-assist[bot] commented on code in PR #19625:
URL: https://github.com/apache/tvm/pull/19625#discussion_r3311552176
##########
src/relax/transform/attach_global_symbol.cc:
##########
@@ -24,15 +24,129 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/module.h>
-#include <tvm/ir/replace_global_vars.h>
+#include <tvm/relax/expr_functor.h>
#include <tvm/relax/struct_info.h>
#include <tvm/relax/transform.h>
#include <tvm/tirx/function.h>
+#include <tvm/tirx/stmt_functor.h>
+
+#include <vector>
namespace tvm {
namespace relax {
namespace transform {
+namespace {
+
+// File-local mutator: replace GlobalVar references inside a relax::Function.
+struct RelaxGvarMutator : ExprMutator {
+ ffi::Map<GlobalVar, GlobalVar> replacements;
+ explicit RelaxGvarMutator(ffi::Map<GlobalVar, GlobalVar> replacements)
+ : replacements(replacements) {}
+
+ using ExprMutator::VisitExpr_;
+ Expr VisitExpr_(const GlobalVarNode* node) override {
+ auto gvar = ffi::GetRef<GlobalVar>(node);
+ return replacements.Get(gvar).value_or(gvar);
+ }
+};
+
+// File-local mutator: replace GlobalVar references inside a tirx::PrimFunc.
+struct TirxGvarMutator : tirx::StmtExprMutator {
+ ffi::Map<GlobalVar, GlobalVar> replacements;
+ explicit TirxGvarMutator(ffi::Map<GlobalVar, GlobalVar> replacements)
+ : replacements(replacements) {}
+
+ PrimExpr VisitExpr_(const tirx::CallNode* node) override {
+ auto call = Downcast<tirx::Call>(tirx::StmtExprMutator::VisitExpr_(node));
+ if (auto old_gvar = call->op.as<GlobalVar>()) {
+ if (auto new_gvar = replacements.Get(old_gvar.value())) {
+ call.CopyOnWrite()->op = new_gvar.value();
+ }
+ }
+ return call;
+ }
+};
+
+// Replace GlobalVar references across all functions in the module.
+// Direct dispatch on function type — no NodeFunctor indirection needed
+// since this file already includes the relax + tirx headers.
+IRModule ReplaceGlobalVarsInModule(IRModule mod, ffi::Map<GlobalVar,
GlobalVar> replacements) {
+ if (replacements.empty()) {
+ return mod;
+ }
+
+ std::vector<GlobalVar> to_remove;
+ IRModule updates;
+
+ for (const auto& [old_gvar, old_func] : mod->functions) {
+ auto new_gvar = replacements.Get(old_gvar).value_or(old_gvar);
+ BaseFunc new_func;
+
+ if (auto* prim_func_node = old_func.as<tirx::PrimFuncNode>()) {
+ auto func = ffi::GetRef<tirx::PrimFunc>(prim_func_node);
+ TirxGvarMutator mutator(replacements);
+ auto new_body = mutator(func->body);
+ if (!new_body.same_as(func->body)) {
+ func.CopyOnWrite()->body = new_body;
+ }
+ // 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;
+ }
+ }
+ }
Review 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).
```c
// 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);
}
}
```
##########
src/relax/transform/attach_global_symbol.cc:
##########
@@ -24,15 +24,129 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/module.h>
-#include <tvm/ir/replace_global_vars.h>
+#include <tvm/relax/expr_functor.h>
#include <tvm/relax/struct_info.h>
#include <tvm/relax/transform.h>
#include <tvm/tirx/function.h>
+#include <tvm/tirx/stmt_functor.h>
+
+#include <vector>
namespace tvm {
namespace relax {
namespace transform {
+namespace {
+
+// File-local mutator: replace GlobalVar references inside a relax::Function.
+struct RelaxGvarMutator : ExprMutator {
+ ffi::Map<GlobalVar, GlobalVar> replacements;
+ explicit RelaxGvarMutator(ffi::Map<GlobalVar, GlobalVar> replacements)
+ : replacements(replacements) {}
+
+ using ExprMutator::VisitExpr_;
+ Expr VisitExpr_(const GlobalVarNode* node) override {
+ auto gvar = ffi::GetRef<GlobalVar>(node);
+ return replacements.Get(gvar).value_or(gvar);
+ }
+};
+
+// File-local mutator: replace GlobalVar references inside a tirx::PrimFunc.
+struct TirxGvarMutator : tirx::StmtExprMutator {
+ ffi::Map<GlobalVar, GlobalVar> replacements;
+ explicit TirxGvarMutator(ffi::Map<GlobalVar, GlobalVar> replacements)
+ : replacements(replacements) {}
+
+ PrimExpr VisitExpr_(const tirx::CallNode* node) override {
+ auto call = Downcast<tirx::Call>(tirx::StmtExprMutator::VisitExpr_(node));
+ if (auto old_gvar = call->op.as<GlobalVar>()) {
+ if (auto new_gvar = replacements.Get(old_gvar.value())) {
+ call.CopyOnWrite()->op = new_gvar.value();
+ }
+ }
+ return call;
+ }
+};
+
+// Replace GlobalVar references across all functions in the module.
+// Direct dispatch on function type — no NodeFunctor indirection needed
+// since this file already includes the relax + tirx headers.
+IRModule ReplaceGlobalVarsInModule(IRModule mod, ffi::Map<GlobalVar,
GlobalVar> replacements) {
+ if (replacements.empty()) {
+ return mod;
+ }
+
+ std::vector<GlobalVar> to_remove;
+ IRModule updates;
+
+ for (const auto& [old_gvar, old_func] : mod->functions) {
+ auto new_gvar = replacements.Get(old_gvar).value_or(old_gvar);
+ BaseFunc new_func;
+
+ if (auto* prim_func_node = old_func.as<tirx::PrimFuncNode>()) {
+ auto func = ffi::GetRef<tirx::PrimFunc>(prim_func_node);
+ TirxGvarMutator mutator(replacements);
+ auto new_body = mutator(func->body);
+ if (!new_body.same_as(func->body)) {
+ func.CopyOnWrite()->body = new_body;
+ }
+ // 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;
+ }
+ }
+ }
+ new_func = func;
+ } else if (auto* relax_func_node = old_func.as<FunctionNode>()) {
+ RelaxGvarMutator mutator(replacements);
+ auto new_relax_func =
+
Downcast<Function>(mutator(Downcast<Function>(ffi::GetRef<Function>(relax_func_node))));
+ // 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;
+ }
+ }
+ }
Review 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.
```c
// 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);
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
