[DAG] allow more select folding for targets that have 'and not' (PR31175)

[spatel]

Summary

The original motivation for this patch comes from wanting to canonicalize more IR to selects and also canonicalizing min/max.

If we're going to do that, we need more backend fixups to undo select codegen when simpler ops will do. I chose AArch64 for the tests because that shows the difference in the simplest way. This should fix:
https://llvm.org/bugs/show_bug.cgi?id=31175

[spatel]

Ping.

[evandro]

LGTM

llvm/trunk/include/llvm/Target/TargetLowering.h

public:
   /// Note that the transform will not occur if Y is known to be a power-of-2
   /// because a mask and compare of a single bit can be handled by inverting the
   /// predicate, for example:
   /// (X & 8) == 8 ---> (X & 8) != 0
   virtual bool hasAndNotCompare(SDValue Y) const {
     return false;
   }
 
+  /// Return true if the target has a bitwise and-not operation:
+  /// X = ~A & B
+  /// This can be used to simplify select or other instructions.
+  virtual bool hasAndNot(SDValue X) const {
+    // If the target has the more complex version of this operation, assume that
+    // it has this operation too.
+    return hasAndNotCompare(X);
+  }
+
   /// \brief Return true if the target wants to use the optimization that
   /// turns ext(promotableInst1(...(promotableInstN(load)))) into
   /// promotedInst1(...(promotedInstN(ext(load)))).
   bool enableExtLdPromotion() const { return EnableExtLdPromotion; }
 
   /// Return true if the target can combine store(extractelement VectorTy,
   /// Idx).
   /// \p Cost[out] gives the cost of that transformation when this is true.

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

bool DAGCombiner::SimplifySelectOps(SDNode *TheSelect, SDValue LHS,
   return false;
 }
 
 /// Try to fold an expression of the form (N0 cond N1) ? N2 : N3 to a shift and
 /// bitwise 'and'.
 SDValue DAGCombiner::foldSelectCCToShiftAnd(const SDLoc &DL, SDValue N0,
                                             SDValue N1, SDValue N2, SDValue N3,
                                             ISD::CondCode CC) {
-  // Check to see if we can perform the "gzip trick", transforming
-  // (select_cc setlt X, 0, A, 0) -> (and (sra X, size(X)-1), A)
+  // If this is a select where the false operand is zero and the compare is a
+  // check of the sign bit, see if we can perform the "gzip trick":
+  // select_cc setlt X, 0, A, 0 -> and (sra X, size(X)-1), A
+  // select_cc setgt X, 0, A, 0 -> and (not (sra X, size(X)-1)), A
   EVT XType = N0.getValueType();
   EVT AType = N2.getValueType();
-  if (!isNullConstant(N3) || CC != ISD::SETLT || !XType.bitsGE(AType))
+  if (!isNullConstant(N3) || !XType.bitsGE(AType))
     return SDValue();
 
-  // (a < 0) ? b : 0
-  // (a < 1) ? a : 0
+  // If the comparison is testing for a positive value, we have to invert
+  // the sign bit mask, so only do that transform if the target has a bitwise
+  // 'and not' instruction (the invert is free).
+  if (CC == ISD::SETGT && TLI.hasAndNot(N2)) {
+    // (X > -1) ? A : 0
+    // (X >  0) ? X : 0 <-- This is canonical signed max.
+    if (!(isAllOnesConstant(N1) || (isNullConstant(N1) && N0 == N2)))
+      return SDValue();
+  } else if (CC == ISD::SETLT) {
+    // (X <  0) ? A : 0
+    // (X <  1) ? X : 0 <-- This is un-canonicalized signed min.
     if (!(isNullConstant(N1) || (isOneConstant(N1) && N0 == N2)))
       return SDValue();
+  } else {
+    return SDValue();
+  }
 
   // and (sra X, size(X)-1), A -> "and (srl X, C2), A" iff A is a single-bit
   // constant.
   EVT ShiftAmtTy = getShiftAmountTy(N0.getValueType());
   auto *N2C = dyn_cast<ConstantSDNode>(N2.getNode());
   if (N2C && ((N2C->getAPIntValue() & (N2C->getAPIntValue() - 1)) == 0)) {
     unsigned ShCt = XType.getSizeInBits() - N2C->getAPIntValue().logBase2() - 1;
     SDValue ShiftAmt = DAG.getConstant(ShCt, DL, ShiftAmtTy);
     SDValue Shift = DAG.getNode(ISD::SRL, DL, XType, N0, ShiftAmt);
     AddToWorklist(Shift.getNode());
 
     if (XType.bitsGT(AType)) {
       Shift = DAG.getNode(ISD::TRUNCATE, DL, AType, Shift);
       AddToWorklist(Shift.getNode());
     }
 
+    if (CC == ISD::SETGT)
+      Shift = DAG.getNOT(DL, Shift, AType);
+
     return DAG.getNode(ISD::AND, DL, AType, Shift, N2);
   }
 
   SDValue ShiftAmt = DAG.getConstant(XType.getSizeInBits() - 1, DL, ShiftAmtTy);
   SDValue Shift = DAG.getNode(ISD::SRA, DL, XType, N0, ShiftAmt);
   AddToWorklist(Shift.getNode());
 
   if (XType.bitsGT(AType)) {
     Shift = DAG.getNode(ISD::TRUNCATE, DL, AType, Shift);
     AddToWorklist(Shift.getNode());
   }
 
+  if (CC == ISD::SETGT)
+    Shift = DAG.getNOT(DL, Shift, AType);
+
   return DAG.getNode(ISD::AND, DL, AType, Shift, N2);
 }
 
 /// Simplify an expression of the form (N0 cond N1) ? N2 : N3
 /// where 'cond' is the comparison specified by CC.
 SDValue DAGCombiner::SimplifySelectCC(const SDLoc &DL, SDValue N0, SDValue N1,
                                       SDValue N2, SDValue N3, ISD::CondCode CC,
                                       bool NotExtCompare) {

llvm/trunk/test/CodeGen/AArch64/arm64-icmp-opt.ll

 ; RUN: llc < %s -mtriple=arm64-eabi | FileCheck %s
 
 ; Optimize (x > -1) to (x >= 0) etc.
 ; Optimize (cmp (add / sub), 0): eliminate the subs used to update flag
 ;   for comparison only
 ; rdar://10233472
 
 define i32 @t1(i64 %a) {
 ; CHECK-LABEL: t1:
 ; CHECK:       // BB#0:
-; CHECK-NEXT:    cmp x0, #0
-; CHECK-NEXT:    cset w0, ge
+; CHECK-NEXT:    lsr x8, x0, #63
+; CHECK-NEXT:    eor w0, w8, #0x1
 ; CHECK-NEXT:    ret
 ;
   %cmp = icmp sgt i64 %a, -1
   %conv = zext i1 %cmp to i32
   ret i32 %conv
 }

llvm/trunk/test/CodeGen/AArch64/selectcc-to-shiftand.ll

 ; CHECK-NEXT:    and w0, w0, w0, asr #31
 ; CHECK-NEXT:    ret
 ;
   %tmp = icmp slt i32 %a, 1
   %min = select i1 %tmp, i32 %a, i32 0
   ret i32 %min
 }
 
-; FIXME: Flipping the comparison condition can be handled by getting the bitwise not of the sign mask.
+; Flipping the comparison condition can be handled by getting the bitwise not of the sign mask.
 
 ; Compare if positive and select of constants where one constant is zero.
 
 define i32 @pos_sel_constants(i32 %a) {
 ; CHECK-LABEL: pos_sel_constants:
 ; CHECK:       // BB#0:
-; CHECK-NEXT:    cmp w0, #0
 ; CHECK-NEXT:    mov w8, #5
-; CHECK-NEXT:    csel w0, w8, wzr, ge
+; CHECK-NEXT:    bic w0, w8, w0, asr #31
 ; CHECK-NEXT:    ret
 ;
   %tmp.1 = icmp sgt i32 %a, -1
   %retval = select i1 %tmp.1, i32 5, i32 0
   ret i32 %retval
 }
 
 ; Compare if positive and select of constants where one constant is zero and the other is a single bit.
 
 define i32 @pos_sel_special_constant(i32 %a) {
 ; CHECK-LABEL: pos_sel_special_constant:
 ; CHECK:       // BB#0:
-; CHECK-NEXT:    cmp w0, #0
-; CHECK-NEXT:    cset w8, ge
-; CHECK-NEXT:    lsl w0, w8, #9
+; CHECK-NEXT:    orr w8, wzr, #0x200
+; CHECK-NEXT:    bic w0, w8, w0, lsr #22
 ; CHECK-NEXT:    ret
 ;
   %tmp.1 = icmp sgt i32 %a, -1
   %retval = select i1 %tmp.1, i32 512, i32 0
   ret i32 %retval
 }
 
 ; Compare if positive and select variable or zero.
 
 define i32 @pos_sel_variable_and_zero(i32 %a, i32 %b) {
 ; CHECK-LABEL: pos_sel_variable_and_zero:
 ; CHECK:       // BB#0:
-; CHECK-NEXT:    cmp w0, #0
-; CHECK-NEXT:    csel w0, w1, wzr, ge
+; CHECK-NEXT:    bic w0, w1, w0, asr #31
 ; CHECK-NEXT:    ret
 ;
   %tmp.1 = icmp sgt i32 %a, -1
   %retval = select i1 %tmp.1, i32 %b, i32 0
   ret i32 %retval
 }
 
 ; Compare if not negative or zero and select the same variable as being compared: smax(a, 0).
 
 define i32 @not_neg_sel_same_variable(i32 %a) {
 ; CHECK-LABEL: not_neg_sel_same_variable:
 ; CHECK:       // BB#0:
-; CHECK-NEXT:    cmp w0, #0
-; CHECK-NEXT:    csel w0, w0, wzr, gt
+; CHECK-NEXT:    bic w0, w0, w0, asr #31
 ; CHECK-NEXT:    ret
 ;
   %tmp = icmp sgt i32 %a, 0
   %min = select i1 %tmp, i32 %a, i32 0
   ret i32 %min
 }
 
 ; https://llvm.org/bugs/show_bug.cgi?id=31175
 
 ; ret = (x-y) > 0 ? x-y : 0
 define i32 @PR31175(i32 %x, i32 %y) {
 ; CHECK-LABEL: PR31175:
 ; CHECK:       // BB#0:
 ; CHECK-NEXT:    sub w8, w0, w1
-; CHECK-NEXT:    cmp w8, #0
-; CHECK-NEXT:    csel w0, w8, wzr, gt
+; CHECK-NEXT:    bic w0, w8, w8, asr #31
 ; CHECK-NEXT:    ret
 ;
   %sub = sub nsw i32 %x, %y
   %cmp = icmp sgt i32 %sub, 0
   %sel = select i1 %cmp, i32 %sub, i32 0
   ret i32 %sel
 }