Skip to content
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

Merged
merged 22 commits into from
Aug 9, 2021

Conversation

mikepapadim
Copy link
Contributor

@mikepapadim
Copy link
Contributor Author

@jroesch can this merged?

@@ -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()) {
Copy link
Member

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 ?

Copy link
Contributor

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.

Copy link
Member

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?

Copy link
Contributor

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?

Copy link
Member

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 😸

Copy link
Contributor

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? :)

@YuchenJin
Copy link
Contributor

As a side note, this PR just naively replaces the old CompileEngine in the VM compiler with the lowering functions in the new TECompiler. This is the first step. Ultimately, we want to invoke the LowerTE pass once in TECompiler and hide those lowering functions.

@junrushao
Copy link
Member

Hey @mikepapadim @YuchenJin what's the status of this PR? Shall we move forward?

@YuchenJin
Copy link
Contributor

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. :)

Copy link
Member

@Mousius Mousius left a 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!

@junrushao junrushao merged commit 005f682 into apache:main Aug 9, 2021
@junrushao
Copy link
Member

Thank you guys for the effort!

mehrdadh pushed a commit to mehrdadh/tvm that referenced this pull request Aug 11, 2021
* [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>
ylc pushed a commit to ylc/tvm that referenced this pull request Sep 29, 2021
* [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>
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022
* [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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants