From 4aa4158875c1bf0c8370f82e851c38e08f686225 Mon Sep 17 00:00:00 2001 From: Victor Perez Date: Wed, 7 Aug 2024 16:43:25 +0100 Subject: [PATCH 1/4] [MLIR][GPU-LLVM] Add GPU to LLVM-SPV address space mapping 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 --- .../GPUCommon/AttrToSPIRVConverter.h | 18 +++++++++++++++ .../GPUToLLVMSPV/GPUToLLVMSPVPass.h | 5 ++++ .../GPUCommon/AttrToSPIRVConverter.cpp | 23 +++++++++++++++++++ mlir/lib/Conversion/GPUCommon/CMakeLists.txt | 1 + .../Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp | 23 ++++++++++++++++--- .../GPUToLLVMSPV/gpu-to-llvm-spv.mlir | 12 ++++++++++ 6 files changed, 79 insertions(+), 3 deletions(-) create mode 100644 mlir/include/mlir/Conversion/GPUCommon/AttrToSPIRVConverter.h create mode 100644 mlir/lib/Conversion/GPUCommon/AttrToSPIRVConverter.cpp diff --git a/mlir/include/mlir/Conversion/GPUCommon/AttrToSPIRVConverter.h b/mlir/include/mlir/Conversion/GPUCommon/AttrToSPIRVConverter.h new file mode 100644 index 0000000000000..e1dcedba094c0 --- /dev/null +++ b/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 +#include + +namespace mlir { +spirv::StorageClass addressSpaceToStorageClass(gpu::AddressSpace addressSpace); +} // namespace mlir + +#endif // MLIR_CONVERSION_GPUCOMMON_ATTRTOSPIRVCONVERTER_H_ diff --git a/mlir/include/mlir/Conversion/GPUToLLVMSPV/GPUToLLVMSPVPass.h b/mlir/include/mlir/Conversion/GPUToLLVMSPV/GPUToLLVMSPVPass.h index e156c3093e21b..20c9f9c394525 100644 --- a/mlir/include/mlir/Conversion/GPUToLLVMSPV/GPUToLLVMSPVPass.h +++ b/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_ diff --git a/mlir/lib/Conversion/GPUCommon/AttrToSPIRVConverter.cpp b/mlir/lib/Conversion/GPUCommon/AttrToSPIRVConverter.cpp new file mode 100644 index 0000000000000..fcabad32974ac --- /dev/null +++ b/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 + +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 diff --git a/mlir/lib/Conversion/GPUCommon/CMakeLists.txt b/mlir/lib/Conversion/GPUCommon/CMakeLists.txt index 55bd168b83a0a..ce914c0ea3dd8 100644 --- a/mlir/lib/Conversion/GPUCommon/CMakeLists.txt +++ b/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 diff --git a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp index 36e4a6a38a68e..de0848579a8e6 100644 --- a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp +++ b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp @@ -6,6 +6,8 @@ // //===----------------------------------------------------------------------===// +#include "mlir/Conversion/GPUCommon/AttrToSPIRVConverter.h" +#include "mlir/Conversion/GPUCommon/GPUCommonPass.h" #include "mlir/Conversion/GPUToLLVMSPV/GPUToLLVMSPVPass.h" #include "../GPUCommon/GPUOpsLowering.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,14 @@ struct GPUToLLVMSPVConversionPass final //===----------------------------------------------------------------------===// namespace mlir { +namespace { +static unsigned +storageClassToOCLAddressSpace(spirv::StorageClass storageClass) { + constexpr spirv::ClientAPI clientAPI = spirv::ClientAPI::OpenCL; + return storageClassToAddressSpace(clientAPI, storageClass); +} +} // namespace + void populateGpuToLLVMSPVConversionPatterns(LLVMTypeConverter &typeConverter, RewritePatternSet &patterns) { patterns.add, LaunchConfigOpConversion, LaunchConfigOpConversion>(typeConverter); - constexpr spirv::ClientAPI clientAPI = spirv::ClientAPI::OpenCL; MLIRContext *context = &typeConverter.getContext(); unsigned privateAddressSpace = - storageClassToAddressSpace(clientAPI, spirv::StorageClass::Function); + storageClassToOCLAddressSpace(spirv::StorageClass::Function); unsigned localAddressSpace = - storageClassToAddressSpace(clientAPI, spirv::StorageClass::Workgroup); + storageClassToOCLAddressSpace(spirv::StorageClass::Workgroup); OperationName llvmFuncOpName(LLVM::LLVMFuncOp::getOperationName(), context); StringAttr kernelBlockSizeAttributeName = LLVM::LLVMFuncOp::getReqdWorkGroupSizeAttrName(llvmFuncOpName); @@ -366,4 +376,11 @@ void populateGpuToLLVMSPVConversionPatterns(LLVMTypeConverter &typeConverter, LLVM::CConv::SPIR_KERNEL, LLVM::CConv::SPIR_FUNC, /*encodeWorkgroupAttributionsAsArguments=*/true}); } + +void populateGpuMemorySpaceAttributeConversions(TypeConverter &typeConverter) { + populateGpuMemorySpaceAttributeConversions( + typeConverter, [](gpu::AddressSpace space) -> unsigned { + return storageClassToOCLAddressSpace(addressSpaceToStorageClass(space)); + }); +} } // namespace mlir diff --git a/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir b/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir index 8e133288b832b..1d6af1a3a9c36 100644 --- a/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir +++ b/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir @@ -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>, %arg1: memref>, %arg2: memref>) { + gpu.return + } +} From 30bf24d1b8ea57e2ec18a5bb6fc4f06c3ca95ebc Mon Sep 17 00:00:00 2001 From: Victor Perez Date: Fri, 9 Aug 2024 17:21:51 +0100 Subject: [PATCH 2/4] Format includes --- mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp index de0848579a8e6..f9231ab612d62 100644 --- a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp +++ b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp @@ -6,11 +6,11 @@ // //===----------------------------------------------------------------------===// -#include "mlir/Conversion/GPUCommon/AttrToSPIRVConverter.h" -#include "mlir/Conversion/GPUCommon/GPUCommonPass.h" #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" From 19fd231d6b184ff88b12577e0c9800c501d9363d Mon Sep 17 00:00:00 2001 From: Victor Perez Date: Sat, 10 Aug 2024 16:38:53 +0100 Subject: [PATCH 3/4] Use address spaces in workgroup and private attributions --- mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir b/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir index 1d6af1a3a9c36..ec4f4a304d507 100644 --- a/mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir +++ b/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>, %arg3: memref<16xi16, #gpu.address_space>) 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>, %arg3: memref<16xi16, #gpu.address_space>) 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>, %arg5: memref<16xindex, #gpu.address_space>) + private(%arg6: memref<32xi32, #gpu.address_space>, %arg7: memref<32xindex, #gpu.address_space>) kernel { gpu.return } From f29be7411fb3ad71d240dcdbfadea393df2edaca Mon Sep 17 00:00:00 2001 From: Victor Perez Date: Mon, 12 Aug 2024 11:03:15 +0100 Subject: [PATCH 4/4] Further simplify --- mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp index f9231ab612d62..ced4236402923 100644 --- a/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp +++ b/mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp @@ -346,9 +346,10 @@ struct GPUToLLVMSPVConversionPass final namespace mlir { namespace { static unsigned -storageClassToOCLAddressSpace(spirv::StorageClass storageClass) { +gpuAddressSpaceToOCLAddressSpace(gpu::AddressSpace addressSpace) { constexpr spirv::ClientAPI clientAPI = spirv::ClientAPI::OpenCL; - return storageClassToAddressSpace(clientAPI, storageClass); + return storageClassToAddressSpace(clientAPI, + addressSpaceToStorageClass(addressSpace)); } } // namespace @@ -362,9 +363,9 @@ void populateGpuToLLVMSPVConversionPatterns(LLVMTypeConverter &typeConverter, LaunchConfigOpConversion>(typeConverter); MLIRContext *context = &typeConverter.getContext(); unsigned privateAddressSpace = - storageClassToOCLAddressSpace(spirv::StorageClass::Function); + gpuAddressSpaceToOCLAddressSpace(gpu::AddressSpace::Private); unsigned localAddressSpace = - storageClassToOCLAddressSpace(spirv::StorageClass::Workgroup); + gpuAddressSpaceToOCLAddressSpace(gpu::AddressSpace::Workgroup); OperationName llvmFuncOpName(LLVM::LLVMFuncOp::getOperationName(), context); StringAttr kernelBlockSizeAttributeName = LLVM::LLVMFuncOp::getReqdWorkGroupSizeAttrName(llvmFuncOpName); @@ -378,9 +379,7 @@ void populateGpuToLLVMSPVConversionPatterns(LLVMTypeConverter &typeConverter, } void populateGpuMemorySpaceAttributeConversions(TypeConverter &typeConverter) { - populateGpuMemorySpaceAttributeConversions( - typeConverter, [](gpu::AddressSpace space) -> unsigned { - return storageClassToOCLAddressSpace(addressSpaceToStorageClass(space)); - }); + populateGpuMemorySpaceAttributeConversions(typeConverter, + gpuAddressSpaceToOCLAddressSpace); } } // namespace mlir