diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td index 3766eb16e9429..187ac9aa18aac 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", ""> { + 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..f2f010a771b77 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 @@ -78,6 +79,7 @@ add_mlir_dialect_library(MLIRGPUTransforms MLIRSupport MLIRTransformUtils MLIRVectorDialect + MLIRXeVMDialect ) add_subdirectory(TransformOps) diff --git a/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp b/mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp new file mode 100644 index 0000000000000..e9cf4939a13b8 --- /dev/null +++ b/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 { + 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()) { + std::pair> 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); + 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 83733db400798..d1112f7411aae 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,21 @@ 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, +// 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] { +// 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] { } } 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