[OpenMP][NFCI] Organize offload entry logic

This moves the offload entry logic into classes and provides convenient
accessors. No functional change intended but we can now print all
offload entries (and later look them up), tested via
`OMPTARGET_DUMP_OFFLOAD_ENTRIES=<device_no>`.

10 files changed, 179 insertions, 55 deletions

diff --git a/openmp/libomptarget/include/DeviceImage.h b/openmp/libomptarget/include/DeviceImage.h
index 369bf75..465bf97 100644
--- a/openmp/libomptarget/include/DeviceImage.h
+++ b/openmp/libomptarget/include/DeviceImage.h
@@ -12,28 +12,39 @@
 #ifndef OMPTARGET_DEVICE_IMAGE_H
 #define OMPTARGET_DEVICE_IMAGE_H
 
+#include "OffloadEntry.h"
 #include "Shared/APITypes.h"
 
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/iterator.h"
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/Object/OffloadBinary.h"
 
+#include <memory>
+
 class DeviceImageTy {
 
   std::unique_ptr<llvm::object::OffloadBinary> Binary;
+  llvm::SmallVector<std::unique_ptr<OffloadEntryTy>> OffloadEntries;
 
+  __tgt_bin_desc *BinaryDesc;
   __tgt_device_image Image;
   __tgt_image_info ImageInfo;
 
 public:
-  DeviceImageTy(__tgt_device_image &Image);
+  DeviceImageTy(__tgt_bin_desc &BinaryDesc, __tgt_device_image &Image);
 
   __tgt_device_image &getExecutableImage() { return Image; }
   __tgt_image_info &getImageInfo() { return ImageInfo; }
+  __tgt_bin_desc &getBinaryDesc() { return *BinaryDesc; }
 
   llvm::StringRef
   getArch(llvm::StringRef DefaultArch = llvm::StringRef()) const {
     return ImageInfo.Arch ? ImageInfo.Arch : DefaultArch;
   }
+
+  auto entries() { return llvm::make_pointee_range(OffloadEntries); }
 };
 
 #endif // OMPTARGET_DEVICE_IMAGE_H
diff --git a/openmp/libomptarget/include/OffloadEntry.h b/openmp/libomptarget/include/OffloadEntry.h
new file mode 100644
index 0000000..f645fe8
--- /dev/null
+++ b/openmp/libomptarget/include/OffloadEntry.h
@@ -0,0 +1,48 @@
+//===-- OffloadEntry.h - Representation of offload entries ------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef OMPTARGET_OFFLOAD_ENTRY_H
+#define OMPTARGET_OFFLOAD_ENTRY_H
+
+#include "Shared/APITypes.h"
+
+#include "omptarget.h"
+
+#include "llvm/ADT/StringRef.h"
+
+class DeviceImageTy;
+
+class OffloadEntryTy {
+  DeviceImageTy &DeviceImage;
+  __tgt_offload_entry &OffloadEntry;
+
+public:
+  OffloadEntryTy(DeviceImageTy &DeviceImage, __tgt_offload_entry &OffloadEntry)
+      : DeviceImage(DeviceImage), OffloadEntry(OffloadEntry) {}
+
+  bool isGlobal() const { return getSize() != 0; }
+  size_t getSize() const { return OffloadEntry.size; }
+
+  void *getAddress() const { return OffloadEntry.addr; }
+  llvm::StringRef getName() const { return OffloadEntry.name; }
+  const char *getNameAsCStr() const { return OffloadEntry.name; }
+  __tgt_bin_desc *getBinaryDescription() const;
+
+  bool isCTor() { return hasFlags(OMP_DECLARE_TARGET_CTOR); }
+  bool isDTor() { return hasFlags(OMP_DECLARE_TARGET_DTOR); }
+  bool isLink() { return hasFlags(OMP_DECLARE_TARGET_LINK); }
+
+  bool hasFlags(OpenMPOffloadingDeclareTargetFlags Flags) {
+    return Flags & OffloadEntry.flags;
+  }
+};
+
+#endif // OMPTARGET_OFFLOAD_ENTRY_H
diff --git a/openmp/libomptarget/include/PluginManager.h b/openmp/libomptarget/include/PluginManager.h
index 3c1f96a..e5a41a0 100644
--- a/openmp/libomptarget/include/PluginManager.h
+++ b/openmp/libomptarget/include/PluginManager.h
@@ -40,6 +40,10 @@ struct PluginAdaptorTy {
   /// Return the number of devices available to this plugin.
   int32_t getNumDevices() const { return NumberOfDevices; }
 
+  /// Add all offload entries described by \p DI to the devices managed by this
+  /// plugin.
+  void addOffloadEntries(DeviceImageTy &DI);
+
   /// RTL index, index is the number of devices of other RTLs that were
   /// registered before, i.e. the OpenMP index of the first device to be
   /// registered with this RTL.
@@ -89,8 +93,8 @@ struct PluginManager {
   /// RTLs identified on the host
   PluginAdaptorManagerTy RTLs;
 
-  void addDeviceImage(__tgt_device_image &TgtDeviceImage) {
-    DeviceImages.emplace_back(std::make_unique<DeviceImageTy>(TgtDeviceImage));
+  void addDeviceImage(__tgt_bin_desc &TgtBinDesc, __tgt_device_image &TgtDeviceImage) {
+    DeviceImages.emplace_back(std::make_unique<DeviceImageTy>(TgtBinDesc, TgtDeviceImage));
   }
 
   /// Iterate over all device images registered with this plugin.
diff --git a/openmp/libomptarget/include/device.h b/openmp/libomptarget/include/device.h
index 6602ee0..05ed654 100644
--- a/openmp/libomptarget/include/device.h
+++ b/openmp/libomptarget/include/device.h
@@ -19,15 +19,20 @@
 #include <cstring>
 #include <list>
 #include <map>
+#include <memory>
 #include <mutex>
 #include <set>
 
 #include "ExclusiveAccess.h"
+#include "OffloadEntry.h"
 #include "omptarget.h"
 #include "rtl.h"
 
 #include "OpenMP/Mapping.h"
 
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+
 // Forward declarations.
 struct PluginAdaptorTy;
 struct __tgt_bin_desc;
@@ -48,7 +53,7 @@ struct DeviceTy {
 
   bool IsInit;
   std::once_flag InitFlag;
-  bool HasPendingGlobals;
+  bool HasMappedGlobalData = false;
 
   /// Host data to device map type with a wrapper key indirection that allows
   /// concurrent modification of the entries without invalidating the underlying
@@ -223,12 +228,21 @@ struct DeviceTy {
   int32_t destroyEvent(void *Event);
   /// }
 
+  /// Register \p Entry as an offload entry that is avalable on this device.
+  void addOffloadEntry(OffloadEntryTy &Entry);
+
+  /// Print all offload entries to stderr.
+  void dumpOffloadEntries();
+
 private:
   // Call to RTL
   void init(); // To be called only via DeviceTy::initOnce()
 
   /// Deinitialize the device (and plugin).
   void deinit();
+
+  /// All offload entries available on this device.
+  llvm::DenseMap<llvm::StringRef, OffloadEntryTy *> DeviceOffloadEntries;
 };
 
 extern bool deviceIsReady(int DeviceNum);
diff --git a/openmp/libomptarget/src/DeviceImage.cpp b/openmp/libomptarget/src/DeviceImage.cpp
index 727d276..910e190 100644
--- a/openmp/libomptarget/src/DeviceImage.cpp
+++ b/openmp/libomptarget/src/DeviceImage.cpp
@@ -10,14 +10,27 @@
 
 #include "DeviceImage.h"
 
+#include "OffloadEntry.h"
 #include "Shared/APITypes.h"
 #include "Shared/Debug.h"
 #include "Shared/Utils.h"
 
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Error.h"
+#include <memory>
+
+__tgt_bin_desc *OffloadEntryTy::getBinaryDescription() const {
+  return &DeviceImage.getBinaryDesc();
+}
+
+DeviceImageTy::DeviceImageTy(__tgt_bin_desc &BinaryDesc,
+                             __tgt_device_image &TgtDeviceImage)
+    : BinaryDesc(&BinaryDesc), Image(TgtDeviceImage) {
+
+  for (__tgt_offload_entry &Entry :
+       llvm::make_range(Image.EntriesBegin, Image.EntriesEnd))
+    OffloadEntries.emplace_back(std::make_unique<OffloadEntryTy>(*this, Entry));
 
-DeviceImageTy::DeviceImageTy(__tgt_device_image &TgtDeviceImage)
-    : Image(TgtDeviceImage) {
   llvm::StringRef ImageStr(
       static_cast<char *>(Image.ImageStart),
       llvm::omp::target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
diff --git a/openmp/libomptarget/src/PluginManager.cpp b/openmp/libomptarget/src/PluginManager.cpp
index 82cca5e..ec29a3d 100644
--- a/openmp/libomptarget/src/PluginManager.cpp
+++ b/openmp/libomptarget/src/PluginManager.cpp
@@ -69,6 +69,14 @@ PluginAdaptorTy::PluginAdaptorTy(const std::string &Name) : Name(Name) {
   DP("Registered '%s' with %d devices!\n", Name.c_str(), NumberOfDevices);
 }
 
+void PluginAdaptorTy::addOffloadEntries(DeviceImageTy &DI) {
+  for (int32_t I = 0; I < NumberOfDevices; ++I) {
+    DeviceTy &Device = *PM->Devices[DeviceOffset + I];
+    for (OffloadEntryTy &Entry : DI.entries())
+      Device.addOffloadEntry(Entry);
+  }
+}
+
 void PluginManager::init() {
   DP("Loading RTLs...\n");
 
diff --git a/openmp/libomptarget/src/device.cpp b/openmp/libomptarget/src/device.cpp
index feb5d64..d3481d4 100644
--- a/openmp/libomptarget/src/device.cpp
+++ b/openmp/libomptarget/src/device.cpp
@@ -11,9 +11,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "device.h"
+#include "OffloadEntry.h"
 #include "OpenMP/OMPT/Callback.h"
 #include "OpenMP/OMPT/Interface.h"
 #include "PluginManager.h"
+#include "Shared/APITypes.h"
+#include "Shared/Debug.h"
 #include "omptarget.h"
 #include "private.h"
 #include "rtl.h"
@@ -61,7 +64,7 @@ int HostDataToTargetTy::addEventIfNecessary(DeviceTy &Device,
 
 DeviceTy::DeviceTy(PluginAdaptorTy *RTL)
     : DeviceID(-1), RTL(RTL), RTLDeviceID(-1), IsInit(false), InitFlag(),
-      HasPendingGlobals(false), PendingCtorsDtors(), PendingGlobalsMtx() {}
+      PendingCtorsDtors(), PendingGlobalsMtx() {}
 
 DeviceTy::~DeviceTy() {
   if (DeviceID == -1 || !(getInfoLevel() & OMP_INFOTYPE_DUMP_TABLE))
@@ -807,3 +810,52 @@ bool deviceIsReady(int DeviceNum) {
 
   return true;
 }
+
+void DeviceTy::addOffloadEntry(OffloadEntryTy &Entry) {
+  std::lock_guard<decltype(PendingGlobalsMtx)> Lock(PendingGlobalsMtx);
+  DeviceOffloadEntries[Entry.getName()] = &Entry;
+  if (Entry.isGlobal())
+    return;
+
+  if (Entry.isCTor()) {
+    DP("Adding ctor " DPxMOD " to the pending list.\n",
+       DPxPTR(Entry.getAddress()));
+    MESSAGE("WARNING: Calling deprecated constructor for entry %s will be "
+            "removed in a future release \n",
+            Entry.getNameAsCStr());
+    PendingCtorsDtors[Entry.getBinaryDescription()].PendingCtors.push_back(
+        Entry.getAddress());
+  } else if (Entry.isDTor()) {
+    // Dtors are pushed in reverse order so they are executed from end
+    // to beginning when unregistering the library!
+    DP("Adding dtor " DPxMOD " to the pending list.\n",
+       DPxPTR(Entry.getAddress()));
+    MESSAGE("WARNING: Calling deprecated destructor for entry %s will be "
+            "removed in a future release \n",
+            Entry.getNameAsCStr());
+    PendingCtorsDtors[Entry.getBinaryDescription()].PendingDtors.push_front(
+        Entry.getAddress());
+  }
+
+  if (Entry.isLink()) {
+    MESSAGE(
+        "WARNING: The \"link\" attribute is not yet supported for entry: %s!\n",
+        Entry.getNameAsCStr());
+  }
+}
+
+void DeviceTy::dumpOffloadEntries() {
+  fprintf(stderr, "Device %i offload entries:\n", DeviceID);
+  for (auto &It : DeviceOffloadEntries) {
+    const char *Kind = "kernel";
+    if (It.second->isCTor())
+      Kind = "constructor";
+    else if (It.second->isDTor())
+      Kind = "destructor";
+    else if (It.second->isLink())
+      Kind = "link";
+    else if (It.second->isGlobal())
+      Kind = "global var.";
+    fprintf(stderr, "  %11s: %s\n", Kind, It.second->getNameAsCStr());
+  }
+}
diff --git a/openmp/libomptarget/src/omptarget.cpp b/openmp/libomptarget/src/omptarget.cpp
index ee221c9..1fcadc0 100644
--- a/openmp/libomptarget/src/omptarget.cpp
+++ b/openmp/libomptarget/src/omptarget.cpp
@@ -16,6 +16,7 @@
 #include "OpenMP/OMPT/Callback.h"
 #include "OpenMP/OMPT/Interface.h"
 #include "PluginManager.h"
+#include "Shared/EnvironmentVar.h"
 #include "device.h"
 #include "private.h"
 #include "rtl.h"
@@ -128,6 +129,9 @@ static uint64_t getPartialStructRequiredAlignment(void *HstPtrBase) {
 
 /// Map global data and execute pending ctors
 static int initLibrary(DeviceTy &Device) {
+  if (Device.HasMappedGlobalData)
+    return OFFLOAD_SUCCESS;
+
   /*
    * Map global data
    */
@@ -276,7 +280,12 @@ static int initLibrary(DeviceTy &Device) {
     if (AsyncInfo.synchronize() != OFFLOAD_SUCCESS)
       return OFFLOAD_FAIL;
   }
-  Device.HasPendingGlobals = false;
+  Device.HasMappedGlobalData = true;
+
+  static Int32Envar DumpOffloadEntries =
+      Int32Envar("OMPTARGET_DUMP_OFFLOAD_ENTRIES", -1);
+  if (DumpOffloadEntries.get() == DeviceId)
+    Device.dumpOffloadEntries();
 
   return OFFLOAD_SUCCESS;
 }
@@ -374,7 +383,7 @@ bool checkDeviceAndCtors(int64_t &DeviceID, ident_t *Loc) {
   {
     std::lock_guard<decltype(Device.PendingGlobalsMtx)> LG(
         Device.PendingGlobalsMtx);
-    if (Device.HasPendingGlobals && initLibrary(Device) != OFFLOAD_SUCCESS) {
+    if (initLibrary(Device) != OFFLOAD_SUCCESS) {
       REPORT("Failed to init globals on device %" PRId64 "\n", DeviceID);
       handleTargetOutcome(false, Loc);
       return true;
diff --git a/openmp/libomptarget/src/rtl.cpp b/openmp/libomptarget/src/rtl.cpp
index d114396..f81baaa 100644
--- a/openmp/libomptarget/src/rtl.cpp
+++ b/openmp/libomptarget/src/rtl.cpp
@@ -26,6 +26,7 @@
 #include <cassert>
 #include <cstdlib>
 #include <cstring>
+#include <memory>
 #include <mutex>
 #include <string>
 
@@ -89,55 +90,12 @@ static void registerImageIntoTranslationTable(TranslationTable &TT,
   }
 }
 
-////////////////////////////////////////////////////////////////////////////////
-// Functionality for registering Ctors/Dtors
-
-static void registerGlobalCtorsDtorsForImage(__tgt_bin_desc *Desc,
-                                             __tgt_device_image *Img,
-                                             PluginAdaptorTy *RTL) {
-
-  for (int32_t I = 0; I < RTL->NumberOfDevices; ++I) {
-    DeviceTy &Device = *PM->Devices[RTL->DeviceOffset + I];
-    Device.PendingGlobalsMtx.lock();
-    Device.HasPendingGlobals = true;
-    for (__tgt_offload_entry *Entry = Img->EntriesBegin;
-         Entry != Img->EntriesEnd; ++Entry) {
-      // Globals are not callable and use a different set of flags.
-      if (Entry->size != 0)
-        continue;
-
-      if (Entry->flags & OMP_DECLARE_TARGET_CTOR) {
-        DP("Adding ctor " DPxMOD " to the pending list.\n",
-           DPxPTR(Entry->addr));
-        Device.PendingCtorsDtors[Desc].PendingCtors.push_back(Entry->addr);
-        MESSAGE("WARNING: Calling deprecated constructor for entry %s will be "
-                "removed in a future release \n",
-                Entry->name);
-      } else if (Entry->flags & OMP_DECLARE_TARGET_DTOR) {
-        // Dtors are pushed in reverse order so they are executed from end
-        // to beginning when unregistering the library!
-        DP("Adding dtor " DPxMOD " to the pending list.\n",
-           DPxPTR(Entry->addr));
-        Device.PendingCtorsDtors[Desc].PendingDtors.push_front(Entry->addr);
-        MESSAGE("WARNING: Calling deprecated destructor for entry %s will be "
-                "removed in a future release \n",
-                Entry->name);
-      }
-
-      if (Entry->flags & OMP_DECLARE_TARGET_LINK) {
-        DP("The \"link\" attribute is not yet supported!\n");
-      }
-    }
-    Device.PendingGlobalsMtx.unlock();
-  }
-}
-
 void PluginAdaptorManagerTy::registerLib(__tgt_bin_desc *Desc) {
   PM->RTLsMtx.lock();
 
   // Extract the exectuable image and extra information if availible.
   for (int32_t i = 0; i < Desc->NumDeviceImages; ++i)
-    PM->addDeviceImage(Desc->DeviceImages[i]);
+    PM->addDeviceImage(*Desc, Desc->DeviceImages[i]);
 
   // Register the images with the RTLs that understand them, if any.
   for (DeviceImageTy &DI : PM->deviceImages()) {
@@ -189,8 +147,8 @@ void PluginAdaptorManagerTy::registerLib(__tgt_bin_desc *Desc) {
       PM->TrlTblMtx.unlock();
       FoundRTL = &R;
 
-      // Load ctors/dtors for static objects
-      registerGlobalCtorsDtorsForImage(Desc, Img, FoundRTL);
+      // Register all offload entries with the devices handled by the plugin.
+      R.addOffloadEntries(DI);
 
       // if an RTL was found we are done - proceed to register the next image
       break;
diff --git a/openmp/libomptarget/test/offloading/ctor_dtor.cpp b/openmp/libomptarget/test/offloading/ctor_dtor.cpp
index 46e9dd4..a1f6d01 100644
--- a/openmp/libomptarget/test/offloading/ctor_dtor.cpp
+++ b/openmp/libomptarget/test/offloading/ctor_dtor.cpp
@@ -1,5 +1,12 @@
 // RUN: %libomptarget-compilexx-run-and-check-generic
 // RUN: %libomptarget-compileoptxx-run-and-check-generic
+// RUN: %libomptarget-compilexx-generic && \
+// RUN: env OMPTARGET_DUMP_OFFLOAD_ENTRIES=0 %libomptarget-run-generic 2>&1 | \
+// RUN: %fcheck-generic --check-prefix=DUMP
+//
+// DUMP:     Device 0 offload entries:
+// DUMP-DAG:   global var.: s
+// DUMP-DAG:        kernel: __omp_offloading_16_{{.*}}_main_
 //
 #include <cstdio>
 struct S {