[SVE][AArch64] Enable first active true vector combine for INTRINSIC_WO_CHAIN

[Allen]

Summary

WHILELO/LS insn is used very important for SVE loop, and itself
is a flag-setting operation, so add it.

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14551

Hi @Allen, don't you also have to check the first operand to see if it matches the intrinsic, i.e. aarch64_sve_whilelo or aarch64_sve_whilels?

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

102

Since you also want to support the DAG combine with whilels is it worth having a test for that too?

[Allen]

update to check the aarch64_sve_whilelo or aarch64_sve_whilels

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14551

I guess I should have asked this a couple of patches ago but being provocative what's the value in having any restriction here? I recall something about allowing hoisting from loops but is this really a necessary consideration at this phase of compilation? Are there other things to consider?

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14551

So I think the MachineLICM pass runs after this phase, right? And if you introduce flag-setting operations they don't get hoisted out. Obviously if you had flag-setting operations to start with then it's no worse. This is what happens when you dup a i1 into a predicate register, for example see LowerSPLAT_VECTOR where we create a whilelo that then never gets hoisted out of the loop.

I'm happy to be proved wrong though as I may have missed something!

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14551

Sure, that makes sense. I guess what I'm asking is what's the scenario where we are extracting from a loop invariant predicate vector that hasn't already been hoisted at the IR level.

[paulwalker-arm]

Whether it's worth having the CC restriction aside, the patch look good to me. I've mentioned some refactoring to consider.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14539–14540

This name is less meaningful after this change.

14549–14552

Perhaps it's worth extracting this into a isPredicateCCSettingOp helper function (or lambda) as eventually we'll want to add all such operations if we're going to persist with this restriction.

[Allen]

update according review

1. define a new helper function isPredicateCCSettingOp
2. add all other type whileXX operations
3. add related cases, and sve->sve2

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static bool hasPairwiseAdd(unsigned Opcode, EVT VT, bool FullFP16) {
   default:
     return false;
   }
 }
 
 static SDValue getPTest(SelectionDAG &DAG, EVT VT, SDValue Pg, SDValue Op,
                         AArch64CC::CondCode Cond);
 
+static bool isPredicateCCSettingOp(SDValue N) {
+  if ((N.getOpcode() == ISD::SETCC) ||
+      (N.getOpcode() == ISD::INTRINSIC_WO_CHAIN &&
+       (N.getConstantOperandVal(0) == Intrinsic::aarch64_sve_whilege ||
+        N.getConstantOperandVal(0) == Intrinsic::aarch64_sve_whilegt ||
+        N.getConstantOperandVal(0) == Intrinsic::aarch64_sve_whilehi ||
+        N.getConstantOperandVal(0) == Intrinsic::aarch64_sve_whilehs ||
+        N.getConstantOperandVal(0) == Intrinsic::aarch64_sve_whilele ||
+        N.getConstantOperandVal(0) == Intrinsic::aarch64_sve_whilelo ||
+        N.getConstantOperandVal(0) == Intrinsic::aarch64_sve_whilels ||
+        N.getConstantOperandVal(0) == Intrinsic::aarch64_sve_whilelt)))
+    return true;
+
+  return false;
+}
+
 // 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();
+  SDValue N0 = N->getOperand(0);
+  EVT VT = N0.getValueType();
 
-  if (!VT.isScalableVector() || VT.getVectorElementType() != MVT::i1)
+  if (!VT.isScalableVector() || VT.getVectorElementType() != MVT::i1 ||
+      !isNullConstant(N->getOperand(1)))
     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)
+  // flag-setting operation.
+  if (!isPredicateCCSettingOp(N0))
     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);
+  return getPTest(DAG, N->getValueType(0), Pg, N0, AArch64CC::FIRST_ACTIVE);
 }
 
 // Materialize : Idx = (add (mul vscale, NumEls), -1)
 //               i1 = extract_vector_elt t37, Constant:i64<Idx>
 //     ... into: "ptrue p, all" + PTEST
 static SDValue
 performLastTrueTestVectorCombine(SDNode *N,
                                  TargetLowering::DAGCombinerInfo &DCI,

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

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple=aarch64-linux-unknown -mattr=+sve -o - < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-linux-unknown -mattr=+sve2 -o - < %s | FileCheck %s
 
 define <vscale x 8 x i1> @not_icmp_sle_nxv8i16(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b) {
 ; CHECK-LABEL: not_icmp_sle_nxv8i16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    cmpgt p0.h, p0/z, z0.h, z1.h
 ; CHECK-NEXT:    ret
   %icmp = icmp sle <vscale x 8 x i16> %a, %b
; CHECK-NEXT:    ret
   %vcond = fcmp oeq <vscale x 4 x float> %a, %b
   %vscale = call i64 @llvm.vscale.i64()
   %shl2 = shl nuw nsw i64 %vscale, 2
   %idx = add nuw nsw i64 %shl2, -1
   %bit = extractelement <vscale x 4 x i1> %vcond, i64 %idx
   ret i1 %bit
 }
 
+define i1 @whilege_first(i64 %next, i64 %end) {
+; CHECK-LABEL: whilege_first:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilege p0.s, x0, x1
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %predicate = call <vscale x 4 x i1> @llvm.aarch64.sve.whilege.nxv4i1.i64(i64 %next, i64 %end)
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
+
+define i1 @whilegt_first(i64 %next, i64 %end) {
+; CHECK-LABEL: whilegt_first:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilegt p0.s, x0, x1
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %predicate = call <vscale x 4 x i1> @llvm.aarch64.sve.whilegt.nxv4i1.i64(i64 %next, i64 %end)
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
+
+define i1 @whilehi_first(i64 %next, i64 %end) {
+; CHECK-LABEL: whilehi_first:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilehi p0.s, x0, x1
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %predicate = call <vscale x 4 x i1> @llvm.aarch64.sve.whilehi.nxv4i1.i64(i64 %next, i64 %end)
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
+
+define i1 @whilehs_first(i64 %next, i64 %end) {
+; CHECK-LABEL: whilehs_first:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilehs p0.s, x0, x1
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %predicate = call <vscale x 4 x i1> @llvm.aarch64.sve.whilehs.nxv4i1.i64(i64 %next, i64 %end)
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
+
+define i1 @whilele_first(i64 %next, i64 %end) {
+; CHECK-LABEL: whilele_first:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilele p0.s, x0, x1
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %predicate = call <vscale x 4 x i1> @llvm.aarch64.sve.whilele.nxv4i1.i64(i64 %next, i64 %end)
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
+
+define i1 @whilelo_first(i64 %next, i64 %end) {
+; CHECK-LABEL: whilelo_first:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilelo p0.s, x0, x1
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %predicate = call <vscale x 4 x i1> @llvm.aarch64.sve.whilelo.nxv4i1.i64(i64 %next, i64 %end)
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
+
+define i1 @whilels_first(i64 %next, i64 %end) {
+; CHECK-LABEL: whilels_first:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilels p0.s, x0, x1
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %predicate = call <vscale x 4 x i1> @llvm.aarch64.sve.whilels.nxv4i1.i64(i64 %next, i64 %end)
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
+
+define i1 @whilelt_first(i64 %next, i64 %end) {
+; CHECK-LABEL: whilelt_first:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilelt p0.s, x0, x1
+; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    ret
+  %predicate = call <vscale x 4 x i1> @llvm.aarch64.sve.whilelt.nxv4i1.i64(i64 %next, i64 %end)
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
 
 declare i64 @llvm.vscale.i64()
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilege.nxv4i1.i64(i64, i64)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilegt.nxv4i1.i64(i64, i64)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilehi.nxv4i1.i64(i64, i64)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilehs.nxv4i1.i64(i64, i64)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilele.nxv4i1.i64(i64, i64)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilelo.nxv4i1.i64(i64, i64)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilels.nxv4i1.i64(i64, i64)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilelt.nxv4i1.i64(i64, i64)