diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp index 7584b1707522..e3dacfff16aa 100644 --- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp +++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp @@ -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>(); diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/CMakeLists.txt b/mlir/lib/Conversion/AMDGPUToROCDL/CMakeLists.txt index 43cbc0732516..25cceb9d1d37 100644 --- a/mlir/lib/Conversion/AMDGPUToROCDL/CMakeLists.txt +++ b/mlir/lib/Conversion/AMDGPUToROCDL/CMakeLists.txt @@ -11,11 +11,12 @@ add_mlir_conversion_library(MLIRAMDGPUToROCDL Core LINK_LIBS PUBLIC - MLIRLLVMCommonConversion - MLIRLLVMDialect - MLIRROCDLDialect MLIRAMDGPUDialect MLIRAMDGPUUtils + MLIRGPUToGPURuntimeTransforms + MLIRLLVMCommonConversion + MLIRLLVMDialect MLIRPass + MLIRROCDLDialect MLIRTransforms ) diff --git a/mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir b/mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir index e48c94195ea5..30578517be1c 100644 --- a/mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir +++ b/mlir/test/Conversion/AMDGPUToROCDL/load_lds.mlir @@ -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> + // 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> + 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>) {