[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,24 @@ 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.
+            unsigned FIOffset = VA.getLocMemOffset();
+            int FI = MFI.CreateFixedObject(VA.getLocVT().getSizeInBits() / 8,
+                                           FIOffset, true);
+
+            // Create load nodes to retrieve arguments from the stack.
+            SDValue FIN = DAG.getFrameIndex(FI, PtrVT);
+            SDValue Load = DAG.getLoad(VA.getLocVT(), dl, Chain, FIN,
+                                       MachinePointerInfo::getFixedStack(
+                                           DAG.getMachineFunction(), FI));
+            InVals.push_back(
+                MoveToHPR(dl, DAG, VA.getLocVT(), VA.getValVT(), Load));
+
           } else {
             unsigned FIOffset = VA.getLocMemOffset();
             int FI = MFI.CreateFixedObject(VA.getLocVT().getSizeInBits()/8,
@@ -15257,10 +15275,19 @@ static SDValue PerformVMOVhrCombine(SDNode *N,
   // fold (VMOVhr (load x)) -> (load (f16*)x)
   if (LoadSDNode *LN0 = dyn_cast<LoadSDNode>(Op0)) {
     if (LN0->hasOneUse() && LN0->isUnindexed() &&
-        LN0->getMemoryVT() == MVT::i16) {
-      SDValue Load =
-          DCI.DAG.getLoad(N->getValueType(0), SDLoc(N), LN0->getChain(),
-                          LN0->getBasePtr(), LN0->getMemOperand());
+        (LN0->getMemoryVT() == MVT::i16 || LN0->getMemoryVT() == MVT::i32)) {
+      SDValue Addr = LN0->getBasePtr();
+      unsigned PtrOffset = 0;
+      if (DCI.DAG.getDataLayout().isBigEndian() &&
+          LN0->getMemoryVT() == MVT::i32) {
+        PtrOffset = 2;
+        Addr = DCI.DAG.getObjectPtrOffset(SDLoc(N), Addr,
+                                          TypeSize::getFixed(PtrOffset));
+      }
+      SDValue Load = DCI.DAG.getLoad(
+          N->getValueType(0), SDLoc(N), LN0->getChain(), Addr,
+          LN0->getPointerInfo().getWithOffset(PtrOffset), LN0->getAlign(),
+          LN0->getMemOperand()->getFlags(), LN0->getAAInfo());
       DCI.DAG.ReplaceAllUsesOfValueWith(SDValue(N, 0), Load.getValue(0));
       DCI.DAG.ReplaceAllUsesOfValueWith(Op0.getValue(1), Load.getValue(1));
       return Load;