From 8901cfc04a5960e61638c99e62b3f1853ef25e76 Mon Sep 17 00:00:00 2001 From: zxy844288792 Date: Fri, 14 May 2021 04:06:49 +0000 Subject: [PATCH 1/2] add removeUnusedFunctions pass in vm memoryopt --- src/relay/backend/vm/compiler.cc | 3 +++ tests/python/relay/test_vm.py | 24 ++++++++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc index 832cc0ee3891..ad23e132c486 100644 --- a/src/relay/backend/vm/compiler.cc +++ b/src/relay/backend/vm/compiler.cc @@ -978,6 +978,9 @@ void VMCompiler::Lower(IRModule mod, const TargetsMap& targets, const tvm::Targe transform::Sequential MemoryOpt(tvm::Target host_target, TargetsMap targets) { Array pass_seqs; + // Remove unused functions + Array entry_functions{"main"}; + pass_seqs.push_back(transform::RemoveUnusedFunctions(entry_functions)); // Manifest the allocations. pass_seqs.push_back(transform::ManifestAlloc(host_target, targets)); diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py index 8f518695e6b4..f1c9326f759f 100644 --- a/tests/python/relay/test_vm.py +++ b/tests/python/relay/test_vm.py @@ -29,6 +29,7 @@ from tvm.contrib import utils from tvm import rpc import tvm.testing +from tvm.relay.transform import InferType def check_result(args, expected_result, mod=None): @@ -185,6 +186,29 @@ def test_multiple_ifs(): res = vmobj_to_list(vm.evaluate()(False)) assert res == [1, 0] +@tvm.testing.uses_gpu +def test_unused_function(): + cond = relay.const(True) + mod = tvm.IRModule() + then_name = relay.GlobalVar("times_2") + # define unused function + else_name = relay.GlobalVar("times_3") + t1 = relay.TensorType((2,2), dtype="float32") + x1 = relay.var("x1", t1, dtype="float32") + x2 = relay.var("x2", t1, dtype="float32") + f2 = relay.multiply(x1, relay.const(2.0)) + f3 = relay.multiply(x2, relay.const(3.0)) + mod[then_name] = relay.Function([x1], f2) + mod[else_name] = relay.Function([x2], f3) + mod = InferType()(mod) + x3 = relay.var("x3", t1, dtype="float32") + # put unused function in else branch + f = relay.If(cond, then_name(x3), else_name(x3)) + mod["main"] = relay.Function([x3], f) + x_data = np.random.rand(2, 2).astype("float32") + y_data = x_data * 2 + + check_result([x_data], y_data, mod=mod) @tvm.testing.uses_gpu def test_simple_call(): From 302a50905d9386f7aa3fca8d5d78ca1e8ad56c58 Mon Sep 17 00:00:00 2001 From: zxy844288792 Date: Fri, 14 May 2021 04:17:16 +0000 Subject: [PATCH 2/2] fix lint --- tests/python/relay/test_vm.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py index f1c9326f759f..9f861a2e7b54 100644 --- a/tests/python/relay/test_vm.py +++ b/tests/python/relay/test_vm.py @@ -186,6 +186,7 @@ def test_multiple_ifs(): res = vmobj_to_list(vm.evaluate()(False)) assert res == [1, 0] + @tvm.testing.uses_gpu def test_unused_function(): cond = relay.const(True) @@ -193,7 +194,7 @@ def test_unused_function(): then_name = relay.GlobalVar("times_2") # define unused function else_name = relay.GlobalVar("times_3") - t1 = relay.TensorType((2,2), dtype="float32") + t1 = relay.TensorType((2, 2), dtype="float32") x1 = relay.var("x1", t1, dtype="float32") x2 = relay.var("x2", t1, dtype="float32") f2 = relay.multiply(x1, relay.const(2.0)) @@ -202,7 +203,7 @@ def test_unused_function(): mod[else_name] = relay.Function([x2], f3) mod = InferType()(mod) x3 = relay.var("x3", t1, dtype="float32") - # put unused function in else branch + # put unused function in else branch f = relay.If(cond, then_name(x3), else_name(x3)) mod["main"] = relay.Function([x3], f) x_data = np.random.rand(2, 2).astype("float32") @@ -210,6 +211,7 @@ def test_unused_function(): check_result([x_data], y_data, mod=mod) + @tvm.testing.uses_gpu def test_simple_call(): mod = tvm.IRModule({})