[llvm] [AArch64][SVE] Rework VECTOR_COMPRESS lowering (PR #171162)

[Gaëtan Bossu via llvm-commits]

================
@@ -7445,60 +7448,43 @@ static SDValue convertFromSVEContainerType(SDLoc DL, SDValue Vec, EVT VecVT,
   return Vec;
 }
 
-SDValue AArch64TargetLowering::LowerVECTOR_COMPRESS(SDValue Op,
-                                                    SelectionDAG &DAG) const {
+SDValue AArch64TargetLowering::LowerFixedLengthVectorCompressToSVE(
+    SDValue Op, SelectionDAG &DAG) const {
   SDLoc DL(Op);
-  SDValue Vec = Op.getOperand(0);
-  SDValue Mask = Op.getOperand(1);
-  SDValue Passthru = Op.getOperand(2);
-  EVT VecVT = Vec.getValueType();
-  EVT MaskVT = Mask.getValueType();
-  EVT ElmtVT = VecVT.getVectorElementType();
-  const bool IsFixedLength = VecVT.isFixedLengthVector();
-  const bool HasPassthru = !Passthru.isUndef();
-  unsigned MinElmts = VecVT.getVectorElementCount().getKnownMinValue();
-  EVT FixedVecVT = MVT::getVectorVT(ElmtVT.getSimpleVT(), MinElmts);
+  EVT VT = Op.getValueType();
 
-  assert(VecVT.isVector() && "Input to VECTOR_COMPRESS must be vector.");
+  EVT ContainerVT = getContainerForFixedLengthVector(DAG, VT);
+  SDValue Vec = convertToScalableVector(DAG, ContainerVT, Op.getOperand(0));
+  SDValue Mask = convertFixedMaskToScalableVector(Op.getOperand(1), DAG);
+  SDValue Passthru =
+      convertToScalableVector(DAG, ContainerVT, Op.getOperand(2));
 
-  if (!Subtarget->isSVEAvailable())
-    return SDValue();
+  SDValue Result =
+      DAG.getNode(ISD::VECTOR_COMPRESS, DL, ContainerVT, Vec, Mask, Passthru);
+  return convertFromScalableVector(DAG, VT, Result);
+}
 
-  if (IsFixedLength && VecVT.getSizeInBits().getFixedValue() > 128)
-    return SDValue();
+SDValue AArch64TargetLowering::LowerVECTOR_COMPRESS(SDValue Op,
+                                                    SelectionDAG &DAG) const {
+  SDLoc DL(Op);
+  EVT VT = Op.getValueType();
+  if (VT.isFixedLengthVector())
+    return LowerFixedLengthVectorCompressToSVE(Op, DAG);
 
-  // Only <vscale x {4|2} x {i32|i64}> supported for compact.
-  if (MinElmts != 2 && MinElmts != 4)
+  if (!Subtarget->isSVEAvailable())
     return SDValue();
 
-  // We can use the SVE register containing the NEON vector in its lowest bits.
-  if (IsFixedLength) {
-    EVT ScalableVecVT =
-        MVT::getScalableVectorVT(ElmtVT.getSimpleVT(), MinElmts);
-    EVT ScalableMaskVT = MVT::getScalableVectorVT(
-        MaskVT.getVectorElementType().getSimpleVT(), MinElmts);
-
-    Vec = DAG.getNode(ISD::INSERT_SUBVECTOR, DL, ScalableVecVT,
-                      DAG.getUNDEF(ScalableVecVT), Vec,
-                      DAG.getConstant(0, DL, MVT::i64));
-    Mask = DAG.getNode(ISD::INSERT_SUBVECTOR, DL, ScalableMaskVT,
-                       DAG.getUNDEF(ScalableMaskVT), Mask,
-                       DAG.getConstant(0, DL, MVT::i64));
-    Mask = DAG.getNode(ISD::TRUNCATE, DL,
-                       ScalableMaskVT.changeVectorElementType(MVT::i1), Mask);
-    Passthru = DAG.getNode(ISD::INSERT_SUBVECTOR, DL, ScalableVecVT,
-                           DAG.getUNDEF(ScalableVecVT), Passthru,
-                           DAG.getConstant(0, DL, MVT::i64));
+  SDValue Vec = Op.getOperand(0);
+  SDValue Mask = Op.getOperand(1);
+  SDValue Passthru = Op.getOperand(2);
+  EVT MaskVT = Mask.getValueType();
 
-    VecVT = Vec.getValueType();
-    MaskVT = Mask.getValueType();
-  }
+  assert(VT.isVector() && "Input to VECTOR_COMPRESS must be vector.");
 
   // Get legal type for compact instruction
-  EVT ContainerVT = getSVEContainerType(VecVT);
+  EVT ContainerVT = getSVEContainerType(VT);
 
-  // Convert to 32 or 64 bits for smaller types, as these are the only supported
-  // sizes for compact.
+  // Convert to a packed 32/64-bit SVE vector of the same element count as VT.
----------------
gbossu wrote:

How do we know that `ContainerVT` will be a vector of 32bit or 64bit elements? I don't see anything that obviously enforces it.

https://github.com/llvm/llvm-project/pull/171162