[mlir][spirv] Fix lookup logic spirv.target_env for gpu.module

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

@@ -385,6 +385,10 @@ LogicalResult GPUModuleConversion::matchAndRewrite(
   if (auto attr = moduleOp->getAttrOfType<spirv::TargetEnvAttr>(
           spirv::getTargetEnvAttrName()))
     spvModule->setAttr(spirv::getTargetEnvAttrName(), attr);
+  for (auto targetAttr : moduleOp.getTargetsAttr())
+    if (auto spirvTargetEnvAttr =
+            llvm::dyn_cast<spirv::TargetEnvAttr>(targetAttr))
+      spvModule->setAttr(spirv::getTargetEnvAttrName(), spirvTargetEnvAttr);
 
   rewriter.eraseOp(moduleOp);
   return success();

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp

@@ -48,19 +48,38 @@ struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase<GPUToSPIRVPass> {
   void runOnOperation() override;
 
 private:
+  spirv::TargetEnvAttr lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp);
+  spirv::TargetEnvAttr lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp);
   bool mapMemorySpace;
 };
 
+spirv::TargetEnvAttr
+GPUToSPIRVPass::lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp) {
+  for (auto &targetAttr : moduleOp.getTargetsAttr())
+    if (auto spirvTargetEnvAttr =
+            llvm::dyn_cast<spirv::TargetEnvAttr>(targetAttr))
+      return spirvTargetEnvAttr;
+
+  return {};
+}
+
+spirv::TargetEnvAttr
+GPUToSPIRVPass::lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp) {
+  if (auto 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);
   };
@@ -86,7 +105,7 @@ void GPUToSPIRVPass::runOnOperation() {
   // TargetEnv attributes.
   for (Operation *gpuModule : gpuModules) {
     spirv::TargetEnvAttr targetAttr =
-        spirv::lookupTargetEnvOrDefault(gpuModule);
+        lookupTargetEnvOrDefault(llvm::cast<gpu::GPUModuleOp>(gpuModule));
 
     // Map MemRef memory space to SPIR-V storage class first if requested.
     if (mapMemorySpace) {

mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir

@@ -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
+
+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
+    }
+  }
+}