[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 8d439cfedf622281bc26d6e82da37aeb210e5364 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Wed, 4 Jun 2025 16:24:48 -0400
Subject: [PATCH 1/2] Update costs for AArch64

---
 .../AArch64/AArch64TargetTransformInfo.cpp    | 35 +++++++++++++++----
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index 68aec80f07e1d..9a00bf2dbc924 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,10 +487,19 @@ InstructionCost AArch64TTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
 
   if (Idx == ImmIdx) {
     int NumConstants = (BitSize + 63) / 64;
-    InstructionCost Cost = AArch64TTIImpl::getIntImmCost(Imm, Ty, CostKind);
-    return (Cost <= NumConstants * TTI::TCC_Basic)
-               ? static_cast<int>(TTI::TCC_Free)
-               : Cost;
+    InstructionCost Cost;
+    if ((Opcode == Instruction::Add || Opcode == Instruction::Sub ||
+         Opcode == Instruction::ICmp) &&
+        BitSize <= 64) {
+      // Add/Sub/ICmp immediates can be flipped.
+      // Also they have different requirements as to fitting in an immediate
+      // than others.
+      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) ? TTI::TCC_Free : Cost;
   }
   return AArch64TTIImpl::getIntImmCost(Imm, Ty, CostKind);
 }
@@ -511,9 +534,7 @@ AArch64TTIImpl::getIntImmCostIntrin(Intrinsic::ID IID, unsigned Idx,
     if (Idx == 1) {
       int NumConstants = (BitSize + 63) / 64;
       InstructionCost Cost = AArch64TTIImpl::getIntImmCost(Imm, Ty, CostKind);
-      return (Cost <= NumConstants * TTI::TCC_Basic)
-                 ? static_cast<int>(TTI::TCC_Free)
-                 : Cost;
+      return (Cost <= NumConstants * TTI::TCC_Basic) ? TTI::TCC_Free : Cost;
     }
     break;
   case Intrinsic::experimental_stackmap:

>From ea8ab29aa55b45527c069c0d9ce6e067071b6321 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Thu, 5 Jun 2025 11:09:41 -0400
Subject: [PATCH 2/2] Division by a constant

---
 .../Target/AArch64/AArch64TargetTransformInfo.cpp   | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index 9a00bf2dbc924..897d8b76866bc 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
@@ -454,16 +454,21 @@ InstructionCost AArch64TTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
   case Instruction::Add:
   case Instruction::Sub:
   case Instruction::Mul:
-  case Instruction::UDiv:
-  case Instruction::SDiv:
-  case Instruction::URem:
-  case Instruction::SRem:
   case Instruction::And:
   case Instruction::Or:
   case Instruction::Xor:
   case Instruction::ICmp:
     ImmIdx = 1;
     break;
+
+  // Division by a constant can be turned into multiplication, but only if we
+  // know it's constant. So it's not so much that the immediate is cheap (it's
+  // not), but that the alternative is worse.
+  case Instruction::UDiv:
+  case Instruction::SDiv:
+  case Instruction::URem:
+  case Instruction::SRem:
+
   // Always return TCC_Free for the shift value of a shift instruction.
   case Instruction::Shl:
   case Instruction::LShr: