[Mlir-commits] [mlir] b68fe86 - [mlir] Prepare convert-gpu-to-spirv for OpenCL support (#69941)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Nov 5 18:56:57 PST 2023
Author: Sang Ik Lee
Date: 2023-11-05T18:56:53-08:00
New Revision: b68fe8699feceadfaef75ed686828252aab6b08d
URL: https://github.com/llvm/llvm-project/commit/b68fe8699feceadfaef75ed686828252aab6b08d
DIFF: https://github.com/llvm/llvm-project/commit/b68fe8699feceadfaef75ed686828252aab6b08d.diff
LOG: [mlir] Prepare convert-gpu-to-spirv for OpenCL support (#69941)
This includes a couple of changes to pass behavior for OpenCL kernels.
Vulkan shaders are not impacted by the changes.
1. SPIR-V module is placed inside GPU module. This change is required for
gpu-module-to-binary to work correctly as it expects kernel function to
be inside the GPU module.
2. A dummy func.func with same kernel name as gpu.func is created. GPU
compilation pipeline defers lowering of gpu launch kernel op. Since
spirv.func is not directly tied to gpu launch kernel, a dummy func.func
is required to avoid legalization issues.
3. Use correct mapping when mapping MemRef memory space to SPIR-V
storage class for OpenCL kernels.
Added:
Modified:
mlir/include/mlir/Conversion/Passes.td
mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index c2f90b8984b97fb..8b64f6dbe741b93 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -556,7 +556,10 @@ def ConvertGPUToSPIRV : Pass<"convert-gpu-to-spirv", "ModuleOp"> {
to control the set and binding if wanted.
}];
let constructor = "mlir::createConvertGPUToSPIRVPass()";
- let dependentDialects = ["spirv::SPIRVDialect"];
+ let dependentDialects = [
+ "func::FuncDialect",
+ "spirv::SPIRVDialect",
+ ];
let options = [
Option<"use64bitIndex", "use-64bit-index",
"bool", /*default=*/"false",
diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
index 272e3de8723aeb6..ae89774239b58c1 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
@@ -17,6 +17,7 @@
#include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h"
#include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h"
#include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
@@ -54,22 +55,47 @@ void GPUToSPIRVPass::runOnOperation() {
SmallVector<Operation *, 1> gpuModules;
OpBuilder builder(context);
+
+ auto targetEnvSupportsKernelCapability = [](gpu::GPUModuleOp moduleOp) {
+ Operation *gpuModule = moduleOp.getOperation();
+ auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
+ spirv::TargetEnv targetEnv(targetAttr);
+ return targetEnv.allows(spirv::Capability::Kernel);
+ };
+
module.walk([&](gpu::GPUModuleOp moduleOp) {
// Clone each GPU kernel module for conversion, given that the GPU
// launch op still needs the original GPU kernel module.
- builder.setInsertionPoint(moduleOp.getOperation());
+ // For Vulkan Shader capabilities, we insert the newly converted SPIR-V
+ // module right after the original GPU module, as that's the expectation of
+ // the in-tree Vulkan runner.
+ // For OpenCL Kernel capabilities, we insert the newly converted SPIR-V
+ // module inside the original GPU module, as that's the expectaion of the
+ // normal GPU compilation pipeline.
+ if (targetEnvSupportsKernelCapability(moduleOp)) {
+ builder.setInsertionPoint(moduleOp.getBody(),
+ moduleOp.getBody()->begin());
+ } else {
+ builder.setInsertionPoint(moduleOp.getOperation());
+ }
gpuModules.push_back(builder.clone(*moduleOp.getOperation()));
});
// Run conversion for each module independently as they can have
diff erent
// TargetEnv attributes.
for (Operation *gpuModule : gpuModules) {
+ spirv::TargetEnvAttr targetAttr =
+ spirv::lookupTargetEnvOrDefault(gpuModule);
+
// Map MemRef memory space to SPIR-V storage class first if requested.
if (mapMemorySpace) {
std::unique_ptr<ConversionTarget> target =
spirv::getMemorySpaceToStorageClassTarget(*context);
spirv::MemorySpaceToStorageClassMap memorySpaceMap =
- spirv::mapMemorySpaceToVulkanStorageClass;
+ targetEnvSupportsKernelCapability(
+ dyn_cast<gpu::GPUModuleOp>(gpuModule))
+ ? spirv::mapMemorySpaceToOpenCLStorageClass
+ : spirv::mapMemorySpaceToVulkanStorageClass;
spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);
RewritePatternSet patterns(context);
@@ -79,7 +105,6 @@ void GPUToSPIRVPass::runOnOperation() {
return signalPassFailure();
}
- auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
std::unique_ptr<ConversionTarget> target =
SPIRVConversionTarget::get(targetAttr);
@@ -108,6 +133,25 @@ void GPUToSPIRVPass::runOnOperation() {
if (failed(applyFullConversion(gpuModule, *target, std::move(patterns))))
return signalPassFailure();
}
+
+ // For OpenCL, the gpu.func op in the original gpu.module op needs to be
+ // replaced with an empty func.func op with the same arguments as the gpu.func
+ // op. The func.func op needs gpu.kernel attribute set.
+ module.walk([&](gpu::GPUModuleOp moduleOp) {
+ if (targetEnvSupportsKernelCapability(moduleOp)) {
+ moduleOp.walk([&](gpu::GPUFuncOp funcOp) {
+ builder.setInsertionPoint(funcOp);
+ auto newFuncOp = builder.create<func::FuncOp>(
+ funcOp.getLoc(), funcOp.getName(), funcOp.getFunctionType());
+ auto entryBlock = newFuncOp.addEntryBlock();
+ builder.setInsertionPointToEnd(entryBlock);
+ builder.create<func::ReturnOp>(funcOp.getLoc());
+ newFuncOp->setAttr(gpu::GPUDialect::getKernelFuncAttrName(),
+ builder.getUnitAttr());
+ funcOp.erase();
+ });
+ }
+ });
}
} // namespace
diff --git a/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir b/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
index 0aa50cc1e25294d..4b8d17cd6449389 100644
--- a/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
@@ -12,6 +12,8 @@ module attributes {
// CHECK-SAME: {{%.*}}: !spirv.ptr<!spirv.array<12 x f32>, CrossWorkgroup>
// CHECK-NOT: spirv.interface_var_abi
// CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
+ // CHECK-LABEL: func.func @basic_module_structure
+ // CHECK-SAME: attributes {gpu.kernel}
gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>) kernel
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
gpu.return
@@ -45,6 +47,8 @@ module attributes {
// CHECK-SAME: {{%.*}}: !spirv.ptr<!spirv.array<12 x f32>, CrossWorkgroup>
// CHECK-NOT: spirv.interface_var_abi
// CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
+ // CHECK-LABEL: func.func @basic_module_structure
+ // CHECK-SAME: attributes {gpu.kernel}
gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>) kernel
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
gpu.return
More information about the Mlir-commits
mailing list