From 3c3e0e556fd768897d003fffa9e2308de982d4ae Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev@redking.me.uk>
Date: Fri, 5 Apr 2024 15:04:03 +0100
Subject: [PATCH] [TTI] getCommonMaskedMemoryOpCost - add AddressSpace argument

Still defaults to AddressSpace = 0, but will allow us to hook up getMaskedMemoryOpCost calls once we have test coverage.
---
 llvm/include/llvm/CodeGen/BasicTTIImpl.h | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/BasicTTIImpl.h b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
index 241529e22538a7..42d8f74fd427fb 100644
--- a/llvm/include/llvm/CodeGen/BasicTTIImpl.h
+++ b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
@@ -209,7 +209,8 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
                                               Align Alignment,
                                               bool VariableMask,
                                               bool IsGatherScatter,
-                                              TTI::TargetCostKind CostKind) {
+                                              TTI::TargetCostKind CostKind,
+                                              unsigned AddressSpace = 0) {
     // We cannot scalarize scalable vectors, so return Invalid.
     if (isa<ScalableVectorType>(DataTy))
       return InstructionCost::getInvalid();
@@ -230,8 +231,9 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
                   CostKind, -1, nullptr, nullptr)
             : 0;
     InstructionCost LoadCost =
-        VF * (AddrExtractCost + getMemoryOpCost(Opcode, VT->getElementType(),
-                                                Alignment, 0, CostKind));
+        VF *
+        (AddrExtractCost + getMemoryOpCost(Opcode, VT->getElementType(),
+                                           Alignment, AddressSpace, CostKind));
 
     // Next, compute the cost of packing the result in a vector.
     InstructionCost PackingCost =
@@ -1369,6 +1371,7 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
   InstructionCost getMaskedMemoryOpCost(unsigned Opcode, Type *DataTy,
                                         Align Alignment, unsigned AddressSpace,
                                         TTI::TargetCostKind CostKind) {
+    // TODO: Pass on AddressSpace when we have test coverage.
     return getCommonMaskedMemoryOpCost(Opcode, DataTy, Alignment, true, false,
                                        CostKind);
   }