[AMD] Sink the 2nd tt.load after local_load's #4823

zhanglx13 · 2024-09-28T02:16:15Z

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

third_party/amd/lib/TritonAMDGPUTransforms/ReorderInstructions.cpp

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

antiagainst

Some final nits.

test/TritonGPU/amd/amd-sched-2nd-load.mlir

third_party/amd/lib/TritonAMDGPUTransforms/ReorderInstructions.cpp

Keep two positive and two negative cases for single dot. And one for two dot.

This reverts commit 664ac51.

…#4920)

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

…#4920)

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

This PR adds more restrictions about when should we apply the sched-load optimizations and un-revert #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

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

…#4920)

…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

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

…#4920)

…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

…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)

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

…#4920)

…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)

* [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]>

…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) (cherry picked from commit f6053a3)

…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) (cherry picked from commit f6053a3) (cherry picked from commit 72d1575)

@lezcano

* [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) (cherry picked from commit 4d70942) (cherry picked from commit 2f8aacc) * [AMD] Restructure ReorderInstructions pass (triton-lang#4998) (cherry picked from commit 86a2ac7) (cherry picked from commit 4c7d56e) (cherry picked from commit 0529343) * [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) (cherry picked from commit d0142d3) (cherry picked from commit 9f2b69b) * [AMD] Add missing dependency to TritonAMDGPUIR (triton-lang#5053) TritonAMDGPUTransforms now depends on it. (cherry picked from commit 0b443ce) (cherry picked from commit 37cec47) (cherry picked from commit 1ab334d) * [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) (cherry picked from commit ca8842c) (cherry picked from commit 3a3902d) * [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) (cherry picked from commit f6053a3) (cherry picked from commit 72d1575) * [BACKEND][NVIDIA] Add Lowering for Shared-to-MMAv3-DotOp Copy (triton-lang#5009) Allows for upcasting in DotOp encoding in RF. This lowering path is not currently in use; pending triton-lang#5003 (cherry picked from commit cfddb09) (cherry picked from commit f8c2c30) (cherry picked from commit 73ef337) * [AMD] Add initial support for scaled_dot(mxfp8, fp8) (triton-lang#4994) This commit adds initial support for scaled_dot with mxfp8 LHS and fp8 RHS. It supports both mfma32 and mfma16 intrinsic variants. Right now we are missing software emulation for `Float8E4M3FN` type, so this only enables for `Float8E5M2`. (cherry picked from commit 3549db8) (cherry picked from commit efe0ec4) (cherry picked from commit 010fe45) * [Frontend][Backend] Implement support for scale_dot(-, bf16) (triton-lang#4996) In the passing we also improve a few other things: - Now `scaled_dot` accepts both uint8/uint16 fp8/bf16 as inputs (before you had to cast it to uint8, which was weird when extending it to bf16). - Add `scaled_dot` to the docs and improve the docs overall (have not render them, might need a few further tweaks) (cherry picked from commit 23c9ec1) (cherry picked from commit 675758b) (cherry picked from commit 4e04af0) * [BACKEND] Improve detection of register to register conversion (triton-lang#4991) Specifically, it fixes problems when `srcLayout` and `dstLayout` have different number of registers but the same number of not free registers. We solved the problem by padding free registers to either `srcLayout` or `dstLayout`, but this can be improved by fixing the `invertAndCompose` function. (cherry picked from commit 15c5e55) (cherry picked from commit 6537eb6) (cherry picked from commit 4ca5013) * [BACKEND] Replace `isMmaToDotShortcut` with linear layout based logic (triton-lang#4951) This PR removes the legacy `isMmaToDotShortcut` and its associated shortcut conversion. (cherry picked from commit 1d5fdfe) (cherry picked from commit fc6d96b) (cherry picked from commit 9f67c54) * [BACKEND]Fix DotOperand(Ampere) LinearLayoutConversion (triton-lang#5038) We also clean a bit `TritonGPU/IR/Dialect.cpp` using some auxiliary functions to make the intentions a bit clearer. We add a few asserts in the `LinearLayoutConversion` to make sure it's clear why we do certain things here and there. We also kill `getCvtOrder`, as it was not used anywhere (cherry picked from commit 56584c4) (cherry picked from commit 276d182) (cherry picked from commit 72651c2) * [BACKEND] Fix uses of getOrder(DotOperand(Nvidia) and MMA(Nvidia)) (triton-lang#5055) We use `getOrder` very liberally throughout the codebase, when we really meant to use `getThreadOrder`. This is an issue with the input layout is an `DotOperand(mma(opIdx=1))`, where the thread order and the matrix order are opposite. Found this to be an issue when a PR changed the `getOrder` of `DotOperand(Hopper)` to an incorrect one and CI still passed! The issue here is that the LLVM lowering for wgmma and the LinearLayout does not use `getOrder`, but there are many other subsystems do, and many heuristics would be getting an incorrect order, and potentially be disabled. This is particularly problematic for `DotOperand(opIdx=1)` in nvidia hardware, as `getThreadOrder` and `getOrder` are different! While doing so we: - Audit most (all?) the calls to `getOrder(dotOperand)`. It turns out that most of them really meant `getThreadOrder` - Fix the ordering methods of `SliceEncodingAttr` to be consistent - Move the implementation of `getWarpOrder` to the Attr classes, because of OOP The test strategy was to add `llvm::report_fatal_error("Testing");` within `getOrder(nvidiaMma)` and `getOrder(DotOperand(nvidiaMma))` and triaging all errors that were raised in CI. (cherry picked from commit 38a11b8) (cherry picked from commit 8412154) (cherry picked from commit a569c3e) * [AMD] Reland instruction scheduling hint changes (triton-lang#4940) This commit relands triton-lang#4819 with the following fixes: * Changed to a better way to mark opIdx for loads * Replaced temlate-based `rewindUnaryOps` to use regular for-loops. The new way is more robust and can handle other unary ops automatically. * Replaced `instr.sched.barriers` using the ones from `rocdl` dialect from the MLIR upstream * Extended lit tests (cherry picked from commit ee5876c) (cherry picked from commit 8dd9226) (cherry picked from commit aed3efc) * [AMD] Enable scaled_dot(-, bf16) (triton-lang#5029) (cherry picked from commit f062540) (cherry picked from commit ca75b5f) (cherry picked from commit 98149dd) * [AMD] Add support for scaled_dot(mxfp4, -) (triton-lang#5034) This commit adds support for mxfp4 typed A tensor for sacled dot in the AMD backend. We moved the `convertMxfp4x2ToBf16x2` impl from NVIDIA side to a common path to reuse. (cherry picked from commit edc5c5c) (cherry picked from commit ac9f0d0) (cherry picked from commit c0710dc) * [BACKEND] Minor Bugfixes for SharedToDotOperand MMAv3 (triton-lang#5030) Two bugfixes following triton-lang#5009. - When `BLOCK_M=64` and `num_warps > 4`, the order of warps for DotOpEncoded tensor should be M-major instead of N-major, since WGMMA expects the 4 warps in each warp group to be stacked along the M dimension. - Should use `mmaBitwidth` instead of `bitwidth` when calculating `numRep` in `SharedToDotOperandMMAv2OrV3`. This was missed in a bad rebase. @lezcano I encountered these bugs when attempting to locally test the [DotOp hoisting PR](triton-lang#5003) after rebasing (they normally would be caught by `test_core.py` but that path was not yet enabled in the last PR). With these fixes added, I was able to successfully validate against pytorch. (cherry picked from commit e82dfd9) (cherry picked from commit 5287a68) (cherry picked from commit 8d70247) (cherry picked from commit 302de9d) * [BACKEND] Get rid of unpack/pack I32 (triton-lang#5044) - Removed functions related to unpacking and packing I32 values. - Updated utilities to handle conversion of mxfp4 values without packing/unpacking I32. - Move the register value ordering logic from the element-wise operation lowering to the dot operation lowering. - Use linear layout to handle conversions between almost all distributed layouts. - Clean up data loading and mma computation involving `repN`, `repK`, and `repM`. (cherry picked from commit 1cf7b1b) (cherry picked from commit 376fe7e) (cherry picked from commit 2141a4e) (cherry picked from commit d0e4abc) * Consolidate `getOrder` as "element order" and implement `getRepOrder` for general and NVIDIA layouts (triton-lang#5089) This partially reverts commit 38a11b8. Supersedes triton-lang#5085 It also documents that we are implicitly choosing a way to tile a full tensor depending on the layout. See triton-lang#5085 (comment) (cherry picked from commit 57643b3) (cherry picked from commit ffb2032) (cherry picked from commit a11c6be) (cherry picked from commit 9d080b4) --------- 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]> Co-authored-by: Gary Geng <[email protected]> Co-authored-by: Mario Lezcano Casado <[email protected]> Co-authored-by: Keren Zhou <[email protected]> Co-authored-by: ravil-mobile <[email protected]>

zhanglx13 force-pushed the test_sched-load branch 3 times, most recently from e58042b to 6995935 Compare October 5, 2024 15:31

zhanglx13 force-pushed the test_sched-load branch 2 times, most recently from 3a8bc98 to d4a4224 Compare October 13, 2024 21:12

zhanglx13 changed the title ~~[WIP] Sink the 2nd tt.load after local_load's~~ [AMD] Sink the 2nd tt.load after local_load's Oct 13, 2024

zhanglx13 requested a review from antiagainst October 13, 2024 21:15

antiagainst requested changes Oct 13, 2024

View reviewed changes

zhanglx13 added 3 commits October 13, 2024 20:15

Sink the 2nd tt.load after local_load's

5511fb6

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

refactor code according to review

606c4db

Adding lit tests for sched-2nd-load

71855ad

zhanglx13 force-pushed the test_sched-load branch from d4a4224 to 71855ad Compare October 14, 2024 04:14

antiagainst approved these changes Oct 14, 2024

View reviewed changes

test/TritonGPU/amd/amd-sched-2nd-load.mlir Show resolved Hide resolved

third_party/amd/lib/TritonAMDGPUTransforms/ReorderInstructions.cpp Outdated Show resolved Hide resolved

zhanglx13 added 2 commits October 14, 2024 14:43

Auto generate the lit tests for various tile sizes

0baea7d

Keep two positive and two negative cases for single dot. And one for two dot.

Address review comment

532686f

antiagainst approved these changes Oct 14, 2024

View reviewed changes

antiagainst marked this pull request as ready for review October 14, 2024 20:57

antiagainst requested a review from ptillet as a code owner October 14, 2024 20:57

antiagainst merged commit 664ac51 into triton-lang:main Oct 14, 2024

ptillet added a commit that referenced this pull request Oct 16, 2024

Revert "[AMD] Sink the 2nd tt.load after local_load's (#4823)"

b6a7c33

This reverts commit 664ac51.

antiagainst pushed a commit that referenced this pull request Oct 16, 2024

[AMD] unrevert #4901; revert #4823 (#4920)

8fb7342

alexsamardzic pushed a commit to alexsamardzic/triton that referenced this pull request Oct 16, 2024

[AMD] unrevert triton-lang#4901; revert triton-lang#4823 (triton-lang…

2eecac6

…#4920)

zhanglx13 added a commit to ROCm/triton that referenced this pull request Oct 16, 2024

[AMD] Sink the 2nd tt.load after local_load's (triton-lang#4823)

3e9371a

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

zhanglx13 mentioned this pull request Oct 16, 2024

[AMD] Reland sinking the 2nd tt.load after local_load's #4935

Merged

lezcano mentioned this pull request Oct 18, 2024

[Backend] Pipeline scale_dot #4950

Merged

jtang10 pushed a commit to ROCm/triton that referenced this pull request Oct 21, 2024

[AMD] Sink the 2nd tt.load after local_load's (triton-lang#4823)

c15f5ce

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

jtang10 pushed a commit to ROCm/triton that referenced this pull request Oct 21, 2024

[AMD] unrevert triton-lang#4901; revert triton-lang#4823 (triton-lang…

9820ed5

…#4920)

zhanglx13 added a commit to ROCm/triton that referenced this pull request Oct 23, 2024

[AMD] Sink the 2nd tt.load after local_load's (triton-lang#4823)

2924a07

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

zhanglx13 added a commit to ROCm/triton that referenced this pull request Oct 28, 2024

[AMD] Sink the 2nd tt.load after local_load's (triton-lang#4823)

10747f2

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

zhanglx13 added a commit to ROCm/triton that referenced this pull request Oct 30, 2024

[AMD] Sink the 2nd tt.load after local_load's (triton-lang#4823)

3088c5d

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

Luosuu pushed a commit to Luosuu/triton that referenced this pull request Nov 13, 2024

[AMD] Sink the 2nd tt.load after local_load's (triton-lang#4823)

18c6462

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

Luosuu pushed a commit to Luosuu/triton that referenced this pull request Nov 13, 2024

[AMD] unrevert triton-lang#4901; revert triton-lang#4823 (triton-lang…

8511af3

…#4920)

guacamoleo pushed a commit to guacamoleo/triton that referenced this pull request Nov 14, 2024

[AMD] Sink the 2nd tt.load after local_load's (triton-lang#4823)

1dec4fa

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

guacamoleo pushed a commit to guacamoleo/triton that referenced this pull request Nov 14, 2024

[AMD] unrevert triton-lang#4901; revert triton-lang#4823 (triton-lang…

c7e277c

…#4920)

bertmaher pushed a commit to bertmaher/triton that referenced this pull request Dec 10, 2024

[AMD] Sink the 2nd tt.load after local_load's (triton-lang#4823)

8eb8b5a

This helps backend to interleave global load and mfma instructions and can reduce global load issue latency.

bertmaher pushed a commit to bertmaher/triton that referenced this pull request Dec 10, 2024

[AMD] unrevert triton-lang#4901; revert triton-lang#4823 (triton-lang…

f3433e1

…#4920)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[AMD] Sink the 2nd tt.load after local_load's #4823

[AMD] Sink the 2nd tt.load after local_load's #4823

Uh oh!

zhanglx13 commented Sep 28, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

antiagainst left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

[AMD] Sink the 2nd tt.load after local_load's #4823

[AMD] Sink the 2nd tt.load after local_load's #4823

Uh oh!

Conversation

zhanglx13 commented Sep 28, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

antiagainst left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants