1 //===-- AMDGPULowerKernelArguments.cpp ------------------------------------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 /// \file This pass replaces accesses to kernel arguments with loads from
11 /// offsets from the kernarg base pointer.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #include "AMDGPU.h"
16 #include "AMDGPUSubtarget.h"
17 #include "AMDGPUTargetMachine.h"
18 #include "llvm/ADT/StringRef.h"
19 #include "llvm/Analysis/Loads.h"
20 #include "llvm/CodeGen/Passes.h"
21 #include "llvm/CodeGen/TargetPassConfig.h"
22 #include "llvm/IR/Attributes.h"
23 #include "llvm/IR/BasicBlock.h"
24 #include "llvm/IR/Constants.h"
25 #include "llvm/IR/DerivedTypes.h"
26 #include "llvm/IR/Function.h"
27 #include "llvm/IR/IRBuilder.h"
28 #include "llvm/IR/InstrTypes.h"
29 #include "llvm/IR/Instruction.h"
30 #include "llvm/IR/Instructions.h"
31 #include "llvm/IR/LLVMContext.h"
32 #include "llvm/IR/MDBuilder.h"
33 #include "llvm/IR/Metadata.h"
34 #include "llvm/IR/Operator.h"
35 #include "llvm/IR/Type.h"
36 #include "llvm/IR/Value.h"
37 #include "llvm/Pass.h"
38 #include "llvm/Support/Casting.h"
39 
40 #define DEBUG_TYPE "amdgpu-lower-kernel-arguments"
41 
42 using namespace llvm;
43 
44 namespace {
45 
46 class AMDGPULowerKernelArguments : public FunctionPass{
47 public:
48   static char ID;
49 
50   AMDGPULowerKernelArguments() : FunctionPass(ID) {}
51 
52   bool runOnFunction(Function &F) override;
53 
54   void getAnalysisUsage(AnalysisUsage &AU) const override {
55     AU.addRequired<TargetPassConfig>();
56     AU.setPreservesAll();
57  }
58 };
59 
60 } // end anonymous namespace
61 
62 bool AMDGPULowerKernelArguments::runOnFunction(Function &F) {
63   CallingConv::ID CC = F.getCallingConv();
64   if (CC != CallingConv::AMDGPU_KERNEL || F.arg_empty())
65     return false;
66 
67   auto &TPC = getAnalysis<TargetPassConfig>();
68 
69   const TargetMachine &TM = TPC.getTM<TargetMachine>();
70   const GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(F);
71   LLVMContext &Ctx = F.getParent()->getContext();
72   const DataLayout &DL = F.getParent()->getDataLayout();
73   BasicBlock &EntryBlock = *F.begin();
74   IRBuilder<> Builder(&*EntryBlock.begin());
75 
76   const unsigned KernArgBaseAlign = 16; // FIXME: Increase if necessary
77   const uint64_t BaseOffset = ST.getExplicitKernelArgOffset(F);
78 
79   unsigned MaxAlign;
80   // FIXME: Alignment is broken broken with explicit arg offset.;
81   const uint64_t TotalKernArgSize = ST.getKernArgSegmentSize(F, MaxAlign);
82   if (TotalKernArgSize == 0)
83     return false;
84 
85   CallInst *KernArgSegment =
86     Builder.CreateIntrinsic(Intrinsic::amdgcn_kernarg_segment_ptr, nullptr,
87                             F.getName() + ".kernarg.segment");
88 
89   KernArgSegment->addAttribute(AttributeList::ReturnIndex, Attribute::NonNull);
90   KernArgSegment->addAttribute(AttributeList::ReturnIndex,
91     Attribute::getWithDereferenceableBytes(Ctx, TotalKernArgSize));
92 
93   unsigned AS = KernArgSegment->getType()->getPointerAddressSpace();
94   uint64_t ExplicitArgOffset = 0;
95 
96   for (Argument &Arg : F.args()) {
97     Type *ArgTy = Arg.getType();
98     unsigned Align = DL.getABITypeAlignment(ArgTy);
99     unsigned Size = DL.getTypeSizeInBits(ArgTy);
100     unsigned AllocSize = DL.getTypeAllocSize(ArgTy);
101 
102     uint64_t EltOffset = alignTo(ExplicitArgOffset, Align) + BaseOffset;
103     ExplicitArgOffset = alignTo(ExplicitArgOffset, Align) + AllocSize;
104 
105     if (Arg.use_empty())
106       continue;
107 
108     if (PointerType *PT = dyn_cast<PointerType>(ArgTy)) {
109       // FIXME: Hack. We rely on AssertZext to be able to fold DS addressing
110       // modes on SI to know the high bits are 0 so pointer adds don't wrap. We
111       // can't represent this with range metadata because it's only allowed for
112       // integer types.
113       if (PT->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS &&
114           ST.getGeneration() == AMDGPUSubtarget::SOUTHERN_ISLANDS)
115         continue;
116 
117       // FIXME: We can replace this with equivalent alias.scope/noalias
118       // metadata, but this appears to be a lot of work.
119       if (Arg.hasNoAliasAttr())
120         continue;
121     }
122 
123     VectorType *VT = dyn_cast<VectorType>(ArgTy);
124     bool IsV3 = VT && VT->getNumElements() == 3;
125     VectorType *V4Ty = nullptr;
126 
127     int64_t AlignDownOffset = alignDown(EltOffset, 4);
128     int64_t OffsetDiff = EltOffset - AlignDownOffset;
129     unsigned AdjustedAlign = MinAlign(KernArgBaseAlign, AlignDownOffset);
130 
131     Value *ArgPtr;
132     if (Size < 32 && !ArgTy->isAggregateType()) { // FIXME: Handle aggregate types
133       // Since we don't have sub-dword scalar loads, avoid doing an extload by
134       // loading earlier than the argument address, and extracting the relevant
135       // bits.
136       //
137       // Additionally widen any sub-dword load to i32 even if suitably aligned,
138       // so that CSE between different argument loads works easily.
139 
140       ArgPtr = Builder.CreateConstInBoundsGEP1_64(
141         KernArgSegment,
142         AlignDownOffset,
143         Arg.getName() + ".kernarg.offset.align.down");
144       ArgPtr = Builder.CreateBitCast(ArgPtr,
145                                      Builder.getInt32Ty()->getPointerTo(AS),
146                                      ArgPtr->getName() + ".cast");
147     } else {
148       ArgPtr = Builder.CreateConstInBoundsGEP1_64(
149         KernArgSegment,
150         AlignDownOffset,
151         Arg.getName() + ".kernarg.offset");
152       ArgPtr = Builder.CreateBitCast(ArgPtr, ArgTy->getPointerTo(AS),
153                                      ArgPtr->getName() + ".cast");
154     }
155 
156     if (IsV3 && Size >= 32) {
157       V4Ty = VectorType::get(VT->getVectorElementType(), 4);
158       // Use the hack that clang uses to avoid SelectionDAG ruining v3 loads
159       ArgPtr = Builder.CreateBitCast(ArgPtr, V4Ty->getPointerTo(AS));
160     }
161 
162     LoadInst *Load = Builder.CreateAlignedLoad(ArgPtr, AdjustedAlign);
163     Load->setMetadata(LLVMContext::MD_invariant_load, MDNode::get(Ctx, {}));
164 
165     MDBuilder MDB(Ctx);
166 
167     if (isa<PointerType>(ArgTy)) {
168       if (Arg.hasNonNullAttr())
169         Load->setMetadata(LLVMContext::MD_nonnull, MDNode::get(Ctx, {}));
170 
171       uint64_t DerefBytes = Arg.getDereferenceableBytes();
172       if (DerefBytes != 0) {
173         Load->setMetadata(
174           LLVMContext::MD_dereferenceable,
175           MDNode::get(Ctx,
176                       MDB.createConstant(
177                         ConstantInt::get(Builder.getInt64Ty(), DerefBytes))));
178       }
179 
180       uint64_t DerefOrNullBytes = Arg.getDereferenceableOrNullBytes();
181       if (DerefOrNullBytes != 0) {
182         Load->setMetadata(
183           LLVMContext::MD_dereferenceable_or_null,
184           MDNode::get(Ctx,
185                       MDB.createConstant(ConstantInt::get(Builder.getInt64Ty(),
186                                                           DerefOrNullBytes))));
187       }
188 
189       unsigned ParamAlign = Arg.getParamAlignment();
190       if (ParamAlign != 0) {
191         Load->setMetadata(
192           LLVMContext::MD_align,
193           MDNode::get(Ctx,
194                       MDB.createConstant(ConstantInt::get(Builder.getInt64Ty(),
195                                                           ParamAlign))));
196       }
197     }
198 
199     // TODO: Convert noalias arg to !noalias
200 
201     if (Size < 32 && !ArgTy->isAggregateType()) {
202       Value *ExtractBits = OffsetDiff == 0 ?
203         Load : Builder.CreateLShr(Load, OffsetDiff * 8);
204 
205       IntegerType *ArgIntTy = Builder.getIntNTy(Size);
206       Value *Trunc = Builder.CreateTrunc(ExtractBits, ArgIntTy);
207       Value *NewVal = Builder.CreateBitCast(Trunc, ArgTy,
208                                             Arg.getName() + ".load");
209       Arg.replaceAllUsesWith(NewVal);
210     } else if (IsV3) {
211       Value *Shuf = Builder.CreateShuffleVector(Load, UndefValue::get(V4Ty),
212                                                 {0, 1, 2},
213                                                 Arg.getName() + ".load");
214       Arg.replaceAllUsesWith(Shuf);
215     } else {
216       Load->setName(Arg.getName() + ".load");
217       Arg.replaceAllUsesWith(Load);
218     }
219   }
220 
221   KernArgSegment->addAttribute(
222     AttributeList::ReturnIndex,
223     Attribute::getWithAlignment(Ctx, std::max(KernArgBaseAlign, MaxAlign)));
224 
225   return true;
226 }
227 
228 INITIALIZE_PASS_BEGIN(AMDGPULowerKernelArguments, DEBUG_TYPE,
229                       "AMDGPU Lower Kernel Arguments", false, false)
230 INITIALIZE_PASS_END(AMDGPULowerKernelArguments, DEBUG_TYPE, "AMDGPU Lower Kernel Arguments",
231                     false, false)
232 
233 char AMDGPULowerKernelArguments::ID = 0;
234 
235 FunctionPass *llvm::createAMDGPULowerKernelArgumentsPass() {
236   return new AMDGPULowerKernelArguments();
237 }
238