[OpenMP] Add new OpenMP globalization functions to library info

[jhuber6]

Summary

The changes to globalization introduced in D97680 created two new functions to
push / pop shareably memory on the GPU, __kmpc_alloc_shared and
__kmpc_free_shared. This patch adds these new runtime functions to the
library info so they can be used by the HeapToStack attributor interface. This
optimization replaces malloc / free pairs with stack memory if legal.

[tianshilei1992]

LG.

[jhuber6]

Adding __kmpc_free_shared as a valid free built-in function so it works properly with AAHeapToStack.

llvm/include/llvm/Analysis/TargetLibraryInfo.def

 TLI_DEFINE_ENUM_INTERNAL(expl_finite)
 TLI_DEFINE_STRING_INTERNAL("__expl_finite")
 /// int __isoc99_scanf (const char *format, ...)
 TLI_DEFINE_ENUM_INTERNAL(dunder_isoc99_scanf)
 TLI_DEFINE_STRING_INTERNAL("__isoc99_scanf")
 /// int __isoc99_sscanf(const char *s, const char *format, ...)
 TLI_DEFINE_ENUM_INTERNAL(dunder_isoc99_sscanf)
 TLI_DEFINE_STRING_INTERNAL("__isoc99_sscanf")
+/// void __kmpc_alloc_shared(size_t nbyte);
+TLI_DEFINE_ENUM_INTERNAL(__kmpc_alloc_shared)
+TLI_DEFINE_STRING_INTERNAL("__kmpc_alloc_shared")
+/// void __kmpc_free_shared(void *ptr);
+TLI_DEFINE_ENUM_INTERNAL(__kmpc_free_shared)
+TLI_DEFINE_STRING_INTERNAL("__kmpc_free_shared")
 /// double __log10_finite(double x);
 TLI_DEFINE_ENUM_INTERNAL(log10_finite)
 TLI_DEFINE_STRING_INTERNAL("__log10_finite")
 /// float __log10f_finite(float x);
 TLI_DEFINE_ENUM_INTERNAL(log10f_finite)
 TLI_DEFINE_STRING_INTERNAL("__log10f_finite")
 /// long double __log10l_finite(long double x);
 TLI_DEFINE_ENUM_INTERNAL(log10l_finite)

llvm/lib/Analysis/MemoryBuiltins.cpp

struct AllocFnsTy {
   unsigned NumParams;
   // First and Second size parameters (or -1 if unused)
   int FstParam, SndParam;
 };
 
 // FIXME: certain users need more information. E.g., SimplifyLibCalls needs to
 // know which functions are nounwind, noalias, nocapture parameters, etc.
 static const std::pair<LibFunc, AllocFnsTy> AllocationFnData[] = {
   {LibFunc_malloc,              {MallocLike,  1, 0,  -1}},
   {LibFunc_vec_malloc,          {MallocLike,  1, 0,  -1}},
   {LibFunc_valloc,              {MallocLike,  1, 0,  -1}},
   {LibFunc_Znwj,                {OpNewLike,   1, 0,  -1}}, // new(unsigned int)
   {LibFunc_ZnwjRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new(unsigned int, nothrow)
   {LibFunc_ZnwjSt11align_val_t, {OpNewLike,   2, 0,  -1}}, // new(unsigned int, align_val_t)
   {LibFunc_ZnwjSt11align_val_tRKSt9nothrow_t, // new(unsigned int, align_val_t, nothrow)
                                 {MallocLike,  3, 0,  -1}},
   {LibFunc_Znwm,                {OpNewLike,   1, 0,  -1}}, // new(unsigned long)
static const std::pair<LibFunc, AllocFnsTy> AllocationFnData[] = {
   {LibFunc_aligned_alloc,       {AlignedAllocLike, 2, 1,  -1}},
   {LibFunc_memalign,            {AlignedAllocLike, 2, 1,  -1}},
   {LibFunc_calloc,              {CallocLike,  2, 0,   1}},
   {LibFunc_vec_calloc,          {CallocLike,  2, 0,   1}},
   {LibFunc_realloc,             {ReallocLike, 2, 1,  -1}},
   {LibFunc_vec_realloc,         {ReallocLike, 2, 1,  -1}},
   {LibFunc_reallocf,            {ReallocLike, 2, 1,  -1}},
   {LibFunc_strdup,              {StrDupLike,  1, -1, -1}},
-  {LibFunc_strndup,             {StrDupLike,  2, 1,  -1}}
+  {LibFunc_strndup,             {StrDupLike,  2, 1,  -1}},
+  {LibFunc___kmpc_alloc_shared, {MallocLike, 1, 0, -1}}
   // TODO: Handle "int posix_memalign(void **, size_t, size_t)"
 };
 
 static const Function *getCalledFunction(const Value *V, bool LookThroughBitCast,
                                          bool &IsNoBuiltin) {
   // Don't care about intrinsics in this case.
   if (isa<IntrinsicInst>(V))
     return nullptr;

llvm/lib/Analysis/TargetLibraryInfo.cpp

if (T.isAMDGPU())
     TLI.disableAllFunctions();
 
   // There are no library implementations of memcpy and memset for AMD gpus and
   // these can be difficult to lower in the backend.
   if (T.isAMDGPU()) {
     TLI.setUnavailable(LibFunc_memcpy);
     TLI.setUnavailable(LibFunc_memset);
     TLI.setUnavailable(LibFunc_memset_pattern16);
+    TLI.setAvailable(llvm::LibFunc___kmpc_alloc_shared);
+    TLI.setAvailable(llvm::LibFunc___kmpc_free_shared);
     return;
   }
 
   // memset_pattern16 is only available on iOS 3.0 and Mac OS X 10.5 and later.
   // All versions of watchOS support it.
   if (T.isMacOSX()) {
     // available IO unlocked variants on Mac OS X
     TLI.setAvailable(LibFunc_getc_unlocked);
if (T.isNVPTX()) {
     // TODO: We could enable the following two according to [0] but we haven't
     //       done an evaluation wrt. the performance implications.
     // [0]
     // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#dynamic-global-memory-allocation-and-operations
     //
     //    TLI.setAvailable(llvm::LibFunc_memcpy);
     //    TLI.setAvailable(llvm::LibFunc_memset);
 
+    TLI.setAvailable(llvm::LibFunc___kmpc_alloc_shared);
+    TLI.setAvailable(llvm::LibFunc___kmpc_free_shared);
   } else {
     TLI.setUnavailable(LibFunc_nvvm_reflect);
   }
 
   // These vec_malloc/free routines are only available on AIX.
   if (!T.isOSAIX()) {
     TLI.setUnavailable(LibFunc_vec_calloc);
     TLI.setUnavailable(LibFunc_vec_malloc);
return NumParams == 5 && FTy.getParamType(0)->isPointerTy() &&
            FTy.getParamType(4)->isPointerTy() &&
            FTy.getReturnType()->isIntegerTy(32);
 
   case LibFunc_setitimer:
     return (NumParams == 3 && FTy.getParamType(1)->isPointerTy() &&
             FTy.getParamType(2)->isPointerTy());
   case LibFunc_system:
     return (NumParams == 1 && FTy.getParamType(0)->isPointerTy());
+  case LibFunc___kmpc_alloc_shared:
   case LibFunc_malloc:
   case LibFunc_vec_malloc:
     return (NumParams == 1 && FTy.getReturnType()->isPointerTy());
   case LibFunc_memcmp:
     return (NumParams == 3 && FTy.getReturnType()->isIntegerTy(32) &&
             FTy.getParamType(0)->isPointerTy() &&
             FTy.getParamType(1)->isPointerTy());
 
bool TargetLibraryInfoImpl::isValidProtoForLibFunc(const FunctionType &FTy,
   case LibFunc_funlockfile:
   case LibFunc_getc:
   case LibFunc_getc_unlocked:
   case LibFunc_getlogin_r:
   case LibFunc_mkdir:
   case LibFunc_mktime:
   case LibFunc_times:
   case LibFunc_vec_free:
+  case LibFunc___kmpc_free_shared:
     return (NumParams != 0 && FTy.getParamType(0)->isPointerTy());
 
   case LibFunc_fopen:
     return (NumParams == 2 && FTy.getReturnType()->isPointerTy() &&
             FTy.getParamType(0)->isPointerTy() &&
             FTy.getParamType(1)->isPointerTy());
   case LibFunc_fork:
     return (NumParams == 0 && FTy.getReturnType()->isIntegerTy(32));

llvm/unittests/Analysis/TargetLibraryInfoTest.cpp

parseAssembly(
       "declare float @exp10f(float)\n"
       "declare x86_fp80 @exp10l(x86_fp80)\n"
 
       // These printf variants have the same prototype as the non-'i' versions.
       "declare i32 @fiprintf(%struct*, i8*, ...)\n"
       "declare i32 @iprintf(i8*, ...)\n"
       "declare i32 @siprintf(i8*, i8*, ...)\n"
 
       // __small_printf variants have the same prototype as the non-'i' versions.
       "declare i32 @__small_fprintf(%struct*, i8*, ...)\n"
       "declare i32 @__small_printf(i8*, ...)\n"
       "declare i32 @__small_sprintf(i8*, i8*, ...)\n"
 
       "declare i32 @htonl(i32)\n"
       "declare i16 @htons(i16)\n"
       "declare i32 @ntohl(i32)\n"
       "declare i16 @ntohs(i16)\n"
parseAssembly(
       "declare float @__sinhf_finite(float)\n"
       "declare x86_fp80 @__sinhl_finite(x86_fp80)\n"
 
       // These functions are aix vec allocation/free routines
       "declare i8* @vec_calloc(i64, i64)\n"
       "declare i8* @vec_malloc(i64)\n"
       "declare i8* @vec_realloc(i8*, i64)\n"
       "declare void @vec_free(i8*)\n"
+
+      // These functions are OpenMP Offloading allocation / free routines
+      "declare i8* @__kmpc_alloc_shared(i64)\n"
+      "declare void @__kmpc_free_shared(i8*)\n"
       );
 
   for (unsigned FI = 0; FI != LibFunc::NumLibFuncs; ++FI) {
     LibFunc LF = (LibFunc)FI;
     // Make sure everything is available; we're not testing target defaults.
     TLII.setAvailable(LF);
     Function *F = M->getFunction(TLI.getName(LF));
     EXPECT_TRUE(isLibFunc(F, LF));
   }
 }