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

silee2 · 2025-07-07T18:37:17Z

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

Co-authored-by: Artem Kroviakov artem.kroviakov@intel.com

llvmbot · 2025-07-07T18:37:45Z

@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Sang Ik Lee (silee2)

Changes

Co-authored-by: Artem Kroviakov artem.kroviakov@intel.com

Full diff: https://github.com/llvm/llvm-project/pull/147372.diff

4 Files Affected:

(modified) mlir/include/mlir/Dialect/GPU/Transforms/Passes.td (+34)
(modified) mlir/lib/Dialect/GPU/CMakeLists.txt (+1)
(modified) mlir/test/Dialect/LLVMIR/attach-targets.mlir (+8-4)
(modified) mlir/test/lib/Dialect/GPU/CMakeLists.txt (+1)

diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
index 3766eb16e9429..b6fc0a1375fd8 100644
--- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
+++ b/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
diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt
index 4862d1f722785..6e7d622b910f2 100644
--- a/mlir/lib/Dialect/GPU/CMakeLists.txt
+++ b/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
 
diff --git a/mlir/test/Dialect/LLVMIR/attach-targets.mlir b/mlir/test/Dialect/LLVMIR/attach-targets.mlir
index 83733db400798..09dbcc5f28075 100644
--- a/mlir/test/Dialect/LLVMIR/attach-targets.mlir
+++ b/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">] {
 }
 }
diff --git a/mlir/test/lib/Dialect/GPU/CMakeLists.txt b/mlir/test/lib/Dialect/GPU/CMakeLists.txt
index 4ca5974ed5a49..418c884dc03b3 100644
--- a/mlir/test/lib/Dialect/GPU/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/GPU/CMakeLists.txt
@@ -29,6 +29,7 @@ set(LIBS
   MLIRTranslateLib
   MLIRVectorDialect
   MLIRVectorToLLVMPass
+  MLIRXeVMDialect
   )
 
 add_mlir_library(MLIRGPUTestPasses

silee2 · 2025-07-07T20:39:03Z

@rengolin @adam-smnk @Jianhui-Li

rengolin

A few comments but overall looks good to me, thanks!

I'll let others review and approve.

rengolin · 2025-07-07T21:17:27Z

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

+  auto target = builder.getAttr<mlir::xevm::XeVMTargetAttr>(
+      optLevel, triple, chip, getFlags(builder),
+      filesToLink.empty() ? nullptr : builder.getStrArrayAttr(filesToLink));
+  llvm::Regex matcher(moduleMatcher);


If the name is empty, do you just attach to all GPU modules?

I didn't get your question.
filesToLink is files to link, like external spir-v library files. Name would be empty in most cases.
Behavior is the same as nvvm-attach-target.

If the name is empty, do you just attach to all GPU modules?

I think so.
When moduleMatcher is empty, nvvm info is attached to all GPU module.

Now after updates, this should be the behavior here as well.

rengolin · 2025-07-07T21:20:48Z

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

+      filesToLink.empty() ? nullptr : builder.getStrArrayAttr(filesToLink));
+  llvm::Regex matcher(moduleMatcher);
+  // Check if the name of the module matches.
+  auto gpuModule = cast<gpu::GPUModuleOp>(getOperation());


Both NV and ROC similar methods iterate through regions and basic blocks, but your definition above specified mlir::gpu::GPUModuleOp. Is that because it's assumed that this will only run on the GPU modules (and thus needs no search)? Or are we just not covering the general case in the tests?

Both NV and ROC also assume that the attachment logic will only be applied to GPU modules:

llvm-project/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp

Line 87 in df6ae45

for (auto module : block.getOps<gpu::GPUModuleOp>()) {

as indicated by their pass description itself:

llvm-project/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

Line 112 in a1ea9e6

This pass searches for all GPU Modules in the immediate regions and attaches

This PR decided to rely on the MLIR pass infrastructure for traversal.

If you think we should follow nv/roc, please consider responding to the comment below.

Updated logic to match nvvm, rocdl and other attach-target passes.

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

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

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

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

…. Update comment.

charithaintc

LGTM

charithaintc · 2025-07-09T20:19:01Z

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

+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<mlir::xevm::XeVMDialect>();


I believe no need to use explicit namespaces because they already imported?

Removed mlir::

charithaintc · 2025-07-09T20:21:13Z

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

+  SmallVector<NamedAttribute, 3> flags;
+  // Tokenize and set the optional command line options.
+  if (!cmdOptions.empty()) {
+    auto options = gpu::TargetOptions::tokenizeCmdOptions(cmdOptions);


please mention the type and avoid using auto because output type is not obvious in the context (not a cast).

Refer to LLVM coding manuals for auto usage.

Replaced auto with actual type std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>

chencha3

LGTM. Thanks for your efforts

Add xevm-attach-target transform pass.

f2fd80a

Co-authored-by: Artem Kroviakov artem.kroviakov@intel.com

llvmbot added mlir:llvm mlir:gpu mlir labels Jul 7, 2025

silee2 added 2 commits July 7, 2025 18:42

Add missing file.

b82816a

Add XeVM dialect as a dependency for xevm-attach-target pass.

62db26d

chencha3 mentioned this pull request Jul 7, 2025

[XeGPU] uArch definition (PR 1/N) mshahneo/llvm-project#2

Open

rengolin reviewed Jul 7, 2025

View reviewed changes

rengolin requested review from adam-smnk and Jianhui-Li July 7, 2025 21:22

adam-smnk reviewed Jul 8, 2025

View reviewed changes

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

39b3471

…. Update comment.

adam-smnk approved these changes Jul 9, 2025

View reviewed changes

charithaintc approved these changes Jul 9, 2025

View reviewed changes

Address reviewer comments.

17f139a

chencha3 approved these changes Jul 10, 2025

View reviewed changes

Jianhui-Li approved these changes Jul 10, 2025

View reviewed changes

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

Are you sure you want to change the base?

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

Conversation

silee2 commented Jul 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Jul 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

silee2 commented Jul 7, 2025

Uh oh!

rengolin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

charithaintc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chencha3 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

silee2 commented Jul 7, 2025 •

edited

Loading

llvmbot commented Jul 7, 2025 •

edited

Loading