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