[Relay] Replace compile engine with TE compiler in the VM #8501
Conversation
mikepapadim
force-pushed
the
vm_te_migration
branch
from
be95b0e
to
452815d
Compare
mikepapadim
requested review from
comaniac,
icemelon,
jroesch,
junrushao,
tqchen,
yzhliu and
zhiics
as code owners
mikepapadim
force-pushed
the
vm_te_migration
branch
from
8e3d7ff
to
aed5d3b
Compare
mikepapadim
requested review from
anijain2305,
jwfromm,
MarisaKirisame,
mbrookhart,
slyubomirsky,
vinx13 and
wweic
as code owners
…hen lowering in TEcompiler
YuchenJin
force-pushed
the
vm_te_migration
branch
from
097406a
to
a30cf6a
Compare
|
@jroesch can this merged? |
src/relay/backend/te_compiler.cc
Outdated
| @@ -349,6 +350,9 @@ class LowerTensorExpr : public ExprMutator { | |||
| Map<GlobalVar, tir::PrimFunc> prim_fns; | |||
|
|
|||
| for (auto prim_fn : ext_func->funcs->functions) { | |||
| if (prim_fn.second->GetAttr<String>(attr::kCompiler).defined()) { | |||
There was a problem hiding this comment.
Just querying, when is this untrue for an external function that's gone via Lower ?
There was a problem hiding this comment.
Hi @Mousius, After calling TECompilerImpl::Lower() on an external function, the TECompiler will just encapsulate it into CachedFunc and return, and the external function is not lowered. It will be lowered by the external codegen with the TECompilerNode::LowerExternalFunctions() later. I added a line ir_module->Add(global_var, key->source_func); to LowerInternal, so this is untrue when some function in the funcs of the CachedFuncNode is not lowered yet.
There was a problem hiding this comment.
Hi @YuchenJin, thanks for getting back to me, the reason I asked was because this loop is gated by the kCompiler attribute:
if (func->GetAttr<String>(attr::kCompiler).defined()) {
And then when it goes into LowerInternal it'll end up in this block (with your alterations):
if (key->source_func->GetAttr<String>(attr::kCompiler).defined()) {
auto ir_module = IRModule();
const auto name_node = key->source_func->GetAttr<String>(tvm::attr::kGlobalSymbol);
ICHECK(name_node.defined()) << "External function has not been attached a name yet.";
auto func_name = GetUniqueName(name_node.value(), &name_map_);
auto target = Target("ext_dev");
auto global_var = GlobalVar(func_name);
global_var->checked_type_ = key->source_func->checked_type();
ir_module->Add(global_var, key->source_func);
value->cached_func = CachedFunc(target, global_var, {}, {}, te::Schedule(), {}, ir_module);
return value;
}
Which means when we enter this loop, the only function in the IRModule should be a function taken from key->source_func which has the attr on it?
There was a problem hiding this comment.
Thanks @Mousius for elaborating it! I think you are right, and we can probably remove this loop because the IRModule should not contain a PrimFunc anyway for external functions, and we don't need to check it. What do you think?
There was a problem hiding this comment.
Let's try removing the loop, if tests still pass we should be confident it wasn't necessary 😸
There was a problem hiding this comment.
Yeah, tests still pass after removing it. @jroesch, could you review this PR? :)
|
As a side note, this PR just naively replaces the old |
|
Hey @mikepapadim @YuchenJin what's the status of this PR? Shall we move forward? |
Hey @junrushao1994, thanks for taking a look! @Mousius asked a good question and I think we can probably improve the code a little bit. Let's wait @Mousius for his reply. :) |
There was a problem hiding this comment.
@junrushao1994 this makes sense to me 😸 thanks for the update @YuchenJin!
|
Thank you guys for the effort! |
* [VM] Add imports to new TE in VM compiler * [VM] Add comments to compile engine usages * [VM] Replace depreceated CachedFunc of compile_engine with TE_compiler * [VM] rm compiler engine compiler.cc * [VM] Replace compile engine with TECompiler in memory allocator * [VM] Add relay interface to te_compiler * [Relay] Fix linting errors * Move TEcompiler to VMCompilerContext; add global func into IRmodule when lowering in TEcompiler * add back the check * skip the check for ext func in tecompiler * skip tvm::build for external functions * trigger ci * retrigger ci * retrigger ci * remove the unnecessary loop in tecompiler Co-authored-by: YuchenJin <yuchenj@cs.washington.edu>
* [VM] Add imports to new TE in VM compiler * [VM] Add comments to compile engine usages * [VM] Replace depreceated CachedFunc of compile_engine with TE_compiler * [VM] rm compiler engine compiler.cc * [VM] Replace compile engine with TECompiler in memory allocator * [VM] Add relay interface to te_compiler * [Relay] Fix linting errors * Move TEcompiler to VMCompilerContext; add global func into IRmodule when lowering in TEcompiler * add back the check * skip the check for ext func in tecompiler * skip tvm::build for external functions * trigger ci * retrigger ci * retrigger ci * remove the unnecessary loop in tecompiler Co-authored-by: YuchenJin <yuchenj@cs.washington.edu>
* [VM] Add imports to new TE in VM compiler * [VM] Add comments to compile engine usages * [VM] Replace depreceated CachedFunc of compile_engine with TE_compiler * [VM] rm compiler engine compiler.cc * [VM] Replace compile engine with TECompiler in memory allocator * [VM] Add relay interface to te_compiler * [Relay] Fix linting errors * Move TEcompiler to VMCompilerContext; add global func into IRmodule when lowering in TEcompiler * add back the check * skip the check for ext func in tecompiler * skip tvm::build for external functions * trigger ci * retrigger ci * retrigger ci * remove the unnecessary loop in tecompiler Co-authored-by: YuchenJin <yuchenj@cs.washington.edu>
This is WIP following the RFC in https://discuss.tvm.apache.org/t/rfc-relay-tecompiler-rewrite-existing-compile-engine-to-match-updated-compiler-flow/9233.
It replaces the Compile_engine with TE_compiler.
@jroesch @YuchenJin