[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 ↗

(On Diff #412670)

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 ↗

(On Diff #412670)

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

51 ↗

(On Diff #412670)

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

15368

You could use Idx->isZero() here.

18317

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

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

14294

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.

14318

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

14319

Unwanted space?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case ISD::FADD:
     return (FullFP16 && VT == MVT::f16) || VT == MVT::f32 || VT == MVT::f64;
   case ISD::ADD:
     return VT == MVT::i64;
   default:
     return false;
   }
 }
 
-static SDValue performExtractVectorEltCombine(SDNode *N, SelectionDAG &DAG) {
+static SDValue getPTest(SelectionDAG &DAG, EVT VT, SDValue Pg, SDValue Op,
+                        AArch64CC::CondCode Cond);
+
+// 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);
+  // Make sure PTEST can be legalised with illegal types.
+  if (!Subtarget->hasSVE() || DCI.isBeforeLegalize())
+    return SDValue();
+
+  SDValue SetCC = N->getOperand(0);
+  EVT VT = SetCC.getValueType();
+
+  if (!VT.isScalableVector() || VT.getVectorElementType() != MVT::i1)
+    return SDValue();
+
+  // Restricted the DAG combine to only cases where we're extracting from a
+  // flag-setting operation
+  auto *Idx = dyn_cast<ConstantSDNode>(N->getOperand(1));
+  if (!Idx || !Idx->isZero() || SetCC.getOpcode() != ISD::SETCC)
+    return SDValue();
+
+  // Extracts of lane 0 for SVE can be expressed as PTEST(Op, FIRST) ? 1 : 0
+  SelectionDAG &DAG = DCI.DAG;
+  SDValue Pg = getPTrue(DAG, SDLoc(N), VT, AArch64SVEPredPattern::all);
+  return getPTest(DAG, N->getValueType(0), Pg, SetCC, AArch64CC::FIRST_ACTIVE);
+}
+
+static SDValue
+performExtractVectorEltCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
+                               const AArch64Subtarget *Subtarget) {
+  assert(N->getOpcode() == ISD::EXTRACT_VECTOR_ELT);
+  if (SDValue Res = performFirstTrueTestVectorCombine(N, DCI, Subtarget))
+    return Res;
+
+  SelectionDAG &DAG = DCI.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();
 
static SDValue combineSVEReductionInt(SDNode *N, unsigned Opc,
   // containing the reduction result, which we'll now extract.
   SDValue Zero = DAG.getConstant(0, DL, MVT::i64);
   return DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, N->getValueType(0), Reduce,
                      Zero);
 }
 
 static SDValue combineSVEReductionFP(SDNode *N, unsigned Opc,
                                      SelectionDAG &DAG) {
   SDLoc DL(N);
 
   SDValue Pred = N->getOperand(1);
   SDValue VecToReduce = N->getOperand(2);
 
   EVT ReduceVT = VecToReduce.getValueType();
   SDValue Reduce = DAG.getNode(Opc, DL, ReduceVT, Pred, VecToReduce);
 
   // SVE reductions set the whole vector register with the first element
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   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);
+    return performExtractVectorEltCombine(N, DCI, Subtarget);
   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:
   case ISD::INTRINSIC_W_CHAIN:

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

 ; CHECK-NEXT:    fcmgt p0.s, p0/z, z1.s, z0.s
 ; CHECK-NEXT:    ret
   %icmp = fcmp uge <vscale x 4 x float> %a, %b
   %tmp = insertelement <vscale x 4 x i1> undef, i1 true, i32 0
   %ones = shufflevector <vscale x 4 x i1> %tmp, <vscale x 4 x i1> undef, <vscale x 4 x i32> zeroinitializer
   %not = xor <vscale x 4 x i1> %icmp, %ones
   ret <vscale x 4 x i1> %not
 }
+
+define i1 @foo(<vscale x 4 x float> %a, <vscale x 4 x float> %b) {
+; CHECK-LABEL: foo:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    fcmeq p1.s, p0/z, z0.s, z1.s
+; CHECK-NEXT:    ptest p0, p1.b
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %vcond = fcmp oeq <vscale x 4 x float> %a, %b
+  %bit = extractelement <vscale x 4 x i1> %vcond, i64 0
+  ret i1 %bit
+}
+