1//===-- NVPTXLowerArgs.cpp - Lower arguments ------------------------------===//
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//
10// Arguments to kernel and device functions are passed via param space,
11// which imposes certain restrictions:
12// http://docs.nvidia.com/cuda/parallel-thread-execution/#state-spaces
13//
14// Kernel parameters are read-only and accessible only via ld.param
15// instruction, directly or via a pointer.
16//
17// Device function parameters are directly accessible via
18// ld.param/st.param, but taking the address of one returns a pointer
19// to a copy created in local space which *can't* be used with
20// ld.param/st.param.
21//
22// Copying a byval struct into local memory in IR allows us to enforce
23// the param space restrictions, gives the rest of IR a pointer w/o
24// param space restrictions, and gives us an opportunity to eliminate
25// the copy.
26//
27// Pointer arguments to kernel functions need more work to be lowered:
28//
29// 1. Convert non-byval pointer arguments of CUDA kernels to pointers in the
30// global address space. This allows later optimizations to emit
31// ld.global.*/st.global.* for accessing these pointer arguments. For
32// example,
33//
34// define void @foo(float* %input) {
35// %v = load float, float* %input, align 4
36// ...
37// }
38//
39// becomes
40//
41// define void @foo(float* %input) {
42// %input2 = addrspacecast float* %input to float addrspace(1)*
43// %input3 = addrspacecast float addrspace(1)* %input2 to float*
44// %v = load float, float* %input3, align 4
45// ...
46// }
47//
48// Later, NVPTXInferAddressSpaces will optimize it to
49//
50// define void @foo(float* %input) {
51// %input2 = addrspacecast float* %input to float addrspace(1)*
52// %v = load float, float addrspace(1)* %input2, align 4
53// ...
54// }
55//
56// 2. Convert byval kernel parameters to pointers in the param address space
57// (so that NVPTX emits ld/st.param). Convert pointers *within* a byval
58// kernel parameter to pointers in the global address space. This allows
59// NVPTX to emit ld/st.global.
60//
61// struct S {
62// int *x;
63// int *y;
64// };
65// __global__ void foo(S s) {
66// int *b = s.y;
67// // use b
68// }
69//
70// "b" points to the global address space. In the IR level,
71//
72// define void @foo(ptr byval %input) {
73// %b_ptr = getelementptr {ptr, ptr}, ptr %input, i64 0, i32 1
74// %b = load ptr, ptr %b_ptr
75// ; use %b
76// }
77//
78// becomes
79//
80// define void @foo({i32*, i32*}* byval %input) {
81// %b_param = addrspacecat ptr %input to ptr addrspace(101)
82// %b_ptr = getelementptr {ptr, ptr}, ptr addrspace(101) %b_param, i64 0, i32 1
83// %b = load ptr, ptr addrspace(101) %b_ptr
84// %b_global = addrspacecast ptr %b to ptr addrspace(1)
85// ; use %b_generic
86// }
87//
88// Create a local copy of kernel byval parameters used in a way that *might* mutate
89// the parameter, by storing it in an alloca. Mutations to "grid_constant" parameters
90// are undefined behaviour, and don't require local copies.
91//
92// define void @foo(ptr byval(%struct.s) align 4 %input) {
93// store i32 42, ptr %input
94// ret void
95// }
96//
97// becomes
98//
99// define void @foo(ptr byval(%struct.s) align 4 %input) #1 {
100// %input1 = alloca %struct.s, align 4
101// %input2 = addrspacecast ptr %input to ptr addrspace(101)
102// %input3 = load %struct.s, ptr addrspace(101) %input2, align 4
103// store %struct.s %input3, ptr %input1, align 4
104// store i32 42, ptr %input1, align 4
105// ret void
106// }
107//
108// If %input were passed to a device function, or written to memory,
109// conservatively assume that %input gets mutated, and create a local copy.
110//
111// Convert param pointers to grid_constant byval kernel parameters that are
112// passed into calls (device functions, intrinsics, inline asm), or otherwise
113// "escape" (into stores/ptrtoints) to the generic address space, using the
114// `nvvm.ptr.param.to.gen` intrinsic, so that NVPTX emits cvta.param
115// (available for sm70+)
116//
117// define void @foo(ptr byval(%struct.s) %input) {
118// ; %input is a grid_constant
119// %call = call i32 @escape(ptr %input)
120// ret void
121// }
122//
123// becomes
124//
125// define void @foo(ptr byval(%struct.s) %input) {
126// %input1 = addrspacecast ptr %input to ptr addrspace(101)
127// ; the following intrinsic converts pointer to generic. We don't use an addrspacecast
128// ; to prevent generic -> param -> generic from getting cancelled out
129// %input1.gen = call ptr @llvm.nvvm.ptr.param.to.gen.p0.p101(ptr addrspace(101) %input1)
130// %call = call i32 @escape(ptr %input1.gen)
131// ret void
132// }
133//
134// TODO: merge this pass with NVPTXInferAddressSpaces so that other passes don't
135// cancel the addrspacecast pair this pass emits.
136//===----------------------------------------------------------------------===//
137
138#include "MCTargetDesc/NVPTXBaseInfo.h"
139#include "NVPTX.h"
140#include "NVPTXTargetMachine.h"
141#include "NVPTXUtilities.h"
142#include "llvm/ADT/STLExtras.h"
143#include "llvm/ADT/SmallVectorExtras.h"
144#include "llvm/Analysis/PtrUseVisitor.h"
145#include "llvm/Analysis/ValueTracking.h"
146#include "llvm/CodeGen/TargetPassConfig.h"
147#include "llvm/IR/Attributes.h"
148#include "llvm/IR/Function.h"
149#include "llvm/IR/IRBuilder.h"
150#include "llvm/IR/InstIterator.h"
151#include "llvm/IR/Instructions.h"
152#include "llvm/IR/IntrinsicInst.h"
153#include "llvm/IR/IntrinsicsNVPTX.h"
154#include "llvm/IR/Type.h"
155#include "llvm/InitializePasses.h"
156#include "llvm/Pass.h"
157#include "llvm/Support/Debug.h"
158#include "llvm/Support/ErrorHandling.h"
159#include "llvm/Support/NVPTXAddrSpace.h"
160#include <numeric>
161#include <queue>
162
163#define DEBUG_TYPE "nvptx-lower-args"
164
165using namespace llvm;
166
167namespace {
168class NVPTXLowerArgsLegacyPass : public FunctionPass {
169 bool runOnFunction(Function &F) override;
170
171public:
172 static char ID; // Pass identification, replacement for typeid
173 NVPTXLowerArgsLegacyPass() : FunctionPass(ID) {}
174 StringRef getPassName() const override {
175 return "Lower pointer arguments of CUDA kernels";
176 }
177 void getAnalysisUsage(AnalysisUsage &AU) const override {
178 AU.addRequired<TargetPassConfig>();
179 }
180};
181} // namespace
182
183char NVPTXLowerArgsLegacyPass::ID = 1;
184
185INITIALIZE_PASS_BEGIN(NVPTXLowerArgsLegacyPass, "nvptx-lower-args",
186 "Lower arguments (NVPTX)", false, false)
187INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
188INITIALIZE_PASS_END(NVPTXLowerArgsLegacyPass, "nvptx-lower-args",
189 "Lower arguments (NVPTX)", false, false)
190
191// =============================================================================
192// If the function had a byval struct ptr arg, say foo(ptr byval(%struct.x) %d),
193// and we can't guarantee that the only accesses are loads,
194// then add the following instructions to the first basic block:
195//
196// %temp = alloca %struct.x, align 8
197// %tempd = addrspacecast ptr %d to ptr addrspace(101)
198// %tv = load %struct.x, ptr addrspace(101) %tempd
199// store %struct.x %tv, ptr %temp, align 8
200//
201// The above code allocates some space in the stack and copies the incoming
202// struct from param space to local space.
203// Then replace all occurrences of %d by %temp.
204//
205// In case we know that all users are GEPs or Loads, replace them with the same
206// ones in parameter AS, so we can access them using ld.param.
207// =============================================================================
208
209/// Recursively convert the users of a param to the param address space.
210static void convertToParamAS(ArrayRef<Use *> OldUses, Value *Param) {
211 struct IP {
212 Use *OldUse;
213 Value *NewParam;
214 };
215
216 const auto CloneInstInParamAS = [](const IP &I) -> Value * {
217 auto *OldInst = cast<Instruction>(Val: I.OldUse->getUser());
218 if (auto *LI = dyn_cast<LoadInst>(Val: OldInst)) {
219 LI->setOperand(i_nocapture: 0, Val_nocapture: I.NewParam);
220 return LI;
221 }
222 if (auto *GEP = dyn_cast<GetElementPtrInst>(Val: OldInst)) {
223 SmallVector<Value *, 4> Indices(GEP->indices());
224 auto *NewGEP = GetElementPtrInst::Create(
225 PointeeType: GEP->getSourceElementType(), Ptr: I.NewParam, IdxList: Indices, NameStr: GEP->getName(),
226 InsertBefore: GEP->getIterator());
227 NewGEP->setNoWrapFlags(GEP->getNoWrapFlags());
228 return NewGEP;
229 }
230 if (auto *BC = dyn_cast<BitCastInst>(Val: OldInst)) {
231 auto *NewBCType = PointerType::get(C&: BC->getContext(), AddressSpace: ADDRESS_SPACE_PARAM);
232 return BitCastInst::Create(BC->getOpcode(), S: I.NewParam, Ty: NewBCType,
233 Name: BC->getName(), InsertBefore: BC->getIterator());
234 }
235 if (auto *ASC = dyn_cast<AddrSpaceCastInst>(Val: OldInst)) {
236 assert(ASC->getDestAddressSpace() == ADDRESS_SPACE_PARAM);
237 (void)ASC;
238 // Just pass through the argument, the old ASC is no longer needed.
239 return I.NewParam;
240 }
241 if (auto *MI = dyn_cast<MemTransferInst>(Val: OldInst)) {
242 if (MI->getRawSource() == I.OldUse->get()) {
243 // convert to memcpy/memmove from param space.
244 IRBuilder<> Builder(OldInst);
245 Intrinsic::ID ID = MI->getIntrinsicID();
246
247 CallInst *B = Builder.CreateMemTransferInst(
248 IntrID: ID, Dst: MI->getRawDest(), DstAlign: MI->getDestAlign(), Src: I.NewParam,
249 SrcAlign: MI->getSourceAlign(), Size: MI->getLength(), isVolatile: MI->isVolatile());
250 for (unsigned I : {0, 1})
251 if (uint64_t Bytes = MI->getParamDereferenceableBytes(i: I))
252 B->addDereferenceableParamAttr(i: I, Bytes);
253 return B;
254 }
255 }
256
257 llvm_unreachable("Unsupported instruction");
258 };
259
260 auto ItemsToConvert =
261 map_to_vector(C&: OldUses, F: [=](Use *U) -> IP { return {.OldUse: U, .NewParam: Param}; });
262 SmallVector<Instruction *> InstructionsToDelete;
263
264 while (!ItemsToConvert.empty()) {
265 IP I = ItemsToConvert.pop_back_val();
266 Value *NewInst = CloneInstInParamAS(I);
267 Instruction *OldInst = cast<Instruction>(Val: I.OldUse->getUser());
268
269 if (NewInst && NewInst != OldInst) {
270 // We've created a new instruction. Queue users of the old instruction to
271 // be converted and the instruction itself to be deleted. We can't delete
272 // the old instruction yet, because it's still in use by a load somewhere.
273 for (Use &U : OldInst->uses())
274 ItemsToConvert.push_back(Elt: {.OldUse: &U, .NewParam: NewInst});
275
276 InstructionsToDelete.push_back(Elt: OldInst);
277 }
278 }
279
280 // Now we know that all argument loads are using addresses in parameter space
281 // and we can finally remove the old instructions in generic AS. Instructions
282 // scheduled for removal should be processed in reverse order so the ones
283 // closest to the load are deleted first. Otherwise they may still be in use.
284 // E.g if we have Value = Load(BitCast(GEP(arg))), InstructionsToDelete will
285 // have {GEP,BitCast}. GEP can't be deleted first, because it's still used by
286 // the BitCast.
287 for (Instruction *I : llvm::reverse(C&: InstructionsToDelete))
288 I->eraseFromParent();
289}
290
291static Align setByValParamAlign(Argument *Arg, const NVPTXTargetLowering *TLI) {
292 Function *F = Arg->getParent();
293 Type *ByValType = Arg->getParamByValType();
294 const DataLayout &DL = F->getDataLayout();
295
296 const Align OptimizedAlign =
297 TLI->getFunctionParamOptimizedAlign(F, ArgTy: ByValType, DL);
298 const Align CurrentAlign = Arg->getParamAlign().valueOrOne();
299
300 if (CurrentAlign >= OptimizedAlign)
301 return CurrentAlign;
302
303 LLVM_DEBUG(dbgs() << "Try to use alignment " << OptimizedAlign.value()
304 << " instead of " << CurrentAlign.value() << " for " << *Arg
305 << '\n');
306
307 Arg->removeAttr(Kind: Attribute::Alignment);
308 Arg->addAttr(Attr: Attribute::getWithAlignment(Context&: F->getContext(), Alignment: OptimizedAlign));
309
310 return OptimizedAlign;
311}
312
313// Adjust alignment of arguments passed byval in .param address space. We can
314// increase alignment of such arguments in a way that ensures that we can
315// effectively vectorize their loads. We should also traverse all loads from
316// byval pointer and adjust their alignment, if those were using known offset.
317// Such alignment changes must be conformed with parameter store and load in
318// NVPTXTargetLowering::LowerCall.
319static void propagateAlignmentToLoads(Value *Val, Align NewAlign,
320 const DataLayout &DL) {
321 struct Load {
322 LoadInst *Inst;
323 uint64_t Offset;
324 };
325
326 struct LoadContext {
327 Value *InitialVal;
328 uint64_t Offset;
329 };
330
331 SmallVector<Load> Loads;
332 std::queue<LoadContext> Worklist;
333 Worklist.push(x: {.InitialVal: Val, .Offset: 0});
334
335 while (!Worklist.empty()) {
336 LoadContext Ctx = Worklist.front();
337 Worklist.pop();
338
339 for (User *CurUser : Ctx.InitialVal->users()) {
340 if (auto *I = dyn_cast<LoadInst>(Val: CurUser))
341 Loads.push_back(Elt: {.Inst: I, .Offset: Ctx.Offset});
342 else if (isa<BitCastInst>(Val: CurUser) || isa<AddrSpaceCastInst>(Val: CurUser))
343 Worklist.push(x: {.InitialVal: cast<Instruction>(Val: CurUser), .Offset: Ctx.Offset});
344 else if (auto *I = dyn_cast<GetElementPtrInst>(Val: CurUser)) {
345 APInt OffsetAccumulated =
346 APInt::getZero(numBits: DL.getIndexSizeInBits(AS: ADDRESS_SPACE_PARAM));
347
348 if (!I->accumulateConstantOffset(DL, Offset&: OffsetAccumulated))
349 continue;
350
351 uint64_t OffsetLimit = -1;
352 uint64_t Offset = OffsetAccumulated.getLimitedValue(Limit: OffsetLimit);
353 assert(Offset != OffsetLimit && "Expect Offset less than UINT64_MAX");
354
355 Worklist.push(x: {.InitialVal: I, .Offset: Ctx.Offset + Offset});
356 }
357 }
358 }
359
360 for (Load &CurLoad : Loads) {
361 Align NewLoadAlign = commonAlignment(A: NewAlign, Offset: CurLoad.Offset);
362 Align CurLoadAlign = CurLoad.Inst->getAlign();
363 CurLoad.Inst->setAlignment(std::max(a: NewLoadAlign, b: CurLoadAlign));
364 }
365}
366
367// Create a call to the nvvm_internal_addrspace_wrap intrinsic and set the
368// alignment of the return value based on the alignment of the argument.
369static CallInst *createNVVMInternalAddrspaceWrap(IRBuilder<> &IRB,
370 Argument &Arg) {
371 CallInst *ArgInParam =
372 IRB.CreateIntrinsic(ID: Intrinsic::nvvm_internal_addrspace_wrap,
373 Types: {IRB.getPtrTy(AddrSpace: ADDRESS_SPACE_PARAM), Arg.getType()},
374 Args: &Arg, FMFSource: {}, Name: Arg.getName() + ".param");
375
376 if (MaybeAlign ParamAlign = Arg.getParamAlign())
377 ArgInParam->addRetAttr(
378 Attr: Attribute::getWithAlignment(Context&: ArgInParam->getContext(), Alignment: *ParamAlign));
379
380 Arg.addAttr(Attr: Attribute::get(Context&: Arg.getContext(), Kind: "nvvm.grid_constant"));
381 Arg.addAttr(Kind: Attribute::ReadOnly);
382
383 return ArgInParam;
384}
385
386namespace {
387struct ArgUseChecker : PtrUseVisitor<ArgUseChecker> {
388 using Base = PtrUseVisitor<ArgUseChecker>;
389 // Set of phi/select instructions using the Arg
390 SmallPtrSet<Instruction *, 4> Conditionals;
391
392 ArgUseChecker(const DataLayout &DL) : PtrUseVisitor(DL) {}
393
394 PtrInfo visitArgPtr(Argument &A) {
395 assert(A.getType()->isPointerTy());
396 IntegerType *IntIdxTy = cast<IntegerType>(Val: DL.getIndexType(PtrTy: A.getType()));
397 IsOffsetKnown = false;
398 Offset = APInt(IntIdxTy->getBitWidth(), 0);
399 PI.reset();
400
401 LLVM_DEBUG(dbgs() << "Checking Argument " << A << "\n");
402 // Enqueue the uses of this pointer.
403 enqueueUsers(I&: A);
404
405 // Visit all the uses off the worklist until it is empty.
406 // Note that unlike PtrUseVisitor we intentionally do not track offsets.
407 // We're only interested in how we use the pointer.
408 while (!(Worklist.empty() || PI.isAborted())) {
409 UseToVisit ToVisit = Worklist.pop_back_val();
410 U = ToVisit.UseAndIsOffsetKnown.getPointer();
411 Instruction *I = cast<Instruction>(Val: U->getUser());
412 LLVM_DEBUG(dbgs() << "Processing " << *I << "\n");
413 Base::visit(I);
414 }
415 if (PI.isEscaped())
416 LLVM_DEBUG(dbgs() << "Argument pointer escaped: " << *PI.getEscapingInst()
417 << "\n");
418 else if (PI.isAborted())
419 LLVM_DEBUG(dbgs() << "Pointer use needs a copy: " << *PI.getAbortingInst()
420 << "\n");
421 LLVM_DEBUG(dbgs() << "Traversed " << Conditionals.size()
422 << " conditionals\n");
423 return PI;
424 }
425
426 void visitStoreInst(StoreInst &SI) {
427 // Storing the pointer escapes it.
428 if (U->get() == SI.getValueOperand())
429 return PI.setEscapedAndAborted(&SI);
430
431 PI.setAborted(&SI);
432 }
433
434 void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
435 // ASC to param space are no-ops and do not need a copy
436 if (ASC.getDestAddressSpace() != ADDRESS_SPACE_PARAM)
437 return PI.setEscapedAndAborted(&ASC);
438 Base::visitAddrSpaceCastInst(ASC);
439 }
440
441 void visitPtrToIntInst(PtrToIntInst &I) { Base::visitPtrToIntInst(I); }
442
443 void visitPHINodeOrSelectInst(Instruction &I) {
444 assert(isa<PHINode>(I) || isa<SelectInst>(I));
445 enqueueUsers(I);
446 Conditionals.insert(Ptr: &I);
447 }
448 // PHI and select just pass through the pointers.
449 void visitPHINode(PHINode &PN) { visitPHINodeOrSelectInst(I&: PN); }
450 void visitSelectInst(SelectInst &SI) { visitPHINodeOrSelectInst(I&: SI); }
451
452 // memcpy/memmove are OK when the pointer is source. We can convert them to
453 // AS-specific memcpy.
454 void visitMemTransferInst(MemTransferInst &II) {
455 if (*U == II.getRawDest())
456 PI.setAborted(&II);
457 }
458
459 void visitMemSetInst(MemSetInst &II) { PI.setAborted(&II); }
460}; // struct ArgUseChecker
461
462void copyByValParam(Function &F, Argument &Arg) {
463 LLVM_DEBUG(dbgs() << "Creating a local copy of " << Arg << "\n");
464 Type *ByValType = Arg.getParamByValType();
465 const DataLayout &DL = F.getDataLayout();
466 IRBuilder<> IRB(&F.getEntryBlock().front());
467 AllocaInst *AllocA = IRB.CreateAlloca(Ty: ByValType, ArraySize: nullptr, Name: Arg.getName());
468 // Set the alignment to alignment of the byval parameter. This is because,
469 // later load/stores assume that alignment, and we are going to replace
470 // the use of the byval parameter with this alloca instruction.
471 AllocA->setAlignment(
472 Arg.getParamAlign().value_or(u: DL.getPrefTypeAlign(Ty: ByValType)));
473 Arg.replaceAllUsesWith(V: AllocA);
474
475 Value *ArgInParamAS = createNVVMInternalAddrspaceWrap(IRB, Arg);
476
477 // Be sure to propagate alignment to this load; LLVM doesn't know that NVPTX
478 // addrspacecast preserves alignment. Since params are constant, this load
479 // is definitely not volatile.
480 const auto ArgSize = *AllocA->getAllocationSize(DL);
481 IRB.CreateMemCpy(Dst: AllocA, DstAlign: AllocA->getAlign(), Src: ArgInParamAS, SrcAlign: AllocA->getAlign(),
482 Size: ArgSize);
483}
484} // namespace
485
486static bool argIsProcessed(Argument *Arg) {
487 if (Arg->use_empty())
488 return true;
489
490 // If the argument is already wrapped, it was processed by this pass before.
491 if (Arg->hasOneUse())
492 if (const auto *II = dyn_cast<IntrinsicInst>(Val: *Arg->user_begin()))
493 if (II->getIntrinsicID() == Intrinsic::nvvm_internal_addrspace_wrap)
494 return true;
495
496 return false;
497}
498
499static void handleByValParam(const NVPTXTargetMachine &TM, Argument *Arg) {
500 Function *F = Arg->getParent();
501 assert(isKernelFunction(*F));
502 const NVPTXSubtarget *ST = TM.getSubtargetImpl(*F);
503
504 const DataLayout &DL = F->getDataLayout();
505 IRBuilder<> IRB(&F->getEntryBlock().front());
506
507 if (argIsProcessed(Arg))
508 return;
509
510 const Align NewArgAlign = setByValParamAlign(Arg, TLI: ST->getTargetLowering());
511
512 // (1) First check the easy case, if were able to trace through all the uses
513 // and we can convert them all to param AS, then we'll do this.
514 ArgUseChecker AUC(DL);
515 ArgUseChecker::PtrInfo PI = AUC.visitArgPtr(A&: *Arg);
516 const bool ArgUseIsReadOnly = !(PI.isEscaped() || PI.isAborted());
517 if (ArgUseIsReadOnly && AUC.Conditionals.empty()) {
518 // Convert all loads and intermediate operations to use parameter AS and
519 // skip creation of a local copy of the argument.
520 SmallVector<Use *, 16> UsesToUpdate(llvm::make_pointer_range(Range: Arg->uses()));
521 Value *ArgInParamAS = createNVVMInternalAddrspaceWrap(IRB, Arg&: *Arg);
522 for (Use *U : UsesToUpdate)
523 convertToParamAS(OldUses: U, Param: ArgInParamAS);
524
525 propagateAlignmentToLoads(Val: ArgInParamAS, NewAlign: NewArgAlign, DL);
526 return;
527 }
528
529 // (2) If the argument is grid constant, we get to use the pointer directly.
530 if (ST->hasCvtaParam() && (ArgUseIsReadOnly || isParamGridConstant(*Arg))) {
531 LLVM_DEBUG(dbgs() << "Using non-copy pointer to " << *Arg << "\n");
532
533 // Cast argument to param address space. Because the backend will emit the
534 // argument already in the param address space, we need to use the noop
535 // intrinsic, this had the added benefit of preventing other optimizations
536 // from folding away this pair of addrspacecasts.
537 Instruction *ArgInParamAS = createNVVMInternalAddrspaceWrap(IRB, Arg&: *Arg);
538
539 // Cast param address to generic address space.
540 Value *GenericArg = IRB.CreateAddrSpaceCast(
541 V: ArgInParamAS, DestTy: IRB.getPtrTy(AddrSpace: ADDRESS_SPACE_GENERIC),
542 Name: Arg->getName() + ".gen");
543
544 Arg->replaceAllUsesWith(V: GenericArg);
545
546 // Do not replace Arg in the cast to param space
547 ArgInParamAS->setOperand(i: 0, Val: Arg);
548 return;
549 }
550
551 // (3) Otherwise we have to create a copy of the argument in local memory.
552 copyByValParam(F&: *F, Arg&: *Arg);
553}
554
555static void markPointerAsAS(Value *Ptr, const unsigned AS) {
556 if (Ptr->getType()->getPointerAddressSpace() != ADDRESS_SPACE_GENERIC)
557 return;
558
559 // Deciding where to emit the addrspacecast pair.
560 BasicBlock::iterator InsertPt;
561 if (Argument *Arg = dyn_cast<Argument>(Val: Ptr)) {
562 // Insert at the functon entry if Ptr is an argument.
563 InsertPt = Arg->getParent()->getEntryBlock().begin();
564 } else {
565 // Insert right after Ptr if Ptr is an instruction.
566 InsertPt = ++cast<Instruction>(Val: Ptr)->getIterator();
567 assert(InsertPt != InsertPt->getParent()->end() &&
568 "We don't call this function with Ptr being a terminator.");
569 }
570
571 Instruction *PtrInGlobal = new AddrSpaceCastInst(
572 Ptr, PointerType::get(C&: Ptr->getContext(), AddressSpace: AS), Ptr->getName(), InsertPt);
573 Value *PtrInGeneric = new AddrSpaceCastInst(PtrInGlobal, Ptr->getType(),
574 Ptr->getName(), InsertPt);
575 // Replace with PtrInGeneric all uses of Ptr except PtrInGlobal.
576 Ptr->replaceAllUsesWith(V: PtrInGeneric);
577 PtrInGlobal->setOperand(i: 0, Val: Ptr);
578}
579
580static void markPointerAsGlobal(Value *Ptr) {
581 markPointerAsAS(Ptr, AS: ADDRESS_SPACE_GLOBAL);
582}
583
584static void handleIntToPtr(Value &V) {
585 if (!all_of(Range: V.users(), P: [](User *U) { return isa<IntToPtrInst>(Val: U); }))
586 return;
587
588 SmallVector<User *, 16> UsersToUpdate(V.users());
589 for (User *U : UsersToUpdate)
590 markPointerAsGlobal(Ptr: U);
591}
592
593// =============================================================================
594// Main function for this pass.
595// =============================================================================
596static bool runOnKernelFunction(const NVPTXTargetMachine &TM, Function &F) {
597 // Copying of byval aggregates + SROA may result in pointers being loaded as
598 // integers, followed by intotoptr. We may want to mark those as global, too,
599 // but only if the loaded integer is used exclusively for conversion to a
600 // pointer with inttoptr.
601 if (TM.getDrvInterface() == NVPTX::CUDA) {
602 // Mark pointers in byval structs as global.
603 for (auto &I : instructions(F)) {
604 auto *LI = dyn_cast<LoadInst>(Val: &I);
605 if (!LI)
606 continue;
607
608 if (LI->getType()->isPointerTy() || LI->getType()->isIntegerTy()) {
609 Value *UO = getUnderlyingObject(V: LI->getPointerOperand());
610 if (Argument *Arg = dyn_cast<Argument>(Val: UO)) {
611 if (Arg->hasByValAttr()) {
612 // LI is a load from a pointer within a byval kernel parameter.
613 if (LI->getType()->isPointerTy())
614 markPointerAsGlobal(Ptr: LI);
615 else
616 handleIntToPtr(V&: *LI);
617 }
618 }
619 }
620 }
621
622 for (Argument &Arg : F.args())
623 if (Arg.getType()->isIntegerTy())
624 handleIntToPtr(V&: Arg);
625 }
626
627 LLVM_DEBUG(dbgs() << "Lowering kernel args of " << F.getName() << "\n");
628 for (Argument &Arg : F.args())
629 if (Arg.hasByValAttr())
630 handleByValParam(TM, Arg: &Arg);
631
632 return true;
633}
634
635// Device functions only need to copy byval args into local memory.
636static bool runOnDeviceFunction(const NVPTXTargetMachine &TM, Function &F) {
637 LLVM_DEBUG(dbgs() << "Lowering function args of " << F.getName() << "\n");
638
639 const NVPTXTargetLowering *TLI = TM.getSubtargetImpl()->getTargetLowering();
640 const DataLayout &DL = F.getDataLayout();
641
642 for (Argument &Arg : F.args())
643 if (Arg.hasByValAttr()) {
644 const Align NewArgAlign = setByValParamAlign(Arg: &Arg, TLI);
645 propagateAlignmentToLoads(Val: &Arg, NewAlign: NewArgAlign, DL);
646 }
647
648 return true;
649}
650
651static bool processFunction(Function &F, NVPTXTargetMachine &TM) {
652 return isKernelFunction(F) ? runOnKernelFunction(TM, F)
653 : runOnDeviceFunction(TM, F);
654}
655
656bool NVPTXLowerArgsLegacyPass::runOnFunction(Function &F) {
657 auto &TM = getAnalysis<TargetPassConfig>().getTM<NVPTXTargetMachine>();
658 return processFunction(F, TM);
659}
660FunctionPass *llvm::createNVPTXLowerArgsPass() {
661 return new NVPTXLowerArgsLegacyPass();
662}
663
664static bool copyFunctionByValArgs(Function &F) {
665 LLVM_DEBUG(dbgs() << "Creating a copy of byval args of " << F.getName()
666 << "\n");
667 bool Changed = false;
668 if (isKernelFunction(F)) {
669 for (Argument &Arg : F.args())
670 if (Arg.hasByValAttr() && !isParamGridConstant(Arg)) {
671 copyByValParam(F, Arg);
672 Changed = true;
673 }
674 }
675 return Changed;
676}
677
678PreservedAnalyses NVPTXCopyByValArgsPass::run(Function &F,
679 FunctionAnalysisManager &AM) {
680 return copyFunctionByValArgs(F) ? PreservedAnalyses::none()
681 : PreservedAnalyses::all();
682}
683
684PreservedAnalyses NVPTXLowerArgsPass::run(Function &F,
685 FunctionAnalysisManager &AM) {
686 auto &NTM = static_cast<NVPTXTargetMachine &>(TM);
687 bool Changed = processFunction(F, TM&: NTM);
688 return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
689}
690