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 DenseMap<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.getContext()), 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, 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 Ty: SwLDSMetadataStructType, C: SwLDSMetadata, IdxList: GEPIdx, NW: true);
559 Value *Offset = IRB.CreateLoad(Ty: Int32Ty, Ptr: GEP);
560 Value *BasePlusOffset =
561 IRB.CreateInBoundsGEP(Ty: IRB.getInt8Ty(), Ptr: SwLDS, IdxList: {Offset});
562 LLVM_DEBUG(GV->printAsOperand(dbgs() << "Sw LDS Lowering, Replacing LDS ",
563 false));
564 replacesUsesOfGlobalInFunction(Func, GV, Replacement: BasePlusOffset);
565 }
566 };
567 ReplaceLDSGlobalUses(DirectAccess.StaticLDSGlobals);
568 ReplaceLDSGlobalUses(IndirectAccess.StaticLDSGlobals);
569 ReplaceLDSGlobalUses(DirectAccess.DynamicLDSGlobals);
570 ReplaceLDSGlobalUses(IndirectAccess.DynamicLDSGlobals);
571}
572
573void AMDGPUSwLowerLDS::updateMallocSizeForDynamicLDS(
574 Function *Func, Value **CurrMallocSize, Value *HiddenDynLDSSize,
575 SetVector<GlobalVariable *> &DynamicLDSGlobals) {
576 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
577 Type *Int32Ty = IRB.getInt32Ty();
578
579 GlobalVariable *SwLDS = LDSParams.SwLDS;
580 GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
581 assert(SwLDS && SwLDSMetadata);
582 StructType *MetadataStructType =
583 cast<StructType>(Val: SwLDSMetadata->getValueType());
584 unsigned MaxAlignment = SwLDS->getAlign().valueOrOne().value();
585 Value *MaxAlignValue = IRB.getInt32(C: MaxAlignment);
586 Value *MaxAlignValueMinusOne = IRB.getInt32(C: MaxAlignment - 1);
587
588 for (GlobalVariable *DynGV : DynamicLDSGlobals) {
589 auto &Indices = LDSParams.LDSToReplacementIndicesMap[DynGV];
590 // Update the Offset metadata.
591 Constant *Index0 = ConstantInt::get(Ty: Int32Ty, V: 0);
592 Constant *Index1 = ConstantInt::get(Ty: Int32Ty, V: Indices[1]);
593
594 Constant *Index2Offset = ConstantInt::get(Ty: Int32Ty, V: 0);
595 auto *GEPForOffset = IRB.CreateInBoundsGEP(
596 Ty: MetadataStructType, Ptr: SwLDSMetadata, IdxList: {Index0, Index1, Index2Offset});
597
598 IRB.CreateStore(Val: *CurrMallocSize, Ptr: GEPForOffset);
599 // Update the size and Aligned Size metadata.
600 Constant *Index2Size = ConstantInt::get(Ty: Int32Ty, V: 1);
601 auto *GEPForSize = IRB.CreateInBoundsGEP(Ty: MetadataStructType, Ptr: SwLDSMetadata,
602 IdxList: {Index0, Index1, Index2Size});
603
604 Value *CurrDynLDSSize = IRB.CreateLoad(Ty: Int32Ty, Ptr: HiddenDynLDSSize);
605 IRB.CreateStore(Val: CurrDynLDSSize, Ptr: GEPForSize);
606 Constant *Index2AlignedSize = ConstantInt::get(Ty: Int32Ty, V: 2);
607 auto *GEPForAlignedSize = IRB.CreateInBoundsGEP(
608 Ty: MetadataStructType, Ptr: SwLDSMetadata, IdxList: {Index0, Index1, Index2AlignedSize});
609
610 Value *AlignedDynLDSSize =
611 IRB.CreateAdd(LHS: CurrDynLDSSize, RHS: MaxAlignValueMinusOne);
612 AlignedDynLDSSize = IRB.CreateUDiv(LHS: AlignedDynLDSSize, RHS: MaxAlignValue);
613 AlignedDynLDSSize = IRB.CreateMul(LHS: AlignedDynLDSSize, RHS: MaxAlignValue);
614 IRB.CreateStore(Val: AlignedDynLDSSize, Ptr: GEPForAlignedSize);
615
616 // Update the Current Malloc Size
617 *CurrMallocSize = IRB.CreateAdd(LHS: *CurrMallocSize, RHS: AlignedDynLDSSize);
618 }
619}
620
621static DebugLoc getOrCreateDebugLoc(const Instruction *InsertBefore,
622 DISubprogram *SP) {
623 assert(InsertBefore);
624 if (InsertBefore->getDebugLoc())
625 return InsertBefore->getDebugLoc();
626 if (SP)
627 return DILocation::get(Context&: SP->getContext(), Line: SP->getLine(), Column: 1, Scope: SP);
628 return DebugLoc();
629}
630
631void AMDGPUSwLowerLDS::getLDSMemoryInstructions(
632 Function *Func, SetVector<Instruction *> &LDSInstructions) {
633 for (BasicBlock &BB : *Func) {
634 for (Instruction &Inst : BB) {
635 if (LoadInst *LI = dyn_cast<LoadInst>(Val: &Inst)) {
636 if (LI->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
637 LDSInstructions.insert(X: &Inst);
638 } else if (StoreInst *SI = dyn_cast<StoreInst>(Val: &Inst)) {
639 if (SI->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
640 LDSInstructions.insert(X: &Inst);
641 } else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(Val: &Inst)) {
642 if (RMW->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
643 LDSInstructions.insert(X: &Inst);
644 } else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(Val: &Inst)) {
645 if (XCHG->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
646 LDSInstructions.insert(X: &Inst);
647 } else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Val: &Inst)) {
648 if (ASC->getSrcAddressSpace() == AMDGPUAS::LOCAL_ADDRESS &&
649 ASC->getDestAddressSpace() == AMDGPUAS::FLAT_ADDRESS)
650 LDSInstructions.insert(X: &Inst);
651 } else if (AnyMemIntrinsic *MI = dyn_cast<AnyMemIntrinsic>(Val: &Inst)) {
652 if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
653 LDSInstructions.insert(X: &Inst);
654 } else if (auto *MTI = dyn_cast<AnyMemTransferInst>(Val: MI)) {
655 if (MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
656 LDSInstructions.insert(X: &Inst);
657 }
658 } else
659 continue;
660 }
661 }
662}
663
664Value *AMDGPUSwLowerLDS::getTranslatedGlobalMemoryPtrOfLDS(Value *LoadMallocPtr,
665 Value *LDSPtr) {
666 assert(LDSPtr && "Invalid LDS pointer operand");
667 Type *LDSPtrType = LDSPtr->getType();
668 LLVMContext &Ctx = M.getContext();
669 const DataLayout &DL = M.getDataLayout();
670 Type *IntTy = DL.getIntPtrType(C&: Ctx, AddressSpace: AMDGPUAS::LOCAL_ADDRESS);
671 if (auto *VecPtrTy = dyn_cast<VectorType>(Val: LDSPtrType)) {
672 // Handle vector of pointers
673 ElementCount NumElements = VecPtrTy->getElementCount();
674 IntTy = VectorType::get(ElementType: IntTy, EC: NumElements);
675 }
676 Value *GepIndex = IRB.CreatePtrToInt(V: LDSPtr, DestTy: IntTy);
677 return IRB.CreateInBoundsGEP(Ty: IRB.getInt8Ty(), Ptr: LoadMallocPtr, IdxList: {GepIndex});
678}
679
680void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
681 Function *Func, Value *LoadMallocPtr,
682 SetVector<Instruction *> &LDSInstructions) {
683 LLVM_DEBUG(dbgs() << "Translating LDS memory operations to global memory : "
684 << Func->getName());
685 for (Instruction *Inst : LDSInstructions) {
686 IRB.SetInsertPoint(Inst);
687 if (LoadInst *LI = dyn_cast<LoadInst>(Val: Inst)) {
688 Value *LIOperand = LI->getPointerOperand();
689 Value *Replacement =
690 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: LIOperand);
691 LoadInst *NewLI =
692 IRB.CreateLoad(Ty: LI->getType(), Ptr: Replacement, Props: LI->getProperties());
693 AsanInfo.Instructions.insert(X: NewLI);
694 LI->replaceAllUsesWith(V: NewLI);
695 LI->eraseFromParent();
696 } else if (StoreInst *SI = dyn_cast<StoreInst>(Val: Inst)) {
697 Value *SIOperand = SI->getPointerOperand();
698 Value *Replacement =
699 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: SIOperand);
700 StoreInst *NewSI = IRB.CreateStore(Val: SI->getValueOperand(), Ptr: Replacement,
701 Props: SI->getProperties());
702 AsanInfo.Instructions.insert(X: NewSI);
703 SI->replaceAllUsesWith(V: NewSI);
704 SI->eraseFromParent();
705 } else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(Val: Inst)) {
706 Value *RMWPtrOperand = RMW->getPointerOperand();
707 Value *RMWValOperand = RMW->getValOperand();
708 Value *Replacement =
709 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: RMWPtrOperand);
710 AtomicRMWInst *NewRMW = IRB.CreateAtomicRMW(
711 Op: RMW->getOperation(), Ptr: Replacement, Val: RMWValOperand, Align: RMW->getAlign(),
712 Ordering: RMW->getOrdering(), SSID: RMW->getSyncScopeID());
713 NewRMW->setVolatile(RMW->isVolatile());
714 AsanInfo.Instructions.insert(X: NewRMW);
715 RMW->replaceAllUsesWith(V: NewRMW);
716 RMW->eraseFromParent();
717 } else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(Val: Inst)) {
718 Value *XCHGPtrOperand = XCHG->getPointerOperand();
719 Value *Replacement =
720 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: XCHGPtrOperand);
721 AtomicCmpXchgInst *NewXCHG = IRB.CreateAtomicCmpXchg(
722 Ptr: Replacement, Cmp: XCHG->getCompareOperand(), New: XCHG->getNewValOperand(),
723 Align: XCHG->getAlign(), SuccessOrdering: XCHG->getSuccessOrdering(),
724 FailureOrdering: XCHG->getFailureOrdering(), SSID: XCHG->getSyncScopeID());
725 NewXCHG->setVolatile(XCHG->isVolatile());
726 AsanInfo.Instructions.insert(X: NewXCHG);
727 XCHG->replaceAllUsesWith(V: NewXCHG);
728 XCHG->eraseFromParent();
729 } else if (AnyMemIntrinsic *MI = dyn_cast<AnyMemIntrinsic>(Val: Inst)) {
730 Value *NewDest = MI->getRawDest();
731 if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
732 NewDest = getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: NewDest);
733 CallInst *NewMI = nullptr;
734 if (AnyMemSetInst *MSI = dyn_cast<AnyMemSetInst>(Val: MI)) {
735 if (MI->isAtomic()) {
736 NewMI = IRB.CreateElementUnorderedAtomicMemSet(
737 Ptr: NewDest, Val: MSI->getValue(), Size: MSI->getLength(),
738 Alignment: MSI->getDestAlign().valueOrOne(), ElementSize: MSI->getElementSizeInBytes());
739 } else {
740 NewMI = IRB.CreateMemSet(Ptr: NewDest, Val: MSI->getValue(), Size: MSI->getLength(),
741 Align: MSI->getDestAlign(),
742 isVolatile: cast<MemSetInst>(Val: MI)->isVolatile());
743 }
744 } else if (AnyMemTransferInst *MTI = dyn_cast<AnyMemTransferInst>(Val: MI)) {
745 Value *NewSrc = MTI->getRawSource();
746 if (MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
747 NewSrc = getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: NewSrc);
748 if (MI->isAtomic()) {
749 if (MI->getIntrinsicID() ==
750 Intrinsic::memmove_element_unordered_atomic) {
751 NewMI = IRB.CreateElementUnorderedAtomicMemMove(
752 Dst: NewDest, DstAlign: MTI->getDestAlign().valueOrOne(), Src: NewSrc,
753 SrcAlign: MTI->getSourceAlign().valueOrOne(), Size: MTI->getLength(),
754 ElementSize: MTI->getElementSizeInBytes());
755 } else {
756 NewMI = IRB.CreateElementUnorderedAtomicMemCpy(
757 Dst: NewDest, DstAlign: MTI->getDestAlign().valueOrOne(), Src: NewSrc,
758 SrcAlign: MTI->getSourceAlign().valueOrOne(), Size: MTI->getLength(),
759 ElementSize: MTI->getElementSizeInBytes());
760 }
761 } else {
762 NewMI = IRB.CreateMemTransferInst(
763 IntrID: MI->getIntrinsicID(), Dst: NewDest, DstAlign: MTI->getDestAlign(), Src: NewSrc,
764 SrcAlign: MTI->getSourceAlign(), Size: MTI->getLength(),
765 isVolatile: cast<MemTransferInst>(Val: MI)->isVolatile());
766 }
767 } else
768 reportFatalUsageError(reason: "Unimplemented LDS lowering memory intrinsic");
769 AsanInfo.Instructions.insert(X: NewMI);
770 MI->replaceAllUsesWith(V: NewMI);
771 MI->eraseFromParent();
772 } else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Val: Inst)) {
773 Value *AIOperand = ASC->getPointerOperand();
774 Value *Replacement =
775 getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr: AIOperand);
776 Value *NewAI = IRB.CreateAddrSpaceCast(V: Replacement, DestTy: ASC->getType());
777 // Note: No need to add the instruction to AsanInfo instructions to be
778 // instrumented list. FLAT_ADDRESS ptr would have been already
779 // instrumented by asan pass prior to this pass.
780 ASC->replaceAllUsesWith(V: NewAI);
781 ASC->eraseFromParent();
782 } else
783 report_fatal_error(reason: "Unimplemented LDS lowering instruction");
784 }
785}
786
787void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
788 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
789 Type *Int64Ty = IRB.getInt64Ty();
790 Type *VoidTy = IRB.getVoidTy();
791 FunctionCallee AsanPoisonRegion = M.getOrInsertFunction(
792 Name: "__asan_poison_region",
793 T: FunctionType::get(Result: VoidTy, Params: {Int64Ty, Int64Ty}, isVarArg: false));
794
795 auto RedzonesVec = LDSParams.RedzoneOffsetAndSizeVector;
796 size_t VecSize = RedzonesVec.size();
797 for (unsigned i = 0; i < VecSize; i++) {
798 auto &RedzonePair = RedzonesVec[i];
799 uint64_t RedzoneOffset = RedzonePair.first;
800 uint64_t RedzoneSize = RedzonePair.second;
801 Value *RedzoneAddrOffset = IRB.CreateInBoundsGEP(
802 Ty: IRB.getInt8Ty(), Ptr: MallocPtr, IdxList: {IRB.getInt64(C: RedzoneOffset)});
803 Value *RedzoneAddress = IRB.CreatePtrToInt(V: RedzoneAddrOffset, DestTy: Int64Ty);
804 IRB.CreateCall(Callee: AsanPoisonRegion,
805 Args: {RedzoneAddress, IRB.getInt64(C: RedzoneSize)});
806 }
807}
808
809void AMDGPUSwLowerLDS::lowerKernelLDSAccesses(Function *Func,
810 DomTreeUpdater &DTU) {
811 LLVM_DEBUG(dbgs() << "Sw Lowering Kernel LDS for : " << Func->getName());
812 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
813 auto &Ctx = M.getContext();
814 auto *PrevEntryBlock = &Func->getEntryBlock();
815 SetVector<Instruction *> LDSInstructions;
816 getLDSMemoryInstructions(Func, LDSInstructions);
817 const DataLayout &DL = M.getDataLayout();
818
819 // Create malloc block.
820 auto *MallocBlock = BasicBlock::Create(Context&: Ctx, Name: "Malloc", Parent: Func, InsertBefore: PrevEntryBlock);
821
822 // Create WIdBlock block which has instructions related to selection of
823 // {0,0,0} indiex work item in the work group.
824 auto *WIdBlock = BasicBlock::Create(Context&: Ctx, Name: "WId", Parent: Func, InsertBefore: MallocBlock);
825
826 // Move constant-size allocas from the original entry block to the new entry
827 // block (WIdBlock) so they remain static allocas. Splice the leading cluster
828 // in bulk, then move any stragglers that are interleaved with other
829 // instructions.
830 auto SplitIt = PrevEntryBlock->getFirstNonPHIOrDbgOrAlloca();
831 WIdBlock->splice(ToIt: WIdBlock->end(), FromBB: PrevEntryBlock, FromBeginIt: PrevEntryBlock->begin(),
832 FromEndIt: SplitIt);
833 for (Instruction &I : make_early_inc_range(Range&: *PrevEntryBlock))
834 if (auto *AI = dyn_cast<AllocaInst>(Val: &I))
835 if (isa<ConstantInt>(Val: AI->getArraySize()))
836 AI->moveBefore(BB&: *WIdBlock, I: WIdBlock->end());
837
838 IRB.SetInsertPoint(TheBB: WIdBlock, IP: WIdBlock->end());
839 DebugLoc FirstDL =
840 getOrCreateDebugLoc(InsertBefore: &*PrevEntryBlock->begin(), SP: Func->getSubprogram());
841 IRB.SetCurrentDebugLocation(FirstDL);
842 Value *WIdx = IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_workitem_id_x, Args: {});
843 Value *WIdy = IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_workitem_id_y, Args: {});
844 Value *WIdz = IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_workitem_id_z, Args: {});
845 Value *XYOr = IRB.CreateOr(LHS: WIdx, RHS: WIdy);
846 Value *XYZOr = IRB.CreateOr(LHS: XYOr, RHS: WIdz);
847 Value *WIdzCond = IRB.CreateICmpEQ(LHS: XYZOr, RHS: IRB.getInt32(C: 0));
848
849 // All work items will branch to PrevEntryBlock except {0,0,0} index
850 // work item which will branch to malloc block.
851 IRB.CreateCondBr(Cond: WIdzCond, True: MallocBlock, False: PrevEntryBlock);
852
853 // Malloc block
854 IRB.SetInsertPoint(TheBB: MallocBlock, IP: MallocBlock->begin());
855
856 // If Dynamic LDS globals are accessed by the kernel,
857 // Get the size of dyn lds from hidden dyn_lds_size kernel arg.
858 // Update the corresponding metadata global entries for this dyn lds global.
859 GlobalVariable *SwLDS = LDSParams.SwLDS;
860 GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
861 assert(SwLDS && SwLDSMetadata);
862 StructType *MetadataStructType =
863 cast<StructType>(Val: SwLDSMetadata->getValueType());
864 Type *Int32Ty = IRB.getInt32Ty();
865 Type *Int64Ty = IRB.getInt64Ty();
866
867 SetVector<GlobalVariable *> UniqueLDSGlobals;
868 auto GetUniqueLDSGlobals = [&](SetVector<GlobalVariable *> &LDSGlobals) {
869 for (auto &GV : LDSGlobals) {
870 if (is_contained(Range&: UniqueLDSGlobals, Element: GV))
871 continue;
872 UniqueLDSGlobals.insert(X: GV);
873 }
874 };
875
876 GetUniqueLDSGlobals(LDSParams.DirectAccess.StaticLDSGlobals);
877 GetUniqueLDSGlobals(LDSParams.IndirectAccess.StaticLDSGlobals);
878 // The metadata global always has an item for the SwLDS pointer itself, so
879 // there is at least one static item and the last one ends the static region.
880 unsigned LastStaticLDSIdx = UniqueLDSGlobals.size();
881 UniqueLDSGlobals.clear();
882
883 auto *GEPForEndStaticLDSOffset =
884 IRB.CreateInBoundsGEP(Ty: MetadataStructType, Ptr: SwLDSMetadata,
885 IdxList: {ConstantInt::get(Ty: Int32Ty, V: 0),
886 ConstantInt::get(Ty: Int32Ty, V: LastStaticLDSIdx),
887 ConstantInt::get(Ty: Int32Ty, V: 0)});
888
889 auto *GEPForEndStaticLDSSize =
890 IRB.CreateInBoundsGEP(Ty: MetadataStructType, Ptr: SwLDSMetadata,
891 IdxList: {ConstantInt::get(Ty: Int32Ty, V: 0),
892 ConstantInt::get(Ty: Int32Ty, V: LastStaticLDSIdx),
893 ConstantInt::get(Ty: Int32Ty, V: 2)});
894
895 Value *EndStaticLDSOffset = IRB.CreateLoad(Ty: Int32Ty, Ptr: GEPForEndStaticLDSOffset);
896 Value *EndStaticLDSSize = IRB.CreateLoad(Ty: Int32Ty, Ptr: GEPForEndStaticLDSSize);
897 Value *CurrMallocSize = IRB.CreateAdd(LHS: EndStaticLDSOffset, RHS: EndStaticLDSSize);
898
899 if (LDSParams.SwDynLDS) {
900 if (!(AMDGPU::getAMDHSACodeObjectVersion(M) >= AMDGPU::AMDHSA_COV5))
901 report_fatal_error(
902 reason: "Dynamic LDS size query is only supported for CO V5 and later.");
903 // Get size from hidden dyn_lds_size argument of kernel
904 Value *ImplicitArg =
905 IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_implicitarg_ptr, Args: {});
906 Value *HiddenDynLDSSize = IRB.CreateInBoundsGEP(
907 Ty: ImplicitArg->getType(), Ptr: ImplicitArg,
908 IdxList: {ConstantInt::get(Ty: Int64Ty, COV5_HIDDEN_DYN_LDS_SIZE_ARG)});
909 UniqueLDSGlobals.clear();
910 GetUniqueLDSGlobals(LDSParams.DirectAccess.DynamicLDSGlobals);
911 GetUniqueLDSGlobals(LDSParams.IndirectAccess.DynamicLDSGlobals);
912 updateMallocSizeForDynamicLDS(Func, CurrMallocSize: &CurrMallocSize, HiddenDynLDSSize,
913 DynamicLDSGlobals&: UniqueLDSGlobals);
914 }
915
916 CurrMallocSize = IRB.CreateZExt(V: CurrMallocSize, DestTy: Int64Ty);
917
918 // Create a call to malloc function which does device global memory allocation
919 // with size equals to all LDS global accesses size in this kernel.
920 Value *ReturnAddress = IRB.CreateIntrinsic(
921 ID: Intrinsic::returnaddress, OverloadTypes: IRB.getPtrTy(AddrSpace: DL.getProgramAddressSpace()),
922 Args: {IRB.getInt32(C: 0)});
923 FunctionCallee MallocFunc = M.getOrInsertFunction(
924 Name: StringRef("__asan_malloc_impl"),
925 T: FunctionType::get(Result: Int64Ty, Params: {Int64Ty, Int64Ty}, isVarArg: false));
926 Value *RAPtrToInt = IRB.CreatePtrToInt(V: ReturnAddress, DestTy: Int64Ty);
927 Value *MallocCall = IRB.CreateCall(Callee: MallocFunc, Args: {CurrMallocSize, RAPtrToInt});
928
929 Value *MallocPtr =
930 IRB.CreateIntToPtr(V: MallocCall, DestTy: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS));
931
932 // Create store of malloc to new global
933 IRB.CreateStore(Val: MallocPtr, Ptr: SwLDS);
934
935 // Create calls to __asan_poison_region to poison redzones.
936 poisonRedzones(Func, MallocPtr);
937
938 // Create branch to PrevEntryBlock
939 IRB.CreateBr(Dest: PrevEntryBlock);
940
941 // Create wave-group barrier at the starting of Previous entry block
942 Type *Int1Ty = IRB.getInt1Ty();
943 IRB.SetInsertPoint(TheBB: PrevEntryBlock, IP: PrevEntryBlock->begin());
944 auto *XYZCondPhi = IRB.CreatePHI(Ty: Int1Ty, NumReservedValues: 2, Name: "xyzCond");
945 XYZCondPhi->addIncoming(V: IRB.getInt1(V: 0), BB: WIdBlock);
946 XYZCondPhi->addIncoming(V: IRB.getInt1(V: 1), BB: MallocBlock);
947
948 IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_s_barrier, Args: {});
949
950 // Load malloc pointer from Sw LDS.
951 Value *LoadMallocPtr =
952 IRB.CreateLoad(Ty: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), Ptr: SwLDS);
953
954 // Replace All uses of LDS globals with new LDS pointers.
955 replaceKernelLDSAccesses(Func);
956
957 // Replace Memory Operations on LDS with corresponding
958 // global memory pointers.
959 translateLDSMemoryOperationsToGlobalMemory(Func, LoadMallocPtr,
960 LDSInstructions);
961
962 auto *CondFreeBlock = BasicBlock::Create(Context&: Ctx, Name: "CondFree", Parent: Func);
963 auto *FreeBlock = BasicBlock::Create(Context&: Ctx, Name: "Free", Parent: Func);
964 auto *EndBlock = BasicBlock::Create(Context&: Ctx, Name: "End", Parent: Func);
965 for (BasicBlock &BB : *Func) {
966 if (!BB.empty()) {
967 if (ReturnInst *RI = dyn_cast<ReturnInst>(Val: &BB.back())) {
968 RI->eraseFromParent();
969 IRB.SetInsertPoint(TheBB: &BB, IP: BB.end());
970 IRB.CreateBr(Dest: CondFreeBlock);
971 }
972 }
973 }
974
975 // Cond Free Block
976 IRB.SetInsertPoint(TheBB: CondFreeBlock, IP: CondFreeBlock->begin());
977 IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_s_barrier, Args: {});
978 IRB.CreateCondBr(Cond: XYZCondPhi, True: FreeBlock, False: EndBlock);
979
980 // Free Block
981 IRB.SetInsertPoint(TheBB: FreeBlock, IP: FreeBlock->begin());
982
983 // Free the previously allocate device global memory.
984 FunctionCallee AsanFreeFunc = M.getOrInsertFunction(
985 Name: StringRef("__asan_free_impl"),
986 T: FunctionType::get(Result: IRB.getVoidTy(), Params: {Int64Ty, Int64Ty}, isVarArg: false));
987 Value *ReturnAddr = IRB.CreateIntrinsic(
988 ID: Intrinsic::returnaddress, OverloadTypes: IRB.getPtrTy(AddrSpace: DL.getProgramAddressSpace()),
989 Args: IRB.getInt32(C: 0));
990 Value *RAPToInt = IRB.CreatePtrToInt(V: ReturnAddr, DestTy: Int64Ty);
991 Value *MallocPtrToInt = IRB.CreatePtrToInt(V: LoadMallocPtr, DestTy: Int64Ty);
992 IRB.CreateCall(Callee: AsanFreeFunc, Args: {MallocPtrToInt, RAPToInt});
993
994 IRB.CreateBr(Dest: EndBlock);
995
996 // End Block
997 IRB.SetInsertPoint(TheBB: EndBlock, IP: EndBlock->begin());
998 IRB.CreateRetVoid();
999 // Update the DomTree with corresponding links to basic blocks.
1000 DTU.applyUpdates(Updates: {{DominatorTree::Insert, WIdBlock, MallocBlock},
1001 {DominatorTree::Insert, MallocBlock, PrevEntryBlock},
1002 {DominatorTree::Insert, CondFreeBlock, FreeBlock},
1003 {DominatorTree::Insert, FreeBlock, EndBlock}});
1004}
1005
1006Constant *AMDGPUSwLowerLDS::getAddressesOfVariablesInKernel(
1007 Function *Func, SetVector<GlobalVariable *> &Variables) {
1008 Type *Int32Ty = IRB.getInt32Ty();
1009 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
1010
1011 GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
1012 assert(SwLDSMetadata);
1013 auto *SwLDSMetadataStructType =
1014 cast<StructType>(Val: SwLDSMetadata->getValueType());
1015 ArrayType *KernelOffsetsType =
1016 ArrayType::get(ElementType: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), NumElements: Variables.size());
1017
1018 SmallVector<Constant *> Elements;
1019 for (auto *GV : Variables) {
1020 auto It = LDSParams.LDSToReplacementIndicesMap.find(Val: GV);
1021 if (It == LDSParams.LDSToReplacementIndicesMap.end()) {
1022 Elements.push_back(
1023 Elt: PoisonValue::get(T: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS)));
1024 continue;
1025 }
1026 auto &Indices = It->second;
1027 Constant *GEPIdx[] = {ConstantInt::get(Ty: Int32Ty, V: Indices[0]),
1028 ConstantInt::get(Ty: Int32Ty, V: Indices[1]),
1029 ConstantInt::get(Ty: Int32Ty, V: Indices[2])};
1030 Constant *GEP = ConstantExpr::getGetElementPtr(Ty: SwLDSMetadataStructType,
1031 C: SwLDSMetadata, IdxList: GEPIdx, NW: true);
1032 Elements.push_back(Elt: GEP);
1033 }
1034 return ConstantArray::get(T: KernelOffsetsType, V: Elements);
1035}
1036
1037void AMDGPUSwLowerLDS::buildNonKernelLDSBaseTable(
1038 NonKernelLDSParameters &NKLDSParams) {
1039 // Base table will have single row, with elements of the row
1040 // placed as per kernel ID. Each element in the row corresponds
1041 // to addresss of "SW LDS" global of the kernel.
1042 auto &Kernels = NKLDSParams.OrderedKernels;
1043 if (Kernels.empty())
1044 return;
1045 const size_t NumberKernels = Kernels.size();
1046 ArrayType *AllKernelsOffsetsType =
1047 ArrayType::get(ElementType: IRB.getPtrTy(AddrSpace: AMDGPUAS::LOCAL_ADDRESS), NumElements: NumberKernels);
1048 std::vector<Constant *> OverallConstantExprElts(NumberKernels);
1049 for (size_t i = 0; i < NumberKernels; i++) {
1050 Function *Func = Kernels[i];
1051 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
1052 OverallConstantExprElts[i] = LDSParams.SwLDS;
1053 }
1054 Constant *init =
1055 ConstantArray::get(T: AllKernelsOffsetsType, V: OverallConstantExprElts);
1056 NKLDSParams.LDSBaseTable = new GlobalVariable(
1057 M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, init,
1058 "llvm.amdgcn.sw.lds.base.table", nullptr, GlobalValue::NotThreadLocal,
1059 AMDGPUAS::GLOBAL_ADDRESS);
1060 GlobalValue::SanitizerMetadata MD;
1061 MD.NoAddress = true;
1062 NKLDSParams.LDSBaseTable->setSanitizerMetadata(MD);
1063}
1064
1065void AMDGPUSwLowerLDS::buildNonKernelLDSOffsetTable(
1066 NonKernelLDSParameters &NKLDSParams) {
1067 // Offset table will have multiple rows and columns.
1068 // Rows are assumed to be from 0 to (n-1). n is total number
1069 // of kernels accessing the LDS through non-kernels.
1070 // Each row will have m elements. m is the total number of
1071 // unique LDS globals accessed by non-kernels.
1072 // Each element in the row correspond to the address of
1073 // the replacement of LDS global done by that particular kernel.
1074 auto &Variables = NKLDSParams.OrdereLDSGlobals;
1075 auto &Kernels = NKLDSParams.OrderedKernels;
1076 if (Variables.empty() || Kernels.empty())
1077 return;
1078 const size_t NumberVariables = Variables.size();
1079 const size_t NumberKernels = Kernels.size();
1080
1081 ArrayType *KernelOffsetsType =
1082 ArrayType::get(ElementType: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), NumElements: NumberVariables);
1083
1084 ArrayType *AllKernelsOffsetsType =
1085 ArrayType::get(ElementType: KernelOffsetsType, NumElements: NumberKernels);
1086 std::vector<Constant *> overallConstantExprElts(NumberKernels);
1087 for (size_t i = 0; i < NumberKernels; i++) {
1088 Function *Func = Kernels[i];
1089 overallConstantExprElts[i] =
1090 getAddressesOfVariablesInKernel(Func, Variables);
1091 }
1092 Constant *Init =
1093 ConstantArray::get(T: AllKernelsOffsetsType, V: overallConstantExprElts);
1094 NKLDSParams.LDSOffsetTable = new GlobalVariable(
1095 M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, Init,
1096 "llvm.amdgcn.sw.lds.offset.table", nullptr, GlobalValue::NotThreadLocal,
1097 AMDGPUAS::GLOBAL_ADDRESS);
1098 GlobalValue::SanitizerMetadata MD;
1099 MD.NoAddress = true;
1100 NKLDSParams.LDSOffsetTable->setSanitizerMetadata(MD);
1101}
1102
1103void AMDGPUSwLowerLDS::lowerNonKernelLDSAccesses(
1104 Function *Func, SetVector<GlobalVariable *> &LDSGlobals,
1105 NonKernelLDSParameters &NKLDSParams) {
1106 // Replace LDS access in non-kernel with replacement queried from
1107 // Base table and offset from offset table.
1108 LLVM_DEBUG(dbgs() << "Sw LDS lowering, lower non-kernel access for : "
1109 << Func->getName());
1110 auto InsertAt = Func->getEntryBlock().getFirstNonPHIOrDbgOrAlloca();
1111 IRB.SetInsertPoint(InsertAt);
1112
1113 // Get LDS memory instructions.
1114 SetVector<Instruction *> LDSInstructions;
1115 getLDSMemoryInstructions(Func, LDSInstructions);
1116
1117 auto *KernelId = IRB.CreateIntrinsic(ID: Intrinsic::amdgcn_lds_kernel_id, Args: {});
1118 GlobalVariable *LDSBaseTable = NKLDSParams.LDSBaseTable;
1119 GlobalVariable *LDSOffsetTable = NKLDSParams.LDSOffsetTable;
1120 auto &OrdereLDSGlobals = NKLDSParams.OrdereLDSGlobals;
1121 Value *BaseGEP = IRB.CreateInBoundsGEP(
1122 Ty: LDSBaseTable->getValueType(), Ptr: LDSBaseTable, IdxList: {IRB.getInt32(C: 0), KernelId});
1123 Value *BaseLoad =
1124 IRB.CreateLoad(Ty: IRB.getPtrTy(AddrSpace: AMDGPUAS::LOCAL_ADDRESS), Ptr: BaseGEP);
1125 Value *LoadMallocPtr =
1126 IRB.CreateLoad(Ty: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), Ptr: BaseLoad);
1127
1128 for (GlobalVariable *GV : LDSGlobals) {
1129 const auto *GVIt = llvm::find(Range&: OrdereLDSGlobals, Val: GV);
1130 assert(GVIt != OrdereLDSGlobals.end());
1131 uint32_t GVOffset = std::distance(first: OrdereLDSGlobals.begin(), last: GVIt);
1132
1133 Value *OffsetGEP = IRB.CreateInBoundsGEP(
1134 Ty: LDSOffsetTable->getValueType(), Ptr: LDSOffsetTable,
1135 IdxList: {IRB.getInt32(C: 0), KernelId, IRB.getInt32(C: GVOffset)});
1136 Value *OffsetLoad =
1137 IRB.CreateLoad(Ty: IRB.getPtrTy(AddrSpace: AMDGPUAS::GLOBAL_ADDRESS), Ptr: OffsetGEP);
1138 Value *Offset = IRB.CreateLoad(Ty: IRB.getInt32Ty(), Ptr: OffsetLoad);
1139 Value *BasePlusOffset =
1140 IRB.CreateInBoundsGEP(Ty: IRB.getInt8Ty(), Ptr: BaseLoad, IdxList: {Offset});
1141 LLVM_DEBUG(dbgs() << "Sw LDS Lowering, Replace non-kernel LDS for "
1142 << GV->getName());
1143 replacesUsesOfGlobalInFunction(Func, GV, Replacement: BasePlusOffset);
1144 }
1145 translateLDSMemoryOperationsToGlobalMemory(Func, LoadMallocPtr,
1146 LDSInstructions);
1147}
1148
1149static void reorderStaticDynamicIndirectLDSSet(KernelLDSParameters &LDSParams) {
1150 // Sort Static, dynamic LDS globals which are either
1151 // direct or indirect access on basis of name.
1152 auto &DirectAccess = LDSParams.DirectAccess;
1153 auto &IndirectAccess = LDSParams.IndirectAccess;
1154 LDSParams.DirectAccess.StaticLDSGlobals = sortByName(
1155 V: std::vector<GlobalVariable *>(DirectAccess.StaticLDSGlobals.begin(),
1156 DirectAccess.StaticLDSGlobals.end()));
1157 LDSParams.DirectAccess.DynamicLDSGlobals = sortByName(
1158 V: std::vector<GlobalVariable *>(DirectAccess.DynamicLDSGlobals.begin(),
1159 DirectAccess.DynamicLDSGlobals.end()));
1160 LDSParams.IndirectAccess.StaticLDSGlobals = sortByName(
1161 V: std::vector<GlobalVariable *>(IndirectAccess.StaticLDSGlobals.begin(),
1162 IndirectAccess.StaticLDSGlobals.end()));
1163 LDSParams.IndirectAccess.DynamicLDSGlobals = sortByName(
1164 V: std::vector<GlobalVariable *>(IndirectAccess.DynamicLDSGlobals.begin(),
1165 IndirectAccess.DynamicLDSGlobals.end()));
1166}
1167
1168void AMDGPUSwLowerLDS::initAsanInfo() {
1169 // Get Shadow mapping scale and offset.
1170 unsigned LongSize =
1171 M.getDataLayout().getPointerSizeInBits(AS: AMDGPUAS::GLOBAL_ADDRESS);
1172 uint64_t Offset;
1173 int Scale;
1174 bool OrShadowOffset;
1175 llvm::getAddressSanitizerParams(TargetTriple: M.getTargetTriple(), LongSize, IsKasan: false, ShadowBase: &Offset,
1176 MappingScale: &Scale, OrShadowOffset: &OrShadowOffset);
1177 AsanInfo.Scale = Scale;
1178 AsanInfo.Offset = Offset;
1179}
1180
1181static bool hasFnWithSanitizeAddressAttr(FunctionVariableMap &LDSAccesses) {
1182 for (auto &K : LDSAccesses) {
1183 Function *F = K.first;
1184 if (!F)
1185 continue;
1186 if (F->hasFnAttribute(Kind: Attribute::SanitizeAddress))
1187 return true;
1188 }
1189 return false;
1190}
1191
1192bool AMDGPUSwLowerLDS::run() {
1193 bool Changed = false;
1194
1195 CallGraph CG = CallGraph(M);
1196
1197 Changed |=
1198 eliminateGVConstantExprUsesFromAllInstructions(M, Filter: isLDSVariableToLower);
1199
1200 // Get all the direct and indirect access of LDS for all the kernels.
1201 GVUsesInfoTy LDSUsesInfo = getTransitiveUsesOfLDSForLowering(CG, M);
1202
1203 // Flag to decide whether to lower all the LDS accesses
1204 // based on sanitize_address attribute.
1205 bool LowerAllLDS = hasFnWithSanitizeAddressAttr(LDSAccesses&: LDSUsesInfo.DirectAccess) ||
1206 hasFnWithSanitizeAddressAttr(LDSAccesses&: LDSUsesInfo.IndirectAccess);
1207
1208 if (!LowerAllLDS)
1209 return Changed;
1210
1211 // Utility to group LDS access into direct, indirect, static and dynamic.
1212 auto PopulateKernelStaticDynamicLDS = [&](FunctionVariableMap &LDSAccesses,
1213 bool DirectAccess) {
1214 for (auto &K : LDSAccesses) {
1215 Function *F = K.first;
1216 if (!F || K.second.empty())
1217 continue;
1218
1219 assert(isKernel(*F));
1220
1221 // Only inserts if key isn't already in the map.
1222 FuncLDSAccessInfo.KernelToLDSParametersMap.insert(
1223 KV: {F, KernelLDSParameters()});
1224
1225 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[F];
1226 if (!DirectAccess)
1227 FuncLDSAccessInfo.KernelsWithIndirectLDSAccess.insert(X: F);
1228 for (GlobalVariable *GV : K.second) {
1229 if (!DirectAccess) {
1230 if (AMDGPU::isDynamicLDS(GV: *GV))
1231 LDSParams.IndirectAccess.DynamicLDSGlobals.insert(X: GV);
1232 else
1233 LDSParams.IndirectAccess.StaticLDSGlobals.insert(X: GV);
1234 FuncLDSAccessInfo.AllNonKernelLDSAccess.insert(X: GV);
1235 } else {
1236 if (AMDGPU::isDynamicLDS(GV: *GV))
1237 LDSParams.DirectAccess.DynamicLDSGlobals.insert(X: GV);
1238 else
1239 LDSParams.DirectAccess.StaticLDSGlobals.insert(X: GV);
1240 }
1241 }
1242 }
1243 };
1244
1245 PopulateKernelStaticDynamicLDS(LDSUsesInfo.DirectAccess, true);
1246 PopulateKernelStaticDynamicLDS(LDSUsesInfo.IndirectAccess, false);
1247
1248 // Get address sanitizer scale.
1249 initAsanInfo();
1250
1251 for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
1252 Function *Func = K.first;
1253 auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
1254 if (LDSParams.DirectAccess.StaticLDSGlobals.empty() &&
1255 LDSParams.DirectAccess.DynamicLDSGlobals.empty() &&
1256 LDSParams.IndirectAccess.StaticLDSGlobals.empty() &&
1257 LDSParams.IndirectAccess.DynamicLDSGlobals.empty())
1258 continue;
1259
1260 removeFnAttrFromReachable(
1261 CG, KernelRoot: Func,
1262 FnAttrs: {"amdgpu-no-workitem-id-x", "amdgpu-no-workitem-id-y",
1263 "amdgpu-no-workitem-id-z", "amdgpu-no-heap-ptr"});
1264 if (!LDSParams.IndirectAccess.StaticLDSGlobals.empty() ||
1265 !LDSParams.IndirectAccess.DynamicLDSGlobals.empty())
1266 removeFnAttrFromReachable(CG, KernelRoot: Func, FnAttrs: {"amdgpu-no-lds-kernel-id"});
1267 reorderStaticDynamicIndirectLDSSet(LDSParams);
1268 buildSwLDSGlobal(Func);
1269 buildSwDynLDSGlobal(Func);
1270 populateSwMetadataGlobal(Func);
1271 populateSwLDSAttributeAndMetadata(Func);
1272 populateLDSToReplacementIndicesMap(Func);
1273 DomTreeUpdater DTU(DTCallback(*Func), DomTreeUpdater::UpdateStrategy::Lazy);
1274 lowerKernelLDSAccesses(Func, DTU);
1275 Changed = true;
1276 }
1277
1278 // Get the Uses of LDS from non-kernels.
1279 getUsesOfLDSByNonKernels();
1280
1281 // Get non-kernels with LDS ptr as argument and called by kernels.
1282 getNonKernelsWithLDSArguments(CG);
1283
1284 // Lower LDS accesses in non-kernels.
1285 if (!FuncLDSAccessInfo.NonKernelToLDSAccessMap.empty() ||
1286 !FuncLDSAccessInfo.NonKernelsWithLDSArgument.empty()) {
1287 NonKernelLDSParameters NKLDSParams;
1288 NKLDSParams.OrderedKernels = getOrderedIndirectLDSAccessingKernels(
1289 Kernels&: FuncLDSAccessInfo.KernelsWithIndirectLDSAccess);
1290 NKLDSParams.OrdereLDSGlobals = getOrderedNonKernelAllLDSGlobals(
1291 Variables&: FuncLDSAccessInfo.AllNonKernelLDSAccess);
1292 buildNonKernelLDSBaseTable(NKLDSParams);
1293 buildNonKernelLDSOffsetTable(NKLDSParams);
1294 for (auto &K : FuncLDSAccessInfo.NonKernelToLDSAccessMap) {
1295 Function *Func = K.first;
1296 DenseSet<GlobalVariable *> &LDSGlobals = K.second;
1297 SetVector<GlobalVariable *> OrderedLDSGlobals = sortByName(
1298 V: std::vector<GlobalVariable *>(LDSGlobals.begin(), LDSGlobals.end()));
1299 lowerNonKernelLDSAccesses(Func, LDSGlobals&: OrderedLDSGlobals, NKLDSParams);
1300 }
1301 for (Function *Func : FuncLDSAccessInfo.NonKernelsWithLDSArgument) {
1302 auto &K = FuncLDSAccessInfo.NonKernelToLDSAccessMap;
1303 if (K.contains(Val: Func))
1304 continue;
1305 SetVector<llvm::GlobalVariable *> Vec;
1306 lowerNonKernelLDSAccesses(Func, LDSGlobals&: Vec, NKLDSParams);
1307 }
1308 Changed = true;
1309 }
1310
1311 if (!Changed)
1312 return Changed;
1313
1314 for (auto &GV : make_early_inc_range(Range: M.globals())) {
1315 if (AMDGPU::isLDSVariableToLower(GV)) {
1316 // probably want to remove from used lists
1317 GV.removeDeadConstantUsers();
1318 if (GV.use_empty())
1319 GV.eraseFromParent();
1320 }
1321 }
1322
1323 if (AsanInstrumentLDS) {
1324 SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
1325 for (Instruction *Inst : AsanInfo.Instructions) {
1326 SmallVector<InterestingMemoryOperand, 1> InterestingOperands;
1327 getInterestingMemoryOperands(M, I: Inst, Interesting&: InterestingOperands);
1328 llvm::append_range(C&: OperandsToInstrument, R&: InterestingOperands);
1329 }
1330 for (auto &Operand : OperandsToInstrument) {
1331 Value *Addr = Operand.getPtr();
1332 instrumentAddress(M, IRB, OrigIns: Operand.getInsn(), InsertBefore: Operand.getInsn(), Addr,
1333 Alignment: Operand.Alignment.valueOrOne(), TypeStoreSize: Operand.TypeStoreSize,
1334 IsWrite: Operand.IsWrite, SizeArgument: nullptr, UseCalls: false, Recover: false, Scale: AsanInfo.Scale,
1335 Offset: AsanInfo.Offset);
1336 Changed = true;
1337 }
1338 }
1339
1340 return Changed;
1341}
1342
1343class AMDGPUSwLowerLDSLegacy : public ModulePass {
1344public:
1345 static char ID;
1346 AMDGPUSwLowerLDSLegacy() : ModulePass(ID) {}
1347 bool runOnModule(Module &M) override;
1348 void getAnalysisUsage(AnalysisUsage &AU) const override {
1349 AU.addPreserved<DominatorTreeWrapperPass>();
1350 }
1351};
1352} // namespace
1353
1354char AMDGPUSwLowerLDSLegacy::ID = 0;
1355char &llvm::AMDGPUSwLowerLDSLegacyPassID = AMDGPUSwLowerLDSLegacy::ID;
1356
1357INITIALIZE_PASS_BEGIN(AMDGPUSwLowerLDSLegacy, "amdgpu-sw-lower-lds",
1358 "AMDGPU Software lowering of LDS", false, false)
1359INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
1360INITIALIZE_PASS_END(AMDGPUSwLowerLDSLegacy, "amdgpu-sw-lower-lds",
1361 "AMDGPU Software lowering of LDS", false, false)
1362
1363bool AMDGPUSwLowerLDSLegacy::runOnModule(Module &M) {
1364 // AddressSanitizer pass adds "nosanitize_address" module flag if it has
1365 // instrumented the IR. Return early if the flag is not present.
1366 if (!M.getModuleFlag(Key: "nosanitize_address"))
1367 return false;
1368 DominatorTreeWrapperPass *const DTW =
1369 getAnalysisIfAvailable<DominatorTreeWrapperPass>();
1370 auto DTCallback = [&DTW](Function &F) -> DominatorTree * {
1371 return DTW ? &DTW->getDomTree() : nullptr;
1372 };
1373
1374 AMDGPUSwLowerLDS SwLowerLDSImpl(M, DTCallback);
1375 bool IsChanged = SwLowerLDSImpl.run();
1376 return IsChanged;
1377}
1378
1379ModulePass *llvm::createAMDGPUSwLowerLDSLegacyPass() {
1380 return new AMDGPUSwLowerLDSLegacy();
1381}
1382
1383PreservedAnalyses AMDGPUSwLowerLDSPass::run(Module &M,
1384 ModuleAnalysisManager &AM) {
1385 // AddressSanitizer pass adds "nosanitize_address" module flag if it has
1386 // instrumented the IR. Return early if the flag is not present.
1387 if (!M.getModuleFlag(Key: "nosanitize_address"))
1388 return PreservedAnalyses::all();
1389 auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
1390 auto DTCallback = [&FAM](Function &F) -> DominatorTree * {
1391 return &FAM.getResult<DominatorTreeAnalysis>(IR&: F);
1392 };
1393 AMDGPUSwLowerLDS SwLowerLDSImpl(M, DTCallback);
1394 bool IsChanged = SwLowerLDSImpl.run();
1395 if (!IsChanged)
1396 return PreservedAnalyses::all();
1397
1398 PreservedAnalyses PA;
1399 PA.preserve<DominatorTreeAnalysis>();
1400 return PA;
1401}
1402