Reland [mlir][amdgpu] Add common gpu mem space conversions to convert-amdgpu-to-rocdl (#171599)

Reland https://github.com/llvm/llvm-project/pull/171543

Added missing GPU lib `MLIRGPUToGPURuntimeTransforms`.
This commit is contained in:
Ivan Butygin
2025-12-10 17:33:51 +03:00
committed by GitHub
parent 8c49509512
commit c9c4e6eb58
3 changed files with 57 additions and 5 deletions

View File

@@ -8,6 +8,7 @@
#include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h"
#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
@@ -2730,6 +2731,18 @@ struct ConvertAMDGPUToROCDLPass
});
populateAMDGPUToROCDLConversionPatterns(converter, patterns, *maybeChipset);
populateGpuMemorySpaceAttributeConversions(
converter, [](gpu::AddressSpace space) {
switch (space) {
case gpu::AddressSpace::Global:
return 1;
case gpu::AddressSpace::Workgroup:
return 3;
case gpu::AddressSpace::Private:
return 5;
}
llvm_unreachable("unknown address space enum value");
});
LLVMConversionTarget target(getContext());
target.addIllegalDialect<::mlir::amdgpu::AMDGPUDialect>();
target.addLegalDialect<::mlir::LLVM::LLVMDialect>();

View File

@@ -11,11 +11,12 @@ add_mlir_conversion_library(MLIRAMDGPUToROCDL
Core
LINK_LIBS PUBLIC
MLIRLLVMCommonConversion
MLIRLLVMDialect
MLIRROCDLDialect
MLIRAMDGPUDialect
MLIRAMDGPUUtils
MLIRGPUToGPURuntimeTransforms
MLIRLLVMCommonConversion
MLIRLLVMDialect
MLIRPass
MLIRROCDLDialect
MLIRTransforms
)

View File

@@ -1,5 +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
// 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
@@ -43,6 +43,44 @@ func.func @global_load_to_rocdl_f32(%global : memref<128x72xf32, #gpu_global_add
func.return
}
// CHECK-LABEL: func @global_load_to_rocdl_wg_mem
// CHECK-SAME: (%[[ARG0:.*]]: memref<128x72xf32>)
func.func @global_load_to_rocdl_wg_mem(%global : memref<128x72xf32>) {
%c0 = arith.constant 0 : index
%c12 = arith.constant 12 : index
%c32 = arith.constant 32 : index
%alloc = memref.alloc() : memref<64x64xf32, #gpu.address_space<workgroup>>
// CHECK: %[[GLOBAL_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[IC0:.*]] = builtin.unrealized_conversion_cast %c0 : index to i64
// CHECK: %[[C12:.*]] = arith.constant 12 : index
// CHECK: %[[IC12:.*]] = builtin.unrealized_conversion_cast %[[C12]]
// CHECK: %[[C32:.*]] = arith.constant 32 : index
// CHECK: %[[IC32:.*]] = builtin.unrealized_conversion_cast %[[C32]]
// CHECK: %[[ALLOC:.*]] = memref.alloc()
// CHECK: %[[LDS_DESC:.*]] = builtin.unrealized_conversion_cast
// CHECK: %[[GLOBAL_BASE:.*]] = llvm.extractvalue %[[GLOBAL_DESC]][1]
// CHECK: %[[C72:.*]] = llvm.mlir.constant(72 : index) : i64
// CHECK: %[[MUL:.*]] = llvm.mul %[[IC12]], %[[C72]] : i64
// CHECK: %[[SRC_OFFSET:.*]] = llvm.add %[[MUL]], %[[IC0]] : i64
// CHECK: %[[GLOBAL_PTR:.*]] = llvm.getelementptr %[[GLOBAL_BASE]][%[[SRC_OFFSET]]]
// CHECK: %[[LDS_BASE:.*]] = llvm.extractvalue %[[LDS_DESC]][1]
// CHECK: %[[C64:.*]] = llvm.mlir.constant(64 : index) : i64
// CHECK: %[[MUL_2:.*]] = llvm.mul %[[IC32]], %[[C64]] : i64
// CHECK: %[[DST_OFFSET:.*]] = llvm.add %[[MUL_2]], %[[IC0]] : i64
// CHECK: %[[LDS_PTR:.*]] = llvm.getelementptr %[[LDS_BASE]][%[[DST_OFFSET]]]
// CHECK: rocdl.load.to.lds %[[GLOBAL_PTR]], %[[LDS_PTR]], 4
amdgpu.gather_to_lds %global[%c12, %c0], %alloc[%c32, %c0]
: f32, memref<128x72xf32>, memref<64x64xf32, #gpu.address_space<workgroup>>
func.return
}
// CHECK-LABEL: func @global_load_to_rocdl_i8
// CHECK-SAME: (%[[ARG0:.*]]: memref<128x72xi8, 1>)
func.func @global_load_to_rocdl_i8(%global : memref<128x72xi8, #gpu_global_addrspace>) {