[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 369bf75979af..465bf970ef17 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 000000000000..f645fe81db2d
--- /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 3c1f96a15841..e5a41a0f329d 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 6602ee052ddd..05ed6546557a 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 727d2768220e..910e1907dcfe 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 82cca5e0ef8f..ec29a3db956b 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 feb5d64190e5..d3481d42af96 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 ee221c9041a8..1fcadc018f72 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 d1143969e48e..f81baaa40bfd 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 46e9dd46a356..a1f6d01754e5 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 {