Skip to content

[Triton] Add canonicalization patterns for tt.reduce #2591

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.

Sign up for GitHub

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

Jump to bottom
Closed
wants to merge 4 commits into from
Closed

[Triton] Add canonicalization patterns for tt.reduce #2591

wants to merge 4 commits into from

Conversation

@victor-eds
Copy link
Contributor

@victor-eds victor-eds commented Oct 30, 2024 •
edited
Loading

Add canonicalization pattern replacing tt.reduce with axis of size 1 with a tt.reshape operation.

This may enable further optimizations and simplifications.

@victor-eds
[Triton] Add canonicalization patterns for tt.reduce
5f64bb7 
Add canonicalization pattern replacing `tt.reduce` with axis of size 1 with a `tt.reshape` operation.

This may leverage further optimizations and simplifications.

Signed-off-by: victor-eds <[email protected]>
@victor-eds victor-eds requested review from a team, chengjunlu and etiotto October 30, 2024 12:36
@victor-eds victor-eds self-assigned this Oct 30, 2024
@victor-eds
Copy link
Contributor Author

victor-eds commented Oct 30, 2024

-tritonintelgpu-optimize-reduction-locality may generate this kind of operations. Running canonicalization right after it simplifies code.

sommerlukas
sommerlukas reviewed Oct 30, 2024
View reviewed changes
Copy link
Contributor

@sommerlukas sommerlukas left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would this also be a candidate for upstreaming?

@victor-eds
Copy link
Contributor Author

victor-eds commented Oct 30, 2024

This enables further simplification of code generated with #2266

@victor-eds
Copy link
Contributor Author

victor-eds commented Oct 30, 2024

Would this also be a candidate for upstreaming?

Yes, I added the label for that

@etiotto
Copy link
Contributor

etiotto commented Oct 30, 2024

Would this also be a candidate for upstreaming?

Yes, I added the label for that

Question: do you plan to upstream this first or shall we land it downstream first and do that separately?

@victor-eds
Copy link
Contributor Author

victor-eds commented Oct 30, 2024

Would this also be a candidate for upstreaming?

Yes, I added the label for that

Question: do you plan to upstream this first or shall we land it downstream first and do that separately?

Upstream and get the changes when rebasing on upstream. No rush for this change.

@vlad-penkin vlad-penkin linked an issue Oct 30, 2024 that may be closed by this pull request
Port "sub-group transpose reduction" to default path #2266
Closed
etiotto
etiotto reviewed Oct 30, 2024
View reviewed changes
Copy link
Contributor

@etiotto etiotto left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@victor-eds victor-eds marked this pull request as draft October 31, 2024 08:45
@victor-eds
Copy link
Contributor Author

victor-eds commented Oct 31, 2024

Upstream PR: triton-lang/triton#5024

@victor-eds
Copy link
Contributor Author

victor-eds commented Nov 4, 2024

Closing as won't fix for now as it isn't that high priority and there are some issues with the reshape instruction itself that make this complicated.

@victor-eds victor-eds closed this Nov 4, 2024
@victor-eds victor-eds deleted the reduce-canonicalize branch November 4, 2024 09:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@etiotto etiotto etiotto left review comments

@chengjunlu chengjunlu Awaiting requested review from chengjunlu

+1 more reviewer

@sommerlukas sommerlukas sommerlukas approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

@victor-eds victor-eds

Labels

codegen: attention upstream: triton

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Port "sub-group transpose reduction" to default path

3 participants

@victor-eds @etiotto @sommerlukas