1//===-- XeVMAttachTarget.cpp - Attach an XeVM target ----------------------===//
2//
3// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements the `GpuXeVMAttachTarget` pass, attaching `#xevm.target`
10// attributes to GPU modules.
11//
12//===----------------------------------------------------------------------===//
13
14#include "mlir/Dialect/GPU/Transforms/Passes.h"
15
16#include "mlir/Dialect/GPU/IR/GPUDialect.h"
17#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
18#include "mlir/IR/Builders.h"
19#include "mlir/Pass/Pass.h"
20#include "llvm/Support/Regex.h"
21
22namespace mlir {
23#define GEN_PASS_DEF_GPUXEVMATTACHTARGET
24#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
25} // namespace mlir
26
27using namespace mlir;
28using namespace mlir::xevm;
29
30namespace {
31struct XeVMAttachTarget
32 : public mlir::impl::GpuXeVMAttachTargetBase<XeVMAttachTarget> {
33 using Base::Base;
34
35 DictionaryAttr getFlags(OpBuilder &builder) const;
36
37 void runOnOperation() override;
38
39 void getDependentDialects(DialectRegistry &registry) const override {
40 registry.insert<xevm::XeVMDialect>();
41 }
42};
43} // namespace
44
45DictionaryAttr XeVMAttachTarget::getFlags(OpBuilder &builder) const {
46 SmallVector<NamedAttribute, 3> flags;
47 // Tokenize and set the optional command line options.
48 if (!cmdOptions.empty()) {
49 std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>> options =
50 gpu::TargetOptions::tokenizeCmdOptions(cmdOptions);
51 if (!options.second.empty()) {
52 llvm::SmallVector<mlir::Attribute> xevmOptionAttrs;
53 for (const char *opt : options.second) {
54 xevmOptionAttrs.emplace_back(
55 Args: mlir::StringAttr::get(context: builder.getContext(), bytes: StringRef(opt)));
56 }
57 flags.push_back(Elt: builder.getNamedAttr(
58 name: "cmd-options",
59 val: mlir::ArrayAttr::get(context: builder.getContext(), value: xevmOptionAttrs)));
60 }
61 }
62
63 if (!flags.empty())
64 return builder.getDictionaryAttr(value: flags);
65 return nullptr;
66}
67
68void XeVMAttachTarget::runOnOperation() {
69 OpBuilder builder(&getContext());
70 ArrayRef<std::string> libs(linkLibs);
71 SmallVector<StringRef> filesToLink(libs);
72 auto target = builder.getAttr<xevm::XeVMTargetAttr>(
73 args&: optLevel, args&: triple, args&: chip, args: getFlags(builder),
74 args: filesToLink.empty() ? nullptr : builder.getStrArrayAttr(values: filesToLink));
75 llvm::Regex matcher(moduleMatcher);
76 for (Region &region : getOperation()->getRegions())
77 for (Block &block : region.getBlocks())
78 for (auto module : block.getOps<gpu::GPUModuleOp>()) {
79 // Check if the name of the module matches.
80 if (!moduleMatcher.empty() && !matcher.match(String: module.getName()))
81 continue;
82 // Create the target array.
83 SmallVector<Attribute> targets;
84 if (std::optional<ArrayAttr> attrs = module.getTargets())
85 targets.append(in_start: attrs->getValue().begin(), in_end: attrs->getValue().end());
86 targets.push_back(Elt: target);
87 // Remove any duplicate targets.
88 targets.erase(CS: llvm::unique(R&: targets), CE: targets.end());
89 // Update the target attribute array.
90 module.setTargetsAttr(builder.getArrayAttr(value: targets));
91 }
92}
93

source code of mlir/lib/Dialect/GPU/Transforms/XeVMAttachTarget.cpp