Skip to content

Conversation

@quic-samanara
Copy link

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, barith.maxf(a, b)
  • select(cmpf olt a, b), a, barith.minf(a, b)
  • select((cmpf ogt a, b) || cmpf une a, a), a, barith.maximumf(a, b)
  • select((cmpf olt a, b) || cmpf une a, a), a, barith.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.

@quic-samanara
Copy link
Author

@microsoft-github-policy-service agree company="Qualcomm"

@quic-samanara quic-samanara marked this pull request as ready for review October 31, 2025 16:07
@quic-samanara
Copy link
Author

quic-samanara commented Nov 11, 2025

@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?
Thanks!

@nhat-nguyen
Copy link
Contributor

@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: MinMaxConverter in ConversionPatterns.hpp. could we merge these somehow?

@quic-samanara quic-samanara force-pushed the samanara/reduceop-minmax-folding branch from 0a1646b to 8211e44 Compare November 11, 2025 19:55
@quic-samanara
Copy link
Author

@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: MinMaxConverter in ConversionPatterns.hpp. could we merge these somehow?

@nhat-nguyen, thanks for the quick review. I tried to leverage the MinMaxConverter infrastructure and perform the rewrite under the required constraints there. Let me know how that looks. Thanks!

@quic-samanara quic-samanara force-pushed the samanara/reduceop-minmax-folding branch from 8211e44 to 69a7c39 Compare November 11, 2025 20:54
Address review comments to reuse code
@quic-samanara quic-samanara force-pushed the samanara/reduceop-minmax-folding branch from 69a7c39 to b8011b4 Compare November 18, 2025 22:43
@quic-samanara
Copy link
Author

Hi @nhat-nguyen, addressed some of the comments you had:

  • Introduced a foldCmpToMinMax helper that centralizes the mapping from arith.cmpf predicates to the appropriate Min/Max operations, and reuse it from the generic cmp+select pattern and the NaN-aware folding logic.
  • Make the cmp->OrIOp->select folding logic operate directly on arith.select relaxing the triton::ReduceOp restriction.
  • Simplify the control flow in matchAndRewrite by
    • The direct cmp->select case is handled "inline".
    • Restricting redundant checking to fire the complex NaN-aware fold pattern.

Hope this addresses some of your concerns.

@quic-samanara
Copy link
Author

Hi, pinging here to check if more changes are needed. Thanks!

@nhat-nguyen
Copy link
Contributor

@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

@quic-samanara
Copy link
Author

@nhat-nguyen, thanks a lot for the review. Addressed the last two comments you pointed.

@red1bluelost
Copy link
Contributor

Could you add a lit test for this.

@quic-samanara quic-samanara force-pushed the samanara/reduceop-minmax-folding branch from c19470b to f1efef1 Compare December 3, 2025 02:06
@quic-samanara
Copy link
Author

Could you add a lit test for this.

Thanks for reviewing @red1bluelost, addressed your comments and added a LIT test.

@red1bluelost
Copy link
Contributor

No more PRs are being accepted on this repo anymore. Please see #367

@quic-samanara
Copy link
Author

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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants