gemini-code-assist[bot] commented on code in PR #19959:
URL: https://github.com/apache/tvm/pull/19959#discussion_r3529174159
##########
src/relax/ir/expr_functor.cc:
##########
@@ -544,19 +544,36 @@ PrimExpr ExprMutatorBase::VisitTypePrimExprField(const
PrimExpr& expr) {
// ==================
// ExprMutator
+namespace {
+
+template <typename TMap>
+void RedirectVarRemapTargets(TMap* var_remap, const Var& old_target, const
Var& new_target) {
+ if (old_target.same_as(new_target)) {
+ return;
+ }
+ for (auto& entry : *var_remap) {
+ if (entry.second.same_as(old_target)) {
+ entry.second = new_target;
+ }
+ }
+}
+
+} // namespace
Review Comment:
### Efficiency Issue: $O(N^2)$ Complexity in Variable Remapping\n\nThe
helper function "RedirectVarRemapTargets" performs a linear scan over the
entire "var_remap" map to update any entries pointing to "old_target". Since
this is called during "ReEmitBinding" and "VisitBinding_" (which are executed
for every binding in a block), this introduces an $O(N^2)$ complexity
bottleneck, where $N$ is the number of bindings. For large models with
thousands of bindings, this can cause significant compilation overhead.\n\n####
Suggested Alternative: Transitive Lookup\nInstead of eagerly updating all
targets via a linear scan, we can resolve the remappings transitively on
lookup. This completely eliminates the need for "RedirectVarRemapTargets" and
reduces the remapping overhead to $O(1)$ amortized time.\n\nWe can define a
helper method in "ExprMutator" (e.g., in "include/tvm/relax/expr_functor.h" or
as a protected method):\n\n```cpp\nVar ExprMutator::GetVarRemap(const Var& var)
const {\n Var
current = var;\n while (true) {\n auto it = var_remap_.find(current);\n
if (it == var_remap_.end() || it->second.same_as(current)) {\n break;\n
}\n current = it->second;\n }\n return current;\n}\n```\n\nAnd then use
"GetVarRemap" in "VisitExpr_", "LookupBinding", and other lookup sites instead
of direct "var_remap_.find(var)".
--
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]
