[OpenMP] Change target memory tests to use allocators

[jhuber6]

Summary

The target allocators have been supported for NVPTX offloading for
awhile. The tests should use the allocators instead of calling the
functions manually. Also the comments indicating these being a preview
should be removed.

[jdoerfert]

openmp/libomptarget/test/api/omp_device_managed_memory.c

11

This sounds so confusing to me. Shared mem alloc, is not using shared/local memory but managed memory, right?
Just to make sure. Maybe leave a comment.

openmp/libomptarget/test/api/omp_host_pinned_memory.c

33

Keep the old tests around please. No need to drop support for the function just yet.

[jhuber6]

openmp/libomptarget/test/api/omp_device_managed_memory.c

11

That's how they were described in the original patch, it's supposed to be "shared" between the host and device I guess.

openmp/libomptarget/test/api/omp_host_pinned_memory.c

33

But the allocators use the function internally, this tests both that function and the allocator

[jhuber6]

Making suggested changes.

[jdoerfert]

LG

openmp/libomptarget/test/api/omp_device_managed_memory.c

 // RUN: %libomptarget-compile-run-and-check-nvptx64-nvidia-cuda
 // REQUIRES: nvptx64-nvidia-cuda
 
 #include <omp.h>
 #include <stdio.h>
 
-void *llvm_omp_target_alloc_shared(size_t, int);
-
 int main() {
   const int N = 64;
-  const int device = omp_get_default_device();
 
-  int *shared_ptr = llvm_omp_target_alloc_shared(N * sizeof(int), device);
+  int *shared_ptr =
+      omp_alloc(N * sizeof(int), llvm_omp_target_shared_mem_alloc);
 
-#pragma omp target teams distribute parallel for device(device) \
-            is_device_ptr(shared_ptr)
+#pragma omp target teams distribute parallel for is_device_ptr(shared_ptr)
   for (int i = 0; i < N; ++i) {
     shared_ptr[i] = 1;
   }
 
   int sum = 0;
   for (int i = 0; i < N; ++i)
     sum += shared_ptr[i];
 
-  omp_target_free(shared_ptr, device);
   // CHECK: PASS
   if (sum == N)
-    printf ("PASS\n");
+    printf("PASS\n");
+
+  omp_free(shared_ptr, llvm_omp_target_shared_mem_alloc);
 }

openmp/libomptarget/test/api/omp_device_memory.c

+// RUN: %libomptarget-compile-run-and-check-nvptx64-nvidia-cuda
+// REQUIRES: nvptx64-nvidia-cuda
+
+#include <omp.h>
+#include <stdio.h>
+
+int main() {
+  const int N = 64;
+
+  int *device_ptr =
+      omp_alloc(N * sizeof(int), llvm_omp_target_device_mem_alloc);
+
+#pragma omp target teams distribute parallel for is_device_ptr(device_ptr)
+  for (int i = 0; i < N; ++i) {
+    device_ptr[i] = 1;
+  }
+
+  int sum = 0;
+#pragma omp target reduction(+ : sum) is_device_ptr(device_ptr)
+  for (int i = 0; i < N; ++i)
+    sum += device_ptr[i];
+
+  // CHECK: PASS
+  if (sum == N)
+    printf("PASS\n");
+
+  omp_free(device_ptr, llvm_omp_target_device_mem_alloc);
+}

openmp/libomptarget/test/api/omp_host_pinned_memory.c

 // RUN: %libomptarget-compile-run-and-check-nvptx64-nvidia-cuda
 // REQUIRES: nvptx64-nvidia-cuda
 
 #include <omp.h>
 #include <stdio.h>
 
-// Allocate pinned memory on the host
-void *llvm_omp_target_alloc_host(size_t, int);
-
 int main() {
   const int N = 64;
-  const int device = omp_get_default_device();
-  const int host = omp_get_initial_device();
 
-  int *hst_ptr = llvm_omp_target_alloc_host(N * sizeof(int), device);
+  int *hst_ptr = omp_alloc(N * sizeof(int), llvm_omp_target_host_mem_alloc);
 
   for (int i = 0; i < N; ++i)
     hst_ptr[i] = 2;
 
-#pragma omp target teams distribute parallel for device(device) \
-           map(tofrom:hst_ptr[0 : N])
+#pragma omp target teams distribute parallel for map(tofrom : hst_ptr [0:N])
   for (int i = 0; i < N; ++i)
     hst_ptr[i] -= 1;
 
   int sum = 0;
   for (int i = 0; i < N; ++i)
     sum += hst_ptr[i];
 
-  omp_target_free(hst_ptr, device);
+  omp_free(hst_ptr, llvm_omp_target_shared_mem_alloc);
   // CHECK: PASS
   if (sum == N)
-    printf ("PASS\n");
+    printf("PASS\n");
 }

openmp/runtime/src/dllexports

kmp_set_disp_num_buffers                    890
     omp_default_mem_alloc                  DATA
     omp_large_cap_mem_alloc                DATA
     omp_const_mem_alloc                    DATA
     omp_high_bw_mem_alloc                  DATA
     omp_low_lat_mem_alloc                  DATA
     omp_cgroup_mem_alloc                   DATA
     omp_pteam_mem_alloc                    DATA
     omp_thread_mem_alloc                   DATA
-    # Preview of target memory support
+
     llvm_omp_target_host_mem_alloc         DATA
     llvm_omp_target_shared_mem_alloc       DATA
     llvm_omp_target_device_mem_alloc       DATA
 
     omp_default_mem_space                  DATA
     omp_large_cap_mem_space                DATA
     omp_const_mem_space                    DATA
     omp_high_bw_mem_space                  DATA
     omp_low_lat_mem_space                  DATA
-    # Preview of target memory support
+
     llvm_omp_target_host_mem_space         DATA
     llvm_omp_target_shared_mem_space       DATA
     llvm_omp_target_device_mem_space       DATA
 
 %ifndef stub
     # Ordinals between 900 and 999 are reserved
 
     # Ordinals between 1000 and 1999 are reserved

openmp/runtime/src/include/omp.h.var

#   if defined(_WIN32)
     extern __KMP_IMP omp_allocator_handle_t const omp_default_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_large_cap_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_const_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_high_bw_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_low_lat_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_cgroup_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_pteam_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_thread_mem_alloc;
-    /* Preview of target memory support */
+
     extern __KMP_IMP omp_allocator_handle_t const llvm_omp_target_host_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const llvm_omp_target_shared_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const llvm_omp_target_device_mem_alloc;
 
     typedef omp_uintptr_t omp_memspace_handle_t;
     extern __KMP_IMP omp_memspace_handle_t const omp_default_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const omp_large_cap_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const omp_const_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const omp_high_bw_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const omp_low_lat_mem_space;
-    /* Preview of target memory support */
+
     extern __KMP_IMP omp_memspace_handle_t const llvm_omp_target_host_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const llvm_omp_target_shared_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const llvm_omp_target_device_mem_space;
 #   else
 #       if __cplusplus >= 201103
     typedef enum omp_allocator_handle_t : omp_uintptr_t
 #       else
     typedef enum omp_allocator_handle_t
 #       endif
     {
       omp_null_allocator = 0,
       omp_default_mem_alloc = 1,
       omp_large_cap_mem_alloc = 2,
       omp_const_mem_alloc = 3,
       omp_high_bw_mem_alloc = 4,
       omp_low_lat_mem_alloc = 5,
       omp_cgroup_mem_alloc = 6,
       omp_pteam_mem_alloc = 7,
       omp_thread_mem_alloc = 8,
-      /* Preview of target memory support */
       llvm_omp_target_host_mem_alloc = 100,
       llvm_omp_target_shared_mem_alloc = 101,
       llvm_omp_target_device_mem_alloc = 102,
       KMP_ALLOCATOR_MAX_HANDLE = UINTPTR_MAX
     } omp_allocator_handle_t;
 #       if __cplusplus >= 201103
     typedef enum omp_memspace_handle_t : omp_uintptr_t
 #       else
     typedef enum omp_memspace_handle_t
 #       endif
     {
       omp_default_mem_space = 0,
       omp_large_cap_mem_space = 1,
       omp_const_mem_space = 2,
       omp_high_bw_mem_space = 3,
       omp_low_lat_mem_space = 4,
-      /* Preview of target memory support */
       llvm_omp_target_host_mem_space = 100,
       llvm_omp_target_shared_mem_space = 101,
       llvm_omp_target_device_mem_space = 102,
       KMP_MEMSPACE_MAX_HANDLE = UINTPTR_MAX
     } omp_memspace_handle_t;
 #   endif
     extern omp_allocator_handle_t __KAI_KMPC_CONVENTION omp_init_allocator(omp_memspace_handle_t m,
                                                        int ntraits, omp_alloctrait_t traits[]);

openmp/runtime/src/include/omp_lib.h.var

!
       integer(kind=omp_allocator_handle_kind)omp_low_lat_mem_alloc
       parameter(omp_low_lat_mem_alloc=5)
       integer(kind=omp_allocator_handle_kind)omp_cgroup_mem_alloc
       parameter(omp_cgroup_mem_alloc=6)
       integer(kind=omp_allocator_handle_kind)omp_pteam_mem_alloc
       parameter(omp_pteam_mem_alloc=7)
       integer(kind=omp_allocator_handle_kind)omp_thread_mem_alloc
       parameter(omp_thread_mem_alloc=8)
-      ! Preview of target memory support
+
       integer(omp_allocator_handle_kind)llvm_omp_target_host_mem_alloc
       parameter(llvm_omp_target_host_mem_alloc=100)
       integer(omp_allocator_handle_kind)llvm_omp_target_shared_mem_alloc
       parameter(llvm_omp_target_shared_mem_alloc=101)
       integer(omp_allocator_handle_kind)llvm_omp_target_device_mem_alloc
       parameter(llvm_omp_target_device_mem_alloc=102)
 
       integer(kind=omp_memspace_handle_kind)omp_default_mem_space
       parameter(omp_default_mem_space=0)
       integer(kind=omp_memspace_handle_kind)omp_large_cap_mem_space
       parameter(omp_large_cap_mem_space=1)
       integer(kind=omp_memspace_handle_kind)omp_const_mem_space
       parameter(omp_const_mem_space=2)
       integer(kind=omp_memspace_handle_kind)omp_high_bw_mem_space
       parameter(omp_high_bw_mem_space=3)
       integer(kind=omp_memspace_handle_kind)omp_low_lat_mem_space
       parameter(omp_low_lat_mem_space=4)
-      ! Preview of target memory support
+
       integer(omp_memspace_handle_kind)llvm_omp_target_host_mem_space
       parameter(llvm_omp_target_host_mem_space=100)
       integer(omp_memspace_handle_kind)llvm_omp_target_shared_mem_space
       parameter(llvm_omp_target_shared_mem_space=101)
       integer(omp_memspace_handle_kind)llvm_omp_target_device_mem_space
       parameter(llvm_omp_target_device_mem_space=102)
 
       integer(kind=omp_pause_resource_kind)omp_pause_resume

openmp/runtime/src/include/omp_lib.f90.var

!
         integer (kind=omp_allocator_handle_kind), parameter :: omp_default_mem_alloc = 1
         integer (kind=omp_allocator_handle_kind), parameter :: omp_large_cap_mem_alloc = 2
         integer (kind=omp_allocator_handle_kind), parameter :: omp_const_mem_alloc = 3
         integer (kind=omp_allocator_handle_kind), parameter :: omp_high_bw_mem_alloc = 4
         integer (kind=omp_allocator_handle_kind), parameter :: omp_low_lat_mem_alloc = 5
         integer (kind=omp_allocator_handle_kind), parameter :: omp_cgroup_mem_alloc = 6
         integer (kind=omp_allocator_handle_kind), parameter :: omp_pteam_mem_alloc = 7
         integer (kind=omp_allocator_handle_kind), parameter :: omp_thread_mem_alloc = 8
-        ! Preview of target memory support
+
         integer (kind=omp_allocator_handle_kind), parameter :: llvm_omp_target_host_mem_alloc = 100
         integer (kind=omp_allocator_handle_kind), parameter :: llvm_omp_target_shared_mem_alloc = 101
         integer (kind=omp_allocator_handle_kind), parameter :: llvm_omp_target_device_mem_alloc = 102
 
         integer (kind=omp_memspace_handle_kind), parameter :: omp_default_mem_space = 0
         integer (kind=omp_memspace_handle_kind), parameter :: omp_large_cap_mem_space = 1
         integer (kind=omp_memspace_handle_kind), parameter :: omp_const_mem_space = 2
         integer (kind=omp_memspace_handle_kind), parameter :: omp_high_bw_mem_space = 3
         integer (kind=omp_memspace_handle_kind), parameter :: omp_low_lat_mem_space = 4
-        ! Preview of target memory support
+
         integer (kind=omp_memspace_handle_kind), parameter :: llvm_omp_target_host_mem_space = 100
         integer (kind=omp_memspace_handle_kind), parameter :: llvm_omp_target_shared_mem_space = 101
         integer (kind=omp_memspace_handle_kind), parameter :: llvm_omp_target_device_mem_space = 102
 
         integer (kind=omp_pause_resource_kind), parameter :: omp_pause_resume = 0
         integer (kind=omp_pause_resource_kind), parameter :: omp_pause_soft = 1
         integer (kind=omp_pause_resource_kind), parameter :: omp_pause_hard = 2
 

openmp/runtime/src/kmp.h

 #define omp_atv_default ((omp_uintptr_t)-1)
 
 typedef void *omp_memspace_handle_t;
 extern omp_memspace_handle_t const omp_default_mem_space;
 extern omp_memspace_handle_t const omp_large_cap_mem_space;
 extern omp_memspace_handle_t const omp_const_mem_space;
 extern omp_memspace_handle_t const omp_high_bw_mem_space;
 extern omp_memspace_handle_t const omp_low_lat_mem_space;
-// Preview of target memory support
 extern omp_memspace_handle_t const llvm_omp_target_host_mem_space;
 extern omp_memspace_handle_t const llvm_omp_target_shared_mem_space;
 extern omp_memspace_handle_t const llvm_omp_target_device_mem_space;
 
 typedef struct {
   omp_alloctrait_key_t key;
   omp_uintptr_t value;
 } omp_alloctrait_t;
 
 typedef void *omp_allocator_handle_t;
 extern omp_allocator_handle_t const omp_null_allocator;
 extern omp_allocator_handle_t const omp_default_mem_alloc;
 extern omp_allocator_handle_t const omp_large_cap_mem_alloc;
 extern omp_allocator_handle_t const omp_const_mem_alloc;
 extern omp_allocator_handle_t const omp_high_bw_mem_alloc;
 extern omp_allocator_handle_t const omp_low_lat_mem_alloc;
 extern omp_allocator_handle_t const omp_cgroup_mem_alloc;
 extern omp_allocator_handle_t const omp_pteam_mem_alloc;
 extern omp_allocator_handle_t const omp_thread_mem_alloc;
-// Preview of target memory support
 extern omp_allocator_handle_t const llvm_omp_target_host_mem_alloc;
 extern omp_allocator_handle_t const llvm_omp_target_shared_mem_alloc;
 extern omp_allocator_handle_t const llvm_omp_target_device_mem_alloc;
 extern omp_allocator_handle_t const kmp_max_mem_alloc;
 extern omp_allocator_handle_t __kmp_def_allocator;
 
 // end of duplicate type definitions from omp.h
 #endif

openmp/runtime/src/kmp_alloc.cpp

 static void **mk_hbw_interleave;
 static void **mk_hbw_preferred;
 static void **mk_hugetlb;
 static void **mk_hbw_hugetlb;
 static void **mk_hbw_preferred_hugetlb;
 static void **mk_dax_kmem;
 static void **mk_dax_kmem_all;
 static void **mk_dax_kmem_preferred;
-// Preview of target memory support
 static void *(*kmp_target_alloc_host)(size_t size, int device);
 static void *(*kmp_target_alloc_shared)(size_t size, int device);
 static void *(*kmp_target_alloc_device)(size_t size, int device);
 static void *(*kmp_target_free)(void *ptr, int device);
 static bool __kmp_target_mem_available;
 #define KMP_IS_TARGET_MEM_SPACE(MS)                                            \
   (MS == llvm_omp_target_host_mem_space ||                                     \
    MS == llvm_omp_target_shared_mem_space ||                                   \
#if KMP_OS_UNIX && KMP_DYNAMIC_LIB
   mk_hugetlb = NULL;
   mk_hbw_hugetlb = NULL;
   mk_hbw_preferred_hugetlb = NULL;
   mk_dax_kmem = NULL;
   mk_dax_kmem_all = NULL;
   mk_dax_kmem_preferred = NULL;
 #endif
 }
-// Preview of target memory support
+
 void __kmp_init_target_mem() {
   *(void **)(&kmp_target_alloc_host) = KMP_DLSYM("llvm_omp_target_alloc_host");
   *(void **)(&kmp_target_alloc_shared) =
       KMP_DLSYM("llvm_omp_target_alloc_shared");
   *(void **)(&kmp_target_alloc_device) =
       KMP_DLSYM("llvm_omp_target_alloc_device");
   *(void **)(&kmp_target_free) = KMP_DLSYM("omp_target_free");
   __kmp_target_mem_available = kmp_target_alloc_host &&

openmp/runtime/src/kmp_global.cpp

 omp_allocator_handle_t const omp_low_lat_mem_alloc =
     (omp_allocator_handle_t const)5;
 omp_allocator_handle_t const omp_cgroup_mem_alloc =
     (omp_allocator_handle_t const)6;
 omp_allocator_handle_t const omp_pteam_mem_alloc =
     (omp_allocator_handle_t const)7;
 omp_allocator_handle_t const omp_thread_mem_alloc =
     (omp_allocator_handle_t const)8;
-// Preview of target memory support
+
 omp_allocator_handle_t const llvm_omp_target_host_mem_alloc =
     (omp_allocator_handle_t const)100;
 omp_allocator_handle_t const llvm_omp_target_shared_mem_alloc =
     (omp_allocator_handle_t const)101;
 omp_allocator_handle_t const llvm_omp_target_device_mem_alloc =
     (omp_allocator_handle_t const)102;
 omp_allocator_handle_t const kmp_max_mem_alloc =
     (omp_allocator_handle_t const)1024;
 omp_allocator_handle_t __kmp_def_allocator = omp_default_mem_alloc;
 
 omp_memspace_handle_t const omp_default_mem_space =
     (omp_memspace_handle_t const)0;
 omp_memspace_handle_t const omp_large_cap_mem_space =
     (omp_memspace_handle_t const)1;
 omp_memspace_handle_t const omp_const_mem_space =
     (omp_memspace_handle_t const)2;
 omp_memspace_handle_t const omp_high_bw_mem_space =
     (omp_memspace_handle_t const)3;
 omp_memspace_handle_t const omp_low_lat_mem_space =
     (omp_memspace_handle_t const)4;
-// Preview of target memory support
+
 omp_memspace_handle_t const llvm_omp_target_host_mem_space =
     (omp_memspace_handle_t const)100;
 omp_memspace_handle_t const llvm_omp_target_shared_mem_space =
     (omp_memspace_handle_t const)101;
 omp_memspace_handle_t const llvm_omp_target_device_mem_space =
     (omp_memspace_handle_t const)102;
 
 /* This check ensures that the compiler is passing the correct data type for the

openmp/runtime/src/kmp_stub.cpp

 omp_allocator_handle_t const omp_low_lat_mem_alloc =
     (omp_allocator_handle_t const)5;
 omp_allocator_handle_t const omp_cgroup_mem_alloc =
     (omp_allocator_handle_t const)6;
 omp_allocator_handle_t const omp_pteam_mem_alloc =
     (omp_allocator_handle_t const)7;
 omp_allocator_handle_t const omp_thread_mem_alloc =
     (omp_allocator_handle_t const)8;
-// Preview of target memory support
 omp_allocator_handle_t const llvm_omp_target_host_mem_alloc =
     (omp_allocator_handle_t const)100;
 omp_allocator_handle_t const llvm_omp_target_shared_mem_alloc =
     (omp_allocator_handle_t const)101;
 omp_allocator_handle_t const llvm_omp_target_device_mem_alloc =
     (omp_allocator_handle_t const)102;
 
 omp_memspace_handle_t const omp_default_mem_space =
     (omp_memspace_handle_t const)0;
 omp_memspace_handle_t const omp_large_cap_mem_space =
     (omp_memspace_handle_t const)1;
 omp_memspace_handle_t const omp_const_mem_space =
     (omp_memspace_handle_t const)2;
 omp_memspace_handle_t const omp_high_bw_mem_space =
     (omp_memspace_handle_t const)3;
 omp_memspace_handle_t const omp_low_lat_mem_space =
     (omp_memspace_handle_t const)4;
-// Preview of target memory support
 omp_memspace_handle_t const llvm_omp_target_host_mem_space =
     (omp_memspace_handle_t const)100;
 omp_memspace_handle_t const llvm_omp_target_shared_mem_space =
     (omp_memspace_handle_t const)101;
 omp_memspace_handle_t const llvm_omp_target_device_mem_space =
     (omp_memspace_handle_t const)102;
 #endif /* KMP_OS_WINDOWS */