[llvm] [ARM][AArch64] Allow the CSE to take into consideration uses of the carry and overflow flags in ARM and AArch64 (PR #150803)

[via llvm-commits]

https://github.com/AZero13 updated https://github.com/llvm/llvm-project/pull/150803

>From d8d04af42c2ed7895bd66dcfb1bfa9533d93bc11 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sat, 26 Jul 2025 18:01:24 -0400
Subject: [PATCH] [ARM][AArch64] Allow the CSE to take into consideration uses
 of the carry and overflow flags in ARM and AArch64

On both of these platforms, we know that the cmp will not stomp on these flags and overwrite them if doing so would be poison, or in ANDS case, it will always have the V flag cleared during an ANDS.
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp | 23 +++++++++++++++++---
 llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp     | 11 +++++++---
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 8685d7a04ac9c..e916c934fd26c 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1884,6 +1884,10 @@ static bool isSUBSRegImm(unsigned Opcode) {
   return Opcode == AArch64::SUBSWri || Opcode == AArch64::SUBSXri;
 }
 
+static bool isANDSRegImm(unsigned Opcode) {
+  return Opcode == AArch64::ANDSWri || Opcode == AArch64::ANDSXri;
+}
+
 /// Check if CmpInstr can be substituted by MI.
 ///
 /// CmpInstr can be substituted:
@@ -1904,7 +1908,8 @@ static bool canInstrSubstituteCmpInstr(MachineInstr &MI, MachineInstr &CmpInstr,
   assert(sForm(MI) != AArch64::INSTRUCTION_LIST_END);
 
   const unsigned CmpOpcode = CmpInstr.getOpcode();
-  if (!isADDSRegImm(CmpOpcode) && !isSUBSRegImm(CmpOpcode))
+  if (!isADDSRegImm(CmpOpcode) && !isSUBSRegImm(CmpOpcode) &&
+      !isANDSRegImm(CmpOpcode))
     return false;
 
   assert((CmpInstr.getOperand(2).isImm() &&
@@ -1912,7 +1917,17 @@ static bool canInstrSubstituteCmpInstr(MachineInstr &MI, MachineInstr &CmpInstr,
          "Caller guarantees that CmpInstr compares with constant 0");
 
   std::optional<UsedNZCV> NZVCUsed = examineCFlagsUse(MI, CmpInstr, TRI);
-  if (!NZVCUsed || NZVCUsed->C)
+  if (!NZVCUsed)
+    return false;
+
+  // CmpInstr is either 'ADDS %vreg, 0' or 'SUBS %vreg, 0', and MI is either
+  // '%vreg = add ...' or '%vreg = sub ...'.
+  // Condition flag C is used to indicate unsigned overflow.
+  // 1) MI and CmpInstr set N and C to the same value if Cmp is an adds
+  // 2) ADDS x, 0, always sets C to 0.
+  // In practice we should not really get here, as an unsigned comparison with 0
+  // should have been optimized out anyway, but just in case.
+  if (NZVCUsed->C && !isADDSRegImm(CmpOpcode))
     return false;
 
   // CmpInstr is either 'ADDS %vreg, 0' or 'SUBS %vreg, 0', and MI is either
@@ -1921,7 +1936,9 @@ static bool canInstrSubstituteCmpInstr(MachineInstr &MI, MachineInstr &CmpInstr,
   // 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))
+  // ANDS also always sets V to 0.
+  if (NZVCUsed->V && !MI.getFlag(MachineInstr::NoSWrap) &&
+      !isANDSRegImm(Opcode))
     return false;
 
   AccessKind AccessToCheck = AK_Write;
diff --git a/llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp b/llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp
index 50217c3a047df..c49ed7a7cc9bd 100644
--- a/llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp
+++ b/llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp
@@ -3089,15 +3089,20 @@ bool ARMBaseInstrInfo::optimizeCompareInstr(
           break;
         case ARMCC::HS: // C
         case ARMCC::LO: // C
-        case ARMCC::VS: // V
-        case ARMCC::VC: // V
         case ARMCC::HI: // C Z
         case ARMCC::LS: // C Z
+          // The instruction uses the C bit which is not safe.
+          return false;
+        case ARMCC::VS: // V
+        case ARMCC::VC: // V
         case ARMCC::GE: // N V
         case ARMCC::LT: // N V
         case ARMCC::GT: // Z N V
         case ARMCC::LE: // Z N V
-          // The instruction uses the V bit or C bit which is not safe.
+          // We MAY be able to do this if signed overflow is poison.
+          if (I->getFlag(MachineInstr::NoSWrap))
+            break;
+          // The instruction uses the V bit which is not safe.
           return false;
         }
       }