Skip to content
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

[AMD] Restructure ReorderInstructions pass #4998

Merged
merged 3 commits into from
Oct 28, 2024
Merged

[AMD] Restructure ReorderInstructions pass #4998

merged 3 commits into from
Oct 28, 2024

Conversation

antiagainst
Copy link
Collaborator

@antiagainst antiagainst commented Oct 26, 2024 •
edited
Loading

This commit restructures the ReorderInstructions pass to be
more modular and cleaner with utility functions and guard
rule applications against their intended usage rather than
always being globally on. Along the way, dropped some
unnecessary local_load sink given it's later hoisted again.

Fixes https://github.com/ROCm/triton-internal/issues/280

@antiagainst
[AMD] NFC: Restructure ReorderInstructions pass
92ad7b5 
This commit restructures the ReorderInstructions pass to be
more modular and cleaner with utility functions and guard
rule applications against their intended usage rather than
always being globally on.
@antiagainst
Drop local_load sink logic given it's reverted in a next step
3255955
@antiagainst
Add guard for logic that appliable to matmul
c6a6a7f
@antiagainst antiagainst marked this pull request as ready for review October 27, 2024 00:22
@antiagainst
Copy link
Collaborator Author

@zhanglx13: first commit is basically NFC with some IR mutation robustness improvements. the second and third commit is meaningful changes. i'd suggest to view the diff with "Hide whitespace" toggle on.

sjw36
sjw36 reviewed Oct 28, 2024
View reviewed changes
Copy link
Contributor

@sjw36 sjw36 left a comment

Choose a reason for hiding this comment

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

LGTM!
The comment below doesn't change functionality, and will be deprecated by sched-v3 anyway.

third_party/amd/lib/TritonAMDGPUTransforms/ReorderInstructions.cpp
static bool isPureMatmulProblem(ModuleOp moduleOp) {
for (auto forOp : moduleOp.getOps<scf::ForOp>()) {
int counter = 0;
forOp.walk([&counter](triton::DotOp dotOp) { ++counter; });
Copy link
Contributor

Choose a reason for hiding this comment

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

This would trigger for parent of nested loops also..

@antiagainst antiagainst merged commit 86a2ac7 into triton-lang:main Oct 28, 2024
7 checks passed
@antiagainst antiagainst deleted the amd-reorder-pass branch October 28, 2024 23:31
AlexAUT pushed a commit to AlexAUT/triton that referenced this pull request Oct 29, 2024
@antiagainst @AlexAUT
[AMD] Restructure ReorderInstructions pass (triton-lang#4998)
dffb97d
Luosuu pushed a commit to Luosuu/triton that referenced this pull request Nov 13, 2024
@antiagainst @Luosuu
[AMD] Restructure ReorderInstructions pass (triton-lang#4998)
e9655c2
guacamoleo pushed a commit to guacamoleo/triton that referenced this pull request Nov 14, 2024
@antiagainst @guacamoleo
[AMD] Restructure ReorderInstructions pass (triton-lang#4998)
bd74bd7
jataylo pushed a commit to jataylo/triton that referenced this pull request Nov 18, 2024
@antiagainst @jataylo
[AMD] Restructure ReorderInstructions pass (triton-lang#4998)
bbd72b7 
(cherry picked from commit 86a2ac7)
@jataylo jataylo mentioned this pull request Nov 19, 2024
Merged
antiagainst added a commit that referenced this pull request Dec 4, 2024
@jataylo @joviliast
@antiagainst @zhanglx13 @Jokeren @binarman @knwng @jungpark-mlir @peterbell10 @htyu
[AMD] release/3.2.x AMD perf cherry picks (#5191)
2d8093c 
Cherry pick list:
- #4925
- #5053 
- #5019 
- #5002 
- #4935 - required additional cherry picks #4991 and #4951
- #4998 
- #4925 
- #5281 
- #5308 
- All previous LLVM hash PRs before #5308

---------

Co-authored-by: Ilya V <[email protected]>
Co-authored-by: Lei Zhang <[email protected]>
Co-authored-by: Lixun Zhang <[email protected]>
Co-authored-by: Keren Zhou <[email protected]>
Co-authored-by: Alexander Efimov <[email protected]>
Co-authored-by: Kyle Wang <[email protected]>
Co-authored-by: Jungwook Park <[email protected]>
Co-authored-by: peterbell10 <[email protected]>
Co-authored-by: Hongtao Yu <[email protected]>
jataylo pushed a commit to jataylo/triton that referenced this pull request Dec 5, 2024
@antiagainst @jataylo
[AMD] Restructure ReorderInstructions pass (triton-lang#4998)
218d8b0 
(cherry picked from commit 86a2ac7)
@jataylo jataylo mentioned this pull request Dec 5, 2024
Merged
antiagainst added a commit that referenced this pull request Dec 5, 2024
@jataylo @jungpark-mlir
@peterbell10 @htyu @antiagainst @joviliast @knwng
[AMD] rc/3.2.x cherry picks (#5347)
7e401df 
Reverts #5191 due to some mlir errors in pytorch unit tests

Smaller set of cherry picks:
- #5308 (and previous LLVM upgrades)
- #5281 
- #4925 
- #5053 
- #5019 
- #4998

---------

Co-authored-by: Jungwook Park <[email protected]>
Co-authored-by: peterbell10 <[email protected]>
Co-authored-by: Hongtao Yu <[email protected]>
Co-authored-by: Lei Zhang <[email protected]>
Co-authored-by: Ilya V <[email protected]>
Co-authored-by: Kyle Wang <[email protected]>
jataylo pushed a commit to jataylo/triton that referenced this pull request Dec 11, 2024
@antiagainst @jataylo
[AMD] Restructure ReorderInstructions pass (triton-lang#4998)
0c35781 
(cherry picked from commit 86a2ac7)
@jataylo jataylo mentioned this pull request Dec 12, 2024
Closed
jataylo pushed a commit to jataylo/triton that referenced this pull request Dec 13, 2024
@antiagainst @jataylo
[AMD] Restructure ReorderInstructions pass (triton-lang#4998)
4c7d56e 
(cherry picked from commit 86a2ac7)
@jataylo jataylo mentioned this pull request Dec 13, 2024
Merged
jataylo added a commit to ROCm/triton that referenced this pull request Dec 13, 2024
@jataylo @joviliast
@antiagainst @knwng @zhanglx13
[CP] AMD Performance cherry picks (#682)
deee2b1 
* [AMD] Emit vectorized 16-bit float LLVM atomic ops (triton-lang#4925)

In the case of 16 bit floats operands for tt::AtomicRMWOp, construct
only one LLVM::AtomicRMWOp but use vector of elements.
Such approach allows to generate packed intrinsics and process 2
elements at once.
Added a lit test for f16 vectorized case.

(cherry picked from commit 78c8054)

* [AMD] Restructure ReorderInstructions pass (triton-lang#4998)

(cherry picked from commit 86a2ac7)

* [AMD] Support warp-level reduction with DPP (triton-lang#5019)

This commit adds support for warp-level reduction
with DPP instructions, which can improve performance.

See https://gpuopen.com/learn/amd-gcn-assembly-cross-lane-operations/

(cherry picked from commit 21119e3)

* [AMD] Add missing dependency to TritonAMDGPUIR (triton-lang#5053)

TritonAMDGPUTransforms now depends on it.

(cherry picked from commit 0b443ce)

* [AMD] Support warp-level reduction with DPP (triton-lang#5019)

This commit adds support for warp-level reduction
with DPP instructions, which can improve performance.

See https://gpuopen.com/learn/amd-gcn-assembly-cross-lane-operations/

(cherry picked from commit 21119e3)

* [AMD] Use DPP to accelerate 16-bit floats (triton-lang#5072)

In the case of unpaired f16 elements utilize DPP instructions to
accelerate atomics. Here is an algorithm of lowering
`tt::atomicRmwOp(%ptr, %val, %mask)`:

0. Group thread by pairs. Master thread is (tid % 2 == 0);
1. All the threads send `%val` to `(tid - 1)` thread via `dppUpdateOp
shl`, so all the masters recieve value from secondary threads;
2. Take into account parity in the `%mask` value, build CF structures
according to it;
3. Generate `llvm::atomicRmwOp` in the threads enabled by `%mask` value;
4. All the threads send result of generated operation to `(tid + 1)`
thread via `dppUpdateOp shl`, so all secondary thread also recieve their
result.

DPP approach has ~5% perf improvment so use this one in the
case target arch supports DPP.

Signed-off-by: Ilya Veselov <[email protected]>
(cherry picked from commit bab3470)

* [AMD] Reland sinking the 2nd tt.load after local_load's (triton-lang#4935)

This PR adds more restrictions about when should we apply
the sched-load optimizations and un-revert
triton-lang#4823.

We will only apply the optimization when all of the following is
satisfied:
1. pureMatmulProblem, i.e. 1 `tt.dot` in the main loop
2. two `tt.load`s in the main loop
3. 2nd `tt.load` is ahead of the `tt.dot`
4. 1st user of 2nd `tt.load` is after the `tt.dot`
5. tile size is large enough, i.e. nonKDim >= 128 and kDim >= 64

(cherry picked from commit 4f6f768)

---------

Co-authored-by: Ilya V <[email protected]>
Co-authored-by: Lei Zhang <[email protected]>
Co-authored-by: Kyle Wang <[email protected]>
Co-authored-by: Lixun Zhang <[email protected]>
jataylo added a commit to jataylo/triton that referenced this pull request Dec 18, 2024
@jataylo
[AMD] release/3.2.x AMD perf cherry picks (triton-lang#5191)
93f31aa 
Cherry pick list:
- triton-lang#4925
- triton-lang#5053
- triton-lang#5019
- triton-lang#5002
- triton-lang#4935 - required additional cherry picks triton-lang#4991 and triton-lang#4951
- triton-lang#4998
- triton-lang#4925
- triton-lang#5281
- triton-lang#5308
- All previous LLVM hash PRs before triton-lang#5308

---------

Co-authored-by: Ilya V <[email protected]>
Co-authored-by: Lei Zhang <[email protected]>
Co-authored-by: Lixun Zhang <[email protected]>
Co-authored-by: Keren Zhou <[email protected]>
Co-authored-by: Alexander Efimov <[email protected]>
Co-authored-by: Kyle Wang <[email protected]>
Co-authored-by: Jungwook Park <[email protected]>
Co-authored-by: peterbell10 <[email protected]>
Co-authored-by: Hongtao Yu <[email protected]>
(cherry picked from commit 2d8093c)
jataylo added a commit to jataylo/triton that referenced this pull request Dec 18, 2024
@jataylo
[AMD] rc/3.2.x cherry picks (triton-lang#5347)
44e18cc 
Reverts triton-lang#5191 due to some mlir errors in pytorch unit tests

Smaller set of cherry picks:
- triton-lang#5308 (and previous LLVM upgrades)
- triton-lang#5281
- triton-lang#4925
- triton-lang#5053
- triton-lang#5019
- triton-lang#4998

---------

Co-authored-by: Jungwook Park <[email protected]>
Co-authored-by: peterbell10 <[email protected]>
Co-authored-by: Hongtao Yu <[email protected]>
Co-authored-by: Lei Zhang <[email protected]>
Co-authored-by: Ilya V <[email protected]>
Co-authored-by: Kyle Wang <[email protected]>
(cherry picked from commit 7e401df)
bertmaher pushed a commit that referenced this pull request Dec 19, 2024
@jataylo @joviliast
@antiagainst @zhanglx13 @Jokeren @binarman @knwng @jungpark-mlir @peterbell10 @htyu @bertmaher
[AMD] release/3.2.x AMD perf cherry picks (#5191)
74abcd4 
Cherry pick list:
- #4925
- #5053
- #5019
- #5002
- #4935 - required additional cherry picks #4991 and #4951
- #4998
- #4925
- #5281
- #5308
- All previous LLVM hash PRs before #5308

---------

Co-authored-by: Ilya V <[email protected]>
Co-authored-by: Lei Zhang <[email protected]>
Co-authored-by: Lixun Zhang <[email protected]>
Co-authored-by: Keren Zhou <[email protected]>
Co-authored-by: Alexander Efimov <[email protected]>
Co-authored-by: Kyle Wang <[email protected]>
Co-authored-by: Jungwook Park <[email protected]>
Co-authored-by: peterbell10 <[email protected]>
Co-authored-by: Hongtao Yu <[email protected]>
bertmaher pushed a commit that referenced this pull request Dec 19, 2024
@jataylo @jungpark-mlir
@peterbell10 @htyu @antiagainst @joviliast @knwng @bertmaher
[AMD] rc/3.2.x cherry picks (#5347)
1f4daa3 
Reverts #5191 due to some mlir errors in pytorch unit tests

Smaller set of cherry picks:
- #5308 (and previous LLVM upgrades)
- #5281 
- #4925 
- #5053 
- #5019 
- #4998

---------

Co-authored-by: Jungwook Park <[email protected]>
Co-authored-by: peterbell10 <[email protected]>
Co-authored-by: Hongtao Yu <[email protected]>
Co-authored-by: Lei Zhang <[email protected]>
Co-authored-by: Ilya V <[email protected]>
Co-authored-by: Kyle Wang <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@sjw36 sjw36 sjw36 approved these changes

@zhanglx13 zhanglx13 zhanglx13 approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@antiagainst @zhanglx13 @sjw36