[MLIR][GPU-LLVM] Add GPU to LLVM-SPV address space mapping (#102621)

Implement mapping:

- `global`: 1
- `workgroup`: 3
- `private`: 0

Add `addressSpaceToStorageClass`, mapping GPU address spaces to SPIR-V
storage classes to be able to use SPIR-V's
`storageClassToAddressSpace`, mapping SPIR-V storage classes to LLVM
address spaces according to our mapping above *by definition*.

---------

Signed-off-by: Victor Perez <[email protected]>

Author: victor-eds

mlir/include/mlir/Conversion/GPUCommon/AttrToSPIRVConverter.h

@@ -0,0 +1,18 @@
+//===- AttrToSPIRVConverter.h - GPU attributes conversion to SPIR-V - C++ -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_CONVERSION_GPUCOMMON_ATTRTOSPIRVCONVERTER_H_
+#define MLIR_CONVERSION_GPUCOMMON_ATTRTOSPIRVCONVERTER_H_
+
+#include <mlir/Dialect/GPU/IR/GPUDialect.h>
+#include <mlir/Dialect/SPIRV/IR/SPIRVEnums.h>
+
+namespace mlir {
+spirv::StorageClass addressSpaceToStorageClass(gpu::AddressSpace addressSpace);
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_GPUCOMMON_ATTRTOSPIRVCONVERTER_H_

mlir/include/mlir/Conversion/GPUToLLVMSPV/GPUToLLVMSPVPass.h

@@ -16,12 +16,17 @@ class DialectRegistry;
 class LLVMTypeConverter;
 class RewritePatternSet;
 class Pass;
+class TypeConverter;
 
 #define GEN_PASS_DECL_CONVERTGPUOPSTOLLVMSPVOPS
 #include "mlir/Conversion/Passes.h.inc"
 
 void populateGpuToLLVMSPVConversionPatterns(LLVMTypeConverter &converter,
                                             RewritePatternSet &patterns);
+
+/// Populates memory space attribute conversion rules for lowering
+/// gpu.address_space to integer values.
+void populateGpuMemorySpaceAttributeConversions(TypeConverter &typeConverter);
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUTOLLVMSPV_GPUTOLLVMSPVPASS_H_

mlir/lib/Conversion/GPUCommon/AttrToSPIRVConverter.cpp

@@ -0,0 +1,23 @@
+//===- AttrToSPIRVConverter.cpp - GPU attributes conversion to SPIR-V - C++===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <mlir/Conversion/GPUCommon/AttrToSPIRVConverter.h>
+
+namespace mlir {
+spirv::StorageClass addressSpaceToStorageClass(gpu::AddressSpace addressSpace) {
+  switch (addressSpace) {
+  case gpu::AddressSpace::Global:
+    return spirv::StorageClass::CrossWorkgroup;
+  case gpu::AddressSpace::Workgroup:
+    return spirv::StorageClass::Workgroup;
+  case gpu::AddressSpace::Private:
+    return spirv::StorageClass::Private;
+  }
+  llvm_unreachable("Unhandled storage class");
+}
+} // namespace mlir

mlir/lib/Conversion/GPUCommon/CMakeLists.txt

@@ -15,6 +15,7 @@ if (MLIR_ENABLE_ROCM_CONVERSIONS)
 endif()
 
 add_mlir_conversion_library(MLIRGPUToGPURuntimeTransforms
+  AttrToSPIRVConverter.cpp
   GPUToLLVMConversion.cpp
   GPUOpsLowering.cpp
 

mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp

@@ -9,6 +9,8 @@
 #include "mlir/Conversion/GPUToLLVMSPV/GPUToLLVMSPVPass.h"
 
 #include "../GPUCommon/GPUOpsLowering.h"
+#include "mlir/Conversion/GPUCommon/AttrToSPIRVConverter.h"
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
@@ -328,6 +330,7 @@ struct GPUToLLVMSPVConversionPass final
                         gpu::ReturnOp, gpu::ShuffleOp, gpu::ThreadIdOp>();
 
     populateGpuToLLVMSPVConversionPatterns(converter, patterns);
+    populateGpuMemorySpaceAttributeConversions(converter);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
@@ -341,6 +344,15 @@ struct GPUToLLVMSPVConversionPass final
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
+namespace {
+static unsigned
+gpuAddressSpaceToOCLAddressSpace(gpu::AddressSpace addressSpace) {
+  constexpr spirv::ClientAPI clientAPI = spirv::ClientAPI::OpenCL;
+  return storageClassToAddressSpace(clientAPI,
+                                    addressSpaceToStorageClass(addressSpace));
+}
+} // namespace
+
 void populateGpuToLLVMSPVConversionPatterns(LLVMTypeConverter &typeConverter,
                                             RewritePatternSet &patterns) {
   patterns.add<GPUBarrierConversion, GPUReturnOpLowering, GPUShuffleConversion,
@@ -349,12 +361,11 @@ void populateGpuToLLVMSPVConversionPatterns(LLVMTypeConverter &typeConverter,
                LaunchConfigOpConversion<gpu::BlockDimOp>,
                LaunchConfigOpConversion<gpu::ThreadIdOp>,
                LaunchConfigOpConversion<gpu::GlobalIdOp>>(typeConverter);
-  constexpr spirv::ClientAPI clientAPI = spirv::ClientAPI::OpenCL;
   MLIRContext *context = &typeConverter.getContext();
   unsigned privateAddressSpace =
-      storageClassToAddressSpace(clientAPI, spirv::StorageClass::Function);
+      gpuAddressSpaceToOCLAddressSpace(gpu::AddressSpace::Private);
   unsigned localAddressSpace =
-      storageClassToAddressSpace(clientAPI, spirv::StorageClass::Workgroup);
+      gpuAddressSpaceToOCLAddressSpace(gpu::AddressSpace::Workgroup);
   OperationName llvmFuncOpName(LLVM::LLVMFuncOp::getOperationName(), context);
   StringAttr kernelBlockSizeAttributeName =
       LLVM::LLVMFuncOp::getReqdWorkGroupSizeAttrName(llvmFuncOpName);
@@ -366,4 +377,9 @@ void populateGpuToLLVMSPVConversionPatterns(LLVMTypeConverter &typeConverter,
           LLVM::CConv::SPIR_KERNEL, LLVM::CConv::SPIR_FUNC,
           /*encodeWorkgroupAttributionsAsArguments=*/true});
 }
+
+void populateGpuMemorySpaceAttributeConversions(TypeConverter &typeConverter) {
+  populateGpuMemorySpaceAttributeConversions(typeConverter,
+                                             gpuAddressSpaceToOCLAddressSpace);
+}
 } // namespace mlir

mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir

@@ -444,7 +444,7 @@ gpu.module @kernels {
 // CHECK:           %[[VAL_17:.*]] = llvm.insertvalue %[[VAL_15]], %[[VAL_16]][0]
 // CHECK:           llvm.insertvalue %[[VAL_15]], %[[VAL_17]][1]
   gpu.func @kernel_with_private_attributions()
-      private(%arg2: memref<32xf32>, %arg3: memref<16xi16>)
+      private(%arg2: memref<32xf32, #gpu.address_space<private>>, %arg3: memref<16xi16, #gpu.address_space<private>>)
       kernel {
     gpu.return
   }
@@ -471,7 +471,7 @@ gpu.module @kernels {
 // CHECK:           %[[VAL_42:.*]] = llvm.insertvalue %[[VAL_30]], %[[VAL_41]][0]
 // CHECK:           llvm.insertvalue %[[VAL_30]], %[[VAL_42]][1]
   gpu.func @kernel_with_workgoup_attributions()
-      workgroup(%arg2: memref<32xf32, 3>, %arg3: memref<16xi16, 3>)
+      workgroup(%arg2: memref<32xf32, #gpu.address_space<workgroup>>, %arg3: memref<16xi16, #gpu.address_space<workgroup>>)
       kernel {
     gpu.return
   }
@@ -491,8 +491,8 @@ gpu.module @kernels {
 // CHECK-64:        %[[VAL_92:.*]] = llvm.alloca %[[VAL_91]] x i64 : (i64) -> !llvm.ptr
 // CHECK-32:        %[[VAL_92:.*]] = llvm.alloca %[[VAL_91]] x i32 : (i64) -> !llvm.ptr
   gpu.func @kernel_with_both_attributions()
-      workgroup(%arg4: memref<8xf32, 3>, %arg5: memref<16xindex, 3>)
-      private(%arg6: memref<32xi32>, %arg7: memref<32xindex>)
+      workgroup(%arg4: memref<8xf32, #gpu.address_space<workgroup>>, %arg5: memref<16xindex, #gpu.address_space<workgroup>>)
+      private(%arg6: memref<32xi32, #gpu.address_space<private>>, %arg7: memref<32xindex, #gpu.address_space<private>>)
       kernel {
     gpu.return
   }
@@ -503,3 +503,15 @@ gpu.module @kernels {
     gpu.return
   }
 }
+
+// -----
+
+gpu.module @kernels {
+// CHECK-LABEL:   llvm.func spir_funccc @address_spaces(
+// CHECK-SAME:                                          {{.*}}: !llvm.ptr<1>
+// CHECK-SAME:                                          {{.*}}: !llvm.ptr<3>
+// CHECK-SAME:                                          {{.*}}: !llvm.ptr
+  gpu.func @address_spaces(%arg0: memref<f32, #gpu.address_space<global>>, %arg1: memref<f32, #gpu.address_space<workgroup>>, %arg2: memref<f32, #gpu.address_space<private>>) {
+    gpu.return
+  }
+}