[OpenMP] Support saving globals for RecordReplay

[ggeorgakoudis]

Summary

This patch adds storing globals besides explicitly allocated device memory for the record-and-replay functionality of OpenMP offloading kernels. Credits to @jdoerfert for the implementation.

[ggeorgakoudis]

openmp/libomptarget/include/rtl.h

77

Is this needed?

openmp/libomptarget/include/device.h

struct DeviceTy {
   /// is fulfilled.
   int32_t waitEvent(void *Event, AsyncInfoTy &AsyncInfo);
 
   /// Synchronize the event. It is expected to block the thread.
   int32_t syncEvent(void *Event);
 
   /// Destroy the event.
   int32_t destroyEvent(void *Event);
+
+  /// Update recorded globals to their values described in \p Globals.
+  int32_t initializeRecordedGlobals(const void *Globals, int32_t Size);
   /// }
 
 private:
   // Call to RTL
   void init(); // To be called only via DeviceTy::initOnce()
 
   /// Deinitialize the device (and plugin).
   void deinit();

openmp/libomptarget/include/omptarget.h

 void __tgt_target_nowait_query(void **AsyncHandle);
 
 /// Executes a target kernel by replaying recorded kernel arguments and
 /// device memory.
 int __tgt_target_kernel_replay(ident_t *Loc, int64_t DeviceId, void *HostPtr,
                                void *DeviceMemory, int64_t DeviceMemorySize,
                                void **TgtArgs, ptrdiff_t *TgtOffsets,
                                int32_t NumArgs, int32_t NumTeams,
-                               int32_t ThreadLimit, uint64_t LoopTripCount);
+                               int32_t ThreadLimit, uint64_t LoopTripCount,
+                               const void *Globals, int64_t GlobalsSize);
 
 void __tgt_set_info_flag(uint32_t);
 
 int __tgt_print_device_info(int64_t DeviceId);
 #ifdef __cplusplus
 }
 #endif
 
 #ifdef __cplusplus
 #define EXTERN extern "C"
 #else
 #define EXTERN extern
 #endif
 
 #endif // _OMPTARGET_H_

openmp/libomptarget/include/rtl.h

struct RTLInfoTy {
   typedef void(set_info_flag_ty)(uint32_t);
   typedef int32_t(create_event_ty)(int32_t, void **);
   typedef int32_t(record_event_ty)(int32_t, void *, __tgt_async_info *);
   typedef int32_t(wait_event_ty)(int32_t, void *, __tgt_async_info *);
   typedef int32_t(sync_event_ty)(int32_t, void *);
   typedef int32_t(destroy_event_ty)(int32_t, void *);
   typedef int32_t(release_async_info_ty)(int32_t, __tgt_async_info *);
   typedef int32_t(init_async_info_ty)(int32_t, __tgt_async_info **);
-  typedef int64_t(init_device_into_ty)(int64_t, __tgt_device_info *,
+  typedef int64_t(init_device_into_ty)(int32_t, __tgt_device_info *,
                                        const char **);
   typedef int32_t(data_lock_ty)(int32_t, void *, int64_t, void **);
   typedef int32_t(data_unlock_ty)(int32_t, void *);
+  typedef int32_t(initialize_recorded_globals_ty)(int32_t, const void *,
+                                                  int32_t);
 
   int32_t Idx = -1;             // 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.
   int32_t NumberOfDevices = -1; // Number of devices this RTL deals with.
 
   std::unique_ptr<llvm::sys::DynamicLibrary> LibraryHandler;
#endif
   wait_event_ty *wait_event = nullptr;
   sync_event_ty *sync_event = nullptr;
   destroy_event_ty *destroy_event = nullptr;
   init_async_info_ty *init_async_info = nullptr;
   init_device_into_ty *init_device_info = nullptr;
   release_async_info_ty *release_async_info = nullptr;
   data_lock_ty *data_lock = nullptr;
   data_unlock_ty *data_unlock = nullptr;
+  initialize_recorded_globals_ty *initialize_recorded_globals = nullptr;
 
   // Are there images associated with this RTL.
   bool IsUsed = false;
 
   llvm::DenseSet<const __tgt_device_image *> UsedImages;
 
   // Mutex for thread-safety when calling RTL interface functions.
   // It is easier to enforce thread-safety at the libomptarget level,

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

operator __tgt_target_table *() {
       if (Entries.empty())
         return nullptr;
       return &TTTablePtr;
     }
 
   private:
     __tgt_target_table TTTablePtr;
     llvm::SmallVector<__tgt_offload_entry> Entries;
+
+  public:
+    using const_iterator = decltype(Entries)::const_iterator;
+    const_iterator begin() const { return Entries.begin(); }
+    const_iterator end() const { return Entries.end(); }
   };
 
   /// 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;
public:
   /// Get a memory buffer reference to the whole image.
   MemoryBufferRef getMemoryBuffer() const {
     return MemoryBufferRef(StringRef((const char *)getStart(), getSize()),
                            "Image");
   }
 
   /// Get a reference to the offload entry table for the image.
   OffloadEntryTableTy &getOffloadEntryTable() { return OffloadEntryTable; }
+  const OffloadEntryTableTy &getOffloadEntryTable() const {
+    return OffloadEntryTable;
+  }
 };
 
 /// Class implementing common functionalities of offload kernels. Each plugin
 /// should define the specific kernel class, derive from this generic one, and
 /// implement the necessary virtual function members.
 struct GenericKernelTy {
   /// Construct a kernel with a name and a execution mode.
   GenericKernelTy(const char *Name, OMPTgtExecModeFlags ExecutionMode)
struct GenericKernelTy {
   virtual Error launchImpl(GenericDeviceTy &GenericDevice, uint32_t NumThreads,
                            uint64_t NumBlocks, uint32_t DynamicMemorySize,
                            int32_t NumKernelArgs, void *KernelArgs,
                            AsyncInfoWrapperTy &AsyncInfoWrapper) const = 0;
 
   /// Get the kernel name.
   const char *getName() const { return Name; }
 
+  /// Get the kernel name.
+  const DeviceImageTy &getImage() const {
+    assert(ImagePtr && "Kernel is not initialized!");
+    return *ImagePtr;
+  }
+
   /// Indicate whether an execution mode is valid.
   static bool isValidExecutionMode(OMPTgtExecModeFlags ExecutionMode) {
     switch (ExecutionMode) {
     case OMP_TGT_EXEC_MODE_SPMD:
     case OMP_TGT_EXEC_MODE_GENERIC:
     case OMP_TGT_EXEC_MODE_GENERIC_SPMD:
       return true;
     }
private:
   }
 
   /// The kernel name.
   const char *Name;
 
   /// The execution flags of the kernel.
   OMPTgtExecModeFlags ExecutionMode;
 
+  /// The image that contains this kernel.
+  DeviceImageTy *ImagePtr = nullptr;
+
 protected:
   /// The dynamic memory size reserved for executing the kernel.
   uint32_t DynamicMemorySize;
 
   /// The preferred number of threads to run the kernel.
   uint32_t PreferredNumThreads;
 
   /// The maximum number of threads which the kernel could leverage.
struct GenericDeviceTy : public DeviceAllocatorTy {
   virtual std::string getComputeUnitKind() const { return "unknown"; }
 
   /// Post processing after jit backend. The ownership of \p MB will be taken.
   virtual Expected<std::unique_ptr<MemoryBuffer>>
   doJITPostProcessing(std::unique_ptr<MemoryBuffer> MB) const {
     return std::move(MB);
   }
 
+  /// Initialize recorded globals encoded by \p GlobalEncoding.
+  Error initializeRecordedGlobals(const void *GlobalEncoding, int32_t Size);
+
 private:
   /// Register offload entry for global variable.
   Error registerGlobalOffloadEntry(DeviceImageTy &DeviceImage,
                                    const __tgt_offload_entry &GlobalEntry,
                                    __tgt_offload_entry &DeviceEntry);
 
   /// Register offload entry for kernel function.
   Error registerKernelOffloadEntry(DeviceImageTy &DeviceImage,

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

 //===- PluginInterface.cpp - Target independent plugin device interface ---===//
 //
 // 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 "PluginInterface.h"
 #include "Debug.h"
 #include "GlobalHandler.h"
 #include "JIT.h"
+#include "Utilities.h"
 #include "elf_common.h"
 #include "omptarget.h"
 #include "omptargetplugin.h"
 
+#include "llvm/ADT/StringMap.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"
 
 #include <cstdint>
 #include <limits>
 
 using namespace llvm;
 using namespace omp;
 using namespace target;
 using namespace plugin;
 
 GenericPluginTy *Plugin::SpecificPlugin = nullptr;
 
 // TODO: Fix any thread safety issues for multi-threaded kernel recording.
 struct RecordReplayTy {
 private:
+  static constexpr int ALIGN = 16;
+
   // Memory pointers for recording, replaying memory.
-  void *MemoryStart;
-  void *MemoryPtr;
-  size_t MemorySize;
-  GenericDeviceTy *Device;
+  void *MemoryStart = nullptr;
+  void *AlignedMemoryStart = nullptr;
+  void *MemoryPtr = nullptr;
+  size_t MemorySize = 0;
+  GenericDeviceTy *Device = nullptr;
+
   std::mutex AllocationLock;
 
   // Environment variables for record and replay.
   // Enables recording kernels if set.
   BoolEnvar OMPX_RecordKernel;
   // Enables replaying a kernel if set.
   BoolEnvar OMPX_ReplayKernel;
   // Enables saving the device memory kernel output post execution if set.
   BoolEnvar OMPX_ReplaySaveOutput;
   // Sets the maximum to pre-allocate device memory.
   UInt32Envar OMPX_DeviceMemorySize;
 
   // Record/replay pre-allocates the largest possible device memory using the
   // default kind.
   // TODO: Expand allocation to include other kinds (device, host, shared) and
   // possibly use a MemoryManager to track (de-)allocations for
   // storing/retrieving when recording/replaying.
   Error preallocateDeviceMemory() {
     // Pre-allocate memory on device. Starts with 64GB and subtracts in steps
     // of 1GB until allocation succeeds.
-    const size_t MAX_MEMORY_ALLOCATION =
+    const size_t MaxMemoryAllocation =
         OMPX_DeviceMemorySize * 1024 * 1024 * 1024ULL;
     constexpr size_t STEP = 1024 * 1024 * 1024ULL;
     MemoryStart = nullptr;
-    for (size_t Try = MAX_MEMORY_ALLOCATION; Try > 0; Try -= STEP) {
+    size_t Size = MaxMemoryAllocation;
+    for (; Size > 0; Size -= STEP) {
       MemoryStart =
-          Device->allocate(Try, /* HstPtr */ nullptr, TARGET_ALLOC_DEFAULT);
+          Device->allocate(Size, /* HstPtr */ nullptr, TARGET_ALLOC_DEFAULT);
       if (MemoryStart)
         break;
     }
 
-    if (!MemoryStart)
+    if (!MemoryStart && MaxMemoryAllocation)
       return Plugin::error("Allocating record/replay memory");
 
-    MemoryPtr = MemoryStart;
+    // Align the memory at two times the step size to avoid mismatch in the
+    // beginning of the memory region.
+    AlignedMemoryStart = alignPtr(MemoryStart, (2 * STEP));
+
+    INFO(OMP_INFOTYPE_PLUGIN_KERNEL, Device->getDeviceId(),
+         "Allocated %" PRIu64
+         " bytes at %p for record and replay, aligned it to %p\n",
+         Size, MemoryStart, AlignedMemoryStart);
+
+    MemoryPtr = AlignedMemoryStart;
     MemorySize = 0;
 
     return Plugin::success();
   }
 
-  void dumpDeviceMemory(StringRef Filename,
-                        AsyncInfoWrapperTy &AsyncInfoWrapper) {
+  void dumpGlobals(StringRef Filename, const DeviceImageTy &Image) {
+    int32_t Size = 0;
+
+    for (auto &OffloadEntry : Image.getOffloadEntryTable()) {
+      if (!OffloadEntry.size)
+        continue;
+      Size += std::strlen(OffloadEntry.name) + /* '\0' */ 1 +
+              /* OffloadEntry.size value */ sizeof(uint32_t) +
+              OffloadEntry.size;
+    }
+
+    ErrorOr<std::unique_ptr<WritableMemoryBuffer>> GlobalsMB =
+        WritableMemoryBuffer::getNewUninitMemBuffer(Size);
+    if (!GlobalsMB)
+      report_fatal_error("Error creating MemoryBuffer for globals memory");
+
+    void *BufferPtr = GlobalsMB.get()->getBufferStart();
+    for (auto &OffloadEntry : Image.getOffloadEntryTable()) {
+      if (!OffloadEntry.size)
+        continue;
+
+      int32_t NameLength = std::strlen(OffloadEntry.name) + 1;
+      memcpy(BufferPtr, OffloadEntry.name, NameLength);
+      BufferPtr = advanceVoidPtr(BufferPtr, NameLength);
+
+      *((uint32_t *)(BufferPtr)) = OffloadEntry.size;
+      BufferPtr = advanceVoidPtr(BufferPtr, sizeof(uint32_t));
+
+      auto Err = Plugin::success();
+      {
+        AsyncInfoWrapperTy AsyncInfoWrapper(Err, *Device, nullptr);
+        if (auto Err =
+                Device->dataRetrieve(BufferPtr, OffloadEntry.addr,
+                                     OffloadEntry.size, AsyncInfoWrapper))
+          report_fatal_error("Error retrieving data for global");
+      }
+      if (Err)
+        report_fatal_error("Error retrieving data for global");
+      BufferPtr = advanceVoidPtr(BufferPtr, OffloadEntry.size);
+    }
+    assert(BufferPtr == GlobalsMB->get()->getBufferEnd() &&
+           "Buffer over/under-filled.");
+    assert(Size == getPtrDiff(BufferPtr, GlobalsMB->get()->getBufferStart()) &&
+           "Buffer size mismatch");
+
+    StringRef GlobalsMemory(GlobalsMB.get()->getBufferStart(), Size);
+    std::error_code EC;
+    raw_fd_ostream OS(Filename, EC);
+    OS << GlobalsMemory;
+    OS.close();
+  }
+
+  void dumpDeviceMemory(StringRef Filename) {
     ErrorOr<std::unique_ptr<WritableMemoryBuffer>> DeviceMemoryMB =
         WritableMemoryBuffer::getNewUninitMemBuffer(MemorySize);
     if (!DeviceMemoryMB)
       report_fatal_error("Error creating MemoryBuffer for device memory");
 
-    auto Err = Device->dataRetrieve(DeviceMemoryMB.get()->getBufferStart(),
-                                    MemoryStart, MemorySize, AsyncInfoWrapper);
+    if (MemorySize) {
+      auto Err = Plugin::success();
+      {
+        AsyncInfoWrapperTy AsyncInfoWrapper(Err, *Device, nullptr);
+        if (auto Err = Device->dataRetrieve(
+                DeviceMemoryMB.get()->getBufferStart(), AlignedMemoryStart,
+                MemorySize, AsyncInfoWrapper))
+          report_fatal_error("Error retrieving data for target pointer");
+      }
       if (Err)
         report_fatal_error("Error retrieving data for target pointer");
+    }
 
     StringRef DeviceMemory(DeviceMemoryMB.get()->getBufferStart(), MemorySize);
     std::error_code EC;
     raw_fd_ostream OS(Filename, EC);
     if (EC)
       report_fatal_error("Error dumping memory to file " + Filename + " :" +
                          EC.message());
     OS << DeviceMemory;
if (auto TgtImageBitcode = Image.getTgtImageBitcode()) {
           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) {
+  void saveKernelInputInfo(const char *Name, const DeviceImageTy &Image,
+                           void **ArgPtrs, ptrdiff_t *ArgOffsets,
+                           int32_t NumArgs, uint64_t NumTeamsClause,
+                           uint32_t ThreadLimitClause, uint64_t LoopTripCount) {
     json::Object JsonKernelInfo;
     JsonKernelInfo["Name"] = Name;
     JsonKernelInfo["NumArgs"] = NumArgs;
     JsonKernelInfo["NumTeamsClause"] = NumTeamsClause;
     JsonKernelInfo["ThreadLimitClause"] = ThreadLimitClause;
     JsonKernelInfo["LoopTripCount"] = LoopTripCount;
     JsonKernelInfo["DeviceMemorySize"] = MemorySize;
     JsonKernelInfo["DeviceId"] = Device->getDeviceId();
 
     json::Array JsonArgPtrs;
     for (int I = 0; I < NumArgs; ++I)
       JsonArgPtrs.push_back((intptr_t)ArgPtrs[I]);
     JsonKernelInfo["ArgPtrs"] = json::Value(std::move(JsonArgPtrs));
 
     json::Array JsonArgOffsets;
     for (int I = 0; I < NumArgs; ++I)
       JsonArgOffsets.push_back(ArgOffsets[I]);
     JsonKernelInfo["ArgOffsets"] = json::Value(std::move(JsonArgOffsets));
 
     Twine KernelName(Name);
     Twine MemoryFilename = KernelName + ".memory";
-    dumpDeviceMemory(MemoryFilename.str(), AsyncInfoWrapper);
+    dumpDeviceMemory(MemoryFilename.str());
+
+    Twine GlobalUpdatesFilename = KernelName + ".globals";
+    dumpGlobals(GlobalUpdatesFilename.str(), Image);
 
     Twine JsonFilename = KernelName + ".json";
     std::error_code EC;
     raw_fd_ostream JsonOS(JsonFilename.str(), EC);
     if (EC)
       report_fatal_error("Error saving kernel json file : " +
                          StringRef(EC.message()));
     JsonOS << json::Value(std::move(JsonKernelInfo));
     JsonOS.close();
   }
 
-  void saveKernelOutputInfo(const char *Name,
-                            AsyncInfoWrapperTy &AsyncInfoWrapper) {
+  void saveKernelOutputInfo(const char *Name) {
     Twine OutputFilename =
         Twine(Name) + (isRecording() ? ".original.output" : ".replay.output");
-    dumpDeviceMemory(OutputFilename.str(), AsyncInfoWrapper);
+    dumpDeviceMemory(OutputFilename.str());
   }
 
   void *alloc(uint64_t Size) {
-    assert(MemoryStart && "Expected memory has been pre-allocated");
+    assert(Size && MemoryStart && "Expected memory has been pre-allocated");
     void *Alloc = nullptr;
-    constexpr int ALIGN = 16;
     // Assumes alignment is a power of 2.
     int64_t AlignedSize = Size + (ALIGN - 1) & (~(ALIGN - 1));
     std::lock_guard<std::mutex> LG(AllocationLock);
     Alloc = MemoryPtr;
-    MemoryPtr = (char *)MemoryPtr + AlignedSize;
+    MemoryPtr = advanceVoidPtr(MemoryPtr, AlignedSize);
     MemorySize += AlignedSize;
     return Alloc;
   }
 
   Error init(GenericDeviceTy *Device) {
     this->Device = Device;
     return preallocateDeviceMemory();
   }
 
+  Error initializeRecordedGlobals(DeviceImageTy &Image, const void *BufferPtr,
+                                  int32_t Size) {
+    StringMap<const __tgt_offload_entry *> OffloadEntryMap;
+    for (auto &OffloadEntry : Image.getOffloadEntryTable())
+      OffloadEntryMap[OffloadEntry.name] = &OffloadEntry;
+
+    GenericGlobalHandlerTy &GHandler = Plugin::get().getGlobalHandler();
+
+    const void *BufferEndPtr = advanceVoidPtr(BufferPtr, Size);
+    while (BufferPtr != BufferEndPtr) {
+      StringRef Name((const char *)BufferPtr);
+      BufferPtr = advanceVoidPtr(BufferPtr, Name.size() + 1);
+
+      uint32_t RecordedSize = *((const uint32_t *)(BufferPtr));
+      BufferPtr = advanceVoidPtr(BufferPtr, sizeof(uint32_t));
+
+      GlobalTy ImageGlobal(Name.str(), 0);
+      if (auto Err =
+              GHandler.getGlobalMetadataFromImage(*Device, Image, ImageGlobal))
+        return Err;
+
+      // Verify the global has the expected size to avoid memory corruption.
+      if (RecordedSize != ImageGlobal.getSize())
+        report_fatal_error(
+            "Recorded global " + Name +
+            " has unexpected size: " + std::to_string(ImageGlobal.getSize()) +
+            " vs " + std::to_string(RecordedSize));
+
+      // Set the host pointer of the global such that we read the value from the
+      // recording.
+      ImageGlobal.setPtr(const_cast<void *>(BufferPtr));
+      if (auto Err = GHandler.writeGlobalToDevice(*Device, Image, ImageGlobal))
+        return Err;
+      BufferPtr = advanceVoidPtr(BufferPtr, ImageGlobal.getSize());
+    }
+    return Plugin::success();
+  }
+
   void deinit() { Device->free(MemoryStart); }
 
 } RecordReplay;
 
 AsyncInfoWrapperTy::~AsyncInfoWrapperTy() {
   // If we used a local async info object we want synchronous behavior.
   // In that case, and assuming the current status code is OK, we will
   // synchronize explicitly when the object is deleted.
   if (AsyncInfoPtr == &LocalAsyncInfo && LocalAsyncInfo.Queue && !Err)
     Err = Device.synchronize(&LocalAsyncInfo);
 }
 
 Error GenericKernelTy::init(GenericDeviceTy &GenericDevice,
                             DeviceImageTy &Image) {
+  ImagePtr = &Image;
   PreferredNumThreads = getDefaultNumThreads(GenericDevice);
 
   MaxNumThreads = GenericDevice.getThreadLimit();
 
   DynamicMemorySize = GenericDevice.getDynamicMemorySize();
 
   if (RecordReplay.isRecording())
     RecordReplay.saveImage(Name, Image);
Error GenericDeviceTy::dataExchange(const void *SrcPtr, GenericDeviceTy &DstDev,
   Err = dataExchangeImpl(SrcPtr, DstDev, DstPtr, Size, AsyncInfoWrapper);
   return Err;
 }
 
 Error GenericDeviceTy::runTargetTeamRegion(
     void *EntryPtr, void **ArgPtrs, ptrdiff_t *ArgOffsets, int32_t NumArgs,
     uint64_t NumTeamsClause, uint32_t ThreadLimitClause, uint64_t LoopTripCount,
     __tgt_async_info *AsyncInfo) {
-  auto Err = Plugin::success();
-  AsyncInfoWrapperTy AsyncInfoWrapper(Err, *this, AsyncInfo);
+
+  bool RecordKernelOutput = RecordReplay.isRecordingOrReplaying() &&
+                            RecordReplay.isSaveOutputEnabled();
 
   GenericKernelTy &GenericKernel =
       *reinterpret_cast<GenericKernelTy *>(EntryPtr);
 
-  if (RecordReplay.isRecording())
+  // Saving kernel input information in recording mode is synchronous as is the
+  // initialization of the device memory.
+  // TODO: If we delay the write of the files we could do this asynchronous:
+  //        - schedule the memory H2D transfers,
+  //        - schedule the kernel info D2H transfers,
+  //        - schedule the kernel,
+  //        - schedule the memory D2H transfers (iff the output is saved),
+  //        - synchronize,
+  //        - write all the files.
+  if (RecordReplay.isRecording()) {
+    // Ensure we are done transfering memory from the device before we issue D2H
+    // copies for the kernel info.
+    if (AsyncInfo && AsyncInfo->Queue)
+      if (auto Err = synchronize(AsyncInfo))
+        return Err;
     RecordReplay.saveKernelInputInfo(
-        GenericKernel.getName(), ArgPtrs, ArgOffsets, NumArgs, NumTeamsClause,
-        ThreadLimitClause, LoopTripCount, AsyncInfoWrapper);
+        GenericKernel.getName(), GenericKernel.getImage(), ArgPtrs, ArgOffsets,
+        NumArgs, NumTeamsClause, ThreadLimitClause, LoopTripCount);
+  }
+
+  // Block to synchronize the launch iff the kernel output is recorded.
+  auto Err = Plugin::success();
+  {
+    AsyncInfoWrapperTy AsyncInfoWrapper(
+        Err, *this, RecordKernelOutput ? nullptr : AsyncInfo);
 
-  Err =
-      GenericKernel.launch(*this, ArgPtrs, ArgOffsets, NumArgs, NumTeamsClause,
-                           ThreadLimitClause, LoopTripCount, AsyncInfoWrapper);
-
-  if (RecordReplay.isRecordingOrReplaying() &&
-      RecordReplay.isSaveOutputEnabled())
-    RecordReplay.saveKernelOutputInfo(GenericKernel.getName(),
-                                      AsyncInfoWrapper);
+    Err = GenericKernel.launch(*this, ArgPtrs, ArgOffsets, NumArgs,
+                               NumTeamsClause, ThreadLimitClause, LoopTripCount,
+  }
+
+  // If we are recording the output the kernel was already synchronized and is
+  // known to be done by now. If not, the AsyncInfoWrapper will not synchronize
+  // if AsyncInfo is a proper object.
+  if (RecordKernelOutput)
+    RecordReplay.saveKernelOutputInfo(GenericKernel.getName());
 
   return Err;
 }
 
 Error GenericDeviceTy::initAsyncInfo(__tgt_async_info **AsyncInfoPtr) {
   assert(AsyncInfoPtr && "Invalid async info");
 
   *AsyncInfoPtr = new __tgt_async_info();
 
   auto Err = Plugin::success();
   AsyncInfoWrapperTy AsyncInfoWrapper(Err, *this, *AsyncInfoPtr);
   Err = initAsyncInfoImpl(AsyncInfoWrapper);
   return Err;
 }
 
 Error GenericDeviceTy::initDeviceInfo(__tgt_device_info *DeviceInfo) {
   assert(DeviceInfo && "Invalid device info");
 
   return initDeviceInfoImpl(DeviceInfo);
 }
 
+Error GenericDeviceTy::initializeRecordedGlobals(const void *GlobalEncoding,
+                                                 int32_t Size) {
+  assert(GlobalEncoding && Size && "Expected encoded globals");
+  assert(LoadedImages.size() == 1);
+  return RecordReplay.initializeRecordedGlobals(*LoadedImages.back(),
+                                                GlobalEncoding, Size);
+}
+
 Error GenericDeviceTy::printInfo() {
   // TODO: Print generic information here
   return printInfoImpl();
 }
 
 Error GenericDeviceTy::createEvent(void **EventPtrStorage) {
   return createEventImpl(EventPtrStorage);
 }
if (Err) {
     REPORT("Failure to initialize device info at " DPxMOD " on device %d: %s\n",
            DPxPTR(DeviceInfo), DeviceId, toString(std::move(Err)).data());
     return OFFLOAD_FAIL;
   }
 
   return OFFLOAD_SUCCESS;
 }
 
+int32_t __tgt_rtl_initialize_recorded_globals(int32_t DeviceId,
+                                              const void *GlobalsEncoding,
+                                              int32_t Size) {
+  auto Err = Plugin::get().getDevice(DeviceId).initializeRecordedGlobals(
+      GlobalsEncoding, Size);
+  if (Err)
+    REPORT("Failure to initialize recorded globlas on device %d: %s\n",
+           DeviceId, toString(std::move(Err)).data());
+
+  return (bool)Err;
+}
+
 #ifdef __cplusplus
 }
 #endif

openmp/libomptarget/src/device.cpp

 
 int32_t DeviceTy::destroyEvent(void *Event) {
   if (RTL->create_event)
     return RTL->destroy_event(RTLDeviceID, Event);
 
   return OFFLOAD_SUCCESS;
 }
 
+int32_t DeviceTy::initializeRecordedGlobals(const void *Globals, int32_t Size) {
+  assert(RTL->initialize_recorded_globals &&
+         "Plugin does not support initialization of recorded globals.");
+  return RTL->initialize_recorded_globals(RTLDeviceID, Globals, Size);
+}
+
 /// Check whether a device has an associated RTL and initialize it if it's not
 /// already initialized.
 bool deviceIsReady(int DeviceNum) {
   DP("Checking whether device %d is ready.\n", DeviceNum);
   // Devices.size() can only change while registering a new
   // library, so try to acquire the lock of RTLs' mutex.
   size_t DevicesSize;
   {

openmp/libomptarget/src/interface.cpp

 ///                    execution.
 /// \param LoopTripCount The pre-recorded value of the loop tripcount, if any.
 /// \return OMP_TGT_SUCCESS on success, OMP_TGT_FAIL on failure.
 EXTERN int __tgt_target_kernel_replay(ident_t *Loc, int64_t DeviceId,
                                       void *HostPtr, void *DeviceMemory,
                                       int64_t DeviceMemorySize, void **TgtArgs,
                                       ptrdiff_t *TgtOffsets, int32_t NumArgs,
                                       int32_t NumTeams, int32_t ThreadLimit,
-                                      uint64_t LoopTripCount) {
+                                      uint64_t LoopTripCount,
+                                      const void *UpdatedGlobals,
+                                      int64_t UpdatedGlobalsSize) {
 
   if (checkDeviceAndCtors(DeviceId, Loc)) {
     DP("Not offloading to device %" PRId64 "\n", DeviceId);
     return OMP_TGT_FAIL;
   }
   DeviceTy &Device = *PM->Devices[DeviceId];
 
   AsyncInfoTy AsyncInfo(Device);
   int Rc = target_replay(Loc, Device, HostPtr, DeviceMemory, DeviceMemorySize,
                          TgtArgs, TgtOffsets, NumArgs, NumTeams, ThreadLimit,
-                         LoopTripCount, AsyncInfo);
+                         LoopTripCount, UpdatedGlobals, UpdatedGlobalsSize,
+                         AsyncInfo);
   if (Rc == OFFLOAD_SUCCESS)
     Rc = AsyncInfo.synchronize();
   handleTargetOutcome(Rc == OFFLOAD_SUCCESS, Loc);
   assert(Rc == OFFLOAD_SUCCESS &&
          "__tgt_target_kernel_replay unexpected failure!");
   return OMP_TGT_SUCCESS;
 }
 

openmp/libomptarget/src/omptarget.cpp

 
 /// Executes a kernel using pre-recorded information for loading to
 /// device memory to launch the target kernel with the pre-recorded
 /// configuration.
 int target_replay(ident_t *Loc, DeviceTy &Device, void *HostPtr,
                   void *DeviceMemory, int64_t DeviceMemorySize, void **TgtArgs,
                   ptrdiff_t *TgtOffsets, int32_t NumArgs, int32_t NumTeams,
                   int32_t ThreadLimit, uint64_t LoopTripCount,
+                  const void *GlobalsEncoding, int32_t GlobalsEncodingSize,
                   AsyncInfoTy &AsyncInfo) {
   int32_t DeviceId = Device.DeviceID;
   TableMap *TM = getTableMap(HostPtr);
   // Fail if the table map fails to find the target kernel pointer for the
   // provided host pointer.
   if (!TM) {
     REPORT("Host ptr " DPxMOD " does not have a matching target pointer.\n",
            DPxPTR(HostPtr));
int target_replay(ident_t *Loc, DeviceTy &Device, void *HostPtr,
   assert(TargetTable && "Global data has not been mapped\n");
 
   // Retrieve the target kernel pointer, allocate and store the recorded device
   // memory data, and launch device execution.
   void *TgtEntryPtr = TargetTable->EntriesBegin[TM->Index].addr;
   DP("Launching target execution %s with pointer " DPxMOD " (index=%d).\n",
      TargetTable->EntriesBegin[TM->Index].name, DPxPTR(TgtEntryPtr), TM->Index);
 
+  if (GlobalsEncodingSize)
+    Device.initializeRecordedGlobals(GlobalsEncoding, GlobalsEncodingSize);
+
+  if (DeviceMemorySize) {
     void *TgtPtr = Device.allocData(DeviceMemorySize, /* HstPtr */ nullptr,
                                     TARGET_ALLOC_DEFAULT);
     Device.submitData(TgtPtr, DeviceMemory, DeviceMemorySize, AsyncInfo);
+  }
 
   int Ret =
       Device.runTeamRegion(TgtEntryPtr, TgtArgs, TgtOffsets, NumArgs, NumTeams,
                            ThreadLimit, LoopTripCount, AsyncInfo);
 
   if (Ret != OFFLOAD_SUCCESS) {
     REPORT("Executing target region abort target.\n");
     return OFFLOAD_FAIL;
   }
 
   return OFFLOAD_SUCCESS;
 }

openmp/libomptarget/src/private.h

extern int target(ident_t *Loc, DeviceTy &Device, void *HostPtr, int32_t ArgNum,
                   void **ArgMappers, int32_t TeamNum, int32_t ThreadLimit,
                   uint64_t Tripcount, int IsTeamConstruct,
                   AsyncInfoTy &AsyncInfo);
 
 extern int target_replay(ident_t *Loc, DeviceTy &Device, void *HostPtr,
                          void *DeviceMemory, int64_t DeviceMemorySize,
                          void **TgtArgs, ptrdiff_t *TgtOffsets, int32_t NumArgs,
                          int32_t NumTeams, int32_t ThreadLimit,
-                         uint64_t LoopTripCount, AsyncInfoTy &AsyncInfo);
+                         uint64_t LoopTripCount, const void *UpdatedGlobals,
+                         int32_t UpdatedGlobalsSize, AsyncInfoTy &AsyncInfo);
 
 extern void handleTargetOutcome(bool Success, ident_t *Loc);
 extern bool checkDeviceAndCtors(int64_t &DeviceID, ident_t *Loc);
 extern void *targetAllocExplicit(size_t Size, int DeviceNum, int Kind,
                                  const char *Name);
 extern void targetFreeExplicit(void *DevicePtr, int DeviceNum, int Kind,
                                const char *Name);
 extern void *targetLockExplicit(void *HostPtr, size_t Size, int DeviceNum,

openmp/libomptarget/src/rtl.cpp

#endif
   *((void **)&RTL.init_async_info) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_init_async_info");
   *((void **)&RTL.init_device_info) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_init_device_info");
   *((void **)&RTL.data_lock) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_data_lock");
   *((void **)&RTL.data_unlock) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_data_unlock");
+  *((void **)&RTL.initialize_recorded_globals) =
+      DynLibrary->getAddressOfSymbol("__tgt_rtl_initialize_recorded_globals");
 
   RTL.LibraryHandler = std::move(DynLibrary);
 
   // Successfully loaded
   return true;
 }
 
 ////////////////////////////////////////////////////////////////////////////////

openmp/libomptarget/tools/kernelreplay/llvm-omp-kernel-replay.cpp

int main(int argc, char **argv) {
   DeviceImage.EntriesEnd = &KernelEntry + 1;
 
   __tgt_bin_desc Desc;
   Desc.NumDeviceImages = 1;
   Desc.HostEntriesBegin = &KernelEntry;
   Desc.HostEntriesEnd = &KernelEntry + 1;
   Desc.DeviceImages = &DeviceImage;
 
+  // Read in the entire file (hence the "as stream"), since we might lock the
+  // memory.
   ErrorOr<std::unique_ptr<MemoryBuffer>> DeviceMemoryMB =
-      MemoryBuffer::getFile(KernelEntryName + ".memory", /* isText */ false,
-                            /* RequiresNullTerminator */ false);
+      MemoryBuffer::getFileAsStream(KernelEntryName + ".memory");
   if (!DeviceMemoryMB)
     report_fatal_error("Error reading the kernel input device memory.");
 
+  ErrorOr<std::unique_ptr<MemoryBuffer>> GlobalsMB =
+      MemoryBuffer::getFileAsStream(KernelEntryName + ".globals");
+  if (!GlobalsMB)
+    report_fatal_error("Error reading the globals state file.");
+
   setenv("LIBOMPTARGET_REPLAY", "1", 1);
   if (VerifyOpt || SaveOutputOpt)
     setenv("LIBOMPTARGET_RR_SAVE_OUTPUT", "1", 1);
 
   auto DeviceMemorySizeJson =
       JsonKernelInfo->getAsObject()->getInteger("DeviceMemorySize");
   // Set device memory size to the ceiling of GB granularity.
   uint64_t DeviceMemorySize =
       std::ceil(DeviceMemorySizeJson.value() / (1024.0 * 1024.0 * 1024.0));
 
-  setenv("LIBOMPTARGET_RR_DEVMEM_SIZE",
-         std::to_string(DeviceMemorySize).c_str(), 1);
+  // TODO: Do not explicitly set device memory size to avoid messing the
+  // aligned memory start workaround. In the genral case, we will need
+  // to implement custom memory mapping to avoid misaligned memory
+  // addresses between recording and replaying a kernel.
+  //  setenv("LIBOMPTARGET_RR_DEVMEM_SIZE",
+  //         std::to_string(DeviceMemorySize).c_str(), 1);
 
   auto DeviceIdJson = JsonKernelInfo->getAsObject()->getInteger("DeviceId");
   // TODO: Print warning if the user overrides the device id in the json file.
   int32_t DeviceId = (DeviceIdOpt > -1 ? DeviceIdOpt : DeviceIdJson.value());
 
   // TODO: do we need requires?
   //__tgt_register_requires(/* Flags */1);
 
   __tgt_init_all_rtls();
 
   __tgt_register_lib(&Desc);
 
   __tgt_target_kernel_replay(
       /* Loc */ nullptr, DeviceId, KernelEntry.addr,
       (void *)DeviceMemoryMB.get()->getBuffer().data(),
       DeviceMemoryMB.get()->getBufferSize(), TgtArgs.data(),
       TgtArgOffsets.data(), NumArgs.value(), NumTeams, NumThreads,
-      LoopTripCount.value());
+      LoopTripCount.value(), GlobalsMB.get()->getBufferStart(),
+      GlobalsMB.get()->getBufferSize());
 
   if (VerifyOpt) {
     ErrorOr<std::unique_ptr<MemoryBuffer>> OriginalOutputMB =
         MemoryBuffer::getFile(KernelEntryName + ".original.output",
                               /* isText */ false,
                               /* RequiresNullTerminator */ false);
     if (!OriginalOutputMB)
       report_fatal_error("Error reading the kernel original output file, make "