[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

14805

Probably better to avoid the uint64_t and use APInt directly.

14806

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

static SDValue performAddUADDVCombine(SDNode *N, SelectionDAG &DAG) {
   SDLoc DL(N);
   SDValue AddVal = DAG.getNode(ISD::ADD, DL, ValVT, Val1, Val2);
   return DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, VT,
                      DAG.getNode(AArch64ISD::UADDV, DL, ValVT, AddVal),
                      DAG.getConstant(0, DL, MVT::i64));
 }
 
 /// Perform the scalar expression combine in the form of:
-///   CSEL (c, 1, cc) + b => CSINC(b+c, b, cc)
+///   CSEL(c, 1, cc) + b => CSINC(b+c, b, cc)
+///   CSNEG(c, -1, cc) + b => CSINC(b+c, b, cc)
 static SDValue performAddCSelIntoCSinc(SDNode *N, SelectionDAG &DAG) {
   EVT VT = N->getValueType(0);
   if (!VT.isScalarInteger() || N->getOpcode() != ISD::ADD)
     return SDValue();
 
-  SDValue CSel = N->getOperand(0);
+  SDValue LHS = N->getOperand(0);
   SDValue RHS = N->getOperand(1);
 
   // Handle commutivity.
-  if (CSel.getOpcode() != AArch64ISD::CSEL) {
-    std::swap(CSel, RHS);
-    if (CSel.getOpcode() != AArch64ISD::CSEL) {
+  if (LHS.getOpcode() != AArch64ISD::CSEL &&
+      LHS.getOpcode() != AArch64ISD::CSNEG) {
+    std::swap(LHS, RHS);
+    if (LHS.getOpcode() != AArch64ISD::CSEL &&
+        LHS.getOpcode() != AArch64ISD::CSNEG) {
       return SDValue();
     }
   }
 
-  if (!CSel.hasOneUse())
+  if (!LHS.hasOneUse())
     return SDValue();
 
   AArch64CC::CondCode AArch64CC =
-      static_cast<AArch64CC::CondCode>(CSel.getConstantOperandVal(2));
+      static_cast<AArch64CC::CondCode>(LHS.getConstantOperandVal(2));
 
-  // The CSEL should include a const one operand.
-  ConstantSDNode *CTVal = dyn_cast<ConstantSDNode>(CSel.getOperand(0));
-  ConstantSDNode *CFVal = dyn_cast<ConstantSDNode>(CSel.getOperand(1));
-  if (!CTVal || !CFVal || (!CTVal->isOne() && !CFVal->isOne()))
+  // 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();
 
-  // switch CSEL (1, c, cc)  to CSEL (c, 1, !cc)
-  if (CTVal->isOne() && !CFVal->isOne()) {
+  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()) {
+    APInt C = -1 * CFVal->getAPIntValue();
+    CTVal = cast<ConstantSDNode>(DAG.getConstant(C, DL, VT));
+    CFVal = cast<ConstantSDNode>(DAG.getAllOnesConstant(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->isOne() && "Unexpected constant value");
+  assert(((LHS.getOpcode() == AArch64ISD::CSEL && CFVal->isOne()) ||
+          (LHS.getOpcode() == AArch64ISD::CSNEG && CFVal->isAllOnes())) &&
+         "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);
+  SDValue Cmp = LHS.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)

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

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-- -o - < %s | FileCheck %s
 
 ; Verify that we can fold csneg/csel into csinc instruction.
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-unknown-linux-gnu"
 
 ; char csinc1 (char a, char b) { return !a ? b+1 : b+3; }
-define dso_local i8 @csinc1(i8 %a, i8 %b) local_unnamed_addr #0 {
+define i8 @csinc1(i8 %a, i8 %b) local_unnamed_addr #0 {
 ; CHECK-LABEL: csinc1:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    tst w0, #0xff
 ; CHECK-NEXT:    add w8, w1, #3
 ; CHECK-NEXT:    csinc w0, w8, w1, ne
 ; CHECK-NEXT:    ret
 entry:
   %tobool.not = icmp eq i8 %a, 0
   %cond.v = select i1 %tobool.not, i8 1, i8 3
   %cond = add i8 %cond.v, %b
   ret i8 %cond
 }
 
 ; short csinc2 (short a, short b) { return !a ? b+1 : b+3; }
-define dso_local i16 @csinc2(i16 %a, i16 %b) local_unnamed_addr #0 {
+define i16 @csinc2(i16 %a, i16 %b) local_unnamed_addr #0 {
 ; CHECK-LABEL: csinc2:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    tst w0, #0xffff
 ; CHECK-NEXT:    add w8, w1, #3
 ; CHECK-NEXT:    csinc w0, w8, w1, ne
 ; CHECK-NEXT:    ret
 entry:
   %tobool.not = icmp eq i16 %a, 0
   %cond.v = select i1 %tobool.not, i16 1, i16 3
   %cond = add i16 %cond.v, %b
   ret i16 %cond
 }
 
 ; int csinc3 (int a, int b) { return !a ? b+1 : b+3; }
-define dso_local i32 @csinc3(i32 %a, i32 %b) local_unnamed_addr #0 {
+define i32 @csinc3(i32 %a, i32 %b) local_unnamed_addr #0 {
 ; CHECK-LABEL: csinc3:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    cmp w0, #0
 ; CHECK-NEXT:    add w8, w1, #3
 ; 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 3
   %cond = add nsw i32 %cond.v, %b
   ret i32 %cond
 }
 
 ; long long csinc4 (long long a, long long b) { return !a ? b+1 : b+3; }
-define dso_local i64 @csinc4(i64 %a, i64 %b) local_unnamed_addr #0 {
+define i64 @csinc4(i64 %a, i64 %b) local_unnamed_addr #0 {
 ; CHECK-LABEL: csinc4:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    cmp x0, #0
 ; CHECK-NEXT:    add x8, x1, #3
 ; CHECK-NEXT:    csinc x0, x8, x1, ne
 ; CHECK-NEXT:    ret
 entry:
   %tobool.not = icmp eq i64 %a, 0
   %cond.v = select i1 %tobool.not, i64 1, i64 3
   %cond = add nsw i64 %cond.v, %b
   ret i64 %cond
 }
+
+; long long csinc8 (long long a, long long b) { return a ? b-1 : b+1; }
+define i64 @csinc8(i64 %a, i64 %b) {
+; CHECK-LABEL: csinc8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    sub x8, x1, #1
+; CHECK-NEXT:    cmp x0, #0
+; CHECK-NEXT:    csinc x0, x8, x1, ne
+; CHECK-NEXT:    ret
+entry:
+  %tobool.not = icmp eq i64 %a, 0
+  %cond.v = select i1 %tobool.not, i64 1, i64 -1
+  %cond = add nsw i64 %cond.v, %b
+  ret i64 %cond
+}
+
+; long long csinc9 (long long a, long long b) { return a ? b+1 : b-1; }
+define i64 @csinc9(i64 %a, i64 %b) {
+; CHECK-LABEL: csinc9:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    sub x8, x1, #1
+; CHECK-NEXT:    cmp x0, #0
+; CHECK-NEXT:    csinc x0, x8, x1, eq
+; CHECK-NEXT:    ret
+entry:
+  %tobool.not = icmp eq i64 %a, 0
+  %cond.v = select i1 %tobool.not, i64 -1, i64 1
+  %cond = add nsw i64 %cond.v, %b
+  ret i64 %cond
+}

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

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-- -o - < %s | FileCheck %s
 
 ; Verify that we can fold csneg/csel into csinc instruction.
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-unknown-linux-gnu"
 
 ; int csinc1 (int a, int b) { return !a ? b+3 : b+1; }
-define dso_local i32 @csinc1(i32 %a, i32 %b) {
+define i32 @csinc1(i32 %a, i32 %b) {
 ; CHECK-LABEL: csinc1:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    cmp w0, #0
 ; CHECK-NEXT:    add w8, w1, #3
 ; 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 3, i32 1
   %cond = add nsw i32 %cond.v, %b
   ret i32 %cond
 }
 
 ; int csinc2 (int a, int b) { return a ? b+3 : b+1; }
-define dso_local i32 @csinc2(i32 %a, i32 %b) {
+define i32 @csinc2(i32 %a, i32 %b) {
 ; CHECK-LABEL: csinc2:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    cmp w0, #0
 ; CHECK-NEXT:    add w8, w1, #3
 ; 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 3
   %cond = add nsw i32 %cond.v, %b
   ret i32 %cond
 }
 
 ; int csinc3 (int a, int b) { return !a ? b+1 : b-3; }
-define dso_local i32 @csinc3(i32 %a, i32 %b) {
+define i32 @csinc3(i32 %a, i32 %b) {
 ; CHECK-LABEL: csinc3:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sub w8, w1, #3
 ; CHECK-NEXT:    cmp w0, #0
 ; 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 -3
   %cond = add nsw i32 %cond.v, %b
   ret i32 %cond
 }
 
 ; int csinc4 (int a, int b) { return a ? b+1 : b-3; }
-define dso_local i32 @csinc4(i32 %a, i32 %b) {
+define i32 @csinc4(i32 %a, i32 %b) {
 ; CHECK-LABEL: csinc4:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sub w8, w1, #3
 ; CHECK-NEXT:    cmp w0, #0
 ; 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 -3, i32 1
   %cond = add nsw i32 %cond.v, %b
   ret i32 %cond
 }
 
 ; int csinc5 (int a, int b) { return a ? b+1 : b-4095; }
-define dso_local i32 @csinc5(i32 %a, i32 %b) {
+define i32 @csinc5(i32 %a, i32 %b) {
 ; CHECK-LABEL: csinc5:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sub w8, w1, #4095
 ; CHECK-NEXT:    cmp w0, #0
 ; 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 -4095, i32 1
   %cond = add nsw i32 %cond.v, %b
   ret i32 %cond
 }
 
 ; int csinc6 (int a, int b) { return a ? b+1 : b-4096; }
-define dso_local i32 @csinc6(i32 %a, i32 %b) {
+define i32 @csinc6(i32 %a, i32 %b) {
 ; CHECK-LABEL: csinc6:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sub w8, w1, #1, lsl #12 // =4096
 ; CHECK-NEXT:    cmp w0, #0
 ; 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 -4096, i32 1
   %cond = add nsw i32 %cond.v, %b
   ret i32 %cond
 }
 
 ; prevent larger constants (the add laid after csinc)
 ; int csinc7 (int a, int b) { return a ? b+1 : b-4097; }
-define dso_local i32 @csinc7(i32 %a, i32 %b) {
+define i32 @csinc7(i32 %a, i32 %b) {
 ; CHECK-LABEL: csinc7:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    cmp w0, #0
 ; CHECK-NEXT:    mov w8, #-4097
 ; CHECK-NEXT:    csinc w8, w8, wzr, eq
 ; 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 i32 @csinc8(i32 %a, i32 %b) {
+; CHECK-LABEL: csinc8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    sub w8, w1, #1
+; CHECK-NEXT:    cmp w0, #0
+; 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 i32 @csinc9(i32 %a, i32 %b) {
+; CHECK-LABEL: csinc9:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    sub w8, w1, #1
+; CHECK-NEXT:    cmp w0, #0
+; 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
+}