-
Notifications
You must be signed in to change notification settings - Fork 11.4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM % nits, but let's give @antiagainst a chance to take a look.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for adding support for this! Overall LGTM; just a few inline comments.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we use the lambda function in the above? This is duplicating the logic there.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice. Just one final request.
std::unique_ptr<ConversionTarget> target = | ||
spirv::getMemorySpaceToStorageClassTarget(*context); | ||
spirv::MemorySpaceToStorageClassMap memorySpaceMap = | ||
spirv::mapMemorySpaceToVulkanStorageClass; | ||
(memoryModel == spirv::MemoryModel::OpenCL) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Here we can also use the above targetEnvSupportsKernelCapability
no? Something like:
spirv::MemorySpaceToStorageClassMap memorySpaceMap =
targetEnvSupportsKernelCapability(gpuModule) ? ...
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
@@ -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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is this gpu.func->func.func necessary? Can you just keep the original gpu.func instead?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM. Thanks!
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.