From c903d476127d62eae6a73b4923cb862cf04fb775 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 12 Feb 2021 13:47:01 +0900 Subject: [PATCH 1/3] [VM] Move param bind to OptimizeModule --- src/relay/backend/vm/compiler.cc | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc index 7861502965a8..49af95afff0f 100644 --- a/src/relay/backend/vm/compiler.cc +++ b/src/relay/backend/vm/compiler.cc @@ -892,15 +892,6 @@ void VMCompiler::SetParam(const std::string& name, runtime::NDArray data_in) { } void VMCompiler::Lower(IRModule mod, const TargetsMap& targets, const tvm::Target& target_host) { - if (params_.size()) { - BaseFunc base_func = mod->Lookup("main"); - ICHECK(base_func->IsInstance()) - << "VM compiler expects to compile relay::Function"; - auto f = relay::backend::BindParamsByName(Downcast(base_func), params_); - auto gvar = mod->GetGlobalVar("main"); - mod->Add(gvar, f); - } - exec_ = make_object(); targets_ = targets; target_host_ = target_host; @@ -1007,6 +998,15 @@ transform::Sequential MemoryOpt(tvm::Target host_target, TargetsMap targets) { IRModule VMCompiler::OptimizeModule(const IRModule& mod, const TargetsMap& targets, const Target& target_host) { + if (params_.size()) { + BaseFunc base_func = mod->Lookup("main"); + ICHECK(base_func->IsInstance()) + << "VM compiler expects to compile relay::Function"; + auto f = relay::backend::BindParamsByName(Downcast(base_func), params_); + auto gvar = mod->GetGlobalVar("main"); + mod->Add(gvar, f); + } + Array pass_seqs; Array entry_functions{"main"}; pass_seqs.push_back(transform::RemoveUnusedFunctions(entry_functions)); From 6f5b6965006837104a6bcf12a4c1c57cd42f4f43 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Sat, 13 Feb 2021 06:06:26 +0900 Subject: [PATCH 2/3] add test to verify the number of free vars after opt --- tests/python/relay/test_vm.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py index 6958010176e3..975070ad1aaa 100644 --- a/tests/python/relay/test_vm.py +++ b/tests/python/relay/test_vm.py @@ -678,6 +678,10 @@ def test_vm_optimize(): comp = relay.vm.VMCompiler() opt_mod, _ = comp.optimize(mod, target="llvm", params=params) + free_vars = relay.analysis.free_vars(opt_mod["main"].body) + # Paremeters should all be bound, so the only free var is data + assert len(free_vars) == 1 + @tvm.testing.uses_gpu def test_loop_free_var(): From 15385a45a703b207ef01459228ecd23dc2700e38 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Sat, 13 Feb 2021 16:20:34 +0900 Subject: [PATCH 3/3] remove const from OptimizeModule --- src/relay/backend/vm/compiler.cc | 2 +- src/relay/backend/vm/compiler.h | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc index 49af95afff0f..7697b59437f0 100644 --- a/src/relay/backend/vm/compiler.cc +++ b/src/relay/backend/vm/compiler.cc @@ -996,7 +996,7 @@ transform::Sequential MemoryOpt(tvm::Target host_target, TargetsMap targets) { return transform::Sequential(pass_seqs); } -IRModule VMCompiler::OptimizeModule(const IRModule& mod, const TargetsMap& targets, +IRModule VMCompiler::OptimizeModule(IRModule mod, const TargetsMap& targets, const Target& target_host) { if (params_.size()) { BaseFunc base_func = mod->Lookup("main"); diff --git a/src/relay/backend/vm/compiler.h b/src/relay/backend/vm/compiler.h index 56965c544701..615a8181b387 100644 --- a/src/relay/backend/vm/compiler.h +++ b/src/relay/backend/vm/compiler.h @@ -125,8 +125,7 @@ class VMCompiler : public runtime::ModuleNode { * * \return The optimized IRModule. */ - IRModule OptimizeModule(const IRModule& mod, const TargetsMap& targets, - const Target& target_host); + IRModule OptimizeModule(IRModule mod, const TargetsMap& targets, const Target& target_host); /*! * \brief Populate the global function names in a map where the value is used