Add xevm-attach-target transform pass.

Co-authored-by: Artem Kroviakov [email protected]

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", "mlir::gpu::GPUModuleOp"> {
+  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
 

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,16 @@ 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}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>, #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">, #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">]  {
 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