Canonicalize triton::ReduceOp select(cmpf) patterns to min/max operations#361
Canonicalize triton::ReduceOp select(cmpf) patterns to min/max operations#361quic-samanara wants to merge 5 commits intomicrosoft:mainfrom
Conversation
|
@microsoft-github-policy-service agree company="Qualcomm" |
|
@nhat-nguyen, sorry for the explicit ping but I did not have access to request a formal review. Could you be pleaase review the PR / or request review from appropriate stakeholders? |
|
@quic-samanara thanks for the pr! it's been a while since i last looked at this pass, but there's a somewhat similar pattern: |
quic-samanara
force-pushed
the
samanara/reduceop-minmax-folding
branch
from
0a1646b to
8211e44
Compare
@nhat-nguyen, thanks for the quick review. I tried to leverage the |
quic-samanara
force-pushed
the
samanara/reduceop-minmax-folding
branch
from
8211e44 to
69a7c39
Compare
Address review comments to reuse code
quic-samanara
force-pushed
the
samanara/reduceop-minmax-folding
branch
from
69a7c39 to
b8011b4
Compare
|
Hi @nhat-nguyen, addressed some of the comments you had:
Hope this addresses some of your concerns. |
|
Hi, pinging here to check if more changes are needed. Thanks! |
|
@quic-samanara looks good to me, just left a few minor comments. i don't have permission to run the workflow or approve anymore though. pinging @bmyerz0 @red1bluelost for help |
to output parameter.
point is restored after pattern rewrites.
|
@nhat-nguyen, thanks a lot for the review. Addressed the last two comments you pointed. |
|
Could you add a lit test for this. |
quic-samanara
force-pushed
the
samanara/reduceop-minmax-folding
branch
from
c19470b to
f1efef1
Compare
Thanks for reviewing @red1bluelost, addressed your comments and added a LIT test. |
|
No more PRs are being accepted on this repo anymore. Please see #367 |
@red1bluelost, oh, I see. Does that mean there is going to be no development from now on? Is someone else taking the ownership? |
This PR introduces a pattern rewrite that simplifies triton::ReduceOp bodies containing arith.select operations based on floating-point comparisons. The transformation replaces select-based logic with equivalent min/max operations, improving IR canonicalization and enabling further optimization. Supported Transformations: select(cmpf ogt a, b), a, b → arith.maxf(a, b) select(cmpf olt a, b), a, b → arith.minf(a, b) select((cmpf ogt a, b) || cmpf une a, a), a, b → arith.maximumf(a, b) select((cmpf olt a, b) || cmpf une a, a), a, b → arith.minimumf(a, b) Applying this transformation before the ReduceOp converter runs allows the natural lowering of tt.reduce to linalg.reduce; the motivational use case was from a torch-inductor spit out Triton kernel which consisted of this complex pattern in the reduction body and did not allow the ReduceOp to get lowered to the linalg dialect. PR cloned from microsoft/triton-shared#361
3 participants
This PR introduces a pattern rewrite that simplifies
triton::ReduceOpbodies containingarith.selectoperations based on floating-point comparisons. The transformation replaces select-based logic with equivalent min/max operations, improving IR canonicalization and enabling further optimization.Supported Transformations:
select(cmpf ogt a, b), a, b→arith.maxf(a, b)select(cmpf olt a, b), a, b→arith.minf(a, b)select((cmpf ogt a, b) || cmpf une a, a), a, b→arith.maximumf(a, b)select((cmpf olt a, b) || cmpf une a, a), a, b→arith.minimumf(a, b)Applying this transformation before the
ReduceOpconverter runs allows the natural lowering oftt.reducetolinalg.reduce; the motivational use case was from a torch-inductor spit out Triton kernel which consisted of this complex pattern in the reduction body and did not allow theReduceOpto get lowered to thelinalgdialect.