1//===- NVVMAttachTarget.cpp - Attach an NVVM target -----------------------===//
2//
3// Part of the LLVM Project, 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 `GpuNVVMAttachTarget` pass, attaching `#nvvm.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/NVVMDialect.h"
18#include "mlir/IR/Builders.h"
19#include "mlir/Pass/Pass.h"
20#include "mlir/Target/LLVM/NVVM/Target.h"
21#include "llvm/Support/Regex.h"
22
23namespace mlir {
24#define GEN_PASS_DEF_GPUNVVMATTACHTARGET
25#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
26} // namespace mlir
27
28using namespace mlir;
29using namespace mlir::NVVM;
30
31namespace {
32struct NVVMAttachTarget
33 : public impl::GpuNVVMAttachTargetBase<NVVMAttachTarget> {
34 using Base::Base;
35
36 DictionaryAttr getFlags(OpBuilder &builder) const;
37
38 void runOnOperation() override;
39
40 void getDependentDialects(DialectRegistry &registry) const override {
41 registry.insert<NVVM::NVVMDialect>();
42 }
43};
44} // namespace
45
46DictionaryAttr NVVMAttachTarget::getFlags(OpBuilder &builder) const {
47 UnitAttr unitAttr = builder.getUnitAttr();
48 SmallVector<NamedAttribute, 3> flags;
49 auto addFlag = [&](StringRef flag) {
50 flags.push_back(Elt: builder.getNamedAttr(name: flag, val: unitAttr));
51 };
52 if (fastFlag)
53 addFlag("fast");
54 if (ftzFlag)
55 addFlag("ftz");
56
57 // Tokenize and set the optional command line options.
58 if (!cmdOptions.empty()) {
59 auto options = gpu::TargetOptions::tokenizeCmdOptions(cmdOptions);
60 if (!options.second.empty()) {
61 llvm::SmallVector<mlir::Attribute> nvvmOptionAttrs;
62 for (const char *opt : options.second) {
63 nvvmOptionAttrs.emplace_back(
64 mlir::StringAttr::get(builder.getContext(), StringRef(opt)));
65 }
66 flags.push_back(builder.getNamedAttr(
67 "ptxas-cmd-options",
68 mlir::ArrayAttr::get(builder.getContext(), nvvmOptionAttrs)));
69 }
70 }
71
72 if (!flags.empty())
73 return builder.getDictionaryAttr(flags);
74 return nullptr;
75}
76
77void NVVMAttachTarget::runOnOperation() {
78 OpBuilder builder(&getContext());
79 ArrayRef<std::string> libs(linkLibs);
80 SmallVector<StringRef> filesToLink(libs);
81 auto target = builder.getAttr<NVVMTargetAttr>(
82 optLevel, triple, chip, features, getFlags(builder),
83 filesToLink.empty() ? nullptr : builder.getStrArrayAttr(filesToLink));
84 llvm::Regex matcher(moduleMatcher);
85 for (Region &region : getOperation()->getRegions())
86 for (Block &block : region.getBlocks())
87 for (auto module : block.getOps<gpu::GPUModuleOp>()) {
88 // Check if the name of the module matches.
89 if (!moduleMatcher.empty() && !matcher.match(module.getName()))
90 continue;
91 // Create the target array.
92 SmallVector<Attribute> targets;
93 if (std::optional<ArrayAttr> attrs = module.getTargets())
94 targets.append(attrs->getValue().begin(), attrs->getValue().end());
95 targets.push_back(target);
96 // Remove any duplicate targets.
97 targets.erase(llvm::unique(targets), targets.end());
98 // Update the target attribute array.
99 module.setTargetsAttr(builder.getArrayAttr(targets));
100 }
101}
102

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