[DAGCombiner][AArch64] Enhance to fold CSNEG into CSINC instruction

[Allen]

Summary

Perform the scalar expression combine in the form of:

CSNEG(1, c, cc) + b  =>  cc  ? b+1 : b-c => CSINC(b-c, b, !cc)
CSNEG(c, -1, cc) + b =>  cc  ? b+c : b+1 => CSINC(b+c, b, cc)

Fix https://github.com/llvm/llvm-project/issues/53071

[Allen]

@dmgreen can you help to give some advice, thanks!

[dmgreen]

The code is very similar to the performAddCSelIntoCSinc method. Can they share the same code, with the differences for CSNeg accounted for?

[Allen]

In D119105#3304447, @dmgreen wrote:

The code is very similar to the performAddCSelIntoCSinc method. Can they share the same code, with the differences for CSNeg accounted for?

Thanks for you comment again! I has tried it before, it seems much complicated, I'll glad to update if you think it's better to share ?

  // The CSEL should include a const one operand, and the CSNEG should include
  // One or NegOne operand.
  ConstantSDNode *CTVal = dyn_cast<ConstantSDNode>(LHS.getOperand(0));
  ConstantSDNode *CFVal = dyn_cast<ConstantSDNode>(LHS.getOperand(1));
  if (!CTVal || !CFVal)
    return SDValue();

  if (!(LHS.getOpcode() == AArch64ISD::CSEL &&
        (CTVal->isOne() || CFVal->isOne())) ||
      !(LHS.getOpcode() == AArch64ISD::CSNEG &&
        (CTVal->isOne() || CFVal->isAllOnes())))
    return SDValue();

  // Switch CSEL(1, c, cc) to CSEL(c, 1, !cc)
  if (LHS.getOpcode() == AArch64ISD::CSEL &&
      CTVal->isOne() && !CFVal->isOne()) {
    std::swap(CTVal, CFVal);
    AArch64CC = AArch64CC::getInvertedCondCode(AArch64CC);
  }

  SDLoc DL(N);
  // Switch CSNEG(1, c, cc) to CSNEG(-c, -1, !cc)
  if (LHS.getOpcode() == AArch64ISD::CSNEG &&
    CTVal->isOne() && !CFVal->isAllOnes()) {
    int64_t C = -1 * CFVal->getSExtValue();
    CTVal = cast<ConstantSDNode>(DAG.getConstant(C, DL, VT));
    CFVal = cast<ConstantSDNode>(DAG.getConstant(-1, DL, VT));
    AArch64CC = AArch64CC::getInvertedCondCode(AArch64CC);
  }

...

  assert((LHS.getOpcode() == AArch64ISD::CSEL && CFVal->isOne() ||
          LHS.getOpcode() == AArch64ISD::CSNEG && CFVal->isAllOnes()) &&
         "Unexpected constant value");

[dmgreen]

Hmm. I feel like it shares more code than it needs special cases? Maybe try and give it a go, and we can see if it looks too bad.

(I had in the past had the idea of removing AArch64ISD::CSNEG/CSINC/CSINV - they are mostly used only for constants as far as I understand and could just as well use CSEL directly. Either that or make sure of them more, canonicalizing more patterns to use them. I don't know what the tradeoffs of those are though. Less target dependant nodes can be a good thing, but more aggressively targeting the nodes we have as opposed to leaving it to chance may be good too. That's just an aside though).

[Allen]

share

[Allen]

fix pre-merge format checks

[dmgreen]

Thanks. I think something might not be right in the new logic.
(add y, (csneg 1, c, eq)) is eq ? y+1 : y-c. Which should be (csinc (sub y, c), y, ne)
(add y, (csneg c, -1, eq)) is eq ? y+c : y+1. Which should be (csinc (add y, c), y, eq)
I may be wrong with all these double negatives though, correct me if that looks wrong. I would expect the assembly to be sub w8, w1, #1; csinc w0, w8, w1

[Allen]

Thanks , you are right, good catch. https://godbolt.org/z/nzvanY3GT

[dmgreen]

Thanks, LGTM with a couple of suggestions.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14882

Probably better to avoid the uint64_t and use APInt directly.

14883

This can use DAG.getAllOnesConstant(..)

llvm/test/CodeGen/AArch64/aarch64-isel-csinc.ll

132

Can you add i64 variants of these tests, maybe with different conditions? You can also remove the dso_local.

[Allen]

update according review

[Allen]

add new test case of type i64

[Allen]

rebase to new head

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (CTVal->isOne() && !CFVal->isOne()) {
     AArch64CC = AArch64CC::getInvertedCondCode(AArch64CC);
   }
 
   // It might be neutral for larger constants, as the immediate need to be
   // materialized in a register.
   APInt ADDC = CTVal->getAPIntValue();
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   if (!TLI.isLegalAddImmediate(ADDC.getSExtValue()))
     return SDValue();
 
   assert(CFVal->isOne() && "Unexpected constant value");
 
   SDLoc DL(N);
   SDValue NewNode = DAG.getNode(ISD::ADD, DL, VT, RHS, SDValue(CTVal, 0));
   SDValue CCVal = DAG.getConstant(AArch64CC, DL, MVT::i32);
   SDValue Cmp = CSel.getOperand(3);
 
   return DAG.getNode(AArch64ISD::CSINC, DL, VT, NewNode, RHS, CCVal, Cmp);
 }
 
+/// Perform the scalar expression combine in the form of:
+///   CSNEG(1, c, cc) + b => CSINC(b-c, b, !cc)
+///   CSNEG(c, -1, cc) + b => CSINC(b+c, b, cc)
+static SDValue performAddCSnegIntoCSinc(SDNode *N, SelectionDAG &DAG) {
+  EVT VT = N->getValueType(0);
+  if (!VT.isScalarInteger() || N->getOpcode() != ISD::ADD)
+    return SDValue();
+
+  SDValue CSneg = N->getOperand(0);
+  SDValue RHS = N->getOperand(1);
+
+  // Handle commutivity.
+  if (CSneg.getOpcode() != AArch64ISD::CSNEG) {
+    std::swap(CSneg, RHS);
+    if (CSneg.getOpcode() != AArch64ISD::CSNEG) {
+      return SDValue();
+    }
+  }
+
+  if (!CSneg.hasOneUse())
+    return SDValue();
+
+  AArch64CC::CondCode AArch64CC =
+      static_cast<AArch64CC::CondCode>(CSneg.getConstantOperandVal(2));
+
+  // The CSNEG should include One or NegOne operand.
+  ConstantSDNode *CTVal = dyn_cast<ConstantSDNode>(CSneg.getOperand(0));
+  ConstantSDNode *CFVal = dyn_cast<ConstantSDNode>(CSneg.getOperand(1));
+  if (!CTVal || !CFVal || (!CTVal->isOne() && !CFVal->isAllOnes()))
+    return SDValue();
+
+  SDLoc DL(N);
+  // Switch CSNEG(1, c, cc)  to CSNEG(-c, -1, !cc)
+  if (CTVal->isOne() && !CFVal->isAllOnes()) {
+    int64_t C = -1 * CFVal->getSExtValue();
+    CTVal = cast<ConstantSDNode>(DAG.getConstant(C, DL, VT));
+    CFVal = cast<ConstantSDNode>(DAG.getConstant(-1, DL, VT));
+    AArch64CC = AArch64CC::getInvertedCondCode(AArch64CC);
+  }
+
+  // It might be neutral for larger constants, as the immediate need to be
+  // materialized in a register.
+  APInt ADDC = CTVal->getAPIntValue();
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+  if (!TLI.isLegalAddImmediate(ADDC.getSExtValue()))
+    return SDValue();
+
+  assert(CFVal->isAllOnes() && "Unexpected constant value");
+
+  SDValue NewNode = DAG.getNode(ISD::SUB, DL, VT, RHS, SDValue(CTVal, 0));
+  SDValue CCVal = DAG.getConstant(AArch64CC, DL, MVT::i32);
+  SDValue Cmp = CSneg.getOperand(3);
+
+  return DAG.getNode(AArch64ISD::CSINC, DL, VT, NewNode, RHS, CCVal, Cmp);
+}
+
 // ADD(UDOT(zero, x, y), A) -->  UDOT(A, x, y)
 static SDValue performAddDotCombine(SDNode *N, SelectionDAG &DAG) {
   EVT VT = N->getValueType(0);
   if (N->getOpcode() != ISD::ADD)
     return SDValue();
 
   SDValue Dot = N->getOperand(0);
   SDValue A = N->getOperand(1);
static SDValue performAddSubCombine(SDNode *N,
                                     SelectionDAG &DAG) {
   // Try to change sum of two reductions.
   if (SDValue Val = performAddUADDVCombine(N, DAG))
     return Val;
   if (SDValue Val = performAddDotCombine(N, DAG))
     return Val;
   if (SDValue Val = performAddCSelIntoCSinc(N, DAG))
     return Val;
+  if (SDValue Val = performAddCSnegIntoCSinc(N, DAG))
+    return Val;
 
   return performAddSubLongCombine(N, DCI, DAG);
 }
 
 // Massage DAGs which we can use the high-half "long" operations on into
 // something isel will recognize better. E.g.
 //
 // (aarch64_neon_umull (extract_high vec) (dupv64 scalar)) -->

llvm/test/CodeGen/AArch64/aarch64-isel-csinc.ll

 ; CHECK-NEXT:    add w0, w8, w1
 ; CHECK-NEXT:    ret
 entry:
   %tobool.not = icmp eq i32 %a, 0
   %cond.v = select i1 %tobool.not, i32 -4097, i32 1
   %cond = add nsw i32 %cond.v, %b
   ret i32 %cond
 }
+
+; int csinc8 (int a, int b) { return a ? b-1 : b+1; }
+define dso_local i32 @csinc8(i32 %a, i32 %b) {
+; CHECK-LABEL: csinc8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    cmp w0, #0
+; CHECK-NEXT:    add w8, w1, #1
+; CHECK-NEXT:    csinc w0, w8, w1, ne
+; CHECK-NEXT:    ret
+entry:
+  %tobool.not = icmp eq i32 %a, 0
+  %cond.v = select i1 %tobool.not, i32 1, i32 -1
+  %cond = add nsw i32 %cond.v, %b
+  ret i32 %cond
+}
+
+; int csinc9 (int a, int b) { return a ? b+1 : b-1; }
+define dso_local i32 @csinc9(i32 %a, i32 %b) {
+; CHECK-LABEL: csinc9:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    cmp w0, #0
+; CHECK-NEXT:    add w8, w1, #1
+; CHECK-NEXT:    csinc w0, w8, w1, eq
+; CHECK-NEXT:    ret
+entry:
+  %tobool.not = icmp eq i32 %a, 0
+  %cond.v = select i1 %tobool.not, i32 -1, i32 1
+  %cond = add nsw i32 %cond.v, %b
+  ret i32 %cond
+}