[OpenMP][NFC] Extract device image handling into a class/header (#74129)

6 files changed, 103 insertions, 47 deletions

diff --git a/openmp/libomptarget/include/DeviceImage.h b/openmp/libomptarget/include/DeviceImage.h
new file mode 100644
index 0000000..369bf75
--- /dev/null
+++ b/openmp/libomptarget/include/DeviceImage.h
@@ -0,0 +1,39 @@
+//===-- DeviceImage.h - Representation of the device code/image -*- 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_DEVICE_IMAGE_H
+#define OMPTARGET_DEVICE_IMAGE_H
+
+#include "Shared/APITypes.h"
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Object/OffloadBinary.h"
+
+class DeviceImageTy {
+
+  std::unique_ptr<llvm::object::OffloadBinary> Binary;
+
+  __tgt_device_image Image;
+  __tgt_image_info ImageInfo;
+
+public:
+  DeviceImageTy(__tgt_device_image &Image);
+
+  __tgt_device_image &getExecutableImage() { return Image; }
+  __tgt_image_info &getImageInfo() { return ImageInfo; }
+
+  llvm::StringRef
+  getArch(llvm::StringRef DefaultArch = llvm::StringRef()) const {
+    return ImageInfo.Arch ? ImageInfo.Arch : DefaultArch;
+  }
+};
+
+#endif // OMPTARGET_DEVICE_IMAGE_H
diff --git a/openmp/libomptarget/include/PluginManager.h b/openmp/libomptarget/include/PluginManager.h
index 5c7e137..3c1f96a 100644
--- a/openmp/libomptarget/include/PluginManager.h
+++ b/openmp/libomptarget/include/PluginManager.h
@@ -13,6 +13,7 @@
 #ifndef OMPTARGET_PLUGIN_MANAGER_H
 #define OMPTARGET_PLUGIN_MANAGER_H
 
+#include "DeviceImage.h"
 #include "Shared/APITypes.h"
 #include "Shared/PluginAPI.h"
 #include "Shared/Requirements.h"
@@ -21,11 +22,13 @@
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/DynamicLibrary.h"
 
 #include <cstdint>
 #include <list>
+#include <memory>
 #include <mutex>
 #include <string>
 
@@ -86,9 +89,12 @@ struct PluginManager {
   /// RTLs identified on the host
   PluginAdaptorManagerTy RTLs;
 
-  /// Executable images and information extracted from the input images passed
-  /// to the runtime.
-  std::list<std::pair<__tgt_device_image, __tgt_image_info>> Images;
+  void addDeviceImage(__tgt_device_image &TgtDeviceImage) {
+    DeviceImages.emplace_back(std::make_unique<DeviceImageTy>(TgtDeviceImage));
+  }
+
+  /// Iterate over all device images registered with this plugin.
+  auto deviceImages() { return llvm::make_pointee_range(DeviceImages); }
 
   /// Devices associated with RTLs
   llvm::SmallVector<std::unique_ptr<DeviceTy>> Devices;
@@ -158,6 +164,10 @@ private:
   // List of all plugin adaptors, in use or not.
   std::list<PluginAdaptorTy> PluginAdaptors;
 
+  /// Executable images and information extracted from the input images passed
+  /// to the runtime.
+  llvm::SmallVector<std::unique_ptr<DeviceImageTy>> DeviceImages;
+
   /// The user provided requirements.
   RequirementCollection Requirements;
 };
diff --git a/openmp/libomptarget/src/CMakeLists.txt b/openmp/libomptarget/src/CMakeLists.txt
index ca40ace..7c311f7 100644
--- a/openmp/libomptarget/src/CMakeLists.txt
+++ b/openmp/libomptarget/src/CMakeLists.txt
@@ -22,6 +22,7 @@ add_llvm_library(omptarget
   rtl.cpp
   LegacyAPI.cpp
   PluginManager.cpp
+  DeviceImage.cpp
 
   OpenMP/Mapping.cpp
   OpenMP/InteropAPI.cpp
diff --git a/openmp/libomptarget/src/DeviceImage.cpp b/openmp/libomptarget/src/DeviceImage.cpp
new file mode 100644
index 0000000..727d276
--- /dev/null
+++ b/openmp/libomptarget/src/DeviceImage.cpp
@@ -0,0 +1,41 @@
+//===-- DeviceImage.cpp - Representation of the device code/image ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+
+#include "DeviceImage.h"
+
+#include "Shared/APITypes.h"
+#include "Shared/Debug.h"
+#include "Shared/Utils.h"
+
+#include "llvm/Support/Error.h"
+
+DeviceImageTy::DeviceImageTy(__tgt_device_image &TgtDeviceImage)
+    : Image(TgtDeviceImage) {
+  llvm::StringRef ImageStr(
+      static_cast<char *>(Image.ImageStart),
+      llvm::omp::target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
+
+  auto BinaryOrErr =
+      llvm::object::OffloadBinary::create(llvm::MemoryBufferRef(ImageStr, ""));
+
+  if (!BinaryOrErr) {
+    consumeError(BinaryOrErr.takeError());
+    return;
+  }
+
+  Binary = std::move(*BinaryOrErr);
+  void *Begin = const_cast<void *>(
+      static_cast<const void *>(Binary->getImage().bytes_begin()));
+  void *End = const_cast<void *>(
+      static_cast<const void *>(Binary->getImage().bytes_end()));
+
+  Image = __tgt_device_image{Begin, End, Image.EntriesBegin, Image.EntriesEnd};
+  ImageInfo = __tgt_image_info{Binary->getArch().data()};
+}
diff --git a/openmp/libomptarget/src/omptarget.cpp b/openmp/libomptarget/src/omptarget.cpp
index 9fb6a96..ee221c9 100644
--- a/openmp/libomptarget/src/omptarget.cpp
+++ b/openmp/libomptarget/src/omptarget.cpp
@@ -299,10 +299,8 @@ void handleTargetOutcome(bool Success, ident_t *Loc) {
 
       if (!PM->getNumUsedPlugins()) {
         llvm::SmallVector<llvm::StringRef> Archs;
-        llvm::transform(PM->Images, std::back_inserter(Archs),
-                        [](const auto &X) {
-                          return !X.second.Arch ? "empty" : X.second.Arch;
-                        });
+        llvm::transform(PM->deviceImages(), std::back_inserter(Archs),
+                        [](const auto &X) { return X.getArch("empty"); });
         FAILURE_MESSAGE(
             "No images found compatible with the installed hardware. ");
         fprintf(stderr, "Found (%s)\n", llvm::join(Archs, ",").c_str());
diff --git a/openmp/libomptarget/src/rtl.cpp b/openmp/libomptarget/src/rtl.cpp
index afe47d6..d114396 100644
--- a/openmp/libomptarget/src/rtl.cpp
+++ b/openmp/libomptarget/src/rtl.cpp
@@ -12,6 +12,7 @@
 
 #include "llvm/Object/OffloadBinary.h"
 
+#include "DeviceImage.h"
 #include "OpenMP/OMPT/Callback.h"
 #include "PluginManager.h"
 #include "device.h"
@@ -131,52 +132,18 @@ static void registerGlobalCtorsDtorsForImage(__tgt_bin_desc *Desc,
   }
 }
 
-static __tgt_device_image getExecutableImage(__tgt_device_image *Image) {
-  StringRef ImageStr(static_cast<char *>(Image->ImageStart),
-                     static_cast<char *>(Image->ImageEnd) -
-                         static_cast<char *>(Image->ImageStart));
-  auto BinaryOrErr =
-      object::OffloadBinary::create(MemoryBufferRef(ImageStr, ""));
-  if (!BinaryOrErr) {
-    consumeError(BinaryOrErr.takeError());
-    return *Image;
-  }
-
-  void *Begin = const_cast<void *>(
-      static_cast<const void *>((*BinaryOrErr)->getImage().bytes_begin()));
-  void *End = const_cast<void *>(
-      static_cast<const void *>((*BinaryOrErr)->getImage().bytes_end()));
-
-  return {Begin, End, Image->EntriesBegin, Image->EntriesEnd};
-}
-
-static __tgt_image_info getImageInfo(__tgt_device_image *Image) {
-  StringRef ImageStr(static_cast<char *>(Image->ImageStart),
-                     static_cast<char *>(Image->ImageEnd) -
-                         static_cast<char *>(Image->ImageStart));
-  auto BinaryOrErr =
-      object::OffloadBinary::create(MemoryBufferRef(ImageStr, ""));
-  if (!BinaryOrErr) {
-    consumeError(BinaryOrErr.takeError());
-    return __tgt_image_info{};
-  }
-
-  return __tgt_image_info{(*BinaryOrErr)->getArch().data()};
-}
-
 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->Images.emplace_back(getExecutableImage(&Desc->DeviceImages[i]),
-                            getImageInfo(&Desc->DeviceImages[i]));
+    PM->addDeviceImage(Desc->DeviceImages[i]);
 
   // Register the images with the RTLs that understand them, if any.
-  for (auto &ImageAndInfo : PM->Images) {
+  for (DeviceImageTy &DI : PM->deviceImages()) {
     // Obtain the image and information that was previously extracted.
-    __tgt_device_image *Img = &ImageAndInfo.first;
-    __tgt_image_info *Info = &ImageAndInfo.second;
+    __tgt_device_image *Img = &DI.getExecutableImage();
+    __tgt_image_info *Info = &DI.getImageInfo();
 
     PluginAdaptorTy *FoundRTL = nullptr;
 
@@ -243,9 +210,9 @@ void PluginAdaptorManagerTy::unregisterLib(__tgt_bin_desc *Desc) {
 
   PM->RTLsMtx.lock();
   // Find which RTL understands each image, if any.
-  for (auto &ImageAndInfo : PM->Images) {
+  for (DeviceImageTy &DI : PM->deviceImages()) {
     // Obtain the image and information that was previously extracted.
-    __tgt_device_image *Img = &ImageAndInfo.first;
+    __tgt_device_image *Img = &DI.getExecutableImage();
 
     PluginAdaptorTy *FoundRTL = NULL;