Rework to make AArch64 specific & handle splitting into legal chunks

Author: HolyMolyCowMan

llvm/lib/Target/AArch64/AArch64InstrGISel.td

@@ -149,6 +149,13 @@ def G_VLSHR : AArch64GenericInstruction {
   let hasSideEffects = 0;
 }
 
+// Float truncation using round to odd
+def G_FPTRUNC_ODD : AArch64GenericInstruction {
+  let OutOperandList = (outs type0:$dst);
+  let InOperandList = (ins type1:$src);
+  let hasSideEffects = false;
+}
+
 // Represents an integer to FP conversion on the FPR bank.
 def G_SITOF : AArch64GenericInstruction {
   let OutOperandList = (outs type0:$dst);

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

@@ -819,7 +819,7 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
           {{s16, s32}, {s16, s64}, {s32, s64}, {v4s16, v4s32}, {v2s32, v2s64}})
       .libcallFor({{s16, s128}, {s32, s128}, {s64, s128}})
       .moreElementsToNextPow2(1)
-      .lowerIf([](const LegalityQuery &Q) {
+      .customIf([](const LegalityQuery &Q) {
         LLT DstTy = Q.Types[0];
         LLT SrcTy = Q.Types[1];
         return SrcTy.isFixedVector() && DstTy.isFixedVector() &&
@@ -1474,6 +1474,10 @@ bool AArch64LegalizerInfo::legalizeCustom(
     return legalizeICMP(MI, MRI, MIRBuilder);
   case TargetOpcode::G_BITCAST:
     return legalizeBitcast(MI, Helper);
+  case TargetOpcode::G_FPTRUNC:
+    // In order to vectorise f16 to f64 properly, we need to use f32 as an
+    // intermediary
+    return legalizeFptrunc(MI, MIRBuilder, MRI);
   }
 
   llvm_unreachable("expected switch to return");
@@ -2400,3 +2404,80 @@ bool AArch64LegalizerInfo::legalizePrefetch(MachineInstr &MI,
   MI.eraseFromParent();
   return true;
 }
+
+bool AArch64LegalizerInfo::legalizeFptrunc(MachineInstr &MI,
+                                           MachineIRBuilder &MIRBuilder,
+                                           MachineRegisterInfo &MRI) const {
+  auto [Dst, DstTy, Src, SrcTy] = MI.getFirst2RegLLTs();
+  assert(SrcTy.isFixedVector() && isPowerOf2_32(SrcTy.getNumElements()) &&
+         "Expected a power of 2 elements");
+
+  LLT s16 = LLT::scalar(16);
+  LLT s32 = LLT::scalar(32);
+  LLT s64 = LLT::scalar(64);
+  LLT v2s16 = LLT::fixed_vector(2, s16);
+  LLT v4s16 = LLT::fixed_vector(4, s16);
+  LLT v2s32 = LLT::fixed_vector(2, s32);
+  LLT v4s32 = LLT::fixed_vector(4, s32);
+  LLT v2s64 = LLT::fixed_vector(2, s64);
+
+  SmallVector<Register> RegsToUnmergeTo;
+  SmallVector<Register> TruncOddDstRegs;
+  SmallVector<Register> RegsToMerge;
+
+  unsigned ElemCount = SrcTy.getNumElements();
+
+  // Find the biggest size chunks we can work with
+  int StepSize = ElemCount % 4 ? 2 : 4;
+
+  // If we have a power of 2 greater than 2, we need to first unmerge into
+  // enough pieces
+  if (ElemCount <= 2)
+    RegsToUnmergeTo.push_back(Src);
+  else {
+    for (unsigned i = 0; i < ElemCount / 2; ++i) {
+      RegsToUnmergeTo.push_back(MRI.createGenericVirtualRegister(v2s64));
+    }
+
+    MIRBuilder.buildUnmerge(RegsToUnmergeTo, Src);
+  }
+
+  // Create all of the round-to-odd instructions and store them
+  for (auto SrcReg : RegsToUnmergeTo) {
+    Register Mid =
+        MIRBuilder.buildInstr(AArch64::G_FPTRUNC_ODD, {v2s32}, {SrcReg})
+            .getReg(0);
+    TruncOddDstRegs.push_back(Mid);
+  }
+
+  // Truncate 4s32 to 4s16 if we can to reduce instruction count, otherwise truncate 2s32 to 2s16.
+  unsigned Index = 0;
+  for (unsigned LoopIter = 0; LoopIter < ElemCount / StepSize; ++LoopIter) {
+    if (StepSize == 4) {
+      Register ConcatDst =
+          MIRBuilder
+              .buildMergeLikeInstr(
+                  {v4s32}, {TruncOddDstRegs[Index++], TruncOddDstRegs[Index++]})
+              .getReg(0);
+
+      RegsToMerge.push_back(
+          MIRBuilder.buildFPTrunc(v4s16, ConcatDst).getReg(0));
+    } else {
+      RegsToMerge.push_back(
+          MIRBuilder.buildFPTrunc(v2s16, TruncOddDstRegs[Index++]).getReg(0));
+    }
+  }
+
+  // If there is only one register, replace the destination
+  if (RegsToMerge.size() == 1) {
+    MRI.replaceRegWith(Dst, RegsToMerge.pop_back_val());
+    MI.eraseFromParent();
+    return true;
+  }
+
+  // Merge the rest of the instructions & replace the register
+  Register Fin = MIRBuilder.buildMergeLikeInstr(DstTy, RegsToMerge).getReg(0);
+  MRI.replaceRegWith(Dst, Fin);
+  MI.eraseFromParent();
+  return true;
+}