[SVE][CodeGen] Fix legalisation for scalable types

[kmclaughlin]

Summary

This patch handles illegal scalable types when lowering IR operations,
addressing several places where the value of isScalableVector() is
ignored.

For types such as <vscale x 8 x i32>, this means splitting the
operations. In this example, we would split it into two
operations of type <vscale x 4 x i32> for the low and high halves.

In cases such as <vscale x 2 x i32>, the elements in the vector
will be promoted. In this case they will be promoted to
i64 (with a vector of type <vscale x 2 x i64>)

[efriedma]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

738

Can we use ElementCount here?

llvm/lib/CodeGen/TargetLoweringBase.cpp

1427

Can we use ElementCount here?

llvm/test/CodeGen/AArch64/llvm-ir-to-intrinsic.ll

107

Maybe also worth adding a testcase for <vscale x 1 x i64>, assuming that doesn't expose anything really tricky.

[kmclaughlin]

• Use ElementCount with getVectorVT

[kmclaughlin]

llvm/test/CodeGen/AArch64/llvm-ir-to-intrinsic.ll

107

Promotion is not possible for <vscale x 1 x i64> since this would result in a <vscale x 1 x i128>, which is also illegal. Instead will need to widen a type such as <vscale x 1 x i64>, which needs some more work. For fixed width vectors, the compiler will scalarise cases such as this (e.g. from 1 x i32 to just i32), which isn't something we can do for scalable vectors because of the runtime scaling.
This has never had much priority because in practice the vectoriser won't pick a VF of 1, so I think we can add support for this at a later point. Currently, tests which use types such as this will trigger the assert added to FoldBUILD_VECTOR in D78636.

[kmclaughlin]

llvm/lib/CodeGen/TargetLoweringBase.cpp

1429

I will create a separate patch to clean this up a bit by adding an overloaded operator>>

[efriedma]

LGTM

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

static void getCopyToPartsVector(SelectionDAG &DAG, const SDLoc &DL,
   assert(NumRegs == NumParts && "Part count doesn't match vector breakdown!");
   NumParts = NumRegs; // Silence a compiler warning.
   assert(RegisterVT == PartVT && "Part type doesn't match vector breakdown!");
 
   unsigned IntermediateNumElts = IntermediateVT.isVector() ?
     IntermediateVT.getVectorNumElements() : 1;
 
   // Convert the vector to the appropriate type if necessary.
-  unsigned DestVectorNoElts = NumIntermediates * IntermediateNumElts;
-
+  auto DestEltCnt = ElementCount(NumIntermediates * IntermediateNumElts,
+                                 ValueVT.isScalableVector());
   EVT BuiltVectorTy = EVT::getVectorVT(
-      *DAG.getContext(), IntermediateVT.getScalarType(), DestVectorNoElts);
+      *DAG.getContext(), IntermediateVT.getScalarType(), DestEltCnt);
   if (ValueVT != BuiltVectorTy) {
     if (SDValue Widened = widenVectorToPartType(DAG, Val, DL, BuiltVectorTy))
       Val = Widened;
 
     Val = DAG.getNode(ISD::BITCAST, DL, BuiltVectorTy, Val);
   }
 
   // Split the vector into intermediate operands.
   SmallVector<SDValue, 8> Ops(NumIntermediates);

llvm/lib/CodeGen/TargetLoweringBase.cpp

 ///
 /// This method returns the number of registers needed, and the VT for each
 /// register.  It also returns the VT and quantity of the intermediate values
 /// before they are promoted/expanded.
 unsigned TargetLoweringBase::getVectorTypeBreakdown(LLVMContext &Context, EVT VT,
                                                 EVT &IntermediateVT,
                                                 unsigned &NumIntermediates,
                                                 MVT &RegisterVT) const {
-  unsigned NumElts = VT.getVectorNumElements();
+  ElementCount EltCnt = VT.getVectorElementCount();
 
   // If there is a wider vector type with the same element type as this one,
   // or a promoted vector type that has the same number of elements which
   // are wider, then we should convert to that legal vector type.
   // This handles things like <2 x float> -> <4 x float> and
   // <4 x i1> -> <4 x i32>.
   LegalizeTypeAction TA = getTypeAction(Context, VT);
-  if (NumElts != 1 && (TA == TypeWidenVector || TA == TypePromoteInteger)) {
+  if (EltCnt.Min != 1 && (TA == TypeWidenVector || TA == TypePromoteInteger)) {
     EVT RegisterEVT = getTypeToTransformTo(Context, VT);
     if (isTypeLegal(RegisterEVT)) {
       IntermediateVT = RegisterEVT;
       RegisterVT = RegisterEVT.getSimpleVT();
       NumIntermediates = 1;
       return 1;
     }
   }
 
   // Figure out the right, legal destination reg to copy into.
   EVT EltTy = VT.getVectorElementType();
 
   unsigned NumVectorRegs = 1;
 
   // FIXME: We don't support non-power-of-2-sized vectors for now.  Ideally we
   // could break down into LHS/RHS like LegalizeDAG does.
-  if (!isPowerOf2_32(NumElts)) {
-    NumVectorRegs = NumElts;
-    NumElts = 1;
+  if (!isPowerOf2_32(EltCnt.Min)) {
+    NumVectorRegs = EltCnt.Min;
+    EltCnt.Min = 1;
   }
 
   // Divide the input until we get to a supported size.  This will always
   // end with a scalar if the target doesn't support vectors.
-  while (NumElts > 1 && !isTypeLegal(
-                                   EVT::getVectorVT(Context, EltTy, NumElts))) {
-    NumElts >>= 1;
+  while (EltCnt.Min > 1 &&
+         !isTypeLegal(EVT::getVectorVT(Context, EltTy, EltCnt))) {
+    EltCnt.Min >>= 1;
     NumVectorRegs <<= 1;
   }
 
   NumIntermediates = NumVectorRegs;
 
-  EVT NewVT = EVT::getVectorVT(Context, EltTy, NumElts);
+  EVT NewVT = EVT::getVectorVT(Context, EltTy, EltCnt);
   if (!isTypeLegal(NewVT))
     NewVT = EltTy;
   IntermediateVT = NewVT;
 
   MVT DestVT = getRegisterType(Context, NewVT);
   RegisterVT = DestVT;
   unsigned NewVTSize = NewVT.getSizeInBits();
 

llvm/test/CodeGen/AArch64/llvm-ir-to-intrinsic.ll

 ; CHECK-LABEL: @sdiv_i64
 ; CHECK-DAG: ptrue p0.d
 ; CHECK-DAG: sdiv z0.d, p0/m, z0.d, z1.d
 ; CHECK-NEXT: ret
   %div = sdiv <vscale x 2 x i64> %a, %b
   ret <vscale x 2 x i64> %div
 }
 
+define <vscale x 8 x i32> @sdiv_split_i32(<vscale x 8 x i32> %a, <vscale x 8 x i32> %b) {
+; CHECK-LABEL: @sdiv_split_i32
+; CHECK-DAG: ptrue p0.s
+; CHECK-DAG: sdiv z0.s, p0/m, z0.s, z2.s
+; CHECK-DAG: sdiv z1.s, p0/m, z1.s, z3.s
+; CHECK-NEXT: ret
+  %div = sdiv <vscale x 8 x i32> %a, %b
+  ret <vscale x 8 x i32> %div
+}
+
+define <vscale x 2 x i32> @sdiv_widen_i32(<vscale x 2 x i32> %a, <vscale x 2 x i32> %b) {
+; CHECK-LABEL: @sdiv_widen_i32
+; CHECK-DAG: ptrue p0.d
+; CHECK-DAG: sxtw z1.d, p0/m, z1.d
+; CHECK-DAG: sxtw z0.d, p0/m, z0.d
+; CHECK-DAG: sdiv z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT: ret
+  %div = sdiv <vscale x 2 x i32> %a, %b
+  ret <vscale x 2 x i32> %div
+}
+
+define <vscale x 4 x i64> @sdiv_split_i64(<vscale x 4 x i64> %a, <vscale x 4 x i64> %b) {
+; CHECK-LABEL: @sdiv_split_i64
+; CHECK-DAG: ptrue p0.d
+; CHECK-DAG: sdiv z0.d, p0/m, z0.d, z2.d
+; CHECK-DAG: sdiv z1.d, p0/m, z1.d, z3.d
+; CHECK-NEXT: ret
+  %div = sdiv <vscale x 4 x i64> %a, %b
+  ret <vscale x 4 x i64> %div
+}
+
 ;
 ; UDIV
 ;
 
 define <vscale x 4 x i32> @udiv_i32(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
 ; CHECK-LABEL: @udiv_i32
 ; CHECK-DAG: ptrue p0.s
 ; CHECK-DAG: udiv z0.s, p0/m, z0.s, z1.s
 ; CHECK-NEXT: ret
   %div = udiv <vscale x 4 x i32> %a, %b
   ret <vscale x 4 x i32> %div
 }
 
 define <vscale x 2 x i64> @udiv_i64(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b) {
 ; CHECK-LABEL: @udiv_i64
 ; CHECK-DAG: ptrue p0.d
 ; CHECK-DAG: udiv z0.d, p0/m, z0.d, z1.d
 ; CHECK-NEXT: ret
   %div = udiv <vscale x 2 x i64> %a, %b
   ret <vscale x 2 x i64> %div
 }
+
+define <vscale x 8 x i32> @udiv_split_i32(<vscale x 8 x i32> %a, <vscale x 8 x i32> %b) {
+; CHECK-LABEL: @udiv_split_i32
+; CHECK-DAG: ptrue p0.s
+; CHECK-DAG: udiv z0.s, p0/m, z0.s, z2.s
+; CHECK-DAG: udiv z1.s, p0/m, z1.s, z3.s
+; CHECK-NEXT: ret
+  %div = udiv <vscale x 8 x i32> %a, %b
+  ret <vscale x 8 x i32> %div
+}
+
+define <vscale x 2 x i32> @udiv_widen_i32(<vscale x 2 x i32> %a, <vscale x 2 x i32> %b) {
+; CHECK-LABEL: @udiv_widen_i32
+; CHECK-DAG: ptrue p0.d
+; CHECK-DAG: and z1.d, z1.d, #0xffffffff
+; CHECK-DAG: and z0.d, z0.d, #0xffffffff
+; CHECK-DAG: udiv z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT: ret
+  %div = udiv <vscale x 2 x i32> %a, %b
+  ret <vscale x 2 x i32> %div
+}
+
+define <vscale x 4 x i64> @udiv_split_i64(<vscale x 4 x i64> %a, <vscale x 4 x i64> %b) {
+; CHECK-LABEL: @udiv_split_i64
+; CHECK-DAG: ptrue p0.d
+; CHECK-DAG: udiv z0.d, p0/m, z0.d, z2.d
+; CHECK-DAG: udiv z1.d, p0/m, z1.d, z3.d
+; CHECK-NEXT: ret
+  %div = udiv <vscale x 4 x i64> %a, %b
+  ret <vscale x 4 x i64> %div
+}