[mlir] Prepare convert-gpu-to-spirv for OpenCL support #69941
Conversation
…ion pipeline for OpenCL kernels. This includes a couple of changes to pass behavior for OpenCL kernels. Vulkan shaders are not impacted by the changes. 1. SPIRV 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.
|
@kuhar @antiagainst Replacing #66445 since a new pass flag was not required. |
|
This PR is one of the components for breaking down #65539 into smaller PRs and use new GPU compilation pipeline. |
|
@antiagainst @kuhar @Hardcode84 @qedawkins Waiting for review. |
There was a problem hiding this comment.
Left some stylistic suggestions. @antiagainst, could you also take a look?
Sort alphabetically. Use if-else instead of ternary expression. List lambda captures explicitly. Replace auto with actually type. Add newline before big block.
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
There was a problem hiding this comment.
LGTM % nits, but let's give @antiagainst a chance to take a look.
| gpuModules.push_back(builder.clone(*moduleOp.getOperation())); | ||
| }); | ||
|
|
||
| // Run conversion for each module independently as they can have different | ||
| // TargetEnv attributes. | ||
| for (Operation *gpuModule : gpuModules) { | ||
| mlir::spirv::TargetEnvAttr targetAttr = |
There was a problem hiding this comment.
Can we use the lambda function in the above? This is duplicating the logic there.
There was a problem hiding this comment.
Now the lambda returns boolean and cannot be used. Also that line is from the original code. I just moved it up a bit since I needed access to targetAttr before memory space mapping.
| std::unique_ptr<ConversionTarget> target = | ||
| spirv::getMemorySpaceToStorageClassTarget(*context); | ||
| spirv::MemorySpaceToStorageClassMap memorySpaceMap = | ||
| spirv::mapMemorySpaceToVulkanStorageClass; | ||
| (memoryModel == spirv::MemoryModel::OpenCL) |
There was a problem hiding this comment.
Here we can also use the above targetEnvSupportsKernelCapability no? Something like:
spirv::MemorySpaceToStorageClassMap memorySpaceMap =
targetEnvSupportsKernelCapability(gpuModule) ? ...
| @@ -108,6 +138,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 | |||
There was a problem hiding this comment.
Is this gpu.func->func.func necessary? Can you just keep the original gpu.func instead?
There was a problem hiding this comment.
Keeping the original gpu.func causes legality check error later in the gpu compile pipeline.
If target attr is set for a gpu.module, gpu-to-llvm pass doesn't lower gpu.launch_func
Instead, it is replaced with another gpu.launch_func that has lowered argument types (llvm ptrs).
If a gpu.func remains as an input to gpu-to-llvm pass, there is an argument mismatch between the new gpu.launch_func and the gpu.func. And an error is fired.
The reason for putting a dummy func.func here is to work around that check.
For func types other than gpu.func, argument types are not checked against gpu.launch_func. A func.func is still need as there will be a symbol check.
There was a problem hiding this comment.
I think the problem is, the lack of spirv support in gpu dialect. For example, gpu.func needs to be able to wrap spirv.func so gpu-to-llvm pass (for the host code) can properly handle the relation between gpu.launch_func and spirv.func.
Basically, using dummy func.func looks little hacky and it'd be also nice if the divergence between Vulkan/OpenCL IR structures could be avoided. However, considering the progress of this commit, we can discuss this later for the future enhancement.
Really appreciate this work and look forward to seeing it merged.
antiagainst
changed the title
|
I've updated the tile and commit message slightly to make it look nicer once land as a git commit. Thanks! |
|
@antiagainst @joker-eph Can someone merge this PR? Thanks in advance. |
Done. Thanks for the contribution again! |
This includes a couple of changes to pass behavior for OpenCL kernels. Vulkan shaders are not impacted by the changes.