[libomptarget] Implement get_device_num for amdgcn, nvptx

[JonChesterfield]

Summary

[libomptarget] Implement get_device_num for amdgcn, nvptx

Adds a field to the existing environment struct which is used for debug_level.

[JonChesterfield]

openmp/libomptarget/plugins/amdgpu/src/rtl.cpp

529–530

I forgot to remove the extra fields from this to match the deviceRTL in trunk when adding amdgpu/src/rtl.cpp. This brings them back in sync.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

83

We have a dedicated header now, in "../../../deviceRTLs/common/device_environment.h".

Options:
1/ keep the copy & paste, plugins & deviceRTL are totally separate
2/ Include from the deviceRTL
3/ Move header to somewhere else, openmp/libomptarget/include looks reasonable

[tianshilei1992]

Overall LGTM.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

83

Option 3 sounds more reasonable. We might have data structures used by both device and host runtime.

[tianshilei1992]

Please also add test case(s) for this. However, it actually depends on the native runtime that which device should be 0, 1, etc., which is probably hard to determine in OpenMP.

[JonChesterfield]

Open to suggestions on testing. Calling it from a target region on amdgcn with multiple GPUs works. I don't have a cuda machine with multiple cards. A machine with one target returns 0 with or without this patch.

Do we have any runtime tests for openmp in tree? Would need a CI machine with GPUs

[tianshilei1992]

In D91713#2403572, @JonChesterfield wrote:

Open to suggestions on testing. Calling it from a target region on amdgcn with multiple GPUs works. I don't have a cuda machine with multiple cards. A machine with one target returns 0 with or without this patch.

Do we have any runtime tests for openmp in tree? Would need a CI machine with GPUs

I wrote a simple program to test this patch and it works as expected. Please add the test case into this patch.

[JonChesterfield]

• Use header from include, add test case

[JonChesterfield]

• Use header from include, add test case
• rebase

[JonChesterfield]

Apologies to the hundreds of reviewers who have just been spammed.

Today I learned that 'arc diff --update' is not relative to the patch one is updating, but is relative to master as exists on github. Also that phabricator will not let one delete some automatically added reviewers before fixing the diff.

[JonChesterfield]

This needs an implementation for x86_64 offloading or a change to the test to only run on nvptx. Intend to do the former.

openmp/libomptarget/deviceRTLs/amdgcn/CMakeLists.txt

set(cuda_sources
   ${devicertl_base_directory}/common/src/sync.cu
   ${devicertl_base_directory}/common/src/task.cu)
 
 set(h_files
   ${CMAKE_CURRENT_SOURCE_DIR}/src/amdgcn_interface.h
   ${CMAKE_CURRENT_SOURCE_DIR}/src/hip_atomics.h
   ${CMAKE_CURRENT_SOURCE_DIR}/src/target_impl.h
   ${devicertl_base_directory}/common/debug.h
-  ${devicertl_base_directory}/common/device_environment.h
   ${devicertl_base_directory}/common/omptarget.h
   ${devicertl_base_directory}/common/omptargeti.h
   ${devicertl_base_directory}/common/state-queue.h
   ${devicertl_base_directory}/common/target_atomic.h
   ${devicertl_base_directory}/common/state-queuei.h
   ${devicertl_base_directory}/common/support.h)
 
 # for both in-tree and out-of-tree build
set(cu_cmd ${AOMP_BINDIR}/clang++
     --cuda-device-only
     -Wno-unused-value
     -x hip
     -nogpulib -nogpuinc
     -O${optimization_level}
     --cuda-gpu-arch=${mcpu}
     ${CUDA_DEBUG}
     -I${CMAKE_CURRENT_SOURCE_DIR}/src
-    -I${devicertl_base_directory})
+    -I${devicertl_base_directory}
+    -I${devicertl_base_directory}/../include)
 
   set(bc1_files)
 
   foreach(file ${ARGN})
     get_filename_component(fname ${file} NAME_WE)
     set(bc1_filename ${fname}.${mcpu}.bc)
 
     add_custom_command(

openmp/libomptarget/deviceRTLs/common/debug.h

 //   WARNING(flag, str, args)      : if warning flag is on, print the warning
 //   WON(flag)                     : return true if warning flag is on
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef _OMPTARGET_NVPTX_DEBUG_H_
 #define _OMPTARGET_NVPTX_DEBUG_H_
 
-#include "common/device_environment.h"
+#include "device_environment.h"
 
 ////////////////////////////////////////////////////////////////////////////////
 // set desired level of debugging
 ////////////////////////////////////////////////////////////////////////////////
 
 #define LD_SET_NONE 0ULL /* none */
 #define LD_SET_ALL -1ULL /* all */
 

openmp/libomptarget/deviceRTLs/common/omptarget.h

 //===---- omptarget.h - OpenMP GPU initialization ---------------- CUDA -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the declarations of all library macros, types,
 // and functions.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef OMPTARGET_H
 #define OMPTARGET_H
 
 #include "target_impl.h"
 #include "common/debug.h"     // debug
 #include "interface.h" // interfaces with omp, compiler, and user
 #include "common/state-queue.h"
 #include "common/support.h"
 
 #define OMPTARGET_NVPTX_VERSION 1.1
 
 // used by the library for the interface with the app

 ////////////////////////////////////////////////////////////////////////////////
 
 ////////////////////////////////////////////////////////////////////////////////
 // global data tables
 ////////////////////////////////////////////////////////////////////////////////
 
 extern DEVICE omptarget_nvptx_SimpleMemoryManager
     omptarget_nvptx_simpleMemoryManager;
+extern DEVICE omptarget_device_environmentTy omptarget_device_environment;
 extern DEVICE SHARED uint32_t usedMemIdx;
 extern DEVICE SHARED uint32_t usedSlotIdx;
 extern DEVICE SHARED uint8_t
     parallelLevel[MAX_THREADS_PER_TEAM / WARPSIZE];
 extern DEVICE SHARED uint16_t threadLimit;
 extern DEVICE SHARED uint16_t threadsInTeam;
 extern DEVICE SHARED uint16_t nThreads;
 extern DEVICE SHARED

openmp/libomptarget/deviceRTLs/common/src/libcall.cu

 //===------------ libcall.cu - OpenMP GPU user calls ------------- CUDA -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the OpenMP runtime functions that can be
 // invoked by the user in an OpenMP region
 //
 //===----------------------------------------------------------------------===//
 
 #include "common/omptarget.h"
 #include "common/target_atomic.h"
+#include "device_environment.h"
 #include "target_impl.h"
 
 EXTERN double omp_get_wtick(void) {
   double rc = __kmpc_impl_get_wtick();
   PRINT(LD_IO, "omp_get_wtick() returns %g\n", rc);
   return rc;
 }
 

 }
 
 EXTERN int omp_get_default_device(void) {
   PRINT0(LD_IO,
          "call omp_get_default_device() is undef on device, returns 0\n");
   return 0;
 }
 
+EXTERN int omp_get_device_num(void) {
+  int32_t num = omptarget_device_environment.device_num;
+  PRINT(LD_IO, "call omp_get_device_num() returns %d\n", num);
+  return num;
+}
+
 EXTERN int omp_get_num_devices(void) {
   PRINT0(LD_IO, "call omp_get_num_devices() is undef on device, returns 0\n");
   return 0;
 }
 
 EXTERN int omp_get_num_teams(void) {
   int rc = GetNumberOfOmpTeams();
   PRINT(LD_IO, "call omp_get_num_teams() returns %d\n", rc);

openmp/libomptarget/deviceRTLs/common/src/omp_data.cu

 //===------------ omp_data.cu - OpenMP GPU objects --------------- CUDA -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the data objects used on the GPU device.
 //
 //===----------------------------------------------------------------------===//
 
 #include "common/omptarget.h"
-#include "common/device_environment.h"
+#include "device_environment.h"
 
 ////////////////////////////////////////////////////////////////////////////////
 // global device environment
 ////////////////////////////////////////////////////////////////////////////////
 
 DEVICE omptarget_device_environmentTy omptarget_device_environment;
 
 ////////////////////////////////////////////////////////////////////////////////

openmp/libomptarget/deviceRTLs/interface.h

 typedef enum omp_proc_bind_t {
   omp_proc_bind_false = 0,
   omp_proc_bind_true = 1,
   omp_proc_bind_master = 2,
   omp_proc_bind_close = 3,
   omp_proc_bind_spread = 4
 } omp_proc_bind_t;
 
 EXTERN double omp_get_wtick(void);
 EXTERN double omp_get_wtime(void);
 
 EXTERN void omp_set_num_threads(int num);
 EXTERN int omp_get_num_threads(void);
 EXTERN int omp_get_max_threads(void);
 EXTERN int omp_get_thread_limit(void);
 EXTERN int omp_get_thread_num(void);
 EXTERN int omp_get_num_procs(void);
 EXTERN int omp_in_parallel(void);
 EXTERN int omp_in_final(void);
 EXTERN void omp_set_dynamic(int flag);
 EXTERN int omp_get_dynamic(void);
 EXTERN void omp_set_nested(int flag);
 EXTERN int omp_get_nested(void);
 EXTERN void omp_set_max_active_levels(int level);
 EXTERN int omp_get_max_active_levels(void);
 EXTERN int omp_get_level(void);
 EXTERN int omp_get_active_level(void);
 EXTERN int omp_get_ancestor_thread_num(int level);
 EXTERN int omp_get_team_size(int level);
 
 EXTERN void omp_init_lock(omp_lock_t *lock);
 EXTERN void omp_init_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_destroy_lock(omp_lock_t *lock);
 EXTERN void omp_destroy_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_set_lock(omp_lock_t *lock);
 EXTERN void omp_set_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_unset_lock(omp_lock_t *lock);
 EXTERN void omp_unset_nest_lock(omp_nest_lock_t *lock);
 EXTERN int omp_test_lock(omp_lock_t *lock);
 EXTERN int omp_test_nest_lock(omp_nest_lock_t *lock);
 
 EXTERN void omp_get_schedule(omp_sched_t *kind, int *modifier);
 EXTERN void omp_set_schedule(omp_sched_t kind, int modifier);
 EXTERN omp_proc_bind_t omp_get_proc_bind(void);
 EXTERN int omp_get_cancellation(void);
 EXTERN void omp_set_default_device(int deviceId);
 EXTERN int omp_get_default_device(void);
+EXTERN int omp_get_device_num(void);
 EXTERN int omp_get_num_devices(void);
 EXTERN int omp_get_num_teams(void);
 EXTERN int omp_get_team_num(void);
 EXTERN int omp_is_initial_device(void);
 EXTERN int omp_get_initial_device(void);
 EXTERN int omp_get_max_task_priority(void);
 
 ////////////////////////////////////////////////////////////////////////////////

openmp/libomptarget/deviceRTLs/nvptx/CMakeLists.txt

if(${LIBOMPTARGET_NVPTX_DEBUG})
     set(CUDA_DEBUG -DOMPTARGET_NVPTX_DEBUG=-1 -g --ptxas-options=-v)
   endif()
 
   # NVPTX runtime library has to be statically linked. Dynamic linking is not
   # yet supported by the CUDA toolchain on the device.
   set(BUILD_SHARED_LIBS OFF)
   set(CUDA_SEPARABLE_COMPILATION ON)
   list(APPEND CUDA_NVCC_FLAGS -I${devicertl_base_directory}
-                              -I${devicertl_nvptx_directory}/src)
+                              -I${devicertl_nvptx_directory}/src
+                              -I${devicertl_base_directory}/../include)
   cuda_add_library(omptarget-nvptx STATIC ${cuda_src_files} ${omp_data_objects}
       OPTIONS ${CUDA_ARCH} ${CUDA_DEBUG} ${MAX_SM_DEFINITION})
 
   # Install device RTL under the lib destination folder.
   install(TARGETS omptarget-nvptx ARCHIVE DESTINATION "${OPENMP_INSTALL_LIBDIR}")
 
   target_link_libraries(omptarget-nvptx ${CUDA_LIBRARIES})
 
if (${LIBOMPTARGET_NVPTX_ENABLE_BCLIB})
     if (NOT ${LIBOMPTARGET_NVPTX_BCLIB_SUPPORTED})
       libomptarget_error_say("Cannot build CUDA LLVM bitcode offloading device RTL!")
     endif()
     libomptarget_say("Building CUDA LLVM bitcode offloading device RTL.")
 
     # Set flags for LLVM Bitcode compilation.
     set(bc_flags ${LIBOMPTARGET_NVPTX_SELECTED_CUDA_COMPILER_FLAGS}
                  -I${devicertl_base_directory}
-                 -I${devicertl_nvptx_directory}/src)
+                 -I${devicertl_nvptx_directory}/src
+                 -I${devicertl_base_directory}/../include)
 
     if(${LIBOMPTARGET_NVPTX_DEBUG})
       set(bc_flags ${bc_flags} -DOMPTARGET_NVPTX_DEBUG=-1)
     else()
       set(bc_flags ${bc_flags} -DOMPTARGET_NVPTX_DEBUG=0)
     endif()
 
     # CUDA 9 header files use the nv_weak attribute which clang is not yet prepared

openmp/libomptarget/include/device_environment.h

 //===---- device_environment.h - OpenMP GPU device environment --- CUDA -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Global device environment
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef _OMPTARGET_DEVICE_ENVIRONMENT_H_
 #define _OMPTARGET_DEVICE_ENVIRONMENT_H_
 
-#include "target_impl.h"
-
 struct omptarget_device_environmentTy {
   int32_t debug_level;
+  int32_t device_num;
 };
 
-extern DEVICE omptarget_device_environmentTy omptarget_device_environment;
-
 #endif

openmp/libomptarget/plugins/CMakeLists.txt

if(CMAKE_SYSTEM_PROCESSOR MATCHES "${tmachine}$")
   endif(LIBOMPTARGET_DEP_LIBELF_FOUND)
 else()
   libomptarget_say("Not building ${tmachine_name} offloading plugin: machine not found in the system.")
 endif()
 endmacro()
 
 add_subdirectory(aarch64)
 add_subdirectory(cuda)
+add_subdirectory(amdgpu)
 add_subdirectory(ppc64)
 add_subdirectory(ppc64le)
 add_subdirectory(ve)
 add_subdirectory(x86_64)
 
 # Make sure the parent scope can see the plugins that will be created.
 set(LIBOMPTARGET_SYSTEM_TARGETS "${LIBOMPTARGET_SYSTEM_TARGETS}" PARENT_SCOPE)
 set(LIBOMPTARGET_TESTED_PLUGINS "${LIBOMPTARGET_TESTED_PLUGINS}" PARENT_SCOPE)
 

openmp/libomptarget/plugins/amdgpu/src/rtl.cpp

 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <dlfcn.h>
 #include <elf.h>
 #include <ffi.h>
 #include <fstream>
 #include <iostream>
 #include <libelf.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
 #include <thread>
 #include <unordered_map>
 #include <vector>
 
~RTLDeviceInfoTy() {
     // atmi_finalize removes access to it
     deviceStateStore.clear();
     KernelArgPoolMap.clear();
     atmi_finalize();
   }
 };
 
 pthread_mutex_t SignalPoolT::mutex = PTHREAD_MUTEX_INITIALIZER;
 
-// TODO: May need to drop the trailing to fields until deviceRTL is updated
-struct omptarget_device_environmentTy {
-  int32_t debug_level; // gets value of envvar LIBOMPTARGET_DEVICE_RTL_DEBUG
-                       // only useful for Debug build of deviceRTLs
-  int32_t num_devices; // gets number of active offload devices
-  int32_t device_num;  // gets a value 0 to num_devices-1
-};
+#include "device_environment.h"
 
 static RTLDeviceInfoTy DeviceInfo;
 
 namespace {
 
 int32_t dataRetrieve(int32_t DeviceId, void *HstPtr, void *TgtPtr, int64_t Size,
                      __tgt_async_info *AsyncInfoPtr) {
   assert(AsyncInfoPtr && "AsyncInfoPtr is nullptr");
__tgt_target_table *__tgt_rtl_load_binary_locked(int32_t device_id,
   // We do not need to set the ELF version because the caller of this function
   // had to do that to decide the right runtime to use
 
   if (!elf_machine_id_is_amdgcn(image)) {
     return NULL;
   }
 
   omptarget_device_environmentTy host_device_env;
-  host_device_env.num_devices = DeviceInfo.NumberOfDevices;
   host_device_env.device_num = device_id;
   host_device_env.debug_level = 0;
 #ifdef OMPTARGET_DEBUG
   if (char *envStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG")) {
     host_device_env.debug_level = std::stoi(envStr);
   }
 #endif
 

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 //===----RTLs/cuda/src/rtl.cpp - Target RTLs Implementation ------- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // RTL for CUDA machine
 //
 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <string>
 #include <vector>
 
 #include "Debug.h"
 #include "omptargetplugin.h"
struct KernelTy {
   int MaxThreadsPerBlock = 0;
 
   KernelTy(CUfunction _Func, int8_t _ExecutionMode)
       : Func(_Func), ExecutionMode(_ExecutionMode) {}
 };
 
 /// Device environment data
 /// Manually sync with the deviceRTL side for now, move to a dedicated header
 /// file later.
-struct omptarget_device_environmentTy {
-  int32_t debug_level;
-};
+#include "device_environment.h"
 
 namespace {
 bool checkResult(CUresult Err, const char *ErrMsg) {
   if (Err == CUDA_SUCCESS)
     return true;
 
   DP("%s", ErrMsg);
   CUDA_ERR_STRING(Err);
__tgt_target_table *loadBinary(const int DeviceId,
     // send device environment data to the device
     {
       omptarget_device_environmentTy DeviceEnv{0};
 
 #ifdef OMPTARGET_DEBUG
       if (const char *EnvStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG"))
         DeviceEnv.debug_level = std::stoi(EnvStr);
 #endif
+      DeviceEnv.device_num = DeviceId;
 
       const char *DeviceEnvName = "omptarget_device_environment";
       CUdeviceptr DeviceEnvPtr;
       size_t CUSize;
 
       Err = cuModuleGetGlobal(&DeviceEnvPtr, &CUSize, Module, DeviceEnvName);
       if (Err == CUDA_SUCCESS) {
         if (CUSize != sizeof(DeviceEnv)) {

openmp/libomptarget/test/offloading/get_device_num.c

+// RUN: %libomptarget-compile-run-and-check-aarch64-unknown-linux-gnu
+// RUN: %libomptarget-compile-run-and-check-powerpc64-ibm-linux-gnu
+// RUN: %libomptarget-compile-run-and-check-powerpc64le-ibm-linux-gnu
+// RUN: %libomptarget-compile-run-and-check-x86_64-pc-linux-gnu
+// RUN: %libomptarget-compile-run-and-check-nvptx64-nvidia-cuda
+
+#include <omp.h>
+#include <stdio.h>
+
+int main() {
+
+  const int num_devices = omp_get_num_devices();
+
+  if (num_devices > 0) {
+    for (unsigned i = 0; i < (unsigned)num_devices; i++) {
+      int num;
+#pragma omp target map(from : num) device(i)
+      num = omp_get_device_num();
+
+      printf("Device %u returned id %u\n", i, num);
+
+      if (num != i) {
+        printf("Fail: Device %u returned id %u\n", i, num);
+        return 1;
+      }
+    }
+  }
+
+  printf("PASS\n");
+  return 0;
+}
+
+// CHECK: PASS