[ARM] Fix calling convention for __fp16 with big-endian

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,