Home | History | Annotate | Download | only in AMDGPU
      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 GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(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   const unsigned KernArgBaseAlign = 16; // FIXME: Increase if necessary
     78   const uint64_t BaseOffset = ST.getExplicitKernelArgOffset(F);
     79 
     80   unsigned MaxAlign;
     81   // FIXME: Alignment is broken broken with explicit arg offset.;
     82   const uint64_t TotalKernArgSize = ST.getKernArgSegmentSize(F, MaxAlign);
     83   if (TotalKernArgSize == 0)
     84     return false;
     85 
     86   CallInst *KernArgSegment =
     87     Builder.CreateIntrinsic(Intrinsic::amdgcn_kernarg_segment_ptr, nullptr,
     88                             F.getName() + ".kernarg.segment");
     89 
     90   KernArgSegment->addAttribute(AttributeList::ReturnIndex, Attribute::NonNull);
     91   KernArgSegment->addAttribute(AttributeList::ReturnIndex,
     92     Attribute::getWithDereferenceableBytes(Ctx, TotalKernArgSize));
     93 
     94   unsigned AS = KernArgSegment->getType()->getPointerAddressSpace();
     95   uint64_t ExplicitArgOffset = 0;
     96 
     97   for (Argument &Arg : F.args()) {
     98     Type *ArgTy = Arg.getType();
     99     unsigned Align = DL.getABITypeAlignment(ArgTy);
    100     unsigned Size = DL.getTypeSizeInBits(ArgTy);
    101     unsigned AllocSize = DL.getTypeAllocSize(ArgTy);
    102 
    103     uint64_t EltOffset = alignTo(ExplicitArgOffset, Align) + BaseOffset;
    104     ExplicitArgOffset = alignTo(ExplicitArgOffset, Align) + AllocSize;
    105 
    106     if (Arg.use_empty())
    107       continue;
    108 
    109     if (PointerType *PT = dyn_cast<PointerType>(ArgTy)) {
    110       // FIXME: Hack. We rely on AssertZext to be able to fold DS addressing
    111       // modes on SI to know the high bits are 0 so pointer adds don't wrap. We
    112       // can't represent this with range metadata because it's only allowed for
    113       // integer types.
    114       if (PT->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS &&
    115           ST.getGeneration() == AMDGPUSubtarget::SOUTHERN_ISLANDS)
    116         continue;
    117 
    118       // FIXME: We can replace this with equivalent alias.scope/noalias
    119       // metadata, but this appears to be a lot of work.
    120       if (Arg.hasNoAliasAttr())
    121         continue;
    122     }
    123 
    124     VectorType *VT = dyn_cast<VectorType>(ArgTy);
    125     bool IsV3 = VT && VT->getNumElements() == 3;
    126     VectorType *V4Ty = nullptr;
    127 
    128     int64_t AlignDownOffset = alignDown(EltOffset, 4);
    129     int64_t OffsetDiff = EltOffset - AlignDownOffset;
    130     unsigned AdjustedAlign = MinAlign(KernArgBaseAlign, AlignDownOffset);
    131 
    132     Value *ArgPtr;
    133     if (Size < 32 && !ArgTy->isAggregateType()) { // FIXME: Handle aggregate types
    134       // Since we don't have sub-dword scalar loads, avoid doing an extload by
    135       // loading earlier than the argument address, and extracting the relevant
    136       // bits.
    137       //
    138       // Additionally widen any sub-dword load to i32 even if suitably aligned,
    139       // so that CSE between different argument loads works easily.
    140 
    141       ArgPtr = Builder.CreateConstInBoundsGEP1_64(
    142         KernArgSegment,
    143         AlignDownOffset,
    144         Arg.getName() + ".kernarg.offset.align.down");
    145       ArgPtr = Builder.CreateBitCast(ArgPtr,
    146                                      Builder.getInt32Ty()->getPointerTo(AS),
    147                                      ArgPtr->getName() + ".cast");
    148     } else {
    149       ArgPtr = Builder.CreateConstInBoundsGEP1_64(
    150         KernArgSegment,
    151         AlignDownOffset,
    152         Arg.getName() + ".kernarg.offset");
    153       ArgPtr = Builder.CreateBitCast(ArgPtr, ArgTy->getPointerTo(AS),
    154                                      ArgPtr->getName() + ".cast");
    155     }
    156 
    157     if (IsV3 && Size >= 32) {
    158       V4Ty = VectorType::get(VT->getVectorElementType(), 4);
    159       // Use the hack that clang uses to avoid SelectionDAG ruining v3 loads
    160       ArgPtr = Builder.CreateBitCast(ArgPtr, V4Ty->getPointerTo(AS));
    161     }
    162 
    163     LoadInst *Load = Builder.CreateAlignedLoad(ArgPtr, AdjustedAlign);
    164     Load->setMetadata(LLVMContext::MD_invariant_load, MDNode::get(Ctx, {}));
    165 
    166     MDBuilder MDB(Ctx);
    167 
    168     if (isa<PointerType>(ArgTy)) {
    169       if (Arg.hasNonNullAttr())
    170         Load->setMetadata(LLVMContext::MD_nonnull, MDNode::get(Ctx, {}));
    171 
    172       uint64_t DerefBytes = Arg.getDereferenceableBytes();
    173       if (DerefBytes != 0) {
    174         Load->setMetadata(
    175           LLVMContext::MD_dereferenceable,
    176           MDNode::get(Ctx,
    177                       MDB.createConstant(
    178                         ConstantInt::get(Builder.getInt64Ty(), DerefBytes))));
    179       }
    180 
    181       uint64_t DerefOrNullBytes = Arg.getDereferenceableOrNullBytes();
    182       if (DerefOrNullBytes != 0) {
    183         Load->setMetadata(
    184           LLVMContext::MD_dereferenceable_or_null,
    185           MDNode::get(Ctx,
    186                       MDB.createConstant(ConstantInt::get(Builder.getInt64Ty(),
    187                                                           DerefOrNullBytes))));
    188       }
    189 
    190       unsigned ParamAlign = Arg.getParamAlignment();
    191       if (ParamAlign != 0) {
    192         Load->setMetadata(
    193           LLVMContext::MD_align,
    194           MDNode::get(Ctx,
    195                       MDB.createConstant(ConstantInt::get(Builder.getInt64Ty(),
    196                                                           ParamAlign))));
    197       }
    198     }
    199 
    200     // TODO: Convert noalias arg to !noalias
    201 
    202     if (Size < 32 && !ArgTy->isAggregateType()) {
    203       Value *ExtractBits = OffsetDiff == 0 ?
    204         Load : Builder.CreateLShr(Load, OffsetDiff * 8);
    205 
    206       IntegerType *ArgIntTy = Builder.getIntNTy(Size);
    207       Value *Trunc = Builder.CreateTrunc(ExtractBits, ArgIntTy);
    208       Value *NewVal = Builder.CreateBitCast(Trunc, ArgTy,
    209                                             Arg.getName() + ".load");
    210       Arg.replaceAllUsesWith(NewVal);
    211     } else if (IsV3) {
    212       Value *Shuf = Builder.CreateShuffleVector(Load, UndefValue::get(V4Ty),
    213                                                 {0, 1, 2},
    214                                                 Arg.getName() + ".load");
    215       Arg.replaceAllUsesWith(Shuf);
    216     } else {
    217       Load->setName(Arg.getName() + ".load");
    218       Arg.replaceAllUsesWith(Load);
    219     }
    220   }
    221 
    222   KernArgSegment->addAttribute(
    223     AttributeList::ReturnIndex,
    224     Attribute::getWithAlignment(Ctx, std::max(KernArgBaseAlign, MaxAlign)));
    225 
    226   return true;
    227 }
    228 
    229 INITIALIZE_PASS_BEGIN(AMDGPULowerKernelArguments, DEBUG_TYPE,
    230                       "AMDGPU Lower Kernel Arguments", false, false)
    231 INITIALIZE_PASS_END(AMDGPULowerKernelArguments, DEBUG_TYPE, "AMDGPU Lower Kernel Arguments",
    232                     false, false)
    233 
    234 char AMDGPULowerKernelArguments::ID = 0;
    235 
    236 FunctionPass *llvm::createAMDGPULowerKernelArgumentsPass() {
    237   return new AMDGPULowerKernelArguments();
    238 }
    239