[OpenMP] Replace pointer comparison with `isSharedMemPtr` check

[jdoerfert]

Summary

The pointer comparison was causing confusion for capture tracking, let's
avoid confusion.

[jhuber6]

LG

openmp/libomptarget/DeviceRTL/include/Utils.h

 template <typename Ty1, typename Ty2> inline Ty1 align_up(Ty1 V, Ty2 Align) {
   return ((V + Ty1(Align) - 1) / Ty1(Align)) * Ty1(Align);
 }
 /// Return \p V aligned "downwards" according to \p Align.
 template <typename Ty1, typename Ty2> inline Ty1 align_down(Ty1 V, Ty2 Align) {
   return V - V % Align;
 }
 
+/// Return true iff \p Ptr is pointing into shared (local) memory (AS(3)).
+bool isSharedMemPtr(void *Ptr);
+
 /// A  pointer variable that has by design an `undef` value. Use with care.
 __attribute__((loader_uninitialized)) static void *const UndefPtr;
 
 #define OMP_LIKELY(EXPR) __builtin_expect((bool)(EXPR), true)
 #define OMP_UNLIKELY(EXPR) __builtin_expect((bool)(EXPR), false)
 
 } // namespace utils
 } // namespace _OMP
 
 #pragma omp end declare target
 
 #endif

openmp/libomptarget/DeviceRTL/src/State.cpp

 //===------ State.cpp - OpenMP State & ICV interface ------------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//
 
 #include "State.h"
 #include "Configuration.h"
 #include "Debug.h"
 #include "Interface.h"
 #include "Synchronization.h"
 #include "Types.h"
+#include "Utils.h"
 
 using namespace _OMP;
 
 #pragma omp begin declare target device_type(nohost)
 
 /// Memory implementation
 ///
 ///{
void *GlobalMemory = memory::allocGlobal(
                     "shared memory stack memory!");
   ASSERT(GlobalMemory != nullptr && "nullptr returned by malloc!");
 
   return GlobalMemory;
 }
 
 void SharedMemorySmartStackTy::pop(void *Ptr, uint32_t Bytes) {
   uint64_t AlignedBytes = utils::align_up(Bytes, Alignment);
-  if (Ptr >= &Data[0] && Ptr < &Data[state::SharedScratchpadSize]) {
+  if (utils::isSharedMemPtr(Ptr)) {
     int TId = mapping::getThreadIdInBlock();
     Usage[TId] -= AlignedBytes;
     return;
   }
   memory::freeGlobal(Ptr, "Slow path shared memory deallocation");
 }
 
 } // namespace

openmp/libomptarget/DeviceRTL/src/Utils.cpp

__attribute__((weak, optnone, cold)) KEEP_ALIVE void keepAlive() {
   __kmpc_get_warp_size();
   __kmpc_barrier_simple_spmd(nullptr, 0);
   __kmpc_barrier_simple_generic(nullptr, 0);
 }
 } // namespace _OMP
 
 namespace impl {
 
+bool isSharedMemPtr(const void *Ptr) { return false; }
 void Unpack(uint64_t Val, uint32_t *LowBits, uint32_t *HighBits);
 uint64_t Pack(uint32_t LowBits, uint32_t HighBits);
 
 /// AMDGCN Implementation
 ///
 ///{
 #pragma omp begin declare variant match(device = {arch(amdgcn)})
 
 void Unpack(uint64_t Val, uint32_t *LowBits, uint32_t *HighBits) {
   static_assert(sizeof(unsigned long) == 8, "");
   *LowBits = (uint32_t)(Val & 0x00000000FFFFFFFFUL);
   *HighBits = (uint32_t)((Val & 0xFFFFFFFF00000000UL) >> 32);
 }
 
 uint64_t Pack(uint32_t LowBits, uint32_t HighBits) {
   return (((uint64_t)HighBits) << 32) | (uint64_t)LowBits;
 }
 
 #pragma omp end declare variant
+///}
 
 /// NVPTX Implementation
 ///
 ///{
 #pragma omp begin declare variant match(                                       \
     device = {arch(nvptx, nvptx64)}, implementation = {extension(match_any)})
 
 void Unpack(uint64_t Val, uint32_t *LowBits, uint32_t *HighBits) {
   uint32_t LowBitsLocal, HighBitsLocal;
   asm("mov.b64 {%0,%1}, %2;"
       : "=r"(LowBitsLocal), "=r"(HighBitsLocal)
       : "l"(Val));
   *LowBits = LowBitsLocal;
   *HighBits = HighBitsLocal;
 }
 
 uint64_t Pack(uint32_t LowBits, uint32_t HighBits) {
   uint64_t Val;
   asm("mov.b64 %0, {%1,%2};" : "=l"(Val) : "r"(LowBits), "r"(HighBits));
   return Val;
 }
 
 #pragma omp end declare variant
+///}
 
 int32_t shuffle(uint64_t Mask, int32_t Var, int32_t SrcLane);
 int32_t shuffleDown(uint64_t Mask, int32_t Var, uint32_t LaneDelta,
                     int32_t Width);
 
 /// AMDGCN Implementation
 ///
 ///{

 int32_t shuffleDown(uint64_t Mask, int32_t Var, uint32_t LaneDelta,
                     int32_t Width) {
   int Self = mapping::getThreadIdInWarp();
   int Index = Self + LaneDelta;
   Index = (int)(LaneDelta + (Self & (Width - 1))) >= Width ? Self : Index;
   return __builtin_amdgcn_ds_bpermute(Index << 2, Var);
 }
 
+bool isSharedMemPtr(const void * Ptr) {
+  return __builtin_amdgcn_is_shared((const __attribute__((address_space(0))) void *)Ptr);
+}
 #pragma omp end declare variant
 ///}
 
 /// NVPTX Implementation
 ///
 ///{
 #pragma omp begin declare variant match(                                       \
     device = {arch(nvptx, nvptx64)}, implementation = {extension(match_any)})
 
 int32_t shuffle(uint64_t Mask, int32_t Var, int32_t SrcLane) {
   return __nvvm_shfl_sync_idx_i32(Mask, Var, SrcLane, 0x1f);
 }
 
 int32_t shuffleDown(uint64_t Mask, int32_t Var, uint32_t Delta, int32_t Width) {
   int32_t T = ((mapping::getWarpSize() - Width) << 8) | 0x1f;
   return __nvvm_shfl_sync_down_i32(Mask, Var, Delta, T);
 }
 
+bool isSharedMemPtr(const void *Ptr) { return __nvvm_isspacep_shared(Ptr); }
+
 #pragma omp end declare variant
+///}
 } // namespace impl
 
 uint64_t utils::pack(uint32_t LowBits, uint32_t HighBits) {
   return impl::Pack(LowBits, HighBits);
 }
 
 void utils::unpack(uint64_t Val, uint32_t &LowBits, uint32_t &HighBits) {
   impl::Unpack(Val, &LowBits, &HighBits);
 }
 
 int32_t utils::shuffle(uint64_t Mask, int32_t Var, int32_t SrcLane) {
   return impl::shuffle(Mask, Var, SrcLane);
 }
 
 int32_t utils::shuffleDown(uint64_t Mask, int32_t Var, uint32_t Delta,
                            int32_t Width) {
   return impl::shuffleDown(Mask, Var, Delta, Width);
 }
 
+bool utils::isSharedMemPtr(void *Ptr) { return impl::isSharedMemPtr(Ptr); }
+
 extern "C" {
 int32_t __kmpc_shuffle_int32(int32_t Val, int16_t Delta, int16_t SrcLane) {
   FunctionTracingRAII();
   return impl::shuffleDown(lanes::All, Val, Delta, SrcLane);
 }
 
 int64_t __kmpc_shuffle_int64(int64_t Val, int16_t Delta, int16_t Width) {
   FunctionTracingRAII();