//===- 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 "Shared/APITypes.h"
#include "Shared/Debug.h"
#include "Shared/Environment.h"

#include "ErrorReporting.h"
#include "GlobalHandler.h"
#include "JIT.h"
#include "Shared/Utils.h"
#include "Utils/ELF.h"
#include "omptarget.h"

#include "llvm/Bitcode/BitcodeReader.h"
#include "llvm/Frontend/OpenMP/OMPConstants.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Signals.h"
#include "llvm/Support/raw_ostream.h"

#include <cstdint>

using namespace llvm;
using namespace omp;
using namespace target;
using namespace plugin;
using namespace error;
using namespace llvm::offload::debug;

AsyncInfoWrapperTy::AsyncInfoWrapperTy(GenericDeviceTy &Device,
                                       __tgt_async_info *AsyncInfoPtr)
    : Device(Device),
      AsyncInfoPtr(AsyncInfoPtr ? AsyncInfoPtr : &LocalAsyncInfo) {}

Error AsyncInfoWrapperTy::synchronize() {
  assert(AsyncInfoPtr && "AsyncInfoWrapperTy already finalized");

  // Synchronize with the async info's operations without releasing the queue.
  return Device.synchronize(AsyncInfoPtr, /*ReleaseQueue=*/false);
}

void AsyncInfoWrapperTy::finalize(Error &Err) {
  assert(AsyncInfoPtr && "AsyncInfoWrapperTy already finalized");

  // If we used a local async info object we want synchronous behavior. In that
  // case, and assuming the current status code is correct, we will synchronize
  // explicitly when the object is deleted. Update the error with the result of
  // the synchronize operation.
  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;

  return initImpl(GenericDevice, Image);
}

Expected<KernelLaunchEnvironmentTy *>
GenericKernelTy::getKernelLaunchEnvironment(
    GenericDeviceTy &GenericDevice, const KernelLaunchArgsTy &LaunchArgs,
    const DynBlockMemConfTy &DynBlockMemConf,
    AsyncInfoWrapperTy &AsyncInfoWrapper, uint32_t NumBlocks0) const {
  // Ctor/Dtor have no arguments, replaying uses the original kernel launch
  // environment, and launches with no reserved dyn_ptr slot (e.g. older
  // compiler versions, or non-OpenMP launches) have nowhere to store one.
  if ((GenericDevice.getRecordReplay() &&
       GenericDevice.getRecordReplay()->isReplaying()) ||
      !LaunchArgs.DynPtrSlot)
    return nullptr;

  const bool NeedsReductionBuffer =
      LaunchArgs.KernelLaunchInfo.ReductionDataSize != 0;
  if (NeedsReductionBuffer && LaunchArgs.OmpABIVersion < OMP_KERNEL_ARG_VERSION)
    return Plugin::error(ErrorCode::INVALID_BINARY,
                         "kernel was built against an older OpenMP "
                         "kernel-launch-environment ABI (v%u); current "
                         "runtime requires v%u for cross-team reductions",
                         LaunchArgs.OmpABIVersion, OMP_KERNEL_ARG_VERSION);
  if (!NeedsReductionBuffer && !LaunchArgs.DynCGroupMem)
    return reinterpret_cast<KernelLaunchEnvironmentTy *>(~0);

  auto AllocOrErr = GenericDevice.dataAlloc(
      sizeof(KernelLaunchEnvironmentTy),
      /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE, /*Alignment=*/0);
  if (!AllocOrErr)
    return AllocOrErr.takeError();

  // Remember to free the memory later.
  AsyncInfoWrapper.freeAllocationAfterSynchronization(
      *AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE);

  /// Use the KLE in the __tgt_async_info to ensure a stable address for the
  /// async data transfer.
  auto &LocalKLE = (*AsyncInfoWrapper).KernelLaunchEnvironment;
  LocalKLE = KernelLaunchEnvironment;

  LocalKLE.DynCGroupMemSize = DynBlockMemConf.Size;
  LocalKLE.DynCGroupMemFbPtr = DynBlockMemConf.FallbackPtr;
  LocalKLE.DynCGroupMemFb = DynBlockMemConf.Fallback;
  LocalKLE.ReductionBuffer = nullptr;

  if (NeedsReductionBuffer) {
    // Use number of teams many buffer elements.
    auto AllocOrErr = GenericDevice.dataAlloc(
        uint64_t(LaunchArgs.KernelLaunchInfo.ReductionDataSize) * NumBlocks0,
        /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE,
        /*Alignment=*/0);
    if (!AllocOrErr)
      return AllocOrErr.takeError();
    LocalKLE.ReductionBuffer = *AllocOrErr;
    // Remember to free the memory later.
    AsyncInfoWrapper.freeAllocationAfterSynchronization(
        *AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE);
  }

  INFO(OMP_INFOTYPE_DATA_TRANSFER, GenericDevice.getDeviceId(),
       "Copying data from host to device, HstPtr=" DPxMOD ", TgtPtr=" DPxMOD
       ", Size=%" PRId64 ", Name=KernelLaunchEnv\n",
       DPxPTR(&LocalKLE), DPxPTR(*AllocOrErr),
       sizeof(KernelLaunchEnvironmentTy));

  auto Err = GenericDevice.dataSubmit(*AllocOrErr, &LocalKLE,
                                      sizeof(KernelLaunchEnvironmentTy),
                                      AsyncInfoWrapper);
  if (Err)
    return Err;
  return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr);
}

Error GenericKernelTy::printLaunchInfo(GenericDeviceTy &GenericDevice,
                                       const KernelLaunchArgsTy &LaunchArgs,
                                       uint32_t NumThreads[3],
                                       uint32_t NumBlocks[3]) const {
  return printLaunchInfoDetails(GenericDevice, LaunchArgs, NumThreads,
                                NumBlocks);
}

Error GenericKernelTy::printLaunchInfoDetails(
    GenericDeviceTy &GenericDevice, const KernelLaunchArgsTy &LaunchArgs,
    uint32_t NumThreads[3], uint32_t NumBlocks[3]) const {
  return Plugin::success();
}

Expected<DynBlockMemConfTy>
GenericKernelTy::prepareBlockMemory(GenericDeviceTy &GenericDevice,
                                    const KernelLaunchArgsTy &LaunchArgs,
                                    uint32_t NumBlocks) const {
  uint32_t MaxBlockMemSize = GenericDevice.getMaxBlockSharedMemSize();
  uint32_t DynBlockMemSize = LaunchArgs.DynCGroupMem;
  uint32_t TotalBlockMemSize = StaticBlockMemSize + DynBlockMemSize;
  uint32_t DynNativeBlockMemSize = DynBlockMemSize;
  void *DynFallbackPtr = nullptr;

  // No enough block memory to cover the static one. Cannot run the kernel.
  if (StaticBlockMemSize > MaxBlockMemSize)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "Static block memory size exceeds maximum");
  // No enough block memory to cover dynamic one, and the fallback is aborting.
  if (static_cast<DynCGroupMemFallbackType>(
          LaunchArgs.Flags.DynCGroupMemFallback) ==
          DynCGroupMemFallbackType::Abort &&
      TotalBlockMemSize > MaxBlockMemSize)
    return Plugin::error(
        ErrorCode::INVALID_ARGUMENT,
        "Requested block memory size (static + dynamic) exceeds maximum");

  DynCGroupMemFallbackType DynFallback = DynCGroupMemFallbackType::None;
  if (DynBlockMemSize && TotalBlockMemSize > MaxBlockMemSize) {
    // Launch without native dynamic block memory.
    DynNativeBlockMemSize = 0;
    DynFallback = static_cast<DynCGroupMemFallbackType>(
        LaunchArgs.Flags.DynCGroupMemFallback);
    if (DynFallback != DynCGroupMemFallbackType::DefaultMem) {
      // Do not provide any memory as fallback.
      DynBlockMemSize = 0;
    } else {
      // Get global memory as fallback.
      auto AllocOrErr = GenericDevice.dataAlloc(
          NumBlocks * DynBlockMemSize,
          /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE,
          /*Alignment=*/0);
      if (!AllocOrErr)
        return AllocOrErr.takeError();
      DynFallbackPtr = *AllocOrErr;
    }
  }
  return DynBlockMemConfTy{DynBlockMemSize, DynNativeBlockMemSize, DynFallback,
                           DynFallbackPtr};
}

Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice,
                              KernelLaunchArgsTy &LaunchArgs,
                              AsyncInfoWrapperTy &AsyncInfoWrapper) const {
  uint32_t EffectiveNumThreads[3] = {LaunchArgs.UserThreadLimit[0],
                                     LaunchArgs.UserThreadLimit[1],
                                     LaunchArgs.UserThreadLimit[2]};
  uint32_t EffectiveNumBlocks[3] = {LaunchArgs.UserNumBlocks[0],
                                    LaunchArgs.UserNumBlocks[1],
                                    LaunchArgs.UserNumBlocks[2]};

  auto DynBlockMemConfOrErr = prepareBlockMemory(
      GenericDevice, LaunchArgs,
      EffectiveNumBlocks[0] * EffectiveNumBlocks[1] * EffectiveNumBlocks[2]);
  if (!DynBlockMemConfOrErr)
    return DynBlockMemConfOrErr.takeError();

  DynBlockMemConfTy &DynBlockMemConf = *DynBlockMemConfOrErr;
  if (DynBlockMemConf.FallbackPtr)
    AsyncInfoWrapper.freeAllocationAfterSynchronization(
        DynBlockMemConf.FallbackPtr, TargetAllocTy::TARGET_ALLOC_DEVICE);

  auto KernelLaunchEnvOrErr =
      getKernelLaunchEnvironment(GenericDevice, LaunchArgs, DynBlockMemConf,
                                 AsyncInfoWrapper, EffectiveNumBlocks[0]);
  if (!KernelLaunchEnvOrErr)
    return KernelLaunchEnvOrErr.takeError();

  // Fill in the kernel launch environment (dyn_ptr) if this launch has a
  // reserved slot for it. When replaying, getKernelLaunchEnvironment()
  // returns null so the recorded value already in the slot is preserved.
  if (LaunchArgs.DynPtrSlot && *KernelLaunchEnvOrErr)
    *LaunchArgs.DynPtrSlot = *KernelLaunchEnvOrErr;

  if (auto Err = printLaunchInfo(GenericDevice, LaunchArgs, EffectiveNumThreads,
                                 EffectiveNumBlocks))
    return Err;

  RecordReplayTy::HandleTy RRHandle;
  RecordReplayTy *RecordReplay = GenericDevice.getRecordReplay();
  if (RecordReplay) {
    // Record replay requires synchronization of any previous operation.
    if (auto Err = AsyncInfoWrapper.synchronize())
      return Err;

    // Record the kernel prologue data before kernel launch.
    auto RRHandleOrErr = RecordReplay->recordPrologue(
        *this, LaunchArgs, EffectiveNumBlocks, EffectiveNumThreads,
        DynBlockMemConf.NativeSize);
    if (!RRHandleOrErr)
      return RRHandleOrErr.takeError();
    RRHandle = *RRHandleOrErr;
  }

  if (auto Err =
          launchImpl(GenericDevice, EffectiveNumThreads, EffectiveNumBlocks,
                     DynBlockMemConf.NativeSize, LaunchArgs, AsyncInfoWrapper))
    return Err;

  if (RecordReplay) {
    // Record replay requires synchronization.
    if (auto Err = AsyncInfoWrapper.synchronize())
      return Err;

    // Record the epilogue data after kernel synchronization.
    return RecordReplay->recordEpilogue(*this, RRHandle);
  }
  return Plugin::success();
}

GenericDeviceTy::GenericDeviceTy(GenericPluginTy &Plugin, int32_t DeviceId,
                                 int32_t NumDevices,
                                 const llvm::omp::GV &OMPGridValues)
    : Plugin(Plugin), OMP_TeamLimit("OMP_TEAM_LIMIT"),
      OMP_NumTeams("OMP_NUM_TEAMS"),
      OMP_TeamsThreadLimit("OMP_TEAMS_THREAD_LIMIT"),
      OMPX_DebugKind("LIBOMPTARGET_DEVICE_RTL_DEBUG"),
      // Do not initialize the following two envars since they depend on the
      // device initialization. These cannot be consulted until the device is
      // initialized correctly. We initialize them in GenericDeviceTy::init().
      OMPX_TargetStackSize(), OMPX_TargetHeapSize(),
      // By default, the initial number of streams and events is 1.
      OMPX_InitialNumStreams("LIBOMPTARGET_NUM_INITIAL_STREAMS", 1),
      OMPX_InitialNumEvents("LIBOMPTARGET_NUM_INITIAL_EVENTS", 1),
      DeviceId(DeviceId), GridValues(OMPGridValues),
      PeerAccesses(NumDevices, PeerAccessState::PENDING), PeerAccessesLock(),
      PinnedAllocs(*this), RPCServer(nullptr) {
  // Conservative fall-back to the plugin's device uid for the case that no real
  // vendor (u)uid will become available later.
  setDeviceUidFromVendorUid(std::to_string(static_cast<uint64_t>(DeviceId)));

  // Envar that indicates whether mapped host buffers should be locked
  // automatically. The possible values are boolean (on/off) and a special:
  //   off:       Mapped host buffers are not locked.
  //   on:        Mapped host buffers are locked in a best-effort approach.
  //              Failure to lock the buffers are silent.
  //   mandatory: Mapped host buffers are always locked and failures to lock
  //              a buffer results in a fatal error.
  StringEnvar OMPX_LockMappedBuffers("LIBOMPTARGET_LOCK_MAPPED_HOST_BUFFERS",
                                     "off");

  bool Enabled;
  if (StringParser::parse(OMPX_LockMappedBuffers.get().data(), Enabled)) {
    // Parsed as a boolean value. Enable the feature if necessary.
    LockMappedBuffers = Enabled;
    IgnoreLockMappedFailures = true;
  } else if (OMPX_LockMappedBuffers.get() == "mandatory") {
    // Enable the feature and failures are fatal.
    LockMappedBuffers = true;
    IgnoreLockMappedFailures = false;
  } else {
    // Disable by default.
    ODBG(OLDT_Alloc) << "Invalid value LIBOMPTARGET_LOCK_MAPPED_HOST_BUFFERS="
                     << OMPX_LockMappedBuffers.get();
    LockMappedBuffers = false;
  }
}

Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
  if (auto Err = initImpl(Plugin))
    return Err;

  // Read and reinitialize the envars that depend on the device initialization.
  // Notice these two envars may change the stack size and heap size of the
  // device, so they need the device properly initialized.
  auto StackSizeEnvarOrErr = UInt64Envar::create(
      "LIBOMPTARGET_STACK_SIZE",
      [this](uint64_t &V) -> Error { return getDeviceStackSize(V); },
      [this](uint64_t V) -> Error { return setDeviceStackSize(V); });
  if (!StackSizeEnvarOrErr)
    return StackSizeEnvarOrErr.takeError();
  OMPX_TargetStackSize = std::move(*StackSizeEnvarOrErr);

  if (hasDeviceHeapSize()) {
    auto HeapSizeEnvarOrErr = UInt64Envar::create(
        "LIBOMPTARGET_HEAP_SIZE",
        [this](uint64_t &V) -> Error { return getDeviceHeapSize(V); },
        [this](uint64_t V) -> Error { return setDeviceHeapSize(V); });
    if (!HeapSizeEnvarOrErr)
      return HeapSizeEnvarOrErr.takeError();
    OMPX_TargetHeapSize = std::move(*HeapSizeEnvarOrErr);
  }

  // Update the maximum number of teams and threads after the device
  // initialization sets the corresponding hardware limit.
  if (OMP_NumTeams > 0)
    GridValues.GV_Max_Teams =
        std::min(GridValues.GV_Max_Teams, uint32_t(OMP_NumTeams));

  if (OMP_TeamsThreadLimit > 0)
    GridValues.GV_Max_WG_Size =
        std::min(GridValues.GV_Max_WG_Size, uint32_t(OMP_TeamsThreadLimit));

  return Plugin::success();
}

Error GenericDeviceTy::unloadBinary(DeviceImageTy *Image) {
  GenericGlobalHandlerTy &Handler = Plugin.getGlobalHandler();
  auto ProfOrErr = Handler.readProfilingGlobals(*this, *Image);
  if (!ProfOrErr)
    return ProfOrErr.takeError();

  if (!ProfOrErr->empty()) {
    // Dump out profdata
    if ((OMPX_DebugKind.get() & uint32_t(DeviceDebugKind::PGODump)) ==
        uint32_t(DeviceDebugKind::PGODump))
      ProfOrErr->dump();

    // Write data to profiling file
    if (auto Err = ProfOrErr->write())
      return Err;
  }

  return unloadBinaryImpl(Image);
}

Error GenericDeviceTy::deinit(GenericPluginTy &Plugin) {
  // Run the global destructors first in case they required the RPC server.
  for (auto &I : LoadedImages) {
    if (auto Err = callGlobalDestructors(Plugin, *I))
      return Err;
  }

  if (RPCServer) {
    if (auto Err = RPCServer->deinitDevice(*this))
      return Err;
    RPCServer = nullptr;
  }

  for (auto &I : LoadedImages)
    if (auto Err = unloadBinary(I))
      return Err;
  LoadedImages.clear();

  if (RecordReplay) {
    if (auto Err = RecordReplay->deinit())
      return Err;
    delete RecordReplay;
    RecordReplay = nullptr;
  }

  return deinitImpl();
}
Expected<DeviceImageTy *>
GenericDeviceTy::loadBinary(GenericPluginTy &Plugin, StringRef InputTgtImage,
                            PluginContextTy *Context) {
  ODBG(OLDT_Init) << "Load data from image "
                  << static_cast<const void *>(InputTgtImage.bytes_begin());

  // An empty image is not a valid binary. Plugins behave differently given
  // empty binaries - e.g. CUDA will map to INVALID_BINARY, while L0 will map to
  // INVALID_SIZE (which is also associated with invalid kernel launch dims
  // etc.), so we guard here for consistent behavior across plugins and API
  // consumers (liboffload and libomptarget).
  if (InputTgtImage.empty())
    return Plugin::error(ErrorCode::INVALID_BINARY,
                         "provided binary image is empty");

  std::unique_ptr<MemoryBuffer> Buffer;
  if (identify_magic(InputTgtImage) == file_magic::bitcode) {
    auto CompiledImageOrErr = Plugin.getJIT().process(InputTgtImage, *this);
    if (!CompiledImageOrErr) {
      return Plugin::error(ErrorCode::COMPILE_FAILURE,
                           CompiledImageOrErr.takeError(),
                           "failure to jit IR image");
    }
    Buffer = std::move(*CompiledImageOrErr);
  } else {
    Buffer = MemoryBuffer::getMemBufferCopy(InputTgtImage);
  }

  // Load the binary and allocate the image object. Use the next available id
  // for the image id, which is the number of previously loaded images.
  auto ImageOrErr =
      loadBinaryImpl(std::move(Buffer), LoadedImages.size(), Context);
  if (!ImageOrErr)
    return ImageOrErr.takeError();
  DeviceImageTy *Image = *ImageOrErr;

  if (identify_magic(InputTgtImage) == file_magic::bitcode)
    Image->setIRImage(MemoryBuffer::getMemBufferCopy(InputTgtImage));

  // Add the image to list.
  LoadedImages.push_back(Image);

  if (auto Err = setupRPCServer(Plugin, *Image))
    return std::move(Err);

  // Call any global constructors present on the device.
  if (auto Err = callGlobalConstructors(Plugin, *Image))
    return std::move(Err);

  // Return the pointer to the table of entries.
  return Image;
}

Error GenericDeviceTy::setupRPCServer(GenericPluginTy &Plugin,
                                      DeviceImageTy &Image) {
  // The plugin either does not need an RPC server or it is unavailable.
  if (!shouldSetupRPCServer())
    return Plugin::success();

  // Check if this device needs to run an RPC server.
  RPCServerTy &Server = Plugin.getRPCServer();
  auto UsingOrErr =
      Server.isDeviceUsingRPC(*this, Plugin.getGlobalHandler(), Image);
  if (!UsingOrErr)
    return UsingOrErr.takeError();

  if (!UsingOrErr.get())
    return Plugin::success();

  if (auto Err = Server.initDevice(*this, Plugin.getGlobalHandler(), Image))
    return Err;

  if (auto Err = Server.startThread())
    return Err;

  RPCServer = &Server;
  ODBG(OLDT_Init) << "Running an RPC server on device " << getDeviceId();
  return Plugin::success();
}

Error PinnedAllocationMapTy::insertEntry(void *HstPtr, void *DevAccessiblePtr,
                                         size_t Size, bool ExternallyLocked) {
  // Insert the new entry into the map.
  auto Res = Allocs.insert({HstPtr, DevAccessiblePtr, Size, ExternallyLocked});
  if (!Res.second)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "cannot insert locked buffer entry");

  // Check whether the next entry overlaps with the inserted entry.
  auto It = std::next(Res.first);
  if (It == Allocs.end())
    return Plugin::success();

  const EntryTy *NextEntry = &(*It);
  if (intersects(NextEntry->HstPtr, NextEntry->Size, HstPtr, Size))
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "partial overlapping not allowed in locked buffers");

  return Plugin::success();
}

Error PinnedAllocationMapTy::eraseEntry(const EntryTy &Entry) {
  // Erase the existing entry. Notice this requires an additional map lookup,
  // but this should not be a performance issue. Using iterators would make
  // the code more difficult to read.
  size_t Erased = Allocs.erase({Entry.HstPtr});
  if (!Erased)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "cannot erase locked buffer entry");
  return Plugin::success();
}

Error PinnedAllocationMapTy::registerEntryUse(const EntryTy &Entry,
                                              void *HstPtr, size_t Size) {
  if (!contains(Entry.HstPtr, Entry.Size, HstPtr, Size))
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "partial overlapping not allowed in locked buffers");

  ++Entry.References;
  return Plugin::success();
}

Expected<bool> PinnedAllocationMapTy::unregisterEntryUse(const EntryTy &Entry) {
  if (Entry.References == 0)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "invalid number of references");

  // Return whether this was the last user.
  return (--Entry.References == 0);
}

Error PinnedAllocationMapTy::registerHostBuffer(void *HstPtr,
                                                void *DevAccessiblePtr,
                                                size_t Size) {
  assert(HstPtr && "Invalid pointer");
  assert(DevAccessiblePtr && "Invalid pointer");
  assert(Size && "Invalid size");

  std::lock_guard<std::shared_mutex> Lock(Mutex);

  // No pinned allocation should intersect.
  const EntryTy *Entry = findIntersecting(HstPtr);
  if (Entry)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "cannot insert entry due to an existing one");

  // Now insert the new entry.
  return insertEntry(HstPtr, DevAccessiblePtr, Size);
}

Error PinnedAllocationMapTy::unregisterHostBuffer(void *HstPtr) {
  assert(HstPtr && "Invalid pointer");

  std::lock_guard<std::shared_mutex> Lock(Mutex);

  const EntryTy *Entry = findIntersecting(HstPtr);
  if (!Entry)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "cannot find locked buffer");

  // The address in the entry should be the same we are unregistering.
  if (Entry->HstPtr != HstPtr)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "unexpected host pointer in locked buffer entry");

  // Unregister from the entry.
  auto LastUseOrErr = unregisterEntryUse(*Entry);
  if (!LastUseOrErr)
    return LastUseOrErr.takeError();

  // There should be no other references to the pinned allocation.
  if (!(*LastUseOrErr))
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "the locked buffer is still being used");

  // Erase the entry from the map.
  return eraseEntry(*Entry);
}

Expected<void *> PinnedAllocationMapTy::registerMemory(void *HstPtr,
                                                       size_t Size,
                                                       bool LockMemory) {
  assert(HstPtr && "Invalid pointer");
  assert(Size && "Invalid size");

  std::lock_guard<std::shared_mutex> Lock(Mutex);

  const EntryTy *Entry = findIntersecting(HstPtr);

  if (Entry) {
    // An already registered intersecting buffer was found. Register a new use.
    if (auto Err = registerEntryUse(*Entry, HstPtr, Size))
      return std::move(Err);

    // Return the device accessible pointer with the correct offset.
    return utils::advancePtr(Entry->DevAccessiblePtr,
                             utils::getPtrDiff(HstPtr, Entry->HstPtr));
  }

  size_t BaseSize;
  void *BaseHstPtr, *BaseDevAccessiblePtr;

  // Check if it was externally pinned by a vendor-specific API.
  auto IsPinnedOrErr = Device.isPinnedPtrImpl(HstPtr, BaseHstPtr,
                                              BaseDevAccessiblePtr, BaseSize);
  if (!IsPinnedOrErr)
    return IsPinnedOrErr.takeError();

  // If pinned, just insert the entry representing the whole pinned buffer.
  if (*IsPinnedOrErr) {
    if (auto Err = insertEntry(BaseHstPtr, BaseDevAccessiblePtr, BaseSize,
                               /*Externallylocked=*/true))
      return std::move(Err);
    return BaseDevAccessiblePtr;
  }

  // Not externally pinned. Do nothing if locking of mapped buffers is disabled.
  if (!LockMemory)
    return nullptr;

  // No intersecting registered allocation found in the map. First, lock the
  // host buffer and retrieve the device accessible pointer.
  auto DevAccessiblePtrOrErr = Device.dataLockImpl(HstPtr, Size);
  if (!DevAccessiblePtrOrErr)
    return DevAccessiblePtrOrErr.takeError();

  // Now insert the new entry into the map.
  if (auto Err = insertEntry(HstPtr, *DevAccessiblePtrOrErr, Size))
    return std::move(Err);

  // Return the device accessible pointer.
  return *DevAccessiblePtrOrErr;
}

Error PinnedAllocationMapTy::unregisterMemory(void *HstPtr, bool UnlockMemory) {
  assert(HstPtr && "Invalid pointer");

  std::lock_guard<std::shared_mutex> Lock(Mutex);

  const EntryTy *Entry = findIntersecting(HstPtr);

  // No entry but automatic locking of mapped buffers is disabled, so
  // nothing to do.
  if (!Entry && !UnlockMemory)
    return Plugin::success();

  if (!Entry)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "cannot find locked buffer");

  // Unregister from the locked buffer. No need to do anything if there are
  // others using the allocation.
  auto LastUseOrErr = unregisterEntryUse(*Entry);
  if (!LastUseOrErr)
    return LastUseOrErr.takeError();

  // No need to do anything if there are others using the allocation.
  if (!(*LastUseOrErr))
    return Plugin::success();

  // This was the last user of the allocation. Unlock the original locked buffer
  // if it was locked by the plugin. Do not unlock it if it was locked by an
  // external entity. Unlock the buffer using the host pointer of the entry.
  if (!Entry->ExternallyLocked)
    if (auto Err = Device.dataUnlockImpl(Entry->HstPtr))
      return Err;

  // Erase the entry from the map.
  return eraseEntry(*Entry);
}

Error GenericDeviceTy::synchronize(__tgt_async_info *AsyncInfo,
                                   bool ReleaseQueue) {
  if (!AsyncInfo)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "invalid async info queue");

  SmallVector<std::pair<void *, TargetAllocTy>> AllocsToDelete{};
  {
    std::lock_guard<std::mutex> AllocationGuard{AsyncInfo->Mutex};

    // This can be false when no work has been added to the AsyncInfo. In which
    // case, the device has nothing to synchronize.
    if (AsyncInfo->Queue)
      if (auto Err = synchronizeImpl(*AsyncInfo, ReleaseQueue))
        return Err;

    std::swap(AllocsToDelete, AsyncInfo->AssociatedAllocations);
  }

  for (auto [Ptr, Kind] : AllocsToDelete)
    if (auto Err = dataDelete(Ptr, Kind))
      return Err;

  return Plugin::success();
}

Error GenericDeviceTy::queryAsync(__tgt_async_info *AsyncInfo,
                                  bool ReleaseQueue,
                                  bool *IsQueueWorkCompleted) {
  if (!AsyncInfo || !AsyncInfo->Queue)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "invalid async info queue");

  bool WorkCompleted = false;
  SmallVector<std::pair<void *, TargetAllocTy>> AllocsToDelete{};

  {
    // Query and collect under the mutex, as synchronize does. Querying outside
    // it would let an operation issued in between have its allocations freed
    // here while it is still using them.
    std::lock_guard<std::mutex> AllocationGuard{AsyncInfo->Mutex};
    if (auto Err = queryAsyncImpl(*AsyncInfo, ReleaseQueue, &WorkCompleted)) {
      if (IsQueueWorkCompleted)
        *IsQueueWorkCompleted = WorkCompleted;
      return Err;
    }

    // A completed query is a completion point like synchronize(), and may be
    // the only one this async info ever gets, so release its allocations here.
    if (WorkCompleted)
      std::swap(AllocsToDelete, AsyncInfo->AssociatedAllocations);
  }

  for (auto [Ptr, Kind] : AllocsToDelete)
    if (auto Err = dataDelete(Ptr, Kind))
      return Err;

  if (IsQueueWorkCompleted)
    *IsQueueWorkCompleted = WorkCompleted;

  return Plugin::success();
}

Error GenericDeviceTy::memoryVAMap(void **Addr, void *VAddr, size_t *RSize) {
  return Plugin::error(ErrorCode::UNSUPPORTED,
                       "device does not support VA Management");
}

Error GenericDeviceTy::memoryVAUnMap(void *VAddr, size_t Size) {
  return Plugin::error(ErrorCode::UNSUPPORTED,
                       "device does not support VA Management");
}

Error GenericDeviceTy::getDeviceMemorySize(uint64_t &DSize) {
  return Plugin::error(
      ErrorCode::UNIMPLEMENTED,
      "missing getDeviceMemorySize implementation (required by RR-heuristic");
}

Expected<void *> GenericDeviceTy::dataAlloc(int64_t Size, void *HostPtr,
                                            TargetAllocTy Kind,
                                            size_t Alignment) {
  void *Alloc = nullptr;

  // TODO Check alignment.
  if (RecordReplay && RecordReplay->isRecordingOrReplaying())
    return RecordReplay->allocate(Size);

  auto AllocOrErr = allocate(Size, HostPtr, Kind, Alignment);
  if (!AllocOrErr)
    return AllocOrErr.takeError();
  Alloc = *AllocOrErr;
  if (!Alloc)
    return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
                         "failed to allocate from device allocator");

  if (Alignment > 0 && !isAddrAligned(Align(Alignment), Alloc)) {
    if (auto Err = free(Alloc, Kind))
      return Err;

    return Plugin::error(ErrorCode::UNSUPPORTED,
                         "device allocator returned a misaligned pointer");
  }

  // Keep track of the allocation stack if we track allocation traces.
  if (OMPX_TrackAllocationTraces) {
    std::string StackTrace;
    llvm::raw_string_ostream OS(StackTrace);
    llvm::sys::PrintStackTrace(OS);

    AllocationTraceInfoTy *ATI = new AllocationTraceInfoTy();
    ATI->AllocationTrace = std::move(StackTrace);
    ATI->DevicePtr = Alloc;
    ATI->HostPtr = HostPtr;
    ATI->Size = Size;
    ATI->Kind = Kind;

    auto AllocationTraceMap = AllocationTraces.getExclusiveAccessor();
    auto *&MapATI = (*AllocationTraceMap)[Alloc];
    ATI->LastAllocationInfo = MapATI;
    MapATI = ATI;
  }

  return Alloc;
}

Error GenericDeviceTy::dataDelete(void *TgtPtr, TargetAllocTy Kind) {
  // Free is a noop when recording or replaying.
  if (RecordReplay && RecordReplay->isRecordingOrReplaying())
    return RecordReplay->deallocate(TgtPtr);

  // Keep track of the deallocation stack if we track allocation traces.
  if (OMPX_TrackAllocationTraces) {
    AllocationTraceInfoTy *ATI = nullptr;
    {
      auto AllocationTraceMap = AllocationTraces.getExclusiveAccessor();
      ATI = (*AllocationTraceMap)[TgtPtr];
    }

    std::string StackTrace;
    llvm::raw_string_ostream OS(StackTrace);
    llvm::sys::PrintStackTrace(OS);

    if (!ATI)
      ErrorReporter::reportDeallocationOfNonAllocatedPtr(TgtPtr, Kind, ATI,
                                                         StackTrace);

    // ATI is not null, thus we can lock it to inspect and modify it further.
    std::lock_guard<std::mutex> LG(ATI->Lock);
    if (!ATI->DeallocationTrace.empty())
      ErrorReporter::reportDeallocationOfDeallocatedPtr(TgtPtr, Kind, ATI,
                                                        StackTrace);

    if (ATI->Kind != Kind)
      ErrorReporter::reportDeallocationOfWrongPtrKind(TgtPtr, Kind, ATI,
                                                      StackTrace);

    ATI->DeallocationTrace = StackTrace;
  }

  if (auto Err = free(TgtPtr, Kind))
    return Err;

  return Plugin::success();
}

Error GenericDeviceTy::dataSubmit(void *TgtPtr, const void *HstPtr,
                                  int64_t Size, __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);

  auto Err = dataSubmitImpl(TgtPtr, HstPtr, Size, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Error GenericDeviceTy::dataRetrieve(void *HstPtr, const void *TgtPtr,
                                    int64_t Size, __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);

  auto Err = dataRetrieveImpl(HstPtr, TgtPtr, Size, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Error GenericDeviceTy::dataMemcpy(void *DstPtr, const void *SrcPtr,
                                  int64_t Size, __tgt_async_info *AsyncInfo) {
  if (Size == 0)
    return Plugin::success();

  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);

  auto Err = dataMemcpyImpl(DstPtr, SrcPtr, Size, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Error GenericDeviceTy::dataExchange(const void *SrcPtr, GenericDeviceTy &DstDev,
                                    void *DstPtr, int64_t Size,
                                    __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);

  auto Err = dataExchangeImpl(SrcPtr, DstDev, DstPtr, Size, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Error GenericDeviceTy::dataFill(void *TgtPtr, const void *PatternPtr,
                                int64_t PatternSize, int64_t Size,
                                __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);
  auto Err =
      dataFillImpl(TgtPtr, PatternPtr, PatternSize, Size, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Error GenericDeviceTy::dataPrefetch(size_t Count, const void **Mems,
                                    const size_t *Sizes, bool ToHost,
                                    __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);
  auto Err = dataPrefetchImpl(Count, Mems, Sizes, ToHost, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Error GenericDeviceTy::launchKernel(void *EntryPtr,
                                    KernelLaunchArgsTy &LaunchArgs,
                                    __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);

  GenericKernelTy &GenericKernel =
      *reinterpret_cast<GenericKernelTy *>(EntryPtr);

  {
    std::string StackTrace;
    if (OMPX_TrackNumKernelLaunches) {
      llvm::raw_string_ostream OS(StackTrace);
      llvm::sys::PrintStackTrace(OS);
    }

    auto KernelTraceInfoRecord = KernelLaunchTraces.getExclusiveAccessor();
    (*KernelTraceInfoRecord)
        .emplace(&GenericKernel, std::move(StackTrace), AsyncInfo);
  }

  auto Err = GenericKernel.launch(*this, LaunchArgs, AsyncInfoWrapper);

  AsyncInfoWrapper.finalize(Err);

  return Err;
}

Error PluginContextTy::initAsyncInfo(GenericDeviceTy &Device,
                                     __tgt_async_info **AsyncInfoPtr) {
  assert(AsyncInfoPtr && "Invalid async info");

  *AsyncInfoPtr = new __tgt_async_info();

  AsyncInfoWrapperTy AsyncInfoWrapper(Device, *AsyncInfoPtr);

  auto Err = initAsyncInfoImpl(Device, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

PluginContextTy::~PluginContextTy() = default;

MemoryManagerTy *
PluginContextTy::getDeviceMemoryManagerFor(GenericDeviceTy &Device,
                                           TargetAllocTy Kind) {
  if (Device.OMPX_TrackAllocationTraces)
    return nullptr;

  if (Kind == TARGET_ALLOC_DEFAULT)
    Kind = TARGET_ALLOC_DEVICE;
  assert((Kind == TARGET_ALLOC_DEVICE || Kind == TARGET_ALLOC_SHARED) &&
         "host allocations are not device-bound");

  size_t Threshold;
  if (Kind == TARGET_ALLOC_DEVICE) {
    auto [EnvThreshold, EnableMM] = MemoryManagerTy::getSizeThresholdFromEnv();
    if (!EnableMM)
      return nullptr;
    Threshold =
        EnvThreshold ? EnvThreshold : Device.getMemoryManagerSizeThreshold();
  } else {
    Threshold = MemoryManagerTy::DefaultSizeThreshold;
  }

  std::pair<GenericDeviceTy *, int> Key{&Device, static_cast<int>(Kind)};

  std::lock_guard<std::mutex> Lock(MemoryManagersMutex);
  auto It = DeviceMemoryManagers.find(Key);
  if (It != DeviceMemoryManagers.end())
    return It->second.get();

  auto Manager = std::make_unique<MemoryManagerTy>(Device, Threshold, Kind);
  auto *Raw = Manager.get();
  DeviceMemoryManagers[Key] = std::move(Manager);
  return Raw;
}

MemoryManagerTy *PluginContextTy::getHostMemoryManager() {
  if (Devices.empty())
    return nullptr;
  if (Devices.front()->OMPX_TrackAllocationTraces)
    return nullptr;

  std::lock_guard<std::mutex> Lock(MemoryManagersMutex);
  if (HostMemoryManager)
    return HostMemoryManager.get();

  HostMemoryManager = std::make_unique<MemoryManagerTy>(
      *Devices.front(), MemoryManagerTy::DefaultSizeThreshold,
      TARGET_ALLOC_HOST);
  return HostMemoryManager.get();
}

Expected<void *> PluginContextTy::allocate(GenericDeviceTy &Device,
                                           int64_t Size, void *HostPtr,
                                           TargetAllocTy Kind,
                                           size_t Alignment) {
  // Record-replay hands out interior pointers into a preallocated slab so
  // recorded kernels can re-execute at their original addresses; the MM pool
  // must be bypassed for those allocations to reach the RR bump allocator.
  if (auto *RR = Device.getRecordReplay(); RR && RR->isRecordingOrReplaying())
    return Device.dataAlloc(Size, HostPtr, Kind, Alignment);

  MemoryManagerTy *MM = (Kind == TARGET_ALLOC_HOST)
                            ? getHostMemoryManager()
                            : getDeviceMemoryManagerFor(Device, Kind);
  if (MM)
    return MM->allocate(Size, HostPtr, Alignment);
  return Device.dataAlloc(Size, HostPtr, Kind, Alignment);
}

Error PluginContextTy::deallocate(void *Ptr) {
  assert(!Devices.empty() && "context constructed without devices");
  auto InfoOrErr = getAllocInfo(Ptr);
  if (!InfoOrErr)
    return InfoOrErr.takeError();
  GenericDeviceTy *OwnerDevice = InfoOrErr->Device;
  if (!OwnerDevice)
    OwnerDevice = Devices.front();
  return deallocate(*OwnerDevice, Ptr, InfoOrErr->Kind);
}

Error PluginContextTy::deallocate(GenericDeviceTy &Device, void *Ptr,
                                  TargetAllocTy Kind) {
  // Symmetric with allocate: record-replay allocations never entered the MM
  // pool, so route their free through dataDelete's RR shortcut.
  if (auto *RR = Device.getRecordReplay(); RR && RR->isRecordingOrReplaying())
    return Device.dataDelete(Ptr, Kind);

  MemoryManagerTy *MM = (Kind == TARGET_ALLOC_HOST)
                            ? getHostMemoryManager()
                            : getDeviceMemoryManagerFor(Device, Kind);
  if (MM)
    return MM->free(Ptr);
  return Device.dataDelete(Ptr, Kind);
}

PluginContextTy &
GenericPluginTy::getDefaultContext(GenericDeviceTy & /*Device*/) {
  assert(DefaultContext && "default context not initialized");
  return *DefaultContext;
}

Expected<std::unique_ptr<PluginContextTy>>
GenericPluginTy::createDefaultPluginContext() {
  return std::make_unique<DefaultPluginContextTy>(*this);
}

Error GenericDeviceTy::enqueueHostCall(void (*Callback)(void *), void *UserData,
                                       __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);

  auto Err = enqueueHostCallImpl(Callback, UserData, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Expected<InfoTreeNode> GenericDeviceTy::obtainInfo() {
  auto InfoOrErr = obtainInfoImpl();
  if (InfoOrErr)
    InfoOrErr->add("UID", getDeviceUid(), "", DeviceInfo::UID);
  return InfoOrErr;
}

Error GenericDeviceTy::printInfo() {
  auto InfoOrErr = obtainInfo();

  // Get the vendor-specific info entries describing the device properties.
  if (auto Err = InfoOrErr.takeError())
    return Err;

  // Print all info entries.
  InfoOrErr->print();

  return Plugin::success();
}

Error GenericDeviceTy::createEvent(void **EventPtrStorage,
                                   bool EnableProfiling) {
  return createEventImpl(EventPtrStorage, EnableProfiling);
}

Error GenericDeviceTy::destroyEvent(void *EventPtr, bool EnableProfiling) {
  return destroyEventImpl(EventPtr, EnableProfiling);
}

Error GenericDeviceTy::recordEvent(void *EventPtr, __tgt_async_info *AsyncInfo,
                                   bool EnableProfiling) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);

  auto Err = recordEventImpl(EventPtr, AsyncInfoWrapper, EnableProfiling);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Error GenericDeviceTy::waitEvent(void *EventPtr, __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);

  auto Err = waitEventImpl(EventPtr, AsyncInfoWrapper);
  AsyncInfoWrapper.finalize(Err);
  return Err;
}

Expected<bool> GenericDeviceTy::hasPendingWork(__tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);
  auto Res = hasPendingWorkImpl(AsyncInfoWrapper);
  if (auto Err = Res.takeError()) {
    AsyncInfoWrapper.finalize(Err);
    return Err;
  }

  auto Err = Plugin::success();
  AsyncInfoWrapper.finalize(Err);
  if (Err)
    return Err;
  return Res;
}

Expected<bool> GenericDeviceTy::isEventComplete(void *Event,
                                                __tgt_async_info *AsyncInfo) {
  AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);
  auto Res = isEventCompleteImpl(Event, AsyncInfoWrapper);
  if (auto Err = Res.takeError()) {
    AsyncInfoWrapper.finalize(Err);
    return Err;
  }

  auto Err = Plugin::success();
  AsyncInfoWrapper.finalize(Err);
  if (Err)
    return Err;
  return Res;
}

Error GenericDeviceTy::syncEvent(void *EventPtr) {
  return syncEventImpl(EventPtr);
}

Expected<float> GenericDeviceTy::getEventElapsedTime(void *StartEventPtr,
                                                     void *EndEventPtr) {
  return getEventElapsedTimeImpl(StartEventPtr, EndEventPtr);
}

bool GenericDeviceTy::useAutoZeroCopy() { return useAutoZeroCopyImpl(); }

Expected<bool> GenericDeviceTy::isAccessiblePtr(const void *Ptr, size_t Size) {
  return isAccessiblePtrImpl(Ptr, Size);
}

void GenericDeviceTy::setDeviceUidFromVendorUid(StringRef VendorUid) {
  DeviceUid = std::string(Plugin.getName()) + "-" + std::string(VendorUid);
}

Error GenericPluginTy::init() {
  if (Initialized)
    return Plugin::success();

  auto NumDevicesOrErr = initImpl();
  if (!NumDevicesOrErr)
    return NumDevicesOrErr.takeError();
  Initialized = true;

  NumDevices = *NumDevicesOrErr;
  if (NumDevices == 0)
    return Plugin::success();

  assert(Devices.size() == 0 && "Plugin already initialized");
  Devices.resize(NumDevices, nullptr);

  GlobalHandler = createGlobalHandler();
  assert(GlobalHandler && "Invalid global handler");

  RPCServer = new RPCServerTy(*this);
  assert(RPCServer && "Invalid RPC server");

  auto DefaultCtxOrErr = createDefaultPluginContext();
  if (!DefaultCtxOrErr)
    return DefaultCtxOrErr.takeError();
  DefaultContext = std::move(*DefaultCtxOrErr);

  return Plugin::success();
}

Error GenericPluginTy::deinit() {
  assert(Initialized && "Plugin was not initialized!");

  // Release context-held resources before the devices that back them.
  DefaultContext.reset();

  // Deinitialize all active devices.
  for (int32_t DeviceId = 0; DeviceId < NumDevices; ++DeviceId) {
    if (Devices[DeviceId]) {
      if (auto Err = deinitDevice(DeviceId))
        return Err;
    }
    assert(!Devices[DeviceId] && "Device was not deinitialized");
  }

  // There is no global handler if no device is available.
  if (GlobalHandler)
    delete GlobalHandler;

  if (RPCServer) {
    if (Error Err = RPCServer->shutDown(*this))
      return Err;
    delete RPCServer;
  }

  // Perform last deinitializations on the plugin.
  if (Error Err = deinitImpl())
    return Err;
  Initialized = false;

  return Plugin::success();
}

Error GenericPluginTy::initDevice(int32_t DeviceId) {
  assert(!Devices[DeviceId] && "Device already initialized");

  // Create the device and save the reference.
  GenericDeviceTy *Device = createDevice(*this, DeviceId, NumDevices);
  assert(Device && "Invalid device");

  // Save the device reference into the list.
  Devices[DeviceId] = Device;

  // Initialize the device and its resources.
  return Device->init(*this);
}

Error GenericPluginTy::deinitDevice(int32_t DeviceId) {
  // The device may be already deinitialized.
  if (Devices[DeviceId] == nullptr)
    return Plugin::success();

  // Deinitialize the device and release its resources.
  if (auto Err = Devices[DeviceId]->deinit(*this))
    return Err;

  // Delete the device and invalidate its reference.
  delete Devices[DeviceId];
  Devices[DeviceId] = nullptr;

  return Plugin::success();
}

Expected<bool> GenericPluginTy::checkELFImage(StringRef Image) const {
  // First check if this image is a regular ELF file.
  if (!utils::elf::isELF(Image))
    return false;

  // Check if this image is an ELF with a matching machine value.
  auto MachineOrErr = utils::elf::checkMachine(Image, getMagicElfBits());
  if (!MachineOrErr)
    return MachineOrErr.takeError();

  return MachineOrErr;
}

Expected<bool> GenericPluginTy::checkBitcodeImage(StringRef Image) const {
  if (identify_magic(Image) != file_magic::bitcode)
    return false;

  LLVMContext Context;
  auto ModuleOrErr = getLazyBitcodeModule(MemoryBufferRef(Image, ""), Context,
                                          /*ShouldLazyLoadMetadata=*/true);
  if (!ModuleOrErr)
    return ModuleOrErr.takeError();
  Module &M = **ModuleOrErr;

  return M.getTargetTriple().getArch() == getTripleArch();
}

int32_t GenericPluginTy::is_initialized() const { return Initialized; }

int32_t GenericPluginTy::isPluginCompatible(StringRef Image) {
  auto HandleError = [&](Error Err) -> bool {
    std::string ErrStr = toString(std::move(Err));
    ODBG(OLDT_Init) << "Failure to check validity of image "
                    << static_cast<const void *>(Image.data()) << ": "
                    << ErrStr;
    return false;
  };
  switch (identify_magic(Image)) {
  case file_magic::elf:
  case file_magic::elf_relocatable:
  case file_magic::elf_executable:
  case file_magic::elf_shared_object:
  case file_magic::elf_core: {
    auto MatchOrErr = checkELFImage(Image);
    if (Error Err = MatchOrErr.takeError())
      return HandleError(std::move(Err));
    return *MatchOrErr;
  }
  case file_magic::bitcode: {
    auto MatchOrErr = checkBitcodeImage(Image);
    if (Error Err = MatchOrErr.takeError())
      return HandleError(std::move(Err));
    return *MatchOrErr;
  }
  default:
    auto MatchOrErr = isImageCompatible(Image);
    if (Error Err = MatchOrErr.takeError())
      return HandleError(std::move(Err));
    return *MatchOrErr;
  }
}

int32_t GenericPluginTy::isDeviceCompatible(int32_t DeviceId, StringRef Image) {
  auto HandleError = [&](Error Err) -> bool {
    std::string ErrStr = toString(std::move(Err));
    ODBG(OLDT_Init) << "Failure to check validity of image "
                    << static_cast<const void *>(Image.data()) << ": "
                    << ErrStr;
    return false;
  };
  switch (identify_magic(Image)) {
  case file_magic::elf:
  case file_magic::elf_relocatable:
  case file_magic::elf_executable:
  case file_magic::elf_shared_object:
  case file_magic::elf_core: {
    auto MatchOrErr = checkELFImage(Image);
    if (Error Err = MatchOrErr.takeError())
      return HandleError(std::move(Err));
    if (!*MatchOrErr)
      return false;

    // Perform plugin-dependent checks for the specific architecture if needed.
    auto CompatibleOrErr = isELFCompatible(DeviceId, Image);
    if (Error Err = CompatibleOrErr.takeError())
      return HandleError(std::move(Err));
    return *CompatibleOrErr;
  }
  case file_magic::bitcode: {
    auto MatchOrErr = checkBitcodeImage(Image);
    if (Error Err = MatchOrErr.takeError())
      return HandleError(std::move(Err));
    return *MatchOrErr;
  }
  default:
    auto MatchOrErr = isImageCompatible(DeviceId, Image);
    if (Error Err = MatchOrErr.takeError())
      return HandleError(std::move(Err));
    return *MatchOrErr;
  }
}

int32_t GenericPluginTy::is_device_initialized(int32_t DeviceId) const {
  return isValidDeviceId(DeviceId) && Devices[DeviceId] != nullptr;
}

int32_t GenericPluginTy::init_device(int32_t DeviceId) {
  auto Err = initDevice(DeviceId);
  if (Err) {
    REPORT() << "Failure to initialize device " << DeviceId << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::number_of_devices() { return getNumDevices(); }

int32_t GenericPluginTy::is_data_exchangable(int32_t SrcDeviceId,
                                             int32_t DstDeviceId) {
  return isDataExchangable(SrcDeviceId, DstDeviceId);
}

int32_t GenericPluginTy::initialize_record_replay(
    int32_t DeviceId, int64_t MemorySize, void *VAddr, bool IsRecord,
    bool IsNative, bool SaveOutput, bool EmitReport, const char *ReportFilename,
    const char *OutputDirPath) {
  GenericDeviceTy &Device = getDevice(DeviceId);

  if (auto Err = Device.initRecordReplay(MemorySize, VAddr, IsRecord, IsNative,
                                         SaveOutput, EmitReport, ReportFilename,
                                         OutputDirPath)) {
    REPORT() << "Failure to initialize RR with " << MemorySize
             << " bytes on device " << DeviceId << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }
  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::load_binary(int32_t DeviceId,
                                     __tgt_device_image *TgtImage,
                                     __tgt_device_binary *Binary) {
  GenericDeviceTy &Device = getDevice(DeviceId);

  StringRef Buffer(reinterpret_cast<const char *>(TgtImage->ImageStart),
                   utils::getPtrDiff(TgtImage->ImageEnd, TgtImage->ImageStart));
  auto ImageOrErr = Device.loadBinary(*this, Buffer, /*Context=*/nullptr);
  if (!ImageOrErr) {
    auto Err = ImageOrErr.takeError();
    REPORT() << "Failure to load binary image " << TgtImage << " on device "
             << DeviceId << ": " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  DeviceImageTy *Image = *ImageOrErr;
  assert(Image != nullptr && "Invalid Image");

  *Binary = __tgt_device_binary{reinterpret_cast<uint64_t>(Image)};

  return OFFLOAD_SUCCESS;
}

void *GenericPluginTy::data_alloc(int32_t DeviceId, int64_t Size, void *HostPtr,
                                  int32_t Kind) {
  auto &Device = getDevice(DeviceId);
  auto AllocOrErr = getDefaultContext(Device).allocate(
      Device, Size, HostPtr, static_cast<TargetAllocTy>(Kind),
      /*Alignment=*/0);
  if (!AllocOrErr) {
    REPORT() << "Failure to allocate device memory: "
             << toString(AllocOrErr.takeError());
    return nullptr;
  }
  assert(*AllocOrErr && "Null pointer upon successful allocation");
  return *AllocOrErr;
}

int32_t GenericPluginTy::data_delete(int32_t DeviceId, void *TgtPtr,
                                     int32_t Kind) {
  auto &Device = getDevice(DeviceId);
  if (auto Err = getDefaultContext(Device).deallocate(
          Device, TgtPtr, static_cast<TargetAllocTy>(Kind))) {
    REPORT() << "Failure to deallocate device pointer " << TgtPtr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }
  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::data_lock(int32_t DeviceId, void *Ptr, int64_t Size,
                                   void **LockedPtr) {
  auto LockedPtrOrErr = getDevice(DeviceId).registerMemory(Ptr, Size);
  if (!LockedPtrOrErr) {
    auto Err = LockedPtrOrErr.takeError();
    REPORT() << "Failure to lock memory " << Ptr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  if (!(*LockedPtrOrErr)) {
    REPORT() << "Failure to lock memory " << Ptr
             << ": obtained a null locked pointer";
    return OFFLOAD_FAIL;
  }
  *LockedPtr = *LockedPtrOrErr;

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::data_unlock(int32_t DeviceId, void *Ptr) {
  auto Err = getDevice(DeviceId).unregisterMemory(Ptr);
  if (Err) {
    REPORT() << "Failure to unlock memory " << Ptr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::data_notify_mapped(int32_t DeviceId, void *HstPtr,
                                            int64_t Size) {
  auto Err = getDevice(DeviceId).notifyDataMapped(HstPtr, Size);
  if (Err) {
    REPORT() << "Failure to notify data mapped " << HstPtr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::data_notify_unmapped(int32_t DeviceId, void *HstPtr) {
  auto Err = getDevice(DeviceId).notifyDataUnmapped(HstPtr);
  if (Err) {
    REPORT() << "Failure to notify data unmapped " << HstPtr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::data_submit(int32_t DeviceId, void *TgtPtr,
                                     void *HstPtr, int64_t Size) {
  return data_submit_async(DeviceId, TgtPtr, HstPtr, Size,
                           /*AsyncInfoPtr=*/nullptr);
}

int32_t GenericPluginTy::data_submit_async(int32_t DeviceId, void *TgtPtr,
                                           void *HstPtr, int64_t Size,
                                           __tgt_async_info *AsyncInfoPtr) {
  auto Err = getDevice(DeviceId).dataSubmit(TgtPtr, HstPtr, Size, AsyncInfoPtr);
  if (Err) {
    REPORT() << "Failure to copy data from host to device. Pointers: host "
             << "= " << HstPtr << ", device = " << TgtPtr << ", size = " << Size
             << ": " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::data_retrieve(int32_t DeviceId, void *HstPtr,
                                       void *TgtPtr, int64_t Size) {
  return data_retrieve_async(DeviceId, HstPtr, TgtPtr, Size,
                             /*AsyncInfoPtr=*/nullptr);
}

int32_t GenericPluginTy::data_retrieve_async(int32_t DeviceId, void *HstPtr,
                                             void *TgtPtr, int64_t Size,
                                             __tgt_async_info *AsyncInfoPtr) {
  auto Err =
      getDevice(DeviceId).dataRetrieve(HstPtr, TgtPtr, Size, AsyncInfoPtr);
  if (Err) {
    REPORT() << "Failure to copy data from device to host. Pointers: host "
             << "= " << HstPtr << ", device = " << TgtPtr << ", size = " << Size
             << ": " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::data_exchange(int32_t SrcDeviceId, void *SrcPtr,
                                       int32_t DstDeviceId, void *DstPtr,
                                       int64_t Size) {
  return data_exchange_async(SrcDeviceId, SrcPtr, DstDeviceId, DstPtr, Size,
                             /*AsyncInfoPtr=*/nullptr);
}

int32_t GenericPluginTy::data_exchange_async(int32_t SrcDeviceId, void *SrcPtr,
                                             int DstDeviceId, void *DstPtr,
                                             int64_t Size,
                                             __tgt_async_info *AsyncInfo) {
  GenericDeviceTy &SrcDevice = getDevice(SrcDeviceId);
  GenericDeviceTy &DstDevice = getDevice(DstDeviceId);
  auto Err = SrcDevice.dataExchange(SrcPtr, DstDevice, DstPtr, Size, AsyncInfo);
  if (Err) {
    REPORT() << "Failure to copy data from device (" << SrcDeviceId
             << ") to device (" << DstDeviceId
             << "). Pointers: host = " << SrcPtr << ", device = " << DstPtr
             << ", size = " << Size << ": " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::launch_kernel(int32_t DeviceId, void *TgtEntryPtr,
                                       KernelLaunchArgsTy &LaunchArgs,
                                       __tgt_async_info *AsyncInfoPtr) {
  auto Err =
      getDevice(DeviceId).launchKernel(TgtEntryPtr, LaunchArgs, AsyncInfoPtr);
  if (Err) {
    REPORT() << "Failure to run target region " << TgtEntryPtr << " in device "
             << DeviceId << ": " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::synchronize(int32_t DeviceId,
                                     __tgt_async_info *AsyncInfoPtr) {
  auto Err = getDevice(DeviceId).synchronize(AsyncInfoPtr);
  if (Err) {
    REPORT() << "Failure to synchronize stream " << AsyncInfoPtr->Queue << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::query_async(int32_t DeviceId,
                                     __tgt_async_info *AsyncInfoPtr) {
  auto Err = getDevice(DeviceId).queryAsync(AsyncInfoPtr);
  if (Err) {
    REPORT() << "Failure to query stream " << AsyncInfoPtr->Queue << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

InfoTreeNode GenericPluginTy::obtain_device_info(int32_t DeviceId) {
  auto InfoOrErr = getDevice(DeviceId).obtainInfo();
  if (auto Err = InfoOrErr.takeError()) {
    REPORT() << "Failure to obtain device " << DeviceId
             << " info: " << toString(std::move(Err));
    return InfoTreeNode{};
  }
  return std::move(*InfoOrErr);
}

void GenericPluginTy::print_device_info(int32_t DeviceId) {
  if (auto Err = getDevice(DeviceId).printInfo())
    REPORT() << "Failure to print device " << DeviceId
             << " info: " << toString(std::move(Err));
}

int32_t GenericPluginTy::create_event(int32_t DeviceId, void **EventPtr) {
  auto Err = getDevice(DeviceId).createEvent(EventPtr);
  if (Err) {
    REPORT() << "Failure to create event: " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::record_event(int32_t DeviceId, void *EventPtr,
                                      __tgt_async_info *AsyncInfoPtr) {
  auto Err = getDevice(DeviceId).recordEvent(EventPtr, AsyncInfoPtr);
  if (Err) {
    REPORT() << "Failure to record event " << EventPtr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::wait_event(int32_t DeviceId, void *EventPtr,
                                    __tgt_async_info *AsyncInfoPtr) {
  auto Err = getDevice(DeviceId).waitEvent(EventPtr, AsyncInfoPtr);
  if (Err) {
    REPORT() << "Failure to wait event " << EventPtr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::sync_event(int32_t DeviceId, void *EventPtr) {
  auto Err = getDevice(DeviceId).syncEvent(EventPtr);
  if (Err) {
    REPORT() << "Failure to synchronize event " << EventPtr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::get_event_elapsed_time(int32_t DeviceId,
                                                void *StartEventPtr,
                                                void *EndEventPtr,
                                                float *ElapsedTime) {
  auto ElapsedTimeOrErr =
      getDevice(DeviceId).getEventElapsedTime(StartEventPtr, EndEventPtr);
  if (!ElapsedTimeOrErr) {
    REPORT() << "Failure to get elapsed time between events " << StartEventPtr
             << " and " << EndEventPtr << ": "
             << toString(ElapsedTimeOrErr.takeError());
    return OFFLOAD_FAIL;
  }

  *ElapsedTime = *ElapsedTimeOrErr;
  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::destroy_event(int32_t DeviceId, void *EventPtr) {
  auto Err = getDevice(DeviceId).destroyEvent(EventPtr);
  if (Err) {
    REPORT() << "Failure to destroy event " << EventPtr << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}

void GenericPluginTy::set_info_flag(uint32_t NewInfoLevel) {
  std::atomic<uint32_t> &InfoLevel = getInfoLevelInternal();
  InfoLevel.store(NewInfoLevel);
}

int32_t GenericPluginTy::use_auto_zero_copy(int32_t DeviceId) {
  return getDevice(DeviceId).useAutoZeroCopy();
}

int32_t GenericPluginTy::is_accessible_ptr(int32_t DeviceId, const void *Ptr,
                                           size_t Size) {
  auto HandleError = [&](Error Err) -> bool {
    std::string ErrStr = toString(std::move(Err));
    ODBG(OLDT_Device) << "Failure while checking accessibility of pointer "
                      << Ptr << " for device " << DeviceId << ": " << ErrStr;
    return false;
  };

  auto AccessibleOrErr = getDevice(DeviceId).isAccessiblePtr(Ptr, Size);
  if (Error Err = AccessibleOrErr.takeError())
    return HandleError(std::move(Err));

  return *AccessibleOrErr;
}

int32_t GenericPluginTy::get_global(__tgt_device_binary Binary, uint64_t Size,
                                    const char *Name, void **DevicePtr) {
  assert(Binary.handle && "Invalid device binary handle");
  DeviceImageTy &Image = *reinterpret_cast<DeviceImageTy *>(Binary.handle);

  GenericDeviceTy &Device = Image.getDevice();

  GlobalTy DeviceGlobal(Name, Size);
  GenericGlobalHandlerTy &GHandler = getGlobalHandler();
  if (auto Err =
          GHandler.getGlobalMetadataFromDevice(Device, Image, DeviceGlobal)) {
    consumeError(std::move(Err));
    return OFFLOAD_FAIL;
  }

  *DevicePtr = DeviceGlobal.getPtr();
  assert(DevicePtr && "Invalid device global's address");

  // Save the loaded globals if we are recording.
  RecordReplayTy *RecordReplay = Device.getRecordReplay();
  if (RecordReplay && RecordReplay->isRecording())
    RecordReplay->addGlobal(Name, Size, *DevicePtr);

  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::get_function(__tgt_device_binary Binary,
                                      const char *Name, void **KernelPtr) {
  assert(Binary.handle && "Invalid device binary handle");
  DeviceImageTy &Image = *reinterpret_cast<DeviceImageTy *>(Binary.handle);

  GenericDeviceTy &Device = Image.getDevice();

  auto KernelOrErr = Device.constructKernel(Name);
  if (Error Err = KernelOrErr.takeError()) {
    REPORT() << "Failure to look up kernel: " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  GenericKernelTy &Kernel = *KernelOrErr;
  if (auto Err = Kernel.init(Device, Image)) {
    REPORT() << "Failure to init kernel: " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  // Note that this is not the kernel's device address.
  *KernelPtr = &Kernel;
  return OFFLOAD_SUCCESS;
}

/// Create OpenMP interop with the given interop context
omp_interop_val_t *
GenericPluginTy::create_interop(int32_t ID, int32_t InteropContext,
                                interop_spec_t *InteropSpec) {
  assert(InteropSpec && "Interop spec is null");
  auto &Device = getDevice(ID);
  auto InteropOrErr = Device.createInterop(InteropContext, *InteropSpec);
  if (!InteropOrErr) {
    REPORT() << "Failure to create interop object for device " << InteropSpec
             << ": " << toString(InteropOrErr.takeError());
    return nullptr;
  }
  return *InteropOrErr;
}

/// Release OpenMP interop object
int32_t GenericPluginTy::release_interop(int32_t ID,
                                         omp_interop_val_t *Interop) {
  assert(Interop && "Interop is null");
  assert(Interop->device_id == ID && "Interop does not match device id");
  auto &Device = getDevice(ID);
  auto Err = Device.releaseInterop(Interop);
  if (Err) {
    REPORT() << "Failure to release interop object " << Interop << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }
  return OFFLOAD_SUCCESS;
}

/// Flush the queue associated with the interop object if necessary
int32_t GenericPluginTy::flush_queue(omp_interop_val_t *Interop) {
  assert(Interop && "Interop is null");
  auto Err = flushQueueImpl(Interop);
  if (Err) {
    REPORT() << "Failure to flush interop object " << Interop
             << " queue: " << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }
  return OFFLOAD_SUCCESS;
}

/// Perform a host synchronization with the queue associated with the interop
/// object and wait for it to complete.
int32_t GenericPluginTy::sync_barrier(omp_interop_val_t *Interop) {
  assert(Interop && "Interop is null");
  auto Err = syncBarrierImpl(Interop);
  if (Err) {
    REPORT() << "Failure to synchronize interop object " << Interop << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }
  return OFFLOAD_SUCCESS;
}

/// Queue an asynchronous barrier in the queue associated with the interop
/// object and return immediately.
int32_t GenericPluginTy::async_barrier(omp_interop_val_t *Interop) {
  assert(Interop && "Interop is null");
  auto Err = asyncBarrierImpl(Interop);
  if (Err) {
    REPORT() << "Failure to queue barrier in interop object " << Interop << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }
  return OFFLOAD_SUCCESS;
}

int32_t GenericPluginTy::data_fence(int32_t DeviceId,
                                    __tgt_async_info *AsyncInfo) {
  auto Err = getDevice(DeviceId).dataFence(AsyncInfo);
  if (Err) {
    REPORT() << "Failure to place data fence on device " << DeviceId << ": "
             << toString(std::move(Err));
    return OFFLOAD_FAIL;
  }

  return OFFLOAD_SUCCESS;
}