[Openmp-commits] [openmp] 0f4b4e8 - [OpenMP] RecordReplay saves bitcode when JIT-ing

[Giorgis Georgakoudis via Openmp-commits]

Author: Giorgis Georgakoudis
Date: 2023-01-18T11:25:25-08:00
New Revision: 0f4b4e8e4d8b3aafba829143e2fd4237620ea4d7

URL: https://github.com/llvm/llvm-project/commit/0f4b4e8e4d8b3aafba829143e2fd4237620ea4d7
DIFF: https://github.com/llvm/llvm-project/commit/0f4b4e8e4d8b3aafba829143e2fd4237620ea4d7.diff

LOG: [OpenMP] RecordReplay saves bitcode when JIT-ing

This patch enables to store bitcode images when JIT is enabled for the record-and-replay functionality (see https://reviews.llvm.org/D138931). Credits to @jdoerfert for refactoring the code.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D141986

Added: 
    

Modified: 
    openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp
    openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp
index d20fbf644ffed..7cc1aa2fcd593 100644
--- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp
@@ -122,7 +122,15 @@ struct RecordReplayTy {
     raw_fd_ostream OS(ImageName.str(), EC);
     if (EC)
       report_fatal_error("Error saving image : " + StringRef(EC.message()));
-    OS << Image.getMemoryBuffer().getBuffer();
+    if (auto TgtImageBitcode = Image.getTgtImageBitcode()) {
+      size_t Size =
+          getPtrDiff(TgtImageBitcode->ImageEnd, TgtImageBitcode->ImageStart);
+      MemoryBufferRef MBR = MemoryBufferRef(
+          StringRef((const char *)TgtImageBitcode->ImageStart, Size), "");
+      OS << MBR.getBuffer();
+    } else {
+      OS << Image.getMemoryBuffer().getBuffer();
+    }
     OS.close();
   }
 
@@ -407,6 +415,8 @@ GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
 
   DeviceImageTy *Image = *ImageOrErr;
   assert(Image != nullptr && "Invalid image");
+  if (InputTgtImage != PostJITImageOrErr.get())
+    Image->setTgtImageBitcode(InputTgtImage);
 
   // Add the image to list.
   LoadedImages.push_back(Image);

diff  --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h
index d65209d3cfbd8..e53c666ad7eff 100644
--- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h
+++ b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h
@@ -113,13 +113,14 @@ class DeviceImageTy {
 
   /// The pointer to the raw __tgt_device_image.
   const __tgt_device_image *TgtImage;
+  const __tgt_device_image *TgtImageBitcode;
 
   /// Table of offload entries.
   OffloadEntryTableTy OffloadEntryTable;
 
 public:
   DeviceImageTy(int32_t Id, const __tgt_device_image *Image)
-      : ImageId(Id), TgtImage(Image) {
+      : ImageId(Id), TgtImage(Image), TgtImageBitcode(nullptr) {
     assert(TgtImage && "Invalid target image");
   }
 
@@ -129,6 +130,14 @@ class DeviceImageTy {
   /// Get the pointer to the raw __tgt_device_image.
   const __tgt_device_image *getTgtImage() const { return TgtImage; }
 
+  void setTgtImageBitcode(const __tgt_device_image *TgtImageBitcode) {
+    this->TgtImageBitcode = TgtImageBitcode;
+  }
+
+  const __tgt_device_image *getTgtImageBitcode() const {
+    return TgtImageBitcode;
+  }
+
   /// Get the image starting address.
   void *getStart() const { return TgtImage->ImageStart; }