[mlir][AMDGPU] Support vector<2xf16> inputs to buffer atomic fadd #108238
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
Extend the lowering of atomic.fadd to support the v2f16 variant avaliable on some AMDGPU chips. Co-authored-by: Giuseppe Rossini <[email protected]>
|
@llvm/pr-subscribers-backend-amdgpu @llvm/pr-subscribers-mlir-amdgpu Author: Krzysztof Drewniak (krzysz00) ChangesExtend the lowering of atomic.fadd to support the v2f16 variant avaliable on some AMDGPU chips. Full diff: https://github.com/llvm/llvm-project/pull/108238.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
index 8a1ef94c853a58..64db4448bc2f2b 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
@@ -253,8 +253,8 @@ def AMDGPU_RawBufferAtomicCmpswapOp :
// Raw buffer atomic floating point add
def AMDGPU_RawBufferAtomicFaddOp :
AMDGPU_Op<"raw_buffer_atomic_fadd", [AllElementTypesMatch<["value", "memref"]>,
- AttrSizedOperandSegments]>,
- Arguments<(ins F32:$value,
+ AttrSizedOperandSegments]>,
+ Arguments<(ins AnyTypeOf<[F32, VectorOfLengthAndType<[2], [F16]>]>:$value,
Arg<AnyMemRef, "buffer to operate on", [MemRead, MemWrite]>:$memref,
Variadic<I32>:$indices,
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 96b433294d258a..fc5dd7c5602129 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -115,15 +115,18 @@ struct RawBufferOpLowering : public ConvertOpToLLVMPattern<GpuOp> {
rewriter.getIntegerType(floatType.getWidth()));
}
if (auto dataVector = dyn_cast<VectorType>(wantedDataType)) {
+ uint32_t vecLen = dataVector.getNumElements();
uint32_t elemBits = dataVector.getElementTypeBitWidth();
- uint32_t totalBits = elemBits * dataVector.getNumElements();
+ uint32_t totalBits = elemBits * vecLen;
+ bool usePackedFp16 =
+ dyn_cast_or_null<RawBufferAtomicFaddOp>(*gpuOp) && vecLen == 2;
if (totalBits > maxVectorOpWidth)
return gpuOp.emitOpError(
"Total width of loads or stores must be no more than " +
Twine(maxVectorOpWidth) + " bits, but we call for " +
Twine(totalBits) +
" bits. This should've been caught in validation");
- if (elemBits < 32) {
+ else if (!usePackedFp16 && elemBits < 32) {
if (totalBits > 32) {
if (totalBits % 32 != 0)
return gpuOp.emitOpError("Load or store of more than 32-bits that "
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/amdgpu-to-rocdl.mlir b/mlir/test/Conversion/AMDGPUToROCDL/amdgpu-to-rocdl.mlir
index 717667c22af800..cc51a8c40942f9 100644
--- a/mlir/test/Conversion/AMDGPUToROCDL/amdgpu-to-rocdl.mlir
+++ b/mlir/test/Conversion/AMDGPUToROCDL/amdgpu-to-rocdl.mlir
@@ -151,6 +151,17 @@ func.func @gpu_gcn_raw_buffer_atomic_fadd_f32(%value: f32, %buf: memref<64xf32>,
func.return
}
+// CHECK-LABEL: func @gpu_gcn_raw_buffer_atomic_fadd_v2f16
+func.func @gpu_gcn_raw_buffer_atomic_fadd_v2f16(%value: vector<2xf16>, %buf: memref<64xf16>, %idx: i32) {
+ // CHECK: %[[numRecords:.*]] = llvm.mlir.constant(128 : i32)
+ // GFX9: %[[flags:.*]] = llvm.mlir.constant(159744 : i32)
+ // RDNA: %[[flags:.*]] = llvm.mlir.constant(822243328 : i32)
+ // CHECK: %[[resource:.*]] = rocdl.make.buffer.rsrc %{{.*}}, %{{.*}}, %[[numRecords]], %[[flags]]
+ // CHECK: rocdl.raw.ptr.buffer.atomic.fadd %{{.*}}, %[[resource]], %{{.*}}, %{{.*}}, %{{.*}} : vector<2xf16>
+ amdgpu.raw_buffer_atomic_fadd {boundsCheck = true} %value -> %buf[%idx] : vector<2xf16> -> memref<64xf16>, i32
+ func.return
+}
+
// CHECK-LABEL: func @gpu_gcn_raw_buffer_atomic_fmax_f32
func.func @gpu_gcn_raw_buffer_atomic_fmax_f32(%value: f32, %buf: memref<64xf32>, %idx: i32) {
// CHECK: %[[numRecords:.*]] = llvm.mlir.constant(256 : i32)
|
arsenm
reviewed
Sep 11, 2024
|
Oh, that's a thing now? Sure! |
|
Update: bf16 support would need me to back out the "bf16 is i16" change that still lives in MLIR from back when |
kuhar
reviewed
Sep 11, 2024
| if (totalBits > maxVectorOpWidth) | ||
| return gpuOp.emitOpError( | ||
| "Total width of loads or stores must be no more than " + | ||
| Twine(maxVectorOpWidth) + " bits, but we call for " + | ||
| Twine(totalBits) + | ||
| " bits. This should've been caught in validation"); | ||
| if (elemBits < 32) { | ||
| else if (!usePackedFp16 && elemBits < 32) { |
There was a problem hiding this comment.
nit: no else after return: https://llvm.org/docs/CodingStandards.html#don-t-use-else-after-a-return
There was a problem hiding this comment.
... I thought I saw an approval, reverting, thanks
krzysz00
added a commit
to krzysz00/llvm-project
that referenced
this pull request
Sep 11, 2024
…fadd (llvm#108238)" This reverts commit 0d48d4d.
kuhar
reviewed
Sep 11, 2024
| uint32_t totalBits = elemBits * dataVector.getNumElements(); | ||
| uint32_t totalBits = elemBits * vecLen; | ||
| bool usePackedFp16 = | ||
| dyn_cast_or_null<RawBufferAtomicFaddOp>(*gpuOp) && vecLen == 2; |
krzysz00
added a commit
to krzysz00/llvm-project
that referenced
this pull request
Sep 11, 2024
Extend the lowering of atomic.fadd to support the v2f16 variant avaliable on some AMDGPU chips. Re-lands llvm#108238 (and addresses review comments from there) Co-authored-by: Giuseppe Rossini <[email protected]>
krzysz00
added a commit
to krzysz00/llvm-project
that referenced
this pull request
Sep 11, 2024
Extend the lowering of atomic.fadd to support the v2f16 variant avaliable on some AMDGPU chips. Re-lands llvm#108238 (and addresses review comments from there) Co-authored-by: Giuseppe Rossini <[email protected]>
krzysz00
added a commit
that referenced
this pull request
Sep 11, 2024
…08286) Extend the lowering of atomic.fadd to support the v2f16 variant avaliable on some AMDGPU chips. Re-lands #108238 (and addresses review comments from there) Co-authored-by: Giuseppe Rossini <[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.
Extend the lowering of atomic.fadd to support the v2f16 variant avaliable on some AMDGPU chips.