1//===-- AMDGPUSwLowerLDS.cpp -----------------------------------------===//
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 pass lowers the local data store, LDS, uses in kernel and non-kernel
10// functions in module to use dynamically allocated global memory.
11// Packed LDS Layout is emulated in the global memory.
12// The lowered memory instructions from LDS to global memory are then
13// instrumented for address sanitizer, to catch addressing errors.
14// This pass only work when address sanitizer has been enabled and has
15// instrumented the IR. It identifies that IR has been instrumented using
16// "nosanitize_address" module flag.
17//
18// Replacement of Kernel LDS accesses:
19// For a kernel, LDS access can be static or dynamic which are direct
20// (accessed within kernel) and indirect (accessed through non-kernels).
21// All these LDS accesses corresponding to kernel will be packed together,
22// where all static LDS accesses will be allocated first and then dynamic
23// LDS follows. The total size with alignment is calculated. A new LDS global
24// will be created for the kernel called "SW LDS" and it will have the
25// attribute "amdgpu-lds-size" attached with value of the size calculated.
26// All the LDS accesses in the module will be replaced by GEP with offset
27// into the "Sw LDS".
28// A new "llvm.amdgcn.<kernel>.dynlds" is created per kernel accessing
29// the dynamic LDS. This will be marked used by kernel and will have
30// MD_absolue_symbol metadata set to total static LDS size, Since dynamic
31// LDS allocation starts after all static LDS allocation.
32//
33// A device global memory equal to the total LDS size will be allocated.
34// At the prologue of the kernel, a single work-item from the
35// work-group, does a "malloc" and stores the pointer of the
36// allocation in "SW LDS".
37//
38// To store the offsets corresponding to all LDS accesses, another global
39// variable is created which will be called "SW LDS metadata" in this pass.
40// - SW LDS Global:
41// It is LDS global of ptr type with name
42// "llvm.amdgcn.sw.lds.<kernel-name>".
43// - Metadata Global:
44// It is of struct type, with n members. n equals the number of LDS
45// globals accessed by the kernel(direct and indirect). Each member of
46// struct is another struct of type {i32, i32, i32}. First member
47// corresponds to offset, second member corresponds to size of LDS global
48// being replaced and third represents the total aligned size. It will
49// have name "llvm.amdgcn.sw.lds.<kernel-name>.md". This global will have
50// an initializer with static LDS related offsets and sizes initialized.
51// But for dynamic LDS related entries, offsets will be initialized to
52// previous static LDS allocation end offset. Sizes for them will be zero
53// initially. These dynamic LDS offset and size values will be updated
54// within the kernel, since kernel can read the dynamic LDS size
55// allocation done at runtime with query to "hidden_dynamic_lds_size"
56// hidden kernel argument.
57//
58// At the epilogue of kernel, allocated memory would be made free by the same
59// single work-item.
60//
61// Replacement of non-kernel LDS accesses:
62// Multiple kernels can access the same non-kernel function.
63// All the kernels accessing LDS through non-kernels are sorted and
64// assigned a kernel-id. All the LDS globals accessed by non-kernels
65// are sorted. This information is used to build two tables:
66// - Base table:
67// Base table will have single row, with elements of the row
68// placed as per kernel ID. Each element in the row corresponds
69// to ptr of "SW LDS" variable created for that kernel.
70// - Offset table:
71// Offset table will have multiple rows and columns.
72// Rows are assumed to be from 0 to (n-1). n is total number
73// of kernels accessing the LDS through non-kernels.
74// Each row will have m elements. m is the total number of
75// unique LDS globals accessed by all non-kernels.
76// Each element in the row correspond to the ptr of
77// the replacement of LDS global done by that particular kernel.
78// A LDS variable in non-kernel will be replaced based on the information
79// from base and offset tables. Based on kernel-id query, ptr of "SW
80// LDS" for that corresponding kernel is obtained from base table.
81// The Offset into the base "SW LDS" is obtained from
82// corresponding element in offset table. With this information, replacement
83// value is obtained.
84//===----------------------------------------------------------------------===//
85
86#include "AMDGPU.h"
87#include "AMDGPUAsanInstrumentation.h"
88#include "AMDGPUMemoryUtils.h"
89#include "llvm/ADT/StringExtras.h"
90#include "llvm/ADT/StringRef.h"
91#include "llvm/Analysis/CallGraph.h"
92#include "llvm/Analysis/DomTreeUpdater.h"
93#include "llvm/IR/Constants.h"
94#include "llvm/IR/DIBuilder.h"
95#include "llvm/IR/DebugInfo.h"
96#include "llvm/IR/DebugInfoMetadata.h"
97#include "llvm/IR/IRBuilder.h"
98#include "llvm/IR/Instructions.h"
99#include "llvm/IR/MDBuilder.h"
100#include "llvm/IR/ReplaceConstant.h"
101#include "llvm/Pass.h"
102#include "llvm/Support/raw_ostream.h"
103#include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
104#include "llvm/Transforms/Utils/ModuleUtils.h"
105
106#include <algorithm>
107
108#define DEBUG_TYPE "amdgpu-sw-lower-lds"
109#define COV5_HIDDEN_DYN_LDS_SIZE_ARG 15
110
111using namespace llvm;
112using namespace AMDGPU;
113
114namespace {
115
116cl::opt<bool>
117 AsanInstrumentLDS("amdgpu-asan-instrument-lds",
118 cl::desc("Run asan instrumentation on LDS instructions "
119 "lowered to global memory"),
120 cl::init(Val: true), cl::Hidden);
121
122using DomTreeCallback = function_ref<DominatorTree *(Function &F)>;
123
124struct LDSAccessTypeInfo {
125 SetVector<GlobalVariable *> StaticLDSGlobals;
126 SetVector<GlobalVariable *> DynamicLDSGlobals;
127};
128
129// Struct to hold all the Metadata required for a kernel
130// to replace a LDS global uses with corresponding offset
131// in to device global memory.
132struct KernelLDSParameters {
133 GlobalVariable *SwLDS = nullptr;
134 GlobalVariable *SwDynLDS = nullptr;
135 GlobalVariable *SwLDSMetadata = nullptr;
136 LDSAccessTypeInfo DirectAccess;
137 LDSAccessTypeInfo IndirectAccess;
138 DenseMap<GlobalVariable *, SmallVector<uint32_t, 3>>
139 LDSToReplacementIndicesMap;
140 uint32_t MallocSize = 0;
141 uint32_t LDSSize = 0;
142 SmallVector<std::pair<uint32_t, uint32_t>, 64> RedzoneOffsetAndSizeVector;
143};
144
145// Struct to store information for creation of offset table
146// for all the non-kernel LDS accesses.
147struct NonKernelLDSParameters {
148 GlobalVariable *LDSBaseTable = nullptr;
149 GlobalVariable *LDSOffsetTable = nullptr;
150 SetVector<Function *> OrderedKernels;
151 SetVector<GlobalVariable *> OrdereLDSGlobals;
152};
153
154struct AsanInstrumentInfo {
155 int Scale = 0;
156 uint32_t Offset = 0;
157 SetVector<Instruction *> Instructions;
158};
159
160struct FunctionsAndLDSAccess {
161 MapVector<Function *, KernelLDSParameters> KernelToLDSParametersMap;
162 SetVector<Function *> KernelsWithIndirectLDSAccess;
163 SetVector<Function *> NonKernelsWithLDSArgument;
164 SetVector<GlobalVariable *> AllNonKernelLDSAccess;
165 FunctionVariableMap NonKernelToLDSAccessMap;
166};
167
168class AMDGPUSwLowerLDS {
169public:
170 AMDGPUSwLowerLDS(Module &Mod, DomTreeCallback Callback)
171 : M(Mod), IRB(M), DTCallback(Callback) {}
172 bool run();
173 void getUsesOfLDSByNonKernels();
174 void getNonKernelsWithLDSArguments(const CallGraph &CG);
175 SetVector<Function *>
176 getOrderedIndirectLDSAccessingKernels(SetVector<Function *> &Kernels);
177 SetVector<GlobalVariable *>
178 getOrderedNonKernelAllLDSGlobals(SetVector<GlobalVariable *> &Variables);
179 void buildSwLDSGlobal(Function *Func);
180 void buildSwDynLDSGlobal(Function *Func);
181 void populateSwMetadataGlobal(Function *Func);
182 void populateSwLDSAttributeAndMetadata(Function *Func);
183 void populateLDSToReplacementIndicesMap(Function *Func);
184 void getLDSMemoryInstructions(Function *Func,
185 SetVector<Instruction *> &LDSInstructions);
186 void replaceKernelLDSAccesses(Function *Func);
187 Value *getTranslatedGlobalMemoryPtrOfLDS(Value *LoadMallocPtr, Value *LDSPtr);
188 Value *getTranslatedLDSPtrOfFlat(Value *LoadMallocPtr, Value *FlatPtr,
189 Type *LDSPtrTy);
190 void translateLDSMemoryOperationsToGlobalMemory(
191 Function *Func, Value *LoadMallocPtr,
192 SetVector<Instruction *> &LDSInstructions);
193 void poisonRedzones(Function *Func, Value *MallocPtr);
194 void lowerKernelLDSAccesses(Function *Func, DomTreeUpdater &DTU);
195 void buildNonKernelLDSOffsetTable(NonKernelLDSParameters &NKLDSParams);
196 void buildNonKernelLDSBaseTable(NonKernelLDSParameters &NKLDSParams);
197 Constant *
198 getAddressesOfVariablesInKernel(Function *Func,
199 SetVector<GlobalVariable *> &Variables);
200 void lowerNonKernelLDSAccesses(Function *Func,
201 SetVector<GlobalVariable *> &LDSGlobals,
202 NonKernelLDSParameters &NKLDSParams);
203 void
204 updateMallocSizeForDynamicLDS(Function *Func, Value **CurrMallocSize,
205 Value *HiddenDynLDSSize,
206 SetVector<GlobalVariable *> &DynamicLDSGlobals);
207 void initAsanInfo();
208
209private:
210 Module &M;
211 IRBuilder<> IRB;
212 DomTreeCallback DTCallback;
213 FunctionsAndLDSAccess FuncLDSAccessInfo;
214 AsanInstrumentInfo AsanInfo;
215};
216
217template <typename T> SetVector<T> sortByName(std::vector<T> &&V) {
218 // Sort the vector of globals or Functions based on their name.
219 // Returns a SetVector of globals/Functions.
220 sort(V, [](const auto *L, const auto *R) {
221 return L->getName() < R->getName();
222 });
223 return {SetVector<T>(llvm::from_range, V)};
224}
225
226SetVector<GlobalVariable *> AMDGPUSwLowerLDS::getOrderedNonKernelAllLDSGlobals(
227 SetVector<GlobalVariable *> &Variables) {
228 // Sort all the non-kernel LDS accesses based on their name.
229 return sortByName(
230 V: std::vector<GlobalVariable *>(Variables.begin(), Variables.end()));
231}
232
233SetVector<Function *> AMDGPUSwLowerLDS::getOrderedIndirectLDSAccessingKernels(
234 SetVector<Function *> &Kernels) {
235 // Sort the non-kernels accessing LDS based on their name.
236 // Also assign a kernel ID metadata based on the sorted order.
237 LLVMContext &Ctx = M.getContext();
238 if (Kernels.size() > UINT32_MAX) {
239 report_fatal_error(reason: "Unimplemented SW LDS lowering for > 2**32 kernels");
240 }
241 SetVector<Function *> OrderedKernels =
242 sortByName(V: std::vector<Function *>(Kernels.begin(), Kernels.end()));
243 for (size_t i = 0; i < Kernels.size(); i++) {
244 Metadata *AttrMDArgs[1] = {
245 ConstantAsMetadata::get(C: IRB.getInt32(C: i)),
246 };
247 Function *Func = OrderedKernels[i];
248 Func->setMetadata(Kind: "llvm.amdgcn.lds.kernel.id",
249 Node: MDNode::get(Context&: Ctx, MDs: AttrMDArgs));
250 }
251 return OrderedKernels;
252}
253
254void AMDGPUSwLowerLDS::getNonKernelsWithLDSArguments(const CallGraph &CG) {
255 // Among the kernels accessing LDS, get list of
256 // Non-kernels to which a call is made and a ptr
257 // to addrspace(3) is passed as argument.
258 for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
259 Function *Func = K.first;
260 const CallGraphNode *CGN = CG[Func];
261 if (!CGN)
262 continue;
263 for (auto &I : *CGN) {
264 CallGraphNode *CallerCGN = I.second;
265 Function *CalledFunc = CallerCGN->getFunction();
266 if (!CalledFunc || CalledFunc->isDeclaration())
267 continue;
268 if (AMDGPU::isKernel(F: *CalledFunc))
269 continue;
270 for (auto AI = CalledFunc->arg_begin(), E = CalledFunc->arg_end();
271 AI != E; ++AI) {
272 Type *ArgTy = (*AI).getType();
273 if (!ArgTy->isPointerTy())
274 continue;
275 if (ArgTy->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
276 continue;
277 FuncLDSAccessInfo.NonKernelsWithLDSArgument.insert(X: CalledFunc);
278 // Also add the Calling function to KernelsWithIndirectLDSAccess list
279 // so that base table of LDS is generated.
280 FuncLDSAccessInfo.KernelsWithIndirectLDSAccess.insert(X: Func);
281 }
282 }
283 }
284}
285
286void AMDGPUSwLowerLDS::getUsesOfLDSByNonKernels() {
287 for (GlobalVariable *GV : FuncLDSAccessInfo.AllNonKernelLDSAccess) {
288 if (!AMDGPU::isLDSVariableToLower(GV: *GV))
289 continue;
290
291 for (User *V : GV->users()) {
292 if (auto *I = dyn_cast<Instruction>(Val: V)) {
293 Function *F = I->getFunction();
294 if (!isKernel(F: *F) && !F->isDeclaration())
295 FuncLDSAccessInfo.NonKernelToLDSAccessMap[F].insert(V: GV);
296 }
297 }
298 }
299}
300
301static void recordLDSAbsoluteAddress(Module &M, GlobalVariable *GV,
302 uint32_t Address) {
303 // Write the specified address into metadata where it can be retrieved by
304 // the assembler. Format is a half open range, [Address Address+1)
305 LLVMContext &Ctx = M.getContext();
306 auto *IntTy = M.getDataLayout().getIntPtrType(C&: Ctx, AddressSpace: AMDGPUAS::LOCAL_ADDRESS);
307 MDBuilder MDB(Ctx);
308 MDNode *MetadataNode = MDB.createRange(Lo: ConstantInt::get(Ty: IntTy, V: Address),
309 Hi: ConstantInt::get(Ty: IntTy, V: Address + 1));
310 GV->setMetadata(KindID: LLVMContext::MD_absolute_symbol, Node: MetadataNode);
311}
312
313static void addLDSSizeAttribute(Function *Func, uint32_t Offset,
314 bool IsDynLDS) {
315 if (Offset != 0) {
316 std::string Buffer;
317 raw_string_ostream SS{Buffer};
318 SS << Offset;
319 if (IsDynLDS)
320 SS << "," << Offset;
321 Func->addFnAttr(Kind: "amdgpu-lds-size", Val: Buffer);
322 }
323}
324
325static void markUsedByKernel(Function *Func, GlobalVariable *SGV) {
326 BasicBlock *Entry = &Func->getEntryBlock();
327 IRBuilder<> Builder(Entry->getFirstNonPHIIt());
328
329 Function *Decl = Intrinsic::getOrInsertDeclaration(M: Func->getParent(),
330 id: Intrinsic::donothing, OverloadTys: {});
331
332 Value *UseInstance[1] = {
333 Builder.CreateConstInBoundsGEP1_32(Ty: SGV->getValueType(), Ptr: SGV, Idx0: 0)};
334
335 Builder.CreateCall(Callee: Decl, Args: {},
336 OpBundles: {OperandBundleDefT<Value *>("ExplicitUse", UseInstance)});
337}
338
339void AMDGPUSwLowerLDS::buildSwLDSGlobal(Function *Func) {
340 // Create new LDS global required for each kernel to store
341 // device global memory pointer.
342 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
343 // Create new global pointer variable
344 LDSParams.SwLDS = new GlobalVariable(
345 M, IRB.getPtrTy(), false, GlobalValue::InternalLinkage,
346 PoisonValue::get(T: IRB.getPtrTy()), "llvm.amdgcn.sw.lds." + Func->getName(),
347 nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, false);
348 GlobalValue::SanitizerMetadata MD;
349 MD.NoAddress = true;
350 LDSParams.SwLDS->setSanitizerMetadata(MD);
351}
352
353void AMDGPUSwLowerLDS::buildSwDynLDSGlobal(Function *Func) {
354 // Create new Dyn LDS global if kernel accesses dyn LDS.
355 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
356 if (LDSParams.DirectAccess.DynamicLDSGlobals.empty() &&
357 LDSParams.IndirectAccess.DynamicLDSGlobals.empty())
358 return;
359 // Create new global pointer variable
360 auto *emptyCharArray = ArrayType::get(ElementType: IRB.getInt8Ty(), NumElements: 0);
361 LDSParams.SwDynLDS = new GlobalVariable(
362 M, emptyCharArray, false, GlobalValue::ExternalLinkage, nullptr,
363 "llvm.amdgcn." + Func->getName() + ".dynlds", nullptr,
364 GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, false);
365 markUsedByKernel(Func, SGV: LDSParams.SwDynLDS);
366 GlobalValue::SanitizerMetadata MD;
367 MD.NoAddress = true;
368 LDSParams.SwDynLDS->setSanitizerMetadata(MD);
369}
370
371void AMDGPUSwLowerLDS::populateSwLDSAttributeAndMetadata(Function *Func) {
372 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
373 bool IsDynLDSUsed = LDSParams.SwDynLDS;
374 uint32_t Offset = LDSParams.LDSSize;
375 recordLDSAbsoluteAddress(M, GV: LDSParams.SwLDS, Address: 0);
376 addLDSSizeAttribute(Func, Offset, IsDynLDS: IsDynLDSUsed);
377 if (LDSParams.SwDynLDS)
378 recordLDSAbsoluteAddress(M, GV: LDSParams.SwDynLDS, Address: Offset);
379}
380
381void AMDGPUSwLowerLDS::populateSwMetadataGlobal(Function *Func) {
382 // Create new metadata global for every kernel and initialize the
383 // start offsets and sizes corresponding to each LDS accesses.
384 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
385 auto &Ctx = M.getContext();
386 auto &DL = M.getDataLayout();
387 std::vector<Type *> Items;
388 Type *Int32Ty = IRB.getInt32Ty();
389 std::vector<Constant *> Initializers;
390 Align MaxAlignment(1);
391 auto UpdateMaxAlignment = [&MaxAlignment, &DL](GlobalVariable *GV) {
392 Align GVAlign = AMDGPU::getAlign(DL, GV);
393 MaxAlignment = std::max(a: MaxAlignment, b: GVAlign);
394 };
395
396 for (GlobalVariable *GV : LDSParams.DirectAccess.StaticLDSGlobals)
397 UpdateMaxAlignment(GV);
398
399 for (GlobalVariable *GV : LDSParams.DirectAccess.DynamicLDSGlobals)
400 UpdateMaxAlignment(GV);
401
402 for (GlobalVariable *GV : LDSParams.IndirectAccess.StaticLDSGlobals)
403 UpdateMaxAlignment(GV);
404
405 for (GlobalVariable *GV : LDSParams.IndirectAccess.DynamicLDSGlobals)
406 UpdateMaxAlignment(GV);
407
408 //{StartOffset, AlignedSizeInBytes}
409 SmallString<128> MDItemStr;
410 raw_svector_ostream MDItemOS(MDItemStr);
411 MDItemOS << "llvm.amdgcn.sw.lds." << Func->getName() << ".md.item";
412
413 StructType *LDSItemTy =
414 StructType::create(Context&: Ctx, Elements: {Int32Ty, Int32Ty, Int32Ty}, Name: MDItemOS.str());
415 uint32_t &MallocSize = LDSParams.MallocSize;
416 SetVector<GlobalVariable *> UniqueLDSGlobals;
417 int AsanScale = AsanInfo.Scale;
418 auto buildInitializerForSwLDSMD =
419 [&](SetVector<GlobalVariable *> &LDSGlobals) {
420 for (auto &GV : LDSGlobals) {
421 if (is_contained(Range&: UniqueLDSGlobals, Element: GV))
422 continue;
423 UniqueLDSGlobals.insert(X: GV);
424
425 Type *Ty = GV->getValueType();
426 const uint64_t SizeInBytes = DL.getTypeAllocSize(Ty);
427 Items.push_back(x: LDSItemTy);
428 Constant *ItemStartOffset = ConstantInt::get(Ty: Int32Ty, V: MallocSize);
429 Constant *SizeInBytesConst = ConstantInt::get(Ty: Int32Ty, V: SizeInBytes);
430 // Get redzone size corresponding a size.
431 const uint64_t RightRedzoneSize =
432 AMDGPU::getRedzoneSizeForGlobal(Scale: AsanScale, SizeInBytes);
433 // Update MallocSize with current size and redzone size.
434 MallocSize += SizeInBytes;
435 if (!AMDGPU::isDynamicLDS(GV: *GV))
436 LDSParams.RedzoneOffsetAndSizeVector.emplace_back(Args&: MallocSize,
437 Args: RightRedzoneSize);
438 MallocSize += RightRedzoneSize;
439 // Align current size plus redzone.
440 uint64_t AlignedSize =
441 alignTo(Size: SizeInBytes + RightRedzoneSize, A: MaxAlignment);
442 Constant *AlignedSizeInBytesConst =
443 ConstantInt::get(Ty: Int32Ty, V: AlignedSize);
444 // Align MallocSize
445 MallocSize = alignTo(Size: MallocSize, A: MaxAlignment);
446 Constant *InitItem =
447 ConstantStruct::get(T: LDSItemTy, V: {ItemStartOffset, SizeInBytesConst,
448 AlignedSizeInBytesConst});
449 Initializers.push_back(x: InitItem);
450 }
451 };
452 SetVector<GlobalVariable *> SwLDSVector;
453 SwLDSVector.insert(X: LDSParams.SwLDS);
454 buildInitializerForSwLDSMD(SwLDSVector);
455 buildInitializerForSwLDSMD(LDSParams.DirectAccess.StaticLDSGlobals);
456 buildInitializerForSwLDSMD(LDSParams.IndirectAccess.StaticLDSGlobals);
457 buildInitializerForSwLDSMD(LDSParams.DirectAccess.DynamicLDSGlobals);
458 buildInitializerForSwLDSMD(LDSParams.IndirectAccess.DynamicLDSGlobals);
459
460 // Update the LDS size used by the kernel.
461 Type *Ty = LDSParams.SwLDS->getValueType();
462 const uint64_t SizeInBytes = DL.getTypeAllocSize(Ty);
463 uint64_t AlignedSize = alignTo(Size: SizeInBytes, A: MaxAlignment);
464 LDSParams.LDSSize = AlignedSize;
465 SmallString<128> MDTypeStr;
466 raw_svector_ostream MDTypeOS(MDTypeStr);
467 MDTypeOS << "llvm.amdgcn.sw.lds." << Func->getName() << ".md.type";
468 StructType *MetadataStructType =
469 StructType::create(Context&: Ctx, Elements: Items, Name: MDTypeOS.str());
470 SmallString<128> MDStr;
471 raw_svector_ostream MDOS(MDStr);
472 MDOS << "llvm.amdgcn.sw.lds." << Func->getName() << ".md";
473 LDSParams.SwLDSMetadata = new GlobalVariable(
474 M, MetadataStructType, false, GlobalValue::InternalLinkage,
475 PoisonValue::get(T: MetadataStructType), MDOS.str(), nullptr,
476 GlobalValue::NotThreadLocal, AMDGPUAS::GLOBAL_ADDRESS, false);
477 Constant *data = ConstantStruct::get(T: MetadataStructType, V: Initializers);
478 LDSParams.SwLDSMetadata->setInitializer(data);
479 assert(LDSParams.SwLDS);
480 // Set the alignment to MaxAlignment for SwLDS.
481 LDSParams.SwLDS->setAlignment(MaxAlignment);
482 if (LDSParams.SwDynLDS)
483 LDSParams.SwDynLDS->setAlignment(MaxAlignment);
484 GlobalValue::SanitizerMetadata MD;
485 MD.NoAddress = true;
486 LDSParams.SwLDSMetadata->setSanitizerMetadata(MD);
487}
488
489void AMDGPUSwLowerLDS::populateLDSToReplacementIndicesMap(Function *Func) {
490 // Fill the corresponding LDS replacement indices for each LDS access
491 // related to this kernel.
492 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
493 SetVector<GlobalVariable *> UniqueLDSGlobals;
494 auto PopulateIndices = [&](SetVector<GlobalVariable *> &LDSGlobals,
495 uint32_t &Idx) {
496 for (auto &GV : LDSGlobals) {
497 if (is_contained(Range&: UniqueLDSGlobals, Element: GV))
498 continue;
499 UniqueLDSGlobals.insert(X: GV);
500 LDSParams.LDSToReplacementIndicesMap[GV] = {0, Idx, 0};
501 ++Idx;
502 }
503 };
504 uint32_t Idx = 0;
505 SetVector<GlobalVariable *> SwLDSVector;
506 SwLDSVector.insert(X: LDSParams.SwLDS);
507 PopulateIndices(SwLDSVector, Idx);
508 PopulateIndices(LDSParams.DirectAccess.StaticLDSGlobals, Idx);
509 PopulateIndices(LDSParams.IndirectAccess.StaticLDSGlobals, Idx);
510 PopulateIndices(LDSParams.DirectAccess.DynamicLDSGlobals, Idx);
511 PopulateIndices(LDSParams.IndirectAccess.DynamicLDSGlobals, Idx);
512}
513
514static void replacesUsesOfGlobalInFunction(Function *Func, GlobalVariable *GV,
515 Value *Replacement) {
516 // Replace all uses of LDS global in this Function with a Replacement.
517 auto ReplaceUsesLambda = [Func](const Use &U) -> bool {
518 auto *V = U.getUser();
519 if (auto *Inst = dyn_cast<Instruction>(Val: V)) {
520 auto *Func1 = Inst->getFunction();
521 if (Func == Func1)
522 return true;
523 }
524 return false;
525 };
526 GV->replaceUsesWithIf(New: Replacement, ShouldReplace: ReplaceUsesLambda);
527}
528
529void AMDGPUSwLowerLDS::replaceKernelLDSAccesses(Function *Func) {
530 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
531 GlobalVariable *SwLDS = LDSParams.SwLDS;
532 assert(SwLDS);
533 GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
534 assert(SwLDSMetadata);
535 StructType *SwLDSMetadataStructType =
536 cast<StructType>(Val: SwLDSMetadata->getValueType());
537 Type *Int32Ty = IRB.getInt32Ty();
538 auto &IndirectAccess = LDSParams.IndirectAccess;
539 auto &DirectAccess = LDSParams.DirectAccess;
540 // Replace all uses of LDS global in this Function with a Replacement.
541 SetVector<GlobalVariable *> UniqueLDSGlobals;
542 auto ReplaceLDSGlobalUses = [&](SetVector<GlobalVariable *> &LDSGlobals) {
543 for (auto &GV : LDSGlobals) {
544 // Do not generate instructions if LDS access is in non-kernel
545 // i.e indirect-access.
546 if ((IndirectAccess.StaticLDSGlobals.contains(key: GV) ||
547 IndirectAccess.DynamicLDSGlobals.contains(key: GV)) &&
548 (!DirectAccess.StaticLDSGlobals.contains(key: GV) &&
549 !DirectAccess.DynamicLDSGlobals.contains(key: GV)))
550 continue;
551 if (is_contained(Range&: UniqueLDSGlobals, Element: GV))
552 continue;
553 UniqueLDSGlobals.insert(X: GV);
554 auto &Indices = LDSParams.LDSToReplacementIndicesMap[GV];
555 assert(Indices.size() == 3);
556 Constant *GEPIdx[] = {ConstantInt::get(Ty: Int32Ty, V: Indices[0]),
557 ConstantInt::get(Ty: Int32Ty, V: Indices[1]),
558 ConstantInt::get(Ty: Int32Ty, V: Indices[2])};
559 Constant *GEP = ConstantExpr::getGetElementPtr(
560 DL: Func->getDataLayout(), Ty: SwLDSMetadataStructType, C: SwLDSMetadata, IdxList: GEPIdx,
561 NW: GEPNoWrapFlags::inBounds());
562 Value *Offset = IRB.CreateLoad(Ty: Int32Ty, Ptr: GEP);
563 Value *BasePlusOffset =
564 IRB.CreateInBoundsGEP(Ty: IRB.getInt8Ty(), Ptr: SwLDS, IdxList: {Offset});
565 LLVM_DEBUG(GV->printAsOperand(dbgs() << "Sw LDS Lowering, Replacing LDS ",
566 false));
567 replacesUsesOfGlobalInFunction(Func, GV, Replacement: BasePlusOffset);
568 }
569 };
570 ReplaceLDSGlobalUses(DirectAccess.StaticLDSGlobals);
571 ReplaceLDSGlobalUses(IndirectAccess.StaticLDSGlobals);
572 ReplaceLDSGlobalUses(DirectAccess.DynamicLDSGlobals);
573 ReplaceLDSGlobalUses(IndirectAccess.DynamicLDSGlobals);
574}
575
576void AMDGPUSwLowerLDS::updateMallocSizeForDynamicLDS(
577 Function *Func, Value **CurrMallocSize, Value *HiddenDynLDSSize,
578 SetVector<GlobalVariable *> &DynamicLDSGlobals) {
579 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
580 Type *Int32Ty = IRB.getInt32Ty();
581
582 GlobalVariable *SwLDS = LDSParams.SwLDS;
583 GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
584 assert(SwLDS && SwLDSMetadata);
585 StructType *MetadataStructType =
586 cast<StructType>(Val: SwLDSMetadata->getValueType());
587 unsigned MaxAlignment = SwLDS->getAlign().valueOrOne().value();
588 Value *MaxAlignValue = IRB.getInt32(C: MaxAlignment);
589 Value *MaxAlignValueMinusOne = IRB.getInt32(C: MaxAlignment - 1);
590
591 for (GlobalVariable *DynGV : DynamicLDSGlobals) {
592 auto &Indices = LDSParams.LDSToReplacementIndicesMap[DynGV];
593 // Update the Offset metadata.
594 Constant *Index0 = ConstantInt::get(Ty: Int32Ty, V: 0);
595 Constant *Index1 = ConstantInt::get(Ty: Int32Ty, V: Indices[1]);
596
597 Constant *Index2Offset = ConstantInt::get(Ty: Int32Ty, V: 0);
598 auto *GEPForOffset = IRB.CreateInBoundsGEP(
599 Ty: MetadataStructType, Ptr: SwLDSMetadata, IdxList: {Index0, Index1, Index2Offset});
600
601 IRB.CreateStore(Val: *CurrMallocSize, Ptr: GEPForOffset);
602 // Update the size and Aligned Size metadata.
603 Constant *Index2Size = ConstantInt::get(Ty: Int32Ty, V: 1);
604 auto *GEPForSize = IRB.CreateInBoundsGEP(Ty: MetadataStructType, Ptr: SwLDSMetadata,
605 IdxList: {Index0, Index1, Index2Size});
606
607 Value *CurrDynLDSSize = IRB.CreateLoad(Ty: Int32Ty, Ptr: HiddenDynLDSSize);
608 IRB.CreateStore(Val: CurrDynLDSSize, Ptr: GEPForSize);
609 Constant *Index2AlignedSize = ConstantInt::get(Ty: Int32Ty, V: 2);
610 auto *GEPForAlignedSize = IRB.CreateInBoundsGEP(
611 Ty: MetadataStructType, Ptr: SwLDSMetadata, IdxList: {Index0, Index1, Index2AlignedSize});
612
613 Value *AlignedDynLDSSize =
614 IRB.CreateAdd(LHS: CurrDynLDSSize, RHS: MaxAlignValueMinusOne);
615 AlignedDynLDSSize = IRB.CreateUDiv(LHS: AlignedDynLDSSize, RHS: MaxAlignValue);
616 AlignedDynLDSSize = IRB.CreateMul(LHS: AlignedDynLDSSize, RHS: MaxAlignValue);
617 IRB.CreateStore(Val: AlignedDynLDSSize, Ptr: GEPForAlignedSize);
618
619 // Update the Current Malloc Size
620 *CurrMallocSize = IRB.CreateAdd(LHS: *CurrMallocSize, RHS: AlignedDynLDSSize);
621 }
622}
623
624static DebugLoc getOrCreateDebugLoc(const Instruction *InsertBefore,
625 DISubprogram *SP) {
626 assert(InsertBefore);
627 if (InsertBefore->getDebugLoc())
628 return InsertBefore->getDebugLoc();
629 if (SP)
630 return DILocation::get(Context&: SP->getContext(), Line: SP->getLine(), Column: 1, Scope: SP);
631 return DebugLoc();
632}
633
634void AMDGPUSwLowerLDS::getLDSMemoryInstructions(
635 Function *Func, SetVector<Instruction *> &LDSInstructions) {
636 for (BasicBlock &BB : *Func) {
637 for (Instruction &Inst : BB) {
638 if (LoadInst *LI = dyn_cast<LoadInst>(Val: &Inst)) {
639 if (LI->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
640 LDSInstructions.insert(X: &Inst);
641 } else if (StoreInst *SI = dyn_cast<StoreInst>(Val: &Inst)) {
642 if (SI->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
643 LDSInstructions.insert(X: &Inst);
644 } else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(Val: &Inst)) {
645 if (RMW->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
646 LDSInstructions.insert(X: &Inst);
647 } else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(Val: &Inst)) {
648 if (XCHG->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
649 LDSInstructions.insert(X: &Inst);
650 } else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Val: &Inst)) {
651 unsigned SrcAS = ASC->getSrcAddressSpace();
652 unsigned DstAS = ASC->getDestAddressSpace();
653 if ((SrcAS == AMDGPUAS::LOCAL_ADDRESS &&
654 DstAS == AMDGPUAS::FLAT_ADDRESS) ||
655 (SrcAS == AMDGPUAS::FLAT_ADDRESS &&
656 DstAS == AMDGPUAS::LOCAL_ADDRESS))
657 LDSInstructions.insert(X: &Inst);
658 } else if (AnyMemIntrinsic *MI = dyn_cast<AnyMemIntrinsic>(Val: &Inst)) {
659 if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
660 LDSInstructions.insert(X: &Inst);
661 } else if (auto *MTI = dyn_cast<AnyMemTransferInst>(Val: MI)) {
662 if (MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
663 LDSInstructions.insert(X: &Inst);
664 }
665 } else
666 continue;
667 }
668 }
669}
670
671Value *AMDGPUSwLowerLDS::getTranslatedGlobalMemoryPtrOfLDS(Value *LoadMallocPtr,
672 Value *LDSPtr) {
673 assert(LDSPtr && "Invalid LDS pointer operand");
674 Type *LDSPtrType = LDSPtr->getType();
675 LLVMContext &Ctx = M.getContext();
676 const DataLayout &DL = M.getDataLayout();
677 Type *IntTy = DL.getIntPtrType(C&: Ctx, AddressSpace: AMDGPUAS::LOCAL_ADDRESS);
678 if (auto *VecPtrTy = dyn_cast<VectorType>(Val: LDSPtrType)) {
679 // Handle vector of pointers
680 ElementCount NumElements = VecPtrTy->getElementCount();
681 IntTy = VectorType::get(ElementType: IntTy, EC: NumElements);
682 }
683 Value *GepIndex = IRB.CreatePtrToInt(V: LDSPtr, DestTy: IntTy);
684 return IRB.CreateInBoundsGEP(Ty: IRB.getInt8Ty(), Ptr: LoadMallocPtr, IdxList: {GepIndex});
685}
686
687// A flat pointer to lowered LDS holds a global address inside the malloc
688// buffer, but LDS values in this pass are offsets into that buffer.
689// Rebase it as inttoptr(trunc(ptrtoaddr(flat) - ptrtoaddr(buf))).
690Value *AMDGPUSwLowerLDS::getTranslatedLDSPtrOfFlat(Value *LoadMallocPtr,
691 Value *FlatPtr,
692 Type *LDSPtrTy) {
693 const DataLayout &DL = M.getDataLayout();
694 Value *FlatAddr = IRB.CreatePtrToAddr(V: FlatPtr);
695 Value *Base = IRB.CreatePtrToAddr(V: LoadMallocPtr);
696 if (auto *VecTy = dyn_cast<VectorType>(Val: FlatAddr->getType()))
697 Base = IRB.CreateVectorSplat(EC: VecTy->getElementCount(), V: Base);
698 Value *Diff = IRB.CreateSub(LHS: FlatAddr, RHS: Base);
699 Value *Offset = IRB.CreateTrunc(V: Diff, DestTy: DL.getAddressType(PtrTy: LDSPtrTy));
700 Value *LDSPtr = IRB.CreateIntToPtr(V: Offset, DestTy: LDSPtrTy);
701 // Flat null is 0 but LDS null is -1, which the rebased offset cannot keep.
702 Constant *LDSNull = ConstantExpr::getAddrSpaceCast(
703 C: Constant::getNullValue(Ty: FlatPtr->getType()), Ty: LDSPtrTy);
704 return IRB.CreateSelect(C: IRB.CreateIsNull(Arg: FlatPtr), True: LDSNull, False: LDSPtr);
705}
706
707void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
708 Function *Func, Value *LoadMallocPtr,
709 SetVector<Instruction *> &LDSInstructions) {
710 LLVM_DEBUG(dbgs() << "Translating LDS memory operations to global memory : "
711 << Func->getName());
712 for (Instruction *Inst : LDSInstructions) {
713 IRB.SetInsertPoint(Inst);
714 if (LoadInst *LI = dyn_cast<LoadInst>(Val: Inst)) {
715 Value *LIOperand = LI->getPointerOperand();
716 Value *Replacement =
717 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: LIOperand);
718 LoadInst *NewLI =
719 IRB.CreateLoad(Ty: LI->getType(), Ptr: Replacement, Props: LI->getProperties());
720 AsanInfo.Instructions.insert(X: NewLI);
721 LI->replaceAllUsesWith(V: NewLI);
722 LI->eraseFromParent();
723 } else if (StoreInst *SI = dyn_cast<StoreInst>(Val: Inst)) {
724 Value *SIOperand = SI->getPointerOperand();
725 Value *Replacement =
726 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: SIOperand);
727 StoreInst *NewSI = IRB.CreateStore(Val: SI->getValueOperand(), Ptr: Replacement,
728 Props: SI->getProperties());
729 AsanInfo.Instructions.insert(X: NewSI);
730 SI->replaceAllUsesWith(V: NewSI);
731 SI->eraseFromParent();
732 } else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(Val: Inst)) {
733 Value *RMWPtrOperand = RMW->getPointerOperand();
734 Value *RMWValOperand = RMW->getValOperand();
735 Value *Replacement =
736 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: RMWPtrOperand);
737 AtomicRMWInst *NewRMW = IRB.CreateAtomicRMW(
738 Op: RMW->getOperation(), Ptr: Replacement, Val: RMWValOperand, Align: RMW->getAlign(),
739 Ordering: RMW->getOrdering(), SSID: RMW->getSyncScopeID());
740 NewRMW->setVolatile(RMW->isVolatile());
741 AsanInfo.Instructions.insert(X: NewRMW);
742 RMW->replaceAllUsesWith(V: NewRMW);
743 RMW->eraseFromParent();
744 } else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(Val: Inst)) {
745 Value *XCHGPtrOperand = XCHG->getPointerOperand();
746 Value *Replacement =
747 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: XCHGPtrOperand);
748 AtomicCmpXchgInst *NewXCHG = IRB.CreateAtomicCmpXchg(
749 Ptr: Replacement, Cmp: XCHG->getCompareOperand(), New: XCHG->getNewValOperand(),
750 Align: XCHG->getAlign(), SuccessOrdering: XCHG->getSuccessOrdering(),
751 FailureOrdering: XCHG->getFailureOrdering(), SSID: XCHG->getSyncScopeID());
752 NewXCHG->setVolatile(XCHG->isVolatile());
753 AsanInfo.Instructions.insert(X: NewXCHG);
754 XCHG->replaceAllUsesWith(V: NewXCHG);
755 XCHG->eraseFromParent();
756 } else if (AnyMemIntrinsic *MI = dyn_cast<AnyMemIntrinsic>(Val: Inst)) {
757 Value *NewDest = MI->getRawDest();
758 if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
759 NewDest = getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: NewDest);
760 CallInst *NewMI = nullptr;
761 if (AnyMemSetInst *MSI = dyn_cast<AnyMemSetInst>(Val: MI)) {
762 if (MI->isAtomic()) {
763 NewMI = IRB.CreateElementUnorderedAtomicMemSet(
764 Ptr: NewDest, Val: MSI->getValue(), Size: MSI->getLength(),
765 Alignment: MSI->getDestAlign().valueOrOne(), ElementSize: MSI->getElementSizeInBytes());
766 } else {
767 NewMI = IRB.CreateMemSet(Ptr: NewDest, Val: MSI->getValue(), Size: MSI->getLength(),
768 Align: MSI->getDestAlign(),
769 isVolatile: cast<MemSetInst>(Val: MI)->isVolatile());
770 }
771 } else if (AnyMemTransferInst *MTI = dyn_cast<AnyMemTransferInst>(Val: MI)) {
772 Value *NewSrc = MTI->getRawSource();
773 if (MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
774 NewSrc = getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: NewSrc);
775 if (MI->isAtomic()) {
776 if (MI->getIntrinsicID() ==
777 Intrinsic::memmove_element_unordered_atomic) {
778 NewMI = IRB.CreateElementUnorderedAtomicMemMove(
779 Dst: NewDest, DstAlign: MTI->getDestAlign().valueOrOne(), Src: NewSrc,
780 SrcAlign: MTI->getSourceAlign().valueOrOne(), Size: MTI->getLength(),
781 ElementSize: MTI->getElementSizeInBytes());
782 } else {
783 NewMI = IRB.CreateElementUnorderedAtomicMemCpy(
784 Dst: NewDest, DstAlign: MTI->getDestAlign().valueOrOne(), Src: NewSrc,
785 SrcAlign: MTI->getSourceAlign().valueOrOne(), Size: MTI->getLength(),
786 ElementSize: MTI->getElementSizeInBytes());
787 }
788 } else {
789 NewMI = IRB.CreateMemTransferInst(
790 IntrID: MI->getIntrinsicID(), Dst: NewDest, DstAlign: MTI->getDestAlign(), Src: NewSrc,
791 SrcAlign: MTI->getSourceAlign(), Size: MTI->getLength(),
792 isVolatile: cast<MemTransferInst>(Val: MI)->isVolatile());
793 }
794 } else
795 reportFatalUsageError(reason: "Unimplemented LDS lowering memory intrinsic");
796 AsanInfo.Instructions.insert(X: NewMI);
797 MI->replaceAllUsesWith(V: NewMI);
798 MI->eraseFromParent();
799 } else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Val: Inst)) {
800 Value *AIOperand = ASC->getPointerOperand();
801 Value *NewAI = nullptr;
802 if (ASC->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
803 NewAI =
804 getTranslatedLDSPtrOfFlat(LoadMallocPtr, FlatPtr: AIOperand, LDSPtrTy: ASC->getType());
805 } else {
806 Value *Replacement =
807 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: AIOperand);
808 // Note: No need to add the instruction to AsanInfo instructions to be
809 // instrumented list. FLAT_ADDRESS ptr would have been already
810 // instrumented by asan pass prior to this pass.
811 NewAI = IRB.CreateAddrSpaceCast(V: Replacement, DestTy: ASC->getType());
812 }
813 ASC->replaceAllUsesWith(V: NewAI);
814 ASC->eraseFromParent();
815 } else
816 report_fatal_error(reason: "Unimplemented LDS lowering instruction");
817 }
818}
819
820void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
821 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
822 Type *Int64Ty = IRB.getInt64Ty();
823 Type *VoidTy = IRB.getVoidTy();
824 FunctionCallee AsanPoisonRegion = M.getOrInsertFunction(
825 Name: "__asan_poison_region",
826 T: FunctionType::get(Result: VoidTy, Params: {Int64Ty, Int64Ty}, isVarArg: false));
827
828 auto RedzonesVec = LDSParams.RedzoneOffsetAndSizeVector;
829 size_t VecSize = RedzonesVec.size();
830 for (unsigned i = 0; i < VecSize; i++) {
831 auto &RedzonePair = RedzonesVec[i];
832 uint64_t RedzoneOffset = RedzonePair.first;
833 uint64_t RedzoneSize = RedzonePair.second;
834 Value *RedzoneAddrOffset = IRB.CreateInBoundsGEP(
835 Ty: IRB.getInt8Ty(), Ptr: MallocPtr, IdxList: {IRB.getInt64(C: RedzoneOffset)});
836 Value *RedzoneAddress = IRB.CreatePtrToInt(V: RedzoneAddrOffset, DestTy: Int64Ty);
837 IRB.CreateCall(Callee: AsanPoisonRegion,
838 Args: {RedzoneAddress, IRB.getInt64(C: RedzoneSize)});
839 }
840}
841
842void AMDGPUSwLowerLDS::lowerKernelLDSAccesses(Function *Func,
843 DomTreeUpdater &DTU) {
844 LLVM_DEBUG(dbgs() << "Sw Lowering Kernel LDS for : " << Func->getName());
845 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
846 auto &Ctx = M.getContext();
847 auto *PrevEntryBlock = &Func->getEntryBlock();
848 SetVector<Instruction *> LDSInstructions;
849 getLDSMemoryInstructions(Func, LDSInstructions);
850 const DataLayout &DL = M.getDataLayout();
851
852 // Create malloc block.
853 auto *MallocBlock = BasicBlock::Create(Context&: Ctx, Name: "Malloc", Parent: Func, InsertBefore: PrevEntryBlock);
854
855 // Create WIdBlock block which has instructions related to selection of
856 // {0,0,0} indiex work item in the work group.
857 auto *WIdBlock = BasicBlock::Create(Context&: Ctx, Name: "WId", Parent: Func, InsertBefore: MallocBlock);
858
859 // Move constant-size allocas from the original entry block to the new entry
860 // block (WIdBlock) so they remain static allocas. Splice the leading cluster
861 // in bulk, then move any stragglers that are interleaved with other
862 // instructions.
863 auto SplitIt = PrevEntryBlock->getFirstNonPHIOrAlloca();
864 WIdBlock->splice(ToIt: WIdBlock->end(), FromBB: PrevEntryBlock, FromBeginIt: PrevEntryBlock->begin(),
865 FromEndIt: SplitIt);
866 for (Instruction &I : make_early_inc_range(Range&: *PrevEntryBlock))
867 if (auto *AI = dyn_cast<AllocaInst>(Val: &I))
868 if (isa<ConstantInt>(Val: AI->getArraySize()))
869 AI->moveBefore(BB&: *WIdBlock, I: WIdBlock->end());
870
871 IRB.SetInsertPoint(WIdBlock->end());
872 DebugLoc FirstDL =
873 getOrCreateDebugLoc(InsertBefore: &*PrevEntryBlock->begin(), SP: Func->getSubprogram());
874 IRB.SetCurrentDebugLocation(FirstDL);
875 Value *WIdx = IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_workitem_id_x, Args: {});
876 Value *WIdy = IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_workitem_id_y, Args: {});
877 Value *WIdz = IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_workitem_id_z, Args: {});
878 Value *XYOr = IRB.CreateOr(LHS: WIdx, RHS: WIdy);
879 Value *XYZOr = IRB.CreateOr(LHS: XYOr, RHS: WIdz);
880 Value *WIdzCond = IRB.CreateICmpEQ(LHS: XYZOr, RHS: IRB.getInt32(C: 0));
881
882 // All work items will branch to PrevEntryBlock except {0,0,0} index
883 // work item which will branch to malloc block.
884 IRB.CreateCondBr(Cond: WIdzCond, True: MallocBlock, False: PrevEntryBlock);
885
886 // Malloc block
887 IRB.SetInsertPoint(MallocBlock->begin());
888
889 // If Dynamic LDS globals are accessed by the kernel,
890 // Get the size of dyn lds from hidden dyn_lds_size kernel arg.
891 // Update the corresponding metadata global entries for this dyn lds global.
892 GlobalVariable *SwLDS = LDSParams.SwLDS;
893 GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
894 assert(SwLDS && SwLDSMetadata);
895 StructType *MetadataStructType =
896 cast<StructType>(Val: SwLDSMetadata->getValueType());
897 Type *Int32Ty = IRB.getInt32Ty();
898 Type *Int64Ty = IRB.getInt64Ty();
899
900 SetVector<GlobalVariable *> UniqueLDSGlobals;
901 auto GetUniqueLDSGlobals = [&](SetVector<GlobalVariable *> &LDSGlobals) {
902 for (auto &GV : LDSGlobals) {
903 if (is_contained(Range&: UniqueLDSGlobals, Element: GV))
904 continue;
905 UniqueLDSGlobals.insert(X: GV);
906 }
907 };
908
909 GetUniqueLDSGlobals(LDSParams.DirectAccess.StaticLDSGlobals);
910 GetUniqueLDSGlobals(LDSParams.IndirectAccess.StaticLDSGlobals);
911 // The metadata global always has an item for the SwLDS pointer itself, so
912 // there is at least one static item and the last one ends the static region.
913 unsigned LastStaticLDSIdx = UniqueLDSGlobals.size();
914 UniqueLDSGlobals.clear();
915
916 auto *GEPForEndStaticLDSOffset =
917 IRB.CreateInBoundsGEP(Ty: MetadataStructType, Ptr: SwLDSMetadata,
918 IdxList: {ConstantInt::get(Ty: Int32Ty, V: 0),
919 ConstantInt::get(Ty: Int32Ty, V: LastStaticLDSIdx),
920 ConstantInt::get(Ty: Int32Ty, V: 0)});
921
922 auto *GEPForEndStaticLDSSize =
923 IRB.CreateInBoundsGEP(Ty: MetadataStructType, Ptr: SwLDSMetadata,
924 IdxList: {ConstantInt::get(Ty: Int32Ty, V: 0),
925 ConstantInt::get(Ty: Int32Ty, V: LastStaticLDSIdx),
926 ConstantInt::get(Ty: Int32Ty, V: 2)});
927
928 Value *EndStaticLDSOffset = IRB.CreateLoad(Ty: Int32Ty, Ptr: GEPForEndStaticLDSOffset);
929 Value *EndStaticLDSSize = IRB.CreateLoad(Ty: Int32Ty, Ptr: GEPForEndStaticLDSSize);
930 Value *CurrMallocSize = IRB.CreateAdd(LHS: EndStaticLDSOffset, RHS: EndStaticLDSSize);
931
932 if (LDSParams.SwDynLDS) {
933 if (!(AMDGPU::getAMDHSACodeObjectVersion(M) >= AMDGPU::AMDHSA_COV5))
934 report_fatal_error(
935 reason: "Dynamic LDS size query is only supported for CO V5 and later.");
936 // Get size from hidden dyn_lds_size argument of kernel
937 Value *ImplicitArg =
938 IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_implicitarg_ptr, Args: {});
939 Value *HiddenDynLDSSize = IRB.CreateInBoundsGEP(
940 Ty: ImplicitArg->getType(), Ptr: ImplicitArg,
941 IdxList: {ConstantInt::get(Ty: Int64Ty, COV5_HIDDEN_DYN_LDS_SIZE_ARG)});
942 UniqueLDSGlobals.clear();
943 GetUniqueLDSGlobals(LDSParams.DirectAccess.DynamicLDSGlobals);
944 GetUniqueLDSGlobals(LDSParams.IndirectAccess.DynamicLDSGlobals);
945 updateMallocSizeForDynamicLDS(Func, CurrMallocSize: &CurrMallocSize, HiddenDynLDSSize,
946 DynamicLDSGlobals&: UniqueLDSGlobals);
947 }
948
949 CurrMallocSize = IRB.CreateZExt(V: CurrMallocSize, DestTy: Int64Ty);
950
951 // Create a call to malloc function which does device global memory allocation
952 // with size equals to all LDS global accesses size in this kernel.
953 Value *ReturnAddress = IRB.CreateIntrinsic(
954 ID: Intrinsic::returnaddress, OverloadTypes: IRB.getPtrTy(AddrSpace: DL.getProgramAddressSpace()),
955 Args: {IRB.getInt32(C: 0)});
956 FunctionCallee MallocFunc = M.getOrInsertFunction(
957 Name: StringRef("__asan_malloc_impl"),
958 T: FunctionType::get(Result: Int64Ty, Params: {Int64Ty, Int64Ty}, isVarArg: false));
959 Value *RAPtrToInt = IRB.CreatePtrToInt(V: ReturnAddress, DestTy: Int64Ty);
960 Value *MallocCall = IRB.CreateCall(Callee: MallocFunc, Args: {CurrMallocSize, RAPtrToInt});
961
962 Value *MallocPtr =
963 IRB.CreateIntToPtr(V: MallocCall, DestTy: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS));
964
965 // Create store of malloc to new global
966 IRB.CreateStore(Val: MallocPtr, Ptr: SwLDS);
967
968 // Create calls to __asan_poison_region to poison redzones.
969 poisonRedzones(Func, MallocPtr);
970
971 // Create branch to PrevEntryBlock
972 IRB.CreateBr(Dest: PrevEntryBlock);
973
974 // Create wave-group barrier at the starting of Previous entry block
975 Type *Int1Ty = IRB.getInt1Ty();
976 IRB.SetInsertPoint(PrevEntryBlock->begin());
977 auto *XYZCondPhi = IRB.CreatePHI(Ty: Int1Ty, NumReservedValues: 2, Name: "xyzCond");
978 XYZCondPhi->addIncoming(V: IRB.getInt1(V: 0), BB: WIdBlock);
979 XYZCondPhi->addIncoming(V: IRB.getInt1(V: 1), BB: MallocBlock);
980
981 IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_s_barrier, Args: {});
982
983 // Load malloc pointer from Sw LDS.
984 Value *LoadMallocPtr =
985 IRB.CreateLoad(Ty: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), Ptr: SwLDS);
986
987 // Replace All uses of LDS globals with new LDS pointers.
988 replaceKernelLDSAccesses(Func);
989
990 // Replace Memory Operations on LDS with corresponding
991 // global memory pointers.
992 translateLDSMemoryOperationsToGlobalMemory(Func, LoadMallocPtr,
993 LDSInstructions);
994
995 auto *CondFreeBlock = BasicBlock::Create(Context&: Ctx, Name: "CondFree", Parent: Func);
996 auto *FreeBlock = BasicBlock::Create(Context&: Ctx, Name: "Free", Parent: Func);
997 auto *EndBlock = BasicBlock::Create(Context&: Ctx, Name: "End", Parent: Func);
998 for (BasicBlock &BB : *Func) {
999 if (!BB.empty()) {
1000 if (ReturnInst *RI = dyn_cast<ReturnInst>(Val: &BB.back())) {
1001 RI->eraseFromParent();
1002 IRB.SetInsertPoint(BB.end());
1003 IRB.CreateBr(Dest: CondFreeBlock);
1004 }
1005 }
1006 }
1007
1008 // Cond Free Block
1009 IRB.SetInsertPoint(CondFreeBlock->begin());
1010 IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_s_barrier, Args: {});
1011 IRB.CreateCondBr(Cond: XYZCondPhi, True: FreeBlock, False: EndBlock);
1012
1013 // Free Block
1014 IRB.SetInsertPoint(FreeBlock->begin());
1015
1016 // Free the previously allocate device global memory.
1017 FunctionCallee AsanFreeFunc = M.getOrInsertFunction(
1018 Name: StringRef("__asan_free_impl"),
1019 T: FunctionType::get(Result: IRB.getVoidTy(), Params: {Int64Ty, Int64Ty}, isVarArg: false));
1020 Value *ReturnAddr = IRB.CreateIntrinsic(
1021 ID: Intrinsic::returnaddress, OverloadTypes: IRB.getPtrTy(AddrSpace: DL.getProgramAddressSpace()),
1022 Args: IRB.getInt32(C: 0));
1023 Value *RAPToInt = IRB.CreatePtrToInt(V: ReturnAddr, DestTy: Int64Ty);
1024 Value *MallocPtrToInt = IRB.CreatePtrToInt(V: LoadMallocPtr, DestTy: Int64Ty);
1025 IRB.CreateCall(Callee: AsanFreeFunc, Args: {MallocPtrToInt, RAPToInt});
1026
1027 IRB.CreateBr(Dest: EndBlock);
1028
1029 // End Block
1030 IRB.SetInsertPoint(EndBlock->begin());
1031 IRB.CreateRetVoid();
1032 // Update the DomTree with corresponding links to basic blocks.
1033 DTU.applyUpdates(Updates: {{DominatorTree::Insert, WIdBlock, MallocBlock},
1034 {DominatorTree::Insert, MallocBlock, PrevEntryBlock},
1035 {DominatorTree::Insert, CondFreeBlock, FreeBlock},
1036 {DominatorTree::Insert, FreeBlock, EndBlock}});
1037}
1038
1039Constant *AMDGPUSwLowerLDS::getAddressesOfVariablesInKernel(
1040 Function *Func, SetVector<GlobalVariable *> &Variables) {
1041 Type *Int32Ty = IRB.getInt32Ty();
1042 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
1043
1044 GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
1045 assert(SwLDSMetadata);
1046 auto *SwLDSMetadataStructType =
1047 cast<StructType>(Val: SwLDSMetadata->getValueType());
1048 ArrayType *KernelOffsetsType =
1049 ArrayType::get(ElementType: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), NumElements: Variables.size());
1050
1051 SmallVector<Constant *> Elements;
1052 for (auto *GV : Variables) {
1053 auto It = LDSParams.LDSToReplacementIndicesMap.find(Val: GV);
1054 if (It == LDSParams.LDSToReplacementIndicesMap.end()) {
1055 Elements.push_back(
1056 Elt: PoisonValue::get(T: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS)));
1057 continue;
1058 }
1059 auto &Indices = It->second;
1060 Constant *GEPIdx[] = {ConstantInt::get(Ty: Int32Ty, V: Indices[0]),
1061 ConstantInt::get(Ty: Int32Ty, V: Indices[1]),
1062 ConstantInt::get(Ty: Int32Ty, V: Indices[2])};
1063 Constant *GEP = ConstantExpr::getGetElementPtr(
1064 DL: Func->getDataLayout(), Ty: SwLDSMetadataStructType, C: SwLDSMetadata, IdxList: GEPIdx,
1065 NW: GEPNoWrapFlags::inBounds());
1066 Elements.push_back(Elt: GEP);
1067 }
1068 return ConstantArray::get(T: KernelOffsetsType, V: Elements);
1069}
1070
1071void AMDGPUSwLowerLDS::buildNonKernelLDSBaseTable(
1072 NonKernelLDSParameters &NKLDSParams) {
1073 // Base table will have single row, with elements of the row
1074 // placed as per kernel ID. Each element in the row corresponds
1075 // to addresss of "SW LDS" global of the kernel.
1076 auto &Kernels = NKLDSParams.OrderedKernels;
1077 if (Kernels.empty())
1078 return;
1079 const size_t NumberKernels = Kernels.size();
1080 ArrayType *AllKernelsOffsetsType =
1081 ArrayType::get(ElementType: IRB.getPtrTy(AddrSpace: AMDGPUAS::LOCAL_ADDRESS), NumElements: NumberKernels);
1082 std::vector<Constant *> OverallConstantExprElts(NumberKernels);
1083 for (size_t i = 0; i < NumberKernels; i++) {
1084 Function *Func = Kernels[i];
1085 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
1086 OverallConstantExprElts[i] = LDSParams.SwLDS;
1087 }
1088 Constant *init =
1089 ConstantArray::get(T: AllKernelsOffsetsType, V: OverallConstantExprElts);
1090 NKLDSParams.LDSBaseTable = new GlobalVariable(
1091 M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, init,
1092 "llvm.amdgcn.sw.lds.base.table", nullptr, GlobalValue::NotThreadLocal,
1093 AMDGPUAS::GLOBAL_ADDRESS);
1094 GlobalValue::SanitizerMetadata MD;
1095 MD.NoAddress = true;
1096 NKLDSParams.LDSBaseTable->setSanitizerMetadata(MD);
1097}
1098
1099void AMDGPUSwLowerLDS::buildNonKernelLDSOffsetTable(
1100 NonKernelLDSParameters &NKLDSParams) {
1101 // Offset table will have multiple rows and columns.
1102 // Rows are assumed to be from 0 to (n-1). n is total number
1103 // of kernels accessing the LDS through non-kernels.
1104 // Each row will have m elements. m is the total number of
1105 // unique LDS globals accessed by non-kernels.
1106 // Each element in the row correspond to the address of
1107 // the replacement of LDS global done by that particular kernel.
1108 auto &Variables = NKLDSParams.OrdereLDSGlobals;
1109 auto &Kernels = NKLDSParams.OrderedKernels;
1110 if (Variables.empty() || Kernels.empty())
1111 return;
1112 const size_t NumberVariables = Variables.size();
1113 const size_t NumberKernels = Kernels.size();
1114
1115 ArrayType *KernelOffsetsType =
1116 ArrayType::get(ElementType: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), NumElements: NumberVariables);
1117
1118 ArrayType *AllKernelsOffsetsType =
1119 ArrayType::get(ElementType: KernelOffsetsType, NumElements: NumberKernels);
1120 std::vector<Constant *> overallConstantExprElts(NumberKernels);
1121 for (size_t i = 0; i < NumberKernels; i++) {
1122 Function *Func = Kernels[i];
1123 overallConstantExprElts[i] =
1124 getAddressesOfVariablesInKernel(Func, Variables);
1125 }
1126 Constant *Init =
1127 ConstantArray::get(T: AllKernelsOffsetsType, V: overallConstantExprElts);
1128 NKLDSParams.LDSOffsetTable = new GlobalVariable(
1129 M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, Init,
1130 "llvm.amdgcn.sw.lds.offset.table", nullptr, GlobalValue::NotThreadLocal,
1131 AMDGPUAS::GLOBAL_ADDRESS);
1132 GlobalValue::SanitizerMetadata MD;
1133 MD.NoAddress = true;
1134 NKLDSParams.LDSOffsetTable->setSanitizerMetadata(MD);
1135}
1136
1137void AMDGPUSwLowerLDS::lowerNonKernelLDSAccesses(
1138 Function *Func, SetVector<GlobalVariable *> &LDSGlobals,
1139 NonKernelLDSParameters &NKLDSParams) {
1140 // Replace LDS access in non-kernel with replacement queried from
1141 // Base table and offset from offset table.
1142 LLVM_DEBUG(dbgs() << "Sw LDS lowering, lower non-kernel access for : "
1143 << Func->getName());
1144 auto InsertAt = Func->getEntryBlock().getFirstNonPHIOrAlloca();
1145 IRB.SetInsertPoint(InsertAt);
1146
1147 // Get LDS memory instructions.
1148 SetVector<Instruction *> LDSInstructions;
1149 getLDSMemoryInstructions(Func, LDSInstructions);
1150
1151 auto *KernelId = IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_lds_kernel_id, Args: {});
1152 GlobalVariable *LDSBaseTable = NKLDSParams.LDSBaseTable;
1153 GlobalVariable *LDSOffsetTable = NKLDSParams.LDSOffsetTable;
1154 auto &OrdereLDSGlobals = NKLDSParams.OrdereLDSGlobals;
1155 Value *BaseGEP = IRB.CreateInBoundsGEP(
1156 Ty: LDSBaseTable->getValueType(), Ptr: LDSBaseTable, IdxList: {IRB.getInt32(C: 0), KernelId});
1157 Value *BaseLoad =
1158 IRB.CreateLoad(Ty: IRB.getPtrTy(AddrSpace: AMDGPUAS::LOCAL_ADDRESS), Ptr: BaseGEP);
1159 Value *LoadMallocPtr =
1160 IRB.CreateLoad(Ty: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), Ptr: BaseLoad);
1161
1162 for (GlobalVariable *GV : LDSGlobals) {
1163 const auto *GVIt = llvm::find(Range&: OrdereLDSGlobals, Val: GV);
1164 assert(GVIt != OrdereLDSGlobals.end());
1165 uint32_t GVOffset = std::distance(first: OrdereLDSGlobals.begin(), last: GVIt);
1166
1167 Value *OffsetGEP = IRB.CreateInBoundsGEP(
1168 Ty: LDSOffsetTable->getValueType(), Ptr: LDSOffsetTable,
1169 IdxList: {IRB.getInt32(C: 0), KernelId, IRB.getInt32(C: GVOffset)});
1170 Value *OffsetLoad =
1171 IRB.CreateLoad(Ty: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), Ptr: OffsetGEP);
1172 Value *Offset = IRB.CreateLoad(Ty: IRB.getInt32Ty(), Ptr: OffsetLoad);
1173 Value *BasePlusOffset =
1174 IRB.CreateInBoundsGEP(Ty: IRB.getInt8Ty(), Ptr: BaseLoad, IdxList: {Offset});
1175 LLVM_DEBUG(dbgs() << "Sw LDS Lowering, Replace non-kernel LDS for "
1176 << GV->getName());
1177 replacesUsesOfGlobalInFunction(Func, GV, Replacement: BasePlusOffset);
1178 }
1179 translateLDSMemoryOperationsToGlobalMemory(Func, LoadMallocPtr,
1180 LDSInstructions);
1181}
1182
1183static void reorderStaticDynamicIndirectLDSSet(KernelLDSParameters &LDSParams) {
1184 // Sort Static, dynamic LDS globals which are either
1185 // direct or indirect access on basis of name.
1186 auto &DirectAccess = LDSParams.DirectAccess;
1187 auto &IndirectAccess = LDSParams.IndirectAccess;
1188 LDSParams.DirectAccess.StaticLDSGlobals = sortByName(
1189 V: std::vector<GlobalVariable *>(DirectAccess.StaticLDSGlobals.begin(),
1190 DirectAccess.StaticLDSGlobals.end()));
1191 LDSParams.DirectAccess.DynamicLDSGlobals = sortByName(
1192 V: std::vector<GlobalVariable *>(DirectAccess.DynamicLDSGlobals.begin(),
1193 DirectAccess.DynamicLDSGlobals.end()));
1194 LDSParams.IndirectAccess.StaticLDSGlobals = sortByName(
1195 V: std::vector<GlobalVariable *>(IndirectAccess.StaticLDSGlobals.begin(),
1196 IndirectAccess.StaticLDSGlobals.end()));
1197 LDSParams.IndirectAccess.DynamicLDSGlobals = sortByName(
1198 V: std::vector<GlobalVariable *>(IndirectAccess.DynamicLDSGlobals.begin(),
1199 IndirectAccess.DynamicLDSGlobals.end()));
1200}
1201
1202void AMDGPUSwLowerLDS::initAsanInfo() {
1203 // Get Shadow mapping scale and offset.
1204 unsigned LongSize =
1205 M.getDataLayout().getPointerSizeInBits(AS: AMDGPUAS::GLOBAL_ADDRESS);
1206 uint64_t Offset;
1207 int Scale;
1208 bool OrShadowOffset;
1209 llvm::getAddressSanitizerParams(Ctx: M.getContext(), TargetTriple: M.getTargetTriple(), LongSize,
1210 IsKasan: false, ShadowBase: &Offset, MappingScale: &Scale, OrShadowOffset: &OrShadowOffset);
1211 AsanInfo.Scale = Scale;
1212 AsanInfo.Offset = Offset;
1213}
1214
1215static bool hasFnWithSanitizeAddressAttr(FunctionVariableMap &LDSAccesses) {
1216 for (auto &K : LDSAccesses) {
1217 Function *F = K.first;
1218 if (!F)
1219 continue;
1220 if (F->hasFnAttribute(Kind: Attribute::SanitizeAddress))
1221 return true;
1222 }
1223 return false;
1224}
1225
1226bool AMDGPUSwLowerLDS::run() {
1227 bool Changed = false;
1228
1229 CallGraph CG = CallGraph(M);
1230
1231 Changed |=
1232 eliminateGVConstantExprUsesFromAllInstructions(M, Filter: isLDSVariableToLower);
1233
1234 // Get all the direct and indirect access of LDS for all the kernels.
1235 GVUsesInfoTy LDSUsesInfo = getTransitiveUsesOfLDSForLowering(CG, M);
1236
1237 // Flag to decide whether to lower all the LDS accesses
1238 // based on sanitize_address attribute.
1239 bool LowerAllLDS = hasFnWithSanitizeAddressAttr(LDSAccesses&: LDSUsesInfo.DirectAccess) ||
1240 hasFnWithSanitizeAddressAttr(LDSAccesses&: LDSUsesInfo.IndirectAccess);
1241
1242 if (!LowerAllLDS)
1243 return Changed;
1244
1245 // Utility to group LDS access into direct, indirect, static and dynamic.
1246 auto PopulateKernelStaticDynamicLDS = [&](FunctionVariableMap &LDSAccesses,
1247 bool DirectAccess) {
1248 for (auto &K : LDSAccesses) {
1249 Function *F = K.first;
1250 if (!F || K.second.empty())
1251 continue;
1252
1253 assert(isKernel(*F));
1254
1255 // Only inserts if key isn't already in the map.
1256 FuncLDSAccessInfo.KernelToLDSParametersMap.insert(
1257 KV: {F, KernelLDSParameters()});
1258
1259 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[F];
1260 if (!DirectAccess)
1261 FuncLDSAccessInfo.KernelsWithIndirectLDSAccess.insert(X: F);
1262 for (GlobalVariable *GV : K.second) {
1263 if (!DirectAccess) {
1264 if (AMDGPU::isDynamicLDS(GV: *GV))
1265 LDSParams.IndirectAccess.DynamicLDSGlobals.insert(X: GV);
1266 else
1267 LDSParams.IndirectAccess.StaticLDSGlobals.insert(X: GV);
1268 FuncLDSAccessInfo.AllNonKernelLDSAccess.insert(X: GV);
1269 } else {
1270 if (AMDGPU::isDynamicLDS(GV: *GV))
1271 LDSParams.DirectAccess.DynamicLDSGlobals.insert(X: GV);
1272 else
1273 LDSParams.DirectAccess.StaticLDSGlobals.insert(X: GV);
1274 }
1275 }
1276 }
1277 };
1278
1279 PopulateKernelStaticDynamicLDS(LDSUsesInfo.DirectAccess, true);
1280 PopulateKernelStaticDynamicLDS(LDSUsesInfo.IndirectAccess, false);
1281
1282 // Get address sanitizer scale.
1283 initAsanInfo();
1284
1285 for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
1286 Function *Func = K.first;
1287 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
1288 if (LDSParams.DirectAccess.StaticLDSGlobals.empty() &&
1289 LDSParams.DirectAccess.DynamicLDSGlobals.empty() &&
1290 LDSParams.IndirectAccess.StaticLDSGlobals.empty() &&
1291 LDSParams.IndirectAccess.DynamicLDSGlobals.empty())
1292 continue;
1293
1294 removeFnAttrFromReachable(
1295 CG, KernelRoot: Func,
1296 FnAttrs: {"amdgpu-no-workitem-id-x", "amdgpu-no-workitem-id-y",
1297 "amdgpu-no-workitem-id-z", "amdgpu-no-heap-ptr"});
1298 if (!LDSParams.IndirectAccess.StaticLDSGlobals.empty() ||
1299 !LDSParams.IndirectAccess.DynamicLDSGlobals.empty())
1300 removeFnAttrFromReachable(CG, KernelRoot: Func, FnAttrs: {"amdgpu-no-lds-kernel-id"});
1301 reorderStaticDynamicIndirectLDSSet(LDSParams);
1302 buildSwLDSGlobal(Func);
1303 buildSwDynLDSGlobal(Func);
1304 populateSwMetadataGlobal(Func);
1305 populateSwLDSAttributeAndMetadata(Func);
1306 populateLDSToReplacementIndicesMap(Func);
1307 DomTreeUpdater DTU(DTCallback(*Func), DomTreeUpdater::UpdateStrategy::Lazy);
1308 lowerKernelLDSAccesses(Func, DTU);
1309 Changed = true;
1310 }
1311
1312 // Get the Uses of LDS from non-kernels.
1313 getUsesOfLDSByNonKernels();
1314
1315 // Get non-kernels with LDS ptr as argument and called by kernels.
1316 getNonKernelsWithLDSArguments(CG);
1317
1318 // Lower LDS accesses in non-kernels.
1319 if (!FuncLDSAccessInfo.NonKernelToLDSAccessMap.empty() ||
1320 !FuncLDSAccessInfo.NonKernelsWithLDSArgument.empty()) {
1321 NonKernelLDSParameters NKLDSParams;
1322 NKLDSParams.OrderedKernels = getOrderedIndirectLDSAccessingKernels(
1323 Kernels&: FuncLDSAccessInfo.KernelsWithIndirectLDSAccess);
1324 NKLDSParams.OrdereLDSGlobals = getOrderedNonKernelAllLDSGlobals(
1325 Variables&: FuncLDSAccessInfo.AllNonKernelLDSAccess);
1326 buildNonKernelLDSBaseTable(NKLDSParams);
1327 buildNonKernelLDSOffsetTable(NKLDSParams);
1328 for (auto &K : FuncLDSAccessInfo.NonKernelToLDSAccessMap) {
1329 Function *Func = K.first;
1330 DenseSet<GlobalVariable *> &LDSGlobals = K.second;
1331 SetVector<GlobalVariable *> OrderedLDSGlobals = sortByName(
1332 V: std::vector<GlobalVariable *>(LDSGlobals.begin(), LDSGlobals.end()));
1333 lowerNonKernelLDSAccesses(Func, LDSGlobals&: OrderedLDSGlobals, NKLDSParams);
1334 }
1335 for (Function *Func : FuncLDSAccessInfo.NonKernelsWithLDSArgument) {
1336 auto &K = FuncLDSAccessInfo.NonKernelToLDSAccessMap;
1337 if (K.contains(Key: Func))
1338 continue;
1339 SetVector<llvm::GlobalVariable *> Vec;
1340 lowerNonKernelLDSAccesses(Func, LDSGlobals&: Vec, NKLDSParams);
1341 }
1342 Changed = true;
1343 }
1344
1345 if (!Changed)
1346 return Changed;
1347
1348 for (auto &GV : make_early_inc_range(Range: M.globals())) {
1349 if (AMDGPU::isLDSVariableToLower(GV)) {
1350 // probably want to remove from used lists
1351 GV.removeDeadConstantUsers();
1352 if (GV.use_empty())
1353 GV.eraseFromParent();
1354 }
1355 }
1356
1357 if (AsanInstrumentLDS) {
1358 SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
1359 for (Instruction *Inst : AsanInfo.Instructions) {
1360 SmallVector<InterestingMemoryOperand, 1> InterestingOperands;
1361 getInterestingMemoryOperands(M, I: Inst, Interesting&: InterestingOperands);
1362 llvm::append_range(C&: OperandsToInstrument, R&: InterestingOperands);
1363 }
1364 for (auto &Operand : OperandsToInstrument) {
1365 Value *Addr = Operand.getPtr();
1366 instrumentAddress(M, IRB, OrigIns: Operand.getInsn(), InsertBefore: Operand.getInsn(), Addr,
1367 Alignment: Operand.Alignment.valueOrOne(), TypeStoreSize: Operand.TypeStoreSize,
1368 IsWrite: Operand.IsWrite, SizeArgument: nullptr, UseCalls: false, Recover: false, Scale: AsanInfo.Scale,
1369 Offset: AsanInfo.Offset);
1370 Changed = true;
1371 }
1372 }
1373
1374 return Changed;
1375}
1376
1377class AMDGPUSwLowerLDSLegacy : public ModulePass {
1378public:
1379 static char ID;
1380 AMDGPUSwLowerLDSLegacy() : ModulePass(ID) {}
1381 bool runOnModule(Module &M) override;
1382 void getAnalysisUsage(AnalysisUsage &AU) const override {
1383 AU.addPreserved<DominatorTreeWrapperPass>();
1384 }
1385};
1386} // namespace
1387
1388char AMDGPUSwLowerLDSLegacy::ID = 0;
1389char &llvm::AMDGPUSwLowerLDSLegacyPassID = AMDGPUSwLowerLDSLegacy::ID;
1390
1391INITIALIZE_PASS_BEGIN(AMDGPUSwLowerLDSLegacy, "amdgpu-sw-lower-lds",
1392 "AMDGPU Software lowering of LDS", false, false)
1393INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
1394INITIALIZE_PASS_END(AMDGPUSwLowerLDSLegacy, "amdgpu-sw-lower-lds",
1395 "AMDGPU Software lowering of LDS", false, false)
1396
1397bool AMDGPUSwLowerLDSLegacy::runOnModule(Module &M) {
1398 // AddressSanitizer pass adds "nosanitize_address" module flag if it has
1399 // instrumented the IR. Return early if the flag is not present.
1400 if (!M.getModuleFlag(Key: "nosanitize_address"))
1401 return false;
1402 DominatorTreeWrapperPass *const DTW =
1403 getAnalysisIfAvailable<DominatorTreeWrapperPass>();
1404 auto DTCallback = [&DTW](Function &F) -> DominatorTree * {
1405 return DTW ? &DTW->getDomTree() : nullptr;
1406 };
1407
1408 AMDGPUSwLowerLDS SwLowerLDSImpl(M, DTCallback);
1409 bool IsChanged = SwLowerLDSImpl.run();
1410 return IsChanged;
1411}
1412
1413ModulePass *llvm::createAMDGPUSwLowerLDSLegacyPass() {
1414 return new AMDGPUSwLowerLDSLegacy();
1415}
1416
1417PreservedAnalyses AMDGPUSwLowerLDSPass::run(Module &M,
1418 ModuleAnalysisManager &AM) {
1419 // AddressSanitizer pass adds "nosanitize_address" module flag if it has
1420 // instrumented the IR. Return early if the flag is not present.
1421 if (!M.getModuleFlag(Key: "nosanitize_address"))
1422 return PreservedAnalyses::all();
1423 auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
1424 auto DTCallback = [&FAM](Function &F) -> DominatorTree * {
1425 return &FAM.getResult<DominatorTreeAnalysis>(IR&: F);
1426 };
1427 AMDGPUSwLowerLDS SwLowerLDSImpl(M, DTCallback);
1428 bool IsChanged = SwLowerLDSImpl.run();
1429 if (!IsChanged)
1430 return PreservedAnalyses::all();
1431
1432 PreservedAnalyses PA;
1433 PA.preserve<DominatorTreeAnalysis>();
1434 return PA;
1435}
1436