[SelectionDAG] Make WidenVecRes_Convert work for scalable vectors.

[sdesmalen]

Summary

Most of the code wasn't yet scalable safe, although most of the
code conceptually just works for scalable vectors. This change
makes the algorithm work on ElementCount, where appropriate,
and leaves the fixed-width only code to use getFixedNumElements.

[david-arm]

llvm/test/CodeGen/AArch64/sve-fcvt.ll

740

Hi @sdesmalen, not sure this is caused by your patch, but this looks a bit odd as we're defining p0.s after it's already been used above.

[david-arm]

LGTM!

llvm/test/CodeGen/AArch64/sve-fcvt.ll

740

My mistake - we're using the input p0 passed to the function. Doh!

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

 }
 
 SDValue DAGTypeLegalizer::WidenVecRes_Convert(SDNode *N) {
   LLVMContext &Ctx = *DAG.getContext();
   SDValue InOp = N->getOperand(0);
   SDLoc DL(N);
 
   EVT WidenVT = TLI.getTypeToTransformTo(Ctx, N->getValueType(0));
-  unsigned WidenNumElts = WidenVT.getVectorNumElements();
+  ElementCount WidenEC = WidenVT.getVectorElementCount();
 
   EVT InVT = InOp.getValueType();
 
   unsigned Opcode = N->getOpcode();
   const SDNodeFlags Flags = N->getFlags();
 
   // Handle the case of ZERO_EXTEND where the promoted InVT element size does
   // not equal that of WidenVT.
   if (N->getOpcode() == ISD::ZERO_EXTEND &&
       getTypeAction(InVT) == TargetLowering::TypePromoteInteger &&
       TLI.getTypeToTransformTo(Ctx, InVT).getScalarSizeInBits() !=
       WidenVT.getScalarSizeInBits()) {
     InOp = ZExtPromotedInteger(InOp);
     InVT = InOp.getValueType();
     if (WidenVT.getScalarSizeInBits() < InVT.getScalarSizeInBits())
       Opcode = ISD::TRUNCATE;
   }
 
   EVT InEltVT = InVT.getVectorElementType();
-  EVT InWidenVT = EVT::getVectorVT(Ctx, InEltVT, WidenNumElts);
-  unsigned InVTNumElts = InVT.getVectorNumElements();
+  EVT InWidenVT = EVT::getVectorVT(Ctx, InEltVT, WidenEC);
+  ElementCount InVTEC = InVT.getVectorElementCount();
 
   if (getTypeAction(InVT) == TargetLowering::TypeWidenVector) {
     InOp = GetWidenedVector(N->getOperand(0));
     InVT = InOp.getValueType();
-    InVTNumElts = InVT.getVectorNumElements();
-    if (InVTNumElts == WidenNumElts) {
+    InVTEC = InVT.getVectorElementCount();
+    if (InVTEC == WidenEC) {
       if (N->getNumOperands() == 1)
         return DAG.getNode(Opcode, DL, WidenVT, InOp);
       return DAG.getNode(Opcode, DL, WidenVT, InOp, N->getOperand(1), Flags);
     }
     if (WidenVT.getSizeInBits() == InVT.getSizeInBits()) {
       // If both input and result vector types are of same width, extend
       // operations should be done with SIGN/ZERO_EXTEND_VECTOR_INREG, which
       // accepts fewer elements in the result than in the input.
       if (Opcode == ISD::ANY_EXTEND)
         return DAG.getNode(ISD::ANY_EXTEND_VECTOR_INREG, DL, WidenVT, InOp);
       if (Opcode == ISD::SIGN_EXTEND)
         return DAG.getNode(ISD::SIGN_EXTEND_VECTOR_INREG, DL, WidenVT, InOp);
       if (Opcode == ISD::ZERO_EXTEND)
         return DAG.getNode(ISD::ZERO_EXTEND_VECTOR_INREG, DL, WidenVT, InOp);
     }
   }
 
   if (TLI.isTypeLegal(InWidenVT)) {
     // Because the result and the input are different vector types, widening
     // the result could create a legal type but widening the input might make
     // it an illegal type that might lead to repeatedly splitting the input
     // and then widening it. To avoid this, we widen the input only if
     // it results in a legal type.
-    if (WidenNumElts % InVTNumElts == 0) {
+    if (WidenEC.isKnownMultipleOf(InVTEC.getKnownMinValue())) {
       // Widen the input and call convert on the widened input vector.
-      unsigned NumConcat = WidenNumElts/InVTNumElts;
+      unsigned NumConcat =
+          WidenEC.getKnownMinValue() / InVTEC.getKnownMinValue();
       SmallVector<SDValue, 16> Ops(NumConcat, DAG.getUNDEF(InVT));
       Ops[0] = InOp;
       SDValue InVec = DAG.getNode(ISD::CONCAT_VECTORS, DL, InWidenVT, Ops);
       if (N->getNumOperands() == 1)
         return DAG.getNode(Opcode, DL, WidenVT, InVec);
       return DAG.getNode(Opcode, DL, WidenVT, InVec, N->getOperand(1), Flags);
     }
 
-    if (InVTNumElts % WidenNumElts == 0) {
+    if (InVTEC.isKnownMultipleOf(WidenEC.getKnownMinValue())) {
       SDValue InVal = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, InWidenVT, InOp,
                                   DAG.getVectorIdxConstant(0, DL));
       // Extract the input and convert the shorten input vector.
       if (N->getNumOperands() == 1)
         return DAG.getNode(Opcode, DL, WidenVT, InVal);
       return DAG.getNode(Opcode, DL, WidenVT, InVal, N->getOperand(1), Flags);
     }
   }
 
   // Otherwise unroll into some nasty scalar code and rebuild the vector.
   EVT EltVT = WidenVT.getVectorElementType();
-  SmallVector<SDValue, 16> Ops(WidenNumElts, DAG.getUNDEF(EltVT));
+  SmallVector<SDValue, 16> Ops(WidenEC.getFixedValue(), DAG.getUNDEF(EltVT));
   // Use the original element count so we don't do more scalar opts than
   // necessary.
   unsigned MinElts = N->getValueType(0).getVectorNumElements();
   for (unsigned i=0; i < MinElts; ++i) {
     SDValue Val = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, InEltVT, InOp,
                               DAG.getVectorIdxConstant(i, DL));
     if (N->getNumOperands() == 1)
       Ops[i] = DAG.getNode(Opcode, DL, EltVT, Val);

llvm/test/CodeGen/AArch64/sve-fcvt.ll

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    fcvt z0.s, p0/m, z0.h
 ; CHECK-NEXT:    ret
   %res = fpext <vscale x 2 x half> %a to <vscale x 2 x float>
   ret <vscale x 2 x float> %res
 }
 
+define <vscale x 3 x float> @fcvts_nxv3f16(<vscale x 3 x half> %a) {
+; CHECK-LABEL: fcvts_nxv3f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    fcvt z0.s, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = fpext <vscale x 3 x half> %a to <vscale x 3 x float>
+  ret <vscale x 3 x float> %res
+}
+
 define <vscale x 4 x float> @fcvts_nxv4f16(<vscale x 4 x half> %a) {
 ; CHECK-LABEL: fcvts_nxv4f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fcvt z0.s, p0/m, z0.h
 ; CHECK-NEXT:    ret
   %res = fpext <vscale x 4 x half> %a to <vscale x 4 x float>
   ret <vscale x 4 x float> %res

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    fcvt z0.h, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = fptrunc <vscale x 2 x float> %a to <vscale x 2 x half>
   ret <vscale x 2 x half> %res
 }
 
+define <vscale x 3 x half> @fcvth_nxv3f32(<vscale x 3 x float> %a) {
+; CHECK-LABEL: fcvth_nxv3f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    fcvt z0.h, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %res = fptrunc <vscale x 3 x float> %a to <vscale x 3 x half>
+  ret <vscale x 3 x half> %res
+}
+
 define <vscale x 4 x half> @fcvth_nxv4f32(<vscale x 4 x float> %a) {
 ; CHECK-LABEL: fcvth_nxv4f32:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fcvt z0.h, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = fptrunc <vscale x 4 x float> %a to <vscale x 4 x half>
   ret <vscale x 4 x half> %res

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fcvtzs z0.s, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = fptosi <vscale x 4 x float> %a to <vscale x 4 x i16>
   ret <vscale x 4 x i16> %res
 }
 
+define <vscale x 7 x i16> @fcvtzs_h_nxv7f16(<vscale x 7 x half> %a) {
+; CHECK-LABEL: fcvtzs_h_nxv7f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    fcvtzs z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = fptosi <vscale x 7 x half> %a to <vscale x 7 x i16>
+  ret <vscale x 7 x i16> %res
+}
+
 define <vscale x 8 x i16> @fcvtzs_h_nxv8f16(<vscale x 8 x half> %a) {
 ; CHECK-LABEL: fcvtzs_h_nxv8f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    fcvtzs z0.h, p0/m, z0.h
 ; CHECK-NEXT:    ret
   %res = fptosi <vscale x 8 x half> %a to <vscale x 8 x i16>
   ret <vscale x 8 x i16> %res

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fcvtzs z0.s, p0/m, z0.h
 ; CHECK-NEXT:    ret
   %res = fptosi <vscale x 4 x half> %a to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %res
 }
 
+define <vscale x 3 x i32> @fcvtzs_s_nxv3f16(<vscale x 3 x half> %a) {
+; CHECK-LABEL: fcvtzs_s_nxv3f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    fcvtzs z0.s, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = fptosi <vscale x 3 x half> %a to <vscale x 3 x i32>
+  ret <vscale x 3 x i32> %res
+}
+
 define <vscale x 4 x i32> @fcvtzs_s_nxv4f32(<vscale x 4 x float> %a) {
 ; CHECK-LABEL: fcvtzs_s_nxv4f32:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fcvtzs z0.s, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = fptosi <vscale x 4 x float> %a to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %res

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fcvtzs z0.s, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = fptosi <vscale x 4 x float> %a to <vscale x 4 x i16>
   ret <vscale x 4 x i16> %res
 }
 
+define <vscale x 7 x i16> @fcvtzu_h_nxv7f16(<vscale x 7 x half> %a) {
+; CHECK-LABEL: fcvtzu_h_nxv7f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    fcvtzu z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = fptoui <vscale x 7 x half> %a to <vscale x 7 x i16>
+  ret <vscale x 7 x i16> %res
+}
+
 define <vscale x 8 x i16> @fcvtzu_h_nxv8f16(<vscale x 8 x half> %a) {
 ; CHECK-LABEL: fcvtzu_h_nxv8f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    fcvtzu z0.h, p0/m, z0.h
 ; CHECK-NEXT:    ret
   %res = fptoui <vscale x 8 x half> %a to <vscale x 8 x i16>
   ret <vscale x 8 x i16> %res

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    fcvtzs z0.d, p0/m, z0.d
 ; CHECK-NEXT:    ret
   %res = fptoui <vscale x 2 x double> %a to <vscale x 2 x i32>
   ret <vscale x 2 x i32> %res
 }
 
+define <vscale x 3 x i32> @fcvtzu_s_nxv3f16(<vscale x 3 x half> %a) {
+; CHECK-LABEL: fcvtzu_s_nxv3f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    fcvtzu z0.s, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = fptoui <vscale x 3 x half> %a to <vscale x 3 x i32>
+  ret <vscale x 3 x i32> %res
+}
+
+define <vscale x 3 x i32> @fcvtzu_s_nxv3f32(<vscale x 3 x float> %a) {
+; CHECK-LABEL: fcvtzu_s_nxv3f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    fcvtzu z0.s, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %res = fptoui <vscale x 3 x float> %a to <vscale x 3 x i32>
+  ret <vscale x 3 x i32> %res
+}
+
 define <vscale x 4 x i32> @fcvtzu_s_nxv4f16(<vscale x 4 x half> %a) {
 ; CHECK-LABEL: fcvtzu_s_nxv4f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fcvtzu z0.s, p0/m, z0.h
 ; CHECK-NEXT:    ret
   %res = fptoui <vscale x 4 x half> %a to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %res

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    scvtf z0.h, p0/m, z0.d
 ; CHECK-NEXT:    ret
   %res = sitofp <vscale x 2 x i64> %a to <vscale x 2 x half>
   ret <vscale x 2 x half> %res
 }
 
+define <vscale x 3 x half> @scvtf_h_nxv3i1(<vscale x 3 x i1> %a) {
+; CHECK-LABEL: scvtf_h_nxv3i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.s, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    scvtf z0.h, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %res = sitofp <vscale x 3 x i1> %a to <vscale x 3 x half>
+  ret <vscale x 3 x half> %res
+}
+
+define <vscale x 3 x half> @scvtf_h_nxv3i16(<vscale x 3 x i16> %a) {
+; CHECK-LABEL: scvtf_h_nxv3i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    scvtf z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = sitofp <vscale x 3 x i16> %a to <vscale x 3 x half>
+  ret <vscale x 3 x half> %res
+}
+
 define <vscale x 4 x half> @scvtf_h_nxv4i1(<vscale x 4 x i1> %a) {
 ; CHECK-LABEL: scvtf_h_nxv4i1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov z0.s, p0/z, #-1 // =0xffffffffffffffff
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    scvtf z0.h, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = sitofp <vscale x 4 x i1> %a to <vscale x 4 x half>

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    scvtf z0.h, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = sitofp <vscale x 4 x i32> %a to <vscale x 4 x half>
   ret <vscale x 4 x half> %res
 }
 
+define <vscale x 7 x half> @scvtf_h_nxv7i1(<vscale x 7 x i1> %a) {
+; CHECK-LABEL: scvtf_h_nxv7i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.h, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    scvtf z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = sitofp <vscale x 7 x i1> %a to <vscale x 7 x half>
+  ret <vscale x 7 x half> %res
+}
+
+define <vscale x 7 x half> @scvtf_h_nxv7i16(<vscale x 7 x i16> %a) {
+; CHECK-LABEL: scvtf_h_nxv7i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    scvtf z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = sitofp <vscale x 7 x i16> %a to <vscale x 7 x half>
+  ret <vscale x 7 x half> %res
+}
+
 define <vscale x 8 x half> @scvtf_h_nxv8i1(<vscale x 8 x i1> %a) {
 ; CHECK-LABEL: scvtf_h_nxv8i1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov z0.h, p0/z, #-1 // =0xffffffffffffffff
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    scvtf z0.h, p0/m, z0.h
 ; CHECK-NEXT:    ret
   %res = sitofp <vscale x 8 x i1> %a to <vscale x 8 x half>

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    scvtf z0.s, p0/m, z0.d
 ; CHECK-NEXT:    ret
   %res = sitofp <vscale x 2 x i64> %a to <vscale x 2 x float>
   ret <vscale x 2 x float> %res
 }
 
+define <vscale x 3 x float> @scvtf_s_nxv3i1(<vscale x 3 x i1> %a) {
+; CHECK-LABEL: scvtf_s_nxv3i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.s, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    scvtf z0.s, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %res = sitofp <vscale x 3 x i1> %a to <vscale x 3 x float>
+  ret <vscale x 3 x float> %res
+}
+
+define <vscale x 3 x float> @scvtf_s_nxv3i32(<vscale x 3 x i32> %a) {
+; CHECK-LABEL: scvtf_s_nxv3i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    scvtf z0.s, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %res = sitofp <vscale x 3 x i32> %a to <vscale x 3 x float>
+  ret <vscale x 3 x float> %res
+}
+
 define <vscale x 4 x float> @scvtf_s_nxv4i1(<vscale x 4 x i1> %a) {
 ; CHECK-LABEL: scvtf_s_nxv4i1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov z0.s, p0/z, #-1 // =0xffffffffffffffff
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    scvtf z0.s, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = sitofp <vscale x 4 x i1> %a to <vscale x 4 x float>

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    ucvtf z0.h, p0/m, z0.d
 ; CHECK-NEXT:    ret
   %res = uitofp <vscale x 2 x i64> %a to <vscale x 2 x half>
   ret <vscale x 2 x half> %res
 }
 
+define <vscale x 3 x half> @ucvtf_h_nxv3i1(<vscale x 3 x i1> %a) {
+; CHECK-LABEL: ucvtf_h_nxv3i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.s, p0/z, #1 // =0x1
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ucvtf z0.h, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %res = uitofp <vscale x 3 x i1> %a to <vscale x 3 x half>
+  ret <vscale x 3 x half> %res
+}
+
+define <vscale x 3 x half> @ucvtf_h_nxv3i16(<vscale x 3 x i16> %a) {
+; CHECK-LABEL: ucvtf_h_nxv3i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ucvtf z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %res = uitofp <vscale x 3 x i16> %a to <vscale x 3 x half>
+  ret <vscale x 3 x half> %res
+}
+
+define <vscale x 3 x half> @ucvtf_h_nxv3i32(<vscale x 3 x i32> %a) {
+; CHECK-LABEL: ucvtf_h_nxv3i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ucvtf z0.h, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %res = uitofp <vscale x 3 x i32> %a to <vscale x 3 x half>
+  ret <vscale x 3 x half> %res
+}
+
 define <vscale x 4 x half> @ucvtf_h_nxv4i1(<vscale x 4 x i1> %a) {
 ; CHECK-LABEL: ucvtf_h_nxv4i1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov z0.s, p0/z, #1 // =0x1
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    ucvtf z0.h, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %res = uitofp <vscale x 4 x i1> %a to <vscale x 4 x half>