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

Add lookup target env in "targets" attr logic to GPUToSPIRV pass

Author: oojahooo

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/lib/Dialect/SPIRV/IR/TargetAndABI.cpp

@@ -184,13 +184,6 @@ spirv::TargetEnvAttr spirv::lookupTargetEnv(Operation *op) {
     if (!op)
       break;
 
-    if (auto arrAttr = op->getAttrOfType<ArrayAttr>("targets")) {
-      for (auto attr : arrAttr)
-        if (auto spirvTargetEnvAttr =
-                llvm::dyn_cast<spirv::TargetEnvAttr>(attr))
-          return spirvTargetEnvAttr;
-    }
-
     if (auto attr = op->getAttrOfType<spirv::TargetEnvAttr>(
             spirv::getTargetEnvAttrName()))
       return attr;