[AArch64] Perform first active true vector combine

[Allen]

Summary

Materialize : i1 = extract_vector_elt t37, Constant:i64<0>

... into: "ptrue p, all" + PTEST

Test bit of lane 0 can use P register directly, and the instruction “pture all”
is loop invariant, which will beneficial to SVE after hoisting out the loop.

[david-arm]

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

486

To be honest, this doesn't look like a win to me. We're introducing control flow where previously we had none. I think we should avoid setting the flags when the original code set none.

I think I'd prefer it if you restricted the DAG combine to only cases where we're extracting from a flag-setting operation, i.e. an extract from a icmp.

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

1

Maybe this is better in something like sve-cmp-folds.ll?

51

This looks quite an odd IR sequence and I don't think will ever be generated by the vectoriser. I imagine this has come from hand-written SVE ACLE code.

[Allen]

add SetCC.getOpcode() != ISD::SETCC, and now the case in file llvm/test/CodeGen/AArch64/sve-extract-element.ll is not touched

[Allen]

merge the test into file llvm/test/CodeGen/AArch64/sve-cmp-folds.ll

[david-arm]

LGTM! Thanks for making the changes @Allen.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

15279

You could use Idx->isZero() here.

18165

performExtractVectorEltCombine looks like a common placeholder for all related combines so can the call to performFirstTrueTestVectorCombine me moved into it?

llvm/test/CodeGen/AArch64/sve-cmp-folds.ll

56–67 ↗

(On Diff #413273)

This test looks rather more involved than the DAGCombine it's trying to verify. Why not just:

define i1 @foo(<vscale x 4 x float> %a, <vscale x 4 x float> %b) #0 {
  %vcond = fcmp oeq <vscale x 4 x float> %a, %b
  %bit = extractelement <vscale x 4 x i1> %vcond, i64 0
  ret i1 %bit
}

[Allen]

Address comments

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14312

Sorry I missed this previously but I think this is bug and should be
if (!Subtarget->hasSVE() || DCI.isBeforeLegalize()) as otherwise you might create a PTEST with illegal types, which cannot be legalised and will instead cause a compiler crash.

14336

I think it's worth having this same assert inside performFirstTrueTestVectorCombine just in case it gets called via a different route in the future.

14337

Unwanted space?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static SDValue performExtractVectorEltCombine(SDNode *N, SelectionDAG &DAG) {
   SDValue N0 = N->getOperand(0), N1 = N->getOperand(1);
   ConstantSDNode *ConstantN1 = dyn_cast<ConstantSDNode>(N1);
 
   EVT VT = N->getValueType(0);
   const bool FullFP16 =
       static_cast<const AArch64Subtarget &>(DAG.getSubtarget()).hasFullFP16();
   bool IsStrict = N0->isStrictFPOpcode();
 
   // Rewrite for pairwise fadd pattern
   //   (f32 (extract_vector_elt
   //           (fadd (vXf32 Other)
   //                 (vector_shuffle (vXf32 Other) undef <1,X,...> )) 0))
   // ->
   //   (f32 (fadd (extract_vector_elt (vXf32 Other) 0)
   //              (extract_vector_elt (vXf32 Other) 1))
   // For strict_fadd we need to make sure the old strict_fadd can be deleted, so
   // we can only do this when it's used only by the extract_vector_elt.
   if (ConstantN1 && ConstantN1->getZExtValue() == 0 &&
       hasPairwiseAdd(N0->getOpcode(), VT, FullFP16) &&
       (!IsStrict || N0.hasOneUse())) {
     SDLoc DL(N0);
     SDValue N00 = N0->getOperand(IsStrict ? 1 : 0);
     SDValue N01 = N0->getOperand(IsStrict ? 2 : 1);
 
     ShuffleVectorSDNode *Shuffle = dyn_cast<ShuffleVectorSDNode>(N01);
     SDValue Other = N00;
 
     // And handle the commutative case.
     if (!Shuffle) {
       Shuffle = dyn_cast<ShuffleVectorSDNode>(N00);
       Other = N01;
     }
 
     if (Shuffle && Shuffle->getMaskElt(0) == 1 &&
         Other == Shuffle->getOperand(0)) {
       SDValue Extract1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, VT, Other,
                                      DAG.getConstant(0, DL, MVT::i64));
       SDValue Extract2 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, VT, Other,
                                      DAG.getConstant(1, DL, MVT::i64));
       if (!IsStrict)
         return DAG.getNode(N0->getOpcode(), DL, VT, Extract1, Extract2);
 
static SDValue getPTest(SelectionDAG &DAG, EVT VT, SDValue Pg, SDValue Op,
 
   // Convert CC to integer based on requested condition.
   // NOTE: Cond is inverted to promote CSEL's removal when it feeds a compare.
   SDValue CC = DAG.getConstant(getInvertedCondCode(Cond), DL, MVT::i32);
   SDValue Res = DAG.getNode(AArch64ISD::CSEL, DL, OutVT, FVal, TVal, CC, Test);
   return DAG.getZExtOrTrunc(Res, DL, VT);
 }
 
+// Materialize : i1 = extract_vector_elt t37, Constant:i64<0>
+// ... into: "ptrue p, all" + PTEST
+static SDValue
+performFirstTrueTestVectorCombine(SDNode *N,
+                                  TargetLowering::DAGCombinerInfo &DCI,
+                                  const AArch64Subtarget *Subtarget) {
+  assert(N->getOpcode() == ISD::EXTRACT_VECTOR_ELT);
+  SelectionDAG &DAG = DCI.DAG;
+
+  if (!Subtarget->hasSVE() || !DCI.isBeforeLegalize())
+    return SDValue();
+
+  SDValue N0 = N->getOperand(0);
+  EVT VT = N0.getValueType();
+
+  if (!VT.isScalableVector() || VT.getVectorElementType() != MVT::i1)
+    return SDValue();
+
+  auto *Idx = dyn_cast<ConstantSDNode>(N->getOperand(1));
+  if (!Idx || Idx->getZExtValue() != 0)
+    return SDValue();
+
+  // Extracts of lane 0 for SVE can be expressed as PTEST(Op, FIRST) ? 1 : 0
+  SDValue Pg = getPTrue(DAG, SDLoc(N), VT, AArch64SVEPredPattern::all);
+  return getPTest(DAG, N->getValueType(0), Pg, N0, AArch64CC::FIRST_ACTIVE);
+}
+
 static SDValue combineSVEReductionInt(SDNode *N, unsigned Opc,
                                       SelectionDAG &DAG) {
   SDLoc DL(N);
 
   SDValue Pred = N->getOperand(1);
   SDValue VecToReduce = N->getOperand(2);
 
   // NOTE: The integer reduction's result type is not always linked to the
case AArch64ISD::VLSHR:
     return performVectorShiftCombine(N, *this, DCI);
   case AArch64ISD::SUNPKLO:
     return performSunpkloCombine(N, DAG);
   case AArch64ISD::BSP:
     return performBSPExpandForSVE(
         N, DAG, Subtarget, useSVEForFixedLengthVectorVT(N->getValueType(0)));
   case ISD::INSERT_VECTOR_ELT:
     return performInsertVectorEltCombine(N, DCI);
-  case ISD::EXTRACT_VECTOR_ELT:
-    return performExtractVectorEltCombine(N, DAG);
+  case ISD::EXTRACT_VECTOR_ELT: {
+    SDValue Res = performFirstTrueTestVectorCombine(N, DCI, Subtarget);
+    if (Res == SDValue())
+      Res = performExtractVectorEltCombine(N, DAG);
+    return Res;
+  }
   case ISD::VECREDUCE_ADD:
     return performVecReduceAddCombine(N, DCI.DAG, Subtarget);
   case AArch64ISD::UADDV:
     return performUADDVCombine(N, DAG);
   case AArch64ISD::SMULL:
   case AArch64ISD::UMULL:
     return tryCombineLongOpWithDup(Intrinsic::not_intrinsic, N, DCI, DAG);
   case ISD::INTRINSIC_VOID:

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

; CHECK-NEXT:    ret
   %b = shufflevector <vscale x 2 x i64> %a, <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer
   %c = extractelement <vscale x 2 x i64> %b, i32 %y
   ret i64 %c
 }
 
 define i1 @test_lane0_16xi1(<vscale x 16 x i1> %a) #0 {
 ; CHECK-LABEL: test_lane0_16xi1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov z0.b, p0/z, #1 // =0x1
-; CHECK-NEXT:    fmov w8, s0
-; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ptrue p1.b
+; CHECK-NEXT:    ptest p1, p0.b
+; CHECK-NEXT:    cset w0, mi
 ; CHECK-NEXT:    ret
   %b = extractelement <vscale x 16 x i1> %a, i32 0
   ret i1 %b
 }
 
 define i1 @test_lane9_8xi1(<vscale x 8 x i1> %a) #0 {
 ; CHECK-LABEL: test_lane9_8xi1:
 ; CHECK:       // %bb.0:

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

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64 -o - %s | FileCheck  %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+define void @add_arrays(i32* nocapture %a, <vscale x 4 x i32> %b, <vscale x 4 x i1> %p0, <vscale x 4 x i64> %step) #0 {
+; CHECK-LABEL: add_arrays:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    cntw x9
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    ptrue p2.s
+; CHECK-NEXT:  .LBB0_1: // %vector.body
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    ld1w { z3.s }, p0/z, [x0, x8, lsl #2]
+; CHECK-NEXT:    add z3.s, z3.s, z0.s
+; CHECK-NEXT:    st1w { z3.s }, p0, [x0, x8, lsl #2]
+; CHECK-NEXT:    add x8, x8, x9
+; CHECK-NEXT:    mov z3.d, x8
+; CHECK-NEXT:    add z4.d, z3.d, z1.d
+; CHECK-NEXT:    add z3.d, z3.d, z2.d
+; CHECK-NEXT:    cmphi p0.d, p1/z, z2.d, z3.d
+; CHECK-NEXT:    cmphi p3.d, p1/z, z1.d, z4.d
+; CHECK-NEXT:    uzp1 p0.s, p3.s, p0.s
+; CHECK-NEXT:    ptest p2, p0.b
+; CHECK-NEXT:    b.mi .LBB0_1
+; CHECK-NEXT:  // %bb.2: // %cleanup
+; CHECK-NEXT:    ret
+entry:
+  br label %vector.body
+
+vector.body:                               ; preds = %vector.body, %entry
+  %index = phi i64 [ 0, %entry ], [ %index.next, %vector.body ]
+  %predicate = phi <vscale x 4 x i1> [ %p0, %entry ], [ %predicate.next, %vector.body ]
+  %addr = getelementptr inbounds i32, i32* %a, i64 %index
+  %addload = bitcast i32* %addr to <vscale x 4 x i32>*
+  %wide.masked.load = call <vscale x 4 x i32> @llvm.masked.load.nxv4i32.p0nxv4i32(<vscale x 4 x i32>* %addload, i32 4, <vscale x 4 x i1> %predicate, <vscale x 4 x i32> undef)
+  %c = add <vscale x 4 x i32> %wide.masked.load, %b
+  %addstore = bitcast i32* %addr to <vscale x 4 x i32>*
+  call void @llvm.masked.store.nxv4i32.p0nxv4i32(<vscale x 4 x i32> %c, <vscale x 4 x i32>* %addstore, i32 4, <vscale x 4 x i1> %predicate)
+  %cond = icmp ult i64 %index, 4294967294
+  call void @llvm.assume(i1 %cond)
+  %0 = call i64 @llvm.vscale.i64()
+  %shl2 = shl nuw nsw i64 %0, 2
+  %index.next = add nuw nsw i64 %index, %shl2
+  %splatinsert = insertelement <vscale x 4 x i64> poison, i64 %index.next, i64 0
+  %splat = shufflevector <vscale x 4 x i64> %splatinsert, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %1 = add nuw <vscale x 4 x i64> %splat, %step
+  %predicate.next = icmp ult <vscale x 4 x i64> %1, %step
+  %test0 = extractelement <vscale x 4 x i1> %predicate.next, i64 0
+  br i1 %test0, label %vector.body, label %cleanup
+
+cleanup:                                 ; preds = %vector.body, %entry
+  ret void
+}
+
+declare void @llvm.assume(i1 noundef)
+declare i64 @llvm.vscale.i64()
+declare <vscale x 4 x i32> @llvm.masked.load.nxv4i32.p0nxv4i32(<vscale x 4 x i32>*, i32 immarg, <vscale x 4 x i1>, <vscale x 4 x i32>)
+declare void @llvm.masked.store.nxv4i32.p0nxv4i32(<vscale x 4 x i32>, <vscale x 4 x i32>*, i32 immarg, <vscale x 4 x i1>)
+
+attributes #0 = { nofree norecurse nosync nounwind uwtable vscale_range(1,16) "target-cpu"="generic" "target-features"="+neon,+sve,+v8a" "unsafe-fp-math"="true" }
+