[ARM][CodeGen] Fixing stack alignment of HFA arguments on AArch32 PCS

[pratlucas]

Summary

Properly complying with AArch32 PCS on the handling of over-aligned HFA
arguments when those are placed on the stack. AAPCS specifies that the
stacked argument address should be adjusted upwards until correctly
aligned for the argument before copying it to memory.

This patch fixes the alignment of these arguments by makign use of the
stack alignment propagated through the alignstack IR argument
attribute during the calling convention lowering for ARM targets.

[pratlucas]

Ping.

clang/lib/CodeGen/TargetInfo.cpp

if (const VectorType *VT = Base->getAs<VectorType>()) {
     if (!getTarget().hasLegalHalfType() && containsAnyFP16Vectors(Ty)) {
       uint64_t Size = getContext().getTypeSize(VT);
       llvm::Type *NewVecTy = llvm::VectorType::get(
           llvm::Type::getInt32Ty(getVMContext()), Size / 32);
       llvm::Type *Ty = llvm::ArrayType::get(NewVecTy, Members);
       return ABIArgInfo::getDirect(Ty, 0, nullptr, false);
     }
   }
-  return ABIArgInfo::getDirect(nullptr, 0, nullptr, false);
+  bool NeedsStackAlignment = getContext().getTypeAlignInChars(Ty) !=
+                             getContext().getTypeAlignInChars(Base);
+  return ABIArgInfo::getDirect(nullptr, /*Offset=*/0, /*Padding=*/nullptr,
+                               /*CanBeFlattened=*/false, NeedsStackAlignment);
 }
 
 ABIArgInfo ARMABIInfo::classifyArgumentType(QualType Ty, bool isVariadic,
                                             unsigned functionCallConv) const {
   // 6.1.2.1 The following argument types are VFP CPRCs:
   //   A single-precision floating-point type (including promoted
   //   half-precision types); A double-precision floating-point type;
   //   A 64-bit or 128-bit containerized vector type; Homogeneous Aggregate
ABIArgInfo ARMABIInfo::classifyArgumentType(QualType Ty, bool isVariadic,
   } else if (getABIKind() == ARMABIInfo::AAPCS16_VFP) {
     // WatchOS does have homogeneous aggregates. Note that we intentionally use
     // this convention even for a variadic function: the backend will use GPRs
     // if needed.
     const Type *Base = nullptr;
     uint64_t Members = 0;
     if (isHomogeneousAggregate(Ty, Base, Members)) {
       assert(Base && Members <= 4 && "unexpected homogeneous aggregate");
+      bool NeedsStackAlignment = getContext().getTypeAlignInChars(Ty) !=
+                                 getContext().getTypeAlignInChars(Base);
       llvm::Type *Ty =
         llvm::ArrayType::get(CGT.ConvertType(QualType(Base, 0)), Members);
-      return ABIArgInfo::getDirect(Ty, 0, nullptr, false);
+      return ABIArgInfo::getDirect(Ty, /*Offset=*/0, /*Padding=*/nullptr,
+                                   /*CanBeFlattened=*/false,
+                                   NeedsStackAlignment);
     }
   }
 
   if (getABIKind() == ARMABIInfo::AAPCS16_VFP &&
       getContext().getTypeSizeInChars(Ty) > CharUnits::fromQuantity(16)) {
     // WatchOS is adopting the 64-bit AAPCS rule on composite types: if they're
     // bigger than 128-bits, they get placed in space allocated by the caller,
     // and a pointer is passed.

clang/test/CodeGen/arm-aapcs-vfp.c

 typedef struct { int x[4]; } struct_four_ints;
 // CHECK: define arm_aapcs_vfpcc void @test_vfp_stack_gpr_split_5(%struct.struct_seventeen_ints* byval(%struct.struct_seventeen_ints) align 4 %a, double %b, double %c, double %d, double %e, double %f, double %g, double %h, double %i, double %j, [4 x i32] %k.coerce)
 void test_vfp_stack_gpr_split_5(struct_seventeen_ints a, double b, double c, double d, double e, double f, double g, double h, double i, double j, struct_four_ints k) {}
 
 // Here, parameter k would need padding to prevent it from being split, but it
 // is passed ByVal (due to being > 64 bytes), so the backend handles this instead.
 void test_vfp_stack_gpr_split_6(double a, double b, double c, double d, double e, double f, double g, double h, double i, int j, struct_seventeen_ints k) {}
 // CHECK: define arm_aapcs_vfpcc void @test_vfp_stack_gpr_split_6(double %a, double %b, double %c, double %d, double %e, double %f, double %g, double %h, double %i, i32 %j, %struct.struct_seventeen_ints* byval(%struct.struct_seventeen_ints) align 4 %k)
+
+// Make sure over-alignment information is propagated to the backend properly
+typedef struct {
+  __attribute__((__aligned__(8))) float v[2];
+} hfa_align;
+// CHECK: define arm_aapcs_vfpcc float @test_hfa_align_arg(%struct.hfa_align alignstack(8) %h1.coerce) #0
+float test_hfa_align_arg(hfa_align h1) {
+  return h1.v[0];
+}
+// CHECK: %call = call arm_aapcs_vfpcc float @test_hfa_align_arg(%struct.hfa_align alignstack(8) %1) #4
+float test_hfa_align_call() {
+  hfa_align h = {1.0, 2.0};
+  return test_hfa_align_arg(h);
+}

llvm/lib/Target/ARM/ARMCallingConv.cpp

static bool CC_ARM_AAPCS_Custom_Aggregate(unsigned ValNo, MVT ValVT,
   // Mark all regs as unavailable (AAPCS rule C.2.vfp for VFP, C.6 for core)
   for (auto Reg : RegList)
     State.AllocateReg(Reg);
 
   // After the first item has been allocated, the rest are packed as tightly as
   // possible. (E.g. an incoming i64 would have starting Align of 8, but we'll
   // be allocating a bunch of i32 slots).
   unsigned RestAlign = std::min(Align, Size);
-
+  if (ArgFlags.getStackAlign()) {
+    const llvm::Align ArgStackAlign(ArgFlags.getStackAlign());
+    Align = std::max(Align, unsigned(ArgStackAlign.value()));
+  }
   for (auto &It : PendingMembers) {
     It.convertToMem(State.AllocateStack(Size, Align));
     State.addLoc(It);
     Align = RestAlign;
   }
 
   // All pending members have now been allocated
   PendingMembers.clear();
 
   // This will be allocated by the last member of the aggregate
   return true;
 }
 
 // Include the table generated calling convention implementations.
 #include "ARMGenCallingConv.inc"

llvm/test/CodeGen/ARM/aapcs-hfa-code.ll

 ; RUN: llc < %s -mtriple=armv7-linux-gnueabihf -o - | FileCheck %s
 ; RUN: llc < %s -mtriple=thumbv7em-none-eabi -mcpu=cortex-m4 | FileCheck %s --check-prefix=CHECK-M4F
 
 target datalayout = "e-m:e-p:32:32-i64:64-v128:64:128-n32-S64"
 
+%struct.hfa_align = type { [2 x float] }
+
 define arm_aapcs_vfpcc void @test_1float({ float } %a) {
   call arm_aapcs_vfpcc void @test_1float({ float } { float 1.0 })
   ret void
 
 ; CHECK-LABEL: test_1float:
 ; CHECK-DAG: vmov.f32 s0, #1.{{0+}}e+00
 ; CHECK: bl test_1float
 

 ; CHECK-M4F: movs [[ONEHI:r[0-9]+]], #0
 ; CHECK-M4F: movt [[ONEHI]], #16368
 ; CHECK-M4F: movs [[ONELO:r[0-9]+]], #0
 ; CHECK-M4F: strd [[ONELO]], [[ONEHI]], [sp, #8]
 ; CHECK-M4F: bl test_1double_misaligned
 
   ret void
 }
+
+; Over-aligned HFA argument placed on register - one element per register
+define arm_aapcs_vfpcc float @test_hfa_align_reg(%struct.hfa_align alignstack(8) %h1.coerce) local_unnamed_addr #3 {
+entry:
+; CHECK-LABEL: test_hfa_align_reg:
+; CHECK-DAG: bx lr
+
+; CHECK-M4F-LABEL: test_hfa_align_reg:
+; CHECK-M4F-DAG: bx lr
+
+  %h1.coerce.fca.0.0.extract = extractvalue %struct.hfa_align %h1.coerce, 0, 0
+  ret float %h1.coerce.fca.0.0.extract
+}
+
+; Call with over-align HFA argument placed on registers - one element per register
+define arm_aapcs_vfpcc float @test_hfa_align_reg_call() local_unnamed_addr #3 {
+entry:
+; CHECK-LABEL: test_hfa_align_reg_call:
+; CHECK-DAG: vmov.f32 s0, #1.000000e+00
+; CHECK-DAG: vmov.f32 s1, #2.000000e+00
+; CHECK-DAG: bl test_hfa_align_reg
+
+; CHECK-M4F-LABEL: test_hfa_align_reg_call:
+; CHECK-M4F-DAG: vmov.f32 s0, #1.000000e+00
+; CHECK-M4F-DAG: vmov.f32 s1, #2.000000e+00
+; CHECK-M4F-DAG: bl test_hfa_align_reg
+
+  %call = call arm_aapcs_vfpcc float @test_hfa_align_reg(%struct.hfa_align alignstack(8) { [2 x float] [float 1.000000e+00, float 2.000000e+00] }) #5
+  ret float %call
+}
+
+; Over-aligned HFA argument placed on the stack - stack round up to alignment
+define arm_aapcs_vfpcc float @test_hfa_align_stack(double %d0, double %d1, double %d2, double %d3, double %d4, double %d5, double %d6, double %d7, float %f1, %struct.hfa_align alignstack(8) %h1.coerce) local_unnamed_addr #3 {
+entry:
+; CHECK-LABEL: test_hfa_align_stack:
+; CHECK-DAG: vldr s0, [sp, #8]
+; CHECK-DAG: bx lr
+
+; CHECK-M4F-LABEL: test_hfa_align_stack:
+; CHECK-M4F-DAG: vldr s0, [sp, #8]
+; CHECK-M4F-DAG: bx lr
+
+  %h1.coerce.fca.0.0.extract = extractvalue %struct.hfa_align %h1.coerce, 0, 0
+  ret float %h1.coerce.fca.0.0.extract
+}
+
+; Call with over-aligned HFA argument placed on the stack - stack round up to alignment
+define arm_aapcs_vfpcc float @test_hfa_align_stack_call() local_unnamed_addr #3 {
+entry:
+; CHECK-LABEL: test_hfa_align_stack_call:
+; CHECK-DAG: sub  sp, sp, #16
+; CHECK-DAG: mov  r0, #1073741824
+; CHECK-DAG: mov  r1, #1065353216
+; CHECK-DAG: str  r1, [sp, #8]
+; CHECK-DAG: str  r0, [sp, #12]
+; CHECK-DAG: bl test_hfa_align_stack
+; CHECK-DAG: add  sp, sp, #16
+
+; CHECK-M4F-LABEL: test_hfa_align_stack_call:
+; CHECK-M4F-DAG: sub  sp, #16
+; CHECK-M4F-DAG: mov.w  r0, #1073741824
+; CHECK-M4F-DAG: mov.w  r1, #1065353216
+; CHECK-M4F-DAG: strd r1, r0, [sp, #8]
+; CHECK-M4F-DAG: bl test_hfa_align_stack
+; CHECK-M4F-DAG: add  sp, #16
+
+  %call = call arm_aapcs_vfpcc float @test_hfa_align_stack(double undef, double undef, double undef, double undef, double undef, double undef, double undef, double undef, float undef, %struct.hfa_align alignstack(8) { [2 x float] [float 1.000000e+00, float 2.000000e+00] }) #5
+  ret float %call
+}
+