[SVE][CodeGen] Fix issues with EXTRACT_SUBVECTOR when using scalable FP vectors

[david-arm]

Summary

In this patch I have fixed two issues:

1. Our SVE tuple get/set intrinsics were using the wrong constant type

for the index passed to EXTRACT_SUBVECTOR. I have fixed this by using the
function SelectionDAG::getVectorIdxConstant to create the value. Also, I
have updated the documentation for EXTRACT_SUBVECTOR describing what type
the constant index should be and we now enforce this when creating the
node.

2. The AArch64 backend was missing the appropriate patterns for

extracting certain subvectors (nxv4f16 and nxv2f32) from legal SVE types.
I have added them as part of this patch.

The only way that I could find to test the new patterns was to use the
SVE tuple get intrinsics, although I realise it looks a bit unusual.
Tests added here:

test/CodeGen/AArch64/sve-extract-subvector.ll

[efriedma]

Our SVE tuple get/set intrinsics were using the wrong constant type for the index passed to EXTRACT_SUBVECTOR. I have fixed this to use an integer pointer type.

Does this make a practical difference? If it matters, we should fix it in legalization; this isn't part of the documented target-independent semantics.

[david-arm]

Hi @efriedma, I'm happy to try and fix up during the legalisation if you think that EXTRACT_SUBVECTOR should accept MVT::i32? However, I'm not sure how I would fix this during legalisation - would this be in DAGTypeLegalizer::PromoteIntOp_EXTRACT_SUBVECTOR? I originally tried adding extract_subvector patterns similar to how it's done for AArch64 and X86, i.e. using a iPTR or i64 type for the index, but the patterns didn't match because the SVE get/set tuple intrinsics were passing MVT::i32. I had a look at every other instance in the codebase where we create EXTRACT_SUBVECTOR operations and they all either call getIntPtrConstant() or getVectorIdxConstant() to create the index type. So it just seemed like this was the way we're expected to do this. Thanks!

[efriedma]

Maybe we should just change the documentation and enforce the rule in SelectionDAG::getNode()? Maybe simpler. I guess it has to be a constant anyway?

Doing it in "legalization" would imply operation legalization, not type legalization, since i32 is a legal type. Probably have to do it in target-specific code, since there isn't any good way to express the rule to LegalizaeDAG.

[efriedma]

LGTM

llvm/include/llvm/CodeGen/ISDOpcodes.h

enum NodeType {
 
   /// EXTRACT_SUBVECTOR(VECTOR, IDX) - Returns a subvector from VECTOR.
   /// Let the result type be T, then IDX represents the starting element number
   /// from which a subvector of type T is extracted. IDX must be a constant
   /// multiple of T's known minimum vector length. If T is a scalable vector,
   /// IDX is first scaled by the runtime scaling factor of T. Elements IDX
   /// through (IDX + num_elements(T) - 1) must be valid VECTOR indices. If this
   /// condition cannot be determined statically but is false at runtime, then
-  /// the result vector is undefined.
+  /// the result vector is undefined. The IDX parameter must be a vector index
+  /// constant type, which for most targets will be an integer pointer type.
   ///
   /// This operation supports extracting a fixed-width vector from a scalable
   /// vector, but not the other way around.
   EXTRACT_SUBVECTOR,
 
   /// VECTOR_SHUFFLE(VEC1, VEC2) - Returns a vector, of the same type as
   /// VEC1/VEC2.  A VECTOR_SHUFFLE node also contains an array of constant int
   /// values that indicate which value (or undef) each result element will

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::EXTRACT_SUBVECTOR:
     assert((VT.isScalableVector() != N1VT.isScalableVector() ||
             VT.getVectorMinNumElements() <= N1VT.getVectorMinNumElements()) &&
            "Extract subvector must be from larger vector to smaller vector!");
     assert(N2C && "Extract subvector index must be a constant");
     assert((VT.isScalableVector() != N1VT.isScalableVector() ||
             (VT.getVectorMinNumElements() + N2C->getZExtValue()) <=
                 N1VT.getVectorMinNumElements()) &&
            "Extract subvector overflow!");
+    assert(N2C->getAPIntValue().getBitWidth() ==
+               TLI->getVectorIdxTy(getDataLayout())
+                   .getSizeInBits()
+                   .getFixedSize() &&
+           "Constant index for EXTRACT_SUBVECTOR has an invalid size");
 
     // Trivial extraction.
     if (VT == N1VT)
       return N1;
 
     // EXTRACT_SUBVECTOR of an UNDEF is an UNDEF.
     if (N1.isUndef())
       return getUNDEF(VT);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case Intrinsic::aarch64_sve_tuple_get: {
       SDLoc DL(N);
       SDValue Chain = N->getOperand(0);
       SDValue Src1 = N->getOperand(2);
       SDValue Idx = N->getOperand(3);
 
       uint64_t IdxConst = cast<ConstantSDNode>(Idx)->getZExtValue();
       EVT ResVT = N->getValueType(0);
       uint64_t NumLanes = ResVT.getVectorElementCount().Min;
+      SDValue ExtIdx = DAG.getVectorIdxConstant(IdxConst * NumLanes, DL);
       SDValue Val =
-          DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ResVT, Src1,
-                      DAG.getConstant(IdxConst * NumLanes, DL, MVT::i32));
+          DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ResVT, Src1, ExtIdx);
       return DAG.getMergeValues({Val, Chain}, DL);
     }
     case Intrinsic::aarch64_sve_tuple_set: {
       SDLoc DL(N);
       SDValue Chain = N->getOperand(0);
       SDValue Tuple = N->getOperand(2);
       SDValue Idx = N->getOperand(3);
       SDValue Vec = N->getOperand(4);
case Intrinsic::aarch64_sve_tuple_set: {
 
       uint64_t NumVecs = TupleLanes / NumLanes;
 
       SmallVector<SDValue, 4> Opnds;
       for (unsigned I = 0; I < NumVecs; ++I) {
         if (I == IdxConst)
           Opnds.push_back(Vec);
         else {
-          Opnds.push_back(
-              DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, Vec.getValueType(), Tuple,
-                          DAG.getConstant(I * NumLanes, DL, MVT::i32)));
+          SDValue ExtIdx = DAG.getVectorIdxConstant(I * NumLanes, DL);
+          Opnds.push_back(DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL,
+                                      Vec.getValueType(), Tuple, ExtIdx));
         }
       }
       SDValue Concat =
           DAG.getNode(ISD::CONCAT_VECTORS, DL, Tuple.getValueType(), Opnds);
       return DAG.getMergeValues({Concat, Chain}, DL);
     }
     case Intrinsic::aarch64_sve_tuple_create2:
     case Intrinsic::aarch64_sve_tuple_create3:

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

def : Pat<(nxv4i1 (extract_subvector (nxv8i1 PPR:$Ps), (i64 0))),
             (ZIP1_PPP_H PPR:$Ps, (PFALSE))>;
   def : Pat<(nxv4i1 (extract_subvector (nxv8i1 PPR:$Ps), (i64 4))),
             (ZIP2_PPP_H PPR:$Ps, (PFALSE))>;
   def : Pat<(nxv8i1 (extract_subvector (nxv16i1 PPR:$Ps), (i64 0))),
             (ZIP1_PPP_B PPR:$Ps, (PFALSE))>;
   def : Pat<(nxv8i1 (extract_subvector (nxv16i1 PPR:$Ps), (i64 8))),
             (ZIP2_PPP_B PPR:$Ps, (PFALSE))>;
 
+  // Extract subvectors from FP SVE vectors
+  def : Pat<(nxv4f16 (extract_subvector (nxv8f16 ZPR:$Zs), (i64 0))),
+            (UUNPKLO_ZZ_S ZPR:$Zs)>;
+  def : Pat<(nxv4f16 (extract_subvector (nxv8f16 ZPR:$Zs), (i64 4))),
+            (UUNPKHI_ZZ_S ZPR:$Zs)>;
+  def : Pat<(nxv2f32 (extract_subvector (nxv4f32 ZPR:$Zs), (i64 0))),
+            (UUNPKLO_ZZ_D ZPR:$Zs)>;
+  def : Pat<(nxv2f32 (extract_subvector (nxv4f32 ZPR:$Zs), (i64 2))),
+            (UUNPKHI_ZZ_D ZPR:$Zs)>;
+
   // Concatenate two predicates.
   def : Pat<(nxv4i1 (concat_vectors nxv2i1:$p1, nxv2i1:$p2)),
             (UZP1_PPP_S $p1, $p2)>;
   def : Pat<(nxv8i1 (concat_vectors nxv4i1:$p1, nxv4i1:$p2)),
             (UZP1_PPP_H $p1, $p2)>;
   def : Pat<(nxv16i1 (concat_vectors nxv8i1:$p1, nxv8i1:$p2)),
             (UZP1_PPP_B $p1, $p2)>;
 

llvm/test/CodeGen/AArch64/sve-extract-subvector.ll

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov z0.d, z1.d
 ; CHECK-NEXT:    ret
   %z0_z1_bc = bitcast <vscale x 32 x i8> %z0_z1 to <vscale x 4 x i64>
   %ext = call <vscale x 2 x i64> @llvm.aarch64.sve.tuple.get.nxv4i64(<vscale x 4 x i64> %z0_z1_bc, i32 1)
   ret <vscale x 2 x i64> %ext
 }
 
+define <vscale x 4 x half> @extract_lo_nxv4f16_nxv8f16(<vscale x 8 x half> %z0) {
+; CHECK-LABEL: extract_lo_nxv4f16_nxv8f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpklo z0.s, z0.h
+; CHECK-NEXT:    ret
+  %ext = call <vscale x 4 x half> @llvm.aarch64.sve.tuple.get.nxv8f16(<vscale x 8 x half> %z0, i32 0)
+  ret <vscale x 4 x half> %ext
+}
+
+define <vscale x 4 x half> @extract_hi_nxv4f16_nxv8f16(<vscale x 8 x half> %z0) {
+; CHECK-LABEL: extract_hi_nxv4f16_nxv8f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpkhi z0.s, z0.h
+; CHECK-NEXT:    ret
+  %ext = call <vscale x 4 x half> @llvm.aarch64.sve.tuple.get.nxv8f16(<vscale x 8 x half> %z0, i32 1)
+  ret <vscale x 4 x half> %ext
+}
+
+define <vscale x 2 x float> @extract_lo_nxv2f32_nxv4f32(<vscale x 4 x float> %z0) {
+; CHECK-LABEL: extract_lo_nxv2f32_nxv4f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpklo z0.d, z0.s
+; CHECK-NEXT:    ret
+  %ext = call <vscale x 2 x float> @llvm.aarch64.sve.tuple.get.nxv4f32(<vscale x 4 x float> %z0, i32 0)
+  ret <vscale x 2 x float> %ext
+}
+
+define <vscale x 2 x float> @extract_hi_nxv2f32_nxv4f32(<vscale x 4 x float> %z0) {
+; CHECK-LABEL: extract_hi_nxv2f32_nxv4f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpkhi z0.d, z0.s
+; CHECK-NEXT:    ret
+  %ext = call <vscale x 2 x float> @llvm.aarch64.sve.tuple.get.nxv4f32(<vscale x 4 x float> %z0, i32 1)
+  ret <vscale x 2 x float> %ext
+}
+
 declare <vscale x 2 x i64> @llvm.aarch64.sve.tuple.get.nxv4i64(<vscale x 4 x i64>, i32)
 declare <vscale x 16 x i8> @llvm.aarch64.sve.tuple.get.nxv32i8(<vscale x 32 x i8>, i32)
+declare <vscale x 2 x float> @llvm.aarch64.sve.tuple.get.nxv4f32(<vscale x 4 x float>, i32)
+declare <vscale x 4 x half> @llvm.aarch64.sve.tuple.get.nxv8f16(<vscale x 8 x half>, i32)