-
Notifications
You must be signed in to change notification settings - Fork 14.4k
[mlir][spirv] Fix lookup logic spirv.target_env
for gpu.module
#147262
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
base: main
Are you sure you want to change the base?
Changes from all commits
5d05b68
71d3d92
ae6ba9c
fa3b44e
7d43c59
40b6f07
0b68019
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -48,19 +48,37 @@ struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase<GPUToSPIRVPass> { | |
void runOnOperation() override; | ||
|
||
private: | ||
spirv::TargetEnvAttr lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp); | ||
spirv::TargetEnvAttr lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp); | ||
Comment on lines
+51
to
+52
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you add documentation comments explaining what these do? |
||
bool mapMemorySpace; | ||
}; | ||
|
||
spirv::TargetEnvAttr | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is there a reason to have those two functions? Why not to call There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thank you for your comment. Since I intentionally did not include Then, I defined a separate However, if adding multiple new functions is considered undesirable, I'm happy to simplify the implementation by having only There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
I don't think it's undesirable. I just felt like removing the extra indirection may offer a nicer solution, but you made some good argument about why it makes sense the way it is. Initially I thought that you could add There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thank you for the thoughtful suggestion and for being open to my perspective! |
||
GPUToSPIRVPass::lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp) { | ||
for (auto &targetAttr : moduleOp.getTargetsAttr()) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Spell out the type here since it's not obvious without an IDE/lsp: https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable |
||
if (auto spirvTargetEnvAttr = dyn_cast<spirv::TargetEnvAttr>(targetAttr)) | ||
return spirvTargetEnvAttr; | ||
|
||
return {}; | ||
} | ||
|
||
spirv::TargetEnvAttr | ||
GPUToSPIRVPass::lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp) { | ||
if (auto targetEnvAttr = lookupTargetEnvInTargets(moduleOp)) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. also here |
||
return targetEnvAttr; | ||
|
||
return spirv::lookupTargetEnvOrDefault(moduleOp); | ||
} | ||
|
||
void GPUToSPIRVPass::runOnOperation() { | ||
MLIRContext *context = &getContext(); | ||
ModuleOp module = getOperation(); | ||
|
||
SmallVector<Operation *, 1> gpuModules; | ||
OpBuilder builder(context); | ||
|
||
auto targetEnvSupportsKernelCapability = [](gpu::GPUModuleOp moduleOp) { | ||
Operation *gpuModule = moduleOp.getOperation(); | ||
auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule); | ||
auto targetEnvSupportsKernelCapability = [this](gpu::GPUModuleOp moduleOp) { | ||
auto targetAttr = lookupTargetEnvOrDefault(moduleOp); | ||
spirv::TargetEnv targetEnv(targetAttr); | ||
return targetEnv.allows(spirv::Capability::Kernel); | ||
}; | ||
|
@@ -86,7 +104,7 @@ void GPUToSPIRVPass::runOnOperation() { | |
// TargetEnv attributes. | ||
for (Operation *gpuModule : gpuModules) { | ||
spirv::TargetEnvAttr targetAttr = | ||
spirv::lookupTargetEnvOrDefault(gpuModule); | ||
lookupTargetEnvOrDefault(cast<gpu::GPUModuleOp>(gpuModule)); | ||
|
||
// Map MemRef memory space to SPIR-V storage class first if requested. | ||
if (mapMemorySpace) { | ||
|
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
@@ -0,0 +1,17 @@ | ||||||
// RUN: mlir-opt --spirv-attach-target='caps=Shader exts=SPV_KHR_storage_buffer_storage_class' --convert-gpu-to-spirv %s -o - | FileCheck %s | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What do you think about running |
||||||
|
||||||
module attributes {gpu.container_module} { | ||||||
// CHECK-LABEL: spirv.module @{{.*}} GLSL450 | ||||||
gpu.module @kernels { | ||||||
// CHECK: spirv.func @load_kernel | ||||||
// CHECK-SAME: %[[ARG:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<48 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}) | ||||||
gpu.func @load_kernel(%arg0: memref<12x4xf32>) kernel attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} { | ||||||
%c0 = arith.constant 0 : index | ||||||
// CHECK: %[[PTR:.*]] = spirv.AccessChain %[[ARG]]{{\[}}{{%.*}}, {{%.*}}{{\]}} | ||||||
// CHECK-NEXT: {{%.*}} = spirv.Load "StorageBuffer" %[[PTR]] : f32 | ||||||
%0 = memref.load %arg0[%c0, %c0] : memref<12x4xf32> | ||||||
// CHECK: spirv.Return | ||||||
gpu.return | ||||||
} | ||||||
} | ||||||
} |
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.
Use the actual type here, it's not obvious outside of an IDE: https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable