[mlir][GPU] Fixes subgroup reduce lowering #141825
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir-gpu Author: Muzammil (Muzammiluddin-Syed-ECE) ChangesFixes issues with subgroup reduce lowering producing erroneous results. Aligning more closely with triton implementation : triton-lang/triton#5019 Full diff: https://github.com/llvm/llvm-project/pull/141825.diff 1 Files Affected:
diff --git a/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp b/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp
index 74face4291353..b7a12763057f0 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp
@@ -447,29 +447,36 @@ createSubgroupDPPReduction(PatternRewriter &rewriter, gpu::SubgroupReduceOp op,
if (ci.clusterSize >= 64) {
if (chipset.majorVersion <= 9) {
// Broadcast 31st lane value to rows 2 and 3.
- // Use row mask to avoid polluting rows 0 and 1.
dpp = rewriter.create<amdgpu::DPPOp>(
loc, res.getType(), res, res, amdgpu::DPPPerm::row_bcast_31,
- rewriter.getUnitAttr(), 0xc, allBanks,
- /*bound_ctrl*/ false);
+ rewriter.getUnitAttr(), 0xf, allBanks,
+ /*bound_ctrl*/ true);
+ res = vector::makeArithReduction(
+ rewriter, loc, gpu::convertReductionKind(mode), dpp, res);
+ // Obtain reduction from last rows, the previous rows are polluted.
+ Value lane63 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(63));
+ res = rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane63);
} else if (chipset.majorVersion <= 12) {
// Assume reduction across 32 lanes has been done.
// Perform final reduction manually by summing values in lane 0 and
// lane 32.
- Value lane0 = rewriter.create<arith::ConstantOp>(
- loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
- Value lane32 = rewriter.create<arith::ConstantOp>(
- loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(32));
- dpp = rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane32);
- res = rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane0);
+ Value lane31 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(31));
+ Value lane63 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(63));
+ lane0 =
+ rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane31);
+ lane32 =
+ rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane63);
+ res = vector::makeArithReduction(
+ rewriter, loc, gpu::convertReductionKind(mode), lane31, lane63);
} else {
return rewriter.notifyMatchFailure(
op, "Subgroup reduce lowering to DPP not currently supported for "
"this device.");
}
- res = vector::makeArithReduction(rewriter, loc,
- gpu::convertReductionKind(mode), res, dpp);
}
assert(res.getType() == input.getType());
return res;
|
|
@llvm/pr-subscribers-mlir Author: Muzammil (Muzammiluddin-Syed-ECE) ChangesFixes issues with subgroup reduce lowering producing erroneous results. Aligning more closely with triton implementation : triton-lang/triton#5019 Full diff: https://github.com/llvm/llvm-project/pull/141825.diff 1 Files Affected:
diff --git a/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp b/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp
index 74face4291353..b7a12763057f0 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp
@@ -447,29 +447,36 @@ createSubgroupDPPReduction(PatternRewriter &rewriter, gpu::SubgroupReduceOp op,
if (ci.clusterSize >= 64) {
if (chipset.majorVersion <= 9) {
// Broadcast 31st lane value to rows 2 and 3.
- // Use row mask to avoid polluting rows 0 and 1.
dpp = rewriter.create<amdgpu::DPPOp>(
loc, res.getType(), res, res, amdgpu::DPPPerm::row_bcast_31,
- rewriter.getUnitAttr(), 0xc, allBanks,
- /*bound_ctrl*/ false);
+ rewriter.getUnitAttr(), 0xf, allBanks,
+ /*bound_ctrl*/ true);
+ res = vector::makeArithReduction(
+ rewriter, loc, gpu::convertReductionKind(mode), dpp, res);
+ // Obtain reduction from last rows, the previous rows are polluted.
+ Value lane63 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(63));
+ res = rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane63);
} else if (chipset.majorVersion <= 12) {
// Assume reduction across 32 lanes has been done.
// Perform final reduction manually by summing values in lane 0 and
// lane 32.
- Value lane0 = rewriter.create<arith::ConstantOp>(
- loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
- Value lane32 = rewriter.create<arith::ConstantOp>(
- loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(32));
- dpp = rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane32);
- res = rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane0);
+ Value lane31 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(31));
+ Value lane63 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(63));
+ lane0 =
+ rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane31);
+ lane32 =
+ rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane63);
+ res = vector::makeArithReduction(
+ rewriter, loc, gpu::convertReductionKind(mode), lane31, lane63);
} else {
return rewriter.notifyMatchFailure(
op, "Subgroup reduce lowering to DPP not currently supported for "
"this device.");
}
- res = vector::makeArithReduction(rewriter, loc,
- gpu::convertReductionKind(mode), res, dpp);
}
assert(res.getType() == input.getType());
return res;
|
Muzammiluddin-Syed-ECE
force-pushed
the
muzasyed/sub2
branch
2 times, most recently
from
a1ecf0e to
7b7e0e5
Compare
Please explain what types of erroneous results can occur, providing an example. |
Signed-off-by: Muzammiluddin Syed <[email protected]>
Muzammiluddin-Syed-ECE
force-pushed
the
muzasyed/sub2
branch
from
7b7e0e5 to
75e849e
Compare
To summarize the debugging situation - and to potentially provide language for the PR description - "the final reduction steps were taken from an implementation of scan, not reduction, causing lanes earlier in the wave to have incorrect results due to masking" |
Hypothetical exampleTo provide an explanation of the issue with the current implementation, let's take the simple example of attempting to perform a sum over 64 lanes where the initial values are as follows (first lane has value 1, and all other lanes have value 0): When performing a sum reduction over these 64 lanes, in the current implementation we perform 6 dpp instructions which in sequential order do the following:
After step 4) the result for the example above looks like this: After step 5) the result looks like this: After step 6) the result looks like this: Note that the correct value here is always 1, yet after the So what this means:
However in the current implementation we do not specifically read the value from one of the correct lanes when returning a final value. In some workloads it seems without this specification, the stale value from the first lane is returned instead. Actual failing testFor a specific example of how the current implementation causes issues, take a look at the IR below which represents an additive reduction over a dynamic dimension. When provided an input of type |
google-yfyang
pushed a commit
to google-yfyang/llvm-project
that referenced
this pull request
May 29, 2025
Fixes the final reduction steps which were taken from an implementation of scan, not reduction, causing lanes earlier in the wave to have incorrect results due to masking. Now aligning more closely with triton implementation : triton-lang/triton#5019 # Hypothetical example To provide an explanation of the issue with the current implementation, let's take the simple example of attempting to perform a sum over 64 lanes where the initial values are as follows (first lane has value 1, and all other lanes have value 0): ``` [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] ``` When performing a sum reduction over these 64 lanes, in the current implementation we perform 6 dpp instructions which in sequential order do the following: 1) sum over clusters of 2 contiguous lanes 2) sum over clusters of 4 contiguous lanes 3) sum over clusters of 8 contiguous lanes 4) sum over an entire row 5) broadcast the result of last lane in each row to the next row and each lane sums current value with incoming value. 5) broadcast the result of the 32nd lane to last two rows and each lane sums current value with incoming value. After step 4) the result for the example above looks like this: ``` [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] ``` After step 5) the result looks like this: ``` [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] ``` After step 6) the result looks like this: ``` [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] ``` Note that the correct value here is always 1, yet after the `dpp.broadcast` ops some lanes have incorrect values. The reason is that for these incorrect lanes, like lanes 0-15 in step 5, the `dpp.broadcast` op doesn't provide them incoming values from other lanes. Instead these lanes are provided either their own values, or 0 (depending on whether `bound_ctrl` is true or false) as values to sum over, either way these values are stale and these lanes shouldn't be used in general. So what this means: - For a subgroup reduce over 32 lanes (like Step 5), the correct result is stored in lanes 16 to 31 - For a subgroup reduce over 64 lanes (like Step 6), the correct result is stored in lanes 32 to 63. However in the current implementation we do not specifically read the value from one of the correct lanes when returning a final value. In some workloads it seems without this specification, the stale value from the first lane is returned instead. # Actual failing test For a specific example of how the current implementation causes issues, take a look at the IR below which represents an additive reduction over a dynamic dimension. ``` !matA = tensor<1x?xf16> !matB = tensor<1xf16> #map = affine_map<(d0, d1) -> (d0, d1)> #map1 = affine_map<(d0, d1) -> (d0)> func.func @only_producer_fusion_multiple_result(%arg0: !matA) -> !matB { %cst_1 = arith.constant 0.000000e+00 : f16 %c2_i64 = arith.constant 2 : i64 %0 = tensor.empty() : !matB %2 = linalg.fill ins(%cst_1 : f16) outs(%0 : !matB) -> !matB %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "reduction"]} ins(%arg0 : !matA) outs(%2 : !matB) { ^bb0(%in: f16, %out: f16): %7 = arith.addf %in, %out : f16 linalg.yield %7 : f16 } -> !matB return %4 : !matB } ``` When provided an input of type `tensor<1x2xf16>` and values `{0, 1}` to perform the reduction over, the value returned is consistently 4. By the same analysis done above, this shows that the returned value is coming from one of these stale lanes and needs to be read instead from one of the lanes storing the correct result. Signed-off-by: Muzammiluddin Syed <[email protected]>
sivan-shani
pushed a commit
to sivan-shani/llvm-project
that referenced
this pull request
Jun 3, 2025
Fixes the final reduction steps which were taken from an implementation of scan, not reduction, causing lanes earlier in the wave to have incorrect results due to masking. Now aligning more closely with triton implementation : triton-lang/triton#5019 # Hypothetical example To provide an explanation of the issue with the current implementation, let's take the simple example of attempting to perform a sum over 64 lanes where the initial values are as follows (first lane has value 1, and all other lanes have value 0): ``` [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] ``` When performing a sum reduction over these 64 lanes, in the current implementation we perform 6 dpp instructions which in sequential order do the following: 1) sum over clusters of 2 contiguous lanes 2) sum over clusters of 4 contiguous lanes 3) sum over clusters of 8 contiguous lanes 4) sum over an entire row 5) broadcast the result of last lane in each row to the next row and each lane sums current value with incoming value. 5) broadcast the result of the 32nd lane to last two rows and each lane sums current value with incoming value. After step 4) the result for the example above looks like this: ``` [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] ``` After step 5) the result looks like this: ``` [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] ``` After step 6) the result looks like this: ``` [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] ``` Note that the correct value here is always 1, yet after the `dpp.broadcast` ops some lanes have incorrect values. The reason is that for these incorrect lanes, like lanes 0-15 in step 5, the `dpp.broadcast` op doesn't provide them incoming values from other lanes. Instead these lanes are provided either their own values, or 0 (depending on whether `bound_ctrl` is true or false) as values to sum over, either way these values are stale and these lanes shouldn't be used in general. So what this means: - For a subgroup reduce over 32 lanes (like Step 5), the correct result is stored in lanes 16 to 31 - For a subgroup reduce over 64 lanes (like Step 6), the correct result is stored in lanes 32 to 63. However in the current implementation we do not specifically read the value from one of the correct lanes when returning a final value. In some workloads it seems without this specification, the stale value from the first lane is returned instead. # Actual failing test For a specific example of how the current implementation causes issues, take a look at the IR below which represents an additive reduction over a dynamic dimension. ``` !matA = tensor<1x?xf16> !matB = tensor<1xf16> #map = affine_map<(d0, d1) -> (d0, d1)> #map1 = affine_map<(d0, d1) -> (d0)> func.func @only_producer_fusion_multiple_result(%arg0: !matA) -> !matB { %cst_1 = arith.constant 0.000000e+00 : f16 %c2_i64 = arith.constant 2 : i64 %0 = tensor.empty() : !matB %2 = linalg.fill ins(%cst_1 : f16) outs(%0 : !matB) -> !matB %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "reduction"]} ins(%arg0 : !matA) outs(%2 : !matB) { ^bb0(%in: f16, %out: f16): %7 = arith.addf %in, %out : f16 linalg.yield %7 : f16 } -> !matB return %4 : !matB } ``` When provided an input of type `tensor<1x2xf16>` and values `{0, 1}` to perform the reduction over, the value returned is consistently 4. By the same analysis done above, this shows that the returned value is coming from one of these stale lanes and needs to be read instead from one of the lanes storing the correct result. Signed-off-by: Muzammiluddin Syed <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fixes the final reduction steps which were taken from an implementation of scan, not reduction, causing lanes earlier in the wave to have incorrect results due to masking.
Now aligning more closely with triton implementation : triton-lang/triton#5019
Hypothetical example
To provide an explanation of the issue with the current implementation, let's take the simple example of attempting to perform a sum over 64 lanes where the initial values are as follows (first lane has value 1, and all other lanes have value 0):
When performing a sum reduction over these 64 lanes, in the current implementation we perform 6 dpp instructions which in sequential order do the following:
After step 4) the result for the example above looks like this:
After step 5) the result looks like this:
After step 6) the result looks like this:
Note that the correct value here is always 1, yet after the
dpp.broadcastops some lanes have incorrect values. The reason is that for these incorrect lanes, like lanes 0-15 in step 5, thedpp.broadcastop doesn't provide them incoming values from other lanes. Instead these lanes are provided either their own values, or 0 (depending on whetherbound_ctrlis true or false) as values to sum over, either way these values are stale and these lanes shouldn't be used in general.So what this means:
However in the current implementation we do not specifically read the value from one of the correct lanes when returning a final value. In some workloads it seems without this specification, the stale value from the first lane is returned instead.
Actual failing test
For a specific example of how the current implementation causes issues, take a look at the IR below which represents an additive reduction over a dynamic dimension.
When provided an input of type
tensor<1x2xf16>and values{0, 1}to perform the reduction over, the value returned is consistently 4. By the same analysis done above, this shows that the returned value is coming from one of these stale lanes and needs to be read instead from one of the lanes storing the correct result.