[tvm] branch unity updated: [Unity] Remove unused local function definitions (#15507)

[kp...@apache.org]

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

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


The following commit(s) were added to refs/heads/unity by this push:
     new bcc9a68b3e [Unity] Remove unused local function definitions (#15507)
bcc9a68b3e is described below

commit bcc9a68b3ef1448049d6c3c9a4be745e8109d834
Author: Eric Lunderberg <Lu...@users.noreply.github.com>
AuthorDate: Thu Aug 10 08:36:19 2023 -0500

    [Unity] Remove unused local function definitions (#15507)
    
    Prior to this commit, the `relax.analysis.remove_all_unused` utility
    only removed variable bindings that occur within a dataflow block.
    This satisfies that the condition that a variable can be removed if is
    unused and if its computation has no side effects.  However, it
    doesn't include any values from a `BindingBlock`.
    
    This commit updates the `RemoveUnusedVars` pass to also inspect
    `BindingBlock` instances for removeable variables.  Currently, this is
    limited to unused local function definitions.
---
 src/relax/ir/binding_rewrite.cc     | 22 ++++++++++++++++-
 tests/python/relax/test_analysis.py | 47 +++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/src/relax/ir/binding_rewrite.cc b/src/relax/ir/binding_rewrite.cc
index 188ed39560..4eec0310fb 100644
--- a/src/relax/ir/binding_rewrite.cc
+++ b/src/relax/ir/binding_rewrite.cc
@@ -245,7 +245,27 @@ class RemoveUnusedVars : public ExprMutator {
   RemoveUnusedVars(Map<Var, Array<Var>> users, Array<Var> fn_outputs)
       : RemoveUnusedVars(GetUnusedVars(users, fn_outputs)) {}
 
-  BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) {
+  BindingBlock VisitBindingBlock_(const BindingBlockNode* block) override {
+    builder_->BeginBindingBlock();
+    for (Binding binding : block->bindings) {
+      bool can_remove = [&]() -> bool {
+        if (!unused_vars.count(binding->var)) {
+          return false;
+        }
+        auto var_binding = binding.as<VarBindingNode>();
+        if (!var_binding) {
+          return false;
+        }
+        return var_binding->value->IsInstance<FunctionNode>();
+      }();
+      if (!can_remove) {
+        VisitBinding(binding);
+      }
+    }
+    return builder_->EndBlock();
+  }
+
+  BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) override {
     auto prev_dfb = GetRef<DataflowBlock>(block);
     builder_->BeginDataflowBlock();
     for (Binding binding : block->bindings) {
diff --git a/tests/python/relax/test_analysis.py b/tests/python/relax/test_analysis.py
index 500a57775e..40bd5146ba 100644
--- a/tests/python/relax/test_analysis.py
+++ b/tests/python/relax/test_analysis.py
@@ -117,6 +117,53 @@ def test_binding_block_remove_all_unused():
     tvm.ir.assert_structural_equal(optimized, GroundTruth["main"])
 
 
+def test_binding_block_remove_all_unused_without_dataflow():
+    @tvm.script.ir_module
+    class IdentityUnused:
+        @R.function
+        def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+            lv0 = x
+            unused0 = R.call_dps_packed("my_sigmoid", (x,), R.Tensor((32, 32), dtype="float32"))
+            unused1 = R.call_dps_packed(
+                "my_dps_func", (unused0,), R.Tensor((32, 32), dtype="float32")
+            )
+            z = R.call_packed("vm.builtin.copy", lv0, sinfo_args=(R.Tensor((32, 32), "float32")))
+            return z
+
+    optimized = remove_all_unused(IdentityUnused["main"])
+
+    GroundTruth = IdentityUnused
+
+    tvm.ir.assert_structural_equal(optimized, GroundTruth["main"])
+
+
+def test_binding_block_remove_all_unused_func_without_dataflow():
+    @tvm.script.ir_module
+    class IdentityUnused:
+        @R.function
+        def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+            lv0 = x
+
+            @R.function
+            def internal_unused_func(A: R.Tensor((32, 32), "float32")) -> R.Tensor:
+                return A
+
+            z = R.call_packed("vm.builtin.copy", lv0, sinfo_args=(R.Tensor((32, 32), "float32")))
+            return z
+
+    optimized = remove_all_unused(IdentityUnused["main"])
+
+    @tvm.script.ir_module
+    class GroundTruth:
+        @R.function
+        def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+            lv0 = x
+            z = R.call_packed("vm.builtin.copy", lv0, sinfo_args=(R.Tensor((32, 32), "float32")))
+            return z
+
+    tvm.ir.assert_structural_equal(optimized, GroundTruth["main"])
+
+
 def test_binding_block_fake_unused_remove_all_unused():
     @tvm.script.ir_module
     class IdentityUnused: