[llvm] [AArch64] Sub, add, and icmp should have the fact they can be negated in the cost (PR #142844)

[via llvm-commits]

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

>From 01fae95b7445311f89bcf055916e7602ba336a46 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Wed, 4 Jun 2025 16:24:48 -0400
Subject: [PATCH] Update costs for AArch64

---
 .../AArch64/AArch64TargetTransformInfo.cpp    | 25 ++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index 68aec80f07e1d..7949a78f3da4b 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
@@ -413,6 +413,20 @@ AArch64TTIImpl::getIntImmCost(const APInt &Imm, Type *Ty,
   return std::max<InstructionCost>(1, Cost);
 }
 
+static bool isLegalArithImmed(uint64_t C) {
+  // Matches AArch64DAGToDAGISel::SelectArithImmed().
+  bool IsLegal = (C >> 12 == 0) || ((C & 0xFFFULL) == 0 && C >> 24 == 0);
+  LLVM_DEBUG(dbgs() << "Is imm " << C
+                    << " legal: " << (IsLegal ? "yes\n" : "no\n"));
+  return IsLegal;
+}
+
+static bool isLegalCmpImmed(APInt C) {
+  // Works for negative immediates too, as it can be written as an ADDS
+  // instruction with a negated immediate.
+  return isLegalArithImmed(C.abs().getZExtValue());
+}
+
 InstructionCost AArch64TTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
                                                   const APInt &Imm, Type *Ty,
                                                   TTI::TargetCostKind CostKind,
@@ -473,7 +487,16 @@ InstructionCost AArch64TTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
 
   if (Idx == ImmIdx) {
     int NumConstants = (BitSize + 63) / 64;
-    InstructionCost Cost = AArch64TTIImpl::getIntImmCost(Imm, Ty, CostKind);
+    InstructionCost Cost;
+    if ((Opcode == Instruction::Add || Opcode == Instruction::Sub ||
+         Opcode == Instruction::ICmp) &&
+        BitSize <= 64) {
+      // Add/Sub/ICmp immediates can be flipped.
+      if (isLegalCmpImmed(Imm))
+        return TTI::TCC_Free;
+      Cost = AArch64TTIImpl::getIntImmCost(Imm.abs(), Ty, CostKind);
+    } else
+      Cost = AArch64TTIImpl::getIntImmCost(Imm, Ty, CostKind);
     return (Cost <= NumConstants * TTI::TCC_Basic)
                ? static_cast<int>(TTI::TCC_Free)
                : Cost;