[Mlir-commits] [mlir] [AMDGPU] [MLIR] Add 96 and 128 bit GatherToLDS for gfx950 (PR #147496)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 8 02:42:03 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Daniel Hernandez-Juarez (dhernandez0)
<details>
<summary>Changes</summary>
This PR adds 96 and 128 gather_to_lds support for gfx950. Updating lowering, verifier and tests.
---
Full diff: https://github.com/llvm/llvm-project/pull/147496.diff
5 Files Affected:
- (modified) mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp (+9-4)
- (modified) mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp (+5-3)
- (added) mlir/test/Conversion/AMDGPUToROCDL/load_lds-gfx950.mlir (+90)
- (modified) mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir (+2-2)
- (modified) mlir/test/Conversion/AMDGPUToROCDL/transpose_load.mlir (+1-1)
``````````diff
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 910fe1b1d93c1..63bfadd4e57af 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -1198,14 +1198,19 @@ struct GatherToLDSOpLowering : public ConvertOpToLLVMPattern<GatherToLDSOp> {
Type transferType = op.getTransferType();
size_t loadWidth = [&]() -> size_t {
if (auto transferVectorType = dyn_cast<VectorType>(transferType)) {
- return transferVectorType.getNumElements() *
- (transferVectorType.getElementTypeBitWidth() / 8);
+ return (transferVectorType.getNumElements() *
+ transferVectorType.getElementTypeBitWidth()) /
+ 8;
}
return transferType.getIntOrFloatBitWidth() / 8;
}();
- // Currently only 1, 2, and 4 byte loads are supported.
- if (loadWidth != 1 && loadWidth != 2 && loadWidth != 4)
+ // Currently only 1, 2, 4, 14 and 16 byte loads are supported.
+ if (loadWidth != 1 && loadWidth != 2 && loadWidth != 4 && loadWidth != 12 &&
+ loadWidth != 16)
+ return op.emitOpError("chipset unsupported element size");
+
+ if (chipset != kGfx950 && (loadWidth == 12 || loadWidth == 16))
return op.emitOpError("chipset unsupported element size");
Value srcPtr =
diff --git a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp
index 4613d14461969..312d26f257d57 100644
--- a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp
+++ b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp
@@ -502,7 +502,7 @@ LogicalResult GatherToLDSOp::verify() {
if (elemType != dstType.getElementType())
return emitOpError("source and destination element types must match");
- // copy type sizes should be 1, 2, or 4 bytes.
+ // copy type sizes should be 1, 2, 4, 12 or 16 bytes.
auto transferType = getTransferType();
size_t transferSize;
if (auto vectorTransfer = dyn_cast<VectorType>(transferType)) {
@@ -511,8 +511,10 @@ LogicalResult GatherToLDSOp::verify() {
} else {
transferSize = transferType.getIntOrFloatBitWidth();
}
- if (transferSize != 8 && transferSize != 16 && transferSize != 32)
- return emitOpError("Transfering type size must be 8, 16, or 32 bits");
+ if (transferSize != 8 && transferSize != 16 && transferSize != 32 &&
+ transferSize != 96 && transferSize != 128)
+ return emitOpError(
+ "Transfering type size must be 8, 16, 32, 96 or 128 bits");
if (!hasGlobalMemorySpace(srcType.getMemorySpace()) &&
!hasFatRawBufferMemorySpace(srcType.getMemorySpace()))
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/load_lds-gfx950.mlir b/mlir/test/Conversion/AMDGPUToROCDL/load_lds-gfx950.mlir
new file mode 100644
index 0000000000000..0bf1054b98270
--- /dev/null
+++ b/mlir/test/Conversion/AMDGPUToROCDL/load_lds-gfx950.mlir
@@ -0,0 +1,90 @@
+// RUN: not mlir-opt %s --split-input-file -convert-amdgpu-to-rocdl=chipset=gfx942 2>&1 | FileCheck %s --check-prefix=GFX942
+// RUN: mlir-opt %s --split-input-file -convert-amdgpu-to-rocdl=chipset=gfx950 | FileCheck %s --check-prefix=GFX950
+
+#gpu_global_addrspace = 1
+#gpu_lds_addrspace = 3
+#amdgpu_fat_buffer_addrspace = 7
+
+// GFX950-LABEL: func @fat_buffer_load_to_rocdl_f96
+// GFX950-SAME: (%[[ARG0:.*]]: memref<128x72xf32, 7>)
+func.func @fat_buffer_load_to_rocdl_f96(%global : memref<128x72xf32, #amdgpu_fat_buffer_addrspace>) {
+ %c0 = arith.constant 0 : index
+ %c12 = arith.constant 12 : index
+ %c32 = arith.constant 32 : index
+ %alloc = memref.alloc() : memref<64x64xf32, #gpu_lds_addrspace>
+ // GFX950: %[[BUFFER_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]
+
+ // GFX950: %[[C0:.*]] = arith.constant 0 : index
+ // GFX950: %[[IC0:.*]] = builtin.unrealized_conversion_cast %c0 : index to i64
+ // GFX950: %[[C12:.*]] = arith.constant 12 : index
+ // GFX950: %[[IC12:.*]] = builtin.unrealized_conversion_cast %[[C12]]
+ // GFX950: %[[C32:.*]] = arith.constant 32 : index
+ // GFX950: %[[IC32:.*]] = builtin.unrealized_conversion_cast %[[C32]]
+
+ // GFX950: %[[ALLOC:.*]] = memref.alloc()
+ // GFX950: %[[LDS_DESC:.*]] = builtin.unrealized_conversion_cast
+ // GFX950: %[[GLOBAL_BASE:.*]] = llvm.extractvalue %[[BUFFER_DESC]][1]
+
+ // GFX950: %[[C72:.*]] = llvm.mlir.constant(72 : index) : i64
+ // GFX950: %[[MUL:.*]] = llvm.mul %[[IC12]], %[[C72]] : i64
+ // GFX950: %[[SRC_OFFSET:.*]] = llvm.add %[[MUL]], %[[IC0]] : i64
+
+ // GFX950: %[[GLOBAL_PTR:.*]] = llvm.getelementptr %[[GLOBAL_BASE]][%[[SRC_OFFSET]]]
+ // GFX950: %[[LDS_BASE:.*]] = llvm.extractvalue %[[LDS_DESC]][1]
+
+ // GFX950: %[[C64:.*]] = llvm.mlir.constant(64 : index) : i64
+ // GFX950: %[[MUL_2:.*]] = llvm.mul %[[IC32]], %[[C64]] : i64
+ // GFX950: %[[DST_OFFSET:.*]] = llvm.add %[[MUL_2]], %[[IC0]] : i64
+
+ // GFX950: %[[LDS_PTR:.*]] = llvm.getelementptr %[[LDS_BASE]][%[[DST_OFFSET]]]
+ // GFX950: rocdl.load.to.lds %[[GLOBAL_PTR]], %[[LDS_PTR]], 12
+ // GFX942: error: 'amdgpu.gather_to_lds' op chipset unsupported element size
+ amdgpu.gather_to_lds %global[%c12, %c0], %alloc[%c32, %c0]
+ : vector<16xf6E3M2FN>, memref<128x72xf32, #amdgpu_fat_buffer_addrspace>, memref<64x64xf32, #gpu_lds_addrspace>
+ func.return
+}
+
+// -----
+
+#gpu_global_addrspace = 1
+#gpu_lds_addrspace = 3
+#amdgpu_fat_buffer_addrspace = 7
+
+// GFX950-LABEL: func @fat_buffer_load_to_rocdl_f128
+// GFX950-SAME: (%[[ARG0:.*]]: memref<128x72xf32, 7>)
+func.func @fat_buffer_load_to_rocdl_f128(%global : memref<128x72xf32, #amdgpu_fat_buffer_addrspace>) {
+ %c0 = arith.constant 0 : index
+ %c12 = arith.constant 12 : index
+ %c32 = arith.constant 32 : index
+ %alloc = memref.alloc() : memref<64x64xf32, #gpu_lds_addrspace>
+ // GFX950: %[[BUFFER_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]
+
+ // GFX950: %[[C0:.*]] = arith.constant 0 : index
+ // GFX950: %[[IC0:.*]] = builtin.unrealized_conversion_cast %c0 : index to i64
+ // GFX950: %[[C12:.*]] = arith.constant 12 : index
+ // GFX950: %[[IC12:.*]] = builtin.unrealized_conversion_cast %[[C12]]
+ // GFX950: %[[C32:.*]] = arith.constant 32 : index
+ // GFX950: %[[IC32:.*]] = builtin.unrealized_conversion_cast %[[C32]]
+
+ // GFX950: %[[ALLOC:.*]] = memref.alloc()
+ // GFX950: %[[LDS_DESC:.*]] = builtin.unrealized_conversion_cast
+ // GFX950: %[[GLOBAL_BASE:.*]] = llvm.extractvalue %[[BUFFER_DESC]][1]
+
+ // GFX950: %[[C72:.*]] = llvm.mlir.constant(72 : index) : i64
+ // GFX950: %[[MUL:.*]] = llvm.mul %[[IC12]], %[[C72]] : i64
+ // GFX950: %[[SRC_OFFSET:.*]] = llvm.add %[[MUL]], %[[IC0]] : i64
+
+ // GFX950: %[[GLOBAL_PTR:.*]] = llvm.getelementptr %[[GLOBAL_BASE]][%[[SRC_OFFSET]]]
+ // GFX950: %[[LDS_BASE:.*]] = llvm.extractvalue %[[LDS_DESC]][1]
+
+ // GFX950: %[[C64:.*]] = llvm.mlir.constant(64 : index) : i64
+ // GFX950: %[[MUL_2:.*]] = llvm.mul %[[IC32]], %[[C64]] : i64
+ // GFX950: %[[DST_OFFSET:.*]] = llvm.add %[[MUL_2]], %[[IC0]] : i64
+
+ // GFX950: %[[LDS_PTR:.*]] = llvm.getelementptr %[[LDS_BASE]][%[[DST_OFFSET]]]
+ // GFX950: rocdl.load.to.lds %[[GLOBAL_PTR]], %[[LDS_PTR]], 16
+ // GFX942: error: 'amdgpu.gather_to_lds' op chipset unsupported element size
+ amdgpu.gather_to_lds %global[%c12, %c0], %alloc[%c32, %c0]
+ : f128, memref<128x72xf32, #amdgpu_fat_buffer_addrspace>, memref<64x64xf32, #gpu_lds_addrspace>
+ func.return
+}
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir b/mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir
index 581346e03b893..fc03d53ea5730 100644
--- a/mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir
+++ b/mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir
@@ -1,4 +1,5 @@
-// RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx942 | FileCheck %s
+// RUN: mlir-opt %s --split-input-file -convert-amdgpu-to-rocdl=chipset=gfx942 | FileCheck %s
+// RUN: mlir-opt %s --split-input-file -convert-amdgpu-to-rocdl=chipset=gfx950 | FileCheck %s
#gpu_global_addrspace = 1
#gpu_lds_addrspace = 3
@@ -118,7 +119,6 @@ func.func @global_load_to_rocdl_vec(%global : memref<128x72xi16, #gpu_global_add
func.return
}
-
// CHECK-LABEL: func @global_load_to_rocdl_dynamic_indices
// CHECK-SAME: (%[[ARG0:.*]]: memref<512xi32, 1>, %[[SRC_IDX:.*]]: index, %[[DST_IDX:.*]]: index)
func.func @global_load_to_rocdl_dynamic_indices(%global : memref<512xi32, #gpu_global_addrspace>, %src_idx : index, %dst_idx : index) {
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/transpose_load.mlir b/mlir/test/Conversion/AMDGPUToROCDL/transpose_load.mlir
index 68799098f1d36..73481e4bced20 100644
--- a/mlir/test/Conversion/AMDGPUToROCDL/transpose_load.mlir
+++ b/mlir/test/Conversion/AMDGPUToROCDL/transpose_load.mlir
@@ -1,5 +1,5 @@
// RUN: mlir-opt %s --split-input-file -convert-amdgpu-to-rocdl=chipset=gfx950 | FileCheck %s
-// RUN: not mlir-opt %s --split-input-file -convert-amdgpu-to-rocdl=chipset=gfx945 2>&1 | FileCheck %s --check-prefix=CHECK-OLD
+// RUN: not mlir-opt %s --split-input-file -convert-amdgpu-to-rocdl=chipset=gfx942 2>&1 | FileCheck %s --check-prefix=CHECK-OLD
// CHECK-LABEL: func @transpose_load_to_rocdl_4xf16
func.func @transpose_load_to_rocdl_4xf16(%idx1 : index, %idx2 : index, %wgmem : memref<128x72xf16, 3>) -> vector<4xf16> {
``````````
</details>
https://github.com/llvm/llvm-project/pull/147496
More information about the Mlir-commits
mailing list