1//===- NVVMProperties.cpp - NVVM annotation utilities ---------------------===//
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 contains NVVM attribute and metadata query utilities.
10//
11//===----------------------------------------------------------------------===//
12
13#include "NVVMProperties.h"
14#include "llvm/ADT/ArrayRef.h"
15#include "llvm/ADT/STLExtras.h"
16#include "llvm/ADT/StringRef.h"
17#include "llvm/IR/Argument.h"
18#include "llvm/IR/Constants.h"
19#include "llvm/IR/Function.h"
20#include "llvm/IR/GlobalValue.h"
21#include "llvm/IR/GlobalVariable.h"
22#include "llvm/IR/Instructions.h"
23#include "llvm/IR/Metadata.h"
24#include "llvm/IR/Module.h"
25#include "llvm/Support/ModRef.h"
26#include "llvm/Support/Mutex.h"
27#include "llvm/Support/NVVMAttributes.h"
28#include <functional>
29#include <map>
30#include <mutex>
31#include <numeric>
32#include <string>
33#include <vector>
34
35using namespace llvm;
36
37namespace {
38using AnnotationValues = std::map<std::string, std::vector<unsigned>>;
39using AnnotationMap = std::map<const GlobalValue *, AnnotationValues>;
40
41struct AnnotationCache {
42 sys::Mutex Lock;
43 std::map<const Module *, AnnotationMap> Cache;
44};
45
46} // namespace
47
48static AnnotationCache &getAnnotationCache() {
49 static AnnotationCache AC;
50 return AC;
51}
52
53// TODO: Replace these legacy nvvm.annotations metadata names with proper
54// function/parameter attributes (like the NVVMAttr:: constants).
55namespace NVVMMetadata {
56constexpr StringLiteral Texture("texture");
57constexpr StringLiteral Surface("surface");
58constexpr StringLiteral Sampler("sampler");
59constexpr StringLiteral ReadOnlyImage("rdoimage");
60constexpr StringLiteral WriteOnlyImage("wroimage");
61constexpr StringLiteral ReadWriteImage("rdwrimage");
62constexpr StringLiteral Managed("managed");
63} // namespace NVVMMetadata
64
65void llvm::clearAnnotationCache(const Module *Mod) {
66 auto &AC = getAnnotationCache();
67 std::lock_guard<sys::Mutex> Guard(AC.Lock);
68 AC.Cache.erase(x: Mod);
69}
70
71static void cacheAnnotationFromMD(const MDNode *MetadataNode,
72 AnnotationValues &RetVal) {
73 auto &AC = getAnnotationCache();
74 std::lock_guard<sys::Mutex> Guard(AC.Lock);
75 assert(MetadataNode && "Invalid mdnode for annotation");
76 assert((MetadataNode->getNumOperands() % 2) == 1 &&
77 "Invalid number of operands");
78 // start index = 1, to skip the global variable key
79 // increment = 2, to skip the value for each property-value pairs
80 for (unsigned I = 1, E = MetadataNode->getNumOperands(); I != E; I += 2) {
81 const MDString *Prop = dyn_cast<MDString>(Val: MetadataNode->getOperand(I));
82 assert(Prop && "Annotation property not a string");
83 std::string Key = Prop->getString().str();
84
85 if (ConstantInt *Val = mdconst::dyn_extract<ConstantInt>(
86 MD: MetadataNode->getOperand(I: I + 1))) {
87 RetVal[Key].push_back(x: Val->getZExtValue());
88 } else {
89 llvm_unreachable("Value operand not a constant int");
90 }
91 }
92}
93
94static void cacheAnnotationFromMD(const Module *M, const GlobalValue *GV) {
95 auto &AC = getAnnotationCache();
96 std::lock_guard<sys::Mutex> Guard(AC.Lock);
97 NamedMDNode *NMD = M->getNamedMetadata(Name: "nvvm.annotations");
98 if (!NMD)
99 return;
100
101 AnnotationValues Tmp;
102 for (unsigned I = 0, E = NMD->getNumOperands(); I != E; ++I) {
103 const MDNode *Elem = NMD->getOperand(i: I);
104 GlobalValue *Entity =
105 mdconst::dyn_extract_or_null<GlobalValue>(MD: Elem->getOperand(I: 0));
106 // entity may be null due to DCE
107 if (!Entity || Entity != GV)
108 continue;
109
110 cacheAnnotationFromMD(MetadataNode: Elem, RetVal&: Tmp);
111 }
112
113 if (Tmp.empty())
114 return;
115
116 AC.Cache[M][GV] = std::move(Tmp);
117}
118
119static std::optional<unsigned> findOneNVVMAnnotation(const GlobalValue *GV,
120 StringRef Prop) {
121 auto &AC = getAnnotationCache();
122 std::lock_guard<sys::Mutex> Guard(AC.Lock);
123 const Module *M = GV->getParent();
124 auto ACIt = AC.Cache.find(x: M);
125 if (ACIt == AC.Cache.end())
126 cacheAnnotationFromMD(M, GV);
127 else if (ACIt->second.find(x: GV) == ACIt->second.end())
128 cacheAnnotationFromMD(M, GV);
129
130 auto &KVP = AC.Cache[M][GV];
131 auto It = KVP.find(x: Prop.str());
132 if (It == KVP.end())
133 return std::nullopt;
134 return It->second[0];
135}
136
137static bool findAllNVVMAnnotation(const GlobalValue *GV, StringRef Prop,
138 std::vector<unsigned> &RetVal) {
139 auto &AC = getAnnotationCache();
140 std::lock_guard<sys::Mutex> Guard(AC.Lock);
141 const Module *M = GV->getParent();
142 auto ACIt = AC.Cache.find(x: M);
143 if (ACIt == AC.Cache.end())
144 cacheAnnotationFromMD(M, GV);
145 else if (ACIt->second.find(x: GV) == ACIt->second.end())
146 cacheAnnotationFromMD(M, GV);
147
148 auto &KVP = AC.Cache[M][GV];
149 auto It = KVP.find(x: Prop.str());
150 if (It == KVP.end())
151 return false;
152 RetVal = It->second;
153 return true;
154}
155
156static bool globalHasNVVMAnnotation(const Value &V, StringRef Prop) {
157 if (const auto *GV = dyn_cast<GlobalValue>(Val: &V))
158 if (const auto Annot = findOneNVVMAnnotation(GV, Prop)) {
159 assert((*Annot == 1) && "Unexpected annotation on a symbol");
160 return true;
161 }
162
163 return false;
164}
165
166static bool argHasNVVMAnnotation(const Value &Val, StringRef Annotation) {
167 if (const auto *Arg = dyn_cast<Argument>(Val: &Val)) {
168 std::vector<unsigned> Annot;
169 if (findAllNVVMAnnotation(GV: Arg->getParent(), Prop: Annotation, RetVal&: Annot) &&
170 is_contained(Range&: Annot, Element: Arg->getArgNo()))
171 return true;
172 }
173 return false;
174}
175
176static std::optional<unsigned> getFnAttrParsedInt(const Function &F,
177 StringRef Attr) {
178 return F.hasFnAttribute(Kind: Attr)
179 ? std::optional(F.getFnAttributeAsParsedInteger(Kind: Attr))
180 : std::nullopt;
181}
182
183static SmallVector<unsigned, 3> getFnAttrParsedVector(const Function &F,
184 StringRef Attr) {
185 SmallVector<unsigned, 3> V;
186 auto &Ctx = F.getContext();
187
188 if (F.hasFnAttribute(Kind: Attr)) {
189 // We expect the attribute value to be of the form "x[,y[,z]]", where x, y,
190 // and z are unsigned values.
191 StringRef S = F.getFnAttribute(Kind: Attr).getValueAsString();
192 for (unsigned I = 0; I < 3 && !S.empty(); I++) {
193 auto [First, Rest] = S.split(Separator: ",");
194 unsigned IntVal;
195 if (First.trim().getAsInteger(Radix: 0, Result&: IntVal))
196 Ctx.emitError(ErrorStr: "can't parse integer attribute " + First + " in " + Attr);
197
198 V.push_back(Elt: IntVal);
199 S = Rest;
200 }
201 }
202 return V;
203}
204
205static std::optional<uint64_t> getVectorProduct(ArrayRef<unsigned> V) {
206 if (V.empty())
207 return std::nullopt;
208
209 return std::accumulate(first: V.begin(), last: V.end(), init: uint64_t(1),
210 binary_op: std::multiplies<uint64_t>{});
211}
212
213PTXOpaqueType llvm::getPTXOpaqueType(const GlobalVariable &GV) {
214 if (findOneNVVMAnnotation(GV: &GV, Prop: NVVMMetadata::Texture))
215 return PTXOpaqueType::Texture;
216 if (findOneNVVMAnnotation(GV: &GV, Prop: NVVMMetadata::Surface))
217 return PTXOpaqueType::Surface;
218 if (findOneNVVMAnnotation(GV: &GV, Prop: NVVMMetadata::Sampler))
219 return PTXOpaqueType::Sampler;
220 return PTXOpaqueType::None;
221}
222
223PTXOpaqueType llvm::getPTXOpaqueType(const Argument &Arg) {
224 if (argHasNVVMAnnotation(Val: Arg, Annotation: NVVMMetadata::Sampler))
225 return PTXOpaqueType::Sampler;
226 if (argHasNVVMAnnotation(Val: Arg, Annotation: NVVMMetadata::ReadOnlyImage))
227 return PTXOpaqueType::Texture;
228 if (argHasNVVMAnnotation(Val: Arg, Annotation: NVVMMetadata::WriteOnlyImage) ||
229 argHasNVVMAnnotation(Val: Arg, Annotation: NVVMMetadata::ReadWriteImage))
230 return PTXOpaqueType::Surface;
231 return PTXOpaqueType::None;
232}
233
234PTXOpaqueType llvm::getPTXOpaqueType(const Value &V) {
235 if (const auto *GV = dyn_cast<GlobalVariable>(Val: &V))
236 return getPTXOpaqueType(GV: *GV);
237 if (const auto *Arg = dyn_cast<Argument>(Val: &V))
238 return getPTXOpaqueType(Arg: *Arg);
239 return PTXOpaqueType::None;
240}
241
242bool llvm::isManaged(const Value &V) {
243 return globalHasNVVMAnnotation(V, Prop: NVVMMetadata::Managed);
244}
245
246SmallVector<unsigned, 3> llvm::getMaxNTID(const Function &F) {
247 return getFnAttrParsedVector(F, Attr: NVVMAttr::MaxNTID);
248}
249
250SmallVector<unsigned, 3> llvm::getReqNTID(const Function &F) {
251 return getFnAttrParsedVector(F, Attr: NVVMAttr::ReqNTID);
252}
253
254SmallVector<unsigned, 3> llvm::getClusterDim(const Function &F) {
255 return getFnAttrParsedVector(F, Attr: NVVMAttr::ClusterDim);
256}
257
258std::optional<uint64_t> llvm::getOverallMaxNTID(const Function &F) {
259 // Note: The semantics here are a bit strange. The PTX ISA states the
260 // following (11.4.2. Performance-Tuning Directives: .maxntid):
261 //
262 // Note that this directive guarantees that the total number of threads does
263 // not exceed the maximum, but does not guarantee that the limit in any
264 // particular dimension is not exceeded.
265 return getVectorProduct(V: getMaxNTID(F));
266}
267
268std::optional<uint64_t> llvm::getOverallReqNTID(const Function &F) {
269 // Note: The semantics here are a bit strange. See getOverallMaxNTID.
270 return getVectorProduct(V: getReqNTID(F));
271}
272
273std::optional<uint64_t> llvm::getOverallClusterRank(const Function &F) {
274 // maxclusterrank and cluster_dim are mutually exclusive.
275 if (const auto ClusterRank = getMaxClusterRank(F))
276 return ClusterRank;
277
278 // Note: The semantics here are a bit strange. See getOverallMaxNTID.
279 return getVectorProduct(V: getClusterDim(F));
280}
281
282std::optional<unsigned> llvm::getMaxClusterRank(const Function &F) {
283 return getFnAttrParsedInt(F, Attr: NVVMAttr::MaxClusterRank);
284}
285
286std::optional<unsigned> llvm::getMinCTASm(const Function &F) {
287 return getFnAttrParsedInt(F, Attr: NVVMAttr::MinCTASm);
288}
289
290std::optional<unsigned> llvm::getMaxNReg(const Function &F) {
291 return getFnAttrParsedInt(F, Attr: NVVMAttr::MaxNReg);
292}
293
294bool llvm::hasBlocksAreClusters(const Function &F) {
295 return F.hasFnAttribute(Kind: NVVMAttr::BlocksAreClusters);
296}
297
298bool llvm::isParamGridConstant(const Argument &Arg) {
299 assert(isKernelFunction(*Arg.getParent()) &&
300 "only kernel arguments can be grid_constant");
301
302 if (!Arg.hasByValAttr())
303 return false;
304
305 // Lowering an argument as a grid_constant violates the byval semantics (and
306 // the C++ API) by reusing the same memory location for the argument across
307 // multiple threads. If an argument doesn't read memory and its address is not
308 // captured (its address is not compared with any value), then the tweak of
309 // the C++ API and byval semantics is unobservable by the program and we can
310 // lower the arg as a grid_constant.
311 if (Arg.onlyReadsMemory()) {
312 const auto CI = Arg.getAttributes().getCaptureInfo();
313 if (!capturesAddress(CC: CI) && !capturesFullProvenance(CC: CI))
314 return true;
315 }
316
317 // "grid_constant" counts argument indices starting from 1
318 return Arg.hasAttribute(Kind: NVVMAttr::GridConstant);
319}
320
321MaybeAlign llvm::getStackAlign(const CallBase &I, unsigned Index) {
322 // First check the alignstack metadata.
323 if (MaybeAlign StackAlign =
324 I.getAttributes().getAttributes(Index).getStackAlignment())
325 return StackAlign;
326
327 // If that is missing, check the legacy nvvm metadata.
328 if (MDNode *AlignNode = I.getMetadata(Kind: "callalign")) {
329 for (int I = 0, N = AlignNode->getNumOperands(); I < N; I++) {
330 if (const auto *CI =
331 mdconst::dyn_extract<ConstantInt>(MD: AlignNode->getOperand(I))) {
332 unsigned V = CI->getZExtValue();
333 if ((V >> 16) == Index)
334 return Align(V & 0xFFFF);
335 if ((V >> 16) > Index)
336 return std::nullopt;
337 }
338 }
339 }
340 return std::nullopt;
341}
342