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