Skip to content

[VM] Move param bind to OptimizeModule #7451

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
masahi merged 3 commits into from
Feb 13, 2021
Merged

[VM] Move param bind to OptimizeModule #7451

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c903d47
[VM] Move param bind to OptimizeModule
masahi Feb 12, 2021
6f5b696
add test to verify the number of free vars after opt
masahi Feb 12, 2021
15385a4
remove const from OptimizeModule
masahi Feb 13, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions src/relay/backend/vm/compiler.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<FunctionNode>())
<< "VM compiler expects to compile relay::Function";
auto f = relay::backend::BindParamsByName(Downcast<Function>(base_func), params_);
auto gvar = mod->GetGlobalVar("main");
mod->Add(gvar, f);
}

exec_ = make_object<Executable>();
targets_ = targets;
target_host_ = target_host;
Expand Down Expand Up @@ -1005,8 +996,17 @@ 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");
ICHECK(base_func->IsInstance<FunctionNode>())
<< "VM compiler expects to compile relay::Function";
auto f = relay::backend::BindParamsByName(Downcast<Function>(base_func), params_);
auto gvar = mod->GetGlobalVar("main");
mod->Add(gvar, f);
}

Array<Pass> pass_seqs;
Array<runtime::String> entry_functions{"main"};
pass_seqs.push_back(transform::RemoveUnusedFunctions(entry_functions));
Expand Down
3 changes: 1 addition & 2 deletions src/relay/backend/vm/compiler.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions tests/python/relay/test_vm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down