102 lines · cpp
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-exception6//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_GPUNVVMATTACHTARGET25#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"26} // namespace mlir27 28using namespace mlir;29using namespace mlir::NVVM;30 31namespace {32struct NVVMAttachTarget33 : 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 ®istry) const override {41 registry.insert<NVVM::NVVMDialect>();42 }43};44} // namespace45 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(builder.getNamedAttr(flag, 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 ®ion : 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