[AArch64][SelectionDAG] Generate subs+csel for usub.sat

Fixes #191488

As this is a regression of
#170076, adds a check to avoid
generic lowering of usub.sat to X - zext(X != 0) in case of aarch64 via
a virtual hook in TargetLowering. All other backends will still receive
generic lowering as implemented in the original patch.

Change-Id: I0a194bcc9e66819c12d0f9179464823301f0d7bf

Author: bojle

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -3595,6 +3595,10 @@ class LLVM_ABI TargetLoweringBase {
     return false;
   }
 
+  /// Should usub.sat(X, 1) prefer the generic lowering X - zext(X != 0) over
+  /// the default overflow/select expansion?
+  virtual bool preferSubOfZextForUsubSatOne(EVT VT) const { return true; }
+
   /// True if target has some particular form of dealing with pointer arithmetic
   /// semantics for pointers with the given value type. False if pointer
   /// arithmetic should not be preserved for passes such as instruction

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -11475,7 +11475,8 @@ SDValue TargetLowering::expandAddSubSat(SDNode *Node, SelectionDAG &DAG) const {
   }
 
   // usub.sat(a, 1) -> sub(a, zext(a != 0))
-  if (Opcode == ISD::USUBSAT && isOneOrOneSplat(RHS)) {
+  if (Opcode == ISD::USUBSAT && isOneOrOneSplat(RHS) &&
+      preferSubOfZextForUsubSatOne(VT)) {
     LHS = DAG.getFreeze(LHS);
     SDValue Zero = DAG.getConstant(0, dl, VT);
     EVT BoolVT = getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(), VT);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -30978,6 +30978,11 @@ bool AArch64TargetLowering::shouldConvertFpToSat(unsigned Op, EVT FPVT,
   return TargetLowering::shouldConvertFpToSat(Op, FPVT, VT);
 }
 
+bool AArch64TargetLowering::preferSubOfZextForUsubSatOne(EVT /*VT*/) const {
+  // See https://github.com/llvm/llvm-project/issues/191488
+  return false;
+}
+
 bool AArch64TargetLowering::preferSelectsOverBooleanArithmetic(EVT VT) const {
   // Expand scalar and SVE operations using selects. Neon vectors prefer sub to
   // avoid vselect becoming bsl / unrolling.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -450,6 +450,8 @@ class AArch64TargetLowering : public TargetLowering {
 
   bool shouldConvertFpToSat(unsigned Op, EVT FPVT, EVT VT) const override;
 
+  bool preferSubOfZextForUsubSatOne(EVT VT) const override;
+
   bool preferSelectsOverBooleanArithmetic(EVT VT) const override;
 
   bool isComplexDeinterleavingSupported() const override;

llvm/test/CodeGen/AArch64/and-mask-removal.ll

@@ -483,9 +483,9 @@ define i64 @pr58109(i8 signext %0) {
 ; CHECK-SD-LABEL: pr58109:
 ; CHECK-SD:       ; %bb.0:
 ; CHECK-SD-NEXT:    add w8, w0, #1
-; CHECK-SD-NEXT:    ands w8, w8, #0xff
-; CHECK-SD-NEXT:    cset w9, ne
-; CHECK-SD-NEXT:    sub w0, w8, w9
+; CHECK-SD-NEXT:    and w8, w8, #0xff
+; CHECK-SD-NEXT:    subs w8, w8, #1
+; CHECK-SD-NEXT:    csel w0, wzr, w8, lo
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: pr58109:

llvm/test/CodeGen/AArch64/usub_sat_plus.ll

@@ -8,6 +8,24 @@ declare i16 @llvm.usub.sat.i16(i16, i16)
 declare i32 @llvm.usub.sat.i32(i32, i32)
 declare i64 @llvm.usub.sat.i64(i64, i64)
 
+define i32 @sat_dec_i32(i32 %x) nounwind {
+; CHECK-SD-LABEL: sat_dec_i32:
+; CHECK-SD:       // %bb.0:
+; CHECK-SD-NEXT:    subs w8, w0, #1
+; CHECK-SD-NEXT:    csel w0, wzr, w8, lo
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: sat_dec_i32:
+; CHECK-GI:       // %bb.0:
+; CHECK-GI-NEXT:    subs w8, w0, #1
+; CHECK-GI-NEXT:    cset w9, lo
+; CHECK-GI-NEXT:    tst w9, #0x1
+; CHECK-GI-NEXT:    csel w0, wzr, w8, ne
+; CHECK-GI-NEXT:    ret
+  %tmp = call i32 @llvm.usub.sat.i32(i32 %x, i32 1)
+  ret i32 %tmp
+}
+
 define i32 @func32(i32 %x, i32 %y, i32 %z) nounwind {
 ; CHECK-SD-LABEL: func32:
 ; CHECK-SD:       // %bb.0: