From 8a75d5b9578a44a0df51df2967d11a2f971ee781 Mon Sep 17 00:00:00 2001 From: "Lee, Sang Ik" Date: Mon, 7 Jul 2025 18:35:00 +0000 Subject: [PATCH 1/5] Add xevm-attach-target transform pass. Co-authored-by: Artem Kroviakov artem.kroviakov@intel.com --- .../mlir/Dialect/GPU/Transforms/Passes.td | 34 +++++++++++++++++++ mlir/lib/Dialect/GPU/CMakeLists.txt | 1 + mlir/test/Dialect/LLVMIR/attach-targets.mlir | 12 ++++--- mlir/test/lib/Dialect/GPU/CMakeLists.txt | 1 + 4 files changed, 44 insertions(+), 4 deletions(-) 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] {...} + ``` + }]; + 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] { // CHECK: @rocdl_module [#rocdl.target] { 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] { +gpu.module @xevm_module { +} // Check the options were added. -// CHECK_OPTS: @options_module_1 [#nvvm.target, #rocdl.target] { +// CHECK_OPTS: @options_module_1 [#nvvm.target, #rocdl.target, #xevm.target] { gpu.module @options_module_1 { } // Check the options were added and that the first target was preserved. -// CHECK_OPTS: @options_module_2 [#nvvm.target, #nvvm.target, #rocdl.target] { +// CHECK_OPTS: @options_module_2 [#nvvm.target, #nvvm.target, #rocdl.target, #xevm.target] { gpu.module @options_module_2 [#nvvm.target] { } } 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 From 75cff5f2da8b3bbc02024bbcf4dd4d1c7f1afea3 Mon Sep 17 00:00:00 2001 From: "Lee, Sang Ik" Date: Mon, 7 Jul 2025 18:42:08 +0000 Subject: [PATCH 2/5] Add missing file. --- .../GPU/Transforms/XeVMAttachTarget.cpp | 88 +++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp diff --git a/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp b/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp new file mode 100644 index 0000000000000..eb84e07479696 --- /dev/null +++ b/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp @@ -0,0 +1,88 @@ +//===-- XeVMAttachTarget.cpp - DESC -----------------------------*- C++ -*-===// +// +// 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 { + using Base::Base; + + DictionaryAttr getFlags(OpBuilder &builder) const; + + void runOnOperation() override; + + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } +}; +} // namespace + +DictionaryAttr XeVMAttachTarget::getFlags(OpBuilder &builder) const { + SmallVector flags; + // Tokenize and set the optional command line options. + if (!cmdOptions.empty()) { + auto options = gpu::TargetOptions::tokenizeCmdOptions(cmdOptions); + if (!options.second.empty()) { + llvm::SmallVector 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 libs(linkLibs); + SmallVector filesToLink(libs); + auto target = builder.getAttr( + 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(getOperation()); + if (!moduleMatcher.empty() && !matcher.match(gpuModule.getName())) + return; + // Create the target array. + SmallVector targets; + if (std::optional 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)); +} From 36edd3f4896848e75791805330ad19fbea43a8af Mon Sep 17 00:00:00 2001 From: "Lee, Sang Ik" Date: Mon, 7 Jul 2025 20:36:53 +0000 Subject: [PATCH 3/5] Add XeVM dialect as a dependency for xevm-attach-target pass. --- mlir/lib/Dialect/GPU/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt index 6e7d622b910f2..f2f010a771b77 100644 --- a/mlir/lib/Dialect/GPU/CMakeLists.txt +++ b/mlir/lib/Dialect/GPU/CMakeLists.txt @@ -79,6 +79,7 @@ add_mlir_dialect_library(MLIRGPUTransforms MLIRSupport MLIRTransformUtils MLIRVectorDialect + MLIRXeVMDialect ) add_subdirectory(TransformOps) From 2abf74d7efba8e2c7337d548f2ecbc22a8886702 Mon Sep 17 00:00:00 2001 From: "Lee, Sang Ik" Date: Wed, 9 Jul 2025 02:56:54 +0000 Subject: [PATCH 4/5] Match pass behavior with other attach target passes. Split check line. Update comment. --- .../mlir/Dialect/GPU/Transforms/Passes.td | 2 +- .../GPU/Transforms/XeVMAttachTarget.cpp | 33 ++++++++++--------- mlir/test/Dialect/LLVMIR/attach-targets.mlir | 9 +++-- 3 files changed, 26 insertions(+), 18 deletions(-) diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td index b6fc0a1375fd8..187ac9aa18aac 100644 --- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td +++ b/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 diff --git a/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp b/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp index eb84e07479696..bdbae80b2a6ce 100644 --- a/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp +++ b/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 libs(linkLibs); SmallVector filesToLink(libs); - auto target = builder.getAttr( + auto target = builder.getAttr( 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(getOperation()); - if (!moduleMatcher.empty() && !matcher.match(gpuModule.getName())) - return; - // Create the target array. - SmallVector targets; - if (std::optional 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 ®ion : getOperation()->getRegions()) + for (Block &block : region.getBlocks()) + for (auto module : block.getOps()) { + // Check if the name of the module matches. + if (!moduleMatcher.empty() && !matcher.match(module.getName())) + continue; + // Create the target array. + SmallVector targets; + if (std::optional 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)); + } } diff --git a/mlir/test/Dialect/LLVMIR/attach-targets.mlir b/mlir/test/Dialect/LLVMIR/attach-targets.mlir index 09dbcc5f28075..d1112f7411aae 100644 --- a/mlir/test/Dialect/LLVMIR/attach-targets.mlir +++ b/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, #rocdl.target, #xevm.target] { +// CHECK_OPTS: @options_module_1 [#nvvm.target, +// CHECK_OPTS-SAME: #rocdl.target, +// CHECK_OPTS-SAME: #xevm.target] { gpu.module @options_module_1 { } // Check the options were added and that the first target was preserved. -// CHECK_OPTS: @options_module_2 [#nvvm.target, #nvvm.target, #rocdl.target, #xevm.target] { +// CHECK_OPTS: @options_module_2 [#nvvm.target, +// CHECK_OPTS-SAME: #nvvm.target, +// CHECK_OPTS-SAME: #rocdl.target, +// CHECK_OPTS-SAME: #xevm.target] { gpu.module @options_module_2 [#nvvm.target] { } } From 07282199b551061e25c429669eba59255225b73d Mon Sep 17 00:00:00 2001 From: "Lee, Sang Ik" Date: Wed, 9 Jul 2025 21:14:59 +0000 Subject: [PATCH 5/5] Address reviewer comments. --- mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp b/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp index bdbae80b2a6ce..e9cf4939a13b8 100644 --- a/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp +++ b/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp @@ -37,7 +37,7 @@ struct XeVMAttachTarget void runOnOperation() override; void getDependentDialects(DialectRegistry ®istry) const override { - registry.insert(); + registry.insert(); } }; } // namespace @@ -46,7 +46,8 @@ DictionaryAttr XeVMAttachTarget::getFlags(OpBuilder &builder) const { SmallVector flags; // Tokenize and set the optional command line options. if (!cmdOptions.empty()) { - auto options = gpu::TargetOptions::tokenizeCmdOptions(cmdOptions); + std::pair> options = + gpu::TargetOptions::tokenizeCmdOptions(cmdOptions); if (!options.second.empty()) { llvm::SmallVector xevmOptionAttrs; for (const char *opt : options.second) {