[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 9a0f0f71a4fc66e80513d1f24e7e17b56bc6b0cc Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sat, 26 Jul 2025 18:01:24 -0400
Subject: [PATCH 1/2] [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..98273bbbda8e5 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) If MI is add/sub with no-unsigned-wrap, it produces a poison value when
+  //    unsigned overflow occurs, so CmpInstr could still be simplified away.
+  if (NZVCUsed->C &&
+      !(isADDSRegImm(CmpOpcode) && MI.getFlag(MachineInstr::NoUWrap)))
     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;
         }
       }

>From 5c8e9205e9ffbf76d367b423e8ef9b10cec861a7 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sat, 26 Jul 2025 18:04:38 -0400
Subject: [PATCH 2/2] Update AArch64InstrInfo.cpp

---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 98273bbbda8e5..fce05dbcc3104 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1924,10 +1924,8 @@ static bool canInstrSubstituteCmpInstr(MachineInstr &MI, MachineInstr &CmpInstr,
   // '%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) If MI is add/sub with no-unsigned-wrap, it produces a poison value when
-  //    unsigned overflow occurs, so CmpInstr could still be simplified away.
-  if (NZVCUsed->C &&
-      !(isADDSRegImm(CmpOpcode) && MI.getFlag(MachineInstr::NoUWrap)))
+  // 2) ADDS x, 0, always sets C to 0.
+  if (NZVCUsed->C && !isADDSRegImm(CmpOpcode))
     return false;
 
   // CmpInstr is either 'ADDS %vreg, 0' or 'SUBS %vreg, 0', and MI is either