[clang][CodeGen] cast addr space of ReturnValue if needed (#154380)

Fixes a bug on AMDGPU targets where a pointer was stored as address
space 5, but then loaded as address space 0.

Issue found as part of [Kokkos](https://github.com/kokkos/kokkos)
testing, specifically `hip.atomics` (see
[core/unit_test/TestAtomics.hpp](https://github.com/kokkos/kokkos/blob/develop/core/unit_test/TestAtomics.hpp)).

Issue was introduced by commit
[39ec9de7c230](https://github.com/llvm/llvm-project/commit/39ec9de7c230)
- [clang][CodeGen] sret args should always point to the alloca AS, so
use that (https://github.com/llvm/llvm-project/pull/114062).

1 files changed, 7 insertions, 0 deletions

diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index ad318f2..fc65199 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -2804,6 +2804,13 @@ public:
     AllocaTracker Tracker;
   };
 
+private:
+  /// If \p Alloca is not in the same address space as \p DestLangAS, insert an
+  /// address space cast and return a new RawAddress based on this value.
+  RawAddress MaybeCastStackAddressSpace(RawAddress Alloca, LangAS DestLangAS,
+                                        llvm::Value *ArraySize = nullptr);
+
+public:
   /// CreateTempAlloca - This creates an alloca and inserts it into the entry
   /// block if \p ArraySize is nullptr, otherwise inserts it at the current
   /// insertion point of the builder. The caller is responsible for setting an