Support for OpenMP-Offload Record Replay

[koparasy]

Summary

Enable record-replay for OpenMP offload kernels. On recording the initialization is performed on device initialization by reading env variables. (This is similar to the way rr used to operate). The primary change takes place in the replay phase with the replay tool explicitly initializing the record-replay functionality.

[jdoerfert]

I like the change to an explicit initialization rather than a env var based one (for the replay tool).
There are minor style changes we can fix in tree, and one issue with a missing sync. We also need tests.
All that said, I think we should get this into 17.

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

173

There is some kind of sync missing here. We should not use an async retrieve to begin with, or wait for it.

[koparasy]

Add synchronize between data transfers

openmp/libomptarget/include/omptarget.h

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);
 
 void __tgt_set_info_flag(uint32_t);
 
 int __tgt_print_device_info(int64_t DeviceId);
+
+int __tgt_activate_record_replay(int64_t DeviceId, uint64_t MemorySize,
+                                 bool IsRecord, bool SaveOutput);
+
 #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 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 *,
                                        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(data_notify_mapped_ty)(int32_t, void *, int64_t);
   typedef int32_t(data_notify_unmapped_ty)(int32_t, void *);
+  typedef int32_t(activate_record_replay_ty)(int32_t, uint64_t, bool, bool);
 
   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
   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;
   data_notify_mapped_ty *data_notify_mapped = nullptr;
   data_notify_unmapped_ty *data_notify_unmapped = nullptr;
+  activate_record_replay_ty *activate_record_replay = 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/JIT.cpp

if (!FOrErr) {
     consumeError(FOrErr.takeError());
     return false;
   }
 
   auto ActualTriple = FOrErr->TheReader.getTargetTriple();
   auto BitcodeTA = Triple(ActualTriple).getArch();
   BitcodeImageMap[Image.ImageStart] = BitcodeTA;
 
+  DP("Is%s IR Image\n", BitcodeTA == TT.getArch() ? " " : " NOT");
+
   return BitcodeTA == TT.getArch();
 }

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;
struct GenericKernelTy {
   virtual Error launchImpl(GenericDeviceTy &GenericDevice, uint32_t NumThreads,
                            uint64_t NumBlocks, KernelArgsTy &KernelArgs,
                            void *Args,
                            AsyncInfoWrapperTy &AsyncInfoWrapper) const = 0;
 
   /// Get the kernel name.
   const char *getName() const { return Name; }
 
+  /// Get the kernel image.
+  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:
   bool isSPMDMode() const { return ExecutionMode == OMP_TGT_EXEC_MODE_SPMD; }
 
   /// 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 preferred number of threads to run the kernel.
   uint32_t PreferredNumThreads;
 
   /// The maximum number of threads which the kernel could leverage.
   uint32_t MaxNumThreads;
 };
 

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

 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 {
+
+  // Describes the state of the record replay mechanism.
+  enum RRStatusTy { RRDeactivated = 0, RRRecording, RRReplaying };
+
 private:
   // Memory pointers for recording, replaying memory.
   void *MemoryStart;
   void *MemoryPtr;
   size_t MemorySize;
   GenericDeviceTy *Device;
   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;
+  RRStatusTy Status;
+  bool ReplaySaveOutput;
+  uint64_t 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() {
+  Error preallocateDeviceMemory(uint64_t DeviceMemorySize) {
     // Pre-allocate memory on device. Starts with 64GB and subtracts in steps
     // of 1GB until allocation succeeds.
-    const size_t MAX_MEMORY_ALLOCATION =
-        OMPX_DeviceMemorySize * 1024 * 1024 * 1024ULL;
+    const size_t MAX_MEMORY_ALLOCATION = DeviceMemorySize;
     constexpr size_t STEP = 1024 * 1024 * 1024ULL;
     MemoryStart = nullptr;
     for (size_t Try = MAX_MEMORY_ALLOCATION; Try > 0; Try -= STEP) {
       MemoryStart =
           Device->allocate(Try, /* HstPtr */ nullptr, TARGET_ALLOC_DEFAULT);
       if (MemoryStart)
         break;
     }
void dumpDeviceMemory(StringRef Filename,
     if (EC)
       report_fatal_error("Error dumping memory to file " + Filename + " :" +
                          EC.message());
     OS << DeviceMemory;
     OS.close();
   }
 
 public:
-  bool isRecording() const { return OMPX_RecordKernel; }
-  bool isReplaying() const { return OMPX_ReplayKernel; }
+  bool isRecording() const { return Status == RRStatusTy::RRRecording; }
+  bool isReplaying() const { return Status == RRStatusTy::RRReplaying; }
   bool isRecordingOrReplaying() const {
-    return (OMPX_RecordKernel || OMPX_ReplayKernel);
+    return (Status != RRStatusTy::RRDeactivated);
   }
-  bool isSaveOutputEnabled() const { return OMPX_ReplaySaveOutput; }
+  void setStatus(RRStatusTy Status) { this->Status = Status; }
+  bool isSaveOutputEnabled() const { return ReplaySaveOutput; }
 
   RecordReplayTy()
-      : OMPX_RecordKernel("LIBOMPTARGET_RECORD"),
-        OMPX_ReplayKernel("LIBOMPTARGET_REPLAY"),
-        OMPX_ReplaySaveOutput("LIBOMPTARGET_RR_SAVE_OUTPUT"),
-        OMPX_DeviceMemorySize("LIBOMPTARGET_RR_DEVMEM_SIZE",
-                              /* Default in GB */ 64) {}
+      : Status(RRStatusTy::RRDeactivated), ReplaySaveOutput(false),
+        DeviceMemorySize(-1) {}
 
-  void saveImage(const char *Name, DeviceImageTy &Image) {
+  void saveImage(const char *Name, const DeviceImageTy &Image) {
     SmallString<128> ImageName = {Name, ".image"};
     std::error_code EC;
     raw_fd_ostream OS(ImageName, EC);
     if (EC)
       report_fatal_error("Error saving image : " + StringRef(EC.message()));
     if (const 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,
+  void dumpGlobals(StringRef Filename, DeviceImageTy &Image,
+                   AsyncInfoWrapperTy &AsyncInfoWrapper) {
+    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();
+      {
+        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 saveKernelInputInfo(const char *Name, DeviceImageTy &Image,
+                           void **ArgPtrs, ptrdiff_t *ArgOffsets,
+                           int32_t NumArgs, uint64_t NumTeamsClause,
+                           uint32_t ThreadLimitClause, uint64_t LoopTripCount,
                            AsyncInfoWrapperTy &AsyncInfoWrapper) {
     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));
 
     SmallString<128> MemoryFilename = {Name, ".memory"};
     dumpDeviceMemory(MemoryFilename, AsyncInfoWrapper);
 
+    SmallString<128> GlobalsFilename = {Name, ".globals"};
+    dumpGlobals(GlobalsFilename, Image, AsyncInfoWrapper);
+
     SmallString<128> JsonFilename = {Name, ".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 *alloc(uint64_t Size) {
     void *Alloc = nullptr;
     constexpr int Alignment = 16;
     // Assumes alignment is a power of 2.
     int64_t AlignedSize = (Size + (Alignment - 1)) & (~(Alignment - 1));
     std::lock_guard<std::mutex> LG(AllocationLock);
     Alloc = MemoryPtr;
     MemoryPtr = (char *)MemoryPtr + AlignedSize;
     MemorySize += AlignedSize;
+    DP("Memory Allocator return " DPxMOD "\n", DPxPTR(Alloc));
     return Alloc;
   }
 
-  Error init(GenericDeviceTy *Device) {
+  Error init(GenericDeviceTy *Device, uint64_t MemSize, RRStatusTy Status,
+             bool SaveOutput) {
     this->Device = Device;
-    return preallocateDeviceMemory();
+    this->Status = Status;
+    this->DeviceMemorySize = MemSize;
+    this->ReplaySaveOutput = SaveOutput;
+
+    if (auto Err = preallocateDeviceMemory(MemSize))
+      return Err;
+
+    INFO(OMP_INFOTYPE_PLUGIN_KERNEL, Device->getDeviceId(),
+         "Record Replay Initialized (%p)"
+         " as starting address, %lu Memory Size"
+         " and set on status %s\n",
+         MemoryStart, MemSize,
+         Status == RRStatusTy::RRRecording ? "Recording" : "Replaying");
+
+    return Plugin::success();
   }
 
   void deinit() { Device->free(MemoryStart); }
 
 } RecordReplay;
 
 AsyncInfoWrapperTy::AsyncInfoWrapperTy(GenericDeviceTy &Device,
                                        __tgt_async_info *AsyncInfoPtr)
if (AsyncInfoPtr == &LocalAsyncInfo && LocalAsyncInfo.Queue && !Err)
     Err = Device.synchronize(&LocalAsyncInfo);
 
   // Invalidate the wrapper object.
   AsyncInfoPtr = nullptr;
 }
 
 Error GenericKernelTy::init(GenericDeviceTy &GenericDevice,
                             DeviceImageTy &Image) {
+
+  ImagePtr = &Image;
+
   PreferredNumThreads = getDefaultNumThreads(GenericDevice);
 
   MaxNumThreads = GenericDevice.getThreadLimit();
 
   return initImpl(GenericDevice, Image);
 }
 
 Error GenericKernelTy::printLaunchInfo(GenericDeviceTy &GenericDevice,
if (OMP_TeamsThreadLimit > 0)
     GridValues.GV_Max_WG_Size =
         std::min(GridValues.GV_Max_WG_Size, uint32_t(OMP_TeamsThreadLimit));
 
   // Enable the memory manager if required.
   auto [ThresholdMM, EnableMM] = MemoryManagerTy::getSizeThresholdFromEnv();
   if (EnableMM)
     MemoryManager = new MemoryManagerTy(*this, ThresholdMM);
 
-  if (RecordReplay.isRecordingOrReplaying())
-    if (auto Err = RecordReplay.init(this))
-      return Err;
-
   return Plugin::success();
 }
 
 Error GenericDeviceTy::deinit() {
   // Delete the memory manager before deinitializing the device. Otherwise,
   // we may delete device allocations after the device is deinitialized.
   if (MemoryManager)
     delete MemoryManager;
Error GenericDeviceTy::launchKernel(void *EntryPtr, void **ArgPtrs,
                                     __tgt_async_info *AsyncInfo) {
   AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);
 
   GenericKernelTy &GenericKernel =
       *reinterpret_cast<GenericKernelTy *>(EntryPtr);
 
   if (RecordReplay.isRecording())
     RecordReplay.saveKernelInputInfo(
-        GenericKernel.getName(), ArgPtrs, ArgOffsets, KernelArgs.NumArgs,
-        KernelArgs.NumTeams[0], KernelArgs.ThreadLimit[0], KernelArgs.Tripcount,
-        AsyncInfoWrapper);
+        GenericKernel.getName(), GenericKernel.getImage(), ArgPtrs, ArgOffsets,
+        KernelArgs.NumArgs, KernelArgs.NumTeams[0], KernelArgs.ThreadLimit[0],
+        KernelArgs.Tripcount, AsyncInfoWrapper);
+
+  if (RecordReplay.isRecording())
+    RecordReplay.saveImage(GenericKernel.getName(), GenericKernel.getImage());
 
   auto Err = GenericKernel.launch(*this, ArgPtrs, ArgOffsets, KernelArgs,
                                   AsyncInfoWrapper);
 
   if (RecordReplay.isRecordingOrReplaying() &&
       RecordReplay.isSaveOutputEnabled())
     RecordReplay.saveKernelOutputInfo(GenericKernel.getName(),
                                       AsyncInfoWrapper);
int64_t __tgt_rtl_init_requires(int64_t RequiresFlags) {
   return RequiresFlags;
 }
 
 int32_t __tgt_rtl_is_data_exchangable(int32_t SrcDeviceId,
                                       int32_t DstDeviceId) {
   return Plugin::get().isDataExchangable(SrcDeviceId, DstDeviceId);
 }
 
+int32_t __tgt_rtl_initialize_record_replay(int32_t DeviceId,
+                                           uint64_t MemorySize, bool isRecord,
+                                           bool SaveOutput) {
+  GenericPluginTy &Plugin = Plugin::get();
+  GenericDeviceTy &Device = Plugin.getDevice(DeviceId);
+  RecordReplayTy::RRStatusTy Status =
+      isRecord ? RecordReplayTy::RRStatusTy::RRRecording
+               : RecordReplayTy::RRStatusTy::RRReplaying;
+
+  if (auto Err = RecordReplay.init(&Device, MemorySize, Status, SaveOutput)) {
+    REPORT("WARNING RR did not intialize RR-properly with %lu bytes"
+           "(Error: %s)\n",
+           MemorySize, toString(std::move(Err)).data());
+    RecordReplay.setStatus(RecordReplayTy::RRStatusTy::RRDeactivated);
+
+    if (!isRecord) {
+      return OFFLOAD_FAIL;
+    }
+  }
+  return OFFLOAD_SUCCESS;
+}
+
 __tgt_target_table *__tgt_rtl_load_binary(int32_t DeviceId,
                                           __tgt_device_image *TgtImage) {
   GenericPluginTy &Plugin = Plugin::get();
   GenericDeviceTy &Device = Plugin.getDevice(DeviceId);
 
   auto TableOrErr = Device.loadBinary(Plugin, TgtImage);
   if (!TableOrErr) {
     auto Err = TableOrErr.takeError();

openmp/libomptarget/src/device.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "device.h"
 #include "omptarget.h"
 #include "private.h"
 #include "rtl.h"
 
+#include "Utilities.h"
+
 #include <cassert>
 #include <climits>
 #include <cstdint>
 #include <cstdio>
 #include <mutex>
 #include <string>
 #include <thread>
 

 void DeviceTy::init() {
   // Make call to init_requires if it exists for this plugin.
   if (RTL->init_requires)
     RTL->init_requires(PM->RTLs.RequiresFlags);
   int32_t Ret = RTL->init_device(RTLDeviceID);
   if (Ret != OFFLOAD_SUCCESS)
     return;
 
+  // Enables recording kernels if set.
+  llvm::omp::target::BoolEnvar OMPX_RecordKernel("LIBOMPTARGET_RECORD", false);
+  if (OMPX_RecordKernel) {
+    // Enables saving the device memory kernel output post execution if set.
+    llvm::omp::target::BoolEnvar OMPX_ReplaySaveOutput(
+        "LIBOMPTARGET_RR_SAVE_OUTPUT", false);
+    // Sets the maximum to pre-allocate device memory.
+    llvm::omp::target::UInt64Envar OMPX_DeviceMemorySize(
+        "LIBOMPTARGET_RR_DEVMEM_SIZE", 16);
+    DP("Activating Record-Replay for Device %d with %lu GB memory\n",
+       RTLDeviceID, OMPX_DeviceMemorySize);
+
+    RTL->activate_record_replay(RTLDeviceID,
+                                OMPX_DeviceMemorySize * 1024 * 1024 * 1024,
+                                true, OMPX_ReplaySaveOutput);
+  }
+
   IsInit = true;
 }
 
 /// Thread-safe method to initialize the device only once.
 int32_t DeviceTy::initOnce() {
   std::call_once(InitFlag, &DeviceTy::init, this);
 
   // At this point, if IsInit is true, then either this thread or some other

openmp/libomptarget/src/exports

global:
     __tgt_target_data_end_nowait_mapper;
     __tgt_target_data_update_nowait_mapper;
     __tgt_target_nowait_mapper;
     __tgt_target_teams_nowait_mapper;
     __tgt_target_kernel;
     __tgt_target_kernel_nowait;
     __tgt_target_nowait_query;
     __tgt_target_kernel_replay;
+    __tgt_activate_record_replay;
     __tgt_mapper_num_components;
     __tgt_push_mapper_component;
     __kmpc_push_target_tripcount;
     __kmpc_push_target_tripcount_mapper;
     omp_get_mapped_ptr;
     omp_get_num_devices;
     omp_get_device_num;
     omp_get_initial_device;

openmp/libomptarget/src/interface.cpp

EXTERN int __tgt_target_kernel(ident_t *Loc, int64_t DeviceId, int32_t NumTeams,
   if (KernelArgs->Flags.NoWait)
     return targetKernel<TaskAsyncInfoWrapperTy>(
         Loc, DeviceId, NumTeams, ThreadLimit, HostPtr, KernelArgs);
   else
     return targetKernel<AsyncInfoTy>(Loc, DeviceId, NumTeams, ThreadLimit,
                                      HostPtr, KernelArgs);
 }
 
+/// Activates the record replay mechanism.
+/// \param DeviceId The device identifier to execute the target region.
+/// \param MemorySize The number of bytes to be (pre-)allocated
+///                   by the bump allocator
+/// /param IsRecord Activates the record replay mechanism in
+///                 'record' mode or 'replay' mode.
+/// /param SaveOutput Store the device memory after kernel
+///                   execution on persistent storage
+EXTERN int __tgt_activate_record_replay(int64_t DeviceId, uint64_t MemorySize,
+                                        bool IsRecord, bool SaveOutput) {
+  if (!deviceIsReady(DeviceId)) {
+    DP("Device %" PRId64 " is not ready\n", DeviceId);
+    return OMP_TGT_FAIL;
+  }
+
+  DeviceTy &Device = *PM->Devices[DeviceId];
+  int Rc = target_activate_rr(Device, MemorySize, IsRecord, SaveOutput);
+  assert(Rc == OFFLOAD_SUCCESS &&
+         "__tgt_activate_record_replay unexpected failure!");
+  return OMP_TGT_SUCCESS;
+};
+
 /// Implements a target kernel entry that replays a pre-recorded kernel.
 /// \param Loc Source location associated with this target region (unused).
 /// \param DeviceId The device identifier to execute the target region.
 /// \param HostPtr A pointer to an address that uniquely identifies the kernel.
 /// \param DeviceMemory A pointer to an array storing device memory data to move
 ///                     prior to kernel execution.
 /// \param DeviceMemorySize The size of the above device memory data in bytes.
 /// \param TgtArgs An array of pointers of the pre-recorded target kernel

openmp/libomptarget/src/omptarget.cpp

if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Failed to process data after launching the kernel.\n");
       return OFFLOAD_FAIL;
     }
   }
 
   return OFFLOAD_SUCCESS;
 }
 
+/// Enables the record replay mechanism by pre-allocating MemorySize
+/// and informing the record-replayer of whether to store the output
+/// in some file.
+int target_activate_rr(DeviceTy &Device, uint64_t MemorySize, bool isRecord,
+                       bool SaveOutput) {
+  return Device.RTL->activate_record_replay(Device.DeviceID, MemorySize,
+                                            isRecord, SaveOutput);
+}
+
 /// 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,
                   AsyncInfoTy &AsyncInfo) {

openmp/libomptarget/src/private.h

extern int targetDataUpdate(ident_t *Loc, DeviceTy &Device, int32_t ArgNum,
                             void **ArgsBase, void **Args, int64_t *ArgSizes,
                             int64_t *ArgTypes, map_var_info_t *ArgNames,
                             void **ArgMappers, AsyncInfoTy &AsyncInfo,
                             bool FromMapper = false);
 
 extern int target(ident_t *Loc, DeviceTy &Device, void *HostPtr,
                   KernelArgsTy &KernelArgs, AsyncInfoTy &AsyncInfo);
 
+extern int target_activate_rr(DeviceTy &Device, uint64_t MemorySize,
+                              bool isRecord, bool SaveOutput);
+
 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);
 
 extern void handleTargetOutcome(bool Success, ident_t *Loc);
 extern bool checkDeviceAndCtors(int64_t &DeviceID, ident_t *Loc);

openmp/libomptarget/src/rtl.cpp

*((void **)&RTL.data_lock) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_data_lock");
   *((void **)&RTL.data_unlock) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_data_unlock");
   *((void **)&RTL.data_notify_mapped) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_data_notify_mapped");
   *((void **)&RTL.data_notify_unmapped) =
       DynLibrary->getAddressOfSymbol("__tgt_rtl_data_notify_unmapped");
 
+  // Record Replay RTL
+  *((void **)&RTL.activate_record_replay) =
+      DynLibrary->getAddressOfSymbol("__tgt_rtl_initialize_record_replay");
+
   RTL.LibraryHandler = std::move(DynLibrary);
 
   // Successfully loaded
   return true;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Functionality for registering libs

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

 //===- llvm-omp-kernel-replay.cpp - Replay OpenMP offload kernel ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This is a command line utility to replay the execution of recorded OpenMP
 // offload kernels.
 //
 //===----------------------------------------------------------------------===//
 
+#include "omptarget.h"
 #include "omptargetplugin.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include <cstdlib>
 
 using namespace llvm;
 
int main(int argc, char **argv) {
   Desc.DeviceImages = &DeviceImage;
 
   ErrorOr<std::unique_ptr<MemoryBuffer>> DeviceMemoryMB =
       MemoryBuffer::getFile(KernelEntryName + ".memory", /* isText */ false,
                             /* RequiresNullTerminator */ false);
   if (!DeviceMemoryMB)
     report_fatal_error("Error reading the kernel input device memory.");
 
-  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);
+  uint64_t DeviceMemorySize = std::ceil(DeviceMemorySizeJson.value());
 
   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);
 
+  int Rc = __tgt_activate_record_replay(DeviceId, DeviceMemorySize, false,
+                                        VerifyOpt);
+  if (Rc != OMP_TGT_SUCCESS) {
+    report_fatal_error("Cannot activate record replay\n");
+  }
+
   __tgt_target_kernel_replay(
       /* Loc */ nullptr, DeviceId, KernelEntry.addr,
       const_cast<char *>(DeviceMemoryMB.get()->getBuffer().data()),
       DeviceMemoryMB.get()->getBufferSize(), TgtArgs.data(),
       TgtArgOffsets.data(), NumArgs.value(), NumTeams, NumThreads,
       LoopTripCount.value());
 
   if (VerifyOpt) {