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] Plumb external codegen target via Target.current() #11432

Merged
merged 3 commits into from
Jun 1, 2022

Conversation

mbs-octoml
Copy link
Contributor

@mbs-octoml mbs-octoml commented May 24, 2022

(See https://discuss.tvm.apache.org/t/byoc-supporting-cutlass-byoc-with-collage/12796/6 for
context, which in turn is part of Collage (https://github.com/apache/tvm-rfcs/blob/main/rfcs/0062-collage.md).

We want both old-style (via relay.ext.$toolchain) and new-style (via "RelayToTIR" Pass
attribute on target kind) external codegen to be able to access the current 'external codegen'
Target instance via Target.current().

  • For old-style, plumb the true Target through TEComplier and push it on the context
    stack before calling relay.ext.$toolchain.

  • For new-style, pass the CompilationConfig to the RelayToTIRTargetHook pass, make the jump from
    "Compiler" attribute value to Target via the new CompilationConfig::FindPrimitiveTargetForKind
    method, and push on the stack before invoking the custom "RelayToTIR" pass.

While working on this discovered RelayToTIRTargetHook was incompatible with the VM's compilation
flow since RelayToTIRTargetHook assumes all "Compiler" attributed functions are inlined. Generalize
it to support both inline and global function styles.

Extend Target::IsExternalCodegen to recognize target kinds with "RelayToTIR" attributes as
external.

Update target hooks unit test to exercise new support for outline-style, picking up the current target,
and compiling via the VM. Also update the example pass to illustrate some of the common gotcha's
when implementing a custom lowering-like pass.

@mbs-octoml mbs-octoml changed the title [Relay] Plumb external codegen target via Target.current() for all ex… [Relay] Plumb external codegen target via Target.current() May 24, 2022
@mbs-octoml
Copy link
Contributor Author

mbs-octoml commented May 24, 2022

Hi @Mousius I think you're probably the only other person who can review this one!

@mbs-octoml mbs-octoml force-pushed the mbs-collage-current-target branch from c65b2c8 to 04085b0 Compare May 24, 2022 18:58
@jwfromm jwfromm requested a review from Mousius May 25, 2022 17:37
@mbs-octoml mbs-octoml force-pushed the mbs-collage-current-target branch from 9753881 to 90aca0d Compare May 26, 2022 19:04
Copy link
Contributor

@jwfromm jwfromm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is very cool work that seems well designed to me. The only thing I'd like to see is a little documentation that points out all the pieces needed to make it work. For now, just having something brief in the python tests might be enough to get interested readers started.

tests/python/relay/test_target_hooks.py Show resolved Hide resolved
@mbs-octoml mbs-octoml force-pushed the mbs-collage-current-target branch from c3fe805 to 70310d8 Compare May 27, 2022 17:49
…ternal codegen paths

(See https://discuss.tvm.apache.org/t/byoc-supporting-cutlass-byoc-with-collage/12796/6 for
context, which in turn is part of Collage (https://github.com/apache/tvm-rfcs/blob/main/rfcs/0062-collage.md).

We want both old-style (via relay.ext.$toolchain) and new-style (via "RelayToTIR" Pass
attribute on target kind) external codegen to be able to access the current 'external codegen'
Target instance via Target.current().

 - For old-style, plumb the true Target through TEComplier and push it on the context
   stack before calling relay.ext.$toolchain.

 - For new-style, pass the CompilationConfig to the RelayToTIRTargetHook pass, make the jump from
   "Compiler" attribute value to Target via the new CompilationConfig::FindPrimitiveTargetForKind
   method, and push on the stack before invoking the custom "RelayToTIR" pass.

While working on this discovered RelayToTIRTargetHook was incompatible with the VM's compilation
flow since RelayToTIRTargetHook assumes all "Compiler" attributed functions are inlined. Generalize
it to support both inline and global function styles.

Extend Target::IsExternalCodegen to recognize target kinds with "RelayToTIR" attributes as
external.

Update target hooks unit test to exercise new support for outline-style, picking up the current target,
and compiling via the VM.
Can't repro tests/python/contrib/test_ethosu/cascader/test_scheduler.py::test_compute_cycles_annotation failure, flake?
@mbs-octoml mbs-octoml force-pushed the mbs-collage-current-target branch from 70310d8 to 673c5ce Compare June 1, 2022 16:31
@mbs-octoml
Copy link
Contributor Author

@tvm-bot merge

@github-actions
Copy link
Contributor

github-actions bot commented Jun 1, 2022

Cannot merge, did not find any approving reviews from users with write access on 673c5ce

@AndrewZhaoLuo AndrewZhaoLuo merged commit 89c0235 into apache:main Jun 1, 2022
@mbs-octoml mbs-octoml deleted the mbs-collage-current-target branch June 1, 2022 19:56
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.

3 participants