[MLIR][GPU] Add xevm-attach-target transform pass. (#147372)

Add xevm-attach-target transform pass and unit-tests.

Co-authored-by: by Sang Ik Lee sang.ik.lee@intel.com.
Co-authored-by: Artem Kroviakov artem.kroviakov@intel.com

Author: silee2

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

@@ -258,4 +258,38 @@ def GpuSPIRVAttachTarget: Pass<"spirv-attach-target", ""> {
   ];
 }
 
+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
+    a XeVM target if the module matches the name specified by the `module` argument.
+
+    Example:
+    ```
+    // File: in.mlir:
+    gpu.module @nvvm_module_1 {...}
+    gpu.module @rocdl_module_2 {...}
+    gpu.module @xevm_module_3 {...}
+    // mlir-opt --xevm-attach-target="module=xevm.* chip=pvc" in.mlir
+    gpu.module @nvvm_module_1 {...}
+    gpu.module @rocdl_module_2 {...}
+    gpu.module @xevm_module_3 [#xevm.target<chip = "pvc">] {...}
+    ```
+  }];
+  let options =
+      [Option<"moduleMatcher", "module", "std::string",
+              /*default=*/[{""}],
+              "Regex used to identify the modules to attach the target to.">,
+       Option<"triple", "triple", "std::string",
+              /*default=*/"\"spirv64-unknown-unknown\"", "Target triple.">,
+       Option<"chip", "chip", "std::string",
+              /*default=*/"\"bmg\"", "Target chip.">,
+       Option<"optLevel", "O", "unsigned",
+              /*default=*/"2", "Optimization level.">,
+       ListOption<"linkLibs", "l", "std::string",
+                  "Extra bitcode libraries paths to link to.">,
+       Option<"cmdOptions", "cmd-options", "std::string",
+              /*default=*/[{""}],
+              "Command line options passed to downstream compiler">];
+}
 #endif // MLIR_DIALECT_GPU_PASSES

mlir/lib/Dialect/GPU/CMakeLists.txt

@@ -44,6 +44,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/ShuffleRewriter.cpp
   Transforms/SubgroupIdRewriter.cpp
   Transforms/SubgroupReduceLowering.cpp
+  Transforms/XeVMAttachTarget.cpp
 
   OBJECT
 
@@ -78,6 +79,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
   MLIRSupport
   MLIRTransformUtils
   MLIRVectorDialect
+  MLIRXeVMDialect
   )
 
 add_subdirectory(TransformOps)

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

@@ -0,0 +1,92 @@
+//===-- 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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the `GpuXeVMAttachTarget` pass, attaching `#xevm.target`
+// attributes to GPU modules.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/Regex.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_GPUXEVMATTACHTARGET
+#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::xevm;
+
+namespace {
+struct XeVMAttachTarget
+    : public mlir::impl::GpuXeVMAttachTargetBase<XeVMAttachTarget> {
+  using Base::Base;
+
+  DictionaryAttr getFlags(OpBuilder &builder) const;
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<xevm::XeVMDialect>();
+  }
+};
+} // namespace
+
+DictionaryAttr XeVMAttachTarget::getFlags(OpBuilder &builder) const {
+  SmallVector<NamedAttribute, 3> flags;
+  // Tokenize and set the optional command line options.
+  if (!cmdOptions.empty()) {
+    std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>> options =
+        gpu::TargetOptions::tokenizeCmdOptions(cmdOptions);
+    if (!options.second.empty()) {
+      llvm::SmallVector<mlir::Attribute> xevmOptionAttrs;
+      for (const char *opt : options.second) {
+        xevmOptionAttrs.emplace_back(
+            mlir::StringAttr::get(builder.getContext(), StringRef(opt)));
+      }
+      flags.push_back(builder.getNamedAttr(
+          "cmd-options",
+          mlir::ArrayAttr::get(builder.getContext(), xevmOptionAttrs)));
+    }
+  }
+
+  if (!flags.empty())
+    return builder.getDictionaryAttr(flags);
+  return nullptr;
+}
+
+void XeVMAttachTarget::runOnOperation() {
+  OpBuilder builder(&getContext());
+  ArrayRef<std::string> libs(linkLibs);
+  SmallVector<StringRef> filesToLink(libs);
+  auto target = builder.getAttr<xevm::XeVMTargetAttr>(
+      optLevel, triple, chip, getFlags(builder),
+      filesToLink.empty() ? nullptr : builder.getStrArrayAttr(filesToLink));
+  llvm::Regex matcher(moduleMatcher);
+  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

@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s --nvvm-attach-target='module=nvvm.* O=3 chip=sm_90' --rocdl-attach-target='module=rocdl.* O=3 chip=gfx90a' | FileCheck %s
-// RUN: mlir-opt %s --nvvm-attach-target='module=options.* O=1 chip=sm_70 fast=true ftz=true' --rocdl-attach-target='module=options.* l=file1.bc,file2.bc wave64=false finite-only=true' | FileCheck %s --check-prefix=CHECK_OPTS
+// RUN: mlir-opt %s --nvvm-attach-target='module=nvvm.* O=3 chip=sm_90' --rocdl-attach-target='module=rocdl.* O=3 chip=gfx90a' --xevm-attach-target='module=xevm.* O=3 chip=pvc' | FileCheck %s
+// RUN: mlir-opt %s --nvvm-attach-target='module=options.* O=1 chip=sm_70 fast=true ftz=true' --rocdl-attach-target='module=options.* l=file1.bc,file2.bc wave64=false finite-only=true' --xevm-attach-target='module=options.* O=1 chip=pvc' | FileCheck %s --check-prefix=CHECK_OPTS
 
 module attributes {gpu.container_module} {
 // Verify the target is appended.
@@ -18,12 +18,21 @@ gpu.module @nvvm_module_3 [#nvvm.target<O = 3, chip = "sm_90">] {
 // CHECK: @rocdl_module [#rocdl.target<O = 3, chip = "gfx90a">] {
 gpu.module @rocdl_module {
 }
+// Verify that other targets are not added as they fail to match the regex, but XeVM does get appended.
+// CHECK: @xevm_module [#xevm.target<O = 3, chip = "pvc">] {
+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"]>]  {
+// 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"]>]  {
+// 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">] {
 }
 }

mlir/test/lib/Dialect/GPU/CMakeLists.txt

@@ -29,6 +29,7 @@ set(LIBS
   MLIRTranslateLib
   MLIRVectorDialect
   MLIRVectorToLLVMPass
+  MLIRXeVMDialect
   )
 
 add_mlir_library(MLIRGPUTestPasses