[offload] Fix kernel record/replay and add extensible mechanism (llvm#190588)

This commit fixes the kernel record replay on both AMD and CUDA devices. It
also re-organizes the record replay code, moves the whole code to separate
files, and makes it extensible to support other record formats (potentially in
the future). The environment variables for controlling the recording have also
been modified.

Author: kevinsala

offload/cmake/OpenMPTesting.cmake

@@ -12,6 +12,7 @@ endif()
 set(OFFLOAD_NOT_EXECUTABLE ${LLVM_TOOLS_BINARY_DIR}/not)
 set(OFFLOAD_DEVICE_INFO_EXECUTABLE ${LLVM_RUNTIME_OUTPUT_INTDIR}/llvm-offload-device-info)
 set(OFFLOAD_TBLGEN_EXECUTABLE ${LLVM_RUNTIME_OUTPUT_INTDIR}/offload-tblgen)
+set(OMP_KERNEL_REPLAY ${LLVM_RUNTIME_OUTPUT_INTDIR}/llvm-omp-kernel-replay)
 
 # Set the information that we know.
 set(OPENMP_TEST_COMPILER_ID "Clang")

offload/include/omptarget.h

@@ -428,17 +428,20 @@ void __tgt_target_nowait_query(void **AsyncHandle);
 /// 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);
+                               const llvm::offloading::EntryTy *Globals,
+                               int32_t NumGlobals, void **TgtArgs,
+                               ptrdiff_t *TgtOffsets, int32_t NumArgs,
+                               int32_t NumTeams, int32_t ThreadLimit,
+                               uint32_t SharedMemorySize,
+                               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,
                                  void *VAddr, bool IsRecord, bool SaveOutput,
-                                 uint64_t &ReqPtrArgOffset);
+                                 bool EmitReport, const char *OutputDirPath);
 
 // Registers a callback for the RPC server. Expects this function type.
 // unsigned callback(rpc::Server::Port *Port, unsigned NumLanes). See the RPC

offload/libomptarget/device.cpp

@@ -89,12 +89,23 @@ llvm::Error DeviceTy::init() {
   // Enables recording kernels if set.
   BoolEnvar OMPX_RecordKernel("LIBOMPTARGET_RECORD", false);
   if (OMPX_RecordKernel) {
-    // Enables saving the device memory kernel output post execution if set.
-    BoolEnvar OMPX_ReplaySaveOutput("LIBOMPTARGET_RR_SAVE_OUTPUT", false);
-
-    uint64_t ReqPtrArgOffset;
-    RTL->initialize_record_replay(RTLDeviceID, 0, nullptr, true,
-                                  OMPX_ReplaySaveOutput, ReqPtrArgOffset);
+    BoolEnvar OMPX_RecordOutput("LIBOMPTARGET_RECORD_OUTPUT", true);
+    Int64Envar OMPX_RecordMemSize("LIBOMPTARGET_RECORD_MEMSIZE",
+                                  8 * 1024 * 1024 * 1024ULL);
+    Int32Envar OMPX_RecordDevice("LIBOMPTARGET_RECORD_DEVICE", 0);
+    StringEnvar OMPX_RecordOutputDir("LIBOMPTARGET_RECORD_DIR", "");
+    BoolEnvar OMPX_EmitRecordReport("LIBOMPTARGET_RECORD_REPORT", false);
+    if (OMPX_RecordDevice != RTLDeviceID)
+      return llvm::Error::success();
+
+    Ret = RTL->initialize_record_replay(
+        RTLDeviceID, OMPX_RecordMemSize, nullptr,
+        /*IsRecord=*/true, /*IsNative=*/true, OMPX_RecordOutput,
+        OMPX_EmitRecordReport, OMPX_RecordOutputDir.get().c_str());
+    if (Ret != OFFLOAD_SUCCESS)
+      return error::createOffloadError(error::ErrorCode::BACKEND_FAILURE,
+                                       "failed to initialize RR in device %d\n",
+                                       DeviceID);
   }
 
   return llvm::Error::success();

offload/libomptarget/interface.cpp

@@ -473,25 +473,32 @@ EXTERN int __tgt_target_kernel(ident_t *Loc, int64_t DeviceId, int32_t NumTeams,
 /// 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
+///                   by the record replay allocator.
 /// /param IsRecord Activates the record replay mechanism in
-///                 'record' mode or 'replay' mode.
+///                 'record' or 'replay' mode.
 /// /param SaveOutput Store the device memory after kernel
-///                   execution on persistent storage
+///                   execution on persistent storage.
+/// /param EmitReport Emit a summary report after the recording.
+/// /param OutputDirPath The output directory where the record replay files
+/// should be stored. An empty string or nullptr indicates the current working
+/// directory should be used.
 EXTERN int __tgt_activate_record_replay(int64_t DeviceId, uint64_t MemorySize,
                                         void *VAddr, bool IsRecord,
-                                        bool SaveOutput,
-                                        uint64_t &ReqPtrArgOffset) {
+                                        bool SaveOutput, bool EmitReport,
+                                        const char *OutputDirPath) {
   assert(PM && "Runtime not initialized");
   OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   auto DeviceOrErr = PM->getDevice(DeviceId);
   if (!DeviceOrErr)
     FATAL_MESSAGE(DeviceId, "%s", toString(DeviceOrErr.takeError()).c_str());
 
-  [[maybe_unused]] int Rc = target_activate_rr(
-      *DeviceOrErr, MemorySize, VAddr, IsRecord, SaveOutput, ReqPtrArgOffset);
-  assert(Rc == OFFLOAD_SUCCESS &&
-         "__tgt_activate_record_replay unexpected failure!");
+  int Rc = target_activate_rr(*DeviceOrErr, MemorySize, VAddr, IsRecord,
+                              SaveOutput, EmitReport, OutputDirPath);
+  if (Rc != OFFLOAD_SUCCESS) {
+    ODBG(ODT_Interface) << "Record replay failed to activate in device "
+                        << DeviceId;
+    return OMP_TGT_FAIL;
+  }
   return OMP_TGT_SUCCESS;
 }
 
@@ -512,12 +519,12 @@ EXTERN int __tgt_activate_record_replay(int64_t DeviceId, uint64_t MemorySize,
 ///                    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) {
+EXTERN int __tgt_target_kernel_replay(
+    ident_t *Loc, int64_t DeviceId, void *HostPtr, void *DeviceMemory,
+    int64_t DeviceMemorySize, const llvm::offloading::EntryTy *Globals,
+    int32_t NumGlobals, void **TgtArgs, ptrdiff_t *TgtOffsets, int32_t NumArgs,
+    int32_t NumTeams, int32_t ThreadLimit, uint32_t SharedMemorySize,
+    uint64_t LoopTripCount) {
   assert(PM && "Runtime not initialized");
   OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   if (checkDevice(DeviceId, Loc)) {
@@ -534,14 +541,18 @@ EXTERN int __tgt_target_kernel_replay(ident_t *Loc, int64_t DeviceId,
                     /*CodePtr=*/OMPT_GET_RETURN_ADDRESS);)
 
   AsyncInfoTy AsyncInfo(*DeviceOrErr);
-  int Rc = target_replay(Loc, *DeviceOrErr, HostPtr, DeviceMemory,
-                         DeviceMemorySize, TgtArgs, TgtOffsets, NumArgs,
-                         NumTeams, ThreadLimit, LoopTripCount, AsyncInfo);
+  int Rc =
+      target_replay(Loc, *DeviceOrErr, HostPtr, DeviceMemory, DeviceMemorySize,
+                    Globals, NumGlobals, TgtArgs, TgtOffsets, NumArgs, NumTeams,
+                    ThreadLimit, SharedMemorySize, LoopTripCount, AsyncInfo);
+
   if (Rc == OFFLOAD_SUCCESS)
     Rc = AsyncInfo.synchronize();
-  handleTargetOutcome(Rc == OFFLOAD_SUCCESS, Loc);
-  assert(Rc == OFFLOAD_SUCCESS &&
-         "__tgt_target_kernel_replay unexpected failure!");
+
+  if (Rc != OFFLOAD_SUCCESS) {
+    ODBG(ODT_Interface) << "Kernel replay failed in device " << DeviceId;
+    return OMP_TGT_FAIL;
+  }
   return OMP_TGT_SUCCESS;
 }
 

offload/libomptarget/omptarget.cpp

@@ -2379,67 +2379,105 @@ int target(ident_t *Loc, DeviceTy &Device, void *HostPtr,
 /// and informing the record-replayer of whether to store the output
 /// in some file.
 int target_activate_rr(DeviceTy &Device, uint64_t MemorySize, void *VAddr,
-                       bool IsRecord, bool SaveOutput,
-                       uint64_t &ReqPtrArgOffset) {
-  return Device.RTL->initialize_record_replay(Device.DeviceID, MemorySize,
-                                              VAddr, IsRecord, SaveOutput,
-                                              ReqPtrArgOffset);
+                       bool IsRecord, bool SaveOutput, bool EmitReport,
+                       const char *OutputDirPath) {
+  return Device.RTL->initialize_record_replay(
+      Device.DeviceID, MemorySize, VAddr, IsRecord,
+      /*IsNative=*/true, SaveOutput, EmitReport, OutputDirPath);
 }
 
 /// 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,
+                  void *DeviceMemory, int64_t DeviceMemorySize,
+                  const llvm::offloading::EntryTy *Globals, int32_t NumGlobals,
+                  void **TgtArgs, ptrdiff_t *TgtOffsets, int32_t NumArgs,
+                  int32_t NumTeams, int32_t ThreadLimit,
+                  uint32_t SharedMemorySize, uint64_t LoopTripCount,
                   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 " << HostPtr
-             << " does not have a matching target pointer.";
-    return OFFLOAD_FAIL;
+  int32_t NumSymbols = NumGlobals + 1;
+
+  struct SymbolDataTy {
+    void *DevPtr = nullptr;
+    TableMap *TM = nullptr;
+    __tgt_target_table *TargetTable = nullptr;
+  };
+  SmallVector<SymbolDataTy> Symbols(NumSymbols);
+
+  for (int32_t I = 0; I < NumSymbols; ++I) {
+    // The first symbol is the kernel entry.
+    void *SymbolHostPtr = (I == 0) ? HostPtr : Globals[I - 1].Address;
+
+    // Get the table map for each symbol.
+    Symbols[I].TM = getTableMap(SymbolHostPtr);
+    if (!Symbols[I].TM) {
+      REPORT() << "Host pointer " << SymbolHostPtr
+               << " does not have a matching target pointer.";
+      return OFFLOAD_FAIL;
+    }
   }
 
-  // Retrieve the target table of offloading entries.
-  __tgt_target_table *TargetTable = nullptr;
+  // Retrieve the target table for each symbol.
   {
     std::lock_guard<std::mutex> TrlTblLock(PM->TrlTblMtx);
     assert(TM->Table->TargetsTable.size() > (size_t)DeviceId &&
            "Not expecting a device ID outside the table's bounds!");
-    TargetTable = TM->Table->TargetsTable[DeviceId];
+    for (auto &S : Symbols) {
+      S.TargetTable = S.TM->Table->TargetsTable[DeviceId];
+      assert(S.TargetTable && "Global data has not been mapped\n");
+    }
   }
-  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].Address;
-  ODBG(ODT_Kernel) << "Launching target execution "
-                   << TargetTable->EntriesBegin[TM->Index].SymbolName
-                   << " with pointer " << TgtEntryPtr << " (index=" << TM->Index
-                   << ").";
+  // Retrieve the device pointers for each symbol.
+  for (auto &S : Symbols)
+    S.DevPtr = S.TargetTable->EntriesBegin[S.TM->Index].Address;
+
+  // Initialize the device memory of each global.
+  for (int32_t I = 0; I < NumGlobals; ++I) {
+    assert(Globals[I].AuxAddr && "Global has no AuxAddr.");
+
+    // Initialize the value of the global in the device.
+    int Ret = Device.submitData(Symbols[I + 1].DevPtr, Globals[I].AuxAddr,
+                                Globals[I].Size, AsyncInfo);
+    if (Ret != OFFLOAD_SUCCESS) {
+      REPORT() << "Failed to submit data to a global.";
+      return OFFLOAD_FAIL;
+    }
+  }
 
   void *TgtPtr = Device.allocData(DeviceMemorySize, /*HstPtr=*/nullptr,
                                   TARGET_ALLOC_DEFAULT);
-  Device.submitData(TgtPtr, DeviceMemory, DeviceMemorySize, AsyncInfo);
+  if (!TgtPtr) {
+    REPORT() << "Failed to allocate device memory.";
+    return OFFLOAD_FAIL;
+  }
+
+  int Ret =
+      Device.submitData(TgtPtr, DeviceMemory, DeviceMemorySize, AsyncInfo);
+  if (Ret != OFFLOAD_SUCCESS) {
+    REPORT() << "Failed to submit data to a global.";
+    return OFFLOAD_FAIL;
+  }
 
   KernelArgsTy KernelArgs{};
   KernelArgs.Version = OMP_KERNEL_ARG_VERSION;
   KernelArgs.NumArgs = NumArgs;
   KernelArgs.Tripcount = LoopTripCount;
   KernelArgs.NumTeams[0] = NumTeams;
+  KernelArgs.NumTeams[1] = 1;
+  KernelArgs.NumTeams[2] = 1;
   KernelArgs.ThreadLimit[0] = ThreadLimit;
+  KernelArgs.ThreadLimit[1] = 1;
+  KernelArgs.ThreadLimit[2] = 1;
+  KernelArgs.DynCGroupMem = SharedMemorySize;
 
-  int Ret = Device.launchKernel(TgtEntryPtr, TgtArgs, TgtOffsets, KernelArgs,
-                                AsyncInfo);
-
+  Ret = Device.launchKernel(Symbols[0].DevPtr, TgtArgs, TgtOffsets, KernelArgs,
+                            AsyncInfo);
   if (Ret != OFFLOAD_SUCCESS) {
-    REPORT() << "Executing target region abort target.";
+    REPORT() << "Failed to launch kernel replay.";
     return OFFLOAD_FAIL;
   }
-
   return OFFLOAD_SUCCESS;
 }

offload/libomptarget/private.h

@@ -27,14 +27,17 @@ extern int target(ident_t *Loc, DeviceTy &Device, void *HostPtr,
                   KernelArgsTy &KernelArgs, AsyncInfoTy &AsyncInfo);
 
 extern int target_activate_rr(DeviceTy &Device, uint64_t MemorySize,
-                              void *ReqAddr, bool isRecord, bool SaveOutput,
-                              uint64_t &ReqPtrArgOffset);
+                              void *ReqAddr, bool IsRecord, bool SaveOutput,
+                              bool EmitReport, const char *OutputDirPath);
 
 extern int target_replay(ident_t *Loc, DeviceTy &Device, void *HostPtr,
                          void *DeviceMemory, int64_t DeviceMemorySize,
-                         void **TgtArgs, ptrdiff_t *TgtOffsets, int32_t NumArgs,
+                         const llvm::offloading::EntryTy *Globals,
+                         int32_t NumGlobals, void **TgtArgs,
+                         ptrdiff_t *TgtOffsets, int32_t NumArgs,
                          int32_t NumTeams, int32_t ThreadLimit,
-                         uint64_t LoopTripCount, AsyncInfoTy &AsyncInfo);
+                         uint32_t SharedMemorySize, uint64_t LoopTripCount,
+                         AsyncInfoTy &AsyncInfo);
 
 extern void handleTargetOutcome(bool Success, ident_t *Loc);
 

offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp

@@ -75,6 +75,13 @@ DLWRAP(hsa_amd_profiling_set_profiler_enabled, 2)
 DLWRAP(hsa_code_object_reader_create_from_memory, 3)
 DLWRAP(hsa_code_object_reader_destroy, 1)
 DLWRAP(hsa_executable_load_agent_code_object, 5)
+DLWRAP(hsa_amd_vmem_address_reserve, 4)
+DLWRAP(hsa_amd_vmem_address_free, 2)
+DLWRAP(hsa_amd_vmem_handle_create, 5)
+DLWRAP(hsa_amd_vmem_handle_release, 1)
+DLWRAP(hsa_amd_vmem_map, 5)
+DLWRAP(hsa_amd_vmem_unmap, 2)
+DLWRAP(hsa_amd_vmem_set_access, 4)
 
 DLWRAP_FINALIZE()
 

offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h

@@ -116,6 +116,13 @@ typedef struct hsa_isa_s {
   uint64_t handle;
 } hsa_isa_t;
 
+typedef enum {
+  HSA_ACCESS_PERMISSION_NONE = 0,
+  HSA_ACCESS_PERMISSION_RO = 1,
+  HSA_ACCESS_PERMISSION_WO = 2,
+  HSA_ACCESS_PERMISSION_RW = 3
+} hsa_access_permission_t;
+
 hsa_status_t hsa_system_get_info(hsa_system_info_t attribute, void *value);
 
 hsa_status_t hsa_agent_get_info(hsa_agent_t agent, hsa_agent_info_t attribute,

offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h

@@ -163,6 +163,20 @@ typedef struct hsa_amd_pointer_info_s {
   size_t sizeInBytes;
 } hsa_amd_pointer_info_t;
 
+typedef enum {
+  MEMORY_TYPE_NONE,
+  MEMORY_TYPE_PINNED,
+} hsa_amd_memory_type_t;
+
+typedef struct hsa_amd_vmem_alloc_handle_s {
+  uint64_t handle;
+} hsa_amd_vmem_alloc_handle_t;
+
+typedef struct hsa_amd_memory_access_desc_s {
+  hsa_access_permission_t permissions;
+  hsa_agent_t agent_handle;
+} hsa_amd_memory_access_desc_t;
+
 hsa_status_t hsa_amd_pointer_info(const void* ptr,
                                           hsa_amd_pointer_info_t* info,
                                           void* (*alloc)(size_t),
@@ -181,6 +195,29 @@ hsa_amd_profiling_get_dispatch_time(hsa_agent_t agent, hsa_signal_t signal,
 hsa_status_t hsa_amd_profiling_set_profiler_enabled(hsa_queue_t *queue,
                                                     int enable);
 
+hsa_status_t hsa_amd_vmem_address_reserve(void **va, size_t size,
+                                          uint64_t address, uint64_t flags);
+
+hsa_status_t hsa_amd_vmem_address_free(void *va, size_t size);
+
+hsa_status_t
+hsa_amd_vmem_handle_create(hsa_amd_memory_pool_t pool, size_t size,
+                           hsa_amd_memory_type_t type, uint64_t flags,
+                           hsa_amd_vmem_alloc_handle_t *memory_handle);
+
+hsa_status_t
+hsa_amd_vmem_handle_release(hsa_amd_vmem_alloc_handle_t memory_handle);
+
+hsa_status_t hsa_amd_vmem_map(void *va, size_t size, size_t in_offset,
+                              hsa_amd_vmem_alloc_handle_t memory_handle,
+                              uint64_t flags);
+
+hsa_status_t hsa_amd_vmem_unmap(void *va, size_t size);
+
+hsa_status_t hsa_amd_vmem_set_access(void *va, size_t size,
+                                     const hsa_amd_memory_access_desc_t *desc,
+                                     size_t desc_cnt);
+
 #ifdef __cplusplus
 }
 #endif