[Unity] Include LegalizeOps in the default relax.build lowering flow #15864
Conversation
Lunderberg
force-pushed
the
unity_default_legalize_ops_when_building
branch
from
16a7df0
to
b821d53
Compare
|
I'm not super sure if it's desirable. Usually after operator lowering, we will apply a series of transformations to reach the end state of the desirable TIR, which we could abstract with |
I think having LegalizeOps as part of I also agree with the idea of keeping |
|
@junrushao Absolutely agreed, there's a wide variety of hand-selected transformations that optimized models require, many of which must be performed after Workflow 1: A user writes a model, applies writes a sequence of transformations that have been hand-selected to optimize that specific model into the desired form, then passes the result into Workflow 2: A user writes a model, then passes the model directly into The goal of this PR is to enable Workflow 2 as a simpler way to get started, either for a new user or when running a new model, without impacting the optimizations performed in Workflow 1. |
|
Philosophically, I like the distinction @quic-sanirudh is drawing between the minimal |
Lunderberg
force-pushed
the
unity_default_legalize_ops_when_building
branch
from
9c0fa78
to
c9762bc
Compare
|
I rebased onto current |
|
I agree with the workflow 2 you've mentioned here, and since |
|
I don’t object either. Happy to merge it in |
|
Sounds good, and thank you! I'm rebasing the commit on top of unity head, as the current CI failures are due to a breakage on |
Prior to this commit, `relax.transform.LegalizeOps` needed to be called prior to `relax.build`. This commit adds `LegalizeOps` to the lowering flow, to simplify the calling steps for an end-user. If the `IRModule` contains no legalizable functions, a second legalization pass has no effect. Some test cases relied on this behavior as an implicit assertion that operator fusion patterns applied. That is, by omitting `LegalizeOps`, a successful compilation `relax.build` would only occur if all legalizable operators have already been removed, and so an incorrect fusion pattern would result in a failure to build the module. While these tests would be better expressed by comparing against an expected fused pattern, updating the tests is outside the scope of this PR. To allow these tests to keep their implicit assertions, a `"relax.transform.apply_legalize_ops"` config can be used to disable the `LegalizeOps` pass.
Lunderberg
force-pushed
the
unity_default_legalize_ops_when_building
branch
from
c9762bc
to
d78b932
Compare
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.
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.
Prior to this commit,
relax.transform.LegalizeOpsneeded to be called prior torelax.build. This commit addsLegalizeOpsto the lowering flow, to simplify the calling steps for an end-user. If theIRModulecontains no legalizable functions, a second legalization pass has no effect.Some test cases relied on this behavior as an implicit assertion that operator fusion patterns applied. That is, by omitting
LegalizeOps, a successful compilationrelax.buildwould only occur if all legalizable operators have already been removed, and so an incorrect fusion pattern would result in a failure to build the module. While these tests would be better expressed by comparing against an expected fused pattern, updating the tests is outside the scope of this PR. To allow these tests to keep their implicit assertions, a"relax.transform.apply_legalize_ops"config can be used to disable theLegalizeOpspass.