[Relax] Apply DefaultGPUSchedule() in default build pipeline #17108
+27
−3
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This is a follow-up to apache#15864, which added `LegalizeOps` to the default Relax build pipeline. Since legalization may produce additional TIR PrimFuncs that require scheduling, the output of `LegalizeOps` typically must also be passed through `tir.transform.DefaultGPUSchedule()`. This PR adds `DefaultGPUSchedule()` to the relax build pipeline to handle these cases. Scheduled PrimFunc have the `"tir.is_scheduled"` attribute set to true, and are ignored by `DefaultGPUSchedule()`. In addition, the `DefaultGPUSchedule` transform has no effect on non-GPU targets. Therefore, this change should only impact `tvm.relax.build` calls that previously resulted in an error due to unscheduled GPU functions, and should not have any impact on existing calls.
|
Seems we discussed this before in the phase ordering discussion in open dev meeting, can't remember the reason why we didn't add this pass in default pipeline, need to take a look at the meeting note |
|
I don't remember that discussion, but I'd be interested in seeing the minutes from it. The main roadblocks to it are ones that I've resolved in previous PRs:
With those fixes applied, I think it's now a pretty big usability benefit to include it in the lowering pipeline. |
|
LGTM - I think it's clear usability gain without significant extra burden on build pipeline. |
MadFunMaker
approved these changes
Aug 2, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This is a follow-up to #15864, which added
LegalizeOpsto the default Relax build pipeline. Since legalization may produce additional TIR PrimFuncs that require scheduling, the output ofLegalizeOpstypically must also be passed throughtir.transform.DefaultGPUSchedule(). This PR addsDefaultGPUSchedule()to the relax build pipeline to handle these cases.Scheduled PrimFunc have the
"tir.is_scheduled"attribute set to true, and are ignored byDefaultGPUSchedule(). In addition, theDefaultGPUScheduletransform has no effect on non-GPU targets. Therefore, this change should only impacttvm.relax.buildcalls that previously resulted in an error due to unscheduled GPU functions, and should not have any impact on existing calls.