-
Notifications
You must be signed in to change notification settings - Fork 83
Canonicalize triton::ReduceOp select(cmpf) patterns to min/max operations #361
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
base: main
Are you sure you want to change the base?
Canonicalize triton::ReduceOp select(cmpf) patterns to min/max operations #361
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: |
0a1646b to
8211e44
Compare
@nhat-nguyen, thanks for the quick review. I tried to leverage the |
8211e44 to
69a7c39
Compare
include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp
Show resolved
Hide resolved
include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp
Outdated
Show resolved
Hide resolved
include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp
Outdated
Show resolved
Hide resolved
Address review comments to reuse code
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! |
include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp
Outdated
Show resolved
Hide resolved
include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp
Outdated
Show resolved
Hide resolved
|
@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. |
include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp
Show resolved
Hide resolved
include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp
Outdated
Show resolved
Hide resolved
include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp
Outdated
Show resolved
Hide resolved
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::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.