OpenMP-Offload Add VA Management for Record-Replay Mechanism

[koparasy]

Summary

Extends the record-replay mechanism to re use the same VA between record and replay runs. This provides consistent memory addresses across runs.

[jdoerfert]

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

63

utils::alignUp(Addr, MaxMemoryAllocation)
Or sth like that.

65

I don't get this. I would at least not do this if Addr < MaxMemoryAllocation. Even if Addr > MaxMemoryAllocation, the space before Addr could very well be allocated already. Would this not be a problem?

71–74

Clang format the commit.

95

Heuristic

316–318

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

647

650

-cuda

651

+ actual number of bytes allocated by this function.
Mention if it can be lower/larger or both

openmp/libomptarget/plugins-nextgen/cuda/dynamic_cuda/cuda.h

114

Remove the comments. We don't have any elsewhere.

openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp

564

utils::roundUp

602–604

983

-memory
s/handlers/handles

openmp/libomptarget/src/interface.cpp

342

alignment

openmp/libomptarget/src/private.h

44

ReqAddr

openmp/libomptarget/include/omptarget.h

int __tgt_target_kernel_replay(ident_t *Loc, int64_t DeviceId, void *HostPtr,
                                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,
+int __tgt_activate_record_replay(int64_t DeviceId, uint64_t MemorySize, void *VAddr,
                                  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);
+  typedef int32_t(activate_record_replay_ty)(int32_t, uint64_t, void *, 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;

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

Error setDeviceStackSize(uint64_t Value) override {
     return Plugin::success();
   }
   Error getDeviceHeapSize(uint64_t &Value) override {
     Value = 0;
     return Plugin::success();
   }
   Error setDeviceHeapSize(uint64_t Value) override { return Plugin::success(); }
 
+  Error getDeviceMemorySize(uint64_t &Value) override {
+    for (AMDGPUMemoryPoolTy *Pool : AllMemoryPools) {
+      if (Pool->isGlobal()){
+        hsa_status_t Status = Pool->getAttrRaw(HSA_AMD_MEMORY_POOL_INFO_SIZE, Value);
+        return Plugin::check(Status, "Error in getDeviceMemorySize %s");
+      }
+    }
+    return Plugin::error("getDeviceMemorySize:: no global pool");
+  }
+
   /// AMDGPU-specific function to get device attributes.
   template <typename Ty> Error getDeviceAttr(uint32_t Kind, Ty &Value) {
     hsa_status_t Status =
         hsa_agent_get_info(Agent, (hsa_agent_info_t)Kind, &Value);
     return Plugin::check(Status, "Error in hsa_agent_get_info: %s");
   }
 
   template <typename Ty>

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

struct GenericDeviceTy : public DeviceAllocatorTy {
   Error synchronize(__tgt_async_info *AsyncInfo);
   virtual Error synchronizeImpl(__tgt_async_info &AsyncInfo) = 0;
 
   /// Query for the completion of the pending operations on the __tgt_async_info
   /// structure in a non-blocking manner.
   Error queryAsync(__tgt_async_info *AsyncInfo);
   virtual Error queryAsyncImpl(__tgt_async_info &AsyncInfo) = 0;
 
+  /// Check whether the architecture
+  virtual bool supportVAManagement() const { return false; }
+
+  /// De-allocates device memory and Unmaps the Virtual Addr
+  virtual Error memoryVAUnMap(void *VAddr, size_t Size);
+
+  /// Allocates \p RSize bytes (rounded up to page size) and hints the cuda driver to map it to \p VAddr.
+  /// The obtained address is stored in \p Addr. At return \p RSize contains the actual size
+  virtual Error memoryVAMap(void **Addr, void *VAddr, size_t *RSize);
+
   /// Allocate data on the device or involving the device.
   Expected<void *> dataAlloc(int64_t Size, void *HostPtr, TargetAllocTy Kind);
 
   /// Deallocate data from the device or involving the device.
   Error dataDelete(void *TgtPtr, TargetAllocTy Kind);
 
   /// Pin host memory to optimize transfers and return the device accessible
   /// pointer that devices should use for memory transfers involving the host
uint32_t getDefaultNumThreads() const {
     return GridValues.GV_Default_WG_Size;
   }
   uint32_t getDefaultNumBlocks() const {
     return GridValues.GV_Default_Num_Teams;
   }
   uint32_t getDynamicMemorySize() const { return OMPX_SharedMemorySize; }
   virtual uint64_t getClockFrequency() const { return CLOCKS_PER_SEC; }
 
+  virtual Error getDeviceMemorySize(uint64_t &DSize);
+
   /// Get target compute unit kind (e.g., sm_80, or gfx908).
   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);
   }

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

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;
+  size_t TotalSize;
   GenericDeviceTy *Device;
   std::mutex AllocationLock;
 
   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(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 = DeviceMemorySize;
+  void *suggestAddress(uint64_t MaxMemoryAllocation){
+    // Get a valid pointer address for this system
+    void *Addr = Device->allocate(1024, /* HstPtr */ nullptr, TARGET_ALLOC_DEFAULT);
+    Device->free(Addr);
+    // Align Address to MaxMemoryAllocation
+    Addr = (void *) (((uintptr_t) Addr + (MaxMemoryAllocation - 1)) & (~(MaxMemoryAllocation - 1)));
+    // Pad Address by MaxMemoryAllocation to guarantee enough space
+    Addr = (void *) ((uintptr_t) Addr - MaxMemoryAllocation);
+    return Addr;
+  }
+
+  Error preAllocateVAMemory(uint64_t MaxMemoryAllocation, void *VAddr){
+    size_t ASize = MaxMemoryAllocation;
+
+    if ( !VAddr && isRecording() ){
+      VAddr = suggestAddress(MaxMemoryAllocation);
+    }
+
+    DP("Request %ld bytes allocated at %p\n", MaxMemoryAllocation, VAddr);
+
+    if ( auto Err = Device->memoryVAMap(&MemoryStart, VAddr, &ASize) )
+      return Err;
+
+    if ( isReplaying() && VAddr != MemoryStart ){
+      return Plugin::error("Record-Replay cannot assign the"
+          "requested recorded address (%p, %p)", VAddr, MemoryStart);
+    }
+
+    INFO(OMP_INFOTYPE_PLUGIN_KERNEL, Device->getDeviceId(),
+       "Allocated %" PRIu64 " bytes at %p for replay.\n", ASize, MemoryStart);
+
+    MemoryPtr = MemoryStart;
+    MemorySize = 0;
+    TotalSize = ASize;
+    return Plugin::success();
+  }
+
+  Error preAllocateHeurustic(uint64_t MaxMemoryAllocation, void *VAddr){
+    const size_t MAX_MEMORY_ALLOCATION = MaxMemoryAllocation;
     constexpr size_t STEP = 1024 * 1024 * 1024ULL;
     MemoryStart = nullptr;
-    for (size_t Try = MAX_MEMORY_ALLOCATION; Try > 0; Try -= STEP) {
+    for (TotalSize = MAX_MEMORY_ALLOCATION; TotalSize> 0; TotalSize -= STEP) {
       MemoryStart =
-          Device->allocate(Try, /* HstPtr */ nullptr, TARGET_ALLOC_DEFAULT);
+          Device->allocate(TotalSize, /* HstPtr */ nullptr, TARGET_ALLOC_DEFAULT);
       if (MemoryStart)
         break;
     }
 
+    INFO(OMP_INFOTYPE_PLUGIN_KERNEL, Device->getDeviceId(),
+       "Allocated %" PRIu64 " bytes at %p for replay.\n", TotalSize, MemoryStart);
+
     if (!MemoryStart)
       return Plugin::error("Allocating record/replay memory");
 
+    if ( VAddr && VAddr != MemoryStart )
+      return Plugin::error("Cannot allocate recorded address");
+
     MemoryPtr = MemoryStart;
     MemorySize = 0;
 
     return Plugin::success();
   }
 
-  void dumpDeviceMemory(StringRef Filename) {
+  Error preallocateDeviceMemory(uint64_t DeviceMemorySize, void *ReqVAddr){
+      if ( Device->supportVAManagement() )
+        return preAllocateVAMemory(DeviceMemorySize, ReqVAddr);
+
+      uint64_t DevMemSize;
+      if ( Device->getDeviceMemorySize(DevMemSize) )
+        return Plugin::error("Cannot determine Device Memory Size");
+
+      return preAllocateHeurustic(DevMemSize, ReqVAddr);
+  }
+
+  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, nullptr);
+
     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;
     OS.close();
   }
 
 public:
   bool isRecording() const { return Status == RRStatusTy::RRRecording; }
   bool isReplaying() const { return Status == RRStatusTy::RRReplaying; }
   bool isRecordingOrReplaying() const {
     return (Status != RRStatusTy::RRDeactivated);
   }
   void setStatus(RRStatusTy Status) { this->Status = Status; }
   bool isSaveOutputEnabled() const { return ReplaySaveOutput; }
 
   RecordReplayTy()
-      : Status(RRStatusTy::RRDeactivated), ReplaySaveOutput(false),
-        DeviceMemorySize(-1) {}
+      : Status(RRStatusTy::RRDeactivated), ReplaySaveOutput(false) {}
 
   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 dumpGlobals(StringRef Filename, DeviceImageTy &Image) {
+  void dumpGlobals(StringRef Filename, 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;
for (auto &OffloadEntry : Image.getOffloadEntryTable()) {
       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, nullptr))
+        if (auto Err =
+                Device->dataRetrieve(BufferPtr, OffloadEntry.addr,
           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) {
+                           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();
+    JsonKernelInfo["BumpAllocVAStart"] = (intptr_t) MemoryStart;
 
     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)
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));
+    DP("Memory Allocator return %p \n", Alloc);
     return Alloc;
   }
 
-  Error init(GenericDeviceTy *Device, uint64_t MemSize, RRStatusTy Status,
+  Error init(GenericDeviceTy *Device, uint64_t MemSize, void *VAddr, RRStatusTy Status,
              bool SaveOutput) {
     this->Device = Device;
     this->Status = Status;
-    this->DeviceMemorySize = MemSize;
     this->ReplaySaveOutput = SaveOutput;
 
-    if (auto Err = preallocateDeviceMemory(MemSize))
+    if (auto Err = preallocateDeviceMemory(MemSize, VAddr))
       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,
+         MemoryStart, TotalSize,
          Status == RRStatusTy::RRRecording ? "Recording" : "Replaying");
 
     return Plugin::success();
   }
 
-  void deinit() { Device->free(MemoryStart); }
+  void deinit() {
+    if ( Device->supportVAManagement()  ){
+      if ( auto Err = Device->memoryVAUnMap(MemoryStart, TotalSize) ){
+          report_fatal_error("Error on releasing virtual memory space");
+      }
+    }
+    else {
+      Device->free(MemoryStart);
+    }
+  }
 
 } RecordReplay;
 
 AsyncInfoWrapperTy::AsyncInfoWrapperTy(GenericDeviceTy &Device,
                                        __tgt_async_info *AsyncInfoPtr)
     : Device(Device),
       AsyncInfoPtr(AsyncInfoPtr ? AsyncInfoPtr : &LocalAsyncInfo) {}
 

 
 Error GenericDeviceTy::queryAsync(__tgt_async_info *AsyncInfo) {
   if (!AsyncInfo || !AsyncInfo->Queue)
     return Plugin::error("Invalid async info queue");
 
   return queryAsyncImpl(*AsyncInfo);
 }
 
+
+Error GenericDeviceTy::memoryVAMap(void **Addr, void *VAddr, size_t *RSize){
+  return Plugin::error("Device does not suppport VA Management");
+}
+
+Error GenericDeviceTy::memoryVAUnMap(void *VAddr, size_t Size){
+  return Plugin::error("Device does not suppport VA Management");
+}
+
+Error GenericDeviceTy::getDeviceMemorySize(uint64_t &DSize) {
+  return Plugin::error("Mising getDeviceMemorySize impelmentation (required by RR-heuristic");
+}
+
 Expected<void *> GenericDeviceTy::dataAlloc(int64_t Size, void *HostPtr,
                                             TargetAllocTy Kind) {
   void *Alloc = nullptr;
 
   if (RecordReplay.isRecordingOrReplaying())
     return RecordReplay.alloc(Size);
 
   switch (Kind) {
if (RecordReplay.isRecording())
     RecordReplay.saveKernelInputInfo(
         GenericKernel.getName(), GenericKernel.getImage(), ArgPtrs, ArgOffsets,
         KernelArgs.NumArgs, KernelArgs.NumTeams[0], KernelArgs.ThreadLimit[0],
         KernelArgs.Tripcount);
 
   if (RecordReplay.isRecording())
     RecordReplay.saveImage(GenericKernel.getName(), GenericKernel.getImage());
 
-  auto Err = GenericKernel.launch(*this, ArgPtrs, ArgOffsets, KernelArgs,
-                                  AsyncInfoWrapper);
+  auto Err = GenericKernel.launch(*this, ArgPtrs, ArgOffsets, KernelArgs, AsyncInfoWrapper);
 
   // 'finalize' here to guarantee next record-replay actions are in-sync
   AsyncInfoWrapper.finalize(Err);
 
   if (RecordReplay.isRecordingOrReplaying() &&
       RecordReplay.isSaveOutputEnabled())
     RecordReplay.saveKernelOutputInfo(GenericKernel.getName());
 

 }
 
 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,
+                                           int64_t MemorySize,
+                                           void *VAddr,
+                                           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)) {
+  if (auto Err = RecordReplay.init(&Device, MemorySize, VAddr, 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;
     }

openmp/libomptarget/plugins-nextgen/cuda/dynamic_cuda/cuda.h

 typedef int CUdevice;
 typedef uintptr_t CUdeviceptr;
 typedef struct CUmod_st *CUmodule;
 typedef struct CUctx_st *CUcontext;
 typedef struct CUfunc_st *CUfunction;
 typedef struct CUstream_st *CUstream;
 typedef struct CUevent_st *CUevent;
 
+typedef unsigned long long CUmemGenericAllocationHandle_v1;
+typedef CUmemGenericAllocationHandle_v1 CUmemGenericAllocationHandle;
+
 #define CU_DEVICE_INVALID ((CUdevice)-2)
 
+typedef enum CUmemAllocationGranularity_flags_enum {
+    CU_MEM_ALLOC_GRANULARITY_MINIMUM     = 0x0,     /**< Minimum required granularity for allocation */
+    CU_MEM_ALLOC_GRANULARITY_RECOMMENDED = 0x1      /**< Recommended granularity for allocation for best performance */
+} CUmemAllocationGranularity_flags;
+
+typedef enum CUmemAccess_flags_enum {
+    CU_MEM_ACCESS_FLAGS_PROT_NONE        = 0x0,  /**< Default, make the address range not accessible */
+    CU_MEM_ACCESS_FLAGS_PROT_READ        = 0x1,  /**< Make the address range read accessible */
+    CU_MEM_ACCESS_FLAGS_PROT_READWRITE   = 0x3,  /**< Make the address range read-write accessible */
+    CU_MEM_ACCESS_FLAGS_PROT_MAX         = 0x7FFFFFFF
+} CUmemAccess_flags;
+
+typedef enum CUmemLocationType_enum {
+    CU_MEM_LOCATION_TYPE_INVALID = 0x0,
+    CU_MEM_LOCATION_TYPE_DEVICE  = 0x1,  /**< Location is a device location, thus id is a device ordinal */
+    CU_MEM_LOCATION_TYPE_MAX     = 0x7FFFFFFF
+} CUmemLocationType;
+
+typedef struct CUmemLocation_st {
+    CUmemLocationType type; /**< Specifies the location type, which modifies the meaning of id. */
+    int id;                 /**< identifier for a given this location's ::CUmemLocationType. */
+} CUmemLocation_v1;
+typedef CUmemLocation_v1 CUmemLocation;
+
+typedef struct CUmemAccessDesc_st {
+    CUmemLocation location;        /**< Location on which the request is to change it's accessibility */
+    CUmemAccess_flags flags;       /**< ::CUmemProt accessibility flags to set on the request */
+} CUmemAccessDesc_v1;
+
+typedef CUmemAccessDesc_v1 CUmemAccessDesc;
+
+typedef enum CUmemAllocationType_enum {
+    CU_MEM_ALLOCATION_TYPE_INVALID = 0x0,
+
+    /** This allocation type is 'pinned', i.e. cannot migrate from its current
+      * location while the application is actively using it
+      */
+    CU_MEM_ALLOCATION_TYPE_PINNED  = 0x1,
+    CU_MEM_ALLOCATION_TYPE_MAX     = 0x7FFFFFFF
+} CUmemAllocationType;
+
+typedef enum CUmemAllocationHandleType_enum {
+    CU_MEM_HANDLE_TYPE_NONE                  = 0x0,  /**< Does not allow any export mechanism. > */
+    CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR = 0x1,  /**< Allows a file descriptor to be used for exporting. Permitted only on POSIX systems. (int) */
+    CU_MEM_HANDLE_TYPE_WIN32                 = 0x2,  /**< Allows a Win32 NT handle to be used for exporting. (HANDLE) */
+    CU_MEM_HANDLE_TYPE_WIN32_KMT             = 0x4,  /**< Allows a Win32 KMT handle to be used for exporting. (D3DKMT_HANDLE) */
+    CU_MEM_HANDLE_TYPE_MAX                   = 0x7FFFFFFF
+} CUmemAllocationHandleType;
+
+typedef struct CUmemAllocationProp_st {
+    /** Allocation type */
+    CUmemAllocationType type;
+    /** requested ::CUmemAllocationHandleType */
+    CUmemAllocationHandleType requestedHandleTypes;
+    /** Location of allocation */
+    CUmemLocation location;
+    /**
+     * Windows-specific POBJECT_ATTRIBUTES required when
+     * ::CU_MEM_HANDLE_TYPE_WIN32 is specified.  This object atributes structure
+     * includes security attributes that define
+     * the scope of which exported allocations may be tranferred to other
+     * processes.  In all other cases, this field is required to be zero.
+     */
+    void *win32HandleMetaData;
+    struct {
+         /**
+         * Allocation hint for requesting compressible memory.
+         * On devices that support Compute Data Compression, compressible
+         * memory can be used to accelerate accesses to data with unstructured
+         * sparsity and other compressible data patterns. Applications are
+         * expected to query allocation property of the handle obtained with
+         * ::cuMemCreate using ::cuMemGetAllocationPropertiesFromHandle to
+         * validate if the obtained allocation is compressible or not. Note that
+         * compressed memory may not be mappable on all devices.
+         */
+         unsigned char compressionType;
+         unsigned char gpuDirectRDMACapable;
+         /** Bitmask indicating intended usage for this allocation */
+         unsigned short usage;
+         unsigned char reserved[4];
+    } allocFlags;
+} CUmemAllocationProp_v1;
+typedef CUmemAllocationProp_v1 CUmemAllocationProp;
+
 typedef enum cudaError_enum {
   CUDA_SUCCESS = 0,
   CUDA_ERROR_INVALID_VALUE = 1,
   CUDA_ERROR_NO_DEVICE = 100,
   CUDA_ERROR_INVALID_HANDLE = 400,
   CUDA_ERROR_NOT_READY = 600,
   CUDA_ERROR_TOO_MANY_PEERS = 711,
 } CUresult;

 CUresult cuCtxSetLimit(CUlimit, size_t);
 
 CUresult cuEventCreate(CUevent *, unsigned int);
 CUresult cuEventRecord(CUevent, CUstream);
 CUresult cuStreamWaitEvent(CUstream, CUevent, unsigned int);
 CUresult cuEventSynchronize(CUevent);
 CUresult cuEventDestroy(CUevent);
 
+CUresult cuMemUnmap(CUdeviceptr ptr, size_t size);
+CUresult cuMemRelease(CUmemGenericAllocationHandle handle);
+CUresult cuMemAddressFree(CUdeviceptr ptr, size_t size);
+CUresult cuMemGetInfo(size_t *free, size_t *total);
+CUresult cuMemAddressReserve(CUdeviceptr *ptr, size_t size, size_t alignment, CUdeviceptr addr, unsigned long long flags);
+CUresult cuMemMap(CUdeviceptr ptr, size_t size, size_t offset, CUmemGenericAllocationHandle handle, unsigned long long flags);
+CUresult cuMemCreate(CUmemGenericAllocationHandle *handle, size_t size, const CUmemAllocationProp *prop, unsigned long long flags);
+CUresult cuMemSetAccess(CUdeviceptr ptr, size_t size, const CUmemAccessDesc *desc, size_t count);
+CUresult cuMemGetAllocationGranularity(size_t *granularity, const CUmemAllocationProp *prop, CUmemAllocationGranularity_flags option);
+
 #endif

openmp/libomptarget/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp

 DLWRAP(cuCtxSetLimit, 2)
 
 DLWRAP(cuEventCreate, 2)
 DLWRAP(cuEventRecord, 2)
 DLWRAP(cuStreamWaitEvent, 3)
 DLWRAP(cuEventSynchronize, 1)
 DLWRAP(cuEventDestroy, 1)
 
+DLWRAP(cuMemUnmap, 2);
+DLWRAP(cuMemRelease, 1);
+DLWRAP(cuMemAddressFree, 2);
+DLWRAP(cuMemGetInfo, 2);
+DLWRAP(cuMemAddressReserve, 5);
+DLWRAP(cuMemMap, 5);
+DLWRAP(cuMemCreate, 4);
+DLWRAP(cuMemSetAccess, 4);
+DLWRAP(cuMemGetAllocationGranularity, 3);
+
 DLWRAP_FINALIZE()
 
 #ifndef DYNAMIC_CUDA_PATH
 #define DYNAMIC_CUDA_PATH "libcuda.so"
 #endif
 
 #ifndef TARGET_NAME
 #define TARGET_NAME CUDA

openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp

Error queryAsyncImpl(__tgt_async_info &AsyncInfo) override {
     // sure the synchronization only works for its own tasks.
     AsyncInfo.Queue = nullptr;
     if (auto Err = CUDAStreamManager.returnResource(Stream))
       return Err;
 
     return Plugin::check(Res, "Error in cuStreamQuery: %s");
   }
 
+  /// CUDA support VA management
+  bool supportVAManagement() const override { return true; }
+
+  /// Allocates \p RSize bytes (rounded up to page size) and hints the cuda driver to map it to \p VAddr.
+  /// The obtained address is stored in \p Addr. At return \p RSize contains the actual size
+  Error memoryVAMap(void **Addr, void *VAddr, size_t *RSize) override {
+    CUdeviceptr DVAddr = reinterpret_cast<CUdeviceptr>(VAddr);
+    auto IHandle = DeviceMMaps.find(DVAddr);
+    size_t Size = *RSize;
+
+    if ( Size == 0 )
+      return Plugin::error("Memory Map Size must be larger than 0");
+
+    // Check if we have already mapped this address
+    if ( IHandle != DeviceMMaps.end() )
+      return Plugin::error("Address already memory mapped");
+
+    CUmemAllocationProp Prop = {};
+    size_t Granularity = 0;
+
+    size_t Free, Total;
+    CUresult Res = cuMemGetInfo(&Free, &Total);
+    if (auto Err = Plugin::check(Res, "Error in cuMemGetInfo: %s"))
+        return std::move(Err);
+
+    if ( Size >= Free ){
+        *Addr = nullptr;
+        return Plugin::error("Canot map memory size larger than the available device memory");
+    }
+
+    // currently NVidia only supports pinned device types
+    Prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+    Prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+
+    Prop.location.id = DeviceId;
+    cuMemGetAllocationGranularity(&Granularity, &Prop,
+                                  CU_MEM_ALLOC_GRANULARITY_MINIMUM);
+    if (auto Err = Plugin::check(Res, "Error in cuMemGetAllocationGranularity: %s"))
+        return std::move(Err);
+
+    if ( Granularity == 0 )
+        return Plugin::error("Wrong device Page size");
+
+    // Ceil to page size.
+    Size = (Size+Granularity-1)/Granularity * Granularity;
+
+    // Create a handler of our allocation
+    CUmemGenericAllocationHandle AHandle;
+    Res = cuMemCreate(&AHandle, Size, &Prop, 0);
+    if (auto Err = Plugin::check(Res, "Error in cuMemCreate: %s"))
+        return std::move(Err);
+
+    CUdeviceptr DevPtr = 0;
+    Res = cuMemAddressReserve(&DevPtr, Size, 0, DVAddr, 0);
+    if (auto Err = Plugin::check(Res, "Error in cuMemAddressReserve: %s"))
+        return std::move(Err);
+
+    Res = cuMemMap(DevPtr, Size, 0, AHandle, 0);
+    if (auto Err = Plugin::check(Res, "Error in cuMemMap: %s"))
+        return std::move(Err);
+
+    CUmemAccessDesc ADesc = {};
+    ADesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+    ADesc.location.id = DeviceId;
+    ADesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+
+    // Sets address
+    Res = cuMemSetAccess(DevPtr, Size, &ADesc, 1);
+    if (auto Err = Plugin::check(Res, "Error in cuMemSetAccess: %s"))
+        return std::move(Err);
+
+    *Addr = reinterpret_cast<void *>(DevPtr);
+    *RSize = Size;
+    DeviceMMaps.insert( { DevPtr, AHandle } );
+    return Plugin::success();
+  }
+
+  /// De-allocates device memory and Unmaps the Virtual Addr
+  Error memoryVAUnMap(void *VAddr, size_t Size) override {
+    CUdeviceptr DVAddr = reinterpret_cast<CUdeviceptr>(VAddr);
+    auto IHandle = DeviceMMaps.find(DVAddr);
+    // Mapping does not exist
+    if ( IHandle == DeviceMMaps.end() ){
+      return Plugin::error("Addr is not MemoryMapped");
+    }
+    CUmemGenericAllocationHandle& AllocHandle = IHandle->second;
+
+    CUresult Res = cuMemUnmap(DVAddr, Size);
+    if (auto Err = Plugin::check(Res, "Error in cuMemUnmap: %s"))
+      return std::move(Err);
+
+    Res = cuMemRelease(AllocHandle);
+    if (auto Err = Plugin::check(Res, "Error in cuMemRelease: %s"))
+      return std::move(Err);
+
+    Res = cuMemAddressFree(DVAddr, Size);
+    if (auto Err = Plugin::check(Res, "Error in cuMemAddressFree: %s"))
+      return std::move(Err);
+
+    DeviceMMaps.erase(IHandle);
+    return Plugin::success();
+  }
+
   Expected<void *> dataLockImpl(void *HstPtr, int64_t Size) override {
     // TODO: Register the buffer as CUDA host memory.
     return HstPtr;
   }
 
   Error dataUnlockImpl(void *HstPtr) override { return Plugin::success(); }
 
   Expected<bool> isPinnedPtrImpl(void *HstPtr, void *&BaseHstPtr,
Error setCtxLimit(CUlimit Kind, uint64_t Value) {
     CUresult Res = cuCtxSetLimit(Kind, Value);
     return Plugin::check(Res, "Error in cuCtxSetLimit: %s");
   }
   Error getCtxLimit(CUlimit Kind, uint64_t &Value) {
     CUresult Res = cuCtxGetLimit(&Value, Kind);
     return Plugin::check(Res, "Error in cuCtxGetLimit: %s");
   }
 
+  Error getDeviceMemorySize(uint64_t &Value) override {
+    CUresult Res = cuDeviceTotalMem(&Value, Device);
+    return Plugin::check(Res , "Error in getDeviceMemorySize %s");
+  }
+
   /// CUDA-specific function to get device attributes.
   Error getDeviceAttr(uint32_t Kind, uint32_t &Value) {
     // TODO: Warn if the new value is larger than the old.
     CUresult Res =
         cuDeviceGetAttribute((int *)&Value, (CUdevice_attribute)Kind, Device);
     return Plugin::check(Res, "Error in cuDeviceGetAttribute: %s");
   }
 
private:
 
   /// The device's context. This context should be set before performing
   /// operations on the device.
   CUcontext Context = nullptr;
 
   /// The CUDA device handler.
   CUdevice Device = CU_DEVICE_INVALID;
 
+  /// The memory memory mapped addresses and their handlers
+  std::unordered_map<CUdeviceptr, CUmemGenericAllocationHandle> DeviceMMaps;
+
   /// The compute capability of the corresponding CUDA device.
   struct ComputeCapabilityTy {
     uint32_t Major;
     uint32_t Minor;
     std::string str() const {
       return "sm_" + std::to_string(Major * 10 + Minor);
     }
   } ComputeCapability;

openmp/libomptarget/src/device.cpp

if (OMPX_RecordKernel) {
     // 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,
+                                nullptr,
                                 true, OMPX_ReplaySaveOutput);
   }
 
   IsInit = true;
 }
 
 /// Thread-safe method to initialize the device only once.
 int32_t DeviceTy::initOnce() {

openmp/libomptarget/src/interface.cpp

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 VAddr The Virtual Address to assing to the initial
+///                               address of the BumpAllocator
 /// /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,
+EXTERN int __tgt_activate_record_replay(int64_t DeviceId, uint64_t MemorySize, void *VAddr,
                                         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);
+  int Rc = target_activate_rr(Device, MemorySize, VAddr, 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.

openmp/libomptarget/src/omptarget.cpp

#endif
   }
 
   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,
+int target_activate_rr(DeviceTy &Device, uint64_t MemorySize, void *VAddr, bool isRecord,
                        bool SaveOutput) {
-  return Device.RTL->activate_record_replay(Device.DeviceID, MemorySize,
+  return Device.RTL->activate_record_replay(Device.DeviceID, MemorySize, VAddr,
                                             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,

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,
+extern int target_activate_rr(DeviceTy &Device, uint64_t MemorySize, void *reqAddr,
                               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);
 

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

int main(int argc, char **argv) {
   for (auto It : *TgtArgsArray)
     TgtArgs.push_back(reinterpret_cast<void *>(It.getAsInteger().value()));
   auto *TgtArgOffsetsArray =
       JsonKernelInfo->getAsObject()->getArray("ArgOffsets");
   for (auto It : *TgtArgOffsetsArray)
     TgtArgOffsets.push_back(
         reinterpret_cast<ptrdiff_t>(It.getAsInteger().value()));
 
+  void *BAllocStart = reinterpret_cast<void*>(JsonKernelInfo->getAsObject()->getInteger("BumpAllocVAStart").value());
+
   __tgt_offload_entry KernelEntry = {nullptr, nullptr, 0, 0, 0};
   std::string KernelEntryName = KernelFunc.value().str();
   KernelEntry.name = const_cast<char *>(KernelEntryName.c_str());
   // Anything non-zero works to uniquely identify the kernel.
   KernelEntry.addr = (void *)0x1;
 
   ErrorOr<std::unique_ptr<MemoryBuffer>> ImageMB =
       MemoryBuffer::getFile(KernelEntryName + ".image", /* isText */ false,
int main(int argc, char **argv) {
   // 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_register_lib(&Desc);
 
-  int Rc = __tgt_activate_record_replay(DeviceId, DeviceMemorySize, false,
+  int Rc = __tgt_activate_record_replay(DeviceId, DeviceMemorySize, BAllocStart, false,
                                         VerifyOpt);
 
   if (Rc != OMP_TGT_SUCCESS) {
     report_fatal_error("Cannot activate record replay\n");
   }
 
   ErrorOr<std::unique_ptr<MemoryBuffer>> DeviceMemoryMB =
       MemoryBuffer::getFile(KernelEntryName + ".memory", /* isText */ false,