[LLVM][CodeGen] Add lowering for scalable vector bfloat operations.

Specifically:
  fabs, fadd, fceil, fdiv, ffloor, fma, fmax, fmaxnm, fmin, fminnm,
  fmul, fnearbyint, fneg, frint, fround, froundeven, fsub, fsqrt &
  ftrunc

Author: paulwalker-arm

llvm/include/llvm/CodeGen/SelectionDAG.h

@@ -1565,6 +1565,12 @@ class SelectionDAG {
   SDValue getSetFPEnv(SDValue Chain, const SDLoc &dl, SDValue Ptr, EVT MemVT,
                       MachineMemOperand *MMO);
 
+  SDValue getExtractSubvector(const SDLoc &DL, EVT VT, SDValue V,
+                              uint64_t Idx) {
+    return getNode(ISD::EXTRACT_SUBVECTOR, DL, VT, V,
+                   getVectorIdxConstant(Idx, DL));
+  }
+
   /// Construct a node to track a Value* through the backend.
   SDValue getSrcValue(const Value *v);
 

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

@@ -191,6 +191,9 @@ class SDValue {
     return getValueType().getSimpleVT();
   }
 
+  /// Return the scalar ValueType of the referenced return value.
+  EVT getScalarValueType() const { return getValueType().getScalarType(); }
+
   /// Returns the size of the value in bits.
   ///
   /// If the value type is a scalable vector type, the scalable property will

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -1663,12 +1663,32 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
     for (auto VT : {MVT::nxv2bf16, MVT::nxv4bf16, MVT::nxv8bf16}) {
       setOperationAction(ISD::BITCAST, VT, Custom);
       setOperationAction(ISD::CONCAT_VECTORS, VT, Custom);
+      setOperationAction(ISD::FCEIL, VT, Custom);
+      setOperationAction(ISD::FDIV, VT, Custom);
+      setOperationAction(ISD::FFLOOR, VT, Custom);
+      setOperationAction(ISD::FMA, VT, Custom);
+      setOperationAction(ISD::FMAXIMUM, VT, Custom);
+      setOperationAction(ISD::FMAXNUM, VT, Custom);
+      setOperationAction(ISD::FMINIMUM, VT, Custom);
+      setOperationAction(ISD::FMINNUM, VT, Custom);
+      setOperationAction(ISD::FNEARBYINT, VT, Custom);
       setOperationAction(ISD::FP_EXTEND, VT, Custom);
       setOperationAction(ISD::FP_ROUND, VT, Custom);
+      setOperationAction(ISD::FRINT, VT, Custom);
+      setOperationAction(ISD::FROUND, VT, Custom);
+      setOperationAction(ISD::FROUNDEVEN, VT, Custom);
+      setOperationAction(ISD::FSQRT, VT, Custom);
+      setOperationAction(ISD::FTRUNC, VT, Custom);
       setOperationAction(ISD::MLOAD, VT, Custom);
       setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
       setOperationAction(ISD::SPLAT_VECTOR, VT, Legal);
       setOperationAction(ISD::VECTOR_SPLICE, VT, Custom);
+
+      if (!Subtarget->hasSVEB16B16()) {
+        setOperationAction(ISD::FADD, VT, Custom);
+        setOperationAction(ISD::FMUL, VT, Custom);
+        setOperationAction(ISD::FSUB, VT, Custom);
+      }
     }
 
     setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::i8, Custom);
@@ -7051,32 +7071,58 @@ SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::UMULO:
     return LowerXALUO(Op, DAG);
   case ISD::FADD:
+    if (Op.getScalarValueType() == MVT::bf16 && !Subtarget->hasSVEB16B16())
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FADD_PRED);
   case ISD::FSUB:
+    if (Op.getScalarValueType() == MVT::bf16 && !Subtarget->hasSVEB16B16())
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FSUB_PRED);
   case ISD::FMUL:
+    if (Op.getScalarValueType() == MVT::bf16 && !Subtarget->hasSVEB16B16())
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FMUL_PRED);
   case ISD::FMA:
+    if (Op.getScalarValueType() == MVT::bf16 && !Subtarget->hasSVEB16B16())
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FMA_PRED);
   case ISD::FDIV:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FDIV_PRED);
   case ISD::FNEG:
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FNEG_MERGE_PASSTHRU);
   case ISD::FCEIL:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FCEIL_MERGE_PASSTHRU);
   case ISD::FFLOOR:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FFLOOR_MERGE_PASSTHRU);
   case ISD::FNEARBYINT:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FNEARBYINT_MERGE_PASSTHRU);
   case ISD::FRINT:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FRINT_MERGE_PASSTHRU);
   case ISD::FROUND:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FROUND_MERGE_PASSTHRU);
   case ISD::FROUNDEVEN:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FROUNDEVEN_MERGE_PASSTHRU);
   case ISD::FTRUNC:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FTRUNC_MERGE_PASSTHRU);
   case ISD::FSQRT:
+    if (Op.getScalarValueType() == MVT::bf16)
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FSQRT_MERGE_PASSTHRU);
   case ISD::FABS:
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FABS_MERGE_PASSTHRU);
@@ -7242,12 +7288,20 @@ SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::SUB:
     return LowerToScalableOp(Op, DAG);
   case ISD::FMAXIMUM:
+    if (Op.getScalarValueType() == MVT::bf16 && !Subtarget->hasSVEB16B16())
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FMAX_PRED);
   case ISD::FMAXNUM:
+    if (Op.getScalarValueType() == MVT::bf16 && !Subtarget->hasSVEB16B16())
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FMAXNM_PRED);
   case ISD::FMINIMUM:
+    if (Op.getScalarValueType() == MVT::bf16 && !Subtarget->hasSVEB16B16())
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FMIN_PRED);
   case ISD::FMINNUM:
+    if (Op.getScalarValueType() == MVT::bf16 && !Subtarget->hasSVEB16B16())
+      return LowerBFloatOp(Op, DAG);
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FMINNM_PRED);
   case ISD::VSELECT:
     return LowerFixedLengthVectorSelectToSVE(Op, DAG);
@@ -28466,6 +28520,40 @@ SDValue AArch64TargetLowering::LowerFixedLengthInsertVectorElt(
   return convertFromScalableVector(DAG, VT, ScalableRes);
 }
 
+// Lower bfloat16 operations by upcasting to float32, performing the operation
+// and then downcasting the result back to bfloat16.
+SDValue AArch64TargetLowering::LowerBFloatOp(SDValue Op,
+                                             SelectionDAG &DAG) const {
+  SDLoc DL(Op);
+  EVT VT = Op.getValueType();
+  assert(isTypeLegal(VT) && VT.isScalableVector() && "Unexpected type!");
+
+  // Split the vector and try again.
+  if (VT == MVT::nxv8bf16) {
+    SmallVector<SDValue, 4> LoOps, HiOps;
+    for (const SDValue &V : Op->op_values()) {
+      LoOps.push_back(DAG.getExtractSubvector(DL, MVT::nxv4bf16, V, 0));
+      HiOps.push_back(DAG.getExtractSubvector(DL, MVT::nxv4bf16, V, 4));
+    }
+
+    unsigned Opc = Op.getOpcode();
+    SDValue SplitOpLo = DAG.getNode(Opc, DL, MVT::nxv4bf16, LoOps);
+    SDValue SplitOpHi = DAG.getNode(Opc, DL, MVT::nxv4bf16, HiOps);
+    return DAG.getNode(ISD::CONCAT_VECTORS, DL, VT, SplitOpLo, SplitOpHi);
+  }
+
+  // Promote to float and try again.
+  EVT PromoteVT = VT.changeVectorElementType(MVT::f32);
+
+  SmallVector<SDValue, 4> Ops;
+  for (const SDValue &V : Op->op_values())
+    Ops.push_back(DAG.getNode(ISD::FP_EXTEND, DL, PromoteVT, V));
+
+  SDValue PromotedOp = DAG.getNode(Op.getOpcode(), DL, PromoteVT, Ops);
+  return DAG.getNode(ISD::FP_ROUND, DL, VT, PromotedOp,
+                     DAG.getIntPtrConstant(0, DL, true));
+}
+
 // Convert vector operation 'Op' to an equivalent predicated operation whereby
 // the original operation's type is used to construct a suitable predicate.
 // NOTE: The results for inactive lanes are undefined.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -1224,6 +1224,7 @@ class AArch64TargetLowering : public TargetLowering {
   SDValue LowerWindowsDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerInlineDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerBFloatOp(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue LowerAVG(SDValue Op, SelectionDAG &DAG, unsigned NewOp) const;
 

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

@@ -663,6 +663,13 @@ let Predicates = [HasSVEorSME] in {
   defm FABS_ZPmZ : sve_int_un_pred_arit_1_fp<0b100, "fabs", AArch64fabs_mt>;
   defm FNEG_ZPmZ : sve_int_un_pred_arit_1_fp<0b101, "fneg", AArch64fneg_mt>;
 
+  foreach VT = [nxv2bf16, nxv4bf16, nxv8bf16] in {
+    def : Pat<(VT (fabs VT:$op)),
+              (AND_ZI $op, (i64 (logical_imm64_XFORM(i64 0x7fff7fff7fff7fff))))>;
+    def : Pat<(VT (fneg VT:$op)),
+              (EOR_ZI $op, (i64 (logical_imm64_XFORM(i64 0x8000800080008000))))>;
+  }
+
   // zext(cmpeq(x, splat(0))) -> cnot(x)
   def : Pat<(nxv16i8 (zext (nxv16i1 (AArch64setcc_z (nxv16i1 (SVEAllActive):$Pg), nxv16i8:$Op2, (SVEDup0), SETEQ)))),
             (CNOT_ZPmZ_B $Op2, $Pg, $Op2)>;

llvm/lib/Target/AArch64/SVEInstrFormats.td

@@ -2299,6 +2299,8 @@ multiclass sve_fp_3op_u_zd_bfloat<bits<3> opc, string asm, SDPatternOperator op>
   def NAME : sve_fp_3op_u_zd<0b00, opc, asm, ZPR16>;
 
   def : SVE_2_Op_Pat<nxv8bf16, op, nxv8bf16, nxv8bf16, !cast<Instruction>(NAME)>;
+  def : SVE_2_Op_Pat<nxv4bf16, op, nxv4bf16, nxv4bf16, !cast<Instruction>(NAME)>;
+  def : SVE_2_Op_Pat<nxv2bf16, op, nxv2bf16, nxv2bf16, !cast<Instruction>(NAME)>;
 }
 
 multiclass sve_fp_3op_u_zd_ftsmul<bits<3> opc, string asm, SDPatternOperator op> {
@@ -9078,6 +9080,8 @@ multiclass sve_fp_bin_pred_bfloat<SDPatternOperator op> {
   def _UNDEF : PredTwoOpPseudo<NAME, ZPR16, FalseLanesUndef>;
 
   def : SVE_3_Op_Pat<nxv8bf16, op, nxv8i1,  nxv8bf16, nxv8bf16, !cast<Pseudo>(NAME # _UNDEF)>;
+  def : SVE_3_Op_Pat<nxv4bf16, op, nxv4i1,  nxv4bf16, nxv4bf16, !cast<Pseudo>(NAME # _UNDEF)>;
+  def : SVE_3_Op_Pat<nxv2bf16, op, nxv2i1,  nxv2bf16, nxv2bf16, !cast<Pseudo>(NAME # _UNDEF)>;
 }
 
 // Predicated pseudo floating point three operand instructions.
@@ -9099,6 +9103,8 @@ multiclass sve_fp_3op_pred_bfloat<SDPatternOperator op> {
   def _UNDEF : PredThreeOpPseudo<NAME, ZPR16, FalseLanesUndef>;
 
   def : SVE_4_Op_Pat<nxv8bf16, op, nxv8i1, nxv8bf16, nxv8bf16, nxv8bf16, !cast<Instruction>(NAME # _UNDEF)>;
+  def : SVE_4_Op_Pat<nxv4bf16, op, nxv4i1, nxv4bf16, nxv4bf16, nxv4bf16, !cast<Instruction>(NAME # _UNDEF)>;
+  def : SVE_4_Op_Pat<nxv2bf16, op, nxv2i1, nxv2bf16, nxv2bf16, nxv2bf16, !cast<Instruction>(NAME # _UNDEF)>;
 }
 
 // Predicated pseudo integer two operand instructions.