[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.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

//        addRegisterClass(VT, &AArch64::ZPRRegClass);
   setOperationAction(ISD::TRUNCATE, VT, Legal);
   setOperationAction(ISD::ZERO_EXTEND, VT, Legal);
   setOperationAction(ISD::EXTRACT_SUBVECTOR, VT, Custom);
   setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
 
   // Lower fixed length vector operations to scalable equivalents.
   setOperationAction(ISD::LOAD, VT, Custom);
   setOperationAction(ISD::STORE, VT, Custom);
+
+  setOperationAction(ISD::ADD, VT, Custom);
+  setOperationAction(ISD::AND, VT, Custom);
+  setOperationAction(ISD::FADD, VT, Custom);
+  setOperationAction(ISD::FDIV, VT, Custom);
+  setOperationAction(ISD::FMAXNUM, VT, Custom);
+  setOperationAction(ISD::FMINNUM, VT, Custom);
+  setOperationAction(ISD::FMUL, VT, Custom);
+  setOperationAction(ISD::FSUB, VT, Custom);
+  setOperationAction(ISD::MUL, VT, Custom);
+  setOperationAction(ISD::OR, VT, Custom);
+  setOperationAction(ISD::SETCC, VT, Custom);
+  setOperationAction(ISD::SHL, VT, Custom);
+  setOperationAction(ISD::SMAX, VT, Custom);
+  setOperationAction(ISD::SMIN, VT, Custom);
+  setOperationAction(ISD::SRA, VT, Custom);
+  setOperationAction(ISD::SRL, VT, Custom);
+  setOperationAction(ISD::SUB, VT, Custom);
+  setOperationAction(ISD::UMAX, VT, Custom);
+  setOperationAction(ISD::UMIN, VT, Custom);
+  setOperationAction(ISD::XOR, VT, Custom);
 }
 
 void AArch64TargetLowering::addDRTypeForNEON(MVT VT) {
   addRegisterClass(VT, &AArch64::FPR64RegClass);
   addTypeForNEON(VT, MVT::v2i32);
 }
 
 void AArch64TargetLowering::addQRTypeForNEON(MVT VT) {
return DAG.getMaskedStore(Store->getChain(), DL,
                             Store->getOffset(),
                             CreatePredicateForFixedVector(DAG, DL, VT),
                             Store->getMemoryVT(),
                             Store->getMemOperand(),
                             Store->getAddressingMode(),
                             Store->isTruncatingStore());
 }
 
+static SDValue LowerFixedVector2OpToSVE(SelectionDAG &DAG, SDValue Op) {
+  SDLoc DL(Op);
+  EVT VT = Op.getValueType();
+  EVT ContainerVT = getContainerForFixedVector(DAG, VT);
+
+  int IntID;
+  switch (Op.getOpcode()) {
+  default: llvm_unreachable("unimplemented opcode");
+  case ISD::ADD:     IntID = Intrinsic::aarch64_sve_add; break;
+  case ISD::AND:     IntID = Intrinsic::aarch64_sve_and; break;
+  case ISD::FADD:    IntID = Intrinsic::aarch64_sve_fadd; break;
+  case ISD::FDIV:    IntID = Intrinsic::aarch64_sve_fdiv; break;
+  case ISD::FMAXNUM: IntID = Intrinsic::aarch64_sve_fmaxnm; break;
+  case ISD::FMINNUM: IntID = Intrinsic::aarch64_sve_fminnm; break;
+  case ISD::FMUL:    IntID = Intrinsic::aarch64_sve_fmul; break;
+  case ISD::FSUB:    IntID = Intrinsic::aarch64_sve_fsub; break;
+  case ISD::MUL:     IntID = Intrinsic::aarch64_sve_mul; break;
+  case ISD::OR:      IntID = Intrinsic::aarch64_sve_orr; break;
+  case ISD::SHL:     IntID = Intrinsic::aarch64_sve_lsl; break;
+  case ISD::SMAX:    IntID = Intrinsic::aarch64_sve_smax; break;
+  case ISD::SMIN:    IntID = Intrinsic::aarch64_sve_smin; break;
+  case ISD::SRA:     IntID = Intrinsic::aarch64_sve_asr; break;
+  case ISD::SRL:     IntID = Intrinsic::aarch64_sve_lsr; break;
+  case ISD::SUB:     IntID = Intrinsic::aarch64_sve_sub; break;
+  case ISD::UMAX:    IntID = Intrinsic::aarch64_sve_umax; break;
+  case ISD::UMIN:    IntID = Intrinsic::aarch64_sve_umin; break;
+  case ISD::XOR:     IntID = Intrinsic::aarch64_sve_eor; break;
+  }
+
+  auto Op1 = ConvertToScalableVector(DAG, ContainerVT, Op->getOperand(0));
+  auto Op2 = ConvertToScalableVector(DAG, ContainerVT, Op->getOperand(1));
+  auto Pg = CreatePredicateForFixedVector(DAG, DL, VT);
+
+  auto Res = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, ContainerVT,
+                         DAG.getConstant(IntID, DL, MVT::i64), Pg, Op1, Op2);
+
+  return ConvertFromScalableVector(DAG, VT, Res);
+}
+
+static SDValue LowerFixedVectorSetccToSVE(SelectionDAG &DAG, SDValue Op) {
+  SDLoc DL(Op);
+  EVT InVT = Op.getOperand(0).getValueType();
+  EVT ContainerVT = getContainerForFixedVector(DAG, InVT);
+
+  // HACK
+  if (!InVT.isInteger())
+    return SDValue();
+
+  auto Op1 = ConvertToScalableVector(DAG, ContainerVT, Op.getOperand(0));
+  auto Op2 = ConvertToScalableVector(DAG, ContainerVT, Op.getOperand(1));
+  auto Pg = CreatePredicateForFixedVector(DAG, DL, InVT);
+
+  int IntID;
+  switch (cast<CondCodeSDNode>(Op.getOperand(2))->get()) {
+  default:
+    llvm_unreachable("Unknown condition code!");
+  case ISD::SETNE:
+    IntID = Intrinsic::aarch64_sve_cmpne;
+    break;
+
+  case ISD::SETEQ:
+    IntID = Intrinsic::aarch64_sve_cmpeq;
+    break;
+
+  case ISD::SETLT:
+    std::swap(Op1, Op2);
+    LLVM_FALLTHROUGH;
+  case ISD::SETGT:
+    IntID = Intrinsic::aarch64_sve_cmpgt;
+    break;
+
+  case ISD::SETLE:
+    std::swap(Op1, Op2);
+    LLVM_FALLTHROUGH;
+  case ISD::SETGE:
+    IntID = Intrinsic::aarch64_sve_cmpge;
+    break;
+
+  case ISD::SETULT:
+    std::swap(Op1, Op2);
+    LLVM_FALLTHROUGH;
+  case ISD::SETUGT:
+    IntID = Intrinsic::aarch64_sve_cmphi;
+    break;
+
+  case ISD::SETULE:
+    std::swap(Op1, Op2);
+    LLVM_FALLTHROUGH;
+  case ISD::SETUGE:
+    IntID = Intrinsic::aarch64_sve_cmphs;
+    break;
+  }
+
+  auto Cmp = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, Pg.getValueType(),
+                         DAG.getConstant(IntID, DL, MVT::i64), Pg, Op1, Op2);
+
+  auto Promote = DAG.getBoolExtOrTrunc(Cmp, DL, ContainerVT, InVT);
+  auto Extract = ConvertFromScalableVector(DAG, InVT, Promote);
+  return DAG.getNode(ISD::TRUNCATE, DL, Op.getValueType(), Extract);
+}
+
 SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
                                               SelectionDAG &DAG) const {
   if (!useSVEForFixedLengthVectors())
     return LowerOperationDefault(Op, DAG);
 
   SDLoc DL(Op);
 
   // If we get here then we have a vector operation whose operands are legal
case ISD::STORE: {
       EVT InVT = cast<StoreSDNode>(Op)->getValue().getValueType();
       if (!InVT.isFixedVector())
         return LowerOperationDefault(Op, DAG);
 
       if (!UseSVEForAll && InVT.getSizeInBits() <= 128)
         return Op; // Use Neon.
       return LowerFixedVectorStoreToSVE(DAG, Op);
     }
+
+    case ISD::ADD:
+    case ISD::AND:
+    case ISD::FADD:
+    case ISD::FDIV:
+    case ISD::FMAXNUM:
+    case ISD::FMINNUM:
+    case ISD::FMUL:
+    case ISD::FSUB:
+    case ISD::MUL:
+    case ISD::OR:
+    case ISD::SHL:
+    case ISD::SMAX:
+    case ISD::SMIN:
+    case ISD::SRA:
+    case ISD::SRL:
+    case ISD::SUB:
+    case ISD::UMAX:
+    case ISD::UMIN:
+    case ISD::XOR: {
+      EVT VT = Op.getValueType();
+      if (!VT.isFixedVector())
+        return LowerOperationDefault(Op, DAG);
+
+      if (!UseSVEForAll && VT.getSizeInBits() <= 128)
+        return Op; // Use Neon.
+      return LowerFixedVector2OpToSVE(DAG, Op);
+    }
+
+    case ISD::SETCC: {
+      if (!Op.getValueType().isFixedVector())
+        return LowerOperationDefault(Op, DAG);
+
+      return LowerFixedVectorSetccToSVE(DAG, Op);
+    }
   }
 }
 
 // Should we expand the build vector with shuffles?
 bool AArch64TargetLowering::shouldExpandBuildVectorWithShuffles(EVT VT,
       unsigned DefinedValues) const {
   if (useSVEForFixedLengthVectors())
     return false;

 static SDValue performANDCombine(SDNode *N,
                                  TargetLowering::DAGCombinerInfo &DCI) {
   SelectionDAG &DAG = DCI.DAG;
   SDValue LHS = N->getOperand(0);
   EVT VT = N->getValueType(0);
   if (!VT.isVector() || !DAG.getTargetLoweringInfo().isTypeLegal(VT))
     return SDValue();
 
+  if (VT.isFixedVector() && VT.getSizeInBits() > 128)
+    return SDValue();
+
   if (VT.isScalableVector())
     return performSVEAndCombine(N, DCI);
 
   BuildVectorSDNode *BVN =
       dyn_cast<BuildVectorSDNode>(N->getOperand(1).getNode());
   if (!BVN)
     return SDValue();
 

llvm/lib/Target/AArch64/SVEInstrFormats.td

multiclass sve_int_dup_imm_pred_zero<string asm> {
   def : InstAlias<"mov $Zd, $Pg/z, $imm",
                   (!cast<Instruction>(NAME # _B) ZPR8:$Zd,  PPRAny:$Pg, cpy_imm8_opt_lsl_i8:$imm), 1>;
   def : InstAlias<"mov $Zd, $Pg/z, $imm",
                   (!cast<Instruction>(NAME # _H) ZPR16:$Zd, PPRAny:$Pg, cpy_imm8_opt_lsl_i16:$imm), 1>;
   def : InstAlias<"mov $Zd, $Pg/z, $imm",
                   (!cast<Instruction>(NAME # _S) ZPR32:$Zd, PPRAny:$Pg, cpy_imm8_opt_lsl_i32:$imm), 1>;
   def : InstAlias<"mov $Zd, $Pg/z, $imm",
                   (!cast<Instruction>(NAME # _D) ZPR64:$Zd, PPRAny:$Pg, cpy_imm8_opt_lsl_i64:$imm), 1>;
+
+  def : Pat<(nxv16i8 (anyext (nxv16i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _B) PPR:$Op, 1, 0)>;
+  def : Pat<(nxv8i16 (anyext (nxv8i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _H) PPR:$Op, 1, 0)>;
+  def : Pat<(nxv4i32 (anyext (nxv4i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _S) PPR:$Op, 1, 0)>;
+  def : Pat<(nxv2i64 (anyext (nxv2i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _D) PPR:$Op, 1, 0)>;
+
+  def : Pat<(nxv16i8 (sext (nxv16i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _B) PPR:$Op, -1, 0)>;
+  def : Pat<(nxv8i16 (sext (nxv8i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _H) PPR:$Op, -1, 0)>;
+  def : Pat<(nxv4i32 (sext (nxv4i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _S) PPR:$Op, -1, 0)>;
+  def : Pat<(nxv2i64 (sext (nxv2i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _D) PPR:$Op, -1, 0)>;
+
+  def : Pat<(nxv16i8 (zext (nxv16i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _B) PPR:$Op, 1, 0)>;
+  def : Pat<(nxv8i16 (zext (nxv8i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _H) PPR:$Op, 1, 0)>;
+  def : Pat<(nxv4i32 (zext (nxv4i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _S) PPR:$Op, 1, 0)>;
+  def : Pat<(nxv2i64 (zext (nxv2i1 PPR:$Op))),
+            (!cast<Instruction>(NAME # _D) PPR:$Op, 1, 0)>;
 }
 
 //===----------------------------------------------------------------------===//
 // SVE Integer Compare - Vectors Group
 //===----------------------------------------------------------------------===//
 
 class sve_int_cmp<bit cmp_1, bits<2> sz8_64, bits<3> opc, string asm,
                   PPRRegOp pprty, ZPRRegOp zprty1, ZPRRegOp zprty2>