(tvm) branch main updated: [Fix] Remove redundant "remove_all_unused" in IPC memory lowering (#16771)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 134f8fd208 [Fix] Remove redundant "remove_all_unused" in IPC memory 
lowering (#16771)
134f8fd208 is described below

commit 134f8fd2081ea958bf23c6d66a2dd464a90fd85c
Author: Ruihang Lai <ruiha...@cs.cmu.edu>
AuthorDate: Sat Mar 23 10:02:58 2024 -0400

    [Fix] Remove redundant "remove_all_unused" in IPC memory lowering (#16771)
    
    This commit removes the redundant invocation of `remove_all_unused`
    function in the GPU IPC memory allocation lowering pass. This is
    because the pass only mutates one call at a time, and thus will
    not introduce new unused bindings.
---
 python/tvm/relax/transform/ipc_allreduce_rewrite.py       | 2 --
 python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py | 2 --
 2 files changed, 4 deletions(-)

diff --git a/python/tvm/relax/transform/ipc_allreduce_rewrite.py 
b/python/tvm/relax/transform/ipc_allreduce_rewrite.py
index 8dc535020b..3e7b005a60 100644
--- a/python/tvm/relax/transform/ipc_allreduce_rewrite.py
+++ b/python/tvm/relax/transform/ipc_allreduce_rewrite.py
@@ -23,7 +23,6 @@ from typing import Dict
 import tvm
 from tvm import relax
 from tvm.ir.module import IRModule
-from tvm.relax.analysis import remove_all_unused
 from tvm.relax.expr import Expr, Var
 from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, 
visitor
 
@@ -138,7 +137,6 @@ class _Rewriter(PyExprMutator):
         for g_var, func in self.mod.functions_items():
             if isinstance(func, relax.Function):
                 updated_func = self.visit_expr(func)
-                updated_func = remove_all_unused(updated_func)
                 self.builder_.update_func(g_var, updated_func)
         return self.builder_.get()
 
diff --git a/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py 
b/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
index 0967e00756..00081f92b1 100644
--- a/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
+++ b/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
@@ -21,7 +21,6 @@ The pass is written in Python for experiment, fast 
development.
 import tvm
 from tvm import relax
 from tvm.ir.module import IRModule
-from tvm.relax.analysis import remove_all_unused
 from tvm.relax.expr import Expr
 from tvm.relax.expr_functor import PyExprMutator, mutator
 
@@ -49,7 +48,6 @@ class _Rewriter(PyExprMutator):
         for g_var, func in self.mod.functions_items():
             if isinstance(func, relax.Function):
                 updated_func = self.visit_expr(func)
-                updated_func = remove_all_unused(updated_func)
                 self.builder_.update_func(g_var, updated_func)
         return self.builder_.get()