[mlir][amdgpu] Move GPU memory spaces conversion to single place (llvm#171876)

Author: Hardcode84

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

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(

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());

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
+  )

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<LLVMTypeConverter &>(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<amdgpu::Chipset> maybeChipset =
       amdgpu::Chipset::parse(getChipset());
   assert(llvm::succeeded(maybeChipset) && "expected valid chipset");