[SVE] Custom ISel for fixed length extract/insert_subvector.

[paulwalker-arm]

Summary

We use extact_subvector and insert_subvector to "cast" between
fixed length and scalable vectors. This patch adds custom c++
based ISel for the following cases:

fixed_vector = ISD::EXTRACT_SUBVECTOR scalable_vector, 0
scalable_vector = ISD::INSERT_SUBVECTOR undef(scalable_vector), fixed_vector, 0

Which result in either EXTRACT_SUBREG/INSERT_SUBREG for NEON sized
vectors or COPY_TO_REGCLASS otherwise.

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3244

I'm not sure I understand the issue here. Is the problem just that for a pattern, you need to write the type of the result? I don't think there's any problem with writing a pattern involving a type that isn't always legal; if the type isn't legal, it just won't match.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

8678

The comment here is really terse, and I'm not sure this is handling all the cases we need to.

In particular, I'm not sure this correctly handles types like nxv2f32.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3244

I believe the issue is the runtime nature of the >128bit fixed length vectors means they are not mapped to any register class, which prevents pattern based matching.

We did investigate using hwmodes but it didn't prove to be a viable solution.

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3244

I believe the issue is the runtime nature of the >128bit fixed length vectors means they are not mapped to any register class

You could change that in AArch64RegisterInfo.td, if you wanted to. I don't think that would cause any issues; it's okay if some of the types in the list aren't legal for all subtargets.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3244

Thanks for the tip, I'll take look and see how it works out.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3244

I tried adding the extra fixed length vector MVTs to ZPRClass but it didn't really work out.

Firstly I needed to fix up a bunch of patterns due to "Could not infer all types in pattern" build errors. This might be down to ZPRClass being sized as 128, which looks especially weird once it starts to contain a lot of MVTs that are known bigger than 128. Ultimately though the failing patterns are easily fixed but I guess potentially burdensome since 99.9% of the patterns shouldn't need to care about fixed length SVE.

After this the extract_subvector patterns still do not match. I've tracked this down to SDTSubVecExtract's usage of SDTCisSubVecOfVec which is not strictly true given the newly expanded definition of extract_subvector. I cannot just update SDTCisSubVecOfVec because it's used by other operations where we don't want to allow a mixture of fixed length and scalable vectors (e.g. concat_vector).

I can add a new variant of SDTCisSubVecOfVec for use by extract_subvector and insert_subvector, but given the minimal usage I'm not sure it's worth it. What do you think?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

8678

I've create a helper function to detect unpacked vector types and used it here so that only fixed length vector extracts from packed scalable types are considered legal.

[paulwalker-arm]

Updated selection comments to fully explain the rational. Updated lowering code to be more caution in what is considered legal.

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3244

I guess it's fine to leave as-is, in that case.

3385

There isn't any corresponding code for INSERT_SUBVECTOR in AArch64ISelLowering?

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3385

Unlike EXTRACT_SUBVECTOR, AArch64ISelLowering didn't have any custom lowering for INSERT_SUBVECTOR for any vector types. For this reason I assumed the defaults were fine. My new usage when lowering fixed length to SVE only uses the legal variants so doesn't introduce any new requirements.

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3366

Do we need the isPackedVectorType check here as well?

3385

In theory, you need exactly the same code.

Realistically, I'm not sure target-independent code will actually ever create an insert_subvector with a non-zero index. I briefly tried grepping through the relevant code, and couldn't find anything. I guess we could ignore the issue for now.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3366

The code in LowerEXTRACT_SUBVECTOR should preclude that case, but I can add some asserts.

[paulwalker-arm]

Added an implementation for LowerINSERT_SUBVECTOR along with more asserts.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3244

Thanks. I'll circle back round to this after we've enough functionality for fixed length to be usable.

3385

To be consistent I've provided an implementation. We're not really in a position to properly test it, although it defaults to expand for all the cases where tests don't exist.

[cameron.mcinally]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3244

No objections, but I do like Eli's idea for the long term. This would be a step in the right direction for passing fixed-width arguments. Having a restricted IR for fixed-width isn't ideal.

llvm/test/CodeGen/AArch64/sve-fixed-length-subvector.ll

26

Nit: could probably mark the loads/stores volatile to avoid the branch.

[paulwalker-arm]

llvm/test/CodeGen/AArch64/sve-fixed-length-subvector.ll

26

I'm not sure how this helps. The reason for the branch is to force a block boundary to ensure the extract_subvector resulting from lowering the load is not combined with the insert_subvector that's created when lowering the store.

[cameron.mcinally]

llvm/test/CodeGen/AArch64/sve-fixed-length-subvector.ll

26

Ok, that makes sense. Let me ask the opposite though -- if the load and store are volatile, will the fixed-width lowering honor the volatile?

x = load volatile *p
y = fneg x
x = load volatile *p
store *p, x

Unlikely a problem considering the loads made it to the backend, but would be good to confirm.

[paulwalker-arm]

llvm/test/CodeGen/AArch64/sve-fixed-length-subvector.ll

26

The resulting masked memory operation takes the same MachineMemOperand as the original fixed length operation, so the volatile flag will be maintained.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

SDNode *N2 = CurDAG->getMachineNode(AArch64::ADDXrr, DL, MVT::i64,
                                       {SDValue(N1, 0), N->getOperand(2)});
   SDNode *N3 = CurDAG->getMachineNode(
       AArch64::ADDG, DL, MVT::i64,
       {SDValue(N2, 0), CurDAG->getTargetConstant(0, DL, MVT::i64),
        CurDAG->getTargetConstant(TagOffset, DL, MVT::i64)});
   ReplaceNode(N, N3);
 }
 
+// NOTE: When targeting fixed length vectors at SVE the range of MVTs is runtime
+// variable, hence this manual selection.
+static SDNode *extractSubReg(SelectionDAG *DAG, EVT VT, SDValue V) {
+  assert(VT.isFixedLengthVector() && "Expected fixed length vector type!");
+  SDLoc DL(V);
+  switch (VT.getSizeInBits()) {
+  case 64: {
+    auto SubReg = DAG->getTargetConstant(AArch64::dsub, DL, MVT::i32);
+    return DAG->getMachineNode(TargetOpcode::EXTRACT_SUBREG, DL, VT, V, SubReg);
+  }
+  case 128: {
+    auto SubReg = DAG->getTargetConstant(AArch64::zsub, DL, MVT::i32);
+    return DAG->getMachineNode(TargetOpcode::EXTRACT_SUBREG, DL, VT, V, SubReg);
+  }
+  default: {
+    auto RC = DAG->getTargetConstant(AArch64::ZPRRegClassID, DL, MVT::i64);
+    return DAG->getMachineNode(TargetOpcode::COPY_TO_REGCLASS, DL, VT, V, RC);
+  }
+  }
+}
+
+// NOTE: When targeting fixed length vectors at SVE the range of MVTs is runtime
+// variable, hence this manual selection.
+static SDNode *insertSubReg(SelectionDAG *DAG, EVT VT, SDValue V) {
+  assert(V.getValueType().isFixedLengthVector() &&
+         "Expected fixed length vector type!");
+  SDLoc DL(V);
+  switch (V.getValueType().getSizeInBits()) {
+  case 64: {
+    auto SubReg = DAG->getTargetConstant(AArch64::dsub, DL, MVT::i32);
+    auto Container = DAG->getMachineNode(TargetOpcode::IMPLICIT_DEF, DL, VT);
+    return DAG->getMachineNode(TargetOpcode::INSERT_SUBREG, DL, VT,
+                               SDValue(Container, 0), V, SubReg);
+  }
+  case 128: {
+    auto SubReg = DAG->getTargetConstant(AArch64::zsub, DL, MVT::i32);
+    auto Container = DAG->getMachineNode(TargetOpcode::IMPLICIT_DEF, DL, VT);
+    return DAG->getMachineNode(TargetOpcode::INSERT_SUBREG, DL, VT,
+                               SDValue(Container, 0), V, SubReg);
+  }
+  default: {
+    auto RC = DAG->getTargetConstant(AArch64::ZPRRegClassID, DL, MVT::i64);
+    return DAG->getMachineNode(TargetOpcode::COPY_TO_REGCLASS, DL, VT, V, RC);
+  }
+  }
+}
+
 void AArch64DAGToDAGISel::Select(SDNode *Node) {
   // If we have a custom node, we already have selected!
   if (Node->isMachineOpcode()) {
     LLVM_DEBUG(errs() << "== "; Node->dump(CurDAG); errs() << "\n");
     Node->setNodeId(-1);
     return;
   }
 
if (tryHighFPExt(Node))
       return;
     break;
 
   case ISD::OR:
     if (tryBitfieldInsertOp(Node))
       return;
     break;
 
+  case ISD::EXTRACT_SUBVECTOR: {
+    // Bail when not a "cast" like extract_subvector.
+    if (cast<ConstantSDNode>(Node->getOperand(1))->getZExtValue() != 0)
+      break;
+
+    // Bail when normal isel can do the job.
+    EVT InVT = Node->getOperand(0).getValueType();
+    if (VT.isScalableVector() || InVT.isFixedLengthVector())
+      break;
+
+    // NOTE: We can only get here when doing fixed length SVE code generation.
+    // We do manual selection because the types involved are not linked to real
+    // registers (despite being legal) and must be coerced into SVE registers.
+
+    ReplaceNode(Node, extractSubReg(CurDAG, VT, Node->getOperand(0)));
+    return;
+  }
+
+  case ISD::INSERT_SUBVECTOR: {
+    // Bail when not a "cast" like insert_subvector.
+    if (cast<ConstantSDNode>(Node->getOperand(2))->getZExtValue() != 0)
+      break;
+    if (!Node->getOperand(0).isUndef())
+      break;
+
+    // Bail when normal isel should do the job.
+    EVT InVT = Node->getOperand(1).getValueType();
+    if (VT.isFixedLengthVector() || InVT.isScalableVector())
+      break;
+
+    // NOTE: We can only get here when doing fixed length SVE code generation.
+    // We do manual selection because the types involved are not linked to real
+    // registers (despite being legal) and must be coerced into SVE registers.
+
+    ReplaceNode(Node, insertSubReg(CurDAG, VT, Node->getOperand(1)));
+    return;
+  }
+
   case ISD::Constant: {
     // Materialize zero constants as copies from WZR/XZR.  This allows
     // the coalescer to propagate these into other instructions.
     ConstantSDNode *ConstNode = cast<ConstantSDNode>(Node);
     if (ConstNode->isNullValue()) {
       if (VT == MVT::i32) {
         SDValue New = CurDAG->getCopyFromReg(
             CurDAG->getEntryNode(), SDLoc(Node), AArch64::WZR, MVT::i32);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::LowerEXTRACT_VECTOR_ELT(SDValue Op,
 
   // For extractions, we just return the result directly.
   return DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ExtrTy, WideVec,
                      Op.getOperand(1));
 }
 
 SDValue AArch64TargetLowering::LowerEXTRACT_SUBVECTOR(SDValue Op,
                                                       SelectionDAG &DAG) const {
-  assert(!Op.getValueType().isScalableVector() &&
-         "Unexpected scalable type for custom lowering EXTRACT_SUBVECTOR");
-
-  EVT VT = Op.getOperand(0).getValueType();
-  SDLoc dl(Op);
-  // Just in case...
-  if (!VT.isVector())
-    return SDValue();
-
-  ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(Op.getOperand(1));
-  if (!Cst)
-    return SDValue();
-  unsigned Val = Cst->getZExtValue();
+  assert(Op.getValueType().isFixedLengthVector() &&
+         "Only cases that extract a fixed length vector are supported!");
 
+  EVT InVT = Op.getOperand(0).getValueType();
+  unsigned Idx = cast<ConstantSDNode>(Op.getOperand(1))->getZExtValue();
   unsigned Size = Op.getValueSizeInBits();
 
+  // Expand cases where EXTRACT_SUBREG cannot be used.
+  if (InVT.isScalableVector())
+    return Idx == 0 ? Op : SDValue();
+
   // This will get lowered to an appropriate EXTRACT_SUBREG in ISel.
-  if (Val == 0)
+  if (Idx == 0 && InVT.getSizeInBits() <= 128)
     return Op;
 
   // If this is extracting the upper 64-bits of a 128-bit vector, we match
   // that directly.
-  if (Size == 64 && Val * VT.getScalarSizeInBits() == 64)
+  if (Size == 64 && Idx * InVT.getScalarSizeInBits() == 64)
     return Op;
 
   return SDValue();
 }
 
 bool AArch64TargetLowering::isShuffleMaskLegal(ArrayRef<int> M, EVT VT) const {
   if (VT.getVectorNumElements() == 4 &&
       (VT.is128BitVector() || VT.is64BitVector())) {

llvm/test/CodeGen/AArch64/sve-fixed-length-subvector.ll

+; RUN: llc -aarch64-sve-vector-bits-min=128  -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefix=NO_SVE
+; RUN: llc -aarch64-sve-vector-bits-min=256  -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK
+; RUN: llc -aarch64-sve-vector-bits-min=384  -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK
+; RUN: llc -aarch64-sve-vector-bits-min=512  -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512
+; RUN: llc -aarch64-sve-vector-bits-min=640  -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512
+; RUN: llc -aarch64-sve-vector-bits-min=768  -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512
+; RUN: llc -aarch64-sve-vector-bits-min=896  -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512
+; RUN: llc -aarch64-sve-vector-bits-min=1024 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024
+; RUN: llc -aarch64-sve-vector-bits-min=1152 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024
+; RUN: llc -aarch64-sve-vector-bits-min=1280 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024
+; RUN: llc -aarch64-sve-vector-bits-min=1408 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024
+; RUN: llc -aarch64-sve-vector-bits-min=1536 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024
+; RUN: llc -aarch64-sve-vector-bits-min=1664 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024
+; RUN: llc -aarch64-sve-vector-bits-min=1792 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024
+; RUN: llc -aarch64-sve-vector-bits-min=1920 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024
+; RUN: llc -aarch64-sve-vector-bits-min=2048 -aarch64-enable-atomic-cfg-tidy=false < %s | FileCheck %s -check-prefixes=CHECK,VBITS_GE_512,VBITS_GE_1024,VBITS_GE_2048
+
+; Test we can code generater patterns of the form:
+;   fixed_length_vector = ISD::EXTRACT_SUBVECTOR scalable_vector, 0
+;   scalable_vector = ISD::INSERT_SUBVECTOR scalable_vector, fixed_length_vector, 0
+;
+; NOTE: Currently shufflevector does not support scalable vectors so it cannot
+; be used to model the above operations.  Instead these tests rely on knowing
+; how fixed length operation are lowered to scalable ones, with multiple blocks
+; ensuring insert/extract sequences are not folded away.
+
+target triple = "aarch64-unknown-linux-gnu"
+
+; Don't use SVE when its registers are no bigger than NEON.
+; NO_SVE-NOT: ptrue
+
+define void @subvector_v8i32(<8 x i32> *%in, <8 x i32>* %out) #0 {
+; CHECK-LABEL: subvector_v8i32:
+; CHECK: ptrue [[PG:p[0-9]+]].s, vl8
+; CHECK: ld1w { [[DATA:z[0-9]+.s]] }, [[PG]]/z, [x0]
+; CHECK: st1w { [[DATA]] }, [[PG]], [x1]
+; CHECK: ret
+  %a = load <8 x i32>, <8 x i32>* %in
+  br label %bb1
+
+bb1:
+  store <8 x i32> %a, <8 x i32>* %out
+  ret void
+}
+
+define void @subvector_v16i32(<16 x i32> *%in, <16 x i32>* %out) #0 {
+; CHECK-LABEL: subvector_v16i32:
+; VBITS_GE_512: ptrue [[PG:p[0-9]+]].s, vl16
+; VBITS_GE_512: ld1w { [[DATA:z[0-9]+.s]] }, [[PG]]/z, [x0]
+; VBITS_GE_512: st1w { [[DATA]] }, [[PG]], [x1]
+; CHECKT: ret
+  %a = load <16 x i32>, <16 x i32>* %in
+  br label %bb1
+
+bb1:
+  store <16 x i32> %a, <16 x i32>* %out
+  ret void
+}
+
+define void @subvector_v32i32(<32 x i32> *%in, <32 x i32>* %out) #0 {
+; CHECK-LABEL: subvector_v32i32:
+; VBITS_GE_1024: ptrue [[PG:p[0-9]+]].s, vl32
+; VBITS_GE_1024: ld1w { [[DATA:z[0-9]+.s]] }, [[PG]]/z, [x0]
+; VBITS_GE_1024: st1w { [[DATA]] }, [[PG]], [x1]
+; CHECK: ret
+  %a = load <32 x i32>, <32 x i32>* %in
+  br label %bb1
+
+bb1:
+  store <32 x i32> %a, <32 x i32>* %out
+  ret void
+}
+
+define void @subvector_v64i32(<64 x i32> *%in, <64 x i32>* %out) #0 {
+; CHECK-LABEL: subvector_v64i32:
+; VBITS_GE_2048: ptrue [[PG:p[0-9]+]].s, vl64
+; VBITS_GE_2048: ld1w { [[DATA:z[0-9]+.s]] }, [[PG]]/z, [x0]
+; VBITS_GE_2048: st1w { [[DATA]] }, [[PG]], [x1]
+; CHECK: ret
+  %a = load <64 x i32>, <64 x i32>* %in
+  br label %bb1
+
+bb1:
+  store <64 x i32> %a, <64 x i32>* %out
+  ret void
+}
+
+attributes #0 = { "target-features"="+sve" }