Match pass behavior with other attach target passes. Split check line. Update comment.

Author: silee2

mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

@@ -258,7 +258,7 @@ def GpuSPIRVAttachTarget: Pass<"spirv-attach-target", ""> {
   ];
 }
 
-def GpuXeVMAttachTarget : Pass<"xevm-attach-target", "mlir::gpu::GPUModuleOp"> {
+def GpuXeVMAttachTarget : Pass<"xevm-attach-target", ""> {
   let summary = "Attaches a XeVM target attribute to a GPU Module.";
   let description = [{
     This pass searches for all GPU Modules in the immediate regions and attaches

mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp

@@ -1,4 +1,4 @@
-//===-- XeVMAttachTarget.cpp - DESC -----------------------------*- C++ -*-===//
+//===-- XeVMAttachTarget.cpp - Attach an XeVM target ----------------------===//
 //
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -68,21 +68,24 @@ void XeVMAttachTarget::runOnOperation() {
   OpBuilder builder(&getContext());
   ArrayRef<std::string> libs(linkLibs);
   SmallVector<StringRef> filesToLink(libs);
-  auto target = builder.getAttr<mlir::xevm::XeVMTargetAttr>(
+  auto target = builder.getAttr<xevm::XeVMTargetAttr>(
       optLevel, triple, chip, getFlags(builder),
       filesToLink.empty() ? nullptr : builder.getStrArrayAttr(filesToLink));
   llvm::Regex matcher(moduleMatcher);
-  // Check if the name of the module matches.
-  auto gpuModule = cast<gpu::GPUModuleOp>(getOperation());
-  if (!moduleMatcher.empty() && !matcher.match(gpuModule.getName()))
-    return;
-  // Create the target array.
-  SmallVector<Attribute> targets;
-  if (std::optional<ArrayAttr> attrs = gpuModule.getTargets())
-    targets.append(attrs->getValue().begin(), attrs->getValue().end());
-  targets.push_back(target);
-  // Remove any duplicate targets.
-  targets.erase(llvm::unique(targets), targets.end());
-  // Update the target attribute array.
-  gpuModule.setTargetsAttr(builder.getArrayAttr(targets));
+  for (Region &region : getOperation()->getRegions())
+    for (Block &block : region.getBlocks())
+      for (auto module : block.getOps<gpu::GPUModuleOp>()) {
+        // Check if the name of the module matches.
+        if (!moduleMatcher.empty() && !matcher.match(module.getName()))
+          continue;
+        // Create the target array.
+        SmallVector<Attribute> targets;
+        if (std::optional<ArrayAttr> attrs = module.getTargets())
+          targets.append(attrs->getValue().begin(), attrs->getValue().end());
+        targets.push_back(target);
+        // Remove any duplicate targets.
+        targets.erase(llvm::unique(targets), targets.end());
+        // Update the target attribute array.
+        module.setTargetsAttr(builder.getArrayAttr(targets));
+      }
 }

mlir/test/Dialect/LLVMIR/attach-targets.mlir

@@ -23,11 +23,16 @@ gpu.module @rocdl_module {
 gpu.module @xevm_module {
 }
 // Check the options were added.
-// CHECK_OPTS: @options_module_1 [#nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>, #xevm.target<O = 1, chip = "pvc">]  {
+// CHECK_OPTS: @options_module_1 [#nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>,
+// CHECK_OPTS-SAME: #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>,
+// CHECK_OPTS-SAME: #xevm.target<O = 1, chip = "pvc">]  {
 gpu.module @options_module_1 {
 }
 // Check the options were added and that the first target was preserved.
-// CHECK_OPTS: @options_module_2 [#nvvm.target<O = 3, chip = "sm_90">, #nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>, #xevm.target<O = 1, chip = "pvc">]  {
+// CHECK_OPTS: @options_module_2 [#nvvm.target<O = 3, chip = "sm_90">,
+// CHECK_OPTS-SAME: #nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>,
+// CHECK_OPTS-SAME: #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>,
+// CHECK_OPTS-SAME: #xevm.target<O = 1, chip = "pvc">]  {
 gpu.module @options_module_2 [#nvvm.target<O = 3, chip = "sm_90">] {
 }
 }