[AArch64][PeepholeOpt]Optimize ALU + compare to flag-setting ALU

[mingmingl]

Summary

The motivating example is in https://godbolt.org/z/45nbdYMK9

• For this example, subs is generated for the good case; sub followed by cmp is generated for the bad case. Since signed overflow is undefined behavior in C/C++ (indicated as nsw flag in LLVM IR), subs should be generated for the good case as well.

This patch relaxes one restriction from "quit optimization when V is used" to "continue if MI produces poison value when signed overflow occurs". This is not meant to be C/C++ specific since it looks at 'NoSWrap' since it looks at MachineInstr flags.

[dmgreen]

I see, I like it. I sometimes worry about poison in the backend but this looks like a good idea.
LGTM. Thanks

[mingmingl]

thanks for reviews! Going to submit this patch and the pre-commit test (D146754) together.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

 /// CmpInstr can be substituted:
 /// - CmpInstr is either 'ADDS %vreg, 0' or 'SUBS %vreg, 0'
 /// - and, MI and CmpInstr are from the same MachineBB
 /// - and, condition flags are not alive in successors of the CmpInstr parent
 /// - and, if MI opcode is the S form there must be no defs of flags between
 ///        MI and CmpInstr
 ///        or if MI opcode is not the S form there must be neither defs of flags
 ///        nor uses of flags between MI and CmpInstr.
-/// - and  C/V flags are not used after CmpInstr
+/// - and, if C/V flags are not used after CmpInstr
+///        or if N flag is used but MI produces poison value if signed overflow
+///        occurs.
 static bool canInstrSubstituteCmpInstr(MachineInstr &MI, MachineInstr &CmpInstr,
                                        const TargetRegisterInfo &TRI) {
+  // NOTE this assertion guarantees that MI.getOpcode() is add or subtraction
+  // that may or may not set flags.
   assert(sForm(MI) != AArch64::INSTRUCTION_LIST_END);
 
   const unsigned CmpOpcode = CmpInstr.getOpcode();
   if (!isADDSRegImm(CmpOpcode) && !isSUBSRegImm(CmpOpcode))
     return false;
 
+  assert((CmpInstr.getOperand(2).isImm() &&
+          CmpInstr.getOperand(2).getImm() == 0) &&
+         "Caller guarantees that CmpInstr compares with constant 0");
+
   std::optional<UsedNZCV> NZVCUsed = examineCFlagsUse(MI, CmpInstr, TRI);
-  if (!NZVCUsed || NZVCUsed->C || NZVCUsed->V)
+  if (!NZVCUsed || NZVCUsed->C)
+    return false;
+
+  // CmpInstr is either 'ADDS %vreg, 0' or 'SUBS %vreg, 0', and MI is either
+  // '%vreg = add ...' or '%vreg = sub ...'.
+  // Condition flag V is used to indicate signed overflow.
+  // 1) MI and CmpInstr set N and V to the same value.
+  // 2) If MI is add/sub with no-signed-wrap, it produces a poison value when
+  //    signed overflow occurs, so CmpInstr could still be simplified away.
+  if (NZVCUsed->V && !MI.getFlag(MachineInstr::NoSWrap))
     return false;
 
   AccessKind AccessToCheck = AK_Write;
   if (sForm(MI) != MI.getOpcode())
     AccessToCheck = AK_All;
   return !areCFlagsAccessedBetweenInstrs(&MI, &CmpInstr, &TRI, AccessToCheck);
 }
 

llvm/test/CodeGen/AArch64/aarch64-icmp-opt.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
 ; RUN: llc -O3 -mtriple=aarch64 %s -o - | FileCheck %s
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-unknown-linux-gnu"
 
 define i32 @sub_icmp_i32(i32 %0, i32 %1) {
 ; CHECK-LABEL: sub_icmp_i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sub w0, w0, w1
-; CHECK-NEXT:    cmp w0, #0
+; CHECK-NEXT:    subs w0, w0, w1
 ; CHECK-NEXT:    b.le .LBB0_2
 ; CHECK-NEXT:  // %bb.1:
 ; CHECK-NEXT:    b _Z2f2i
 ; CHECK-NEXT:  .LBB0_2:
 ; CHECK-NEXT:    b _Z2f1i
   %3 = sub nsw i32 %0, %1
   %4 = icmp slt i32 %3, 1
   br i1 %4, label %5, label %7
9:
   ret i32 %10
 }
 
 
 
 define i64 @sub_icmp_i64(i64 %0, i64 %1) {
 ; CHECK-LABEL: sub_icmp_i64:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sub x0, x0, x1
-; CHECK-NEXT:    cmp x0, #0
+; CHECK-NEXT:    subs x0, x0, x1
 ; CHECK-NEXT:    b.le .LBB1_2
 ; CHECK-NEXT:  // %bb.1:
 ; CHECK-NEXT:    b _Z2f4l
 ; CHECK-NEXT:  .LBB1_2:
 ; CHECK-NEXT:    b _Z2f3l
   %3 = sub nsw i64 %0, %1
   %4 = icmp slt i64 %3, 1
   br i1 %4, label %5, label %7

 9:
   %10 = phi i64 [ %6, %5 ], [ %8, %7 ]
   ret i64 %10
 }
 
 define i64 @add_i64(i64 %0, i64 %1) {
 ; CHECK-LABEL: add_i64:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    add x0, x1, x0
-; CHECK-NEXT:    cmp x0, #0
+; CHECK-NEXT:    adds x0, x1, x0
 ; CHECK-NEXT:    b.le .LBB2_2
 ; CHECK-NEXT:  // %bb.1:
 ; CHECK-NEXT:    b _Z2f4l
 ; CHECK-NEXT:  .LBB2_2:
 ; CHECK-NEXT:    b _Z2f3l
   %3 = add nsw i64 %1, %0
   %4 = icmp slt i64 %3, 1
   br i1 %4, label %5, label %7

 9:
   %10 = phi i64 [ %6, %5 ], [ %8, %7 ]
   ret i64 %10
 }
 
 define i32 @add_i32(i32 %0, i32 %1) {
 ; CHECK-LABEL: add_i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    add w0, w1, w0
-; CHECK-NEXT:    cmp w0, #0
+; CHECK-NEXT:    adds w0, w1, w0
 ; CHECK-NEXT:    b.le .LBB3_2
 ; CHECK-NEXT:  // %bb.1:
 ; CHECK-NEXT:    b _Z2f4l
 ; CHECK-NEXT:  .LBB3_2:
 ; CHECK-NEXT:    b _Z2f3l
   %3 = add nsw i32 %1, %0
   %4 = icmp slt i32 %3, 1
   br i1 %4, label %5, label %7

llvm/test/CodeGen/AArch64/arm64-csel.ll

entry:
   %cond = select i1 %cmp, i32 %sub, i32 %sub3
   ret i32 %cond
 }
 
 ; make sure we can handle branch instruction in optimizeCompare.
 define i32@foo6(i32 %a, i32 %b) nounwind ssp {
 ; CHECK-LABEL: foo6:
 ; CHECK:       // %bb.0: // %common.ret
-; CHECK-NEXT:    sub w8, w0, w1
-; CHECK-NEXT:    cmp w8, #0
+; CHECK-NEXT:    subs w8, w0, w1
 ; CHECK-NEXT:    csinc w0, w8, wzr, le
 ; CHECK-NEXT:    ret
   %sub = sub nsw i32 %a, %b
   %cmp = icmp sgt i32 %sub, 0
   br i1 %cmp, label %l.if, label %l.else
 
 l.if:
   ret i32 1