[OpenMP] RecordReplay saves bitcode when JIT-ing

[ggeorgakoudis]

Summary

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.

[jdoerfert]

LG, nits

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

127

There is a helper for this in utilities now. Getptrdiff or similar

132

Nit: add braces

[ggeorgakoudis]

Fix nits

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

class DeviceImageTy {
   };
 
   /// Image identifier within the corresponding device. Notice that this id is
   /// not unique between different device; they may overlap.
   int32_t ImageId;
 
   /// 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");
   }
 
   /// Get the image identifier within the device.
   int32_t getId() const { return ImageId; }
 
   /// 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; }
 
   /// Get the image size.
   size_t getSize() const {
     return getPtrDiff(TgtImage->ImageEnd, TgtImage->ImageStart);
   }
 

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

RecordReplayTy()
                               /* Default in GB */ 64) {}
 
   void saveImage(const char *Name, DeviceImageTy &Image) {
     Twine ImageName = Twine(Name) + Twine(".image");
     std::error_code EC;
     raw_fd_ostream OS(ImageName.str(), EC);
     if (EC)
       report_fatal_error("Error saving image : " + StringRef(EC.message()));
+    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();
   }
 
   void saveKernelInputInfo(const char *Name, void **ArgPtrs,
                            ptrdiff_t *ArgOffsets, int32_t NumArgs,
                            uint64_t NumTeamsClause, uint32_t ThreadLimitClause,
                            uint64_t LoopTripCount,
                            AsyncInfoWrapperTy &AsyncInfoWrapper) {
GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
   // for the image id, which is the number of previously loaded images.
   auto ImageOrErr =
       loadBinaryImpl(PostJITImageOrErr.get(), LoadedImages.size());
   if (!ImageOrErr)
     return ImageOrErr.takeError();
 
   DeviceImageTy *Image = *ImageOrErr;
   assert(Image != nullptr && "Invalid image");
+  if (InputTgtImage != PostJITImageOrErr.get())
+    Image->setTgtImageBitcode(InputTgtImage);
 
   // Add the image to list.
   LoadedImages.push_back(Image);
 
   // Setup the device environment if needed.
   if (auto Err = setupDeviceEnvironment(Plugin, *Image))
     return std::move(Err);