[Mlir-commits] [mlir] [mlir][spirv] Fix lookup logic `spirv.target_env` for `gpu.module` (PR #147262)

[Jaeho Kim]

================
@@ -48,19 +48,45 @@ 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) {
+  if (ArrayAttr targets = moduleOp.getTargetsAttr()) {
+    for (Attribute targetAttr : targets)
+      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);
----------------
oojahooo wrote:

If I understood correctly, we should test the behavior when the `gpu.module` does not have a `targets` attribute with spirv target env so that the line 78 code is executed.

I believe this case might already be covered by existing tests. For instance, [gpu-to-spirv.mlir](https://github.com/llvm/llvm-project/blob/49d89bc9f43f967051920731f0ec138ca8aaf5ee/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir) tests the behavior when no target env is present at all, and [load-store.mlir](https://github.com/llvm/llvm-project/blob/49d89bc9f43f967051920731f0ec138ca8aaf5ee/mlir/test/Conversion/GPUToSPIRV/load-store.mlir) covers the case where the target env is not attached to the`gpu.module`'s `targets` attribute but to the `module`.

However, if a more specific test is needed, such as one where the `gpu.module` lacks the `targets` attribute but has a target env in its attr-dict, or if other additional tests are necessary, I would be happy to add them based on your feedback.

https://github.com/llvm/llvm-project/pull/147262