[AArch64][SVE] Perform fixed-width predicate AND reduction on SVE predicate vectors.

[sdesmalen]

Summary

Adds a combine for AND to 'bubble up' the sign-extend + extract to the users
of the AND operation. It also adds support in VECREDUCE_AND to perform the
operation on SVE predicate vectors.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
   setTargetDAGCombine(ISD::VSELECT);
 
   setTargetDAGCombine(ISD::INTRINSIC_VOID);
   setTargetDAGCombine(ISD::INTRINSIC_W_CHAIN);
   setTargetDAGCombine(ISD::INSERT_VECTOR_ELT);
   setTargetDAGCombine(ISD::EXTRACT_VECTOR_ELT);
   setTargetDAGCombine(ISD::VECREDUCE_ADD);
   setTargetDAGCombine(ISD::VECREDUCE_OR);
+  setTargetDAGCombine(ISD::VECREDUCE_AND);
   setTargetDAGCombine(ISD::STEP_VECTOR);
 
   setTargetDAGCombine(ISD::MGATHER);
   setTargetDAGCombine(ISD::MSCATTER);
 
   setTargetDAGCombine(ISD::FP_EXTEND);
 
   setTargetDAGCombine(ISD::GlobalAddress);
if (SDValue Pred = findScalablePredicateOperand(N->getOperand(0), DAG)) {
     EVT PromVT = getPromotedVTForPredicate(Pred.getValueType());
     SDValue PredForVL = getPredicateForFixedLengthVector(
         DAG, DL, OpVT.changeVectorElementType(PromVT.getVectorElementType()));
 
     // If not all bits in the scalable vector are defined, we need to
     // manually define these to be 0 or 1.
     if (!TLI.isAllActivePredicate(DAG, PredForVL)) {
       EVT PredVT = PredForVL.getValueType();
+      if (N->getOpcode() == ISD::VECREDUCE_OR)
         Pred = DAG.getNode(ISD::AND, DL, PredVT, Pred, PredForVL);
+      else
+        Pred = DAG.getNode(
+            ISD::OR, DL, PredVT, Pred,
+            DAG.getNode(ISD::XOR, DL, PredVT, PredForVL, PredForVL));
     }
 
     return DAG.getNode(N->getOpcode(), DL, N->getValueType(0), Pred);
   }
 
   return SDValue();
 }
 
static SDValue performANDCombine(SDNode *N,
                                  TargetLowering::DAGCombinerInfo &DCI) {
   SelectionDAG &DAG = DCI.DAG;
   SDValue LHS = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   if (SDValue R = PerformANDCSELCombine(N, DAG))
     return R;
 
+  // Try to perform the operation on SVE predicate vectors, if available.
+  SDValue NewLHS, NewRHS;
+  if (VT.isFixedLengthVector() &&
+      (NewLHS = findScalablePredicateOperand(N->getOperand(0), DAG)) &&
+      (NewRHS = findScalablePredicateOperand(N->getOperand(1), DAG))) {
+    assert(!(isa<ConstantSDNode>(NewLHS) && isa<ConstantSDNode>(NewRHS)) &&
+           "Expected nodes to have been constant folded");
+    EVT ContainerVT = getContainerForFixedLengthVector(DAG, N->getValueType(0));
+    SDValue PredAND =
+        DAG.getNode(ISD::AND, SDLoc(N), NewLHS.getValueType(), NewLHS, NewRHS);
+    SDValue Ext = DAG.getSExtOrTrunc(PredAND, SDLoc(N), ContainerVT);
+    return convertFromScalableVector(DAG, N->getValueType(0), Ext);
+  }
+
   if (!VT.isVector() || !DAG.getTargetLoweringInfo().isTypeLegal(VT))
     return SDValue();
 
   if (VT.isScalableVector())
     return performSVEAndCombine(N, DCI);
 
   // The combining code below works only for NEON vectors. In particular, it
   // does not work for SVE when dealing with vectors wider than 128 bits.
case ISD::INSERT_VECTOR_ELT:
     return performInsertVectorEltCombine(N, DCI);
   case ISD::EXTRACT_VECTOR_ELT:
     return performExtractVectorEltCombine(N, DAG);
   case ISD::VECREDUCE_ADD:
     return performVecReduceAddCombine(N, DCI.DAG, Subtarget);
   case AArch64ISD::UADDV:
     return performUADDVCombine(N, DAG);
   case ISD::VECREDUCE_OR:
+  case ISD::VECREDUCE_AND:
     return performVecreduceAndOrCombine(N, *this, DAG);
   case ISD::INTRINSIC_VOID:
   case ISD::INTRINSIC_W_CHAIN:
     switch (cast<ConstantSDNode>(N->getOperand(1))->getZExtValue()) {
     case Intrinsic::aarch64_sve_prfb_gather_scalar_offset:
       return combineSVEPrefetchVecBaseImmOff(N, DAG, 1 /*=ScalarSizeInBytes*/);
     case Intrinsic::aarch64_sve_prfh_gather_scalar_offset:
       return combineSVEPrefetchVecBaseImmOff(N, DAG, 2 /*=ScalarSizeInBytes*/);

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

 
 define i1 @ptest_and_v16i1_512bit_sve(float* %a, float * %b) vscale_range(4, 4) {
 ; CHECK-LABEL: ptest_and_v16i1_512bit_sve:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    ld1w { z0.s }, p0/z, [x0]
 ; CHECK-NEXT:    ld1w { z1.s }, p0/z, [x1]
 ; CHECK-NEXT:    fcmeq p1.s, p0/z, z0.s, #0.0
+; CHECK-NEXT:    fcmeq p2.s, p0/z, z1.s, #0.0
 ; CHECK-NEXT:    not p1.b, p0/z, p1.b
-; CHECK-NEXT:    fcmeq p0.s, p0/z, z1.s, #0.0
-; CHECK-NEXT:    bic p0.b, p1/z, p1.b, p0.b
-; CHECK-NEXT:    mov z0.s, p0/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    ptrue p0.b, vl16
-; CHECK-NEXT:    uzp1 z0.h, z0.h, z0.h
-; CHECK-NEXT:    uzp1 z0.b, z0.b, z0.b
-; CHECK-NEXT:    andv b0, p0, z0.b
-; CHECK-NEXT:    fmov w8, s0
-; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    bic p1.b, p1/z, p1.b, p2.b
+; CHECK-NEXT:    not p1.b, p0/z, p1.b
+; CHECK-NEXT:    ptest p0, p1.b
+; CHECK-NEXT:    cset w0, eq
 ; CHECK-NEXT:    ret
   %v0 = bitcast float* %a to <16 x float>*
   %v1 = load <16 x float>, <16 x float>* %v0, align 4
   %v2 = fcmp une <16 x float> %v1, zeroinitializer
   %v3 = bitcast float* %b to <16 x float>*
   %v4 = load <16 x float>, <16 x float>* %v3, align 4
   %v5 = fcmp une <16 x float> %v4, zeroinitializer
   %v6 = and <16 x i1> %v2, %v5
   %v7 = call i1 @llvm.vector.reduce.and.i1.v16i1 (<16 x i1> %v6)
   ret i1 %v7
 }
 
 define i1 @ptest_and_v16i1_512bit_min_sve(float* %a, float * %b) vscale_range(4, 0) {
 ; CHECK-LABEL: ptest_and_v16i1_512bit_min_sve:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s, vl16
 ; CHECK-NEXT:    ptrue p2.s
 ; CHECK-NEXT:    ld1w { z0.s }, p0/z, [x0]
 ; CHECK-NEXT:    ld1w { z1.s }, p0/z, [x1]
 ; CHECK-NEXT:    fcmeq p1.s, p0/z, z0.s, #0.0
 ; CHECK-NEXT:    fcmeq p0.s, p0/z, z1.s, #0.0
 ; CHECK-NEXT:    not p1.b, p2/z, p1.b
 ; CHECK-NEXT:    bic p0.b, p1/z, p1.b, p0.b
-; CHECK-NEXT:    mov z0.s, p0/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    ptrue p0.b, vl16
-; CHECK-NEXT:    uzp1 z0.h, z0.h, z0.h
-; CHECK-NEXT:    uzp1 z0.b, z0.b, z0.b
-; CHECK-NEXT:    andv b0, p0, z0.b
-; CHECK-NEXT:    fmov w8, s0
-; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    not p0.b, p2/z, p0.b
+; CHECK-NEXT:    ptest p2, p0.b
+; CHECK-NEXT:    cset w0, eq
 ; CHECK-NEXT:    ret
   %v0 = bitcast float* %a to <16 x float>*
   %v1 = load <16 x float>, <16 x float>* %v0, align 4
   %v2 = fcmp une <16 x float> %v1, zeroinitializer
   %v3 = bitcast float* %b to <16 x float>*
   %v4 = load <16 x float>, <16 x float>* %v3, align 4
   %v5 = fcmp une <16 x float> %v4, zeroinitializer
   %v6 = and <16 x i1> %v2, %v5
   %v7 = call i1 @llvm.vector.reduce.and.i1.v16i1 (<16 x i1> %v6)
   ret i1 %v7
 }
 
 declare i1 @llvm.vector.reduce.and.i1.v16i1(<16 x i1>)