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