[POC][SVE] Allow code generation for fixed length vectorised loops [Patch 2/2].

[paulwalker-arm]

Summary

No expectations for review at this stage unless you are super keen.

This extends the proof of concept introduced by https://reviews.llvm.org/D71760 to show the work required to custom lower the main fixed length vector operations. The general scheme being:

<n x ty> op (<n x ty> op1, <n x ty> op2...

becomes

pg = create_predicate_for(<n x ty>)
new_op1 = convertToSVE(op1)
new_op2 = convertToSVE(op2)
...
return convertFromSVE(sve_op(pg, new_op1, new_op2...

To keep the patch small I've reused the existing intrinsic isel rules. This provides a route to reasonable performance whilst allowing us to start work on immediate packing, condition code handling and better utilisation of reversed instructions and movprfx. The ultimately goal is that the first patch ensures everything can run, whilst this and subsequent patches improve performance.

WARNING: The more operations we lower the more existing code paths will need to be cleaned so they no longer assume that (isTypeLegal(VectorType) implies VectorType.getSizeInBits() >= 128). This can be seen in this patch which required a charge to one of the AArch64 DAGCombines.

[efriedma]

Is there any particular reason to lower ISD::AND on a fixed vector to an intrinsic, as opposed to simply lowering it to an ISD::AND on a scalable vector? The patterns should mostly be there to make it work.

Otherwise, the general approach seems reasonable.

[efriedma]

Err, oops, ignore me; somehow I thought this was new.

[cameron.mcinally]

This seems like a good approach.

Do we have any tests for it? I notice a simple fixed-width argument test blows up (unless I'm making a mistake on my side):

; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve -aarch64-isel-sve-vector-bits=512 < %s | FileCheck %s

define <16 x float> @fadd_float_512b(<16 x float> %a, <16 x float> %b) {
; CHECK-LABEL: fadd_float_512b:
; CHECK:       // %bb.0:
; CHECK-NEXT:    fadd z0.s, z0.s, z1.s
; CHECK-NEXT:    ret
  %fadd = fadd <16 x float> %a, %b
  ret <16 x float> %fadd
}

I have a somewhat hacky solution to the calling convention issues, which I'll share soon. Hoping for a better solution though, since mine is clunky...

[paulwalker-arm]

The handling of larger than NEON fixed length vector function parameters is still up in the air. We cannot break existing NEON support so we're almost certainly going to need a new function attribute for fixed length SVE. For this reason and generally not knowing if the approach this patch is based on will be acceptable I did not prioritise writing tests, instead I focused on execution testing. I'll see about resolving the lack of tests albeit at this stage using pass by reference rather than value.

[cameron.mcinally]

Ah, ok. So you wouldn't see fixed-width vector arguments since the vectorizers works at the loop/block/etc level. That's what I misunderstood.

We have a pre-IR vectorizer and vectorize lib calls, so our needs are a little askew.

What about a scheme where we make the fixed-width vector types proper subregs of the scalable types? That way we can add the fixed-width types to the calling conventions and register classes. Would that be of interest you?

It would be similar to X86's XMM->YMM->ZMM registers (i.e. 128b->256b->512b), except that all of SVE's sub-registers would print the same register name. it's a little hacky, but might polish up ok.

I'll see about resolving the lack of tests albeit at this stage using pass by reference rather than value.

Thanks, but I'm fine without them for now. Up to you...

[paulwalker-arm]

What about a scheme where we make the fixed-width vector types proper subregs of the scalable types? That way we can add the fixed-width types to the calling conventions and register classes. Would that be of interest you?

It would be similar to X86's XMM->YMM->ZMM registers (i.e. 128b->256b->512b), except that all of SVE's sub-registers would print the same register name. it's a little hacky, but might polish

The problem here is what fixed-width vector types would you use? The mapping would be dependent on the target. We did experiment with HwMode but it didn't work out. At this stage I think the simplest option is for the vectorised functions to use scalable argument and then provide a way to insert/extract fixed width vectors into/from them. During code generation these insert/extract operations will become _SUBVECTOR operations that should generate no code. This is akin to the general philosophy I have used for fixed width code generation.

[cameron.mcinally]

In D71767#1954496, @paulwalker-arm wrote:

What about a scheme where we make the fixed-width vector types proper subregs of the scalable types? That way we can add the fixed-width types to the calling conventions and register classes. Would that be of interest you?

It would be similar to X86's XMM->YMM->ZMM registers (i.e. 128b->256b->512b), except that all of SVE's sub-registers would print the same register name. it's a little hacky, but might polish

The problem here is what fixed-width vector types would you use? The mapping would be dependent on the target.

I'm not sure if I'm following. If we don't add a vector type to a register class for a target, then it shouldn't be a problem, I think. We should be able to do that programmatically in ISel since we know the max vector width at compile time. E.g. Don't add the 1024b vector types to the ZPR register class for a 512b wide target (or we could create a ZPR and subreg hierarchy instead).

I've posted a back-of-the-envelope patch so we can look at something concrete, D77224. It definitely has some wart, but I think it could be workable going forward.

[paulwalker-arm]

Nothing to see here, just rebasing for those who want to experiment.

[paulwalker-arm]

@cameron.mcinally this is the patch I mentioned the other day, which contains the nodes where once I've written suitable tests I'll push separate patches for. Anything else is fair game. This patch implements VSELECT but that was just to investigate what we talked about during the previous sync call so I'll ignore it if you're planning to push on with your original work?

[cameron.mcinally]

In D71767#2206947, @paulwalker-arm wrote:

@cameron.mcinally this is the patch I mentioned the other day, which contains the nodes where once I've written suitable tests I'll push separate patches for.

Thanks, Paul. You mentioned that you would be focusing on another project for a few weeks. Would it help if I attempted to cherry-pick some of this Diff into individual patches (with new tests) for you? Or would I be stepping on your toes too much?

Anything else is fair game. This patch implements VSELECT but that was just to investigate what we talked about during the previous sync call so I'll ignore it if you're planning to push on with your original work?

I can carry on with VSELECT. I think I have enough intuition built around it now. I haven't yet looked into the state of the unpack[lo|hi] instructions, but will do so today.

[paulwalker-arm]

In D71767#2207158, @cameron.mcinally wrote:

In D71767#2206947, @paulwalker-arm wrote:

@cameron.mcinally this is the patch I mentioned the other day, which contains the nodes where once I've written suitable tests I'll push separate patches for.

Thanks, Paul. You mentioned that you would be focusing on another project for a few weeks. Would it help if I attempted to cherry-pick some of this Diff into individual patches (with new tests) for you? Or would I be stepping on your toes too much?

That would be great, thanks. I already have patches up for the extends and am currently focusing on setcc, sub and the shifts, which leaves min/max and divides. That said, one area I've not looked at yet are the VECREDUCE_ nodes. I don't anticipate them to be that problematic but having proof of this would be nice. Let me know what you decide.

[cameron.mcinally]

MIN/MAX and DIVs are up my alley. I'll try those. Will check out VREDUCE if I get through the others.

[paulwalker-arm]

rebase

[paulwalker-arm]

With the exception of VSELECT lowering, which is being worked under D85364, everything else is available in master.

clang/lib/Driver/ToolChains/Clang.cpp

if (IndirectBranches)
       CmdArgs.push_back("-mbranch-target-enforce");
   }
 
   // Handle -msve_vector_bits=<bits>
   if (Arg *A = Args.getLastArg(options::OPT_msve_vector_bits_EQ)) {
     StringRef Val = A->getValue();
     const Driver &D = getToolChain().getDriver();
     if (Val.equals("128") || Val.equals("256") || Val.equals("512") ||
-        Val.equals("1024") || Val.equals("2048"))
+        Val.equals("1024") || Val.equals("2048")) {
       CmdArgs.push_back(
           Args.MakeArgString(llvm::Twine("-msve-vector-bits=") + Val));
+
+      CmdArgs.push_back("-mllvm");
+      CmdArgs.push_back(
+            Args.MakeArgString("-aarch64-sve-vector-bits-min=" + Val));
+      // CmdArgs.push_back("-mllvm");
+      // CmdArgs.push_back(
+      //    Args.MakeArgString("-aarch64-sve-vector-bits-max=" + Val));
+    }
     // Silently drop requests for vector-length agnostic code as it's implied.
     else if (!Val.equals("scalable"))
       // Handle the unsupported values passed to msve-vector-bits.
       D.Diag(diag::err_drv_unsupported_option_argument)
           << A->getOption().getName() << Val;
   }
 }
 

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitINSERT_SUBVECTOR(SDNode *N) {
   // -> bitcast(insert_subvector(v, s, c2))
   if ((N0.isUndef() || N0.getOpcode() == ISD::BITCAST) &&
       N1.getOpcode() == ISD::BITCAST) {
     SDValue N0Src = peekThroughBitcasts(N0);
     SDValue N1Src = peekThroughBitcasts(N1);
     EVT N0SrcSVT = N0Src.getValueType().getScalarType();
     EVT N1SrcSVT = N1Src.getValueType().getScalarType();
     if ((N0.isUndef() || N0SrcSVT == N1SrcSVT) &&
-        N0Src.getValueType().isVector() && N1Src.getValueType().isVector()) {
+        N0Src.getValueType().isFixedLengthVector() &&
+        N1Src.getValueType().isFixedLengthVector()) {
       EVT NewVT;
       SDLoc DL(N);
       SDValue NewIdx;
       LLVMContext &Ctx = *DAG.getContext();
       unsigned NumElts = VT.getVectorNumElements();
       unsigned EltSizeInBits = VT.getScalarSizeInBits();
       if ((EltSizeInBits % N1SrcSVT.getSizeInBits()) == 0) {
         unsigned Scale = EltSizeInBits / N1SrcSVT.getSizeInBits();

llvm/lib/Target/AArch64/AArch64ISelLowering.h

private:
   SDValue LowerToScalableOp(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerEXTRACT_SUBVECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINSERT_SUBVECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDIV(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerMUL(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 LowerVSELECT(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;
   SDValue LowerVectorFP_TO_INT(SDValue Op, SelectionDAG &DAG) const;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

void AArch64TargetLowering::addTypeForFixedLengthSVE(MVT VT) {
   setOperationAction(ISD::FSUB, VT, Custom);
   setOperationAction(ISD::LOAD, VT, Custom);
   setOperationAction(ISD::MUL, VT, Custom);
   setOperationAction(ISD::OR, VT, Custom);
   setOperationAction(ISD::SETCC, VT, Custom);
   setOperationAction(ISD::SHL, VT, Custom);
   setOperationAction(ISD::SIGN_EXTEND, VT, Custom);
   setOperationAction(ISD::SIGN_EXTEND_INREG, VT, Custom);
+  setOperationAction(ISD::SMAX, VT, Custom);
+  setOperationAction(ISD::SMIN, VT, Custom);
   setOperationAction(ISD::SRA, VT, Custom);
   setOperationAction(ISD::SRL, VT, Custom);
   setOperationAction(ISD::STORE, VT, Custom);
   setOperationAction(ISD::SUB, VT, Custom);
   setOperationAction(ISD::TRUNCATE, VT, Custom);
+  setOperationAction(ISD::UMAX, VT, Custom);
+  setOperationAction(ISD::UMIN, VT, Custom);
+  setOperationAction(ISD::VSELECT, VT, Custom);
   setOperationAction(ISD::XOR, VT, Custom);
   setOperationAction(ISD::ZERO_EXTEND, VT, Custom);
+
+  if (VT.getVectorElementType() == MVT::i32 ||
+      VT.getVectorElementType() == MVT::i64) {
+    setOperationAction(ISD::SDIV, VT, Custom);
+    setOperationAction(ISD::UDIV, VT, Custom);
+  }
 }
 
 void AArch64TargetLowering::addDRTypeForNEON(MVT VT) {
   addRegisterClass(VT, &AArch64::FPR64RegClass);
   addTypeForNEON(VT, MVT::v2i32);
 }
 
 void AArch64TargetLowering::addQRTypeForNEON(MVT VT) {
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::AND:
     return LowerToScalableOp(Op, DAG);
   case ISD::SUB:
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::SUB_PRED);
   case ISD::FMAXNUM:
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FMAXNM_PRED);
   case ISD::FMINNUM:
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::FMINNM_PRED);
+  case ISD::VSELECT:
+    return LowerVSELECT(Op, DAG);
   }
 }
 
 bool AArch64TargetLowering::useSVEForFixedLengthVectors() const {
   // Prefer NEON unless larger SVE registers are available.
   return Subtarget->hasSVE() && Subtarget->getMinSVEVectorSizeInBits() >= 256;
 }
 
SDValue AArch64TargetLowering::LowerFixedLengthVectorSetccToSVE(
   EVT CmpVT = Pg.getValueType();
   SmallVector<SDValue, 4> CmpOps = {Pg, Op1, Op2, Op.getOperand(2)};
   auto Cmp = DAG.getNode(AArch64ISD::SETCC_MERGE_ZERO, DL, CmpVT, CmpOps);
 
   EVT PromoteVT = ContainerVT.changeTypeToInteger();
   auto Promote = DAG.getBoolExtOrTrunc(Cmp, DL, PromoteVT, InVT);
   return convertFromScalableVector(DAG, Op.getValueType(), Promote);
 }
+
+SDValue AArch64TargetLowering::LowerVSELECT(SDValue Op,
+                                            SelectionDAG &DAG) const {
+  SDLoc DL(Op);
+
+  EVT InVT = Op.getOperand(1).getValueType();
+  EVT ContainerVT = getContainerForFixedLengthVector(DAG, InVT);
+  auto Op1 = convertToScalableVector(DAG, ContainerVT, Op.getOperand(1));
+  auto Op2 = convertToScalableVector(DAG, ContainerVT, Op.getOperand(2));
+
+  // Convert the mask to a predicated (NOTE: We don't need to worry about
+  // inactive lanes since VSELECT is safe when given undefined elements).
+  EVT MaskVT = Op.getOperand(0).getValueType();
+  EVT MaskContainerVT = getContainerForFixedLengthVector(DAG, MaskVT);
+  auto Mask = convertToScalableVector(DAG, MaskContainerVT, Op.getOperand(0));
+  Mask = DAG.getNode(ISD::TRUNCATE, DL,
+                     MaskContainerVT.changeVectorElementType(MVT::i1), Mask);
+
+  auto VSel = DAG.getNode(ISD::VSELECT, DL, ContainerVT, Mask, Op1, Op2);
+  return convertFromScalableVector(DAG, InVT, VSel);
+}