[OpenMP][NFC] Encapsulate Devices.size() (#74010)

llvm · Dec 1, 2023 · 1035cc7 · 1035cc7
1 parent b6cad75
commit 1035cc7
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 16 deletions.
diff --git a/openmp/libomptarget/include/PluginManager.h b/openmp/libomptarget/include/PluginManager.h
@@ -143,6 +143,11 @@ struct PluginManager {
     DelayedBinDesc.clear();
   }
 
+  int getNumDevices() {
+    std::lock_guard<decltype(RTLsMtx)> Lock(RTLsMtx);
+    return Devices.size();
+  }
+
 private:
   bool RTLsLoaded = false;
   llvm::SmallVector<__tgt_bin_desc *> DelayedBinDesc;

diff --git a/openmp/libomptarget/src/api.cpp b/openmp/libomptarget/src/api.cpp
@@ -28,13 +28,11 @@
 
 EXTERN int omp_get_num_devices(void) {
   TIMESCOPE();
-  PM->RTLsMtx.lock();
-  size_t DevicesSize = PM->Devices.size();
-  PM->RTLsMtx.unlock();
+  size_t NumDevices = PM->getNumDevices();
 
-  DP("Call to omp_get_num_devices returning %zd\n", DevicesSize);
+  DP("Call to omp_get_num_devices returning %zd\n", NumDevices);
 
-  return DevicesSize;
+  return NumDevices;
 }
 
 EXTERN int omp_get_device_num(void) {
@@ -112,10 +110,8 @@ EXTERN int omp_target_is_present(const void *Ptr, int DeviceNum) {
     return true;
   }
 
-  PM->RTLsMtx.lock();
-  size_t DevicesSize = PM->Devices.size();
-  PM->RTLsMtx.unlock();
-  if (DevicesSize <= (size_t)DeviceNum) {
+  size_t NumDevices = PM->getNumDevices();
+  if (NumDevices <= (size_t)DeviceNum) {
     DP("Call to omp_target_is_present with invalid device ID, returning "
        "false\n");
     return false;
@@ -562,18 +558,14 @@ EXTERN void *omp_get_mapped_ptr(const void *Ptr, int DeviceNum) {
     return nullptr;
   }
 
-  if (DeviceNum == omp_get_initial_device()) {
+  size_t NumDevices = omp_get_initial_device();
+  if (DeviceNum == NumDevices) {
     REPORT("Device %d is initial device, returning Ptr " DPxMOD ".\n",
            DeviceNum, DPxPTR(Ptr));
     return const_cast<void *>(Ptr);
   }
 
-  int DevicesSize = omp_get_initial_device();
-  {
-    std::lock_guard<std::mutex> LG(PM->RTLsMtx);
-    DevicesSize = PM->Devices.size();
-  }
-  if (DevicesSize <= DeviceNum) {
+  if (NumDevices <= DeviceNum) {
     DP("DeviceNum %d is invalid, returning nullptr.\n", DeviceNum);
     return nullptr;
   }