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

[Konstantinos Parasyris via llvm-commits]

================
@@ -0,0 +1,247 @@
+#include "PluginInterface.h"
+
+#include "Shared/APITypes.h"
+
+#include "ErrorReporting.h"
+#include "Shared/Utils.h"
+
+#include "llvm/Support/Error.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cstdint>
+#include <filesystem>
+#include <functional>
+
+using namespace llvm;
+using namespace omp;
+using namespace target;
+using namespace plugin;
+using namespace error;
+
+Error RecordReplayTy::init(uint64_t MemSize, void *VAddr) {
+  if (!VAddr)
+    VAddr = Device.getSuggestedVirtualAddress();
+
+  auto StartAddrOrErr = Device.allocateWithVirtualAddress(MemSize, VAddr);
+  if (!StartAddrOrErr)
+    return StartAddrOrErr.takeError();
+  if (!*StartAddrOrErr)
+    return Plugin::error(ErrorCode::OUT_OF_RESOURCES, "allocating memory");
+
+  StartAddr = *StartAddrOrErr;
+  TotalSize = MemSize;
+
+  INFO(OMP_INFOTYPE_PLUGIN_KERNEL, Device.getDeviceId(),
+       "Record initialized with starting address %p, "
+       "memory size %lu bytes and status %s\n",
+       StartAddr, TotalSize,
+       Status == StatusTy::Recording ? "recording" : "replaying");
+
+  return Plugin::success();
+}
+
+Error RecordReplayTy::deinit() {
+  if (StartAddr)
+    return Device.deallocateWithVirtualAddress(StartAddr, TotalSize);
+  return Plugin::success();
+}
+
+std::pair<const RecordReplayTy::InstanceTy &, bool>
+RecordReplayTy::registerInstance(StringRef KernelName, uint32_t NumTeams,
+                                 uint32_t NumThreads,
+                                 uint32_t SharedMemorySize) {
+  std::lock_guard<std::mutex> LG(InstancesLock);
+  auto [It, Inserted] =
+      Instances.emplace(KernelName, NumTeams, NumThreads, SharedMemorySize);
+  // Increase the number of occurrences.
+  It->Occurrences += 1;
+  return {*It, Inserted};
+}
+
+void *RecordReplayTy::allocate(uint64_t Size) {
+  assert(StartAddr && "Expected memory has been pre-allocated");
+  constexpr int Alignment = 16;
+  // Assume alignment is a power of 2.
+  int64_t AlignedSize = (Size + (Alignment - 1)) & (~(Alignment - 1));
+
+  std::lock_guard<std::mutex> LG(AllocationLock);
+  void *Alloc = (char *)StartAddr + CurrentSize;
+  CurrentSize += AlignedSize;
+  return Alloc;
+}
+
+Expected<RecordReplayTy::HandleTy> RecordReplayTy::recordPrologue(
+    const GenericKernelTy &Kernel, const KernelArgsTy &KernelArgs,
+    const KernelLaunchParamsTy &LaunchParams, uint32_t NumTeams[3],
+    uint32_t NumThreads[3], uint32_t SharedMemorySize) {
+  if (!isRecordingOrReplaying())
+    return HandleTy{nullptr, false};
+
+  // Register the instance and avoid recording if it is inactive or replaying.
+  auto [Instance, First] = registerInstance(Kernel.getName(), NumTeams[0],
+                                            NumThreads[0], SharedMemorySize);
+
+  HandleTy Handle{&Instance, First};
+  if (isReplaying() || !First)
+    return Handle;
+
+  if (auto Err = recordDescImpl(Kernel, Instance, KernelArgs, LaunchParams))
+    return Err;
+
+  if (auto Err = recordPrologueImpl(Kernel, Instance, KernelArgs, LaunchParams))
+    return Err;
+
+  return Handle;
+}
+
+Error RecordReplayTy::recordEpilogue(const GenericKernelTy &Kernel,
+                                     HandleTy Handle) {
+  if (!shouldRecordEpilogue() || !Handle.Active)
+    return Plugin::success();
+
+  return recordEpilogueImpl(Kernel, *Handle.Instance);
+}
+
+Error NativeRecordReplayTy::recordPrologueImpl(
+    const GenericKernelTy &Kernel, const InstanceTy &Instance,
+    const KernelArgsTy &KernelArgs, const KernelLaunchParamsTy &LaunchParams) {
+  SmallString<128> SnapshotFilename = {Kernel.getName(), ".memory"};
----------------
koparasy wrote:

IIRC kernel names can have more characters than the maximum amount of characters allowed by the OS, we had opted to hash it. 

https://github.com/llvm/llvm-project/pull/190588