[Openmp-commits] [openmp] ec02c34 - [OpenMP] Add additional fields to device environment

[Joseph Huber via Openmp-commits]

Author: Joseph Huber
Date: 2021-09-17T21:25:32-04:00
New Revision: ec02c34b6df1efc643e5889bf3f181d55c4fed38

URL: https://github.com/llvm/llvm-project/commit/ec02c34b6df1efc643e5889bf3f181d55c4fed38
DIFF: https://github.com/llvm/llvm-project/commit/ec02c34b6df1efc643e5889bf3f181d55c4fed38.diff

LOG: [OpenMP] Add additional fields to device environment

This patch adds fields for the device number and number of devices into
the device environment struct and debugging values.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D110004

Added: 
    

Modified: 
    openmp/libomptarget/DeviceRTL/include/Configuration.h
    openmp/libomptarget/DeviceRTL/src/Configuration.cpp
    openmp/libomptarget/deviceRTLs/common/device_environment.h
    openmp/libomptarget/plugins/cuda/src/rtl.cpp

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/DeviceRTL/include/Configuration.h b/openmp/libomptarget/DeviceRTL/include/Configuration.h
index 6fef6b709ba3..11aa5481c0f3 100644
--- a/openmp/libomptarget/DeviceRTL/include/Configuration.h
+++ b/openmp/libomptarget/DeviceRTL/include/Configuration.h
@@ -24,8 +24,12 @@ enum DebugLevel : int32_t { Assertion };
 /// host by omp_get_num_devices.
 uint32_t getNumDevices();
 
+/// Return the number of devices in the system, same number as returned on the
+/// host by omp_get_num_devices.
+uint32_t getDeviceNum();
+
 /// Return the user choosen debug level.
-int32_t getDebugLevel();
+uint32_t getDebugLevel();
 
 bool isDebugMode(DebugLevel Level);
 

diff  --git a/openmp/libomptarget/DeviceRTL/src/Configuration.cpp b/openmp/libomptarget/DeviceRTL/src/Configuration.cpp
index ac67af144897..4e485992cde3 100644
--- a/openmp/libomptarget/DeviceRTL/src/Configuration.cpp
+++ b/openmp/libomptarget/DeviceRTL/src/Configuration.cpp
@@ -18,23 +18,29 @@
 using namespace _OMP;
 
 struct DeviceEnvironmentTy {
-  int32_t DebugLevel;
+  uint32_t DebugLevel;
+  uint32_t NumDevices;
+  uint32_t DeviceNum;
 };
 
 #pragma omp declare target
 
+extern uint32_t __omp_rtl_debug_kind;
+
 // TOOD: We want to change the name as soon as the old runtime is gone.
 DeviceEnvironmentTy CONSTANT(omptarget_device_environment)
     __attribute__((used));
 
-int32_t config::getDebugLevel() {
-  // TODO: Implement libomptarget initialization of DeviceEnvironmentTy
-  return 0;
+uint32_t config::getDebugLevel() {
+  return __omp_rtl_debug_kind & omptarget_device_environment.DebugLevel;
 }
 
 uint32_t config::getNumDevices() {
-  // TODO: Implement libomptarget initialization of DeviceEnvironmentTy
-  return 1;
+  return omptarget_device_environment.NumDevices;
+}
+
+uint32_t config::getDeviceNum() {
+  return omptarget_device_environment.DeviceNum;
 }
 
 bool config::isDebugMode(config::DebugLevel Level) {

diff  --git a/openmp/libomptarget/deviceRTLs/common/device_environment.h b/openmp/libomptarget/deviceRTLs/common/device_environment.h
index d1629f89e53a..5f94567a7885 100644
--- a/openmp/libomptarget/deviceRTLs/common/device_environment.h
+++ b/openmp/libomptarget/deviceRTLs/common/device_environment.h
@@ -17,6 +17,8 @@
 
 struct omptarget_device_environmentTy {
   int32_t debug_level;
+  uint32_t num_devices;
+  uint32_t device_num;
 };
 
 extern omptarget_device_environmentTy omptarget_device_environment;

diff  --git a/openmp/libomptarget/plugins/cuda/src/rtl.cpp b/openmp/libomptarget/plugins/cuda/src/rtl.cpp
index c6f51a5a57bf..f3a810cc4d87 100644
--- a/openmp/libomptarget/plugins/cuda/src/rtl.cpp
+++ b/openmp/libomptarget/plugins/cuda/src/rtl.cpp
@@ -101,6 +101,8 @@ struct KernelTy {
 /// file later.
 struct omptarget_device_environmentTy {
   int32_t debug_level;
+  uint32_t num_devices;
+  uint32_t device_num;
 };
 
 namespace {
@@ -899,7 +901,10 @@ class DeviceRTLTy {
 
     // send device environment data to the device
     {
-      omptarget_device_environmentTy DeviceEnv{0};
+      // TODO: The device ID used here is not the real device ID used by OpenMP.
+      omptarget_device_environmentTy DeviceEnv{
+          0, static_cast<uint32_t>(NumberOfDevices),
+          static_cast<uint32_t>(DeviceId)};
 
 #ifdef OMPTARGET_DEBUG
       if (const char *EnvStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG"))