[ARM] Fix calling convention for __fp16 with big-endian (#126741)

AAPCS32 defines the fp16 and bf16 types as being passed as if they were
extended to 32 bits, with the high 16 bits being unspecified. The
extension is specified as happening as-if it was done in a register,
which means that for big endian targets, the actual value gets passed in
the higher addressed half of the stack slot, instead of the lower
addressed half as for little endian. Previously, for targets with the
fp16 extension, we were passing these types as a 16 bit stack slot,
which worked for little endian because every later stack slot would be
4-byte aligned leaving the 2 byte gap, but was incorrect for big endian.

Author: ostannard

llvm/lib/Target/ARM/ARMCallingConv.cpp

@@ -298,7 +298,8 @@ static bool CustomAssignInRegList(unsigned ValNo, MVT ValVT, MVT LocVT,
 static bool CC_ARM_AAPCS_Custom_f16(unsigned ValNo, MVT ValVT, MVT LocVT,
                                     CCValAssign::LocInfo LocInfo,
                                     ISD::ArgFlagsTy ArgFlags, CCState &State) {
-  // f16 arguments are extended to i32 and assigned to a register in [r0, r3]
+  // f16 and bf16 arguments are extended to i32 and assigned to a register in
+  // [r0, r3].
   return CustomAssignInRegList(ValNo, ValVT, MVT::i32, LocInfo, State,
                                RRegList);
 }
@@ -307,10 +308,25 @@ static bool CC_ARM_AAPCS_VFP_Custom_f16(unsigned ValNo, MVT ValVT, MVT LocVT,
                                         CCValAssign::LocInfo LocInfo,
                                         ISD::ArgFlagsTy ArgFlags,
                                         CCState &State) {
-  // f16 arguments are extended to f32 and assigned to a register in [s0, s15]
+  // f16 and bf16 arguments are extended to f32 and assigned to a register in
+  // [s0, s15].
   return CustomAssignInRegList(ValNo, ValVT, MVT::f32, LocInfo, State,
                                SRegList);
 }
 
+static bool CC_ARM_AAPCS_Common_Custom_f16_Stack(unsigned ValNo, MVT ValVT,
+                                                 MVT LocVT,
+                                                 CCValAssign::LocInfo LocInfo,
+                                                 ISD::ArgFlagsTy ArgFlags,
+                                                 CCState &State) {
+  // f16 and bf16 (if not passed in a register) are assigned to a 32-bit stack
+  // slot, with the most-significant 16 bits unspecified. The 32-bit slot is
+  // important to make sure that the byte ordering is correct for big endian
+  // targets.
+  State.addLoc(CCValAssign::getCustomMem(
+      ValNo, ValVT, State.AllocateStack(4, Align(4)), MVT::i32, LocInfo));
+  return true;
+}
+
 // Include the table generated calling convention implementations.
 #include "ARMGenCallingConv.inc"

llvm/lib/Target/ARM/ARMCallingConv.td

@@ -139,7 +139,8 @@ def CC_ARM_AAPCS_Common : CallingConv<[
 
   CCIfType<[i32], CCIfAlign<"8", CCAssignToStackWithShadow<4, 8, [R0, R1, R2, R3]>>>,
   CCIfType<[i32], CCAssignToStackWithShadow<4, 4, [R0, R1, R2, R3]>>,
-  CCIfType<[f16, bf16, f32], CCAssignToStackWithShadow<4, 4, [Q0, Q1, Q2, Q3]>>,
+  CCIfType<[f32], CCAssignToStackWithShadow<4, 4, [Q0, Q1, Q2, Q3]>>,
+  CCIfType<[f16, bf16], CCCustom<"CC_ARM_AAPCS_Common_Custom_f16_Stack">>,
   CCIfType<[f64], CCAssignToStackWithShadow<8, 8, [Q0, Q1, Q2, Q3]>>,
   CCIfType<[v2f64], CCIfAlign<"16",
            CCAssignToStackWithShadow<16, 16, [Q0, Q1, Q2, Q3]>>>,

llvm/lib/Target/ARM/ARMISelLowering.cpp

@@ -4759,6 +4759,25 @@ SDValue ARMTargetLowering::LowerFormalArguments(
                 VA.getLocMemOffset(), Flags.getByValSize());
             InVals.push_back(DAG.getFrameIndex(FrameIndex, PtrVT));
             CCInfo.nextInRegsParam();
+          } else if (VA.needsCustom() && (VA.getValVT() == MVT::f16 ||
+                                          VA.getValVT() == MVT::bf16)) {
+            // f16 and bf16 values are passed in the least-significant half of
+            // a 4 byte stack slot. This is done as-if the extension was done
+            // in a 32-bit register, so the actual bytes used for the value
+            // differ between little and big endian.
+            assert(VA.getLocVT().getSizeInBits() == 32);
+            unsigned FIOffset = VA.getLocMemOffset();
+            int FI = MFI.CreateFixedObject(VA.getLocVT().getSizeInBits() / 8,
+                                           FIOffset, true);
+
+            SDValue Addr = DAG.getFrameIndex(FI, PtrVT);
+            if (DAG.getDataLayout().isBigEndian())
+              Addr = DAG.getObjectPtrOffset(dl, Addr, TypeSize::getFixed(2));
+
+            InVals.push_back(DAG.getLoad(VA.getValVT(), dl, Chain, Addr,
+                                         MachinePointerInfo::getFixedStack(
+                                             DAG.getMachineFunction(), FI)));
+
           } else {
             unsigned FIOffset = VA.getLocMemOffset();
             int FI = MFI.CreateFixedObject(VA.getLocVT().getSizeInBits()/8,