[OpenMP][NFC] Extract device image handling into a class/header

[jdoerfert]

No description provided.

[llvmbot]

@llvm/pr-subscribers-openmp

Author: Johannes Doerfert (jdoerfert)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/74129.diff

6 Files Affected:

• (added) openmp/libomptarget/include/DeviceImage.h (+39)
• (modified) openmp/libomptarget/include/PluginManager.h (+2-1)
• (modified) openmp/libomptarget/src/CMakeLists.txt (+1)
• (added) openmp/libomptarget/src/DeviceImage.cpp (+40)
• (modified) openmp/libomptarget/src/omptarget.cpp (+1-3)
• (modified) openmp/libomptarget/src/rtl.cpp (+7-40)

diff --git a/openmp/libomptarget/include/DeviceImage.h b/openmp/libomptarget/include/DeviceImage.h
new file mode 100644
index 000000000000000..369bf75979afb6e
--- /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 91298928716b6a8..0813a7d5ee6bfd2 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"
 
@@ -91,7 +92,7 @@ struct PluginManager {
 
   /// Executable images and information extracted from the input images passed
   /// to the runtime.
-  std::list<std::pair<__tgt_device_image, __tgt_image_info>> Images;
+  llvm::SmallVector<DeviceImageTy> Images;
 
   /// Devices associated with RTLs
   llvm::SmallVector<std::unique_ptr<DeviceTy>> Devices;
diff --git a/openmp/libomptarget/src/CMakeLists.txt b/openmp/libomptarget/src/CMakeLists.txt
index ca40ace456458b4..7c311f738ac8eb8 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 000000000000000..86591df24cb630d
--- /dev/null
+++ b/openmp/libomptarget/src/DeviceImage.cpp
@@ -0,0 +1,40 @@
+//===-- 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 &OrigImage) : Image(OrigImage) {
+  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 9fb6a965fe2402b..8ac1c288c4a8d0d 100644
--- a/openmp/libomptarget/src/omptarget.cpp
+++ b/openmp/libomptarget/src/omptarget.cpp
@@ -300,9 +300,7 @@ 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;
-                        });
+                        [](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 52ea76438d79a82..e8929edd9b99c8f 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,39 +132,6 @@ 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::registerRequires(int64_t Flags) {
   // TODO: add more elaborate check.
   // Minimal check: only set requires flags if previous value
@@ -210,14 +178,13 @@ void PluginAdaptorManagerTy::registerLib(__tgt_bin_desc *Desc) {
 
   // 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->Images.emplace_back(Desc->DeviceImages[i]);
 
   // Register the images with the RTLs that understand them, if any.
-  for (auto &ImageAndInfo : PM->Images) {
+  for (DeviceImageTy &DI : PM->Images) {
     // 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;
 
@@ -284,9 +251,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->Images) {
     // Obtain the image and information that was previously extracted.
-    __tgt_device_image *Img = &ImageAndInfo.first;
+    __tgt_device_image *Img = &DI.getExecutableImage();
 
     PluginAdaptorTy *FoundRTL = NULL;
 

[jhuber6]

I want to redesign all of this at some point. All the information we need should be in the ELF header so there's no need for the runtime to need to parse some special binary format. We will however still store it at the section such that tools like llvm-objdump can read what it is.

[jdoerfert]

Agreed.