Skip to content

Canonicalization of triton reduce with select to arith min/max - #21

Merged
nhat-nguyen merged 1 commit into
facebookincubator:mainfrom
mitekoth:mitekoth/select-to-max-nan-canonicalize
Apr 16, 2026
Merged

Canonicalization of triton reduce with select to arith min/max#21
nhat-nguyen merged 1 commit into
facebookincubator:mainfrom
mitekoth:mitekoth/select-to-max-nan-canonicalize

Conversation

@mitekoth

Copy link
Copy Markdown
Contributor

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Apr 16, 2026
@mitekoth

Copy link
Copy Markdown
Contributor Author

@nhat-nguyen This PR was already reviewed in the Microsoft repo. Let me know if any additional changes are needed so that this can be merged asap. Thanks!

@mitekoth

Copy link
Copy Markdown
Contributor Author

Thanks for the review. Please feel free to merge the PR 😄

@nhat-nguyen
nhat-nguyen merged commit 957b4d1 into facebookincubator:main Apr 16, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants