[llvm] [AArch64] Increase scatter overhead on Neoverse-V2 (PR #101296)

[Madhur Amilkanthwar via llvm-commits]

https://github.com/madhur13490 updated https://github.com/llvm/llvm-project/pull/101296

>From 5bc19b3d59dded4756d2cd18d22c6cd727217b37 Mon Sep 17 00:00:00 2001
From: Madhur Amilkanthwar <madhura at nvidia.com>
Date: Thu, 25 Jul 2024 18:30:21 +0530
Subject: [PATCH] [AArch64] Increase scatter overhead on Neoverse-V2

This patch increases scatter overhead on Neoverse-V2 to 13.
This benefits s128 kernel from TSVC_2 test suite.
SPEC 17, RAJAPerf, Sptter are unaffected with this patch.
---
 .../AArch64/AArch64TargetTransformInfo.cpp    | 22 +++++++++++++++++--
 .../LoopVectorize/AArch64/scatter-cost.ll     |  8 +++++++
 2 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopVectorize/AArch64/scatter-cost.ll

diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index 79c0e45e3aa5b..d8c6416849cc5 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
@@ -3389,8 +3389,26 @@ AArch64TTIImpl::getMaskedMemoryOpCost(unsigned Opcode, Type *Src,
   return LT.first;
 }
 
-static unsigned getSVEGatherScatterOverhead(unsigned Opcode) {
+static unsigned getSVEGatherScatterOverhead(unsigned Opcode, AArch64Subtarget::ARMProcFamilyEnum ProcFamily) {
+  assert((Opcode == Instruction::Load || Opcode == Instruction::Store) &&
+          "Should be called on only load or stores.");
   return Opcode == Instruction::Load ? SVEGatherOverhead : SVEScatterOverhead;
+  unsigned Cost = 1;
+  switch(Opcode) {
+    case Instruction::Load:
+      Cost = SVEGatherOverhead;
+      break;
+    case Instruction::Store:
+      if (ProcFamily == AArch64Subtarget::NeoverseV2) {
+        Cost = 13;
+      } else {
+        Cost = SVEScatterOverhead;
+      }
+    break;
+    default:
+      llvm_unreachable("Shouldn't have reached here");
+  }
+  return Cost;
 }
 
 InstructionCost AArch64TTIImpl::getGatherScatterOpCost(
@@ -3424,7 +3442,7 @@ InstructionCost AArch64TTIImpl::getGatherScatterOpCost(
   // Add on an overhead cost for using gathers/scatters.
   // TODO: At the moment this is applied unilaterally for all CPUs, but at some
   // point we may want a per-CPU overhead.
-  MemOpCost *= getSVEGatherScatterOverhead(Opcode);
+  MemOpCost *= getSVEGatherScatterOverhead(Opcode, ST->getProcFamily());
   return LT.first * MemOpCost * getMaxNumElements(LegalVF);
 }
 
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/scatter-cost.ll b/llvm/test/Transforms/LoopVectorize/AArch64/scatter-cost.ll
new file mode 100644
index 0000000000000..8bdae00411a1f
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/scatter-cost.ll
@@ -0,0 +1,8 @@
+; RUN: opt -mtriple aarch64 -mcpu=neoverse-v2 -passes="print<cost-model>" -disable-output | FileCheck %s
+; CHECK: Cost Model: Found an estimated cost of 52 for instruction: call void @llvm.masked.scatter.nxv4f32
+
+define void @masked_scatter_nxv8f32_i64(<vscale x 4 x float> %data, <vscale x 4 x ptr> %b, <vscale x 4 x i64> %V) #0 {
+  call void @llvm.masked.scatter.nxv4f32.nxv4p0(<vscale x 4 x float> %data, <vscale x 4 x ptr> %b, i32 4, <vscale x 4 x i1> shufflevector (<vscale x 4 x i1> insertelement (<vscale x 4 x i1> poison, i1 true, i64 0), <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer))
+  ret void
+}
+