[Mlir-commits] [mlir] 668c964 - [AMDGPU] [MLIR] Add 96 and 128 bit GatherToLDS for gfx950 (#147496)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 9 08:53:29 PDT 2025
Author: Daniel Hernandez-Juarez
Date: 2025-07-09T11:53:26-04:00
New Revision: 668c964282f0850a07976f46182efe745f77b789
URL: https://github.com/llvm/llvm-project/commit/668c964282f0850a07976f46182efe745f77b789
DIFF: https://github.com/llvm/llvm-project/commit/668c964282f0850a07976f46182efe745f77b789.diff
LOG: [AMDGPU] [MLIR] Add 96 and 128 bit GatherToLDS for gfx950 (#147496)
This PR adds 96 and 128 gather_to_lds support for gfx950. Updating
lowering, verifier and tests.
Added:
mlir/test/Conversion/AMDGPUToROCDL/load_lds-gfx950.mlir
Modified:
mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp
mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir
mlir/test/Conversion/AMDGPUToROCDL/transpose_load.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 910fe1b1d93c1..ef35ee208f002 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -1196,18 +1196,23 @@ struct GatherToLDSOpLowering : public ConvertOpToLLVMPattern<GatherToLDSOp> {
// augment it to transfer multiple elements per thread by issuing multiple
// `global_load_lds` instructions.
Type transferType = op.getTransferType();
- size_t loadWidth = [&]() -> size_t {
+ int loadWidth = [&]() -> int {
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, 12 and 16 byte loads are supported.
+ if (!llvm::is_contained({1, 2, 4, 12, 16}, loadWidth))
return op.emitOpError("chipset unsupported element size");
+ if (chipset != kGfx950 && llvm::is_contained({12, 16}, loadWidth))
+ return op.emitOpError("Gather to LDS instructions with 12-byte and "
+ "16-byte load widths are only supported on gfx950");
+
Value srcPtr =
getStridedElementPtr(rewriter, loc, srcMemRefType, adaptor.getSrc(),
(adaptor.getSrcIndices()));
diff --git a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp
index 4613d14461969..acaf6a2f8792a 100644
--- a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp
+++ b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp
@@ -502,17 +502,18 @@ 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;
+ int transferSize;
if (auto vectorTransfer = dyn_cast<VectorType>(transferType)) {
transferSize = vectorTransfer.getNumElements() *
vectorTransfer.getElementTypeBitWidth();
} else {
transferSize = transferType.getIntOrFloatBitWidth();
}
- if (transferSize != 8 && transferSize != 16 && transferSize != 32)
- return emitOpError("Transfering type size must be 8, 16, or 32 bits");
+ if (!llvm::is_contained({8, 16, 32, 96, 128}, transferSize))
+ 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..42fb18006bea4
--- /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 Gather to LDS instructions with 12-byte and 16-byte load widths are only supported on gfx950
+ 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 Gather to LDS instructions with 12-byte and 16-byte load widths are only supported on gfx950
+ 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..77103fa5c25f1 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 -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> {
More information about the Mlir-commits
mailing list