[SVE][CodeGen] Lower SDIV & UDIV to SVE intrinsics

[kmclaughlin]

Summary

This patch maps IR operations for sdiv & udiv to the
@llvm.aarch64.sve.[s|u]div intrinsics.

A ptrue must be created during lowering as the div instructions
have only a predicated form.

Patch contains changes by Andrzej Warzynski.

[efriedma]

I'd prefer to handle legalization in a separate patch from handling legal sdiv/udiv operations, so we actually have some context to discuss the legalization strategy.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7670

If we're going to support these operations, we might as well just add isel patterns; that's what we've been doing for other arithmetic operations.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7670

Just to provide a bit of context to this approach:

For unpredicated ISD nodes for which there is no predicated instruction, the predicate needs to be generated. For scalable vectors this will be a ptrue all, but for fixed-width vectors may take some other predicate such as VL8 for fixed 8 elements.

Rather than creating new predicated AArch64 ISD nodes for each operation such as AArch64ISD::UDIV_PRED, the idea is to reuse the intrinsic layer we already added to support the ACLE - which are predicated and for which we already have the patterns - and map directly onto those.

By doing the expansion in ISelLowering, the patterns stay simple and we can generalise getPtrue method so that it generates the right predicate for any scalable/fixed vector size as done in D71760 avoiding the need to write multiple patterns for different vector lengths.

This patch was meant as the proof of concept of that idea (as discussed in the sync-up call of Apr 2).

llvm/test/CodeGen/AArch64/llvm-ir-to-intrinsic.ll

28

This test should use CHECK-DAG instead of CHECK-NEXT, as the sdiv instructions are independent. (same for some of the other tests)

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7670

Using INTRINSIC_WO_CHAIN is a little annoying; it's hard to read in DAG dumps, and it gives weird error messages if we fail in selection. But there aren't really any other immediate downsides I can think of, vs. doing it the other way (converting the intrinsic to AArch64ISD::UDIV_PRED).

Long-term, we're going to have a target-independent ISD::UDIV_PRED. We probably want to start using those nodes at some point, to get target-independent optimizations. Not sure if that impacts what we want to do right now.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7670

I agree that using INTRINSIC_WO_CHAIN will be a bit annoying for more complicated patterns. The reuse of the intrinsics was merely for convenience because we already have the patterns, and was not a critical part of the design. It shouldn't be a big effort to create AArch64ISD nodes and use these for the intrinsics as well.

If we use AArch64-specific nodes we can implement what's needed now for SVE codegen and I expect we can easily migrate to target-independent nodes when they get added.

[kmclaughlin]

• Removed changes to handle legalisation from this patch (this will be included in a follow up)
• Added AArch64ISD nodes for SDIV_PRED & UDIV_PRED
• Changed LowerDIV to use the new ISD nodes rather than lowering to SVE intrinsics
• Update tests to use CHECK-DAG

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

11354

Whitespace.

[kmclaughlin]

Thank you both for your comments on this patch, @efriedma & @sdesmalen!

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

SDValue DAGTypeLegalizer::SplitVecOp_VSETCC(SDNode *N) {
   assert(N->getValueType(0).isVector() &&
          N->getOperand(0).getValueType().isVector() &&
          "Operand types must be vectors");
   // The result has a legal vector type, but the input needs splitting.
   SDValue Lo0, Hi0, Lo1, Hi1, LoRes, HiRes;
   SDLoc DL(N);
   GetSplitVector(N->getOperand(0), Lo0, Hi0);
   GetSplitVector(N->getOperand(1), Lo1, Hi1);
-  unsigned PartElements = Lo0.getValueType().getVectorNumElements();
-  EVT PartResVT = EVT::getVectorVT(*DAG.getContext(), MVT::i1, PartElements);
-  EVT WideResVT = EVT::getVectorVT(*DAG.getContext(), MVT::i1, 2*PartElements);
+  auto PartEltCnt = Lo0.getValueType().getVectorElementCount();
+  EVT PartResVT = EVT::getVectorVT(*DAG.getContext(), MVT::i1, PartEltCnt);
+  EVT WideResVT = EVT::getVectorVT(*DAG.getContext(), MVT::i1, PartEltCnt*2);
 
   LoRes = DAG.getNode(ISD::SETCC, DL, PartResVT, Lo0, Lo1, N->getOperand(2));
   HiRes = DAG.getNode(ISD::SETCC, DL, PartResVT, Hi0, Hi1, N->getOperand(2));
   SDValue Con = DAG.getNode(ISD::CONCAT_VECTORS, DL, WideResVT, LoRes, HiRes);
 
   EVT OpVT = N->getOperand(0).getValueType();
   ISD::NodeType ExtendCode =
       TargetLowering::getExtendForContent(TLI.getBooleanContents(OpVT));

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

static void getCopyToPartsVector(SelectionDAG &DAG, const SDLoc &DL,
   assert(RegisterVT == PartVT && "Part type doesn't match vector breakdown!");
 
   unsigned IntermediateNumElts = IntermediateVT.isVector() ?
     IntermediateVT.getVectorNumElements() : 1;
 
   // Convert the vector to the appropriate type if necessary.
   unsigned DestVectorNoElts = NumIntermediates * IntermediateNumElts;
 
   EVT BuiltVectorTy = EVT::getVectorVT(
-      *DAG.getContext(), IntermediateVT.getScalarType(), DestVectorNoElts);
+      *DAG.getContext(), IntermediateVT.getScalarType(), DestVectorNoElts,
+      ValueVT.isScalableVector());
   if (ValueVT != BuiltVectorTy) {
     if (SDValue Widened = widenVectorToPartType(DAG, Val, DL, BuiltVectorTy))
       Val = Widened;
 
     Val = DAG.getNode(ISD::BITCAST, DL, BuiltVectorTy, Val);
   }
 
   // Split the vector into intermediate operands.

llvm/lib/CodeGen/TargetLoweringBase.cpp

 /// This method returns the number of registers needed, and the VT for each
 /// register.  It also returns the VT and quantity of the intermediate values
 /// before they are promoted/expanded.
 unsigned TargetLoweringBase::getVectorTypeBreakdown(LLVMContext &Context, EVT VT,
                                                 EVT &IntermediateVT,
                                                 unsigned &NumIntermediates,
                                                 MVT &RegisterVT) const {
   unsigned NumElts = VT.getVectorNumElements();
+  bool IsScalable = VT.isScalableVector();
 
   // If there is a wider vector type with the same element type as this one,
   // or a promoted vector type that has the same number of elements which
   // are wider, then we should convert to that legal vector type.
   // This handles things like <2 x float> -> <4 x float> and
   // <4 x i1> -> <4 x i32>.
   LegalizeTypeAction TA = getTypeAction(Context, VT);
   if (NumElts != 1 && (TA == TypeWidenVector || TA == TypePromoteInteger)) {
unsigned TargetLoweringBase::getVectorTypeBreakdown(LLVMContext &Context, EVT VT,
   // could break down into LHS/RHS like LegalizeDAG does.
   if (!isPowerOf2_32(NumElts)) {
     NumVectorRegs = NumElts;
     NumElts = 1;
   }
 
   // Divide the input until we get to a supported size.  This will always
   // end with a scalar if the target doesn't support vectors.
-  while (NumElts > 1 && !isTypeLegal(
-                                   EVT::getVectorVT(Context, EltTy, NumElts))) {
+  while (NumElts > 1 &&
+         !isTypeLegal(EVT::getVectorVT(Context, EltTy, NumElts, IsScalable))) {
     NumElts >>= 1;
     NumVectorRegs <<= 1;
   }
 
   NumIntermediates = NumVectorRegs;
 
-  EVT NewVT = EVT::getVectorVT(Context, EltTy, NumElts);
+  EVT NewVT = EVT::getVectorVT(Context, EltTy, NumElts, IsScalable);
   if (!isTypeLegal(NewVT))
     NewVT = EltTy;
   IntermediateVT = NewVT;
 
   MVT DestVT = getRegisterType(Context, NewVT);
   RegisterVT = DestVT;
   unsigned NewVTSize = NewVT.getSizeInBits();
 

llvm/lib/Target/AArch64/AArch64ISelLowering.h

private:
   SDValue LowerFLT_ROUNDS_(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINSERT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerEXTRACT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSCALAR_TO_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerBUILD_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVECTOR_SHUFFLE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSPLAT_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDUPQLane(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerDIV(SDValue Op, SelectionDAG &DAG,
+                   unsigned IntrID) const;
   SDValue LowerEXTRACT_SUBVECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVectorSRA_SRL_SHL(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerShiftLeftParts(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerShiftRightParts(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVSETCC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerCTPOP(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerF128Call(SDValue Op, SelectionDAG &DAG,
                         RTLIB::Libcall Call) const;
   SDValue LowerFCOPYSIGN(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFP_EXTEND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFP_ROUND(SDValue Op, SelectionDAG &DAG) const;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasNEON()) {
     setTruncStoreAction(MVT::v4i16, MVT::v4i8, Custom);
   }
 
   if (Subtarget->hasSVE()) {
     // FIXME: Add custom lowering of MLOAD to handle different passthrus (not a
     // splat of 0 or undef) once vector selects supported in SVE codegen. See
     // D68877 for more details.
     for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
-      if (isTypeLegal(VT))
+      if (isTypeLegal(VT)) {
         setOperationAction(ISD::SPLAT_VECTOR, VT, Custom);
+        setOperationAction(ISD::SDIV, VT, Custom);
+        setOperationAction(ISD::UDIV, VT, Custom);
+      }
     }
     setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::i8, Custom);
     setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::i16, Custom);
 
     for (MVT VT : MVT::fp_scalable_vector_valuetypes()) {
       if (isTypeLegal(VT)) {
         setOperationAction(ISD::SPLAT_VECTOR, VT, Custom);
       }
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::BUILD_VECTOR:
     return LowerBUILD_VECTOR(Op, DAG);
   case ISD::VECTOR_SHUFFLE:
     return LowerVECTOR_SHUFFLE(Op, DAG);
   case ISD::SPLAT_VECTOR:
     return LowerSPLAT_VECTOR(Op, DAG);
   case ISD::EXTRACT_SUBVECTOR:
     return LowerEXTRACT_SUBVECTOR(Op, DAG);
+  case ISD::SDIV:
+    return LowerDIV(Op, DAG, Intrinsic::aarch64_sve_sdiv);
+  case ISD::UDIV:
+    return LowerDIV(Op, DAG, Intrinsic::aarch64_sve_udiv);
   case ISD::SRA:
   case ISD::SRL:
   case ISD::SHL:
     return LowerVectorSRA_SRL_SHL(Op, DAG);
   case ISD::SHL_PARTS:
     return LowerShiftLeftParts(Op, DAG);
   case ISD::SRL_PARTS:
   case ISD::SRA_PARTS:
SDValue AArch64TargetLowering::LowerDUPQLane(SDValue Op,
   SDValue SplatIdx64 = DAG.getNode(ISD::SPLAT_VECTOR, DL, MVT::nxv2i64, Idx64);
   SDValue ShuffleMask = DAG.getNode(ISD::ADD, DL, MVT::nxv2i64, SV, SplatIdx64);
 
   // create the vector Val[idx64],Val[idx64+1],Val[idx64],Val[idx64+1],...
   SDValue TBL = DAG.getNode(AArch64ISD::TBL, DL, MVT::nxv2i64, V, ShuffleMask);
   return DAG.getNode(ISD::BITCAST, DL, VT, TBL);
 }
 
+SDValue AArch64TargetLowering::LowerDIV(SDValue Op,
+                                        SelectionDAG &DAG,
+                                        unsigned IntrID) const {
+  EVT VT = Op.getValueType();
+  SDLoc DL(Op);
+
+  assert(Op.getOperand(0).getValueType().isScalableVector() &&
+         Op.getOperand(1).getValueType().isScalableVector() &&
+         "Only scalable vectors are supported");
+
+  auto PredTy = VT.getVectorVT(*DAG.getContext(), MVT::i1,
+                               VT.getVectorNumElements(), true);
+  SDValue Mask = getPTrue(DAG, DL, PredTy, AArch64SVEPredPattern::all);
+
+  return DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, VT,
+                     DAG.getConstant(IntrID, DL, MVT::i64),
+                     Mask, Op.getOperand(0), Op.getOperand(1));
+}
+
 static bool resolveBuildVector(BuildVectorSDNode *BVN, APInt &CnstBits,
                                APInt &UndefBits) {
   EVT VT = BVN->getValueType(0);
   APInt SplatBits, SplatUndef;
   unsigned SplatBitSize;
   bool HasAnyUndefs;
   if (BVN->isConstantSplat(SplatBits, SplatUndef, SplatBitSize, HasAnyUndefs)) {
     unsigned NumSplats = VT.getSizeInBits() / SplatBitSize;
 
case Intrinsic::aarch64_sve_index:
     return LowerSVEIntrinsicIndex(N, DAG);
   case Intrinsic::aarch64_sve_dup:
     return LowerSVEIntrinsicDUP(N, DAG);
   case Intrinsic::aarch64_sve_dup_x:
     return DAG.getNode(ISD::SPLAT_VECTOR, SDLoc(N), N->getValueType(0),
                        N->getOperand(1));
   case Intrinsic::aarch64_sve_ext:
     return LowerSVEIntrinsicEXT(N, DAG);
   case Intrinsic::aarch64_sve_sel:
     return DAG.getNode(ISD::VSELECT, SDLoc(N), N->getValueType(0),
                        N->getOperand(1), N->getOperand(2), N->getOperand(3));
   case Intrinsic::aarch64_sve_cmpeq_wide:
     return tryConvertSVEWideCompare(N, Intrinsic::aarch64_sve_cmpeq,
                                     false, DCI, DAG);
   case Intrinsic::aarch64_sve_cmpne_wide:
     return tryConvertSVEWideCompare(N, Intrinsic::aarch64_sve_cmpne,
                                     false, DCI, DAG);

llvm/test/CodeGen/AArch64/llvm-ir-to-intrinsic.ll

+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+;
+; SDIV
+;
+
+define <vscale x 4 x i32> @sdiv_i32(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
+; CHECK-LABEL: @sdiv_i32
+; CHECK: ptrue p0.s
+; CHECK-NEXT: sdiv z0.s, p0/m, z0.s, z1.s
+; CHECK-NEXT: ret
+  %div = sdiv <vscale x 4 x i32> %a, %b
+  ret <vscale x 4 x i32> %div
+}
+
+define <vscale x 2 x i64> @sdiv_i64(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b) {
+; CHECK-LABEL: @sdiv_i64
+; CHECK: ptrue p0.d
+; CHECK-NEXT: sdiv z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT: ret
+  %div = sdiv <vscale x 2 x i64> %a, %b
+  ret <vscale x 2 x i64> %div
+}
+
+define <vscale x 8 x i32> @sdiv_narrow(<vscale x 8 x i32> %a, <vscale x 8 x i32> %b) {
+; CHECK-LABEL: @sdiv_narrow
+; CHECK: ptrue p0.s
+; CHECK-NEXT: sdiv z0.s, p0/m, z0.s, z2.s
+; CHECK-NEXT: sdiv z1.s, p0/m, z1.s, z3.s
+; CHECK-NEXT: ret
+  %div = sdiv <vscale x 8 x i32> %a, %b
+  ret <vscale x 8 x i32> %div
+}
+
+define <vscale x 2 x i32> @sdiv_widen(<vscale x 2 x i32> %a, <vscale x 2 x i32> %b) {
+; CHECK-LABEL: @sdiv_widen
+; CHECK: ptrue p0.d
+; CHECK-NEXT: sxtw z1.d, p0/m, z1.d
+; CHECK-NEXT: sxtw z0.d, p0/m, z0.d
+; CHECK-NEXT: sdiv z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT: ret
+  %div = sdiv <vscale x 2 x i32> %a, %b
+  ret <vscale x 2 x i32> %div
+}
+
+;
+; UDIV
+;
+
+define <vscale x 4 x i32> @udiv_i32(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
+; CHECK-LABEL: @udiv_i32
+; CHECK: ptrue p0.s
+; CHECK-NEXT: udiv z0.s, p0/m, z0.s, z1.s
+; CHECK-NEXT: ret
+  %div = udiv <vscale x 4 x i32> %a, %b
+  ret <vscale x 4 x i32> %div
+}
+
+define <vscale x 2 x i64> @udiv_i64(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b) {
+; CHECK-LABEL: @udiv_i64
+; CHECK: ptrue p0.d
+; CHECK-NEXT: udiv z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT: ret
+  %div = udiv <vscale x 2 x i64> %a, %b
+  ret <vscale x 2 x i64> %div
+}
+
+define <vscale x 8 x i32> @udiv_narrow(<vscale x 8 x i32> %a, <vscale x 8 x i32> %b) {
+; CHECK-LABEL: @udiv_narrow
+; CHECK: ptrue p0.s
+; CHECK-NEXT: udiv z0.s, p0/m, z0.s, z2.s
+; CHECK-NEXT: udiv z1.s, p0/m, z1.s, z3.s
+; CHECK-NEXT: ret
+  %div = udiv <vscale x 8 x i32> %a, %b
+  ret <vscale x 8 x i32> %div
+}
+
+define <vscale x 2 x i32> @udiv_widen(<vscale x 2 x i32> %a, <vscale x 2 x i32> %b) {
+; CHECK-LABEL: @udiv_widen
+; CHECK: ptrue p0.d
+; CHECK-NEXT: and z1.d, z1.d, #0xffffffff
+; CHECK-NEXT: and z0.d, z0.d, #0xffffffff
+; CHECK-NEXT: udiv z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT: ret
+  %div = udiv <vscale x 2 x i32> %a, %b
+  ret <vscale x 2 x i32> %div
+}