diff --git a/mlir/include/mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h b/mlir/include/mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h index 44aae7b65a69c..4942c39f9745f 100644 --- a/mlir/include/mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h +++ b/mlir/include/mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h @@ -29,6 +29,11 @@ void populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter, RewritePatternSet &patterns, amdgpu::Chipset chipset); +/// Remap common GPU memory spaces (Workgroup, Private, etc) to LLVM address +/// spaces. +void populateCommonAMDGPUTypeAndAttributeConversions( + TypeConverter &typeConverter); + /// Remap AMDGPU memory spaces to LLVM address spaces /// by mapping amdgpu::AddressSpace::fat_raw_buffer to ptr addrspace(7), /// amdgpu::AddressSpace::buffer_rsrc to ptr addrspace(8), and diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp index 91154b846f567..4b1509392aa6f 100644 --- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp +++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp @@ -2727,18 +2727,7 @@ struct ConvertAMDGPUToROCDLPass LLVMTypeConverter converter(ctx); 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"); - }); + populateCommonAMDGPUTypeAndAttributeConversions(converter); LLVMConversionTarget target(getContext()); target.addIllegalDialect<::mlir::amdgpu::AMDGPUDialect>(); target.addLegalDialect<::mlir::LLVM::LLVMDialect>(); @@ -2750,6 +2739,22 @@ struct ConvertAMDGPUToROCDLPass }; } // namespace +void mlir::populateCommonAMDGPUTypeAndAttributeConversions( + TypeConverter &typeConverter) { + populateGpuMemorySpaceAttributeConversions( + typeConverter, [](gpu::AddressSpace space) { + switch (space) { + case gpu::AddressSpace::Global: + return ROCDL::ROCDLDialect::kGlobalMemoryAddressSpace; + case gpu::AddressSpace::Workgroup: + return ROCDL::ROCDLDialect::kSharedMemoryAddressSpace; + case gpu::AddressSpace::Private: + return ROCDL::ROCDLDialect::kPrivateMemoryAddressSpace; + } + llvm_unreachable("unknown address space enum value"); + }); +} + void mlir::populateAMDGPUTypeAndAttributeConversions( TypeConverter &typeConverter) { typeConverter.addTypeAttributeConversion( diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp index c03f3a5d3889c..7d7feb58aa726 100644 --- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp +++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp @@ -349,19 +349,7 @@ struct LowerGpuOpsToROCDLOpsPass final } LLVMTypeConverter converter(ctx, options); - 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"); - return 0; - }); + populateCommonAMDGPUTypeAndAttributeConversions(converter); RewritePatternSet llvmPatterns(ctx); LLVMConversionTarget target(getContext()); diff --git a/mlir/lib/Dialect/GPU/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/GPU/TransformOps/CMakeLists.txt index e5cc0254f1ffe..2d3dfdbfc1136 100644 --- a/mlir/lib/Dialect/GPU/TransformOps/CMakeLists.txt +++ b/mlir/lib/Dialect/GPU/TransformOps/CMakeLists.txt @@ -10,7 +10,7 @@ add_mlir_dialect_library(MLIRGPUTransformOps MLIRGPUTransformOpsIncGen MLIRDeviceMappingInterfacesIncGen MLIRGPUDeviceMapperEnumsGen - + LINK_LIBS PUBLIC MLIRGPUDialect MLIRGPUTransforms @@ -22,7 +22,8 @@ add_mlir_dialect_library(MLIRGPUTransformOps MLIRVectorTransforms # ConversionPatterns - MLIRNVGPUToNVVM + MLIRAMDGPUToROCDL MLIRGPUToNVVMTransforms MLIRGPUToROCDLTransforms - ) + MLIRNVGPUToNVVM + ) diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp index 0a3ef7d5c9890..fdace3b662314 100644 --- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp +++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp @@ -8,6 +8,7 @@ #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h" +#include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h" #include "mlir/Conversion/GPUCommon/GPUCommonPass.h" #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h" #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h" @@ -128,18 +129,7 @@ LogicalResult transform::ApplyGPUSubgroupReduceToNVVMConversionPatternsOp:: void transform::ApplyGPUToROCDLConversionPatternsOp::populatePatterns( TypeConverter &typeConverter, RewritePatternSet &patterns) { auto &llvmTypeConverter = static_cast(typeConverter); - populateGpuMemorySpaceAttributeConversions( - llvmTypeConverter, [](AddressSpace space) { - switch (space) { - case AddressSpace::Global: - return ROCDL::ROCDLDialect::kGlobalMemoryAddressSpace; - case AddressSpace::Workgroup: - return ROCDL::ROCDLDialect::kSharedMemoryAddressSpace; - case AddressSpace::Private: - return ROCDL::ROCDLDialect::kPrivateMemoryAddressSpace; - } - llvm_unreachable("unknown address space enum value"); - }); + populateCommonAMDGPUTypeAndAttributeConversions(llvmTypeConverter); FailureOr maybeChipset = amdgpu::Chipset::parse(getChipset()); assert(llvm::succeeded(maybeChipset) && "expected valid chipset");