AMDGPU: Don't use struct type for argument layout

This was introducing unnecessary padding after the explicit
arguments, depending on the alignment of the total struct type.
Also has the side effect of avoiding creating an extra GEP for
the offset from the base kernel argument to the explicit kernel
argument offset.

llvm-svn: 335999

1 files changed, 40 insertions, 44 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp
index 438d6c8ab32..b06f3014ac6 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp
@@ -74,31 +74,11 @@ bool AMDGPULowerKernelArguments::runOnFunction(Function &F) {
   BasicBlock &EntryBlock = *F.begin();
   IRBuilder<> Builder(&*EntryBlock.begin());
 
-  SmallVector<Type *, 16> ArgTypes;
-  for (Argument &Arg : F.args()) {
-    Type *ArgTy = Arg.getType();
-    unsigned Size = DL.getTypeStoreSizeInBits(ArgTy);
-    bool IsExtArg = Size < 32 && (Arg.hasZExtAttr() || Arg.hasSExtAttr()) &&
-                    !ST.isAmdHsaOS();
-
-    // Clover seems to always pad i8/i16 to i32, but doesn't properly align
-    // them?
-    // Make sure the struct elements have correct size and alignment for ext
-    // args. These seem to be padded up to 4-bytes but not correctly aligned.
-    ArgTypes.push_back(
-      IsExtArg ? ArrayType::get(ArgTy, 32 / Size) : Arg.getType());
-  }
-
-  StructType *ArgStructTy = StructType::create(Ctx, ArgTypes, F.getName());
-  const StructLayout *Layout = DL.getStructLayout(ArgStructTy);
-
-  // Minimum alignment for kern segment is 16.
-  unsigned KernArgBaseAlign = std::max(16u, DL.getABITypeAlignment(ArgStructTy));
+  const unsigned KernArgBaseAlign = 16; // FIXME: Increase if necessary
   const uint64_t BaseOffset = ST.getExplicitKernelArgOffset(F);
 
   // FIXME: Alignment is broken broken with explicit arg offset.;
-  const uint64_t TotalKernArgSize = BaseOffset +
-    ST.getKernArgSegmentSize(F, DL.getTypeAllocSize(ArgStructTy));
+  const uint64_t TotalKernArgSize = ST.getKernArgSegmentSize(F);
   if (TotalKernArgSize == 0)
     return false;
 
@@ -109,23 +89,34 @@ bool AMDGPULowerKernelArguments::runOnFunction(Function &F) {
   KernArgSegment->addAttribute(AttributeList::ReturnIndex, Attribute::NonNull);
   KernArgSegment->addAttribute(AttributeList::ReturnIndex,
     Attribute::getWithDereferenceableBytes(Ctx, TotalKernArgSize));
-  KernArgSegment->addAttribute(AttributeList::ReturnIndex,
-    Attribute::getWithAlignment(Ctx, KernArgBaseAlign));
-
-  Value *KernArgBase = KernArgSegment;
-  if (BaseOffset != 0) {
-    KernArgBase = Builder.CreateConstInBoundsGEP1_64(KernArgBase, BaseOffset);
-    KernArgBaseAlign = MinAlign(KernArgBaseAlign, BaseOffset);
-  }
 
   unsigned AS = KernArgSegment->getType()->getPointerAddressSpace();
-  Value *CastStruct = Builder.CreateBitCast(KernArgBase,
-                                            ArgStructTy->getPointerTo(AS));
+  unsigned MaxAlign = 1;
+  uint64_t ExplicitArgOffset = 0;
+
   for (Argument &Arg : F.args()) {
+    Type *ArgTy = Arg.getType();
+    unsigned Align = DL.getABITypeAlignment(ArgTy);
+    MaxAlign = std::max(Align, MaxAlign);
+    unsigned Size = DL.getTypeSizeInBits(ArgTy);
+    unsigned AllocSize = DL.getTypeAllocSize(ArgTy);
+
+
+    // Clover seems to always pad i8/i16 to i32, but doesn't properly align
+    // them?
+    // Make sure the struct elements have correct size and alignment for ext
+    // args. These seem to be padded up to 4-bytes but not correctly aligned.
+    bool IsExtArg = AllocSize < 32 && (Arg.hasZExtAttr() || Arg.hasSExtAttr()) &&
+                    !ST.isAmdHsaOS();
+    if (IsExtArg)
+      AllocSize = 4;
+
+    uint64_t EltOffset = alignTo(ExplicitArgOffset, Align) + BaseOffset;
+    ExplicitArgOffset = alignTo(ExplicitArgOffset, Align) + AllocSize;
+
     if (Arg.use_empty())
       continue;
 
-    Type *ArgTy = Arg.getType();
     if (PointerType *PT = dyn_cast<PointerType>(ArgTy)) {
       // FIXME: Hack. We rely on AssertZext to be able to fold DS addressing
       // modes on SI to know the high bits are 0 so pointer adds don't wrap. We
@@ -145,10 +136,6 @@ bool AMDGPULowerKernelArguments::runOnFunction(Function &F) {
     bool IsV3 = VT && VT->getNumElements() == 3;
     VectorType *V4Ty = nullptr;
 
-    unsigned Size = DL.getTypeSizeInBits(ArgTy);
-    bool IsExtArg = Size < 32 && (Arg.hasZExtAttr() || Arg.hasSExtAttr()) &&
-                    !ST.isAmdHsaOS();
-    int64_t EltOffset = Layout->getElementOffset(Arg.getArgNo());
     int64_t AlignDownOffset = alignDown(EltOffset, 4);
     int64_t OffsetDiff = EltOffset - AlignDownOffset;
     unsigned AdjustedAlign = MinAlign(KernArgBaseAlign, AlignDownOffset);
@@ -162,19 +149,24 @@ bool AMDGPULowerKernelArguments::runOnFunction(Function &F) {
       // Additionally widen any sub-dword load to i32 even if suitably aligned,
       // so that CSE between different argument loads works easily.
 
-      ArgPtr = Builder.CreateConstGEP1_64(KernArgBase, AlignDownOffset);
-      ArgPtr = Builder.CreateBitCast(
-        ArgPtr,
-        Builder.getInt32Ty()->getPointerTo(AS),
+      ArgPtr = Builder.CreateConstInBoundsGEP1_64(
+        KernArgSegment,
+        AlignDownOffset,
         Arg.getName() + ".kernarg.offset.align.down");
+      ArgPtr = Builder.CreateBitCast(ArgPtr,
+                                     Builder.getInt32Ty()->getPointerTo(AS),
+                                     ArgPtr->getName() + ".cast");
     } else {
-      ArgPtr = Builder.CreateStructGEP(CastStruct, Arg.getArgNo(),
-                                       Arg.getName() + ".kernarg.offset");
+      ArgPtr = Builder.CreateConstInBoundsGEP1_64(
+        KernArgSegment,
+        AlignDownOffset,
+        Arg.getName() + ".kernarg.offset");
+      ArgPtr = Builder.CreateBitCast(ArgPtr, ArgTy->getPointerTo(AS),
+                                     ArgPtr->getName() + ".cast");
     }
 
     assert((!IsExtArg || !IsV3) && "incompatible situation");
 
-
     if (IsV3 && Size >= 32) {
       V4Ty = VectorType::get(VT->getVectorElementType(), 4);
       // Use the hack that clang uses to avoid SelectionDAG ruining v3 loads
@@ -254,6 +246,10 @@ bool AMDGPULowerKernelArguments::runOnFunction(Function &F) {
     }
   }
 
+  KernArgSegment->addAttribute(
+    AttributeList::ReturnIndex,
+    Attribute::getWithAlignment(Ctx, std::max(KernArgBaseAlign, MaxAlign)));
+
   return true;
 }