//===----RTLs/amdgpu/src/rtl.cpp - Target RTLs Implementation ----- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// RTL NextGen for AMDGPU machine
//
//===----------------------------------------------------------------------===//

#include <atomic>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <deque>
#include <functional>
#include <map>
#include <mutex>
#include <string>
#include <system_error>
#include <unistd.h>
#include <unordered_map>

#include "ErrorReporting.h"
#include "Shared/APITypes.h"
#include "Shared/Debug.h"
#include "Shared/Environment.h"
#include "Shared/RefCnt.h"
#include "Shared/Utils.h"
#include "Utils/ELF.h"

#include "GenericProfiler.h"
#include "GlobalHandler.h"
#include "OffloadAPI.h"
#include "PluginInterface.h"
#include "UtilitiesRTL.h"
#include "omptarget.h"

#include "llvm/ADT/SmallString.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/BinaryFormat/ELF.h"
#include "llvm/Frontend/OpenMP/OMPConstants.h"
#include "llvm/Frontend/OpenMP/OMPGridValues.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/FileOutputBuffer.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Program.h"
#include "llvm/Support/Signals.h"
#include "llvm/Support/VersionTuple.h"
#include "llvm/Support/raw_ostream.h"

#if !defined(__BYTE_ORDER__) || !defined(__ORDER_LITTLE_ENDIAN__) ||           \
    !defined(__ORDER_BIG_ENDIAN__)
#error "Missing preprocessor definitions for endianness detection."
#endif

// The HSA headers require these definitions.
#if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__)
#define LITTLEENDIAN_CPU
#elif defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__)
#define BIGENDIAN_CPU
#endif

#if defined(__has_include)
#if __has_include("hsa.h")
#include "hsa.h"
#include "hsa_ext_amd.h"
#elif __has_include("hsa/hsa.h")
#include "hsa/hsa.h"
#include "hsa/hsa_ext_amd.h"
#endif
#else
#include "hsa/hsa.h"
#include "hsa/hsa_ext_amd.h"
#endif

using namespace llvm::offload::debug;
using namespace error;

namespace llvm {
namespace omp {
namespace target {
namespace plugin {

/// Forward declarations for all specialized data structures.
struct AMDGPUKernelTy;
struct AMDGPUDeviceTy;
struct AMDGPUPluginTy;
struct AMDGPUStreamTy;
struct AMDGPUEventTy;
struct AMDGPUStreamManagerTy;
struct AMDGPUEventManagerTy;
struct AMDGPUDeviceImageTy;
struct AMDGPUMemoryManagerTy;
struct AMDGPUMemoryPoolTy;
struct AMDGPUSignalTy;

/// Use to transport information to profiler timing functions.
/// The profiler is captured as a pointer because this outlives the call that
/// schedules the asynchronous action.
struct ProfilingInfoTy {
  GenericProfilerTy *Profiler;
  hsa_agent_t Agent;
  AMDGPUSignalTy *Signal;
  double TicksToTime;
  void *ProfilerSpecificData;
};

static ProfilingInfoTy *getProfilingInfo(void *Data);

static std::pair<uint64_t, uint64_t>
getKernelStartAndEndTime(const ProfilingInfoTy *Args);

static std::pair<uint64_t, uint64_t>
getCopyStartAndEndTime(const ProfilingInfoTy *Args);

static Error timeKernelInNsAsync(void *Data);

static Error timeDataTransferInNsAsync(void *Data) {
  auto Args = getProfilingInfo(Data);
  auto [Start, End] = getCopyStartAndEndTime(Args);
  Args->Profiler->handleDataTransfer(Start, End, Args->ProfilerSpecificData);
  return Plugin::success();
}

static void *
getOrNullProfilerSpecificData(AsyncInfoWrapperTy &AsyncInfoWrapper) {
  __tgt_async_info *AI = AsyncInfoWrapper;
  return AI ? AI->ProfilerData : nullptr;
}

} // namespace plugin
} // namespace target
} // namespace omp
} // namespace llvm

static double setTicksToTime() {
  uint64_t TicksFrequency = 1;
  double TicksToTime = 1.0;

  hsa_status_t Status =
      hsa_system_get_info(HSA_SYSTEM_INFO_TIMESTAMP_FREQUENCY, &TicksFrequency);
  if (Status == HSA_STATUS_SUCCESS)
    TicksToTime = (double)1e9 / (double)TicksFrequency;

  return TicksToTime;
}

static double TicksToTime = 1.0;

static void setHSATicksToTimeConstant() { TicksToTime = setTicksToTime(); }

/// Get the current HSA-based system timestamp in nanoseconds.
/// Called by OmptTracing.cpp (from PluginOmpt) for device time queries.
uint64_t getSystemTimestampInNs() {
  uint64_t TimeStamp = 0;
  hsa_status_t Status =
      hsa_system_get_info(HSA_SYSTEM_INFO_TIMESTAMP, &TimeStamp);
  if (Status != HSA_STATUS_SUCCESS)
    return 0;
  return TimeStamp;
}

/// Enable or disable HSA async copy profiling for OMPT device tracing.
/// Called by OmptTracing.cpp when a tool activates/deactivates device tracing.
/// Full HSA profiling integration (enabling per-copy timing signals) will be
/// wired up in a follow-on commit.
void setOmptAsyncCopyProfile(bool Enable) {}

/// Enable or disable HSA queue kernel profiling for OMPT device tracing.
/// Called by OmptTracing.cpp when a tool activates/deactivates device tracing.
/// Full HSA queue profiling integration will be wired up in a follow-on commit.
void setGlobalOmptKernelProfile(void *Device, int Enable) {}

namespace llvm {
namespace omp {
namespace target {
namespace plugin {

namespace hsa_utils {

/// Adapted from HSA-runtime.
static void storeFence() {
#if __has_builtin(__builtin_ia32_sfence)
  __builtin_ia32_sfence();
#elif __has_builtin(__builtin_arm_dmb)
  __builtin_arm_dmb(0x2); // oshst
#elif defined(__riscv)
  __asm__ __volatile__("fence ow, ow" ::: "memory");
#else
  __atomic_thread_fence(__ATOMIC_SEQ_CST);
#endif
}

/// Adapted from HSA-runtime.
static void memoryFence() {
#if __has_builtin(__builtin_ia32_mfence)
  __builtin_ia32_mfence();
#elif __has_builtin(__builtin_arm_dsb)
  __builtin_arm_dsb(0xf); // sy
#elif defined(__riscv)
  __asm__ __volatile__("fence iorw, iorw" ::: "memory");
#else
  __atomic_thread_fence(__ATOMIC_SEQ_CST);
#endif
}

/// Copied from HSA-runtime.
/// Force a WC flush on PCIe devices by doing a write and then read-back
static void flushPcieWc(void *Ptr, size_t Size) {
  storeFence();
  *((uint8_t *)Ptr + Size - 1) = *((uint8_t *)Ptr + Size - 1);
  memoryFence();
  auto Readback = *(reinterpret_cast<volatile uint8_t *>(Ptr) + Size - 1);
  (void)Readback;
}

/// Copy \p Size bytes from \p Src to the device memory at \p Dst.
/// \p OverXGMI tells whether the host reaches that memory over XGMI.
///
/// Without XGMI, this needs a WC flush. With XGMI, this isn't necessary
/// (according to the code in HSA-runtime).
static Error publishToDeviceMemory(void *Dst, const void *Src, size_t Size,
                                   bool OverXGMI) {
  std::memcpy(Dst, Src, Size);
  if (OverXGMI)
    return Plugin::success();
  hsa_utils::flushPcieWc(Dst, Size);
  return Plugin::success();
}

/// Iterate elements using an HSA iterate function. Do not use this function
/// directly but the specialized ones below instead.
template <typename ElemTy, typename IterFuncTy, typename CallbackTy>
static hsa_status_t iterate(IterFuncTy Func, CallbackTy Cb) {
  auto L = [](ElemTy Elem, void *Data) -> hsa_status_t {
    CallbackTy *Unwrapped = static_cast<CallbackTy *>(Data);
    return (*Unwrapped)(Elem);
  };
  return Func(L, static_cast<void *>(&Cb));
}

/// Iterate elements using an HSA iterate function passing a parameter. Do not
/// use this function directly but the specialized ones below instead.
template <typename ElemTy, typename IterFuncTy, typename IterFuncArgTy,
          typename CallbackTy>
static hsa_status_t iterate(IterFuncTy Func, IterFuncArgTy FuncArg,
                            CallbackTy Cb) {
  auto L = [](ElemTy Elem, void *Data) -> hsa_status_t {
    CallbackTy *Unwrapped = static_cast<CallbackTy *>(Data);
    return (*Unwrapped)(Elem);
  };
  return Func(FuncArg, L, static_cast<void *>(&Cb));
}

/// Iterate elements using an HSA iterate function passing a parameter. Do not
/// use this function directly but the specialized ones below instead.
template <typename Elem1Ty, typename Elem2Ty, typename IterFuncTy,
          typename IterFuncArgTy, typename CallbackTy>
[[maybe_unused]] static hsa_status_t
iterate(IterFuncTy Func, IterFuncArgTy FuncArg, CallbackTy Cb) {
  auto L = [](Elem1Ty Elem1, Elem2Ty Elem2, void *Data) -> hsa_status_t {
    CallbackTy *Unwrapped = static_cast<CallbackTy *>(Data);
    return (*Unwrapped)(Elem1, Elem2);
  };
  return Func(FuncArg, L, static_cast<void *>(&Cb));
}

/// Iterate agents.
template <typename CallbackTy> static Error iterateAgents(CallbackTy Callback) {
  hsa_status_t Status = iterate<hsa_agent_t>(hsa_iterate_agents, Callback);
  return Plugin::check(Status, "error in hsa_iterate_agents: %s");
}

/// Iterate ISAs of an agent.
template <typename CallbackTy>
static Error iterateAgentISAs(hsa_agent_t Agent, CallbackTy Cb) {
  hsa_status_t Status = iterate<hsa_isa_t>(hsa_agent_iterate_isas, Agent, Cb);
  return Plugin::check(Status, "error in hsa_agent_iterate_isas: %s");
}

/// Iterate memory pools of an agent.
template <typename CallbackTy>
static Error iterateAgentMemoryPools(hsa_agent_t Agent, CallbackTy Cb) {
  hsa_status_t Status = iterate<hsa_amd_memory_pool_t>(
      hsa_amd_agent_iterate_memory_pools, Agent, Cb);
  return Plugin::check(Status,
                       "error in hsa_amd_agent_iterate_memory_pools: %s");
}

/// Get the SDMA engine ID for a mem copy operation from src to dst agent.
/// Use a round-robin system to target all available engines. Check availability
/// using hsa_amd_memory_copy_engine_status.
/// Return nullopt if no available engine ID could be found.
static std::optional<hsa_amd_sdma_engine_id_t>
getSdmaEngineId(hsa_agent_t DstAgent, hsa_agent_t SrcAgent) {
  static std::atomic<uint32_t> NextIdx{0u};

  uint32_t AvailMask;
  if (hsa_amd_memory_copy_engine_status(DstAgent, SrcAgent, &AvailMask) ||
      !AvailMask)
    return std::nullopt;

  // Evaluates to 31, aka 0b11111.
  constexpr uint32_t ShiftWidthMask = sizeof(uint32_t) * 8 - 1;
  // The first engine ID we want to check is at NextIdx in the engine mask.
  uint32_t Start = NextIdx.load(std::memory_order_relaxed);
  // Inspect the bits starting from this ID.
  uint32_t FromStart = llvm::rotr(AvailMask, Start);
  // The first non-zero bit is the index of the next available engine. This
  // might be in the wrapped-around part of the mask, so we need to do a "modulo
  // bitwidth", aka & 0b11111. The addition of Start is necessary since we
  // rotated by that amount.
  uint32_t Idx = (Start + llvm::countr_zero(FromStart)) & ShiftWidthMask;
  // The next engine we want to check is at + 1, again wrapped by bitwidth.
  NextIdx.store((Idx + 1u) & ShiftWidthMask, std::memory_order_relaxed);

  // We now have the *index* of the engine. The ID is the bit set at that index.
  return static_cast<hsa_amd_sdma_engine_id_t>(1u << Idx);
}

/// Dispatches an asynchronous memory copy.
/// Enables different SDMA engines for the dispatch in a round-robin fashion.
static Error asyncMemCopy(bool UseMultipleSdmaEngines, void *Dst,
                          hsa_agent_t DstAgent, const void *Src,
                          hsa_agent_t SrcAgent, size_t Size,
                          uint32_t NumDepSignals,
                          const hsa_signal_t *DepSignals,
                          hsa_signal_t CompletionSignal) {
  if (!UseMultipleSdmaEngines) {
    hsa_status_t S =
        hsa_amd_memory_async_copy(Dst, DstAgent, Src, SrcAgent, Size,
                                  NumDepSignals, DepSignals, CompletionSignal);
    return Plugin::check(S, "error in hsa_amd_memory_async_copy: %s");
  }

  std::optional<hsa_amd_sdma_engine_id_t> SdmaEngine =
      getSdmaEngineId(DstAgent, SrcAgent);
  if (!SdmaEngine)
    return asyncMemCopy(/*UseMultipleSdmaEngines=*/false, Dst, DstAgent, Src,
                        SrcAgent, Size, NumDepSignals, DepSignals,
                        CompletionSignal);

  hsa_status_t S = hsa_amd_memory_async_copy_on_engine(
      Dst, DstAgent, Src, SrcAgent, Size, NumDepSignals, DepSignals,
      CompletionSignal, *SdmaEngine, /*force_copy_on_sdma=*/true);
  return Plugin::check(S, "error in hsa_amd_memory_async_copy_on_engine: %s");
}

static Error getTargetTripleAndFeatures(hsa_agent_t Agent,
                                        SmallVector<SmallString<32>> &Targets) {
  auto Err = hsa_utils::iterateAgentISAs(Agent, [&](hsa_isa_t ISA) {
    uint32_t Length;
    hsa_status_t Status;
    Status = hsa_isa_get_info_alt(ISA, HSA_ISA_INFO_NAME_LENGTH, &Length);
    if (Status != HSA_STATUS_SUCCESS)
      return Status;

    llvm::SmallVector<char> ISAName(Length);
    Status = hsa_isa_get_info_alt(ISA, HSA_ISA_INFO_NAME, ISAName.begin());
    if (Status != HSA_STATUS_SUCCESS)
      return Status;

    // The format returned here is a partially malformed triple, e.g.,
    // "amdgcn-amd-amdhsa--gfx90a", or
    // "amdgcn-amd-amdhsa--gfx90a:sramecc+:xnack-". The subtarget is in the
    // position that is supposed to be the object format. Reconstitute the valid
    // part of the triple for parsing, and take the appended subtarget name.
    SmallVector<StringRef, 5> Components;

    llvm::StringRef TripleLikeStr(ISAName.data(), ISAName.size() - 1);
    TripleLikeStr.split(Components, '-', /*MaxSplit=*/4);

    if (Components.size() == 5) {
      llvm::Triple TripleTarget(Components[0], Components[1], Components[2]);
      if (TripleTarget.isAMDGCN() && TripleTarget.getOS() == Triple::AMDHSA)
        Targets.emplace_back(Components[4]);
    }

    return HSA_STATUS_SUCCESS;
  });
  return Err;
}

/// Get the version of the AMD extension that the HSA runtime implements, if the
/// runtime reports it.
static std::optional<llvm::VersionTuple> getAMDExtensionVersion() {
  uint16_t Major, Minor;
  if (hsa_system_get_info(HSA_AMD_SYSTEM_INFO_EXT_VERSION_MAJOR, &Major) !=
          HSA_STATUS_SUCCESS ||
      hsa_system_get_info(HSA_AMD_SYSTEM_INFO_EXT_VERSION_MINOR, &Minor) !=
          HSA_STATUS_SUCCESS)
    return std::nullopt;
  return llvm::VersionTuple(Major, Minor);
}

/// Whether the host can write \p Size bytes at \p Ptr directly: they must lie
/// within an ordinary allocation of this device. On a large BAR system ROCr
/// maps all of those for the host, which HIP relies on as well. Memory mapped
/// through hsa_amd_vmem_* (such as the record-replay slab) is an example for
/// something the host cannot write to directly without explicit access grant
/// (needs to be checked explicitly; not yet handled here).
static bool isHostWritable(hsa_agent_t Agent, const void *Ptr, int64_t Size) {
  hsa_amd_pointer_info_t Info;
  Info.size = sizeof(hsa_amd_pointer_info_t);
  if (hsa_amd_pointer_info(Ptr, &Info, nullptr, nullptr, nullptr) !=
      HSA_STATUS_SUCCESS)
    return false;
  return Info.type == HSA_EXT_POINTER_TYPE_HSA &&
         Info.agentOwner.handle == Agent.handle &&
         utils::getPtrDiff(Ptr, Info.agentBaseAddress) + Size <=
             int64_t(Info.sizeInBytes);
}
} // namespace hsa_utils

/// Utility class representing generic resource references to AMDGPU resources.
template <typename ResourceTy>
struct AMDGPUResourceRef : public GenericDeviceResourceRef {
  /// The underlying handle type for resources.
  using HandleTy = ResourceTy *;

  /// Create an empty reference to an invalid resource.
  AMDGPUResourceRef() : Resource(nullptr) {}

  /// Create a reference to an existing resource.
  AMDGPUResourceRef(HandleTy Resource) : Resource(Resource) {}

  virtual ~AMDGPUResourceRef() {}

  /// Create a new resource and save the reference. The reference must be empty
  /// before calling to this function.
  Error create(GenericDeviceTy &Device) override;

  /// Destroy the referenced resource and invalidate the reference. The
  /// reference must be to a valid resource before calling to this function.
  Error destroy(GenericDeviceTy &Device) override {
    if (!Resource)
      return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                           "destroying an invalid resource");

    if (auto Err = Resource->deinit())
      return Err;

    delete Resource;

    Resource = nullptr;
    return Plugin::success();
  }

  /// Get the underlying resource handle.
  operator HandleTy() const { return Resource; }

private:
  /// The handle to the actual resource.
  HandleTy Resource;
};

/// Class holding an HSA memory pool.
struct AMDGPUMemoryPoolTy {
  /// Create a memory pool from an HSA memory pool.
  AMDGPUMemoryPoolTy(hsa_amd_memory_pool_t MemoryPool)
      : MemoryPool(MemoryPool), GlobalFlags(0) {}

  /// Initialize the memory pool retrieving its properties.
  Error init() {
    if (auto Err = getAttr(HSA_AMD_MEMORY_POOL_INFO_SEGMENT, Segment))
      return Err;

    if (auto Err = getAttr(HSA_AMD_MEMORY_POOL_INFO_GLOBAL_FLAGS, GlobalFlags))
      return Err;

    if (auto Err = getAttr(HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALIGNMENT,
                           PoolAllocationAlignment))
      return Err;

    return getAttr(HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_GRANULE, Granule);
  }

  /// Getter of the HSA memory pool.
  hsa_amd_memory_pool_t get() const { return MemoryPool; }

  /// Indicate the segment which belongs to.
  bool isGlobal() const { return (Segment == HSA_AMD_SEGMENT_GLOBAL); }
  bool isReadOnly() const { return (Segment == HSA_AMD_SEGMENT_READONLY); }
  bool isPrivate() const { return (Segment == HSA_AMD_SEGMENT_PRIVATE); }
  bool isGroup() const { return (Segment == HSA_AMD_SEGMENT_GROUP); }

  /// Indicate if it is fine-grained memory. Valid only for global.
  bool isFineGrained() const {
    assert(isGlobal() && "Not global memory");
    return (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED);
  }

  /// Indicate if it is coarse-grained memory. Valid only for global.
  bool isCoarseGrained() const {
    assert(isGlobal() && "Not global memory");
    return (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_COARSE_GRAINED);
  }

  /// Indicate if it supports storing kernel arguments. Valid only for global.
  bool supportsKernelArgs() const {
    assert(isGlobal() && "Not global memory");
    return (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_KERNARG_INIT);
  }

  /// Get the allocation granularity of the pool.
  size_t getGranule() const { return Granule; }

  /// Get the allocation alignment of the pool.
  size_t getAlignment() const { return PoolAllocationAlignment; }

  /// Allocate memory on the memory pool.
  Error allocate(size_t Size, void **PtrStorage, size_t Alignment) {
    // A non-zero value passed as the Alignment indicates that the user expects
    // the allocation to have a specific alignment. However, the HSA API does
    // not allow users to define alignment. Therefore, the passed alignment is
    // compared with the alignment of the memory allocated using the given pool.
    // If the default alignment is greater than or equal to the alignment
    // requested by the user, it would still meet the user's requirements.
    if (Alignment > PoolAllocationAlignment) {
      return Plugin::error(ErrorCode::UNSUPPORTED,
                           "requested alignment (%lu) larger than maximum "
                           "supported pool alignment (%lu)",
                           Alignment, PoolAllocationAlignment);
    }

    hsa_status_t Status =
        hsa_amd_memory_pool_allocate(MemoryPool, Size, 0, PtrStorage);

    if (Alignment > 0 && !isAddrAligned(Align(Alignment), *PtrStorage)) {
      if (auto FreeErr = deallocate(*PtrStorage)) {
        return Plugin::error(ErrorCode::UNKNOWN,
                             "Failure in deallcation of the incorrectly "
                             "aligned pointer; requested alignemnt: %lu",
                             Alignment);
      }

      return Plugin::error(ErrorCode::UNSUPPORTED,
                           "unsupported alignment size");
    }

    return Plugin::check(Status, "error in hsa_amd_memory_pool_allocate: %s");
  }

  /// Return memory to the memory pool.
  Error deallocate(void *Ptr) {
    hsa_status_t Status = hsa_amd_memory_pool_free(Ptr);
    return Plugin::check(Status, "error in hsa_amd_memory_pool_free: %s");
  }

  /// Returns if the \p Agent can access the memory pool.
  bool canAccess(hsa_agent_t Agent) {
    hsa_amd_memory_pool_access_t Access;
    if (hsa_amd_agent_memory_pool_get_info(
            Agent, MemoryPool, HSA_AMD_AGENT_MEMORY_POOL_INFO_ACCESS, &Access))
      return false;
    return Access != HSA_AMD_MEMORY_POOL_ACCESS_NEVER_ALLOWED;
  }

  /// Returns if the \p Agent reaches the memory pool over XGMI only, i.e.,
  /// without any PCIe (or other) hop along the way.
  bool isLinkedOverXGMI(hsa_agent_t Agent) {
    uint32_t NumHops = 0;
    if (hsa_amd_agent_memory_pool_get_info(
            Agent, MemoryPool, HSA_AMD_AGENT_MEMORY_POOL_INFO_NUM_LINK_HOPS,
            &NumHops) ||
        NumHops == 0)
      return false;
    llvm::SmallVector<hsa_amd_memory_pool_link_info_t> Links(NumHops);
    if (hsa_amd_agent_memory_pool_get_info(
            Agent, MemoryPool, HSA_AMD_AGENT_MEMORY_POOL_INFO_LINK_INFO,
            Links.data()))
      return false;
    return llvm::all_of(Links, [](const hsa_amd_memory_pool_link_info_t &L) {
      return L.link_type == HSA_AMD_LINK_INFO_TYPE_XGMI;
    });
  }

  /// Allow the device to access a specific allocation.
  Error enableAccess(void *Ptr, int64_t Size,
                     const llvm::SmallVector<hsa_agent_t> &Agents) const {
#ifdef OMPTARGET_DEBUG
    for (hsa_agent_t Agent : Agents) {
      hsa_amd_memory_pool_access_t Access;
      if (auto Err =
              getAttr(Agent, HSA_AMD_AGENT_MEMORY_POOL_INFO_ACCESS, Access))
        return Err;

      // The agent is not allowed to access the memory pool in any case. Do not
      // continue because otherwise it result in undefined behavior.
      if (Access == HSA_AMD_MEMORY_POOL_ACCESS_NEVER_ALLOWED)
        return Plugin::error(ErrorCode::INVALID_VALUE,
                             "an agent is not allowed to access a memory pool");
    }
#endif

    // We can access but it is disabled by default. Enable the access then.
    hsa_status_t Status =
        hsa_amd_agents_allow_access(Agents.size(), Agents.data(), nullptr, Ptr);
    return Plugin::check(Status, "error in hsa_amd_agents_allow_access: %s");
  }

  /// Get attribute from the memory pool.
  template <typename Ty>
  Error getAttr(hsa_amd_memory_pool_info_t Kind, Ty &Value) const {
    hsa_status_t Status;
    Status = hsa_amd_memory_pool_get_info(MemoryPool, Kind, &Value);
    return Plugin::check(Status, "error in hsa_amd_memory_pool_get_info: %s");
  }

  template <typename Ty>
  hsa_status_t getAttrRaw(hsa_amd_memory_pool_info_t Kind, Ty &Value) const {
    return hsa_amd_memory_pool_get_info(MemoryPool, Kind, &Value);
  }

  /// Get attribute from the memory pool relating to an agent.
  template <typename Ty>
  Error getAttr(hsa_agent_t Agent, hsa_amd_agent_memory_pool_info_t Kind,
                Ty &Value) const {
    hsa_status_t Status;
    Status =
        hsa_amd_agent_memory_pool_get_info(Agent, MemoryPool, Kind, &Value);
    return Plugin::check(Status,
                         "error in hsa_amd_agent_memory_pool_get_info: %s");
  }

private:
  /// The HSA memory pool.
  hsa_amd_memory_pool_t MemoryPool;

  /// The segment where the memory pool belongs to.
  hsa_amd_segment_t Segment;

  /// The global flags of memory pool. Only valid if the memory pool belongs to
  /// the global segment.
  uint32_t GlobalFlags;

  /// The page size in this memory pool.
  size_t Granule;

  /// The alignment of the buffers allocated by
  /// hsa_amd_memory_pool_allocate(...). This attribute is defined only if
  /// HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALLOWED is set to true. Since
  /// hsa_amd_memory_pool_allocate is the only memory-allocating function that
  /// is used by the memory pool, HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALLOWED
  /// should always be set.
  size_t PoolAllocationAlignment;
};

/// Class that implements a memory manager that gets memory from a specific
/// memory pool.
struct AMDGPUMemoryManagerTy : public DeviceAllocatorTy {

  /// Create an empty memory manager.
  AMDGPUMemoryManagerTy(AMDGPUPluginTy &Plugin)
      : Plugin(Plugin), MemoryPool(nullptr), MemoryManager(nullptr) {}

  /// Initialize the memory manager from a memory pool.
  Error init(AMDGPUMemoryPoolTy &MemoryPool) {
    const uint32_t Threshold = 1 << 30;
    this->MemoryManager = new MemoryManagerTy(*this, Threshold);
    this->MemoryPool = &MemoryPool;
    return Plugin::success();
  }

  /// Deinitialize the memory manager and free its allocations.
  Error deinit() {
    assert(MemoryManager && "Invalid memory manager");

    // Delete and invalidate the memory manager. At this point, the memory
    // manager will deallocate all its allocations.
    delete MemoryManager;
    MemoryManager = nullptr;

    return Plugin::success();
  }

  /// Reuse or allocate memory through the memory manager.
  Error allocate(size_t Size, void **PtrStorage) {
    assert(MemoryManager && "Invalid memory manager");
    assert(PtrStorage && "Invalid pointer storage");

    auto PtrStorageOrErr =
        MemoryManager->allocate(Size, nullptr, /*Alignment=*/0);
    if (!PtrStorageOrErr)
      return PtrStorageOrErr.takeError();

    *PtrStorage = *PtrStorageOrErr;
    if (Size && *PtrStorage == nullptr)
      return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
                           "failure to allocate from AMDGPU memory manager");

    return Plugin::success();
  }

  /// Release an allocation to be reused.
  Error deallocate(void *Ptr) {
    if (MemoryManager->free(Ptr))
      return Plugin::error(ErrorCode::UNKNOWN,
                           "failure to deallocate from AMDGPU memory manager");

    return Plugin::success();
  }

private:
  /// Allocation callback that will be called once the memory manager does not
  /// have more previously allocated buffers.
  Expected<void *> allocate(size_t Size, void *HstPtr, TargetAllocTy Kind,
                            size_t Alignment) override;

  /// Deallocation callback that will be called by the memory manager.
  Error free(void *TgtPtr, TargetAllocTy Kind) override {
    return MemoryPool->deallocate(TgtPtr);
  }

  /// The underlying plugin that owns this memory manager.
  AMDGPUPluginTy &Plugin;

  /// The memory pool used to allocate memory.
  AMDGPUMemoryPoolTy *MemoryPool;

  /// Reference to the actual memory manager.
  MemoryManagerTy *MemoryManager;
};

/// Class implementing the AMDGPU device images' properties.
struct AMDGPUDeviceImageTy : public DeviceImageTy {
  /// Create the AMDGPU image with the id and the target image pointer.
  AMDGPUDeviceImageTy(int32_t ImageId, GenericDeviceTy &Device,
                      std::unique_ptr<MemoryBuffer> &&TgtImage)
      : DeviceImageTy(ImageId, Device, std::move(TgtImage)) {}

  /// Prepare and load the executable corresponding to the image.
  Error loadExecutable(const AMDGPUDeviceTy &Device);

  /// Unload the executable.
  Error unloadExecutable() {
    hsa_status_t Status = hsa_executable_destroy(Executable);
    return Plugin::check(Status, "error in hsa_executable_destroy: %s");
  }

  /// Get the executable.
  hsa_executable_t getExecutable() const { return Executable; }

  /// Get to Code Object Version of the ELF
  uint16_t getELFABIVersion() const { return ELFABIVersion; }

  /// Find an HSA device symbol by its name on the executable.
  Expected<hsa_executable_symbol_t>
  findDeviceSymbol(GenericDeviceTy &Device, StringRef SymbolName) const;

  /// Get additional info for kernel, e.g., register spill counts
  Expected<offloading::amdgpu::AMDGPUKernelMetaData>
  getKernelInfo(StringRef Identifier) const {
    auto It = KernelInfoMap.find(Identifier);

    if (It == KernelInfoMap.end())
      return Plugin::error(ErrorCode::INVALID_BINARY,
                           "could not find metadata for kernel %s",
                           Identifier.str().c_str());

    return It->second;
  }

  /// Return the maximum wavefront size across all known kernels in this image.
  uint32_t getMaxWavefrontSize() const {
    uint32_t Max = 0;
    for (const auto &[Name, Info] : KernelInfoMap)
      if (Info.WavefrontSize !=
          offloading::amdgpu::AMDGPUKernelMetaData::KInvalidValue)
        Max = std::max(Max, Info.WavefrontSize);
    return Max;
  }

private:
  /// The executable loaded on the agent.
  hsa_executable_t Executable;
  StringMap<offloading::amdgpu::AMDGPUKernelMetaData> KernelInfoMap;
  uint16_t ELFABIVersion;
};

/// Class implementing the AMDGPU kernel functionalities which derives from the
/// generic kernel class.
struct AMDGPUKernelTy : public GenericKernelTy {
  /// Create an AMDGPU kernel with a name and an execution mode.
  AMDGPUKernelTy(StringRef Name) : GenericKernelTy(Name) {}

  /// Initialize the AMDGPU kernel.
  Error initImpl(GenericDeviceTy &Device, DeviceImageTy &Image) override {
    AMDGPUDeviceImageTy &AMDImage = static_cast<AMDGPUDeviceImageTy &>(Image);

    // Kernel symbols have a ".kd" suffix.
    std::string KernelName(getName());
    KernelName += ".kd";

    // Find the symbol on the device executable.
    auto SymbolOrErr = AMDImage.findDeviceSymbol(Device, KernelName);
    if (!SymbolOrErr)
      return SymbolOrErr.takeError();

    hsa_executable_symbol_t Symbol = *SymbolOrErr;
    hsa_symbol_kind_t SymbolType;
    hsa_status_t Status;

    // Retrieve different properties of the kernel symbol.
    std::pair<hsa_executable_symbol_info_t, void *> RequiredInfos[] = {
        {HSA_EXECUTABLE_SYMBOL_INFO_TYPE, &SymbolType},
        {HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_OBJECT, &KernelObject},
        {HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE, &ArgsSize},
        {HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_GROUP_SEGMENT_SIZE, &GroupSize},
        {HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_DYNAMIC_CALLSTACK, &DynamicStack},
        {HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_PRIVATE_SEGMENT_SIZE, &PrivateSize}};

    for (auto &Info : RequiredInfos) {
      Status = hsa_executable_symbol_get_info(Symbol, Info.first, Info.second);
      if (auto Err = Plugin::check(
              Status, "error in hsa_executable_symbol_get_info: %s"))
        return Err;
    }

    // Set the static block memory size required by the kernel.
    StaticBlockMemSize = GroupSize;

    // Make sure it is a kernel symbol.
    if (SymbolType != HSA_SYMBOL_KIND_KERNEL)
      return Plugin::error(ErrorCode::INVALID_BINARY,
                           "symbol %s is not a kernel function");

    // TODO: Read the kernel descriptor for the max threads per block. May be
    // read from the image.

    ImplicitArgsSize =
        hsa_utils::getImplicitArgsSize(AMDImage.getELFABIVersion());
    ODBG(OLDT_Module) << "ELFABIVersion: " << AMDImage.getELFABIVersion();

    // Get additional kernel info read from image
    auto KernelInfoOrErr = AMDImage.getKernelInfo(getName());
    if (!KernelInfoOrErr)
      return KernelInfoOrErr.takeError();
    KernelInfo = std::move(*KernelInfoOrErr);

    return Plugin::success();
  }

  /// Launch the AMDGPU kernel function.
  Error launchImpl(GenericDeviceTy &GenericDevice, uint32_t NumThreads[3],
                   uint32_t NumBlocks[3], uint32_t DynBlockMemSize,
                   KernelLaunchArgsTy &LaunchArgs,
                   AsyncInfoWrapperTy &AsyncInfoWrapper,
                   GenericProfilerTy *ProfilerPtr) const override;

  /// Return maximum block size for maximum occupancy
  ///
  /// TODO: This needs to be implemented for amdgpu
  Expected<uint64_t> maxGroupSize(GenericDeviceTy &GenericDevice,
                                  uint64_t DynamicMemSize) const override {
    return Plugin::error(
        ErrorCode::UNSUPPORTED,
        "occupancy calculations for AMDGPU are not yet implemented");
  }

  /// Print more elaborate kernel launch info for AMDGPU
  Error printLaunchInfoDetails(GenericDeviceTy &GenericDevice,
                               const KernelLaunchArgsTy &LaunchArgs,
                               uint32_t NumThreads[3],
                               uint32_t NumBlocks[3]) const override;

  /// Get group and private segment kernel size.
  uint32_t getGroupSize() const { return GroupSize; }
  uint32_t getPrivateSize() const { return PrivateSize; }

  /// Get the HSA kernel object representing the kernel function.
  uint64_t getKernelObject() const { return KernelObject; }

  /// Get the size of implicitargs based on the code object version.
  uint32_t getImplicitArgsSize() const { return ImplicitArgsSize; }

  /// Indicates whether or not we need to set up our own private segment size.
  bool usesDynamicStack() const { return DynamicStack; }

private:
  /// The kernel object to execute.
  uint64_t KernelObject;

  /// The args, group and private segments sizes required by a kernel instance.
  uint32_t ArgsSize;
  uint32_t GroupSize;
  uint32_t PrivateSize;
  bool DynamicStack;

  /// The size of implicit kernel arguments.
  uint32_t ImplicitArgsSize;

  /// Additional Info for the AMD GPU Kernel
  offloading::amdgpu::AMDGPUKernelMetaData KernelInfo;
};

/// Class representing an HSA signal. Signals are used to define dependencies
/// between asynchronous operations: kernel launches and memory transfers.
struct AMDGPUSignalTy {
  /// Create an empty signal.
  AMDGPUSignalTy() : HSASignal({0}), UseCount() {}
  AMDGPUSignalTy(AMDGPUDeviceTy &Device) : HSASignal({0}), UseCount() {}

  /// Initialize the signal with an initial value.
  Error init(uint32_t InitialValue = 1) {
    hsa_status_t Status =
        hsa_amd_signal_create(InitialValue, 0, nullptr, 0, &HSASignal);
    return Plugin::check(Status, "error in hsa_signal_create: %s");
  }

  /// Deinitialize the signal.
  Error deinit() {
    hsa_status_t Status = hsa_signal_destroy(HSASignal);
    return Plugin::check(Status, "error in hsa_signal_destroy: %s");
  }

  /// Wait until the signal gets a zero value.
  Error wait(const uint64_t ActiveTimeout = 0,
             GenericDeviceTy *Device = nullptr) const {
    if (ActiveTimeout) {
      hsa_signal_value_t Got = 1;
      Got = hsa_signal_wait_scacquire(HSASignal, HSA_SIGNAL_CONDITION_EQ, 0,
                                      ActiveTimeout, HSA_WAIT_STATE_ACTIVE);
      if (Got == 0)
        return Plugin::success();
    }

    // If there is an RPC device attached to this stream we run it as a server.
    uint64_t Timeout = UINT64_MAX;
    auto WaitState = HSA_WAIT_STATE_BLOCKED;
    while (hsa_signal_wait_scacquire(HSASignal, HSA_SIGNAL_CONDITION_EQ, 0,
                                     Timeout, WaitState) != 0)
      ;
    return Plugin::success();
  }

  /// Load the value on the signal.
  hsa_signal_value_t load() const {
    return hsa_signal_load_scacquire(HSASignal);
  }

  /// Signal decrementing by one.
  void signal() {
    assert(load() > 0 && "Invalid signal value");
    hsa_signal_subtract_screlease(HSASignal, 1);
  }

  /// Reset the signal value before reusing the signal. Do not call this
  /// function if the signal is being currently used by any watcher, such as a
  /// plugin thread or the HSA runtime.
  void reset() { hsa_signal_store_screlease(HSASignal, 1); }

  /// Increase the number of concurrent uses by \p Amount.
  void increaseUseCount(uint32_t Amount = 1) { UseCount.increase(Amount); }

  /// Decrease the number of concurrent uses by \p Amount and return whether it
  /// became zero.
  bool decreaseUseCount(uint32_t Amount = 1) {
    return UseCount.decrease(Amount);
  }

  hsa_signal_t get() const { return HSASignal; }

private:
  /// The underlying HSA signal.
  hsa_signal_t HSASignal;

  /// Reference counter for tracking the concurrent use count. This is mainly
  /// used for knowing how many streams are using the signal.
  RefCountTy<uint32_t> UseCount;
};

/// Classes for holding AMDGPU signals and managing signals.
using AMDGPUSignalRef = AMDGPUResourceRef<AMDGPUSignalTy>;
using AMDGPUSignalManagerTy = GenericDeviceResourceManagerTy<AMDGPUSignalRef>;

/// Class holding an HSA queue to submit kernel and barrier packets.
struct AMDGPUQueueTy {
  /// Create an empty queue.
  AMDGPUQueueTy() : Queue(nullptr), Mutex(), NumUsers(0) {}

  /// Lazily initialize a new queue belonging to a specific agent.
  Error init(GenericDeviceTy &Device, hsa_agent_t Agent, int32_t QueueSize) {
    if (Queue)
      return Plugin::success();

    hsa_status_t Status =
        hsa_queue_create(Agent, QueueSize, HSA_QUEUE_TYPE_MULTI, callbackError,
                         &Device, UINT32_MAX, UINT32_MAX, &Queue);
    if (auto Err = Plugin::check(Status, "error in hsa_queue_create: %s"))
      return Err;

    // Enable queue profiling from creation time onward, as HIP/ROCclr does.
    // Elapsed-time queries rely on queue-level hardware profiling support to
    // retrieve packet timing.
    Status = hsa_amd_profiling_set_profiler_enabled(Queue, 1);
    if (auto Err = Plugin::check(
            Status, "error in hsa_amd_profiling_set_profiler_enabled: %s"))
      return Err;

    return Plugin::success();
  }

  /// Deinitialize the queue and destroy its resources.
  Error deinit() {
    std::lock_guard<std::mutex> Lock(Mutex);
    if (!Queue)
      return Plugin::success();
    hsa_status_t Status = hsa_queue_destroy(Queue);
    return Plugin::check(Status, "error in hsa_queue_destroy: %s");
  }

  /// Returns the number of streams, this queue is currently assigned to.
  bool getUserCount() const { return NumUsers; }

  /// Returns if the underlying HSA queue is initialized.
  bool isInitialized() { return Queue != nullptr; }

  /// Decrement user count of the queue object.
  void removeUser() { --NumUsers; }

  /// Increase user count of the queue object.
  void addUser() { ++NumUsers; }

  /// Push a kernel launch to the queue. The kernel launch requires an output
  /// signal and can define an optional input signal (nullptr if none).
  Error pushKernelLaunch(const AMDGPUKernelTy &Kernel, void *KernelArgs,
                         uint32_t NumThreads[3], uint32_t NumBlocks[3],
                         uint32_t GroupSize, uint64_t StackSize,
                         AMDGPUSignalTy *OutputSignal,
                         AMDGPUSignalTy *InputSignal) {
    assert(OutputSignal && "Invalid kernel output signal");

    // Lock the queue during the packet publishing process. Notice this blocks
    // the addition of other packets to the queue. The following piece of code
    // should be lightweight; do not block the thread, allocate memory, etc.
    std::lock_guard<std::mutex> Lock(Mutex);
    assert(Queue && "Interacted with a non-initialized queue!");

    // Add a barrier packet before the kernel packet in case there is a pending
    // preceding operation. The barrier packet will delay the processing of
    // subsequent queue's packets until the barrier input signal are satisfied.
    // No need output signal needed because the dependency is already guaranteed
    // by the queue barrier itself.
    if (InputSignal && InputSignal->load())
      if (auto Err = pushBarrierImpl(nullptr, InputSignal))
        return Err;

    // Now prepare the kernel packet.
    uint64_t PacketId;
    hsa_kernel_dispatch_packet_t *Packet = acquirePacket(PacketId);
    assert(Packet && "Invalid packet");

    // The first 32 bits of the packet are written after the other fields
    uint16_t Dims = NumBlocks[2] * NumThreads[2] > 1
                        ? 3
                        : 1 + (NumBlocks[1] * NumThreads[1] != 1);
    uint16_t Setup = UINT16_C(Dims)
                     << HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS;
    Packet->workgroup_size_x = NumThreads[0];
    Packet->workgroup_size_y = NumThreads[1];
    Packet->workgroup_size_z = NumThreads[2];
    Packet->reserved0 = 0;
    Packet->grid_size_x = NumBlocks[0] * NumThreads[0];
    Packet->grid_size_y = NumBlocks[1] * NumThreads[1];
    Packet->grid_size_z = NumBlocks[2] * NumThreads[2];
    Packet->private_segment_size =
        Kernel.usesDynamicStack() ? StackSize : Kernel.getPrivateSize();
    Packet->group_segment_size = GroupSize;
    Packet->kernel_object = Kernel.getKernelObject();
    Packet->kernarg_address = KernelArgs;
    Packet->reserved2 = 0;
    Packet->completion_signal = OutputSignal->get();

    // Publish the packet. Do not modify the packet after this point.
    publishKernelPacket(PacketId, Setup, Packet);

    return Plugin::success();
  }

  /// Push a barrier packet that will wait up to two input signals. All signals
  /// are optional (nullptr if none).
  Error pushBarrier(AMDGPUSignalTy *OutputSignal,
                    const AMDGPUSignalTy *InputSignal1,
                    const AMDGPUSignalTy *InputSignal2) {
    // Lock the queue during the packet publishing process.
    std::lock_guard<std::mutex> Lock(Mutex);
    assert(Queue && "Interacted with a non-initialized queue!");

    // Push the barrier with the lock acquired.
    return pushBarrierImpl(OutputSignal, InputSignal1, InputSignal2);
  }

private:
  /// Push a barrier packet that will wait up to two input signals. Assumes the
  /// the queue lock is acquired.
  Error pushBarrierImpl(AMDGPUSignalTy *OutputSignal,
                        const AMDGPUSignalTy *InputSignal1,
                        const AMDGPUSignalTy *InputSignal2 = nullptr) {
    // Add a queue barrier waiting on both the other stream's operation and the
    // last operation on the current stream (if any).
    uint64_t PacketId;
    hsa_barrier_and_packet_t *Packet =
        (hsa_barrier_and_packet_t *)acquirePacket(PacketId);
    assert(Packet && "Invalid packet");

    Packet->reserved0 = 0;
    Packet->reserved1 = 0;
    Packet->dep_signal[0] = {0};
    Packet->dep_signal[1] = {0};
    Packet->dep_signal[2] = {0};
    Packet->dep_signal[3] = {0};
    Packet->dep_signal[4] = {0};
    Packet->reserved2 = 0;
    Packet->completion_signal = {0};

    // Set input and output dependencies if needed.
    if (OutputSignal)
      Packet->completion_signal = OutputSignal->get();
    if (InputSignal1)
      Packet->dep_signal[0] = InputSignal1->get();
    if (InputSignal2)
      Packet->dep_signal[1] = InputSignal2->get();

    // Publish the packet. Do not modify the packet after this point.
    publishBarrierPacket(PacketId, Packet);

    return Plugin::success();
  }

  /// Acquire a packet from the queue. This call may block the thread if there
  /// is no space in the underlying HSA queue. It may need to wait until the HSA
  /// runtime processes some packets. Assumes the queue lock is acquired.
  hsa_kernel_dispatch_packet_t *acquirePacket(uint64_t &PacketId) {
    // Increase the queue index with relaxed memory order. Notice this will need
    // another subsequent atomic operation with acquire order.
    PacketId = hsa_queue_add_write_index_relaxed(Queue, 1);

    // Wait for the package to be available. Notice the atomic operation uses
    // the acquire memory order.
    while (PacketId - hsa_queue_load_read_index_scacquire(Queue) >= Queue->size)
      ;

    // Return the packet reference.
    const uint32_t Mask = Queue->size - 1; // The size is a power of 2.
    return (hsa_kernel_dispatch_packet_t *)Queue->base_address +
           (PacketId & Mask);
  }

  /// Publish the kernel packet so that the HSA runtime can start processing
  /// the kernel launch. Do not modify the packet once this function is called.
  /// Assumes the queue lock is acquired.
  void publishKernelPacket(uint64_t PacketId, uint16_t Setup,
                           hsa_kernel_dispatch_packet_t *Packet) {
    uint32_t *PacketPtr = reinterpret_cast<uint32_t *>(Packet);

    uint16_t Header = HSA_PACKET_TYPE_KERNEL_DISPATCH << HSA_PACKET_HEADER_TYPE;
    Header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_ACQUIRE_FENCE_SCOPE;
    Header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_RELEASE_FENCE_SCOPE;

    // Publish the packet. Do not modify the package after this point.
    uint32_t HeaderWord = Header | (Setup << 16u);
    __atomic_store_n(PacketPtr, HeaderWord, __ATOMIC_RELEASE);

    // Signal the doorbell about the published packet.
    hsa_signal_store_relaxed(Queue->doorbell_signal, PacketId);
  }

  /// Publish the barrier packet so that the HSA runtime can start processing
  /// the barrier. Next packets in the queue will not be processed until all
  /// barrier dependencies (signals) are satisfied. Assumes the queue is locked
  void publishBarrierPacket(uint64_t PacketId,
                            hsa_barrier_and_packet_t *Packet) {
    uint32_t *PacketPtr = reinterpret_cast<uint32_t *>(Packet);
    uint16_t Setup = 0;
    uint16_t Header = HSA_PACKET_TYPE_BARRIER_AND << HSA_PACKET_HEADER_TYPE;
    Header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_ACQUIRE_FENCE_SCOPE;
    Header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_RELEASE_FENCE_SCOPE;

    // Publish the packet. Do not modify the package after this point.
    uint32_t HeaderWord = Header | (Setup << 16u);
    __atomic_store_n(PacketPtr, HeaderWord, __ATOMIC_RELEASE);

    // Signal the doorbell about the published packet.
    hsa_signal_store_relaxed(Queue->doorbell_signal, PacketId);
  }

  /// Callback that will be called when an error is detected on the HSA queue.
  static void callbackError(hsa_status_t Status, hsa_queue_t *Source,
                            void *Data);

  /// The HSA queue.
  hsa_queue_t *Queue;

  /// Mutex to protect the acquiring and publishing of packets. For the moment,
  /// we need this mutex to prevent publishing packets that are not ready to be
  /// published in a multi-thread scenario. Without a queue lock, a thread T1
  /// could acquire packet P and thread T2 acquire packet P+1. Thread T2 could
  /// publish its packet P+1 (signaling the queue's doorbell) before packet P
  /// from T1 is ready to be processed. That scenario should be invalid. Thus,
  /// we use the following mutex to make packet acquiring and publishing atomic.
  /// TODO: There are other more advanced approaches to avoid this mutex using
  /// atomic operations. We can further investigate it if this is a bottleneck.
  std::mutex Mutex;

  /// The number of streams, this queue is currently assigned to. A queue is
  /// considered idle when this is zero, otherwise: busy.
  uint32_t NumUsers;
};

/// Struct that implements a stream of asynchronous operations for AMDGPU
/// devices. This class relies on signals to implement streams and define the
/// dependencies between asynchronous operations.
struct AMDGPUStreamTy {
public:
  /// Function pointer type for `pushHostCallback`
  using HostFnType = void (*)(void *);

private:
  /// Utility struct holding arguments for async H2H memory copies.
  struct MemcpyArgsTy {
    void *Dst;
    const void *Src;
    size_t Size;
    size_t NumTimes;
  };

  /// Utility struct holding arguments for freeing buffers to memory managers.
  struct ReleaseBufferArgsTy {
    void *Buffer;
    AMDGPUMemoryManagerTy *MemoryManager;
  };

  /// Utility struct holding arguments for releasing signals to signal managers.
  struct ReleaseSignalArgsTy {
    AMDGPUSignalTy *Signal;
    AMDGPUSignalManagerTy *SignalManager;
  };

  using AMDGPUStreamCallbackTy = Error(void *Data);

  /// The stream is composed of N stream's slots. The struct below represents
  /// the fields of each slot. Each slot has a signal and an optional action
  /// function. When appending an HSA asynchronous operation to the stream, one
  /// slot is consumed and used to store the operation's information. The
  /// operation's output signal is set to the consumed slot's signal. If there
  /// is a previous asynchronous operation on the previous slot, the HSA async
  /// operation's input signal is set to the signal of the previous slot. This
  /// way, we obtain a chain of dependent async operations. The action is a
  /// function that will be executed eventually after the operation is
  /// completed, e.g., for releasing a buffer.
  struct StreamSlotTy {
    /// The output signal of the stream operation. May be used by the subsequent
    /// operation as input signal.
    AMDGPUSignalTy *Signal;

    /// The actions that must be performed after the operation's completion. Set
    /// to nullptr when there is no action to perform.
    llvm::SmallVector<AMDGPUStreamCallbackTy *> Callbacks;

    /// Space for the action's arguments. A pointer to these arguments is passed
    /// to the action function. Notice the space of arguments is limited.
    union ActionArgsTy {
      MemcpyArgsTy MemcpyArgs;
      ReleaseBufferArgsTy ReleaseBufferArgs;
      ReleaseSignalArgsTy ReleaseSignalArgs;
      ProfilingInfoTy ProfilerArgs;
      void *CallbackArgs;
    };

    llvm::SmallVector<ActionArgsTy> ActionArgs;

    /// Create an empty slot.
    StreamSlotTy() : Signal(nullptr), Callbacks({}), ActionArgs({}) {}

    /// Schedule a host memory copy action on the slot.
    ///
    /// Num times will repeat the copy that many times, sequentually in the dest
    /// buffer.
    Error schedHostMemoryCopy(void *Dst, const void *Src, size_t Size,
                              size_t NumTimes = 1) {
      Callbacks.emplace_back(memcpyAction);
      ActionArgs.emplace_back().MemcpyArgs =
          MemcpyArgsTy{Dst, Src, Size, NumTimes};
      return Plugin::success();
    }

    /// Schedule a release buffer action on the slot.
    Error schedReleaseBuffer(void *Buffer, AMDGPUMemoryManagerTy &Manager) {
      Callbacks.emplace_back(releaseBufferAction);
      ActionArgs.emplace_back().ReleaseBufferArgs =
          ReleaseBufferArgsTy{Buffer, &Manager};
      return Plugin::success();
    }

    /// Schedule a signal release action on the slot.
    Error schedReleaseSignal(AMDGPUSignalTy *SignalToRelease,
                             AMDGPUSignalManagerTy *SignalManager) {
      Callbacks.emplace_back(releaseSignalAction);
      ActionArgs.emplace_back().ReleaseSignalArgs =
          ReleaseSignalArgsTy{SignalToRelease, SignalManager};
      return Plugin::success();
    }

    /// Schedule kernel timing measurement via the profiler on the slot.
    Error schedProfilerKernelTiming(GenericProfilerTy *ProfilerPtr,
                                    hsa_agent_t Agent,
                                    AMDGPUSignalTy *OutputSignal,
                                    double TicksToTime,
                                    void *ProfilerSpecificData) {
      Callbacks.emplace_back(timeKernelInNsAsync);
      ActionArgs.emplace_back().ProfilerArgs = ProfilingInfoTy{
          ProfilerPtr, Agent, OutputSignal, TicksToTime, ProfilerSpecificData};
      return Plugin::success();
    }

    /// Schedule data transfer timing via the profiler on the slot.
    Error schedProfilerDataTransferTiming(GenericProfilerTy *ProfilerPtr,
                                          hsa_agent_t Agent,
                                          AMDGPUSignalTy *OutputSignal,
                                          double TicksToTime,
                                          void *ProfilerSpecificData) {
      Callbacks.emplace_back(timeDataTransferInNsAsync);
      ActionArgs.emplace_back().ProfilerArgs = ProfilingInfoTy{
          ProfilerPtr, Agent, OutputSignal, TicksToTime, ProfilerSpecificData};
      return Plugin::success();
    }

    /// Register a callback to be called on completion.
    Error schedCallback(AMDGPUStreamCallbackTy *Func, void *Data) {
      Callbacks.emplace_back(Func);
      ActionArgs.emplace_back().CallbackArgs = Data;

      return Plugin::success();
    }

    // Perform the action if needed.
    Error performAction() {
      if (Callbacks.empty())
        return Plugin::success();

      assert(Callbacks.size() == ActionArgs.size() && "Size mismatch");
      for (auto [Callback, ActionArg] : llvm::zip(Callbacks, ActionArgs)) {
        // Perform the action.
        if (Callback == memcpyAction) {
          if (auto Err = memcpyAction(&ActionArg))
            return Err;
        } else if (Callback == releaseBufferAction) {
          if (auto Err = releaseBufferAction(&ActionArg))
            return Err;
        } else if (Callback == releaseSignalAction) {
          if (auto Err = releaseSignalAction(&ActionArg))
            return Err;
        } else if (Callback == timeKernelInNsAsync) {
          if (auto Err = timeKernelInNsAsync(&ActionArg))
            return Err;
        } else if (Callback == timeDataTransferInNsAsync) {
          if (auto Err = timeDataTransferInNsAsync(&ActionArg))
            return Err;
        } else if (Callback) {
          if (auto Err = Callback(ActionArg.CallbackArgs))
            return Err;
        }
      }

      // Invalidate the action.
      Callbacks.clear();
      ActionArgs.clear();

      return Plugin::success();
    }
  };

  /// The device agent where the stream was created.
  hsa_agent_t Agent;

  /// The queue that the stream uses to launch kernels.
  AMDGPUQueueTy *Queue;

  /// The manager of signals to reuse signals.
  AMDGPUSignalManagerTy &SignalManager;

  /// A reference to the associated device.
  GenericDeviceTy &Device;

  /// Array of stream slots. Use std::deque because it can dynamically grow
  /// without invalidating the already inserted elements. For instance, the
  /// std::vector may invalidate the elements by reallocating the internal
  /// array if there is not enough space on new insertions.
  std::deque<StreamSlotTy> Slots;

  /// The next available slot on the queue. This is reset to zero each time the
  /// stream is synchronized. It also indicates the current number of consumed
  /// slots at a given time.
  uint32_t NextSlot;

  /// The synchronization id. This number is increased each time the stream is
  /// synchronized. It is useful to detect if an AMDGPUEventTy points to an
  /// operation that was already finalized in a previous stream sycnhronize.
  uint32_t SyncCycle;

  /// Mutex to protect stream's management.
  mutable std::mutex Mutex;

  /// Timeout hint for HSA actively waiting for signal value to change
  const uint64_t StreamBusyWaitMicroseconds;

  /// Indicate to spread data transfers across all available SDMAs
  bool UseMultipleSdmaEngines;

  struct CallbackDataType {
    HostFnType UserFn;
    void *UserData;
    AMDGPUSignalTy *OutputSignal;
  };
  /// Wrapper function for implementing host callbacks
  static bool callbackWrapper([[maybe_unused]] hsa_signal_value_t Signal,
                              void *UserData) {
    auto CallbackData = reinterpret_cast<CallbackDataType *>(UserData);
    CallbackData->UserFn(CallbackData->UserData);
    CallbackData->OutputSignal->signal();
    delete CallbackData;
    return false;
  }

  /// Return the current number of asynchronous operations on the stream.
  uint32_t size() const { return NextSlot; }

  /// Return the last valid slot on the stream.
  uint32_t last() const { return size() - 1; }

  /// Consume one slot from the stream. Since the stream uses signals on demand
  /// and releases them once the slot is no longer used, the function requires
  /// an idle signal for the new consumed slot.
  std::pair<uint32_t, AMDGPUSignalTy *> consume(AMDGPUSignalTy *OutputSignal) {
    // Double the stream size if needed. Since we use std::deque, this operation
    // does not invalidate the already added slots.
    if (Slots.size() == NextSlot)
      Slots.resize(Slots.size() * 2);

    // Update the next available slot and the stream size.
    uint32_t Curr = NextSlot++;

    // Retrieve the input signal, if any, of the current operation.
    AMDGPUSignalTy *InputSignal = (Curr > 0) ? Slots[Curr - 1].Signal : nullptr;

    // Set the output signal of the current slot.
    Slots[Curr].Signal = OutputSignal;

    return {Curr, InputSignal};
  }

  /// Roll back the last consumed slot after a submission failure so the stream
  /// does not retain a slot for an operation that was never enqueued.
  void rollbackConsumedSlot(uint32_t Slot) {
    assert(NextSlot > 0 && "Cannot roll back an empty stream");
    assert(Slot + 1 == NextSlot && "Can only roll back the last consumed slot");

    Slots[Slot].Signal = nullptr;
    Slots[Slot].Callbacks.clear();
    Slots[Slot].ActionArgs.clear();
    --NextSlot;
  }

  /// Complete all pending post actions and reset the stream after synchronizing
  /// or positively querying the stream.
  Error complete() {
    for (uint32_t Slot = 0; Slot < NextSlot; ++Slot) {
      // Take the post action of the operation if any.
      if (auto Err = Slots[Slot].performAction())
        return Err;

      // Release the slot's signal if possible. Otherwise, another user will.
      if (Slots[Slot].Signal->decreaseUseCount())
        if (auto Err = SignalManager.returnResource(Slots[Slot].Signal))
          return Err;

      Slots[Slot].Signal = nullptr;
    }

    // Reset the stream slots to zero.
    NextSlot = 0;

    // Increase the synchronization id since the stream completed a sync cycle.
    SyncCycle += 1;

    return Plugin::success();
  }

  /// Complete pending post actions until and including the event in target
  /// slot.
  Error completeUntil(uint32_t TargetSlot) {
    for (uint32_t Slot = 0; Slot <= TargetSlot; ++Slot) {
      // Take the post action of the operation if any.
      if (auto Err = Slots[Slot].performAction())
        return Err;
    }

    return Plugin::success();
  }

  /// Make the current stream wait on a specific operation of another stream.
  /// The idea is to make the current stream waiting on two signals: 1) the last
  /// signal of the current stream, and 2) the last signal of the other stream.
  /// Use a barrier packet with two input signals.
  Error waitOnStreamOperation(AMDGPUStreamTy &OtherStream, uint32_t Slot) {
    if (Queue == nullptr)
      return Plugin::error(ErrorCode::INVALID_NULL_POINTER,
                           "target queue was nullptr");

    /// The signal that we must wait from the other stream.
    AMDGPUSignalTy *OtherSignal = OtherStream.Slots[Slot].Signal;

    // Prevent the release of the other stream's signal.
    OtherSignal->increaseUseCount();

    // Retrieve an available signal for the operation's output.
    AMDGPUSignalTy *OutputSignal = nullptr;
    if (auto Err = SignalManager.getResource(OutputSignal))
      return Err;
    OutputSignal->reset();
    OutputSignal->increaseUseCount();

    // Consume stream slot and compute dependencies.
    auto [Curr, InputSignal] = consume(OutputSignal);

    // Setup the post action to release the signal.
    if (auto Err = Slots[Curr].schedReleaseSignal(OtherSignal, &SignalManager))
      return Err;

    // Push a barrier into the queue with both input signals.
    return Queue->pushBarrier(OutputSignal, InputSignal, OtherSignal);
  }

  /// Callback for running a specific asynchronous operation. This callback is
  /// used for hsa_amd_signal_async_handler. The argument is the operation that
  /// should be executed. Notice we use the post action mechanism to codify the
  /// asynchronous operation.
  static bool asyncActionCallback(hsa_signal_value_t Value, void *Args) {
    StreamSlotTy *Slot = reinterpret_cast<StreamSlotTy *>(Args);
    assert(Slot && "Invalid slot");
    assert(Slot->Signal && "Invalid signal");

    // This thread is outside the stream mutex. Make sure the thread sees the
    // changes on the slot.
    std::atomic_thread_fence(std::memory_order_acquire);

    // Perform the operation.
    if (auto Err = Slot->performAction())
      FATAL_MESSAGE(1, "Error performing post action: %s",
                    toString(std::move(Err)).data());

    // Signal the output signal to notify the asynchronous operation finalized.
    Slot->Signal->signal();

    // Unregister callback.
    return false;
  }

  // Callback for host-to-host memory copies. This is an asynchronous action.
  static Error memcpyAction(void *Data) {
    MemcpyArgsTy *Args = reinterpret_cast<MemcpyArgsTy *>(Data);
    assert(Args && "Invalid arguments");
    assert(Args->Dst && "Invalid destination buffer");
    assert(Args->Src && "Invalid source buffer");

    auto *BasePtr = Args->Dst;
    for (size_t I = 0; I < Args->NumTimes; I++) {
      std::memcpy(BasePtr, Args->Src, Args->Size);
      BasePtr = reinterpret_cast<uint8_t *>(BasePtr) + Args->Size;
    }

    return Plugin::success();
  }

  /// Releasing a memory buffer to a memory manager. This is a post completion
  /// action. There are two kinds of memory buffers:
  ///   1. For kernel arguments. This buffer can be freed after receiving the
  ///   kernel completion signal.
  ///   2. For H2D transfers that need pinned memory space for staging. This
  ///   buffer can be freed after receiving the transfer completion signal.
  ///   3. For D2H transfers that need pinned memory space for staging. This
  ///   buffer cannot be freed after receiving the transfer completion signal
  ///   because of the following asynchronous H2H callback.
  ///      For this reason, This action can only be taken at
  ///      AMDGPUStreamTy::complete()
  /// Because of the case 3, all releaseBufferActions are taken at
  /// AMDGPUStreamTy::complete() in the current implementation.
  static Error releaseBufferAction(void *Data) {
    ReleaseBufferArgsTy *Args = reinterpret_cast<ReleaseBufferArgsTy *>(Data);
    assert(Args && "Invalid arguments");
    assert(Args->MemoryManager && "Invalid memory manager");

    // Release the allocation to the memory manager.
    return Args->MemoryManager->deallocate(Args->Buffer);
  }

  /// Releasing a signal object back to SignalManager. This is a post completion
  /// action. This action can only be taken at AMDGPUStreamTy::complete()
  static Error releaseSignalAction(void *Data) {
    ReleaseSignalArgsTy *Args = reinterpret_cast<ReleaseSignalArgsTy *>(Data);
    assert(Args && "Invalid arguments");
    assert(Args->Signal && "Invalid signal");
    assert(Args->SignalManager && "Invalid signal manager");

    // Release the signal if needed.
    if (Args->Signal->decreaseUseCount())
      if (auto Err = Args->SignalManager->returnResource(Args->Signal))
        return Err;

    return Plugin::success();
  }

public:
  /// Create an empty stream associated with a specific device.
  AMDGPUStreamTy(AMDGPUDeviceTy &Device);

  /// Initialize the stream's signals.
  Error init() { return Plugin::success(); }

  /// Deinitialize the stream's signals.
  Error deinit() { return Plugin::success(); }

  /// Push a asynchronous kernel to the stream. The kernel arguments must be
  /// placed in a special allocation for kernel args and must keep alive until
  /// the kernel finalizes. Once the kernel is finished, the stream will release
  /// the kernel args buffer to the specified memory manager.
  Error pushKernelLaunch(const AMDGPUKernelTy &Kernel, void *KernelArgs,
                         uint32_t NumThreads[3], uint32_t NumBlocks[3],
                         uint32_t GroupSize, uint64_t StackSize,
                         AMDGPUMemoryManagerTy &MemoryManager,
                         void *ProfilerSpecificData = nullptr,
                         GenericProfilerTy *ProfilerPtr = nullptr) {
    if (Queue == nullptr)
      return Plugin::error(ErrorCode::INVALID_NULL_POINTER,
                           "target queue was nullptr");

    // Retrieve an available signal for the operation's output.
    AMDGPUSignalTy *OutputSignal = nullptr;
    if (auto Err = SignalManager.getResource(OutputSignal))
      return Err;
    OutputSignal->reset();
    OutputSignal->increaseUseCount();

    std::lock_guard<std::mutex> StreamLock(Mutex);

    // Consume stream slot and compute dependencies.
    auto [Curr, InputSignal] = consume(OutputSignal);

    // Setup the post action to release the kernel args buffer.
    if (auto Err = Slots[Curr].schedReleaseBuffer(KernelArgs, MemoryManager))
      return Err;

    if (ProfilerSpecificData) {
      if (auto Err = Slots[Curr].schedProfilerKernelTiming(
              ProfilerPtr, Agent, OutputSignal, TicksToTime,
              ProfilerSpecificData))
        return Err;
    }

    // If we are running an RPC server we want to wake up the server thread
    // whenever there is a kernel running and let it sleep otherwise.
    if (Device.getRPCServer())
      Device.Plugin.getRPCServer().Thread->notify();

    // Push the kernel with the output signal and an input signal (optional)
    if (auto Err = Queue->pushKernelLaunch(Kernel, KernelArgs, NumThreads,
                                           NumBlocks, GroupSize, StackSize,
                                           OutputSignal, InputSignal))
      return Err;

    // Register a callback to indicate when the kernel is complete.
    if (Device.getRPCServer()) {
      if (auto Err = Slots[Curr].schedCallback(
              [](void *Data) -> llvm::Error {
                GenericPluginTy &Plugin =
                    *reinterpret_cast<GenericPluginTy *>(Data);
                Plugin.getRPCServer().Thread->finish();
                return Error::success();
              },
              &Device.Plugin))
        return Err;
    }
    return Plugin::success();
  }

  /// Push an asynchronous memory copy between pinned memory buffers.
  Error pushPinnedMemoryCopyAsync(void *Dst, const void *Src,
                                  uint64_t CopySize) {
    // Retrieve an available signal for the operation's output.
    AMDGPUSignalTy *OutputSignal = nullptr;
    if (auto Err = SignalManager.getResource(OutputSignal))
      return Err;
    OutputSignal->reset();
    OutputSignal->increaseUseCount();

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

    // Consume stream slot and compute dependencies.
    auto [Curr, InputSignal] = consume(OutputSignal);

    // Issue the async memory copy.
    if (InputSignal && InputSignal->load()) {
      hsa_signal_t InputSignalRaw = InputSignal->get();
      return hsa_utils::asyncMemCopy(UseMultipleSdmaEngines, Dst, Agent, Src,
                                     Agent, CopySize, 1, &InputSignalRaw,
                                     OutputSignal->get());
    }

    return hsa_utils::asyncMemCopy(UseMultipleSdmaEngines, Dst, Agent, Src,
                                   Agent, CopySize, 0, nullptr,
                                   OutputSignal->get());
  }

  /// Push an asynchronous memory copy device-to-host involving an unpinned
  /// memory buffer. The operation consists of a two-step copy from the
  /// device buffer to an intermediate pinned host buffer, and then, to a
  /// unpinned host buffer. Both operations are asynchronous and dependent.
  /// The intermediate pinned buffer will be released to the specified memory
  /// manager once the operation completes.
  Error pushMemoryCopyD2HAsync(void *Dst, const void *Src, void *Inter,
                               uint64_t CopySize,
                               AMDGPUMemoryManagerTy &MemoryManager,
                               void *ProfilerSpecificData = nullptr,
                               GenericProfilerTy *ProfilerPtr = nullptr) {
    // Retrieve available signals for the operation's outputs.
    AMDGPUSignalTy *OutputSignals[2] = {};
    if (auto Err = SignalManager.getResources(/*Num=*/2, OutputSignals))
      return Err;
    for (auto *Signal : OutputSignals) {
      Signal->reset();
      Signal->increaseUseCount();
    }

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

    // Consume stream slot and compute dependencies.
    auto [Curr, InputSignal] = consume(OutputSignals[0]);

    // Setup the post action for releasing the intermediate buffer.
    if (auto Err = Slots[Curr].schedReleaseBuffer(Inter, MemoryManager))
      return Err;

    // Issue the first step: device to host transfer. Avoid defining the input
    // dependency if already satisfied.
    if (InputSignal && InputSignal->load()) {
      hsa_signal_t InputSignalRaw = InputSignal->get();
      if (auto Err = hsa_utils::asyncMemCopy(
              UseMultipleSdmaEngines, Inter, Agent, Src, Agent, CopySize, 1,
              &InputSignalRaw, OutputSignals[0]->get()))
        return Err;
    } else {
      if (auto Err = hsa_utils::asyncMemCopy(UseMultipleSdmaEngines, Inter,
                                             Agent, Src, Agent, CopySize, 0,
                                             nullptr, OutputSignals[0]->get()))
        return Err;
    }

    // Consume another stream slot and compute dependencies.
    std::tie(Curr, InputSignal) = consume(OutputSignals[1]);
    assert(InputSignal && "Invalid input signal");

    // The std::memcpy is done asynchronously using an async handler. We store
    // the function's information in the action but it's not actually an action.
    if (auto Err = Slots[Curr].schedHostMemoryCopy(Dst, Inter, CopySize))
      return Err;

    // Make changes on this slot visible to the async handler's thread.
    std::atomic_thread_fence(std::memory_order_release);

    // Issue the second step: host to host transfer.
    hsa_status_t Status = hsa_amd_signal_async_handler(
        InputSignal->get(), HSA_SIGNAL_CONDITION_EQ, 0, asyncActionCallback,
        (void *)&Slots[Curr]);

    return Plugin::check(Status, "error in hsa_amd_signal_async_handler: %s");
  }

  /// Push an asynchronous memory copy host-to-device involving an unpinned
  /// memory buffer. The operation consists of a two-step copy from the
  /// unpinned host buffer to an intermediate pinned host buffer, and then, to
  /// the pinned host buffer. Both operations are asynchronous and dependent.
  /// The intermediate pinned buffer will be released to the specified memory
  /// manager once the operation completes.
  Error pushMemoryCopyH2DAsync(void *Dst, const void *Src, void *Inter,
                               uint64_t CopySize,
                               AMDGPUMemoryManagerTy &MemoryManager,
                               void *ProfilerSpecificData = nullptr,
                               GenericProfilerTy *ProfilerPtr = nullptr,
                               size_t NumTimes = 1) {
    // Retrieve available signals for the operation's outputs.
    AMDGPUSignalTy *OutputSignals[2] = {};
    if (auto Err = SignalManager.getResources(/*Num=*/2, OutputSignals))
      return Err;
    for (auto *Signal : OutputSignals) {
      Signal->reset();
      Signal->increaseUseCount();
    }

    AMDGPUSignalTy *OutputSignal = OutputSignals[0];

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

    // Consume stream slot and compute dependencies.
    auto [Curr, InputSignal] = consume(OutputSignal);

    // Issue the first step: host to host transfer.
    if (InputSignal && InputSignal->load()) {
      // The std::memcpy is done asynchronously using an async handler. We store
      // the function's information in the action but it is not actually a
      // post action.
      if (auto Err =
              Slots[Curr].schedHostMemoryCopy(Inter, Src, CopySize, NumTimes))
        return Err;

      // Make changes on this slot visible to the async handler's thread.
      std::atomic_thread_fence(std::memory_order_release);

      hsa_status_t Status = hsa_amd_signal_async_handler(
          InputSignal->get(), HSA_SIGNAL_CONDITION_EQ, 0, asyncActionCallback,
          (void *)&Slots[Curr]);

      if (auto Err = Plugin::check(Status,
                                   "error in hsa_amd_signal_async_handler: %s"))
        return Err;

      // Let's use now the second output signal.
      OutputSignal = OutputSignals[1];

      // Consume another stream slot and compute dependencies.
      std::tie(Curr, InputSignal) = consume(OutputSignal);
    } else {
      // All preceding operations completed, copy the memory synchronously.
      auto *InterPtr = Inter;
      for (size_t I = 0; I < NumTimes; I++) {
        std::memcpy(InterPtr, Src, CopySize);
        InterPtr = reinterpret_cast<uint8_t *>(InterPtr) + CopySize;
      }

      // Return the second signal because it will not be used.
      OutputSignals[1]->decreaseUseCount();
      if (auto Err = SignalManager.returnResource(OutputSignals[1]))
        return Err;
    }

    // Setup the post action to release the intermediate pinned buffer.
    if (auto Err = Slots[Curr].schedReleaseBuffer(Inter, MemoryManager))
      return Err;

    // Issue the second step: host to device transfer. Avoid defining the input
    // dependency if already satisfied.
    if (InputSignal && InputSignal->load()) {
      hsa_signal_t InputSignalRaw = InputSignal->get();
      return hsa_utils::asyncMemCopy(UseMultipleSdmaEngines, Dst, Agent, Inter,
                                     Agent, CopySize * NumTimes, 1,
                                     &InputSignalRaw, OutputSignal->get());
    }
    return hsa_utils::asyncMemCopy(UseMultipleSdmaEngines, Dst, Agent, Inter,
                                   Agent, CopySize * NumTimes, 0, nullptr,
                                   OutputSignal->get());
  }

  // AMDGPUDeviceTy is incomplete here, passing the underlying agent instead
  Error pushMemoryCopyD2DAsync(void *Dst, hsa_agent_t DstAgent, const void *Src,
                               hsa_agent_t SrcAgent, uint64_t CopySize,
                               void *ProfilerSpecificData = nullptr,
                               GenericProfilerTy *ProfilerPtr = nullptr) {
    AMDGPUSignalTy *OutputSignal;
    if (auto Err = SignalManager.getResources(/*Num=*/1, &OutputSignal))
      return Err;
    OutputSignal->reset();
    OutputSignal->increaseUseCount();

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

    // Consume stream slot and compute dependencies.
    auto [Curr, InputSignal] = consume(OutputSignal);

    // The agents need to have access to the corresponding memory
    // This is presently only true if the pointers were originally
    // allocated by this runtime or the caller made the appropriate
    // access calls.

    if (InputSignal && InputSignal->load()) {
      hsa_signal_t InputSignalRaw = InputSignal->get();
      return hsa_utils::asyncMemCopy(UseMultipleSdmaEngines, Dst, DstAgent, Src,
                                     SrcAgent, CopySize, 1, &InputSignalRaw,
                                     OutputSignal->get());
    }
    return hsa_utils::asyncMemCopy(UseMultipleSdmaEngines, Dst, DstAgent, Src,
                                   SrcAgent, CopySize, 0, nullptr,
                                   OutputSignal->get());
  }

  Error pushHostCallback(HostFnType Callback, void *UserData) {
    // Retrieve an available signal for the operation's output.
    AMDGPUSignalTy *OutputSignal = nullptr;
    if (auto Err = SignalManager.getResource(OutputSignal))
      return Err;
    OutputSignal->reset();
    OutputSignal->increaseUseCount();

    AMDGPUSignalTy *InputSignal;
    {
      std::lock_guard<std::mutex> Lock(Mutex);

      // Consume stream slot and compute dependencies.
      InputSignal = consume(OutputSignal).second;
    }

    auto *CallbackData = new CallbackDataType{Callback, UserData, OutputSignal};
    if (InputSignal && InputSignal->load()) {
      hsa_status_t Status = hsa_amd_signal_async_handler(
          InputSignal->get(), HSA_SIGNAL_CONDITION_EQ, 0, callbackWrapper,
          CallbackData);

      return Plugin::check(Status, "error in hsa_amd_signal_async_handler: %s");
    }

    // No dependencies - schedule it now.
    // Using a seperate thread because this function should run asynchronously
    // and not block the main thread.
    std::thread([](void *CallbackData) { callbackWrapper(0, CallbackData); },
                CallbackData)
        .detach();
    return Plugin::success();
  }

  /// Synchronize with the stream. The current thread waits until all operations
  /// are finalized and it performs the pending post actions (i.e., releasing
  /// intermediate buffers).
  Error synchronize() {
    std::lock_guard<std::mutex> Lock(Mutex);

    // No need to synchronize anything.
    if (size() == 0)
      return Plugin::success();

    // Wait until all previous operations on the stream have completed.
    if (auto Err =
            Slots[last()].Signal->wait(StreamBusyWaitMicroseconds, &Device))
      return Err;

    // Reset the stream and perform all pending post actions.
    return complete();
  }

  /// Synchronize the stream until the given event. The current thread waits
  /// until the provided event is finalized, and it performs the pending post
  /// actions for that and prior events.
  Error synchronizeOn(AMDGPUEventTy &Event);

  /// Return true if the event from this queue is complete
  Expected<bool> isEventComplete(const AMDGPUEventTy &Event);

  /// Query the stream and complete pending post actions if operations finished.
  /// Return whether all the operations completed. This operation does not block
  /// the calling thread.
  Expected<bool> query() {
    std::lock_guard<std::mutex> Lock(Mutex);

    // No need to query anything.
    if (size() == 0)
      return true;

    // The last operation did not complete yet. Return directly.
    if (Slots[last()].Signal->load())
      return false;

    // Reset the stream and perform all pending post actions.
    if (auto Err = complete())
      return std::move(Err);

    return true;
  }

  const AMDGPUQueueTy *getQueue() const { return Queue; }

  /// Record an event on the stream. If \p EnableProfiling is true, a barrier
  /// marker packet is enqueued to record timestamps.
  Error recordEvent(AMDGPUEventTy &Event,
                    AMDGPUSignalTy *ReusedSignal = nullptr,
                    bool EnableProfiling = false);

  /// Make the stream wait on an event.
  Error waitEvent(const AMDGPUEventTy &Event);

private:
  /// Lightweight sync-only recording.
  Error recordEventSyncOnly(AMDGPUEventTy &Event, AMDGPUSignalTy *ReusedSignal);

  /// Full profiling recording.
  Error recordEventWithProfiling(AMDGPUEventTy &Event,
                                 AMDGPUSignalTy *ReusedSignal);

  friend struct AMDGPUStreamManagerTy;
};

/// Class representing an event on AMDGPU. The event stores the recorded stream
/// point and retained timing state.
struct AMDGPUEventTy {
  /// Create an empty event.
  AMDGPUEventTy(AMDGPUDeviceTy &Device)
      : Device(Device), RecordedStream(nullptr), RecordedSlot(-1),
        RecordedSyncCycle(-1), TimingSignal(nullptr) {}

  /// Initialize and deinitialize.
  Error init() { return resetState(); }
  Error deinit() { return resetState(); }

  /// Clear the current recording and retained timing state, optionally
  /// returning a reusable timing signal.
  Error resetState(AMDGPUSignalTy **ReusableSignalPtr = nullptr) {
    RecordedStream = nullptr;
    RecordedSlot = -1;
    RecordedSyncCycle = -1;
    return releaseTimingSignal(ReusableSignalPtr);
  }

  /// Record the current stream point on the event.
  Error record(AMDGPUStreamTy &Stream, bool EnableProfiling = false) {
    std::lock_guard<std::mutex> Lock(Mutex);

    // Discard the previous recording and retained timing state, reusing the
    // retained timing signal if it becomes available.
    AMDGPUSignalTy *Signal = nullptr;
    if (auto Err = resetState(&Signal))
      return Err;

    RecordedStream = &Stream;

    if (auto Err = Stream.recordEvent(*this, Signal, EnableProfiling)) {
      if (auto ResetErr = resetState())
        return joinErrors(std::move(Err), std::move(ResetErr));
      return Err;
    }

    return Plugin::success();
  }

  /// Make a stream wait on the current event.
  Error wait(AMDGPUStreamTy &Stream) {
    std::lock_guard<std::mutex> Lock(Mutex);

    if (!RecordedStream)
      return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                           "event does not have any recorded stream");

    // Synchronizing the same stream. Do nothing.
    if (RecordedStream == &Stream)
      return Plugin::success();

    // No need to wait anything, the recorded stream already finished the
    // corresponding operation.
    if (RecordedSlot < 0)
      return Plugin::success();

    return Stream.waitEvent(*this);
  }

  Error sync() {
    std::lock_guard<std::mutex> Lock(Mutex);

    if (!RecordedStream)
      return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                           "event does not have any recorded stream");

    // No need to wait on anything, the recorded stream already finished the
    // corresponding operation.
    if (RecordedSlot < 0)
      return Plugin::success();

    return RecordedStream->synchronizeOn(*this);
  }

  /// Return the elapsed time in milliseconds between this event and EndEvent.
  Expected<float> getElapsedTime(AMDGPUEventTy &EndEvent);

protected:
  /// Release the retained timing signal, if any, either back to the signal
  /// manager or through \p ReusableSignalPtr when provided.
  Error releaseTimingSignal(AMDGPUSignalTy **ReusableSignalPtr = nullptr);

  /// The device that owns this event.
  AMDGPUDeviceTy &Device;

  /// The stream registered in this event.
  AMDGPUStreamTy *RecordedStream;

  /// The recorded operation on the recorded stream.
  int64_t RecordedSlot;

  /// The sync cycle when the stream was recorded. Used to detect stale events.
  int64_t RecordedSyncCycle;

  /// The signal of the recorded timing barrier.
  AMDGPUSignalTy *TimingSignal;

  /// Mutex to safely access event fields.
  mutable std::mutex Mutex;

  friend struct AMDGPUStreamTy;
};

Error AMDGPUStreamTy::recordEvent(AMDGPUEventTy &Event,
                                  AMDGPUSignalTy *ReusedSignal,
                                  bool EnableProfiling) {
  if (Queue == nullptr)
    return Plugin::error(ErrorCode::INVALID_NULL_POINTER,
                         "target queue was nullptr");

  if (EnableProfiling)
    return recordEventWithProfiling(Event, ReusedSignal);
  return recordEventSyncOnly(Event, ReusedSignal);
}

Error AMDGPUStreamTy::recordEventSyncOnly(AMDGPUEventTy &Event,
                                          AMDGPUSignalTy *ReusedSignal) {
  std::lock_guard<std::mutex> StreamLock(Mutex);

  if (size() > 0) {
    // Record the synchronize identifier (to detect stale recordings) and
    // the last valid stream's operation.
    Event.RecordedSyncCycle = SyncCycle;
    Event.RecordedSlot = last();

    assert(Event.RecordedSyncCycle >= 0 && "Invalid recorded sync cycle");
    assert(Event.RecordedSlot >= 0 && "Invalid recorded slot");
  } else {
    // The stream is empty, everything already completed, record nothing.
    Event.RecordedSyncCycle = -1;
    Event.RecordedSlot = -1;
  }

  // Return the unused reusable signal back to the manager.
  if (ReusedSignal)
    return SignalManager.returnResource(ReusedSignal);
  return Plugin::success();
}

Error AMDGPUStreamTy::recordEventWithProfiling(AMDGPUEventTy &Event,
                                               AMDGPUSignalTy *ReusedSignal) {
  // One use for the stream slot and one for the event timing signal.
  const uint32_t OutputSignalUses = 2;

  // Reuse the provided signal or retrieve one for the operation's output.
  AMDGPUSignalTy *OutputSignal = ReusedSignal;
  if (!OutputSignal) {
    if (auto Err = SignalManager.getResource(OutputSignal))
      return Err;
  }

  OutputSignal->reset();
  OutputSignal->increaseUseCount(OutputSignalUses);

  std::lock_guard<std::mutex> StreamLock(Mutex);

  // Consume stream slot and compute dependencies.
  auto [Curr, InputSignal] = consume(OutputSignal);

  // Materialize the event as a real marker on the queue. Elapsed-time queries
  // need a packet-backed completion signal to retrieve dispatch timing.
  if (auto Err = Queue->pushBarrier(OutputSignal, InputSignal, nullptr)) {
    rollbackConsumedSlot(Curr);

    if (OutputSignal->decreaseUseCount(OutputSignalUses)) {
      if (auto ReturnErr = SignalManager.returnResource(OutputSignal))
        return joinErrors(std::move(Err), std::move(ReturnErr));
    }

    return Err;
  }

  Event.RecordedSlot = Curr;
  Event.RecordedSyncCycle = SyncCycle;
  Event.TimingSignal = OutputSignal;

  assert(Event.RecordedSyncCycle >= 0 && "Invalid recorded sync cycle");
  assert(Event.RecordedSlot >= 0 && "Invalid recorded slot");

  return Plugin::success();
}

Error AMDGPUStreamTy::waitEvent(const AMDGPUEventTy &Event) {
  // Retrieve the recorded stream on the event.
  AMDGPUStreamTy &RecordedStream = *Event.RecordedStream;

  std::scoped_lock<std::mutex, std::mutex> Lock(Mutex, RecordedStream.Mutex);

  // The recorded stream already completed the operation because the synchronize
  // identifier is already outdated.
  if (RecordedStream.SyncCycle != (uint32_t)Event.RecordedSyncCycle)
    return Plugin::success();

  // Again, the recorded stream already completed the operation, the last
  // operation's output signal is satisfied.
  if (!RecordedStream.Slots[Event.RecordedSlot].Signal->load())
    return Plugin::success();

  // Otherwise, make the current stream wait on the other stream's operation.
  return waitOnStreamOperation(RecordedStream, Event.RecordedSlot);
}

Error AMDGPUStreamTy::synchronizeOn(AMDGPUEventTy &Event) {
  std::lock_guard<std::mutex> Lock(Mutex);

  // If this event was for an older sync cycle, it has already been finalized
  if (Event.RecordedSyncCycle < SyncCycle)
    return Plugin::success();
  assert(Event.RecordedSyncCycle == SyncCycle && "event is from the future?");

  // Wait until the requested slot has completed
  if (auto Err = Slots[Event.RecordedSlot].Signal->wait(
          StreamBusyWaitMicroseconds, &Device))
    return Err;

  // If the event is the last one in the stream, just do a full finalize
  if (Event.RecordedSlot == last())
    return complete();

  // Otherwise, only finalize until the appropriate event
  return completeUntil(Event.RecordedSlot);
}

Expected<bool> AMDGPUStreamTy::isEventComplete(const AMDGPUEventTy &Event) {
  std::lock_guard<std::mutex> Lock(Mutex);
  assert(Event.RecordedStream == this && "event is for a different stream");

  if (Event.RecordedSyncCycle < SyncCycle) {
    return true;
  }
  assert(Event.RecordedSyncCycle == SyncCycle && "event is from the future?");

  return !Slots[Event.RecordedSlot].Signal->load();
}

struct AMDGPUStreamManagerTy final
    : GenericDeviceResourceManagerTy<AMDGPUResourceRef<AMDGPUStreamTy>> {
  using ResourceRef = AMDGPUResourceRef<AMDGPUStreamTy>;
  using ResourcePoolTy = GenericDeviceResourceManagerTy<ResourceRef>;

  AMDGPUStreamManagerTy(GenericDeviceTy &Device, hsa_agent_t HSAAgent)
      : GenericDeviceResourceManagerTy(Device), Device(Device),
        OMPX_QueueTracking("LIBOMPTARGET_AMDGPU_HSA_QUEUE_BUSY_TRACKING", true),
        NextQueue(0), Agent(HSAAgent) {}

  Error init(uint32_t InitialSize, int NumHSAQueues, int HSAQueueSize) {
    Queues = std::vector<AMDGPUQueueTy>(NumHSAQueues);
    QueueSize = HSAQueueSize;
    MaxNumQueues = NumHSAQueues;
    // Initialize one queue eagerly
    if (auto Err = Queues.front().init(Device, Agent, QueueSize))
      return Err;

    return GenericDeviceResourceManagerTy::init(InitialSize);
  }

  /// Deinitialize the resource pool and delete all resources. This function
  /// must be called before the destructor.
  Error deinit() override {
    // De-init all queues
    for (AMDGPUQueueTy &Queue : Queues) {
      if (auto Err = Queue.deinit())
        return Err;
    }

    return GenericDeviceResourceManagerTy::deinit();
  }

  /// Get a single stream from the pool or create new resources.
  virtual Error getResource(AMDGPUStreamTy *&StreamHandle) override {
    return getResourcesImpl(1, &StreamHandle, [this](AMDGPUStreamTy *&Handle) {
      return assignNextQueue(Handle);
    });
  }

  /// Return stream to the pool.
  virtual Error returnResource(AMDGPUStreamTy *StreamHandle) override {
    return returnResourceImpl(StreamHandle, [](AMDGPUStreamTy *Handle) {
      Handle->Queue->removeUser();
      return Plugin::success();
    });
  }

private:
  /// Search for and assign an preferably idle queue to the given Stream. If
  /// there is no queue without current users, choose the queue with the lowest
  /// user count. If utilization is ignored: use round robin selection.
  inline Error assignNextQueue(AMDGPUStreamTy *Stream) {
    // Start from zero when tracking utilization, otherwise: round robin policy.
    uint32_t Index = OMPX_QueueTracking ? 0 : NextQueue++ % MaxNumQueues;

    if (OMPX_QueueTracking) {
      // Find the least used queue.
      for (uint32_t I = 0; I < MaxNumQueues; ++I) {
        // Early exit when an initialized queue is idle.
        if (Queues[I].isInitialized() && Queues[I].getUserCount() == 0) {
          Index = I;
          break;
        }

        // Update the least used queue.
        if (Queues[Index].getUserCount() > Queues[I].getUserCount())
          Index = I;
      }
    }

    // Make sure the queue is initialized, then add user & assign.
    if (auto Err = Queues[Index].init(Device, Agent, QueueSize))
      return Err;
    Queues[Index].addUser();
    Stream->Queue = &Queues[Index];

    return Plugin::success();
  }

  /// The device associated with this stream.
  GenericDeviceTy &Device;

  /// Envar for controlling the tracking of busy HSA queues.
  BoolEnvar OMPX_QueueTracking;

  /// The next queue index to use for round robin selection.
  uint32_t NextQueue;

  /// The queues which are assigned to requested streams.
  std::vector<AMDGPUQueueTy> Queues;

  /// The corresponding device as HSA agent.
  hsa_agent_t Agent;

  /// The maximum number of queues.
  uint32_t MaxNumQueues;

  /// The size of created queues.
  uint32_t QueueSize;
};

/// Abstract class that holds the common members of the actual kernel devices
/// and the host device. Both types should inherit from this class.
struct AMDGenericDeviceTy {
  AMDGenericDeviceTy() {}

  virtual ~AMDGenericDeviceTy() {}

  /// Create all memory pools which the device has access to and classify them.
  Error initMemoryPools() {
    // Retrieve all memory pools from the device agent(s).
    Error Err = retrieveAllMemoryPools();
    if (Err)
      return Err;

    for (AMDGPUMemoryPoolTy *MemoryPool : AllMemoryPools) {
      // Initialize the memory pool and retrieve some basic info.
      Error Err = MemoryPool->init();
      if (Err)
        return Err;

      if (!MemoryPool->isGlobal())
        continue;

      // Classify the memory pools depending on their properties.
      if (MemoryPool->isFineGrained()) {
        FineGrainedMemoryPools.push_back(MemoryPool);
        if (MemoryPool->supportsKernelArgs())
          ArgsMemoryPools.push_back(MemoryPool);
      } else if (MemoryPool->isCoarseGrained()) {
        CoarseGrainedMemoryPools.push_back(MemoryPool);
      }
    }
    return Plugin::success();
  }

  /// Destroy all memory pools.
  Error deinitMemoryPools() {
    for (AMDGPUMemoryPoolTy *Pool : AllMemoryPools)
      delete Pool;

    AllMemoryPools.clear();
    FineGrainedMemoryPools.clear();
    CoarseGrainedMemoryPools.clear();
    ArgsMemoryPools.clear();

    return Plugin::success();
  }

  /// Retrieve and construct all memory pools from the device agent(s).
  virtual Error retrieveAllMemoryPools() = 0;

  /// Get the device agent.
  virtual hsa_agent_t getAgent() const = 0;

protected:
  /// Array of all memory pools available to the host agents.
  llvm::SmallVector<AMDGPUMemoryPoolTy *> AllMemoryPools;

  /// Array of fine-grained memory pools available to the host agents.
  llvm::SmallVector<AMDGPUMemoryPoolTy *> FineGrainedMemoryPools;

  /// Array of coarse-grained memory pools available to the host agents.
  llvm::SmallVector<AMDGPUMemoryPoolTy *> CoarseGrainedMemoryPools;

  /// Array of kernel args memory pools available to the host agents.
  llvm::SmallVector<AMDGPUMemoryPoolTy *> ArgsMemoryPools;
};

/// Class representing a host device, which corresponds to a single HSA host
/// agent. The plugin creates one host device per host agent, and each GPU
/// device uses the host device nearest to it.
struct AMDHostDeviceTy : public AMDGenericDeviceTy {
  /// Create a host device from a host agent.
  AMDHostDeviceTy(AMDGPUPluginTy &Plugin, hsa_agent_t Agent)
      : AMDGenericDeviceTy(), Agent(Agent), ArgsMemoryManager(Plugin),
        PinnedMemoryManager(Plugin) {}

  /// Initialize the host device memory pools and the memory managers for
  /// kernel args and host pinned memory allocations. Return false if the host
  /// agent does not have the required memory pools.
  Expected<bool> init() {
    if (auto Err = initMemoryPools())
      return std::move(Err);

    // The host device cannot be used without these pools. Release the pools
    // and leave the managers uninitialized.
    if (ArgsMemoryPools.empty() || FineGrainedMemoryPools.empty()) {
      if (auto Err = deinitMemoryPools())
        return std::move(Err);
      return false;
    }

    if (auto Err = ArgsMemoryManager.init(getArgsMemoryPool()))
      return std::move(Err);

    if (auto Err = PinnedMemoryManager.init(getFineGrainedMemoryPool()))
      return std::move(Err);

    return true;
  }

  /// Deinitialize memory pools and managers.
  Error deinit() {
    if (auto Err = deinitMemoryPools())
      return Err;

    if (auto Err = ArgsMemoryManager.deinit())
      return Err;

    if (auto Err = PinnedMemoryManager.deinit())
      return Err;

    return Plugin::success();
  }

  /// Retrieve and construct all memory pools from the host agent.
  Error retrieveAllMemoryPools() override {
    return hsa_utils::iterateAgentMemoryPools(
        Agent, [&](hsa_amd_memory_pool_t HSAMemoryPool) {
          AMDGPUMemoryPoolTy *MemoryPool =
              new AMDGPUMemoryPoolTy(HSAMemoryPool);
          AllMemoryPools.push_back(MemoryPool);
          return HSA_STATUS_SUCCESS;
        });
  }

  /// Get the host agent.
  hsa_agent_t getAgent() const override { return Agent; }

  /// Get a memory pool for fine-grained allocations.
  AMDGPUMemoryPoolTy &getFineGrainedMemoryPool() {
    assert(!FineGrainedMemoryPools.empty() && "No fine-grained mempool");
    // Retrieve any memory pool.
    return *FineGrainedMemoryPools[0];
  }

  AMDGPUMemoryPoolTy &getCoarseGrainedMemoryPool() {
    assert(!CoarseGrainedMemoryPools.empty() && "No coarse-grained mempool");
    // Retrieve any memory pool.
    return *CoarseGrainedMemoryPools[0];
  }

  /// Get a memory pool for kernel args allocations.
  AMDGPUMemoryPoolTy &getArgsMemoryPool() {
    assert(!ArgsMemoryPools.empty() && "No kernelargs mempool");
    // Retrieve any memory pool.
    return *ArgsMemoryPools[0];
  }

  /// Getters for kernel args and host pinned memory managers.
  AMDGPUMemoryManagerTy &getArgsMemoryManager() { return ArgsMemoryManager; }
  AMDGPUMemoryManagerTy &getPinnedMemoryManager() {
    return PinnedMemoryManager;
  }

private:
  /// The HSA host agent.
  hsa_agent_t Agent;

  // Memory manager for kernel arguments.
  AMDGPUMemoryManagerTy ArgsMemoryManager;

  // Memory manager for pinned memory.
  AMDGPUMemoryManagerTy PinnedMemoryManager;
};

/// Class implementing the AMDGPU device functionalities which derives from the
/// generic device class.
struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
  // Create an AMDGPU device with a device id and default AMDGPU grid values.
  AMDGPUDeviceTy(GenericPluginTy &Plugin, int32_t DeviceId, int32_t NumDevices,
                 AMDHostDeviceTy &HostDevice, hsa_agent_t Agent)
      : GenericDeviceTy(Plugin, DeviceId, NumDevices, {}), AMDGenericDeviceTy(),
        OMPX_NumQueues("LIBOMPTARGET_AMDGPU_NUM_HSA_QUEUES", 4),
        OMPX_QueueSize("LIBOMPTARGET_AMDGPU_HSA_QUEUE_SIZE", 512),
        OMPX_DefaultTeamsPerCU("LIBOMPTARGET_AMDGPU_TEAMS_PER_CU", 4),
        OMPX_MaxAsyncCopyBytes("LIBOMPTARGET_AMDGPU_MAX_ASYNC_COPY_BYTES",
                               1 * 1024 * 1024), // 1MB
        OMPX_InitialNumSignals("LIBOMPTARGET_AMDGPU_NUM_INITIAL_HSA_SIGNALS",
                               64),
        OMPX_StreamBusyWait("LIBOMPTARGET_AMDGPU_STREAM_BUSYWAIT", 2000000),
        OMPX_UseMultipleSdmaEngines(
            "LIBOMPTARGET_AMDGPU_USE_MULTIPLE_SDMA_ENGINES", false),
        OMPX_ApuMaps("OMPX_APU_MAPS", false), AMDGPUStreamManager(*this, Agent),
        AMDGPUEventManager(*this), AMDGPUSignalManager(*this), Agent(Agent),
        HostDevice(HostDevice) {}

  ~AMDGPUDeviceTy() {}

  /// Initialize the device, its resources and get its properties.
  Error initImpl(GenericPluginTy &Plugin,
                 GenericProfilerTy *ProfilerPtr) override {
    // First setup all the memory pools.
    if (auto Err = initMemoryPools())
      return Err;

    setHSATicksToTimeConstant();

    HostWritesDeviceMemory =
        !CoarseGrainedMemoryPools.empty() &&
        CoarseGrainedMemoryPools[0]->canAccess(HostDevice.getAgent());
    // The writing thread may run on any host agent, so all of them must reach
    // the device over XGMI.
    HostLinkIsXGMI =
        HostWritesDeviceMemory &&
        CoarseGrainedMemoryPools[0]->isLinkedOverXGMI(HostDevice.getAgent());

    char GPUName[64];
    if (auto Err = getDeviceAttr(HSA_AGENT_INFO_NAME, GPUName))
      return Err;
    ComputeUnitKind = GPUName;

    // From the ROCm HSA documentation:
    // Query the UUID of the agent. The value is an Ascii string with a maximum
    // of 21 chars including NUL. The string value consists of two parts: header
    // and body. The header identifies the device type (GPU, CPU, DSP) while the
    // body encodes the UUID as a 16 digit hex string.
    //
    // Agents that do not support UUID will return the string "GPU-XX" or
    // "CPU-XX" or "DSP-XX" depending on their device type.
    char UUID[24] = {0};
    if (auto Err = getDeviceAttr(HSA_AMD_AGENT_INFO_UUID, UUID))
      return Err;
    if (!StringRef(UUID).ends_with("-XX"))
      setDeviceUidFromVendorUid(UUID);

    // Get the wavefront size.
    uint32_t WavefrontSize = 0;
    if (auto Err = getDeviceAttr(HSA_AGENT_INFO_WAVEFRONT_SIZE, WavefrontSize))
      return Err;
    GridValues.GV_Warp_Size = WavefrontSize;

    // Get the frequency of the steady clock. If the attribute is missing
    // assume running on an older libhsa and default to 0, omp_get_wtime
    // will be inaccurate but otherwise programs can still run.
    if (getDeviceAttrRaw(HSA_AMD_AGENT_INFO_TIMESTAMP_FREQUENCY,
                         ClockFrequency) != HSA_STATUS_SUCCESS)
      ClockFrequency = 0;

    // Retrieve the HSA system timestamp frequency for this runtime. A zero
    // value means the frequency is unavailable.
    if (hsa_system_get_info(HSA_SYSTEM_INFO_TIMESTAMP_FREQUENCY,
                            &SystemTimestampFrequency) != HSA_STATUS_SUCCESS)
      SystemTimestampFrequency = 0;

    // Load the grid values depending on the wavefront.
    if (WavefrontSize == 32)
      GridValues = getAMDGPUGridValues<32>();
    else if (WavefrontSize == 64)
      GridValues = getAMDGPUGridValues<64>();
    else
      return Plugin::error(ErrorCode::UNSUPPORTED,
                           "unexpected AMDGPU wavefront %d", WavefrontSize);

    // Get maximum number of workitems per workgroup.
    uint16_t WorkgroupMaxDim[3];
    if (auto Err =
            getDeviceAttr(HSA_AGENT_INFO_WORKGROUP_MAX_DIM, WorkgroupMaxDim))
      return Err;
    GridValues.GV_Max_WG_Size = WorkgroupMaxDim[0];

    // Get maximum number of workgroups.
    hsa_dim3_t GridMaxDim;
    if (auto Err = getDeviceAttr(HSA_AGENT_INFO_GRID_MAX_DIM, GridMaxDim))
      return Err;

    // The HSA grid is expressed in work items, so it limits blocks * threads
    // rather than blocks alone. A block has at least one thread.
    GridValues.GV_Max_Teams = GridMaxDim.x;
    MaxGridSize = GridMaxDim.x;
    if (GridValues.GV_Max_Teams == 0)
      return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                           "maximum number of teams cannot be zero");

    // Compute the default number of teams.
    uint32_t ComputeUnits = 0;
    if (auto Err =
            getDeviceAttr(HSA_AMD_AGENT_INFO_COMPUTE_UNIT_COUNT, ComputeUnits))
      return Err;
    GridValues.GV_Default_Num_Teams = ComputeUnits * OMPX_DefaultTeamsPerCU;

    uint32_t WavesPerCU = 0;
    if (auto Err =
            getDeviceAttr(HSA_AMD_AGENT_INFO_MAX_WAVES_PER_CU, WavesPerCU))
      return Err;
    HardwareParallelism = ComputeUnits * WavesPerCU;

    // Get maximum size of any device queues and maximum number of queues.
    uint32_t MaxQueueSize;
    if (auto Err = getDeviceAttr(HSA_AGENT_INFO_QUEUE_MAX_SIZE, MaxQueueSize))
      return Err;

    uint32_t MaxQueues;
    if (auto Err = getDeviceAttr(HSA_AGENT_INFO_QUEUES_MAX, MaxQueues))
      return Err;

    // Compute the number of queues and their size.
    OMPX_NumQueues = std::max(1U, std::min(OMPX_NumQueues.get(), MaxQueues));
    OMPX_QueueSize = std::min(OMPX_QueueSize.get(), MaxQueueSize);

    // Initialize stream pool.
    if (auto Err = AMDGPUStreamManager.init(OMPX_InitialNumStreams,
                                            OMPX_NumQueues, OMPX_QueueSize))
      return Err;

    // Initialize event pool.
    if (auto Err = AMDGPUEventManager.init(OMPX_InitialNumEvents))
      return Err;

    // Initialize signal pool.
    if (auto Err = AMDGPUSignalManager.init(OMPX_InitialNumSignals))
      return Err;

    // Detect if XNACK is enabled
    SmallVector<SmallString<32>> Targets;
    if (auto Err = hsa_utils::getTargetTripleAndFeatures(Agent, Targets))
      return Err;
    if (!Targets.empty() && Targets[0].str().contains("xnack+"))
      IsXnackEnabled = true;

    // detect if device is an APU.
    if (auto Err = checkIfAPU())
      return Err;

    // Retrieve the size of the group memory.
    for (const auto *Pool : AllMemoryPools) {
      if (Pool->isGroup()) {
        if (auto Err = Pool->getAttr(HSA_AMD_MEMORY_POOL_INFO_SIZE,
                                     MaxBlockSharedMemSize))
          return Err;
        break;
      }
    }

    return Plugin::success();
  }

  /// Suggest a virtual address for device memory mapping.
  void *getSuggestedVirtualAddress() override {
    return reinterpret_cast<void *>(0x1534f7e00000ULL);
  }

  /// Allocate \p Size bytes on the device and hints the backend to map it to
  /// virtual address \p VAddr. The function returns the allocated virtual
  /// address. The memory must be deallocated through
  /// GenericDeviceTy::deallocateWithVirtualAddress().
  Expected<void *> allocateWithVirtualAddress(uint64_t Size,
                                              void *VAddr) override {
    uint64_t ExpectedVAddr = 0;
    if (VAddr != nullptr)
      ExpectedVAddr = static_cast<uint64_t>(reinterpret_cast<uintptr_t>(VAddr));

    // Transparently round up to a multiple of the page size.
    auto *Pool = CoarseGrainedMemoryPools[0];
    Size = llvm::alignTo(Size, (uint64_t)Pool->getGranule());

    // Reserve the virtual address range.
    hsa_status_t Status =
        hsa_amd_vmem_address_reserve(&VAddr, Size, ExpectedVAddr, 0);
    if (auto Err = Plugin::check(Status,
                                 "error in hsa_amd_vmem_address_reserve: %s\n"))
      return Err;

    if (ExpectedVAddr != 0 && reinterpret_cast<void *>(ExpectedVAddr) != VAddr)
      ODBG(OLDT_Alloc)
          << "hsa_amd_vmem_address_reserve reserved device virtual address "
          << VAddr << " instead of " << reinterpret_cast<void *>(ExpectedVAddr);

    // Create a handle of the allocation.
    hsa_amd_vmem_alloc_handle_t Handle;
    Status = hsa_amd_vmem_handle_create(Pool->get(), Size, MEMORY_TYPE_PINNED,
                                        0, &Handle);
    if (auto Err =
            Plugin::check(Status, "error in hsa_amd_vmem_handle_create: %s\n"))
      return Err;

    // Map the virtual address range to the memory allocation.
    Status = hsa_amd_vmem_map(VAddr, Size, 0, Handle, 0);
    if (auto Err = Plugin::check(Status, "error in hsa_amd_vmem_map: %s\n"))
      return Err;

    // Set the memory access properties for the allocation.
    hsa_amd_memory_access_desc_t Desc;
    Desc.agent_handle = Agent;
    Desc.permissions = HSA_ACCESS_PERMISSION_RW;
    Status = hsa_amd_vmem_set_access(VAddr, Size, &Desc, 1);
    if (auto Err =
            Plugin::check(Status, "error in hsa_amd_vmem_set_access: %s\n"))
      return Err;

    // Register the virtual address range in the tracker.
    if (auto Err = VMemTracker.registerReservation(VAddr, Size, Handle))
      return Err;

    return VAddr;
  }

  /// Deallocate device memory \p VAddr, which was allocated through
  /// GenericDeviceTy::allocateWithVirtualAddress(), and unmap the virtual
  /// address range.
  Error deallocateWithVirtualAddress(void *VAddr, uint64_t) override {
    // Unregister the virtual address range and obtain the information about
    // the reservation.
    auto InfoOrErr = VMemTracker.unregisterReservation(VAddr);
    if (!InfoOrErr)
      return InfoOrErr.takeError();

    auto [Size, Handle] = *InfoOrErr;

    hsa_status_t Status = hsa_amd_vmem_unmap(VAddr, Size);
    if (auto Err = Plugin::check(Status, "error in hsa_amd_vmem_unmap: %s\n"))
      return Err;

    Status = hsa_amd_vmem_handle_release(Handle);
    if (auto Err =
            Plugin::check(Status, "error in hsa_amd_vmem_handle_release: %s\n"))
      return Err;

    Status = hsa_amd_vmem_address_free(VAddr, Size);
    return Plugin::check(Status, "error in hsa_amd_vmem_address_free: %s\n");
  }

  Error unloadBinaryImpl(DeviceImageTy *Image) override {
    AMDGPUDeviceImageTy &AMDImage = static_cast<AMDGPUDeviceImageTy &>(*Image);

    // Unload the executable of the image.
    if (auto Err = AMDImage.unloadExecutable())
      return Err;

    // Destroy the associated memory and invalidate the object.
    Plugin.free(Image);
    return Error::success();
  }

  /// Deinitialize the device and release its resources.
  Error deinitImpl() override {
    // Deinitialize the stream and event pools.
    if (auto Err = AMDGPUStreamManager.deinit())
      return Err;

    if (auto Err = AMDGPUEventManager.deinit())
      return Err;

    if (auto Err = AMDGPUSignalManager.deinit())
      return Err;

    // Invalidate agent reference.
    Agent = {0};

    return Plugin::success();
  }

  virtual Error callGlobalConstructors(GenericPluginTy &Plugin,
                                       DeviceImageTy &Image) override {
    return callGlobalCtorDtorCommon(Plugin, Image, /*IsCtor=*/true);
  }

  virtual Error callGlobalDestructors(GenericPluginTy &Plugin,
                                      DeviceImageTy &Image) override {
    return callGlobalCtorDtorCommon(Plugin, Image, /*IsCtor=*/false);
  }

  uint64_t getStreamBusyWaitMicroseconds() const { return OMPX_StreamBusyWait; }

  Expected<std::unique_ptr<MemoryBuffer>>
  doJITPostProcessing(std::unique_ptr<MemoryBuffer> MB) const override {

    // TODO: We should try to avoid materialization but there seems to be no
    // good linker interface w/o file i/o.
    SmallString<128> LinkerInputFilePath;
    std::error_code EC = sys::fs::createTemporaryFile("amdgpu-pre-link-jit",
                                                      "o", LinkerInputFilePath);
    if (EC)
      return Plugin::error(ErrorCode::HOST_IO,
                           "failed to create temporary file for linker");

    // Write the file's contents to the output file.
    Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
        FileOutputBuffer::create(LinkerInputFilePath, MB->getBuffer().size());
    if (!OutputOrErr)
      return OutputOrErr.takeError();
    std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
    llvm::copy(MB->getBuffer(), Output->getBufferStart());
    if (Error E = Output->commit())
      return std::move(E);

    SmallString<128> LinkerOutputFilePath;
    EC = sys::fs::createTemporaryFile("amdgpu-pre-link-jit", "so",
                                      LinkerOutputFilePath);
    if (EC)
      return Plugin::error(ErrorCode::HOST_IO,
                           "failed to create temporary file for linker");

    const auto &ErrorOrPath = sys::findProgramByName("lld");
    if (!ErrorOrPath)
      return createStringError(ErrorCode::HOST_TOOL_NOT_FOUND,
                               "failed to find `lld` on the PATH.");

    std::string LLDPath = ErrorOrPath.get();
    INFO(OMP_INFOTYPE_PLUGIN_KERNEL, getDeviceId(),
         "Using `%s` to link JITed amdgcn output.", LLDPath.c_str());

    std::string MCPU = "-plugin-opt=mcpu=" + getComputeUnitKind();
    StringRef Args[] = {LLDPath,
                        "-flavor",
                        "gnu",
                        "--no-undefined",
                        "-shared",
                        MCPU,
                        "-o",
                        LinkerOutputFilePath.data(),
                        LinkerInputFilePath.data()};

    std::string Error;
    int RC = sys::ExecuteAndWait(LLDPath, Args, std::nullopt, {}, 0, 0, &Error);
    if (RC)
      return Plugin::error(ErrorCode::LINK_FAILURE,
                           "linking optimized bitcode failed: %s",
                           Error.c_str());

    auto BufferOrErr = MemoryBuffer::getFileOrSTDIN(LinkerOutputFilePath);
    if (!BufferOrErr)
      return Plugin::error(ErrorCode::HOST_IO,
                           "failed to open temporary file for lld");

    // Clean up the temporary files afterwards.
    if (sys::fs::remove(LinkerOutputFilePath))
      return Plugin::error(ErrorCode::HOST_IO,
                           "failed to remove temporary output file for lld");
    if (sys::fs::remove(LinkerInputFilePath))
      return Plugin::error(ErrorCode::HOST_IO,
                           "failed to remove temporary input file for lld");

    return std::move(*BufferOrErr);
  }

  /// See GenericDeviceTy::getComputeUnitKind().
  std::string getComputeUnitKind() const override { return ComputeUnitKind; }

  /// Returns the clock frequency for the given AMDGPU device.
  uint64_t getClockFrequency() const override { return ClockFrequency; }

  /// Returns the current HSA system timestamp for profiling.
  uint64_t getDeviceTimeStamp() override { return getSystemTimestampInNs(); }

  /// Returns the HSA system timestamp frequency. Zero means unavailable.
  uint64_t getSystemTimestampFrequency() const {
    return SystemTimestampFrequency;
  }

  /// Allocate and construct an AMDGPU kernel.
  Expected<GenericKernelTy &> constructKernel(StringRef Name) override {
    // Allocate and construct the AMDGPU kernel.
    AMDGPUKernelTy *AMDGPUKernel = Plugin.allocate<AMDGPUKernelTy>();
    if (!AMDGPUKernel)
      return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
                           "failed to allocate memory for AMDGPU kernel");

    new (AMDGPUKernel) AMDGPUKernelTy(Name);

    return *AMDGPUKernel;
  }

  /// Set the current context to this device's context. Do nothing since the
  /// AMDGPU devices do not have the concept of contexts.
  Error setContext() override { return Plugin::success(); }

  /// AMDGPU returns the product of the number of compute units and the waves
  /// per compute unit.
  uint64_t getHardwareParallelism() const override {
    return HardwareParallelism;
  }

  /// See allowsDirectWrite() for the condition a transfer must meet in
  /// addition.
  uint64_t getMaxDirectWriteSize() const override {
    return HostWritesDeviceMemory ? MaxDirectWriteSize : 0;
  }

  /// We want to set up the RPC server for host services to the GPU if it is
  /// available.
  bool shouldSetupRPCServer() const override { return true; }

  /// The RPC interface should have enough space for all available parallelism.
  uint64_t requestedRPCPortCount() const override {
    return getHardwareParallelism();
  }

  /// In cases of mixed wave32 and wave64 code we need to over-alloacte memory.
  uint32_t getRPCNumLanes() const override {
    return MaxWavefrontSize ? MaxWavefrontSize : getWarpSize();
  }

  /// Get the stream of the asynchronous info structure or get a new one.
  Error getStream(AsyncInfoWrapperTy &AsyncInfoWrapper,
                  AMDGPUStreamTy *&Stream) {
    auto WrapperStream =
        AsyncInfoWrapper.getOrInitQueue<AMDGPUStreamTy *>(AMDGPUStreamManager);
    if (!WrapperStream)
      return WrapperStream.takeError();
    Stream = *WrapperStream;
    return Plugin::success();
  }

  /// Load the binary image into the device and allocate an image object.
  Expected<DeviceImageTy *>
  loadBinaryImpl(std::unique_ptr<MemoryBuffer> &&TgtImage, int32_t ImageId,
                 PluginContextTy * /*Context*/) override {
    // Allocate and initialize the image object.
    AMDGPUDeviceImageTy *AMDImage = Plugin.allocate<AMDGPUDeviceImageTy>();
    new (AMDImage) AMDGPUDeviceImageTy(ImageId, *this, std::move(TgtImage));

    // Load the HSA executable.
    if (Error Err = AMDImage->loadExecutable(*this))
      return std::move(Err);

    if (uint32_t WFS = AMDImage->getMaxWavefrontSize())
      MaxWavefrontSize = std::max(MaxWavefrontSize, WFS);

    return AMDImage;
  }

  /// Allocate memory on the device or related to the device.
  Expected<void *> allocate(size_t Size, void *, TargetAllocTy Kind,
                            size_t Alignment) override;

  /// Deallocate memory on the device or related to the device.
  Error free(void *TgtPtr, TargetAllocTy Kind) override {
    if (TgtPtr == nullptr)
      return Plugin::success();

    AMDGPUMemoryPoolTy *MemoryPool = nullptr;
    switch (Kind) {
    case TARGET_ALLOC_DEFAULT:
    case TARGET_ALLOC_DEVICE:
      MemoryPool = CoarseGrainedMemoryPools[0];
      break;
    case TARGET_ALLOC_HOST:
      MemoryPool = &HostDevice.getFineGrainedMemoryPool();
      break;
    case TARGET_ALLOC_SHARED:
      MemoryPool = &HostDevice.getFineGrainedMemoryPool();
      break;
    }

    if (!MemoryPool)
      return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
                           "no memory pool for the specified allocation kind");

    // See allocate() for the registration of host / shared memory as pinned
    // memory.
    if (Kind == TARGET_ALLOC_HOST || Kind == TARGET_ALLOC_SHARED)
      if (auto Err = PinnedAllocs.unregisterHostBuffer(TgtPtr))
        return Err;

    if (auto Err = MemoryPool->deallocate(TgtPtr))
      return Err;

    return Plugin::success();
  }

  /// Synchronize current thread with the pending operations on the async info.
  Error synchronizeImpl(__tgt_async_info &AsyncInfo,
                        bool ReleaseQueue) override {
    AMDGPUStreamTy *Stream =
        reinterpret_cast<AMDGPUStreamTy *>(AsyncInfo.Queue);
    assert(Stream && "Invalid stream");

    if (auto Err = Stream->synchronize())
      return Err;

    // Once the stream is synchronized, return it to stream pool and reset
    // AsyncInfo. This is to make sure the synchronization only works for its
    // own tasks.
    if (ReleaseQueue) {
      AsyncInfo.Queue = nullptr;
      return AMDGPUStreamManager.returnResource(Stream);
    }
    return Plugin::success();
  }

  /// Query for the completion of the pending operations on the async info.
  Error queryAsyncImpl(__tgt_async_info &AsyncInfo, bool ReleaseQueue,
                       bool *IsQueueWorkCompleted) override {
    if (IsQueueWorkCompleted)
      *IsQueueWorkCompleted = false;
    AMDGPUStreamTy *Stream =
        reinterpret_cast<AMDGPUStreamTy *>(AsyncInfo.Queue);
    assert(Stream && "Invalid stream");

    auto CompletedOrErr = Stream->query();
    if (!CompletedOrErr)
      return CompletedOrErr.takeError();

    // Return if it the stream did not complete yet.
    if (!(*CompletedOrErr))
      return Plugin::success();

    if (IsQueueWorkCompleted)
      *IsQueueWorkCompleted = true;
    // Once the stream is completed, return it to stream pool and reset
    // AsyncInfo. This is to make sure the synchronization only works for its
    // own tasks.
    if (ReleaseQueue) {
      AsyncInfo.Queue = nullptr;
      return AMDGPUStreamManager.returnResource(Stream);
    }
    return Plugin::success();
  }

  /// Pin the host buffer and return the device pointer that should be used for
  /// device transfers.
  Expected<void *> dataLockImpl(void *HstPtr, int64_t Size) override {
    void *PinnedPtr = nullptr;

    hsa_status_t Status =
        hsa_amd_memory_lock(HstPtr, Size, nullptr, 0, &PinnedPtr);
    if (auto Err = Plugin::check(Status, "error in hsa_amd_memory_lock: %s\n"))
      return std::move(Err);

    return PinnedPtr;
  }

  /// Unpin the host buffer.
  Error dataUnlockImpl(void *HstPtr) override {
    hsa_status_t Status = hsa_amd_memory_unlock(HstPtr);
    return Plugin::check(Status, "error in hsa_amd_memory_unlock: %s\n");
  }

  /// Check through the HSA runtime whether the \p HstPtr buffer is pinned.
  Expected<bool> isPinnedPtrImpl(void *HstPtr, void *&BaseHstPtr,
                                 void *&BaseDevAccessiblePtr,
                                 size_t &BaseSize) const override {
    hsa_amd_pointer_info_t Info;
    Info.size = sizeof(hsa_amd_pointer_info_t);

    hsa_status_t Status = hsa_amd_pointer_info(
        HstPtr, &Info, /*Allocator=*/nullptr, /*num_agents_accessible=*/nullptr,
        /*accessible=*/nullptr);
    if (auto Err = Plugin::check(Status, "error in hsa_amd_pointer_info: %s"))
      return std::move(Err);

    // The buffer may be locked or allocated through HSA allocators. Assume that
    // the buffer is host pinned if the runtime reports a HSA type.
    if (Info.type != HSA_EXT_POINTER_TYPE_LOCKED &&
        Info.type != HSA_EXT_POINTER_TYPE_HSA)
      return false;

    assert(Info.hostBaseAddress && "Invalid host pinned address");
    assert(Info.agentBaseAddress && "Invalid agent pinned address");
    assert(Info.sizeInBytes > 0 && "Invalid pinned allocation size");

    // Save the allocation info in the output parameters.
    BaseHstPtr = Info.hostBaseAddress;
    BaseDevAccessiblePtr = Info.agentBaseAddress;
    BaseSize = Info.sizeInBytes;

    return true;
  }

  /// Whether a transfer of \p Size bytes should be done via direct H2D
  /// transfer.
  bool allowsDirectWrite(const void *Ptr, int64_t Size) const {
    return Size > 0 && uint64_t(Size) <= getMaxDirectWriteSize() &&
           hsa_utils::isHostWritable(Agent, Ptr, Size);
  }

  /// Submit data to the device (host to device transfer).
  Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
                       AsyncInfoWrapperTy &AsyncInfoWrapper,
                       GenericProfilerTy *ProfilerPtr) override {
    AMDGPUStreamTy *Stream = nullptr;
    void *PinnedPtr = nullptr;

    auto ProfilerSpecificData = getOrNullProfilerSpecificData(AsyncInfoWrapper);

    // A synchronous transfer is not ordered with the work on any queue, so the
    // host can write the memory directly.
    if (AsyncInfoWrapper.isSynchronous() && allowsDirectWrite(TgtPtr, Size))
      return hsa_utils::publishToDeviceMemory(TgtPtr, HstPtr, Size,
                                              HostLinkIsXGMI);

    // Use one-step asynchronous operation when host memory is already pinned.
    if (void *PinnedPtr =
            PinnedAllocs.getDeviceAccessiblePtrFromPinnedBuffer(HstPtr)) {
      if (auto Err = getStream(AsyncInfoWrapper, Stream))
        return Err;
      return Stream->pushPinnedMemoryCopyAsync(TgtPtr, PinnedPtr, Size);
    }

    // For large transfers use synchronous behavior.
    if (Size >= OMPX_MaxAsyncCopyBytes) {
      if (AsyncInfoWrapper.hasQueue())
        if (auto Err = synchronize(AsyncInfoWrapper))
          return Err;

      hsa_status_t Status;
      Status = hsa_amd_memory_lock(const_cast<void *>(HstPtr), Size, nullptr, 0,
                                   &PinnedPtr);
      if (auto Err =
              Plugin::check(Status, "error in hsa_amd_memory_lock: %s\n"))
        return Err;

      AMDGPUSignalTy Signal;
      if (auto Err = Signal.init())
        return Err;

      if (auto Err = hsa_utils::asyncMemCopy(useMultipleSdmaEngines(), TgtPtr,
                                             Agent, PinnedPtr, Agent, Size, 0,
                                             nullptr, Signal.get()))
        return Err;

      if (auto Err = Signal.wait(getStreamBusyWaitMicroseconds()))
        return Err;

      if (auto Err = Signal.deinit())
        return Err;

      Status = hsa_amd_memory_unlock(const_cast<void *>(HstPtr));
      return Plugin::check(Status, "error in hsa_amd_memory_unlock: %s\n");
    }

    // Otherwise, use two-step copy with an intermediate pinned host buffer.
    AMDGPUMemoryManagerTy &PinnedMemoryManager =
        HostDevice.getPinnedMemoryManager();
    if (auto Err = PinnedMemoryManager.allocate(Size, &PinnedPtr))
      return Err;

    if (auto Err = getStream(AsyncInfoWrapper, Stream))
      return Err;

    return Stream->pushMemoryCopyH2DAsync(TgtPtr, HstPtr, PinnedPtr, Size,
                                          PinnedMemoryManager,
                                          ProfilerSpecificData, ProfilerPtr);
  }

  /// Retrieve data from the device (device to host transfer).
  Error dataRetrieveImpl(void *HstPtr, const void *TgtPtr, int64_t Size,
                         AsyncInfoWrapperTy &AsyncInfoWrapper,
                         GenericProfilerTy *ProfilerPtr) override {
    AMDGPUStreamTy *Stream = nullptr;
    void *PinnedPtr = nullptr;

    auto ProfilerSpecificData = getOrNullProfilerSpecificData(AsyncInfoWrapper);

    // Use one-step asynchronous operation when host memory is already pinned.
    if (void *PinnedPtr =
            PinnedAllocs.getDeviceAccessiblePtrFromPinnedBuffer(HstPtr)) {
      if (auto Err = getStream(AsyncInfoWrapper, Stream))
        return Err;

      return Stream->pushPinnedMemoryCopyAsync(PinnedPtr, TgtPtr, Size);
    }

    // For large transfers use synchronous behavior.
    if (Size >= OMPX_MaxAsyncCopyBytes) {
      if (AsyncInfoWrapper.hasQueue())
        if (auto Err = synchronize(AsyncInfoWrapper))
          return Err;

      hsa_status_t Status;
      Status = hsa_amd_memory_lock(const_cast<void *>(HstPtr), Size, nullptr, 0,
                                   &PinnedPtr);
      if (auto Err =
              Plugin::check(Status, "error in hsa_amd_memory_lock: %s\n"))
        return Err;

      AMDGPUSignalTy Signal;
      if (auto Err = Signal.init())
        return Err;

      if (auto Err = hsa_utils::asyncMemCopy(useMultipleSdmaEngines(),
                                             PinnedPtr, Agent, TgtPtr, Agent,
                                             Size, 0, nullptr, Signal.get()))
        return Err;

      if (auto Err = Signal.wait(getStreamBusyWaitMicroseconds()))
        return Err;

      if (auto Err = Signal.deinit())
        return Err;

      Status = hsa_amd_memory_unlock(const_cast<void *>(HstPtr));
      return Plugin::check(Status, "error in hsa_amd_memory_unlock: %s\n");
    }

    // Otherwise, use two-step copy with an intermediate pinned host buffer.
    AMDGPUMemoryManagerTy &PinnedMemoryManager =
        HostDevice.getPinnedMemoryManager();
    if (auto Err = PinnedMemoryManager.allocate(Size, &PinnedPtr))
      return Err;

    if (auto Err = getStream(AsyncInfoWrapper, Stream))
      return Err;

    return Stream->pushMemoryCopyD2HAsync(HstPtr, TgtPtr, PinnedPtr, Size,
                                          PinnedMemoryManager,
                                          ProfilerSpecificData, ProfilerPtr);
  }

  Error dataMemcpyImpl(void *DstPtr, const void *SrcPtr, int64_t Size,
                       AsyncInfoWrapperTy &AsyncInfoWrapper) override {
    struct MemcpyArgsTy {
      void *DstPtr;
      const void *SrcPtr;
      int64_t Size;
    };

    AMDGPUStreamTy *Stream = nullptr;
    if (auto Err = getStream(AsyncInfoWrapper, Stream))
      return Err;

    auto Args =
        std::make_unique<MemcpyArgsTy>(MemcpyArgsTy{DstPtr, SrcPtr, Size});
    if (auto Err = Stream->pushHostCallback(
            [](void *Data) {
              std::unique_ptr<MemcpyArgsTy> Args(
                  static_cast<MemcpyArgsTy *>(Data));
              std::memcpy(Args->DstPtr, Args->SrcPtr, Args->Size);
            },
            Args.get()))
      return Err;

    Args.release();
    return Plugin::success();
  }

  /// Exchange data between two devices within the plugin.
  Error dataExchangeImpl(const void *SrcPtr, GenericDeviceTy &DstGenericDevice,
                         void *DstPtr, int64_t Size,
                         AsyncInfoWrapperTy &AsyncInfoWrapper,
                         GenericProfilerTy *ProfilerPtr) override {
    AMDGPUDeviceTy &DstDevice = static_cast<AMDGPUDeviceTy &>(DstGenericDevice);

    auto ProfilerSpecificData = getOrNullProfilerSpecificData(AsyncInfoWrapper);

    // For large transfers use synchronous behavior.
    if (Size >= OMPX_MaxAsyncCopyBytes) {
      if (AsyncInfoWrapper.hasQueue())
        if (auto Err = synchronize(AsyncInfoWrapper))
          return Err;

      AMDGPUSignalTy Signal;
      if (auto Err = Signal.init())
        return Err;

      if (auto Err = hsa_utils::asyncMemCopy(
              useMultipleSdmaEngines(), DstPtr, DstDevice.getAgent(), SrcPtr,
              getAgent(), (uint64_t)Size, 0, nullptr, Signal.get()))
        return Err;

      if (auto Err = Signal.wait(getStreamBusyWaitMicroseconds()))
        return Err;

      return Signal.deinit();
    }

    AMDGPUStreamTy *Stream = nullptr;
    if (auto Err = getStream(AsyncInfoWrapper, Stream))
      return Err;
    if (Size <= 0)
      return Plugin::success();

    return Stream->pushMemoryCopyD2DAsync(DstPtr, DstDevice.getAgent(), SrcPtr,
                                          getAgent(), (uint64_t)Size,
                                          ProfilerSpecificData, ProfilerPtr);
  }

  Error dataFillImpl(void *TgtPtr, const void *PatternPtr, int64_t PatternSize,
                     int64_t Size,
                     AsyncInfoWrapperTy &AsyncInfoWrapper) override {
    // Fast case, where we can use the 4 byte hsa_amd_memory_fill
    if (Size % 4 == 0 &&
        (PatternSize == 4 || PatternSize == 2 || PatternSize == 1)) {
      uint32_t Pattern;
      if (PatternSize == 1) {
        auto *Byte = reinterpret_cast<const uint8_t *>(PatternPtr);
        Pattern = *Byte | *Byte << 8 | *Byte << 16 | *Byte << 24;
      } else if (PatternSize == 2) {
        auto *Word = reinterpret_cast<const uint16_t *>(PatternPtr);
        Pattern = *Word | (*Word << 16);
      } else if (PatternSize == 4) {
        Pattern = *reinterpret_cast<const uint32_t *>(PatternPtr);
      } else {
        // Shouldn't be here if the pattern size is outwith those values
        llvm_unreachable("Invalid pattern size");
      }

      if (hasPendingWorkImpl(AsyncInfoWrapper)) {
        AMDGPUStreamTy *Stream = nullptr;
        if (auto Err = getStream(AsyncInfoWrapper, Stream))
          return Err;

        struct MemFillArgsTy {
          void *Dst;
          uint32_t Pattern;
          int64_t Size;
        };
        auto *Args = new MemFillArgsTy{TgtPtr, Pattern, Size / 4};
        auto Fill = [](void *Data) {
          MemFillArgsTy *Args = reinterpret_cast<MemFillArgsTy *>(Data);
          assert(Args && "Invalid arguments");

          auto Status =
              hsa_amd_memory_fill(Args->Dst, Args->Pattern, Args->Size);
          delete Args;
          auto Err =
              Plugin::check(Status, "error in hsa_amd_memory_fill: %s\n");
          if (Err) {
            FATAL_MESSAGE(1, "error performing async fill: %s",
                          toString(std::move(Err)).data());
          }
        };

        // hsa_amd_memory_fill doesn't signal completion using a signal, so use
        // the existing host callback logic to handle that instead
        return Stream->pushHostCallback(Fill, Args);
      }
      // If there is no pending work, do the fill synchronously
      auto Status = hsa_amd_memory_fill(TgtPtr, Pattern, Size / 4);
      return Plugin::check(Status, "error in hsa_amd_memory_fill: %s\n");
    }

    // Slow case; allocate an appropriate memory size and enqueue copies
    void *PinnedPtr = nullptr;
    AMDGPUMemoryManagerTy &PinnedMemoryManager =
        HostDevice.getPinnedMemoryManager();
    if (auto Err = PinnedMemoryManager.allocate(Size, &PinnedPtr))
      return Err;

    AMDGPUStreamTy *Stream = nullptr;
    if (auto Err = getStream(AsyncInfoWrapper, Stream))
      return Err;

    return Stream->pushMemoryCopyH2DAsync(
        TgtPtr, PatternPtr, PinnedPtr, PatternSize, PinnedMemoryManager,
        /*ProfilerSpecificData=*/nullptr, /*ProfilerPtr=*/nullptr,
        Size / PatternSize);
  }

  Expected<void *> getNativeHandle() override {
    return reinterpret_cast<void *>(Agent.handle);
  }

  interop_spec_t selectInteropPreference(int32_t InteropType,
                                         int32_t NumPrefers,
                                         interop_spec_t *Prefers) override {
    // TODO: update once targetsync is supported
    if (InteropType == kmp_interop_type_target)
      return interop_spec_t{tgt_fr_hsa, {false, 0}, 0};
    return interop_spec_t{tgt_fr_none, {false, 0}, 0};
  }

  Expected<omp_interop_val_t *>
  createInterop(int32_t InteropType, interop_spec_t &InteropSpec) override {
    auto *Ret = new omp_interop_val_t(
        DeviceId, static_cast<kmp_interop_type_t>(InteropType));
    Ret->fr_id = tgt_fr_hsa;
    Ret->vendor_id = omp_vendor_amd;

    // TODO: implement targetsync support

    Ret->device_info.Platform = nullptr;
    Ret->device_info.Device = reinterpret_cast<void *>(Agent.handle);
    Ret->device_info.Context = nullptr;

    return Ret;
  }

  Error releaseInterop(omp_interop_val_t *Interop) override {
    if (Interop)
      delete Interop;
    return Plugin::success();
  }

  Error enqueueHostCallImpl(AMDGPUStreamTy::HostFnType Callback, void *UserData,
                            AsyncInfoWrapperTy &AsyncInfo) override {
    AMDGPUStreamTy *Stream = nullptr;
    if (auto Err = getStream(AsyncInfo, Stream))
      return Err;

    return Stream->pushHostCallback(Callback, UserData);
  };

  /// Create an event.
  Error createEventImpl(void **EventPtrStorage, bool EnableProfiling) override {
    AMDGPUEventTy **Event = reinterpret_cast<AMDGPUEventTy **>(EventPtrStorage);
    if (auto Err = AMDGPUEventManager.getResource(*Event))
      return Err;
    return (*Event)->resetState();
  }

  /// Destroy a previously created event.
  Error destroyEventImpl(void *EventPtr, bool EnableProfiling) override {
    AMDGPUEventTy *Event = reinterpret_cast<AMDGPUEventTy *>(EventPtr);
    assert(Event && "Invalid event");

    if (auto Err = Event->resetState())
      return Err;

    return AMDGPUEventManager.returnResource(Event);
  }

  /// Record the event.
  Error recordEventImpl(void *EventPtr, AsyncInfoWrapperTy &AsyncInfoWrapper,
                        bool EnableProfiling) override {
    AMDGPUEventTy *Event = reinterpret_cast<AMDGPUEventTy *>(EventPtr);
    assert(Event && "Invalid event");

    AMDGPUStreamTy *Stream = nullptr;
    if (auto Err = getStream(AsyncInfoWrapper, Stream))
      return Err;

    return Event->record(*Stream, EnableProfiling);
  }

  /// Make the stream wait on the event.
  Error waitEventImpl(void *EventPtr,
                      AsyncInfoWrapperTy &AsyncInfoWrapper) override {
    AMDGPUEventTy *Event = reinterpret_cast<AMDGPUEventTy *>(EventPtr);

    AMDGPUStreamTy *Stream = nullptr;
    if (auto Err = getStream(AsyncInfoWrapper, Stream))
      return Err;

    return Event->wait(*Stream);
  }

  Expected<bool> hasPendingWorkImpl(AsyncInfoWrapperTy &AsyncInfo) override {
    auto *Stream = AsyncInfo.getQueueAs<AMDGPUStreamTy *>();
    if (!Stream)
      return false;

    auto Query = Stream->query();
    if (Query)
      return !*Query;
    return Query.takeError();
  }

  Expected<bool> isEventCompleteImpl(void *EventPtr,
                                     AsyncInfoWrapperTy &AsyncInfo) override {
    AMDGPUEventTy *Event = reinterpret_cast<AMDGPUEventTy *>(EventPtr);
    auto *Stream = AsyncInfo.getQueueAs<AMDGPUStreamTy *>();
    return Stream && Stream->isEventComplete(*Event);
  }

  /// Synchronize the current thread with the event.
  Error syncEventImpl(void *EventPtr) override {
    AMDGPUEventTy *Event = reinterpret_cast<AMDGPUEventTy *>(EventPtr);
    return Event->sync();
  }

  /// Get the elapsed time in milliseconds between two events.
  Expected<float> getEventElapsedTimeImpl(void *StartEventPtr,
                                          void *EndEventPtr) override {
    AMDGPUEventTy *StartEvent =
        reinterpret_cast<AMDGPUEventTy *>(StartEventPtr);
    AMDGPUEventTy *EndEvent = reinterpret_cast<AMDGPUEventTy *>(EndEventPtr);

    if (!StartEvent || !EndEvent)
      return Plugin::error(ErrorCode::INVALID_ARGUMENT, "invalid event handle");

    return StartEvent->getElapsedTime(*EndEvent);
  }

  /// Print information about the device.
  Expected<InfoTreeNode> obtainInfoImpl() override {
    char TmpChar[1000];
    const char *TmpCharPtr = "Unknown";
    uint16_t Major, Minor;
    uint32_t TmpUInt, TmpUInt2;
    uint32_t CacheSize[4];
    size_t TmpSt;
    bool TmpBool;
    uint16_t WorkgrpMaxDim[3];
    hsa_dim3_t GridMaxDim;
    hsa_status_t Status, Status2;
    InfoTreeNode Info;

    Status = hsa_system_get_info(HSA_SYSTEM_INFO_VERSION_MAJOR, &Major);
    Status2 = hsa_system_get_info(HSA_SYSTEM_INFO_VERSION_MINOR, &Minor);
    if (Status == HSA_STATUS_SUCCESS && Status2 == HSA_STATUS_SUCCESS)
      Info.add("HSA Runtime Version",
               std::to_string(Major) + "." + std::to_string(Minor), "",
               DeviceInfo::DRIVER_VERSION);
    if (auto Version = hsa_utils::getAMDExtensionVersion())
      Info.add("HSA AMD Extension Version", Version->getAsString());

    Info.add("HSA OpenMP Device Number", DeviceId);

    Status = getDeviceAttrRaw(HSA_AMD_AGENT_INFO_PRODUCT_NAME, TmpChar);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Product Name", TmpChar, "", DeviceInfo::PRODUCT_NAME);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_NAME, TmpChar);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Device Name", TmpChar, "", DeviceInfo::NAME);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_VENDOR_NAME, TmpChar);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Vendor Name", TmpChar, "", DeviceInfo::VENDOR);

    Info.add("Vendor ID", uint64_t{4130}, "", DeviceInfo::VENDOR_ID);

    hsa_machine_model_t MachineModel;
    Status = getDeviceAttrRaw(HSA_AGENT_INFO_MACHINE_MODEL, MachineModel);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Memory Address Size",
               uint64_t{MachineModel == HSA_MACHINE_MODEL_SMALL ? 32u : 64u},
               "bits", DeviceInfo::ADDRESS_BITS);

    hsa_device_type_t DevType;
    Status = getDeviceAttrRaw(HSA_AGENT_INFO_DEVICE, DevType);
    if (Status == HSA_STATUS_SUCCESS) {
      switch (static_cast<int>(DevType)) {
      case HSA_DEVICE_TYPE_CPU:
        TmpCharPtr = "CPU";
        break;
      case HSA_DEVICE_TYPE_GPU:
        TmpCharPtr = "GPU";
        break;
      case HSA_DEVICE_TYPE_DSP:
        TmpCharPtr = "DSP";
        break;
      default:
        TmpCharPtr = "Unknown";
        break;
      }
      Info.add("Device Type", TmpCharPtr);
    }

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_QUEUES_MAX, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Max Queues", TmpUInt);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_QUEUE_MIN_SIZE, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Queue Min Size", TmpUInt);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_QUEUE_MAX_SIZE, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Queue Max Size", TmpUInt);

    // FIXME: This is deprecated according to HSA documentation. But using
    // hsa_agent_iterate_caches and hsa_cache_get_info breaks execution during
    // runtime.
    Status = getDeviceAttrRaw(HSA_AGENT_INFO_CACHE_SIZE, CacheSize);
    if (Status == HSA_STATUS_SUCCESS) {
      auto &Cache = *Info.add("Cache");

      for (int I = 0; I < 4; I++)
        if (CacheSize[I])
          Cache.add("L" + std::to_string(I), CacheSize[I]);
    }

    Status = getDeviceAttrRaw(HSA_AMD_AGENT_INFO_CACHELINE_SIZE, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Cacheline Size", TmpUInt);

    Info.add("Max Shared Memory per Work Group", MaxBlockSharedMemSize, "bytes",
             DeviceInfo::WORK_GROUP_LOCAL_MEM_SIZE);

    Status = getDeviceAttrRaw(HSA_AMD_AGENT_INFO_MAX_CLOCK_FREQUENCY, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Max Clock Freq", TmpUInt, "MHz",
               DeviceInfo::MAX_CLOCK_FREQUENCY);

    Status = getDeviceAttrRaw(HSA_AMD_AGENT_INFO_MEMORY_MAX_FREQUENCY, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Max Memory Clock Freq", TmpUInt, "MHz",
               DeviceInfo::MEMORY_CLOCK_RATE);

    Status = getDeviceAttrRaw(HSA_AMD_AGENT_INFO_COMPUTE_UNIT_COUNT, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Compute Units", TmpUInt, "", DeviceInfo::NUM_COMPUTE_UNITS);

    Status = getDeviceAttrRaw(HSA_AMD_AGENT_INFO_NUM_SIMDS_PER_CU, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("SIMD per CU", TmpUInt);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_FAST_F16_OPERATION, TmpBool);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Fast F16 Operation", TmpBool);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_WAVEFRONT_SIZE, TmpUInt2);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Number of Lanes", TmpUInt2, "", DeviceInfo::NUM_LANES);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_WORKGROUP_MAX_SIZE, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Workgroup Max Size", TmpUInt, "",
               DeviceInfo::MAX_WORK_GROUP_SIZE);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_WORKGROUP_MAX_DIM, WorkgrpMaxDim);
    if (Status == HSA_STATUS_SUCCESS) {
      auto &MaxSize =
          *Info.add("Workgroup Max Size per Dimension", std::monostate{}, "",
                    DeviceInfo::MAX_WORK_GROUP_SIZE_PER_DIMENSION);
      MaxSize.add("x", WorkgrpMaxDim[0]);
      MaxSize.add("y", WorkgrpMaxDim[1]);
      MaxSize.add("z", WorkgrpMaxDim[2]);
    }

    Status = getDeviceAttrRaw(
        (hsa_agent_info_t)HSA_AMD_AGENT_INFO_MAX_WAVES_PER_CU, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS) {
      Info.add("Max Waves Per CU", TmpUInt);
      Info.add("Max Work-item Per CU", TmpUInt * TmpUInt2);
    }

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_GRID_MAX_SIZE, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Grid Max Size", TmpUInt, "", DeviceInfo::MAX_WORK_SIZE);

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_GRID_MAX_DIM, GridMaxDim);
    if (Status == HSA_STATUS_SUCCESS) {
      auto &MaxDim = *Info.add("Grid Max Size per Dimension", std::monostate{},
                               "", DeviceInfo::MAX_WORK_SIZE_PER_DIMENSION);
      MaxDim.add("x", GridMaxDim.x);
      MaxDim.add("y", GridMaxDim.y);
      MaxDim.add("z", GridMaxDim.z);
    }

    Status = getDeviceAttrRaw(HSA_AGENT_INFO_FBARRIER_MAX_SIZE, TmpUInt);
    if (Status == HSA_STATUS_SUCCESS)
      Info.add("Max fbarriers/Workgrp", TmpUInt);

    auto &RootPool = *Info.add("Memory Pools");
    for (AMDGPUMemoryPoolTy *Pool : AllMemoryPools) {
      std::string TmpStr, TmpStr2;

      if (Pool->isGlobal())
        TmpStr = "Global";
      else if (Pool->isReadOnly())
        TmpStr = "ReadOnly";
      else if (Pool->isPrivate())
        TmpStr = "Private";
      else if (Pool->isGroup())
        TmpStr = "Group";
      else
        TmpStr = "Unknown";

      auto &PoolNode = *RootPool.add(std::string("Pool ") + TmpStr);

      if (Pool->isGlobal()) {
        if (Pool->isFineGrained())
          TmpStr2 += "Fine Grained ";
        if (Pool->isCoarseGrained())
          TmpStr2 += "Coarse Grained ";
        if (Pool->supportsKernelArgs())
          TmpStr2 += "Kernarg ";

        PoolNode.add("Flags", TmpStr2);
      }

      Status = Pool->getAttrRaw(HSA_AMD_MEMORY_POOL_INFO_SIZE, TmpSt);
      if (Status == HSA_STATUS_SUCCESS)
        PoolNode.add(
            "Size", TmpSt, "bytes",
            (Pool->isGlobal() && Pool->isCoarseGrained())
                ? std::optional<DeviceInfo>{DeviceInfo::GLOBAL_MEM_SIZE}
                : std::nullopt);

      Status = Pool->getAttrRaw(HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALLOWED,
                                TmpBool);
      if (Status == HSA_STATUS_SUCCESS)
        PoolNode.add("Allocatable", TmpBool);

      PoolNode.add("Runtime Alloc Granule", Pool->getGranule(), "bytes");

      Status = Pool->getAttrRaw(
          HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALIGNMENT, TmpSt);
      if (Status == HSA_STATUS_SUCCESS)
        PoolNode.add("Runtime Alloc Alignment", TmpSt, "bytes");

      Status =
          Pool->getAttrRaw(HSA_AMD_MEMORY_POOL_INFO_ACCESSIBLE_BY_ALL, TmpBool);
      if (Status == HSA_STATUS_SUCCESS)
        PoolNode.add("Accessible by all", TmpBool);
    }

    auto &ISAs = *Info.add("ISAs");
    auto Err = hsa_utils::iterateAgentISAs(getAgent(), [&](hsa_isa_t ISA) {
      Status = hsa_isa_get_info_alt(ISA, HSA_ISA_INFO_NAME, TmpChar);
      if (Status == HSA_STATUS_SUCCESS)
        ISAs.add("Name", TmpChar);

      return Status;
    });

    // Silently consume the error.
    if (Err)
      consumeError(std::move(Err));

    ol_device_fp_capability_flags_t FPFlags =
        OL_DEVICE_FP_CAPABILITY_FLAG_CORRECTLY_ROUNDED_DIVIDE_SQRT |
        OL_DEVICE_FP_CAPABILITY_FLAG_ROUND_TO_NEAREST |
        OL_DEVICE_FP_CAPABILITY_FLAG_ROUND_TO_ZERO |
        OL_DEVICE_FP_CAPABILITY_FLAG_ROUND_TO_INF |
        OL_DEVICE_FP_CAPABILITY_FLAG_INF_NAN |
        OL_DEVICE_FP_CAPABILITY_FLAG_DENORM | OL_DEVICE_FP_CAPABILITY_FLAG_FMA;

    Info.add("Single FP Support", true, "", DeviceInfo::SINGLE_FP_SUPPORT);
    Info.add("Single FP Capabilities", FPFlags, "",
             DeviceInfo::SINGLE_FP_CONFIG);

    Info.add("Double FP Support", true, "", DeviceInfo::DOUBLE_FP_SUPPORT);
    Info.add("Double FP Capabilities", FPFlags, "",
             DeviceInfo::DOUBLE_FP_CONFIG);

    // TODO: Use HSA_AGENT_INFO_FAST_F16_OPERATION to detect FP16 support.
    Info.add("Half FP Support", false, "", DeviceInfo::HALF_FP_SUPPORT);
    Info.add("Half FP Capabilities", ol_device_fp_capability_flags_t{0}, "",
             DeviceInfo::HALF_FP_CONFIG);

    return Info;
  }

  /// Returns true if auto zero-copy the best configuration for the current
  /// arch.
  /// On AMDGPUs, automatic zero-copy is turned on
  /// when running on an APU with XNACK (unified memory) support
  /// enabled. On discrete GPUs, automatic zero-copy is triggered
  /// if the user sets the environment variable OMPX_APU_MAPS=1
  /// and if XNACK is enabled. The rationale is that zero-copy
  /// is the best configuration (performance, memory footprint) on APUs,
  /// while it is often not the best on discrete GPUs.
  /// XNACK can be enabled with a kernel boot parameter or with
  /// the HSA_XNACK environment variable.
  bool useAutoZeroCopyImpl() override {
    return ((IsAPU || OMPX_ApuMaps) && IsXnackEnabled);
  }

  Expected<bool> isAccessiblePtrImpl(const void *Ptr, size_t Size) override {
    hsa_amd_pointer_info_t Info;
    Info.size = sizeof(hsa_amd_pointer_info_t);

    hsa_agent_t *Agents = nullptr;
    uint32_t Count = 0;
    hsa_status_t Status =
        hsa_amd_pointer_info(Ptr, &Info, malloc, &Count, &Agents);

    if (auto Err = Plugin::check(Status, "error in hsa_amd_pointer_info: %s"))
      return std::move(Err);

    // Checks if the pointer is known by HSA and accessible by the device
    for (uint32_t i = 0; i < Count; i++) {
      if (Agents[i].handle == getAgent().handle)
        return Info.sizeInBytes >= Size;
    }

    // If the pointer is unknown to HSA it's assumed a host pointer
    // in that case the device can access it on unified memory support is
    // enabled
    return IsXnackEnabled;
  }

  /// Getters and setters for stack and heap sizes.
  Error getDeviceStackSize(uint64_t &Value) override {
    Value = StackSize;
    return Plugin::success();
  }
  Error setDeviceStackSize(uint64_t Value) override {
    StackSize = Value;
    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 getting device memory size: %s");
      }
    }
    return Plugin::error(ErrorCode::UNSUPPORTED,
                         "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>
  hsa_status_t getDeviceAttrRaw(uint32_t Kind, Ty &Value) {
    return hsa_agent_get_info(Agent, (hsa_agent_info_t)Kind, &Value);
  }

  /// Get the device agent.
  hsa_agent_t getAgent() const override { return Agent; }

  /// Get the host device nearest to the device.
  AMDHostDeviceTy &getHostDevice() { return HostDevice; }

  /// Get the signal manager.
  AMDGPUSignalManagerTy &getSignalManager() { return AMDGPUSignalManager; }

  /// Retrieve and construct all memory pools of the device agent.
  Error retrieveAllMemoryPools() override {
    // Iterate through the available pools of the device agent.
    return hsa_utils::iterateAgentMemoryPools(
        Agent, [&](hsa_amd_memory_pool_t HSAMemoryPool) {
          AMDGPUMemoryPoolTy *MemoryPool =
              Plugin.allocate<AMDGPUMemoryPoolTy>();
          new (MemoryPool) AMDGPUMemoryPoolTy(HSAMemoryPool);
          AllMemoryPools.push_back(MemoryPool);
          return HSA_STATUS_SUCCESS;
        });
  }

  bool useMultipleSdmaEngines() const { return OMPX_UseMultipleSdmaEngines; }

private:
  using AMDGPUEventRef = AMDGPUResourceRef<AMDGPUEventTy>;
  using AMDGPUEventManagerTy = GenericDeviceResourceManagerTy<AMDGPUEventRef>;

  /// Common method to invoke a single threaded constructor or destructor
  /// kernel by name.
  Error callGlobalCtorDtorCommon(GenericPluginTy &Plugin, DeviceImageTy &Image,
                                 bool IsCtor) {
    const char *KernelName =
        IsCtor ? "amdgcn.device.init" : "amdgcn.device.fini";
    // Perform a quick check for the named kernel in the image. The kernel
    // should be created by the 'amdgpu-lower-ctor-dtor' pass.
    GenericGlobalHandlerTy &Handler = Plugin.getGlobalHandler();
    if (!Handler.isSymbolInImage(*this, Image, KernelName))
      return Plugin::success();

    // Allocate and construct the AMDGPU kernel.
    AMDGPUKernelTy AMDGPUKernel(KernelName);
    if (auto Err = AMDGPUKernel.init(*this, Image))
      return Err;

    AsyncInfoWrapperTy AsyncInfoWrapper(*this, nullptr);

    KernelLaunchArgsTy LaunchArgs = {};
    uint32_t NumBlocksAndThreads[3] = {1u, 1u, 1u};
    auto Err = AMDGPUKernel.launchImpl(
        *this, NumBlocksAndThreads, NumBlocksAndThreads, 0, LaunchArgs,
        AsyncInfoWrapper, /*ProfilerPtr=*/nullptr);

    AsyncInfoWrapper.finalize(Err);
    return Err;
  }

  /// Detect if current architecture is an APU.
  Error checkIfAPU() {
    uint8_t MemoryProperties[8];
    if (auto Err = getDeviceAttr(HSA_AMD_AGENT_INFO_MEMORY_PROPERTIES,
                                 MemoryProperties)) {
      IsAPU = false;
      ODBG(OLDT_Init) << "HSA_AMD_AGENT_INFO_MEMORY_PROPERTIES is unavailable, "
                         "assuming not APU";
      return Plugin::success();
    }

    IsAPU = hsa_flag_isset64(MemoryProperties,
                             HSA_AMD_MEMORY_PROPERTY_AGENT_IS_APU);

    return Plugin::success();
  }

  bool checkIfCoarseGrainMemoryNearOrAbove64GB() {
    for (AMDGPUMemoryPoolTy *Pool : AllMemoryPools) {
      if (!Pool->isGlobal() || !Pool->isCoarseGrained())
        continue;
      uint64_t Value;
      hsa_status_t Status =
          Pool->getAttrRaw(HSA_AMD_MEMORY_POOL_INFO_SIZE, Value);
      if (Status != HSA_STATUS_SUCCESS)
        continue;
      constexpr uint64_t Almost64Gig = 0xFF0000000;
      if (Value >= Almost64Gig)
        return true;
    }
    return false; // CoarseGrain pool w/ 64GB or more capacity not found
  }

  size_t getMemoryManagerSizeThreshold() override {
    // Targeting high memory capacity GPUs such as
    // data center GPUs.
    if (checkIfCoarseGrainMemoryNearOrAbove64GB()) {
      // Set GenericDeviceTy::MemoryManager's Threshold to 3GiB,
      // if threshold is not already set by ENV var
      // LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD.
      // This MemoryManager is used for omp_target_alloc(), OpenMP
      // (non-usm) map clause, etc.
      //
      // Ideally, this kind of pooling is best performed at
      // a common level (e.g, user side of HSA) between OpenMP and HIP
      // but that feature does not exist (yet).
      return 3ul * 1024 * 1024 * 1024 /* 3 GiB */;
    }
    return 0;
  }

  /// Envar for controlling the number of HSA queues per device. High number of
  /// queues may degrade performance.
  UInt32Envar OMPX_NumQueues;

  /// Envar for controlling the size of each HSA queue. The size is the number
  /// of HSA packets a queue is expected to hold. It is also the number of HSA
  /// packets that can be pushed into each queue without waiting the driver to
  /// process them.
  UInt32Envar OMPX_QueueSize;

  /// Envar for controlling the default number of teams relative to the number
  /// of compute units (CUs) the device has:
  ///   #default_teams = OMPX_DefaultTeamsPerCU * #CUs.
  UInt32Envar OMPX_DefaultTeamsPerCU;

  /// Envar specifying the maximum size in bytes where the memory copies are
  /// asynchronous operations. Up to this transfer size, the memory copies are
  /// asynchronous operations pushed to the corresponding stream. For larger
  /// transfers, they are synchronous transfers.
  UInt32Envar OMPX_MaxAsyncCopyBytes;

  /// Envar controlling the initial number of HSA signals per device. There is
  /// one manager of signals per device managing several pre-allocated signals.
  /// These signals are mainly used by AMDGPU streams. If needed, more signals
  /// will be created.
  UInt32Envar OMPX_InitialNumSignals;

  /// Environment variables to set the time to wait in active state before
  /// switching to blocked state. The default 2000000 busywaits for 2 seconds
  /// before going into a blocking HSA wait state. The unit for these variables
  /// are microseconds.
  UInt32Envar OMPX_StreamBusyWait;

  /// Use ROCm 5.7 interface for multiple SDMA engines
  BoolEnvar OMPX_UseMultipleSdmaEngines;

  /// Value of OMPX_APU_MAPS env var used to force
  /// automatic zero-copy behavior on non-APU GPUs.
  BoolEnvar OMPX_ApuMaps;

  /// Stream manager for AMDGPU streams.
  AMDGPUStreamManagerTy AMDGPUStreamManager;

  /// Event manager for AMDGPU events.
  AMDGPUEventManagerTy AMDGPUEventManager;

  /// Signal manager for AMDGPU signals.
  AMDGPUSignalManagerTy AMDGPUSignalManager;

  /// The agent handler corresponding to the device.
  hsa_agent_t Agent;

  /// The GPU architecture.
  std::string ComputeUnitKind;

  /// The frequency of the steady clock inside the device.
  uint64_t ClockFrequency;

  /// The HSA system timestamp frequency reported by the runtime. Zero means
  /// unavailable.
  uint64_t SystemTimestampFrequency = 0;

  /// The total number of concurrent work items that can be running on the GPU.
  uint64_t HardwareParallelism;

  /// The largest wavefront size across all loaded images, used for RPC.
  uint32_t MaxWavefrontSize = 0;

  /// Reference to the host device.
  AMDHostDeviceTy &HostDevice;

  /// Whether the host can map this device's memory (large BAR) and write small
  /// transfers into it directly.
  bool HostWritesDeviceMemory = false;

  /// The maximum size of a direct write. A first heuristic, tune as new
  /// insights arise.
  static constexpr uint64_t MaxDirectWriteSize = 2 * 4096;

  /// Whether the host reaches this device's memory over XGMI rather than PCIe.
  bool HostLinkIsXGMI = false;

  /// The current size of the stack that will be used in cases where it could
  /// not be statically determined.
  uint64_t StackSize = 16 * 1024 /* 16 KB */;

  /// Is the plugin associated with an APU?
  bool IsAPU = false;

  /// True is the system is configured with XNACK-Enabled.
  /// False otherwise.
  bool IsXnackEnabled = false;

  /// Tracker for virtual address reservations.
  VMemTrackerTy<hsa_amd_vmem_alloc_handle_t> VMemTracker;
};

Error AMDGPUDeviceImageTy::loadExecutable(const AMDGPUDeviceTy &Device) {
  hsa_code_object_reader_t Reader;
  hsa_status_t Status =
      hsa_code_object_reader_create_from_memory(getStart(), getSize(), &Reader);
  if (auto Err = Plugin::check(
          Status, "error in hsa_code_object_reader_create_from_memory: %s"))
    return Err;

  Status = hsa_executable_create_alt(
      HSA_PROFILE_FULL, HSA_DEFAULT_FLOAT_ROUNDING_MODE_ZERO, "", &Executable);
  if (auto Err =
          Plugin::check(Status, "error in hsa_executable_create_alt: %s"))
    return Err;

  hsa_loaded_code_object_t Object;
  Status = hsa_executable_load_agent_code_object(Executable, Device.getAgent(),
                                                 Reader, "", &Object);
  if (auto Err = Plugin::check(
          Status, "error in hsa_executable_load_agent_code_object: %s"))
    return Err;

  Status = hsa_executable_freeze(Executable, "");
  if (auto Err = Plugin::check(Status, "error in hsa_executable_freeze: %s"))
    return Err;

  uint32_t Result;
  Status = hsa_executable_validate(Executable, &Result);
  if (auto Err = Plugin::check(Status, "error in hsa_executable_validate: %s"))
    return Err;

  if (Result)
    return Plugin::error(ErrorCode::INVALID_BINARY,
                         "loaded HSA executable does not validate");

  Status = hsa_code_object_reader_destroy(Reader);
  if (auto Err =
          Plugin::check(Status, "error in hsa_code_object_reader_destroy: %s"))
    return Err;

  if (auto Err = hsa_utils::readAMDGPUMetaDataFromImage(
          getMemoryBuffer(), KernelInfoMap, ELFABIVersion))
    return Err;

  return Plugin::success();
}

Expected<hsa_executable_symbol_t>
AMDGPUDeviceImageTy::findDeviceSymbol(GenericDeviceTy &Device,
                                      StringRef SymbolName) const {

  AMDGPUDeviceTy &AMDGPUDevice = static_cast<AMDGPUDeviceTy &>(Device);
  hsa_agent_t Agent = AMDGPUDevice.getAgent();

  hsa_executable_symbol_t Symbol;
  hsa_status_t Status = hsa_executable_get_symbol_by_name(
      Executable, SymbolName.data(), &Agent, &Symbol);
  if (auto Err = Plugin::check(
          Status, "error in hsa_executable_get_symbol_by_name(%s): %s",
          SymbolName.data()))
    return std::move(Err);

  return Symbol;
}

template <typename ResourceTy>
Error AMDGPUResourceRef<ResourceTy>::create(GenericDeviceTy &Device) {
  if (Resource)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "creating an existing resource");

  AMDGPUDeviceTy &AMDGPUDevice = static_cast<AMDGPUDeviceTy &>(Device);

  Resource = new ResourceTy(AMDGPUDevice);

  return Resource->init();
}

AMDGPUStreamTy::AMDGPUStreamTy(AMDGPUDeviceTy &Device)
    : Agent(Device.getAgent()), Queue(nullptr),
      SignalManager(Device.getSignalManager()), Device(Device),
      // Initialize the std::deque with some empty positions.
      Slots(32), NextSlot(0), SyncCycle(0),
      StreamBusyWaitMicroseconds(Device.getStreamBusyWaitMicroseconds()),
      UseMultipleSdmaEngines(Device.useMultipleSdmaEngines()) {}

Error AMDGPUEventTy::releaseTimingSignal(AMDGPUSignalTy **ReusableSignalPtr) {
  AMDGPUSignalTy *Signal = TimingSignal;
  TimingSignal = nullptr;

  if (!Signal)
    return Plugin::success();

  if (!Signal->decreaseUseCount())
    return Plugin::success();

  if (ReusableSignalPtr) {
    *ReusableSignalPtr = Signal;
    return Plugin::success();
  }

  return Device.getSignalManager().returnResource(Signal);
}

Expected<float> AMDGPUEventTy::getElapsedTime(AMDGPUEventTy &EndEvent) {
  if (this == &EndEvent) {
    std::lock_guard<std::mutex> Lock(Mutex);

    if (!TimingSignal)
      return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                           "event does not have a recorded timing signal");

    if (TimingSignal->load())
      return Plugin::error(ErrorCode::UNKNOWN, "event timing is not ready");

    return 0.0f;
  }

  const uint64_t TicksPerSecond = Device.getSystemTimestampFrequency();
  if (TicksPerSecond == 0)
    return Plugin::error(ErrorCode::UNSUPPORTED,
                         "HSA system timestamp frequency is unavailable");

  std::scoped_lock<std::mutex, std::mutex> Lock(Mutex, EndEvent.Mutex);

  if (&Device != &EndEvent.Device)
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "events belong to different devices");

  if (!TimingSignal || !EndEvent.TimingSignal)
    return Plugin::error(
        ErrorCode::INVALID_ARGUMENT,
        "one or both events do not have a recorded timing signal");

  if (TimingSignal->load() || EndEvent.TimingSignal->load())
    return Plugin::error(
        ErrorCode::UNKNOWN,
        "timing information is not ready for one or both events");

  hsa_amd_profiling_dispatch_time_t StartTime = {};
  hsa_amd_profiling_dispatch_time_t StopTime = {};

  hsa_status_t Status = hsa_amd_profiling_get_dispatch_time(
      Device.getAgent(), TimingSignal->get(), &StartTime);
  if (auto Err = Plugin::check(
          Status, "error in hsa_amd_profiling_get_dispatch_time: %s"))
    return std::move(Err);

  Status = hsa_amd_profiling_get_dispatch_time(
      EndEvent.Device.getAgent(), EndEvent.TimingSignal->get(), &StopTime);
  if (auto Err = Plugin::check(
          Status, "error in hsa_amd_profiling_get_dispatch_time: %s"))
    return std::move(Err);

  const int64_t DeltaTicks =
      static_cast<int64_t>(StopTime.end) - static_cast<int64_t>(StartTime.end);
  constexpr double MillisecondsPerSecond = 1000.0;

  return static_cast<float>(static_cast<double>(DeltaTicks) *
                            MillisecondsPerSecond /
                            static_cast<double>(TicksPerSecond));
}

/// Class implementing the AMDGPU-specific functionalities of the global
/// handler.
struct AMDGPUGlobalHandlerTy final : public GenericGlobalHandlerTy {
  /// Get the metadata of a global from the device. The name and size of the
  /// global is read from DeviceGlobal and the address of the global is written
  /// to DeviceGlobal.
  Error getGlobalMetadataFromDevice(GenericDeviceTy &Device,
                                    DeviceImageTy &Image,
                                    GlobalTy &DeviceGlobal) override {
    AMDGPUDeviceImageTy &AMDImage = static_cast<AMDGPUDeviceImageTy &>(Image);

    // Find the symbol on the device executable.
    auto SymbolOrErr =
        AMDImage.findDeviceSymbol(Device, DeviceGlobal.getName());
    if (!SymbolOrErr)
      return SymbolOrErr.takeError();

    hsa_executable_symbol_t Symbol = *SymbolOrErr;
    hsa_symbol_kind_t SymbolType;
    hsa_status_t Status;
    uint64_t SymbolAddr;
    uint32_t SymbolSize;

    // Retrieve the type, address and size of the symbol.
    std::pair<hsa_executable_symbol_info_t, void *> RequiredInfos[] = {
        {HSA_EXECUTABLE_SYMBOL_INFO_TYPE, &SymbolType},
        {HSA_EXECUTABLE_SYMBOL_INFO_VARIABLE_ADDRESS, &SymbolAddr},
        {HSA_EXECUTABLE_SYMBOL_INFO_VARIABLE_SIZE, &SymbolSize}};

    for (auto &Info : RequiredInfos) {
      Status = hsa_executable_symbol_get_info(Symbol, Info.first, Info.second);
      if (auto Err = Plugin::check(
              Status, "error in hsa_executable_symbol_get_info: %s"))
        return Err;
    }

    // Check the size of the symbol.
    if (DeviceGlobal.getSize() && SymbolSize != DeviceGlobal.getSize())
      return Plugin::error(
          ErrorCode::INVALID_BINARY,
          "failed to load global '%s' due to size mismatch (%zu != %zu)",
          DeviceGlobal.getName().data(), SymbolSize,
          (size_t)DeviceGlobal.getSize());

    // Store the symbol address and size on the device global metadata.
    DeviceGlobal.setPtr(reinterpret_cast<void *>(SymbolAddr));
    DeviceGlobal.setSize(SymbolSize);

    return Plugin::success();
  }

protected:
  /// Kernels are represented by a kernel descriptor, which is an object named
  /// after the function it describes with a '.kd' suffix.
  std::optional<StringRef> matchSymbol(const ELFSymbolRef &Symbol,
                                       StringRef Name,
                                       SymbolKindTy Kind) override {
    if (Symbol.getELFType() != ELF::STT_OBJECT)
      return std::nullopt;

    StringRef Function = Name;
    bool IsDescriptor =
        Function.consume_back(".kd") && isFunction(Symbol, Function);
    if (IsDescriptor != (Kind == SymbolKindTy::Kernel))
      return std::nullopt;

    return IsDescriptor ? Function : Name;
  }

private:
  /// Returns whether \p Name is a function in the image containing \p Symbol.
  static bool isFunction(const ELFSymbolRef &Symbol, StringRef Name) {
    auto SymbolOrErr = utils::elf::getSymbol(*Symbol.getObject(), Name);
    if (!SymbolOrErr) {
      consumeError(SymbolOrErr.takeError());
      return false;
    }
    return *SymbolOrErr && (*SymbolOrErr)->getELFType() == ELF::STT_FUNC;
  }
};

struct AMDGPUPluginContextTy final : public PluginContextTy {
  using PluginContextTy::PluginContextTy;

  Error initAsyncInfoImpl(GenericDeviceTy &, AsyncInfoWrapperTy &) override {
    // TODO: Implement this function.
    return Plugin::success();
  }

  void trackAllocation(GenericDeviceTy &Device, void *Ptr, int64_t Size,
                       TargetAllocTy Kind) override {
    std::lock_guard<std::mutex> Lock(AllocationsMutex);
    Allocations[Ptr] = {Kind, &Device};
  }

  void untrackAllocation(void *Ptr) override {
    std::lock_guard<std::mutex> Lock(AllocationsMutex);
    Allocations.erase(Ptr);
  }

  Expected<PluginAllocInfoTy> getAllocInfo(const void *Ptr) override;

private:
  // Track each allocation's Kind so we can tell HOST from SHARED — HSA
  // backs both with the same host fine-grained pool.
  // TODO: drop once TARGET_ALLOC_SHARED has its own backing pool.
  struct AllocInfo {
    TargetAllocTy Kind;
    GenericDeviceTy *Device;
  };
  llvm::DenseMap<const void *, AllocInfo> Allocations;
  std::mutex AllocationsMutex;
};

/// Class implementing the AMDGPU-specific functionalities of the plugin.
struct AMDGPUPluginTy final : public GenericPluginTy {
  /// Create an AMDGPU plugin and initialize the AMDGPU driver.
  AMDGPUPluginTy() : GenericPluginTy(getTripleArch()), Initialized(false) {}

  /// This class should not be copied.
  AMDGPUPluginTy(const AMDGPUPluginTy &) = delete;
  AMDGPUPluginTy(AMDGPUPluginTy &&) = delete;

  /// Initialize the plugin and return the number of devices.
  Expected<int32_t> initImpl() override {
    hsa_status_t Status = hsa_init();
    if (Status != HSA_STATUS_SUCCESS) {
      // Cannot call hsa_success_string.
      ODBG(OLDT_Init) << "Failed to initialize AMDGPU's HSA library";
      return 0;
    }

    // The initialization of HSA was successful. It should be safe to call
    // HSA functions from now on, e.g., hsa_shut_down.
    Initialized = true;

    // Register event handler to detect memory errors on the devices.
    Status = hsa_amd_register_system_event_handler(eventHandler, this);
    if (auto Err = Plugin::check(
            Status, "error in hsa_amd_register_system_event_handler: %s"))
      return std::move(Err);

    // List of host (CPU) agents.
    llvm::SmallVector<hsa_agent_t> HostAgents;

    // Count the number of available agents.
    auto Err = hsa_utils::iterateAgents([&](hsa_agent_t Agent) {
      // Get the device type of the agent.
      hsa_device_type_t DeviceType;
      hsa_status_t Status =
          hsa_agent_get_info(Agent, HSA_AGENT_INFO_DEVICE, &DeviceType);
      if (Status != HSA_STATUS_SUCCESS)
        return Status;

      // Classify the agents into kernel (GPU) and host (CPU) kernels.
      if (DeviceType == HSA_DEVICE_TYPE_GPU) {
        // Ensure that the GPU agent supports kernel dispatch packets.
        hsa_agent_feature_t Features;
        Status = hsa_agent_get_info(Agent, HSA_AGENT_INFO_FEATURE, &Features);
        if (Features & HSA_AGENT_FEATURE_KERNEL_DISPATCH)
          KernelAgents.push_back(Agent);
      } else if (DeviceType == HSA_DEVICE_TYPE_CPU) {
        HostAgents.push_back(Agent);
      }
      return HSA_STATUS_SUCCESS;
    });

    if (Err)
      return std::move(Err);

    int32_t NumDevices = KernelAgents.size();
    if (NumDevices == 0) {
      // Do not initialize if there are no devices.
      ODBG(OLDT_Init) << "There are no devices supporting AMDGPU.";
      return 0;
    }

    // There are kernel agents but there is no host agent. That should be
    // treated as an error.
    if (HostAgents.empty())
      return Plugin::error(ErrorCode::BACKEND_FAILURE, "no AMDGPU host agents");

    // Initialize a host device for each host agent. Skip host agents without
    // the required memory pools, e.g., NUMA nodes without memory banks.
    for (hsa_agent_t Agent : HostAgents) {
      AMDHostDeviceTy *HostDevice = allocate<AMDHostDeviceTy>();
      new (HostDevice) AMDHostDeviceTy(*this, Agent);

      auto UsableOrErr = HostDevice->init();
      if (!UsableOrErr)
        return UsableOrErr.takeError();
      if (*UsableOrErr)
        HostDevices.push_back(HostDevice);
    }

    if (HostDevices.empty())
      return Plugin::error(ErrorCode::BACKEND_FAILURE,
                           "no AMDGPU host agents with kernel args and "
                           "fine-grained memory pools");

    return NumDevices;
  }

  /// Deinitialize the plugin.
  Error deinitImpl() override {
    // The HSA runtime was not initialized, so nothing from the plugin was
    // actually initialized.
    if (!Initialized)
      return Plugin::success();

    for (AMDHostDeviceTy *HostDevice : HostDevices)
      if (auto Err = HostDevice->deinit())
        return Err;

    // Finalize the HSA runtime.
    hsa_status_t Status = hsa_shut_down();
    return Plugin::check(Status, "error in hsa_shut_down: %s");
  }

  /// Creates an AMDGPU device.
  GenericDeviceTy *createDevice(GenericPluginTy &Plugin, int32_t DeviceId,
                                int32_t NumDevices) override {
    hsa_agent_t Agent = getKernelAgent(DeviceId);
    return new AMDGPUDeviceTy(Plugin, DeviceId, NumDevices,
                              getNearestHostDevice(Agent), Agent);
  }

  Expected<std::unique_ptr<PluginContextTy>>
  createPluginContext(llvm::ArrayRef<GenericDeviceTy *> Devices) override {
    return std::make_unique<AMDGPUPluginContextTy>(*this, Devices);
  }

  /// Creates an AMDGPU global handler.
  GenericGlobalHandlerTy *createGlobalHandler() override {
    return new AMDGPUGlobalHandlerTy();
  }

  Triple::ArchType getTripleArch() const override { return Triple::amdgpu; }

  const char *getName() const override { return GETNAME(TARGET_NAME); }

  /// Get the ELF code for recognizing the compatible image binary.
  uint16_t getMagicElfBits() const override { return ELF::EM_AMDGPU; }

  /// Check whether the image is compatible with an AMDGPU device.
  Expected<bool> isELFCompatible(uint32_t DeviceId,
                                 StringRef Image) const override {
    // Get the associated architecture and flags from the ELF.
    auto ElfOrErr = ELF64LEObjectFile::create(
        MemoryBufferRef(Image, /*Identifier=*/""), /*InitContent=*/false);
    if (!ElfOrErr)
      return ElfOrErr.takeError();
    std::optional<StringRef> Processor = ElfOrErr->tryGetCPUName();
    if (!Processor)
      return false;

    SmallVector<SmallString<32>> Targets;
    if (auto Err = hsa_utils::getTargetTripleAndFeatures(
            getKernelAgent(DeviceId), Targets))
      return Err;
    for (auto &Target : Targets)
      if (offloading::amdgpu::isImageCompatibleWithEnv(
              *Processor, ElfOrErr->getPlatformFlags(), Target.str()))
        return true;
    return false;
  }

  bool isDataExchangable(int32_t SrcDeviceId, int32_t DstDeviceId) override {
    return true;
  }

  /// Get the host device nearest to \p KernelAgent. If the runtime cannot
  /// report the nearest host agent (e.g., ROCm older than 6.0) or the agent has
  /// no host device, use the first host device.
  AMDHostDeviceTy &getNearestHostDevice(hsa_agent_t KernelAgent) {
    hsa_agent_t NearestCPU;
    if (hsa_agent_get_info(
            KernelAgent,
            static_cast<hsa_agent_info_t>(HSA_AMD_AGENT_INFO_NEAREST_CPU),
            &NearestCPU) == HSA_STATUS_SUCCESS) {
      for (AMDHostDeviceTy *HostDevice : HostDevices)
        if (HostDevice->getAgent().handle == NearestCPU.handle)
          return *HostDevice;
    }
    return *HostDevices.front();
  }

  /// Get the kernel agent with the corresponding agent id.
  hsa_agent_t getKernelAgent(int32_t AgentId) const {
    assert((uint32_t)AgentId < KernelAgents.size() && "Invalid agent id");
    return KernelAgents[AgentId];
  }

  /// Get the list of the available kernel agents.
  const llvm::SmallVector<hsa_agent_t> &getKernelAgents() const {
    return KernelAgents;
  }

  /// Create an HSA signal for the RPC doorbell and return the fields needed
  /// for the GPU to fire interrupts that wake the server thread.
  Error initRPCDoorbell(uint64_t *&Value, uint64_t *&Mailbox,
                        uint32_t &EventID) override {
    // Lazily initialize the RPC signal on first use so we don't leak an HSA
    // signal when no device uses RPC.
    {
      const std::lock_guard<std::mutex> Lock(RPCSignalMutex);
      if (!RPCSignal.get().handle) {
        if (auto Err = RPCSignal.init(0))
          return Err;
      }
    }

    // Pull out the necessary fields to communicate with the signal from the
    // device. These are not exposed but are unlikely to be changed.
    struct AMDSignal {
      int64_t kind;
      int64_t value;
      uint64_t event_mailbox_ptr;
      uint32_t event_id;
    };
    auto *Doorbell = reinterpret_cast<AMDSignal *>(RPCSignal.get().handle);

    // The event ID corresponds do the HSA signal's slot in the interrupt list.
    Value = reinterpret_cast<uint64_t *>(&Doorbell->value);
    Mailbox = reinterpret_cast<uint64_t *>(Doorbell->event_mailbox_ptr);
    EventID = Doorbell->event_id;

    hsa_signal_t Signal = RPCSignal.get();
    getRPCServer().setSleepFunction(
        [Signal]() {
          hsa_signal_wait_scacquire(Signal, HSA_SIGNAL_CONDITION_NE, 0,
                                    /*timeout_hint=*/UINT64_MAX,
                                    HSA_WAIT_STATE_BLOCKED);
        },
        [Signal]() { hsa_signal_store_screlease(Signal, 1); });

    return Plugin::success();
  }

  Error deinitRPCDoorbell() override {
    const std::lock_guard<std::mutex> Lock(RPCSignalMutex);
    if (RPCSignal.get().handle)
      return RPCSignal.deinit();
    return Plugin::success();
  }

private:
  /// Event handler that will be called by ROCr if an event is detected.
  static hsa_status_t eventHandler(const hsa_amd_event_t *Event,
                                   void *PluginPtr) {
    if (Event->event_type != HSA_AMD_GPU_MEMORY_FAULT_EVENT)
      return HSA_STATUS_SUCCESS;

    SmallVector<std::string> Reasons;
    uint32_t ReasonsMask = Event->memory_fault.fault_reason_mask;
    if (ReasonsMask & HSA_AMD_MEMORY_FAULT_PAGE_NOT_PRESENT)
      Reasons.emplace_back("Page not present or supervisor privilege");
    if (ReasonsMask & HSA_AMD_MEMORY_FAULT_READ_ONLY)
      Reasons.emplace_back("Write access to a read-only page");
    if (ReasonsMask & HSA_AMD_MEMORY_FAULT_NX)
      Reasons.emplace_back("Execute access to a page marked NX");
    if (ReasonsMask & HSA_AMD_MEMORY_FAULT_HOST_ONLY)
      Reasons.emplace_back("GPU attempted access to a host only page");
    if (ReasonsMask & HSA_AMD_MEMORY_FAULT_DRAMECC)
      Reasons.emplace_back("DRAM ECC failure");
    if (ReasonsMask & HSA_AMD_MEMORY_FAULT_IMPRECISE)
      Reasons.emplace_back("Can't determine the exact fault address");
    if (ReasonsMask & HSA_AMD_MEMORY_FAULT_SRAMECC)
      Reasons.emplace_back("SRAM ECC failure (ie registers, no fault address)");
    if (ReasonsMask & HSA_AMD_MEMORY_FAULT_HANG)
      Reasons.emplace_back("GPU reset following unspecified hang");

    // If we do not know the reason, say so, otherwise remove the trailing comma
    // and space.
    if (Reasons.empty())
      Reasons.emplace_back("Unknown (" + std::to_string(ReasonsMask) + ")");

    uint32_t Node = -1;
    hsa_agent_get_info(Event->memory_fault.agent, HSA_AGENT_INFO_NODE, &Node);

    AMDGPUPluginTy &Plugin = *reinterpret_cast<AMDGPUPluginTy *>(PluginPtr);
    for (uint32_t I = 0, E = Plugin.getNumDevices();
         Node != uint32_t(-1) && I < E; ++I) {
      AMDGPUDeviceTy &AMDGPUDevice =
          reinterpret_cast<AMDGPUDeviceTy &>(Plugin.getDevice(I));
      auto KernelTraceInfoRecord =
          AMDGPUDevice.KernelLaunchTraces.getExclusiveAccessor();

      uint32_t DeviceNode = -1;
      if (auto Err =
              AMDGPUDevice.getDeviceAttr(HSA_AGENT_INFO_NODE, DeviceNode)) {
        consumeError(std::move(Err));
        continue;
      }
      if (DeviceNode != Node)
        continue;
      void *DevicePtr = (void *)Event->memory_fault.virtual_address;
      std::string S;
      llvm::raw_string_ostream OS(S);
      OS << llvm::format("memory access fault by GPU %" PRIu32
                         " (agent 0x%" PRIx64
                         ") at virtual address %p. Reasons: %s",
                         Node, Event->memory_fault.agent.handle,
                         (void *)Event->memory_fault.virtual_address,
                         llvm::join(Reasons, ", ").c_str());
      ErrorReporter::reportKernelTraces(AMDGPUDevice, *KernelTraceInfoRecord);
      ErrorReporter::reportMemoryAccessError(AMDGPUDevice, DevicePtr, S,
                                             /*Abort*/ true);
    }

    // Abort the execution since we do not recover from this error.
    FATAL_MESSAGE(1,
                  "memory access fault by GPU %" PRIu32 " (agent 0x%" PRIx64
                  ") at virtual address %p. Reasons: %s",
                  Node, Event->memory_fault.agent.handle,
                  (void *)Event->memory_fault.virtual_address,
                  llvm::join(Reasons, ", ").c_str());

    return HSA_STATUS_ERROR;
  }

  /// Indicate whether the HSA runtime was correctly initialized. Even if there
  /// is no available devices this boolean will be true. It indicates whether
  /// we can safely call HSA functions (e.g., hsa_shut_down).
  bool Initialized;

  /// Arrays of the available GPU and CPU agents. These arrays of handles should
  /// not be here but in the AMDGPUDeviceTy structures directly. However, the
  /// HSA standard does not provide API functions to retirve agents directly,
  /// only iterating functions. We cache the agents here for convenience.
  llvm::SmallVector<hsa_agent_t> KernelAgents;

  /// HSA signal used as the RPC doorbell for GPU-to-host interrupts.
  AMDGPUSignalTy RPCSignal;
  std::mutex RPCSignalMutex;

  /// The host devices, one per HSA host agent.
  llvm::SmallVector<AMDHostDeviceTy *> HostDevices;
};

Expected<PluginAllocInfoTy>
AMDGPUPluginContextTy::getAllocInfo(const void *Ptr) {
  // HSA gives the base of the region containing Ptr, so interior pointers
  // resolve to the same tracker entry.
  hsa_amd_pointer_info_t HsaInfo{};
  HsaInfo.size = sizeof(hsa_amd_pointer_info_t);
  hsa_status_t Status = hsa_amd_pointer_info(
      const_cast<void *>(Ptr), &HsaInfo, /*Allocator=*/nullptr,
      /*num_agents_accessible=*/nullptr, /*accessible=*/nullptr);
  if (auto Err = Plugin::check(Status, "error in hsa_amd_pointer_info: %s"))
    return std::move(Err);

  AllocInfo Info;
  {
    std::lock_guard<std::mutex> Lock(AllocationsMutex);
    auto It = Allocations.find(HsaInfo.agentBaseAddress);
    if (It == Allocations.end())
      return Plugin::error(ErrorCode::NOT_FOUND,
                           "pointer is not a known allocation in this context");
    Info = It->second;
  }

  return PluginAllocInfoTy{Info.Device, Info.Kind, HsaInfo.agentBaseAddress,
                           HsaInfo.sizeInBytes};
}

Error AMDGPUKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
                                 uint32_t NumThreads[3], uint32_t NumBlocks[3],
                                 uint32_t DynBlockMemSize,
                                 KernelLaunchArgsTy &LaunchArgs,
                                 AsyncInfoWrapperTy &AsyncInfoWrapper,
                                 GenericProfilerTy *ProfilerPtr) const {
  // Cooperative kernel launch is not yet supported for AMDGPU
  if (LaunchArgs.Flags.Cooperative)
    return Plugin::error(ErrorCode::UNSUPPORTED,
                         "cooperative kernel launch not supported for AMDGPU");

  AMDGPUDeviceTy &AMDGPUDevice = static_cast<AMDGPUDeviceTy &>(GenericDevice);
  AMDGPUMemoryManagerTy &ArgsMemoryManager =
      AMDGPUDevice.getHostDevice().getArgsMemoryManager();

  void *AllArgs = nullptr;
  if (auto Err = ArgsMemoryManager.allocate(ArgsSize, &AllArgs))
    return Err;

  uint64_t StackSize;
  if (auto Err = GenericDevice.getDeviceStackSize(StackSize))
    return Err;

  // Copy explicit arguments.
  size_t ExplicitEnd = 0;
  if (LaunchArgs.Args) {
    const auto &ArgMDs = KernelInfo.ArgMDs;
    uint32_t NumArgs = LaunchArgs.NumArgs;

    if (NumArgs > ArgMDs.size())
      return Plugin::error(
          ErrorCode::INVALID_ARGUMENT,
          "number of arguments (%u) exceeds the number of arguments "
          "expected by the kernel (%zu)",
          NumArgs, ArgMDs.size());

    for (size_t I = 0; I < NumArgs; I++) {
      auto [Offset, Size] = ArgMDs[I];
      std::memcpy(utils::advancePtr(AllArgs, Offset), LaunchArgs.Args[I], Size);
    }

    auto [Offset, Size] = ArgMDs[NumArgs - 1];
    ExplicitEnd = Offset + Size;
  }

  AMDGPUStreamTy *Stream = nullptr;
  if (auto Err = AMDGPUDevice.getStream(AsyncInfoWrapper, Stream))
    return Err;

  uint64_t ImplArgsOffset =
      llvm::alignTo(ExplicitEnd, alignof(hsa_utils::AMDGPUImplicitArgsTy));
  if (ArgsSize > ImplArgsOffset) {
    hsa_utils::AMDGPUImplicitArgsTy *ImplArgs =
        reinterpret_cast<hsa_utils::AMDGPUImplicitArgsTy *>(
            utils::advancePtr(AllArgs, ImplArgsOffset));

    // Set the COV5+ implicit arguments to the appropriate values if present.
    uint64_t ImplArgsSize = ArgsSize - ImplArgsOffset;
    std::memset(ImplArgs, 0, ImplArgsSize);

    using ImplArgsTy = hsa_utils::AMDGPUImplicitArgsTy;
    hsa_utils::initImplArg(ImplArgs, &ImplArgsTy::BlockCountX, ImplArgsSize,
                           NumBlocks[0]);
    hsa_utils::initImplArg(ImplArgs, &ImplArgsTy::BlockCountY, ImplArgsSize,
                           NumBlocks[1]);
    hsa_utils::initImplArg(ImplArgs, &ImplArgsTy::BlockCountZ, ImplArgsSize,
                           NumBlocks[2]);

    hsa_utils::initImplArg(ImplArgs, &ImplArgsTy::GroupSizeX, ImplArgsSize,
                           NumThreads[0]);
    hsa_utils::initImplArg(ImplArgs, &ImplArgsTy::GroupSizeY, ImplArgsSize,
                           NumThreads[1]);
    hsa_utils::initImplArg(ImplArgs, &ImplArgsTy::GroupSizeZ, ImplArgsSize,
                           NumThreads[2]);

    hsa_utils::initImplArg(ImplArgs, &ImplArgsTy::GridDims, ImplArgsSize,
                           NumBlocks[2] * NumThreads[2] > 1
                               ? 3
                               : 1 + (NumBlocks[1] * NumThreads[1] != 1));

    hsa_utils::initImplArg(ImplArgs, &ImplArgsTy::DynamicLdsSize, ImplArgsSize,
                           DynBlockMemSize);
  }

  // HSA requires the group segment size to include both static and dynamic.
  uint32_t TotalBlockMemSize = getStaticBlockMemSize() + DynBlockMemSize;

  auto ProfilerSpecificData = getOrNullProfilerSpecificData(AsyncInfoWrapper);

  // Push the kernel launch into the stream.
  return Stream->pushKernelLaunch(
      *this, AllArgs, NumThreads, NumBlocks, TotalBlockMemSize, StackSize,
      ArgsMemoryManager, ProfilerSpecificData, ProfilerPtr);
}

Error AMDGPUKernelTy::printLaunchInfoDetails(
    GenericDeviceTy &GenericDevice, const KernelLaunchArgsTy &LaunchArgs,
    uint32_t NumThreads[3], uint32_t NumBlocks[3]) const {
  // Only do all this when the output is requested
  if (!(getInfoLevel() & OMP_INFOTYPE_PLUGIN_KERNEL))
    return Plugin::success();

  // General Info
  auto *NumGroups = NumBlocks;
  auto *ThreadsPerGroup = NumThreads;

  // Kernel Arguments Info
  auto ArgNum = LaunchArgs.NumArgs;

  // Details for AMDGPU kernels (read from image)
  // https://www.llvm.org/docs/AMDGPUUsage.html#code-object-v4-metadata
  auto GroupSegmentSize = KernelInfo.GroupSegmentList;
  auto SGPRCount = KernelInfo.SGPRCount;
  auto VGPRCount = KernelInfo.VGPRCount;
  auto SGPRSpillCount = KernelInfo.SGPRSpillCount;
  auto VGPRSpillCount = KernelInfo.VGPRSpillCount;
  auto MaxFlatWorkgroupSize = KernelInfo.MaxFlatWorkgroupSize;

  // Prints additional launch info that contains the following.
  // Num Args: The number of kernel arguments
  // Teams x Thrds: The number of teams and the number of threads actually
  // running.
  // MaxFlatWorkgroupSize: Maximum flat work-group size supported by the
  // kernel in work-items
  // LDS Usage: Amount of bytes used in LDS storage
  // S/VGPR Count: the number of S/V GPRs occupied by the kernel
  // S/VGPR Spill Count: how many S/VGPRs are spilled by the kernel
  INFO(OMP_INFOTYPE_PLUGIN_KERNEL, GenericDevice.getDeviceId(),
       "#Args: %d Teams x Thrds: %4ux%4u (MaxFlatWorkGroupSize: %u) LDS "
       "Usage: %uB #SGPRs/VGPRs: %u/%u #SGPR/VGPR Spills: %u/%u\n",
       ArgNum, NumGroups[0] * NumGroups[1] * NumGroups[2],
       ThreadsPerGroup[0] * ThreadsPerGroup[1] * ThreadsPerGroup[2],
       MaxFlatWorkgroupSize, GroupSegmentSize, SGPRCount, VGPRCount,
       SGPRSpillCount, VGPRSpillCount);

  return Plugin::success();
}

/// Map an HSA status code to the corresponding offload error code.
static ErrorCode getOffloadErrorCode(hsa_status_t ResultCode) {
  switch (ResultCode) {
  case HSA_STATUS_ERROR_INVALID_SYMBOL_NAME:
  case HSA_STATUS_ERROR_INVALID_ISA_NAME:
    return ErrorCode::NOT_FOUND;
  case HSA_STATUS_ERROR_INVALID_CODE_OBJECT:
  case HSA_STATUS_ERROR_INVALID_ISA:
  case HSA_STATUS_ERROR_INCOMPATIBLE_ARGUMENTS:
    return ErrorCode::INVALID_BINARY;
  case HSA_STATUS_ERROR_OUT_OF_RESOURCES:
    return ErrorCode::OUT_OF_RESOURCES;
  case HSA_STATUS_ERROR_NOT_INITIALIZED:
    return ErrorCode::UNINITIALIZED;
  case HSA_STATUS_ERROR_INVALID_ARGUMENT:
  case HSA_STATUS_ERROR_INVALID_ALLOCATION:
  case HSA_STATUS_ERROR_INVALID_AGENT:
  case HSA_STATUS_ERROR_INVALID_REGION:
  case HSA_STATUS_ERROR_INVALID_QUEUE:
  case HSA_STATUS_ERROR_INVALID_INDEX:
    return ErrorCode::INVALID_ARGUMENT;
  default:
    return ErrorCode::UNKNOWN;
  }
}

template <typename... ArgsTy>
static Error Plugin::check(int32_t Code, const char *ErrFmt, ArgsTy... Args) {
  hsa_status_t ResultCode = static_cast<hsa_status_t>(Code);
  if (ResultCode == HSA_STATUS_SUCCESS || ResultCode == HSA_STATUS_INFO_BREAK)
    return Plugin::success();

  const char *Desc = "unknown error";
  hsa_status_t Ret = hsa_status_string(ResultCode, &Desc);
  if (Ret != HSA_STATUS_SUCCESS)
    REPORT() << "Unrecognized " GETNAME(TARGET_NAME) " error code " << Code;

  return Plugin::error(getOffloadErrorCode(ResultCode), ErrFmt, Args..., Desc);
}

Expected<void *> AMDGPUMemoryManagerTy::allocate(size_t Size, void *HstPtr,
                                                 TargetAllocTy Kind,
                                                 size_t Alignment) {
  // Allocate memory from the pool.
  void *Ptr = nullptr;
  if (auto Err = MemoryPool->allocate(Size, &Ptr, Alignment))
    return std::move(Err);

  assert(Ptr && "Invalid pointer");

  // Get a list of agents that can access this memory pool.
  llvm::SmallVector<hsa_agent_t> Agents;
  llvm::copy_if(
      Plugin.getKernelAgents(), std::back_inserter(Agents),
      [&](hsa_agent_t Agent) { return MemoryPool->canAccess(Agent); });

  // Allow all valid kernel agents to access the allocation.
  if (auto Err = MemoryPool->enableAccess(Ptr, Size, Agents))
    return std::move(Err);
  return Ptr;
}

Expected<void *> AMDGPUDeviceTy::allocate(size_t Size, void *,
                                          TargetAllocTy Kind,
                                          size_t Alignment) {
  if (Size == 0)
    return nullptr;

  // Find the correct memory pool.
  AMDGPUMemoryPoolTy *MemoryPool = nullptr;
  switch (Kind) {
  case TARGET_ALLOC_DEFAULT:
  case TARGET_ALLOC_DEVICE:
    MemoryPool = CoarseGrainedMemoryPools[0];
    break;
  case TARGET_ALLOC_HOST:
    MemoryPool = &HostDevice.getFineGrainedMemoryPool();
    break;
  case TARGET_ALLOC_SHARED:
    MemoryPool = &HostDevice.getFineGrainedMemoryPool();
    break;
  }

  if (!MemoryPool)
    return Plugin::error(ErrorCode::UNSUPPORTED,
                         "no memory pool for the specified allocation kind");

  // Allocate from the corresponding memory pool.
  void *Alloc = nullptr;
  if (auto Err = MemoryPool->allocate(Size, &Alloc, Alignment))
    return std::move(Err);

  if (Alloc) {
    // Get a list of agents that can access this memory pool. Inherently
    // necessary for host or shared allocations Also enabled for device memory
    // to allow device to device memcpy
    llvm::SmallVector<hsa_agent_t> Agents;
    llvm::copy_if(static_cast<AMDGPUPluginTy &>(Plugin).getKernelAgents(),
                  std::back_inserter(Agents), [&](hsa_agent_t Agent) {
                    return MemoryPool->canAccess(Agent);
                  });

    // Enable all valid kernel agents to access the buffer.
    if (auto Err = MemoryPool->enableAccess(Alloc, Size, Agents))
      return std::move(Err);

    // Register host / shared memory as pinned memory, so that transfers reading
    // from it can take a device-accessible path.
    if (Kind == TARGET_ALLOC_HOST || Kind == TARGET_ALLOC_SHARED)
      if (auto Err = PinnedAllocs.registerHostBuffer(Alloc, Alloc, Size))
        return std::move(Err);
  }

  return Alloc;
}

void AMDGPUQueueTy::callbackError(hsa_status_t Status, hsa_queue_t *Source,
                                  void *Data) {
  auto &AMDGPUDevice = *reinterpret_cast<AMDGPUDeviceTy *>(Data);

  // Drain any pending RPC work the device pushed before its queue died.
  if (RPCServerTy *RPCServer = AMDGPUDevice.getRPCServer())
    RPCServer->flushDevice(AMDGPUDevice);

  if (Status == HSA_STATUS_ERROR_EXCEPTION) {
    auto KernelTraceInfoRecord =
        AMDGPUDevice.KernelLaunchTraces.getExclusiveAccessor();
    std::function<bool(__tgt_async_info &)> AsyncInfoWrapperMatcher =
        [=](__tgt_async_info &AsyncInfo) {
          auto *Stream = reinterpret_cast<AMDGPUStreamTy *>(AsyncInfo.Queue);
          if (!Stream || !Stream->getQueue())
            return false;
          return Stream->getQueue()->Queue == Source;
        };
    ErrorReporter::reportTrapInKernel(AMDGPUDevice, *KernelTraceInfoRecord,
                                      AsyncInfoWrapperMatcher);
  }

  auto Err = Plugin::check(Status, "received error in queue %p: %s", Source);
  FATAL_MESSAGE(1, "%s", toString(std::move(Err)).data());
}

/// Implementation of profiling helper functions.
static ProfilingInfoTy *getProfilingInfo(void *Data) {
  return reinterpret_cast<ProfilingInfoTy *>(Data);
}

static std::pair<uint64_t, uint64_t>
getKernelStartAndEndTime(const ProfilingInfoTy *Args) {
  hsa_amd_profiling_dispatch_time_t Time = {};
  hsa_status_t Status = hsa_amd_profiling_get_dispatch_time(
      Args->Agent, Args->Signal->get(), &Time);
  if (Status != HSA_STATUS_SUCCESS)
    return {0, 0};
  return {static_cast<uint64_t>(Time.start * Args->TicksToTime),
          static_cast<uint64_t>(Time.end * Args->TicksToTime)};
}

static std::pair<uint64_t, uint64_t>
getCopyStartAndEndTime(const ProfilingInfoTy *Args) {
  hsa_amd_profiling_async_copy_time_t Time = {};
  hsa_status_t Status =
      hsa_amd_profiling_get_async_copy_time(Args->Signal->get(), &Time);
  if (Status != HSA_STATUS_SUCCESS)
    return {0, 0};
  return {static_cast<uint64_t>(Time.start * Args->TicksToTime),
          static_cast<uint64_t>(Time.end * Args->TicksToTime)};
}

static Error timeKernelInNsAsync(void *Data) {
  assert(Data && "Invalid data pointer");
  auto ProfilerInfo = getProfilingInfo(Data);
  assert(ProfilerInfo && "Invalid profiling info");
  assert(ProfilerInfo->ProfilerSpecificData && "Invalid ProfilerSpecificData");

  auto [StartTime, EndTime] = getKernelStartAndEndTime(ProfilerInfo);
  ProfilerInfo->Profiler->handleKernelCompletion(
      StartTime, EndTime, ProfilerInfo->ProfilerSpecificData);
  return Plugin::success();
}

} // namespace plugin
} // namespace target
} // namespace omp
} // namespace llvm

extern "C" {
llvm::omp::target::plugin::GenericPluginTy *createPlugin_amdgpu() {
  return new llvm::omp::target::plugin::AMDGPUPluginTy();
}
}