-
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] Plumb external codegen target via Target.current() #11432
[Relay] Plumb external codegen target via Target.current() #11432
Conversation
Hi @Mousius I think you're probably the only other person who can review this one! |
c65b2c8
to
04085b0
Compare
9753881
to
90aca0d
Compare
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.
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.
c3fe805
to
70310d8
Compare
…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?
70310d8
to
673c5ce
Compare
@tvm-bot merge |
Cannot merge, did not find any approving reviews from users with write access on 673c5ce |
(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.