-
Notifications
You must be signed in to change notification settings - Fork 3.5k
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.
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
[Relay] Replace compile engine with TE compiler in the VM #8501
Conversation
be95b0e
to
452815d
Compare
8e3d7ff
to
aed5d3b
Compare
…hen lowering in TEcompiler
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just querying, when is this untrue for an external function that's gone via Lower
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@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