Skip to content

[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

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open
3 changes: 3 additions & 0 deletions mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -385,6 +385,9 @@ LogicalResult GPUModuleConversion::matchAndRewrite(
if (auto attr = moduleOp->getAttrOfType<spirv::TargetEnvAttr>(
spirv::getTargetEnvAttrName()))
spvModule->setAttr(spirv::getTargetEnvAttrName(), attr);
for (const Attribute &targetAttr : moduleOp.getTargetsAttr())
if (auto spirvTargetEnvAttr = dyn_cast<spirv::TargetEnvAttr>(targetAttr))
spvModule->setAttr(spirv::getTargetEnvAttrName(), spirvTargetEnvAttr);

rewriter.eraseOp(moduleOp);
return success();
Expand Down
32 changes: 28 additions & 4 deletions mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -48,19 +48,43 @@ struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase<GPUToSPIRVPass> {
void runOnOperation() override;

private:
/// Queries the target environment from 'targets' attribute of the given
/// `moduleOp`.
spirv::TargetEnvAttr lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp);

/// Queries the target environment from 'targets' attribute of the given
/// `moduleOp` or returns target environment as returned by
/// `spirv::lookupTargetEnvOrDefault` if not provided by 'targets'.
spirv::TargetEnvAttr lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp);
bool mapMemorySpace;
};

spirv::TargetEnvAttr
GPUToSPIRVPass::lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp) {
for (const Attribute &targetAttr : moduleOp.getTargetsAttr())
if (auto spirvTargetEnvAttr = dyn_cast<spirv::TargetEnvAttr>(targetAttr))
return spirvTargetEnvAttr;

return {};
}

spirv::TargetEnvAttr
GPUToSPIRVPass::lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp) {
if (spirv::TargetEnvAttr targetEnvAttr = lookupTargetEnvInTargets(moduleOp))
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);
};
Expand All @@ -86,7 +110,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) {
Expand Down
17 changes: 17 additions & 0 deletions mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
// RUN: mlir-opt --convert-gpu-to-spirv %s | FileCheck %s

module attributes {gpu.container_module} {
// CHECK-LABEL: spirv.module @{{.*}} GLSL450
gpu.module @kernels [#spirv.target_env<#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>] {
// 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
}
}
}