Merge commit '996e8bbfb71310d9a1b641bc954b69138cf85daa' into develop

projects/rocr-runtime/runtime/hsa-runtime/CMakeLists.txt

@@ -204,6 +204,11 @@ set ( SRCS core/driver/driver.cpp
            libamdhsacode/amd_hsa_code.cpp
            libamdhsacode/amd_core_dump.cpp )
 
+if ( BUILD_THUNK_VIRTIO )
+  list(APPEND SRCS core/driver/virtio/amd_kfd_virtio_driver.cpp)
+  target_compile_definitions(hsa-runtime64 PRIVATE HSAKMT_VIRTIO_ENABLED=1)
+endif()
+
 target_sources( ${CORE_RUNTIME_TARGET} PRIVATE ${SRCS} )
 
 ## Depend on trap handler target.
@@ -302,6 +307,10 @@ target_link_libraries ( ${CORE_RUNTIME_TARGET} PRIVATE elf::elf dl pthread rt )
 # Link to hsakmt-staticdrm target for static library builds
 if( BUILD_SHARED_LIBS )
   target_link_libraries ( ${CORE_RUNTIME_TARGET} PRIVATE hsakmt::hsakmt PkgConfig::drm)
+  if( BUILD_THUNK_VIRTIO )
+    message(STATUS "Building with virtio support")
+    target_link_libraries ( ${CORE_RUNTIME_TARGET} PRIVATE hsakmt_virtio)
+  endif()
   find_package(rocprofiler-register)
   if(rocprofiler-register_FOUND)
     target_compile_definitions(${CORE_RUNTIME_TARGET} PRIVATE HSA_ROCPROFILER_REGISTER=1

projects/rocr-runtime/runtime/hsa-runtime/core/driver/virtio/amd_kfd_virtio_driver.cpp

@@ -0,0 +1,514 @@
+////////////////////////////////////////////////////////////////////////////////
+//
+// The University of Illinois/NCSA
+// Open Source License (NCSA)
+//
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+//
+// Developed by:
+//
+//                 AMD Research and AMD HSA Software Development
+//
+//                 Advanced Micro Devices, Inc.
+//
+//                 www.amd.com
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to
+// deal with the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+//  - Redistributions of source code must retain the above copyright notice,
+//    this list of conditions and the following disclaimers.
+//  - Redistributions in binary form must reproduce the above copyright
+//    notice, this list of conditions and the following disclaimers in
+//    the documentation and/or other materials provided with the distribution.
+//  - Neither the names of Advanced Micro Devices, Inc,
+//    nor the names of its contributors may be used to endorse or promote
+//    products derived from this Software without specific prior written
+//    permission.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+// THE CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
+// OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+// ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS WITH THE SOFTWARE.
+//
+////////////////////////////////////////////////////////////////////////////////
+
+#include "core/inc/amd_virtio_driver.h"
+#include "hsakmt/hsakmt_virtio.h"
+
+#include <link.h>
+#include <vector>
+
+#include "core/inc/amd_gpu_agent.h"
+#include "core/inc/amd_memory_region.h"
+#include "core/inc/runtime.h"
+
+extern r_debug _amdgpu_r_debug;
+
+namespace rocr {
+namespace AMD {
+
+KfdVirtioDriver::KfdVirtioDriver(std::string devnode_name)
+    : core::Driver(core::DriverType::KFD_VIRTIO, std::move(devnode_name)) {}
+
+hsa_status_t KfdVirtioDriver::DiscoverDriver(std::unique_ptr<core::Driver>& driver) {
+  auto tmp_driver = std::unique_ptr<core::Driver>(new KfdVirtioDriver(""));
+
+  if (tmp_driver->Open() == HSA_STATUS_SUCCESS) {
+    driver = std::move(tmp_driver);
+    return HSA_STATUS_SUCCESS;
+  }
+
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::Open() {
+  return vhsaKmtOpenKFD() == HSAKMT_STATUS_SUCCESS ? HSA_STATUS_SUCCESS : HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::Close() {
+  return vhsaKmtCloseKFD() == HSAKMT_STATUS_SUCCESS ? HSA_STATUS_SUCCESS : HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::Init() {
+  HSAKMT_STATUS ret =
+      vhsaKmtRuntimeEnable(&_amdgpu_r_debug, core::Runtime::runtime_singleton_->flag().debug());
+  uint32_t caps_mask = 0;
+
+  if (ret != HSAKMT_STATUS_SUCCESS && ret != HSAKMT_STATUS_NOT_SUPPORTED) return HSA_STATUS_ERROR;
+
+  if (vhsaKmtGetRuntimeCapabilities(&caps_mask) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+  core::Runtime::runtime_singleton_->KfdVersion(
+      ret != HSAKMT_STATUS_NOT_SUPPORTED,
+      !!(caps_mask & HSA_RUNTIME_ENABLE_CAPS_SUPPORTS_CORE_DUMP_MASK));
+
+  if (vhsaKmtGetVersion(&version_) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+  core::Runtime::runtime_singleton_->KfdVersion(version_);
+
+  if (version_.KernelInterfaceMajorVersion == 1 && version_.KernelInterfaceMinorVersion == 0)
+    core::g_use_interrupt_wait = false;
+
+  /* Force disable interrupt wait in VIRTIO driver temporarily */
+  core::g_use_interrupt_wait = false;
+
+  /* Force disable XNACK in VIRTIO driver temporarily */
+  core::Runtime::runtime_singleton_->XnackEnabled(false);
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::ShutDown() {
+  HSAKMT_STATUS ret = vhsaKmtRuntimeDisable();
+  if (ret != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+  ret = vhsaKmtReleaseSystemProperties();
+
+  if (ret != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+  return Close();
+}
+
+hsa_status_t KfdVirtioDriver::QueryKernelModeDriver(core::DriverQuery query) {
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::GetSystemProperties(HsaSystemProperties& sys_props) const {
+  if (vhsaKmtAcquireSystemProperties(&sys_props) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::GetNodeProperties(HsaNodeProperties& node_props,
+                                                uint32_t node_id) const {
+  if (vhsaKmtGetNodeProperties(node_id, &node_props) != HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::GetEdgeProperties(std::vector<HsaIoLinkProperties>& io_link_props,
+                                                uint32_t node_id) const {
+  if (vhsaKmtGetNodeIoLinkProperties(node_id, io_link_props.size(), io_link_props.data()) !=
+      HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::GetMemoryProperties(
+    uint32_t node_id, std::vector<HsaMemoryProperties>& mem_props) const {
+  if (mem_props.empty()) {
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  if (vhsaKmtGetNodeMemoryProperties(node_id, mem_props.size(), mem_props.data()) !=
+      HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::GetCacheProperties(
+    uint32_t node_id, uint32_t processor_id, std::vector<HsaCacheProperties>& cache_props) const {
+  if (cache_props.empty()) {
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  if (vhsaKmtGetNodeCacheProperties(node_id, 0, cache_props.size(), cache_props.data()) !=
+      HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::GetDeviceHandle(uint32_t node_id, void** device_handle) const {
+  assert(device_handle != nullptr);
+
+  if (vhsaKmtGetAMDGPUDeviceHandle(node_id, device_handle) != HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::GetClockCounters(uint32_t node_id,
+                                               HsaClockCounters* clock_counter) const {
+  assert(clock_counter != nullptr);
+
+  if (vhsaKmtGetClockCounters(node_id, clock_counter) != HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::SetTrapHandler(uint32_t node_id, const void* base, uint64_t base_size,
+                                             const void* buffer_base,
+                                             uint64_t buffer_base_size) const {
+  if (vhsaKmtSetTrapHandler(node_id, const_cast<void*>(base), base_size,
+                            const_cast<void*>(buffer_base),
+                            buffer_base_size) != HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::AllocateMemory(const core::MemoryRegion& mem_region,
+                                             core::MemoryRegion::AllocateFlags alloc_flags,
+                                             void** mem, size_t size, uint32_t agent_node_id) {
+  const MemoryRegion& m_region(static_cast<const MemoryRegion&>(mem_region));
+  HsaMemFlags kmt_alloc_flags(m_region.mem_flags());
+  HSAKMT_STATUS ret;
+
+  kmt_alloc_flags.ui32.ExecuteAccess =
+      (alloc_flags & core::MemoryRegion::AllocateExecutable ? 1 : 0);
+  kmt_alloc_flags.ui32.AQLQueueMemory =
+      (alloc_flags & core::MemoryRegion::AllocateDoubleMap ? 1 : 0);
+
+  if (m_region.IsSystem() && (alloc_flags & core::MemoryRegion::AllocateNonPaged)) {
+    kmt_alloc_flags.ui32.NonPaged = 1;
+  }
+
+  if (!m_region.IsLocalMemory() && (alloc_flags & core::MemoryRegion::AllocateMemoryOnly)) {
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  // Allocating a memory handle for virtual memory
+  kmt_alloc_flags.ui32.NoAddress = !!(alloc_flags & core::MemoryRegion::AllocateMemoryOnly);
+
+  // Allocate pseudo fine grain memory
+  kmt_alloc_flags.ui32.CoarseGrain =
+      (alloc_flags & core::MemoryRegion::AllocatePCIeRW ? 0 : kmt_alloc_flags.ui32.CoarseGrain);
+
+  kmt_alloc_flags.ui32.NoSubstitute =
+      (alloc_flags & core::MemoryRegion::AllocatePinned ? 1 : kmt_alloc_flags.ui32.NoSubstitute);
+
+  kmt_alloc_flags.ui32.GTTAccess =
+      (alloc_flags & core::MemoryRegion::AllocateGTTAccess ? 1 : kmt_alloc_flags.ui32.GTTAccess);
+
+  kmt_alloc_flags.ui32.Uncached =
+      (alloc_flags & core::MemoryRegion::AllocateUncached ? 1 : kmt_alloc_flags.ui32.Uncached);
+
+  if (m_region.IsLocalMemory()) {
+    // Allocate physically contiguous memory. AllocateKfdMemory function call
+    // will fail if this flag is not supported in KFD.
+    kmt_alloc_flags.ui32.Contiguous =
+        (alloc_flags & core::MemoryRegion::AllocateContiguous ? 1
+                                                              : kmt_alloc_flags.ui32.Contiguous);
+  }
+
+  //// Only allow using the suballocator for ordinary VRAM.
+  if (m_region.IsLocalMemory() && !kmt_alloc_flags.ui32.NoAddress) {
+    bool subAllocEnabled = !core::Runtime::runtime_singleton_->flag().disable_fragment_alloc();
+    // Avoid modifying executable or queue allocations.
+    bool useSubAlloc = subAllocEnabled;
+    useSubAlloc &= ((alloc_flags & (~core::MemoryRegion::AllocateRestrict)) == 0);
+
+    if (useSubAlloc) {
+      *mem = m_region.fragment_alloc(size);
+
+      if ((alloc_flags & core::MemoryRegion::AllocateAsan)) {
+        // TODO: Implement ASAN support for VIRTIO driver
+        return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+      }
+
+      return HSA_STATUS_SUCCESS;
+    }
+  }
+
+  const uint32_t node_id = (alloc_flags & core::MemoryRegion::AllocateGTTAccess)
+      ? agent_node_id
+      : m_region.owner()->node_id();
+
+  //// Allocate memory.
+  //// If it fails attempt to release memory from the block allocator and retry.
+  ret = vhsaKmtAllocMemory(node_id, size, kmt_alloc_flags, mem);
+  if (ret != HSAKMT_STATUS_SUCCESS) {
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
+  if (*mem == nullptr) {
+    m_region.owner()->Trim();
+    ret = vhsaKmtAllocMemory(node_id, size, kmt_alloc_flags, mem);
+    if (ret != HSAKMT_STATUS_SUCCESS) {
+      return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+    }
+  }
+
+  if (*mem != nullptr) {
+    if (kmt_alloc_flags.ui32.NoAddress) return HSA_STATUS_SUCCESS;
+
+    // Commit the memory.
+    // For system memory, on non-restricted allocation, map it to all GPUs. On
+    // restricted allocation, only CPU is allowed to access by default, so
+    // no need to map
+    // For local memory, only map it to the owning GPU. Mapping to other GPU,
+    // if the access is allowed, is performed on AllowAccess.
+    HsaMemMapFlags map_flag = m_region.map_flags();
+    size_t map_node_count = 1;
+    const uint32_t owner_node_id = m_region.owner()->node_id();
+    const uint32_t* map_node_id = &owner_node_id;
+
+    if (m_region.IsSystem()) {
+      if ((alloc_flags & core::MemoryRegion::AllocateRestrict) == 0) {
+        // Map to all GPU agents.
+        map_node_count = core::Runtime::runtime_singleton_->gpu_ids().size();
+
+        if (map_node_count == 0) {
+          // No need to pin since no GPU in the platform.
+          return HSA_STATUS_SUCCESS;
+        }
+
+        map_node_id = &core::Runtime::runtime_singleton_->gpu_ids()[0];
+      } else {
+        // No need to pin it for CPU exclusive access.
+        return HSA_STATUS_SUCCESS;
+      }
+    }
+
+    uint64_t alternate_va = 0;
+    const bool is_resident =
+        (MakeMemoryResident(*mem, size, &alternate_va, &map_flag, map_node_count, map_node_id) ==
+         HSA_STATUS_SUCCESS);
+
+    const bool require_pinning =
+        (!m_region.full_profile() || m_region.IsLocalMemory() || m_region.IsScratch());
+
+    if (require_pinning && !is_resident) {
+      vhsaKmtFreeMemory(*mem, size);
+      *mem = nullptr;
+      return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+    }
+
+    if ((alloc_flags & core::MemoryRegion::AllocateAsan)) {
+      // TODO: Implement ASAN support for VIRTIO driver
+      return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+    }
+    return HSA_STATUS_SUCCESS;
+  }
+
+  return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+}
+
+hsa_status_t KfdVirtioDriver::FreeMemory(void* mem, size_t size) {
+  MakeMemoryUnresident(mem);
+  return vhsaKmtFreeMemory(mem, size) == HSAKMT_STATUS_SUCCESS ? HSA_STATUS_SUCCESS
+                                                               : HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::AllocateScratchMemory(uint32_t node_id, uint64_t size,
+                                                    void** mem) const {
+  assert(mem != nullptr);
+  assert(size != 0);
+
+  HsaMemFlags flags = {};
+  flags.ui32.Scratch = 1;
+  flags.ui32.HostAccess = 1;
+  void* ptr = nullptr;
+
+  HSAKMT_STATUS ret = vhsaKmtAllocMemory(node_id, size, flags, &ptr);
+  if (ret != HSAKMT_STATUS_SUCCESS || ptr == nullptr) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+
+  *mem = ptr;
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::RegisterMemory(void* ptr, uint64_t size,
+                                             HsaMemFlags mem_flags) const {
+  assert(ptr != nullptr);
+  assert(size != 0);
+
+  if (vhsaKmtRegisterMemoryWithFlags(ptr, size, mem_flags) != HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::DeregisterMemory(void* ptr) const {
+  if (vhsaKmtDeregisterMemory(ptr) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::AvailableMemory(uint32_t node_id, uint64_t* available_size) const {
+  assert(available_size != nullptr);
+
+  if (vhsaKmtAvailableMemory(node_id, available_size) != HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::MakeMemoryResident(const void* mem, size_t size,
+                                                 uint64_t* alternate_va,
+                                                 const HsaMemMapFlags* mem_flags,
+                                                 uint32_t num_nodes, const uint32_t* nodes) const {
+  assert(mem != nullptr);
+  assert(size != 0);
+
+  if (mem_flags == nullptr && nodes == nullptr) {
+    if (vhsaKmtMapMemoryToGPU(const_cast<void*>(mem), size, alternate_va) != HSAKMT_STATUS_SUCCESS)
+      return HSA_STATUS_ERROR;
+  } else if (mem_flags != nullptr && nodes != nullptr) {
+    if (vhsaKmtMapMemoryToGPUNodes(const_cast<void*>(mem), size, alternate_va, *mem_flags,
+                                   num_nodes,
+                                   const_cast<uint32_t*>(nodes)) != HSAKMT_STATUS_SUCCESS)
+      return HSA_STATUS_ERROR;
+  } else {
+    debug_print("Invalid memory flags ptr:%p nodes ptr:%p\n", mem_flags, nodes);
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::MakeMemoryUnresident(const void* mem) const {
+  vhsaKmtUnmapMemoryToGPU(const_cast<void*>(mem));
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::CreateQueue(uint32_t node_id, HSA_QUEUE_TYPE type, uint32_t queue_pct,
+                                          HSA_QUEUE_PRIORITY priority, uint32_t sdma_engine_id,
+                                          void* queue_addr, uint64_t queue_size_bytes,
+                                          HsaEvent* event, HsaQueueResource& queue_resource) const {
+  if (vhsaKmtCreateQueueExt(node_id, type, queue_pct, priority, sdma_engine_id, queue_addr,
+                            queue_size_bytes, event, &queue_resource) != HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::DestroyQueue(HSA_QUEUEID queue_id) const {
+  if (vhsaKmtDestroyQueue(queue_id) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::UpdateQueue(HSA_QUEUEID queue_id, uint32_t queue_percentage,
+                                          HSA_QUEUE_PRIORITY priority, void* queue_mem,
+                                          uint64_t queue_size, HsaEvent* event) const {
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::SetQueueCUMask(HSA_QUEUEID queue_id, uint32_t num_cu_mask,
+                                             uint32_t* cu_mask) const {
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::AllocQueueGWS(HSA_QUEUEID queue_id, uint32_t num_GWS,
+                                            uint32_t* GWS) const {
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::ExportDMABuf(void* mem, size_t size, int* dmabuf_fd, size_t* offset) {
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::ImportDMABuf(int dmabuf_fd, core::Agent& agent,
+                                           core::ShareableHandle& handle) {
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::Map(core::ShareableHandle handle, void* mem, size_t offset,
+                                  size_t size, hsa_access_permission_t perms) {
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::Unmap(core::ShareableHandle handle, void* mem, size_t offset,
+                                    size_t size) {
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::ReleaseShareableHandle(core::ShareableHandle& handle) {
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t KfdVirtioDriver::GetTileConfig(uint32_t node_id, HsaGpuTileConfig* config) const {
+  if (vhsaKmtGetTileConfig(node_id, config) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::SPMAcquire(uint32_t node_id) const { return HSA_STATUS_ERROR; }
+
+hsa_status_t KfdVirtioDriver::SPMRelease(uint32_t node_id) const { return HSA_STATUS_ERROR; }
+
+hsa_status_t KfdVirtioDriver::SPMSetDestBuffer(uint32_t node_id, uint32_t size, uint32_t* timeout,
+                                               uint32_t* size_copied, void* dest,
+                                               bool* is_data_loss) const {
+  return HSA_STATUS_ERROR;
+}
+
+
+hsa_status_t KfdVirtioDriver::OpenSMI(uint32_t node_id, int* fd) const { return HSA_STATUS_ERROR; }
+
+hsa_status_t KfdVirtioDriver::GetWallclockFrequency(uint32_t node_id, uint64_t* frequency) const {
+  assert(frequency != nullptr);
+
+  amdgpu_gpu_info info;
+  amdgpu_device_handle handle;
+  if (GetDeviceHandle(node_id, reinterpret_cast<void**>(&handle)) != HSA_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  if (vamdgpu_query_gpu_info(handle, &info) < 0) return HSA_STATUS_ERROR;
+
+  // Reported by libdrm in KHz.
+  *frequency = uint64_t(info.gpu_counter_freq) * 1000ull;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdVirtioDriver::IsModelEnabled(bool* enable) const {
+  *enable = false;
+  return HSA_STATUS_SUCCESS;
+}
+
+}  // namespace AMD
+}  // namespace rocr

projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp

@@ -332,10 +332,24 @@ hsa_status_t XdnaDriver::AllocQueueGWS(HSA_QUEUEID queue_id, uint32_t num_gws,
   return HSA_STATUS_ERROR_INVALID_QUEUE;
 }
 
-hsa_status_t XdnaDriver::ExportDMABuf(void *mem, size_t size, int *dmabuf_fd,
-                                      size_t *offset) {
-  // Not implemented yet.
-  return HSA_STATUS_ERROR;
+hsa_status_t XdnaDriver::ExportDMABuf(void* mem, size_t size, int* dmabuf_fd, size_t* offset) {
+  auto bo_handle = FindBOHandle(mem);
+  if (!bo_handle.IsValid()) {
+    return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+  }
+
+  drm_prime_handle export_params = {};
+  export_params.handle = bo_handle.handle;
+  export_params.flags = DRM_RDWR;
+  export_params.fd = -1;
+  if (ioctl(fd_, DRM_IOCTL_PRIME_HANDLE_TO_FD, &export_params) < 0) {
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
+  *dmabuf_fd = export_params.fd;
+  *offset = reinterpret_cast<uintptr_t>(mem) - reinterpret_cast<uintptr_t>(bo_handle.vaddr);
+
+  return HSA_STATUS_SUCCESS;
 }
 
 hsa_status_t XdnaDriver::ImportDMABuf(int dmabuf_fd, core::Agent &agent,

projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_blit_sdma.h

@@ -73,11 +73,7 @@ class BlitSdmaBase : public core::Blit {
                                              core::Signal& out_signal) = 0;
 };
 
-// RingIndexTy: 32/64-bit monotonic ring index, counting in bytes.
-// HwIndexMonotonic: true if SDMA HW index is monotonic, false if it wraps at end of ring.
-// SizeToCountOffset: value added to size (in bytes) to form SDMA command count field.
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-class BlitSdma : public BlitSdmaBase {
+template <bool useGCR> class BlitSdma : public BlitSdmaBase {
  public:
   BlitSdma();
 
@@ -163,9 +159,9 @@ class BlitSdma : public BlitSdmaBase {
   /// could be written. NULL if input size is greater than the size of queue
   /// buffer.
 
-  char* AcquireWriteAddress(uint32_t cmd_size, RingIndexTy& curr_index);
+  char* AcquireWriteAddress(uint32_t cmd_size, uint64_t& curr_index);
 
-  void UpdateWriteAndDoorbellRegister(RingIndexTy curr_index, RingIndexTy new_index);
+  void UpdateWriteAndDoorbellRegister(uint64_t curr_index, uint64_t new_index);
 
   /// @brief Updates the Write Register of compute device to the end of
   /// SDMA packet written into queue buffer. The update to Write Register
@@ -178,16 +174,16 @@ class BlitSdma : public BlitSdmaBase {
   /// @param curr_index Index passed back from AcquireWriteAddress.
   ///
   /// @param cmd_size Command packet size in bytes.
-  void ReleaseWriteAddress(RingIndexTy curr_index, uint32_t cmd_size);
+  void ReleaseWriteAddress(uint64_t curr_index, uint32_t cmd_size);
 
   /// @brief Writes NO-OP words into queue buffer in case writing a command
   /// causes the queue buffer to wrap.
   ///
   /// @param curr_index Index to begin padding from.
-  void PadRingToEnd(RingIndexTy curr_index);
+  void PadRingToEnd(uint64_t curr_index);
 
-  uint32_t WrapIntoRing(RingIndexTy index);
-  bool CanWriteUpto(RingIndexTy upto_index);
+  uint32_t WrapIntoRing(uint64_t index);
+  bool CanWriteUpto(uint64_t upto_index);
 
   /// @brief Build fence command
   void BuildFenceCommand(char* fence_command_addr, uint32_t* fence,
@@ -265,8 +261,8 @@ class BlitSdma : public BlitSdmaBase {
   HsaQueueResource queue_resource_;
 
   // Monotonic ring indices, in bytes, tracking written and submitted commands.
-  RingIndexTy cached_reserve_index_;
-  RingIndexTy cached_commit_index_;
+  uint64_t cached_reserve_index_;
+  uint64_t cached_commit_index_;
 
   static const uint32_t linear_copy_command_size_;
 
@@ -314,21 +310,11 @@ class BlitSdma : public BlitSdmaBase {
   size_t min_submission_size_;
 };
 
-// Ring indices are 32-bit.
-// HW ring indices are not monotonic (wrap at end of ring).
-// Count fields of SDMA commands are 0-based.
-typedef BlitSdma<uint32_t, false, 0, false> BlitSdmaV2V3;
 
-// Ring indices are 64-bit.
-// HW ring indices are monotonic (do not wrap at end of ring).
-// Count fields of SDMA commands are 1-based.
-typedef BlitSdma<uint64_t, true, -1, false> BlitSdmaV4;
+typedef BlitSdma<false> BlitSdmaV4;
 
-// Ring indices are 64-bit.
-// HW ring indices are monotonic (do not wrap at end of ring).
-// Count fields of SDMA commands are 1-based.
 // SDMA is connected to gL2.
-typedef BlitSdma<uint64_t, true, -1, true> BlitSdmaV5;
+typedef BlitSdma<true> BlitSdmaV5;
 
 }  // namespace amd
 }  // namespace rocr

projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_cpu_agent.h

@@ -51,6 +51,7 @@
 #include "core/inc/agent.h"
 #include "core/inc/queue.h"
 #include "core/inc/cache.h"
+#include "core/inc/driver.h"
 
 namespace rocr {
 namespace AMD {
@@ -62,7 +63,9 @@ class CpuAgent : public core::Agent {
   // @param [in] node Node id. Each CPU in different socket will get distinct
   // id.
   // @param [in] node_props Node property.
-  CpuAgent(HSAuint32 node, const HsaNodeProperties& node_props);
+  // @param [in] driver_type Driver type. Default is KFD.
+  CpuAgent(HSAuint32 node, const HsaNodeProperties& node_props,
+           core::DriverType driver_type = core::DriverType::KFD);
 
   // @brief CpuAgent destructor.
   ~CpuAgent();

projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_gpu_agent.h

@@ -73,10 +73,11 @@ typedef ScratchCache::ScratchInfo ScratchInfo;
 class GpuAgentInt : public core::Agent {
  public:
   // @brief Constructor
-   GpuAgentInt(uint32_t node_id)
-       : core::Agent(core::Runtime::runtime_singleton_->AgentDriver(
-                         core::DriverType::KFD),
-                     node_id, core::Agent::DeviceType::kAmdGpuDevice) {}
+  // @param [in] node_id Node id.
+  // @param [in] driver_type Driver type. Default is KFD.
+  GpuAgentInt(uint32_t node_id, core::DriverType driver_type)
+      : core::Agent(core::Runtime::runtime_singleton_->AgentDriver(driver_type), node_id,
+      core::Agent::DeviceType::kAmdGpuDevice) {}
 
    // @brief Ensure blits are ready (performance hint).
    virtual void PreloadBlits() {}
@@ -231,7 +232,10 @@ class GpuAgent : public GpuAgentInt {
   // id.
   // @param [in] node_props Node property.
   // @param [in] xnack_mode XNACK mode of device.
-  GpuAgent(HSAuint32 node, const HsaNodeProperties& node_props, bool xnack_mode, uint32_t index);
+  // @param [in] index Index of the GPU device.
+  // @param [in] driver_type Driver type. Default is KFD.
+  GpuAgent(HSAuint32 node, const HsaNodeProperties& node_props, bool xnack_mode, uint32_t index,
+           core::DriverType driver_type = core::DriverType::KFD);
 
   // @brief GPU agent destructor.
   ~GpuAgent();
@@ -721,9 +725,6 @@ class GpuAgent : public GpuAgentInt {
   // @brief Alternative aperture base address. Only on KV.
   uintptr_t ape1_base_;
 
-  // @brief Alternative aperture size. Only on KV.
-  size_t ape1_size_;
-
   // @brief Queue with GWS access.
   struct {
     lazy_ptr<core::Queue> queue_;

projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_gpu_pm4.h

@@ -89,7 +89,7 @@
 #  define PM4_ACQUIRE_MEM_GCR_CNTL_GLV_INV                 (1 << 8)
 #  define PM4_ACQUIRE_MEM_GCR_CNTL_GL1_INV                 (1 << 9)
 #  define PM4_ACQUIRE_MEM_GCR_CNTL_GL2_INV                 (1 << 14)
-
+#  define PM4_ACQUIRE_MEM_GCR_CNTL_GL2_WB                  (1 << 15)
 #define PM4_RELEASE_MEM_DW1_EVENT_INDEX(x)                 (((x) & 0xF) << 8)
 #  define PM4_RELEASE_MEM_EVENT_INDEX_AQL                  0x7
 

projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_virtio_driver.h

@@ -0,0 +1,124 @@
+////////////////////////////////////////////////////////////////////////////////
+//
+// The University of Illinois/NCSA
+// Open Source License (NCSA)
+//
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+//
+// Developed by:
+//
+//                 AMD Research and AMD HSA Software Development
+//
+//                 Advanced Micro Devices, Inc.
+//
+//                 www.amd.com
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to
+// deal with the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+//  - Redistributions of source code must retain the above copyright notice,
+//    this list of conditions and the following disclaimers.
+//  - Redistributions in binary form must reproduce the above copyright
+//    notice, this list of conditions and the following disclaimers in
+//    the documentation and/or other materials provided with the distribution.
+//  - Neither the names of Advanced Micro Devices, Inc,
+//    nor the names of its contributors may be used to endorse or promote
+//    products derived from this Software without specific prior written
+//    permission.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+// THE CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
+// OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+// ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS WITH THE SOFTWARE.
+//
+////////////////////////////////////////////////////////////////////////////////
+
+#ifndef HSA_RUNTIME_CORE_INC_AMD_VIRTIO_DRIVER_H_
+#define HSA_RUNTIME_CORE_INC_AMD_VIRTIO_DRIVER_H_
+
+#include <memory>
+#include <string>
+
+#include "hsakmt/hsakmt.h"
+
+#include "core/inc/driver.h"
+#include "core/inc/memory_region.h"
+
+namespace rocr {
+namespace AMD {
+
+class KfdVirtioDriver final : public core::Driver {
+ public:
+  KfdVirtioDriver(std::string devnode_name);
+
+  static hsa_status_t DiscoverDriver(std::unique_ptr<core::Driver>& driver);
+
+  hsa_status_t Init() override;
+  hsa_status_t ShutDown() override;
+  hsa_status_t QueryKernelModeDriver(core::DriverQuery query) override;
+  hsa_status_t Open() override;
+  hsa_status_t Close() override;
+  hsa_status_t GetSystemProperties(HsaSystemProperties& sys_props) const override;
+  hsa_status_t GetNodeProperties(HsaNodeProperties& node_props, uint32_t node_id) const override;
+  hsa_status_t GetEdgeProperties(std::vector<HsaIoLinkProperties>& io_link_props,
+                                 uint32_t node_id) const override;
+  hsa_status_t GetMemoryProperties(uint32_t node_id,
+                                   std::vector<HsaMemoryProperties>& mem_props) const override;
+  hsa_status_t GetCacheProperties(uint32_t node_id, uint32_t processor_id,
+                                  std::vector<HsaCacheProperties>& cache_props) const override;
+  hsa_status_t GetDeviceHandle(uint32_t node_id, void** device_handle) const;
+  hsa_status_t GetClockCounters(uint32_t node_id, HsaClockCounters* clock_counter) const;
+  hsa_status_t SetTrapHandler(uint32_t node_id, const void* base, uint64_t base_size,
+                              const void* buffer_base, uint64_t buffer_base_size) const;
+  hsa_status_t AllocateMemory(const core::MemoryRegion& mem_region,
+                              core::MemoryRegion::AllocateFlags alloc_flags, void** mem,
+                              size_t size, uint32_t agent_node_id) override;
+  hsa_status_t FreeMemory(void* mem, size_t size) override;
+  hsa_status_t AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const;
+  hsa_status_t RegisterMemory(void* ptr, uint64_t size, HsaMemFlags mem_flags) const override;
+  hsa_status_t DeregisterMemory(void* ptr) const override;
+  hsa_status_t AvailableMemory(uint32_t node_id, uint64_t* available_size) const;
+  hsa_status_t MakeMemoryResident(const void* mem, size_t size, uint64_t* alternate_va,
+                                  const HsaMemMapFlags* mem_flags, uint32_t num_nodes,
+                                  const uint32_t* nodes) const override;
+  hsa_status_t MakeMemoryUnresident(const void* mem) const override;
+  hsa_status_t CreateQueue(uint32_t node_id, HSA_QUEUE_TYPE type, uint32_t queue_pct,
+                           HSA_QUEUE_PRIORITY priority, uint32_t sdma_engine_id, void* queue_addr,
+                           uint64_t queue_size_bytes, HsaEvent* event,
+                           HsaQueueResource& queue_resource) const override;
+  hsa_status_t DestroyQueue(HSA_QUEUEID queue_id) const override;
+  hsa_status_t UpdateQueue(HSA_QUEUEID queue_id, uint32_t queue_percentage,
+                           HSA_QUEUE_PRIORITY priority, void* queue_mem, uint64_t queue_size,
+                           HsaEvent* event) const override;
+  hsa_status_t SetQueueCUMask(HSA_QUEUEID queue_id, uint32_t num_cu_mask,
+                              uint32_t* cu_mask) const override;
+  hsa_status_t AllocQueueGWS(HSA_QUEUEID queue_id, uint32_t num_GWS, uint32_t* GWS) const override;
+  hsa_status_t ExportDMABuf(void* mem, size_t size, int* dmabuf_fd, size_t* offset) override;
+  hsa_status_t ImportDMABuf(int dmabuf_fd, core::Agent& agent,
+                            core::ShareableHandle& handle) override;
+  hsa_status_t Map(core::ShareableHandle handle, void* mem, size_t offset, size_t size,
+                   hsa_access_permission_t perms) override;
+  hsa_status_t Unmap(core::ShareableHandle handle, void* mem, size_t offset, size_t size) override;
+  hsa_status_t ReleaseShareableHandle(core::ShareableHandle& handle) override;
+  hsa_status_t GetTileConfig(uint32_t node_id, HsaGpuTileConfig* config) const;
+  hsa_status_t SPMAcquire(uint32_t node_id) const override;
+  hsa_status_t SPMRelease(uint32_t node_id) const override;
+  hsa_status_t SPMSetDestBuffer(uint32_t node_id, uint32_t size, uint32_t* timeout,
+                                uint32_t* size_copied, void* dest,
+                                bool* is_data_loss) const override;
+  hsa_status_t OpenSMI(uint32_t node_id, int* fd) const override;
+  hsa_status_t GetWallclockFrequency(uint32_t node_id, uint64_t* frequency) const;
+  hsa_status_t IsModelEnabled(bool* enable) const override;
+};
+
+}  // namespace AMD
+}  // namespace rocr
+
+#endif  // HSA_RUNTIME_CORE_INC_AMD_VIRTIO_DRIVER_H_

projects/rocr-runtime/runtime/hsa-runtime/core/inc/driver.h

@@ -58,7 +58,14 @@ class Queue;
 
 enum class DriverQuery { GET_DRIVER_VERSION };
 
-enum class DriverType { XDNA = 0, KFD, NUM_DRIVER_TYPES };
+enum class DriverType {
+  XDNA = 0,
+  KFD,
+#ifdef HSAKMT_VIRTIO_ENABLED
+  KFD_VIRTIO,
+#endif
+  NUM_DRIVER_TYPES
+};
 
 /// @brief Handle for exported / imported memory.
 struct ShareableHandle {

projects/rocr-runtime/runtime/hsa-runtime/core/inc/runtime.h

@@ -510,6 +510,14 @@ class Runtime {
 
   std::vector<std::unique_ptr<Driver>>& AgentDrivers() { return agent_drivers_; }
 
+  static bool IsGPUDriver(DriverType driver_type) {
+    return driver_type == core::DriverType::KFD
+#ifdef HSAKMT_VIRTIO_ENABLED
+        || driver_type == core::DriverType::KFD_VIRTIO
+#endif
+        ;
+  }
+
  protected:
   static void AsyncEventsLoop(void*);
   static void AsyncIPCSockServerConnLoop(void*);
@@ -814,7 +822,6 @@ class Runtime {
   std::map<const void*, AddressHandle> reserved_address_map_;  // Indexed by VA
 
   struct MemoryHandle {
-    MemoryHandle() : region(NULL), size(0), ref_count(0), thunk_handle(NULL), alloc_flag(0) {}
     MemoryHandle(const MemoryRegion* region, size_t size, uint64_t flags_unused,
                  ThunkHandle thunk_handle, MemoryRegion::AllocateFlags alloc_flag)
         : region(region),
@@ -824,19 +831,23 @@ class Runtime {
           thunk_handle(thunk_handle),
           alloc_flag(alloc_flag) {}
 
-    static __forceinline hsa_amd_vmem_alloc_handle_t Convert(void* handle) {
+    static __forceinline hsa_amd_vmem_alloc_handle_t Convert(ThunkHandle handle) {
       hsa_amd_vmem_alloc_handle_t ret_handle = {
           static_cast<uint64_t>(reinterpret_cast<uintptr_t>(handle))};
       return ret_handle;
     }
 
+    static __forceinline ThunkHandle Convert(hsa_amd_vmem_alloc_handle_t handle) {
+      return reinterpret_cast<void*>(handle.handle);
+    }
+
     __forceinline core::Agent* agentOwner() const { return region->owner(); }
 
     const MemoryRegion* region;
     size_t size;
     int ref_count;
     int use_count;
-    ThunkHandle thunk_handle;  // handle returned by hsaKmtAllocMemory(NoAddress = 1)
+    ThunkHandle thunk_handle;  // handle returned by Driver::Allocate(NoAddress = 1)
     MemoryRegion::AllocateFlags alloc_flag;
   };
   std::map<ThunkHandle, MemoryHandle> memory_handle_map_;
@@ -888,11 +899,6 @@ class Runtime {
                             const hsa_amd_memory_access_desc_t *desc,
                             const size_t desc_cnt);
 
-  // Frees runtime memory when the runtime library is unloaded if safe to do so.
-  // Failure to release the runtime indicates an incorrect application but is
-  // common (example: calls library routines at process exit).
-  friend class RuntimeCleanup;
-
   void InitIPCDmaBufSupport();
   bool ipc_dmabuf_supported_;
   int  IPCClientImport(uint32_t conn_handle, uint64_t dmabuf_fd_handle,

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_blit_sdma.cpp

@@ -77,44 +77,33 @@ const size_t BlitSdmaBase::kMaxSingleCopySize = SDMA_PKT_COPY_LINEAR::kMaxSize_;
 const size_t BlitSdmaBase::kMaxSingleFillSize = SDMA_PKT_CONSTANT_FILL::kMaxSize_;
 
 // Initialize size of various sDMA commands use by this module
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::linear_copy_command_size_ = sizeof(SDMA_PKT_COPY_LINEAR);
+template <bool useGCR>
+const uint32_t BlitSdma<useGCR>::linear_copy_command_size_ = sizeof(SDMA_PKT_COPY_LINEAR);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::fill_command_size_ = sizeof(SDMA_PKT_CONSTANT_FILL);
+template <bool useGCR>
+const uint32_t BlitSdma<useGCR>::fill_command_size_ = sizeof(SDMA_PKT_CONSTANT_FILL);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::fence_command_size_ = sizeof(SDMA_PKT_FENCE);
+template <bool useGCR>
+const uint32_t BlitSdma<useGCR>::fence_command_size_ = sizeof(SDMA_PKT_FENCE);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::poll_command_size_ = sizeof(SDMA_PKT_POLL_REGMEM);
+template <bool useGCR>
+const uint32_t BlitSdma<useGCR>::poll_command_size_ = sizeof(SDMA_PKT_POLL_REGMEM);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::flush_command_size_ = sizeof(SDMA_PKT_POLL_REGMEM);
+template <bool useGCR>
+const uint32_t BlitSdma<useGCR>::flush_command_size_ = sizeof(SDMA_PKT_POLL_REGMEM);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::atomic_command_size_ = sizeof(SDMA_PKT_ATOMIC);
+template <bool useGCR>
+const uint32_t BlitSdma<useGCR>::atomic_command_size_ = sizeof(SDMA_PKT_ATOMIC);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::timestamp_command_size_ = sizeof(SDMA_PKT_TIMESTAMP);
+template <bool useGCR>
+const uint32_t BlitSdma<useGCR>::timestamp_command_size_ = sizeof(SDMA_PKT_TIMESTAMP);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::trap_command_size_ = sizeof(SDMA_PKT_TRAP);
+template <bool useGCR> const uint32_t BlitSdma<useGCR>::trap_command_size_ = sizeof(SDMA_PKT_TRAP);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-const uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                        useGCR>::gcr_command_size_ = sizeof(SDMA_PKT_GCR);
+template <bool useGCR> const uint32_t BlitSdma<useGCR>::gcr_command_size_ = sizeof(SDMA_PKT_GCR);
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BlitSdma()
+template <bool useGCR>
+BlitSdma<useGCR>::BlitSdma()
     : agent_(NULL),
       queue_start_addr_(NULL),
       bytes_queued_(0),
@@ -129,12 +118,11 @@ BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BlitSdma()
   std::memset(&queue_resource_, 0, sizeof(queue_resource_));
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::~BlitSdma() {}
+template <bool useGCR> BlitSdma<useGCR>::~BlitSdma() {}
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::Initialize(
-    const core::Agent& agent, bool use_xgmi, size_t linear_copy_size_override, int rec_eng) {
+template <bool useGCR>
+hsa_status_t BlitSdma<useGCR>::Initialize(const core::Agent& agent, bool use_xgmi,
+                                          size_t linear_copy_size_override, int rec_eng) {
   if (queue_start_addr_ != NULL) {
     // Already initialized.
     return HSA_STATUS_SUCCESS;
@@ -201,7 +189,7 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>:
     return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
   }
 
-  cached_reserve_index_ = *reinterpret_cast<RingIndexTy*>(queue_resource_.Queue_write_ptr);
+  cached_reserve_index_ = *reinterpret_cast<uint64_t*>(queue_resource_.Queue_write_ptr);
   cached_commit_index_ = cached_reserve_index_;
 
   if (core::g_use_interrupt_wait) {
@@ -218,9 +206,7 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>:
   return HSA_STATUS_SUCCESS;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::Destroy(
-    const core::Agent& agent) {
+template <bool useGCR> hsa_status_t BlitSdma<useGCR>::Destroy(const core::Agent& agent) {
   // Release all allocated resources and reset them to zero.
 
   if (queue_resource_.QueueId != 0) {
@@ -245,9 +231,8 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>:
   return HSA_STATUS_SUCCESS;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                      useGCR>::SubmitBlockingCommand(const void* cmd, size_t cmd_size,
+template <bool useGCR>
+hsa_status_t BlitSdma<useGCR>::SubmitBlockingCommand(const void* cmd, size_t cmd_size,
                                                      uint64_t size) {
   ScopedAcquire<KernelMutex> lock(&lock_);
 
@@ -278,11 +263,11 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
   return ret;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::SubmitCommand(
-    const void* cmd, size_t cmd_size, uint64_t size, const std::vector<core::Signal*>& dep_signals,
-    core::Signal& out_signal, std::vector<core::Signal*>& gang_signals) {
-
+template <bool useGCR>
+hsa_status_t BlitSdma<useGCR>::SubmitCommand(const void* cmd, size_t cmd_size, uint64_t size,
+                                             const std::vector<core::Signal*>& dep_signals,
+                                             core::Signal& out_signal,
+                                             std::vector<core::Signal*>& gang_signals) {
   uint32_t num_poll_command = 0;
 
   // Cached copy of dep_signals[i]->LoadRelaxed
@@ -355,7 +340,7 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>:
   // Add space for acquire or release Hdp flush command
   uint32_t flush_cmd_size = 0;
   if (core::Runtime::runtime_singleton_->flag().enable_sdma_hdp_flush()) {
-    if ((HwIndexMonotonic) && (hdp_flush_support_)) {
+    if (hdp_flush_support_) {
       flush_cmd_size = flush_command_size_;
     }
   }
@@ -368,7 +353,7 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>:
   const uint32_t pad_size = total_command_size < min_submission_size_ ?
                             min_submission_size_ - total_command_size : 0;
 
-  RingIndexTy curr_index;
+  uint64_t curr_index;
   char* command_addr;
   uint64_t prior_bytes, post_bytes;
   {
@@ -426,7 +411,7 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>:
 
   // Issue a Hdp flush cmd
   if (core::Runtime::runtime_singleton_->flag().enable_sdma_hdp_flush()) {
-    if ((HwIndexMonotonic) && (hdp_flush_support_)) {
+    if (hdp_flush_support_) {
       BuildHdpFlushCommand(command_addr);
       command_addr += flush_command_size_;
       bytes_written_[wrapped_index] = prior_bytes;
@@ -542,9 +527,8 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>:
   return HSA_STATUS_SUCCESS;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                      useGCR>::SubmitLinearCopyCommand(void* dst, const void* src, size_t size) {
+template <bool useGCR>
+hsa_status_t BlitSdma<useGCR>::SubmitLinearCopyCommand(void* dst, const void* src, size_t size) {
   // Break the copy into multiple copy operation incase the copy size exceeds
   // the SDMA linear copy limit.
   const size_t max_copy_size = max_single_linear_copy_size_ ? max_single_linear_copy_size_ :
@@ -557,9 +541,8 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
   return SubmitBlockingCommand(&buff[0], buff.size() * sizeof(SDMA_PKT_COPY_LINEAR), size);
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                      useGCR>::SubmitLinearCopyCommand(void* dst, const void* src, size_t size,
+template <bool useGCR>
+hsa_status_t BlitSdma<useGCR>::SubmitLinearCopyCommand(void* dst, const void* src, size_t size,
                                                        std::vector<core::Signal*>& dep_signals,
                                                        core::Signal& out_signal,
                                                        std::vector<core::Signal*>& gang_signals) {
@@ -577,9 +560,8 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
                        out_signal, gang_signals);
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t
-BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::SubmitCopyRectCommand(
+template <bool useGCR>
+hsa_status_t BlitSdma<useGCR>::SubmitCopyRectCommand(
     const hsa_pitched_ptr_t* dst, const hsa_dim3_t* dst_offset, const hsa_pitched_ptr_t* src,
     const hsa_dim3_t* src_offset, const hsa_dim3_t* range, std::vector<core::Signal*>& dep_signals,
     core::Signal& out_signal) {
@@ -653,9 +635,8 @@ BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::SubmitCopyRe
                        out_signal, gang_signals);
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-                      useGCR>::SubmitLinearFillCommand(void* ptr, uint32_t value, size_t count) {
+template <bool useGCR>
+hsa_status_t BlitSdma<useGCR>::SubmitLinearFillCommand(void* ptr, uint32_t value, size_t count) {
   const size_t size = count * sizeof(uint32_t);
 
   const uint32_t num_fill_command = (size + kMaxSingleFillSize - 1) / kMaxSingleFillSize;
@@ -666,15 +647,12 @@ hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
   return SubmitBlockingCommand(&buff[0], buff.size() * sizeof(SDMA_PKT_CONSTANT_FILL), size);
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-hsa_status_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::EnableProfiling(
-    bool enable) {
+template <bool useGCR> hsa_status_t BlitSdma<useGCR>::EnableProfiling(bool enable) {
   return HSA_STATUS_SUCCESS;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-char* BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::AcquireWriteAddress(
-    uint32_t cmd_size, RingIndexTy& curr_index) {
+template <bool useGCR>
+char* BlitSdma<useGCR>::AcquireWriteAddress(uint32_t cmd_size, uint64_t& curr_index) {
   // Ring is full when all but one byte is written.
   if (cmd_size >= kQueueSize) {
     return nullptr;
@@ -692,7 +670,7 @@ char* BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::Acquir
     }
 
     // Check whether the engine has finished using this region.
-    const RingIndexTy new_index = curr_index + cmd_size;
+    const uint64_t new_index = curr_index + cmd_size;
 
     if (CanWriteUpto(new_index) == false) {
       // Wait for read index to move and try again.
@@ -713,10 +691,8 @@ char* BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::Acquir
   return nullptr;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-              useGCR>::UpdateWriteAndDoorbellRegister(RingIndexTy curr_index,
-                                                      RingIndexTy new_index) {
+template <bool useGCR>
+void BlitSdma<useGCR>::UpdateWriteAndDoorbellRegister(uint64_t curr_index, uint64_t new_index) {
   while (true) {
     // Make sure that the address before ::curr_index is already released.
     // Otherwise the CP may read invalid packets.
@@ -725,21 +701,19 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
         // TODO: remove when sdma wpointer issue is resolved.
         // Wait until the SDMA engine finish processing all packets before
         // updating the wptr and doorbell.
-        while (WrapIntoRing(*reinterpret_cast<RingIndexTy*>(queue_resource_.Queue_read_ptr)) !=
+        while (WrapIntoRing(*reinterpret_cast<uint64_t*>(queue_resource_.Queue_read_ptr)) !=
                WrapIntoRing(curr_index)) {
           os::YieldThread();
         }
       }
 
       // Update write pointer and doorbell register.
-      *reinterpret_cast<RingIndexTy*>(queue_resource_.Queue_write_ptr) =
-          (HwIndexMonotonic ? new_index : WrapIntoRing(new_index));
+      *reinterpret_cast<uint64_t*>(queue_resource_.Queue_write_ptr) = new_index;
 
       // Ensure write pointer is visible to GPU before doorbell.
       std::atomic_thread_fence(std::memory_order_release);
 
-      *reinterpret_cast<RingIndexTy*>(queue_resource_.Queue_DoorBell) =
-          (HwIndexMonotonic ? new_index : WrapIntoRing(new_index));
+      *reinterpret_cast<uint64_t*>(queue_resource_.Queue_DoorBell) = new_index;
 
       atomic::Store(&cached_commit_index_, new_index, std::memory_order_release);
       break;
@@ -750,9 +724,8 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
   }
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::ReleaseWriteAddress(
-    RingIndexTy curr_index, uint32_t cmd_size) {
+template <bool useGCR>
+void BlitSdma<useGCR>::ReleaseWriteAddress(uint64_t curr_index, uint32_t cmd_size) {
   if (cmd_size > kQueueSize) {
     assert(false && "cmd_addr is outside the queue buffer range");
     return;
@@ -761,11 +734,9 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::Release
   UpdateWriteAndDoorbellRegister(curr_index, curr_index + cmd_size);
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::PadRingToEnd(
-    RingIndexTy curr_index) {
+template <bool useGCR> void BlitSdma<useGCR>::PadRingToEnd(uint64_t curr_index) {
   // Reserve region from here to the end of the ring.
-  RingIndexTy new_index = curr_index + (kQueueSize - WrapIntoRing(curr_index));
+  uint64_t new_index = curr_index + (kQueueSize - WrapIntoRing(curr_index));
 
   // Check whether the engine has finished using this region.
   if (CanWriteUpto(new_index) == false) {
@@ -786,37 +757,22 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::PadRing
   }
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-uint32_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::WrapIntoRing(
-    RingIndexTy index) {
+template <bool useGCR> uint32_t BlitSdma<useGCR>::WrapIntoRing(uint64_t index) {
   return index & (kQueueSize - 1);
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-bool BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::CanWriteUpto(
-    RingIndexTy upto_index) {
+template <bool useGCR> bool BlitSdma<useGCR>::CanWriteUpto(uint64_t upto_index) {
   // Get/calculate the monotonic read index.
-  RingIndexTy hw_read_index = *reinterpret_cast<RingIndexTy*>(queue_resource_.Queue_read_ptr);
-  RingIndexTy read_index;
-
-  if (HwIndexMonotonic) {
-    read_index = hw_read_index;
-  } else {
-    // Calculate distance from commit index to HW read index.
-    // Commit index is always < kQueueSize away from HW read index.
-    RingIndexTy commit_index = atomic::Load(&cached_commit_index_, std::memory_order_relaxed);
-    RingIndexTy dist_to_read_index = WrapIntoRing(commit_index - hw_read_index);
-    read_index = commit_index - dist_to_read_index;
-  }
+  uint64_t hw_read_index = *reinterpret_cast<uint64_t*>(queue_resource_.Queue_read_ptr);
 
   // Check whether the read pointer has passed the given index.
   // At most we can submit (kQueueSize - 1) bytes at a time.
-  return (upto_index - read_index) < kQueueSize;
+  return (upto_index - hw_read_index) < kQueueSize;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildFenceCommand(
-    char* fence_command_addr, uint32_t* fence, uint32_t fence_value) {
+template <bool useGCR>
+void BlitSdma<useGCR>::BuildFenceCommand(char* fence_command_addr, uint32_t* fence,
+                                         uint32_t fence_value) {
   assert(fence_command_addr != NULL);
   SDMA_PKT_FENCE* packet_addr =
       reinterpret_cast<SDMA_PKT_FENCE*>(fence_command_addr);
@@ -836,9 +792,9 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildFe
   packet_addr->DATA_UNION.data = fence_value;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildCopyCommand(
-    char* cmd_addr, uint32_t num_copy_command, void* dst, const void* src, size_t size) {
+template <bool useGCR>
+void BlitSdma<useGCR>::BuildCopyCommand(char* cmd_addr, uint32_t num_copy_command, void* dst,
+                                        const void* src, size_t size) {
   size_t cur_size = 0;
   const size_t max_copy_size = max_single_linear_copy_size_ ? max_single_linear_copy_size_ :
                                                               kMaxSingleCopySize;
@@ -858,9 +814,9 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildCo
     packet_addr->HEADER_UNION.sub_op = SDMA_SUBOP_COPY_LINEAR;
 
     if (max_copy_size == (1 << 30) -1)
-      packet_addr->COUNT_UNION.count_ext.count = copy_size + SizeToCountOffset;
+      packet_addr->COUNT_UNION.count_ext.count = copy_size - 1; /* count is 1-based */
     else
-      packet_addr->COUNT_UNION.count.count = copy_size + SizeToCountOffset;
+      packet_addr->COUNT_UNION.count.count = copy_size - 1; /* count is 1-based */
 
     packet_addr->SRC_ADDR_LO_UNION.src_addr_31_0 = ptrlow32(cur_src);
     packet_addr->SRC_ADDR_HI_UNION.src_addr_63_32 = ptrhigh32(cur_src);
@@ -881,11 +837,12 @@ Elements are coded by the log2 of the element size in bytes (ie. element 0=1 byt
 This routine breaks a large rect into tiles that can be handled by hardware.  Pitches and offsets
 must be representable in terms of elements in all tiles of the copy.
 */
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildCopyRectCommand(
-    const std::function<void*(size_t)>& append, const hsa_pitched_ptr_t* dst,
-    const hsa_dim3_t* dst_offset, const hsa_pitched_ptr_t* src, const hsa_dim3_t* src_offset,
-    const hsa_dim3_t* range) {
+template <bool useGCR>
+void BlitSdma<useGCR>::BuildCopyRectCommand(const std::function<void*(size_t)>& append,
+                                            const hsa_pitched_ptr_t* dst,
+                                            const hsa_dim3_t* dst_offset,
+                                            const hsa_pitched_ptr_t* src,
+                                            const hsa_dim3_t* src_offset, const hsa_dim3_t* range) {
   // Returns the index of the first set bit (ie log2 of the largest power of 2 that evenly divides
   // width), the largest element that perfectly covers width.
   // width | 16 ensures that we don't return a higher element than is supported and avoids
@@ -1029,9 +986,9 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildCo
   }
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildFillCommand(
-    char* cmd_addr, uint32_t num_fill_command, void* ptr, uint32_t value, size_t count) {
+template <bool useGCR>
+void BlitSdma<useGCR>::BuildFillCommand(char* cmd_addr, uint32_t num_fill_command, void* ptr,
+                                        uint32_t value, size_t count) {
   char* cur_ptr = reinterpret_cast<char*>(ptr);
   const uint32_t maxDwordCount = kMaxSingleFillSize / sizeof(uint32_t);
   SDMA_PKT_CONSTANT_FILL* packet_addr = reinterpret_cast<SDMA_PKT_CONSTANT_FILL*>(cmd_addr);
@@ -1050,7 +1007,8 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildFi
 
     packet_addr->DATA_UNION.src_data_31_0 = value;
 
-    packet_addr->COUNT_UNION.count = (fill_count + SizeToCountOffset) * sizeof(uint32_t);
+    /* count is 1-based */
+    packet_addr->COUNT_UNION.count = (fill_count - 1) * sizeof(uint32_t);
 
     packet_addr++;
     cur_ptr += fill_count * sizeof(uint32_t);
@@ -1059,9 +1017,8 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildFi
   assert(count == 0 && "SDMA fill command count error.");
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildPollCommand(
-    char* cmd_addr, void* addr, uint32_t reference) {
+template <bool useGCR>
+void BlitSdma<useGCR>::BuildPollCommand(char* cmd_addr, void* addr, uint32_t reference) {
   SDMA_PKT_POLL_REGMEM* packet_addr =
       reinterpret_cast<SDMA_PKT_POLL_REGMEM*>(cmd_addr);
 
@@ -1081,9 +1038,8 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildPo
   packet_addr->DW5_UNION.retry_count = 0xfff;  // Retry forever.
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-              useGCR>::BuildAtomicDecrementCommand(char* cmd_addr, void* addr) {
+template <bool useGCR>
+void BlitSdma<useGCR>::BuildAtomicDecrementCommand(char* cmd_addr, void* addr) {
   SDMA_PKT_ATOMIC* packet_addr = reinterpret_cast<SDMA_PKT_ATOMIC*>(cmd_addr);
 
   memset(packet_addr, 0, sizeof(SDMA_PKT_ATOMIC));
@@ -1098,9 +1054,8 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
   packet_addr->SRC_DATA_HI_UNION.src_data_63_32 = 0xffffffff;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
-              useGCR>::BuildGetGlobalTimestampCommand(char* cmd_addr, void* write_address) {
+template <bool useGCR>
+void BlitSdma<useGCR>::BuildGetGlobalTimestampCommand(char* cmd_addr, void* write_address) {
   SDMA_PKT_TIMESTAMP* packet_addr =
       reinterpret_cast<SDMA_PKT_TIMESTAMP*>(cmd_addr);
 
@@ -1113,9 +1068,7 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset,
   packet_addr->ADDR_HI_UNION.addr_63_32 = ptrhigh32(write_address);
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildTrapCommand(
-    char* cmd_addr, uint32_t event_id) {
+template <bool useGCR> void BlitSdma<useGCR>::BuildTrapCommand(char* cmd_addr, uint32_t event_id) {
   SDMA_PKT_TRAP* packet_addr =
       reinterpret_cast<SDMA_PKT_TRAP*>(cmd_addr);
 
@@ -1125,17 +1078,13 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildTr
   packet_addr->INT_CONTEXT_UNION.int_ctx = event_id;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildHdpFlushCommand(
-    char* cmd_addr) {
+template <bool useGCR> void BlitSdma<useGCR>::BuildHdpFlushCommand(char* cmd_addr) {
   assert(cmd_addr != NULL);
   SDMA_PKT_POLL_REGMEM* addr = reinterpret_cast<SDMA_PKT_POLL_REGMEM*>(cmd_addr);
   memcpy(addr, &hdp_flush_cmd, flush_command_size_);
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildGCRCommand(
-    char* cmd_addr, bool invalidate) {
+template <bool useGCR> void BlitSdma<useGCR>::BuildGCRCommand(char* cmd_addr, bool invalidate) {
   assert(cmd_addr != NULL);
   assert(useGCR && "Unsupported SDMA command - GCR.");
   SDMA_PKT_GCR* addr = reinterpret_cast<SDMA_PKT_GCR*>(cmd_addr);
@@ -1154,25 +1103,16 @@ void BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::BuildGC
   addr->WORD2_UNION.GCR_CONTROL_GL2_RANGE = 0;
 }
 
-template <typename RingIndexTy, bool HwIndexMonotonic, int SizeToCountOffset, bool useGCR>
-uint64_t BlitSdma<RingIndexTy, HwIndexMonotonic, SizeToCountOffset, useGCR>::PendingBytes() {
-  RingIndexTy commit = atomic::Load(&cached_commit_index_, std::memory_order_acquire);
-  RingIndexTy hw_read_index = *reinterpret_cast<RingIndexTy*>(queue_resource_.Queue_read_ptr);
-  RingIndexTy read;
-  if (HwIndexMonotonic) {
-    read = hw_read_index;
-  } else {
-    RingIndexTy dist_to_read_index = WrapIntoRing(commit - hw_read_index);
-    read = commit - dist_to_read_index;
-  }
+template <bool useGCR> uint64_t BlitSdma<useGCR>::PendingBytes() {
+  uint64_t commit = atomic::Load(&cached_commit_index_, std::memory_order_acquire);
+  uint64_t hw_read_index = *reinterpret_cast<uint64_t*>(queue_resource_.Queue_read_ptr);
 
-  if (commit == read) return 0;
-  return bytes_queued_ - bytes_written_[WrapIntoRing(read)];
+  if (commit == hw_read_index) return 0;
+  return bytes_queued_ - bytes_written_[WrapIntoRing(hw_read_index)];
 }
 
-template class BlitSdma<uint32_t, false, 0, false>;
-template class BlitSdma<uint64_t, true, -1, false>;
-template class BlitSdma<uint64_t, true, -1, true>;
+template class BlitSdma<false>;
+template class BlitSdma<true>;
 
 }  // namespace amd
 }  // namespace rocr

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_cpu_agent.cpp

@@ -54,10 +54,9 @@
 
 namespace rocr {
 namespace AMD {
-CpuAgent::CpuAgent(HSAuint32 node, const HsaNodeProperties &node_props)
-    : core::Agent(
-          core::Runtime::runtime_singleton_->AgentDriver(core::DriverType::KFD),
-          node, kAmdCpuDevice),
+CpuAgent::CpuAgent(HSAuint32 node, const HsaNodeProperties& node_props,
+                   core::DriverType driver_type)
+    : core::Agent(core::Runtime::runtime_singleton_->AgentDriver(driver_type), node, kAmdCpuDevice),
       properties_(node_props) {
   InitRegionList();
 

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp

@@ -93,8 +93,8 @@ namespace AMD {
 const uint64_t CP_DMA_DATA_TRANSFER_CNT_MAX = (1 << 26);
 
 GpuAgent::GpuAgent(HSAuint32 node, const HsaNodeProperties& node_props, bool xnack_mode,
-                   uint32_t index)
-    : GpuAgentInt(node),
+                   uint32_t index, core::DriverType driver_type)
+    : GpuAgentInt(node, driver_type),
       properties_(node_props),
       current_coherency_type_(HSA_AMD_COHERENCY_TYPE_COHERENT),
       scratch_used_large_(0),
@@ -106,7 +106,6 @@ GpuAgent::GpuAgent(HSAuint32 node, const HsaNodeProperties& node_props, bool xna
       memory_max_frequency_(0),
       enum_index_(index),
       ape1_base_(0),
-      ape1_size_(0),
       pending_copy_req_ref_(0),
       pending_copy_stat_check_ref_(0),
       sdma_blit_used_mask_(0),
@@ -712,10 +711,6 @@ core::Blit* GpuAgent::CreateBlitSdma(bool use_xgmi, int rec_eng) {
   const size_t copy_size_overrides[2] = {0x3fffff, 0x3fffffff};
 
   switch (isa_->GetMajorVersion()) {
-    case 7:
-    case 8:
-      sdma = new BlitSdmaV2V3();
-      break;
     case 9:
       sdma = new BlitSdmaV4();
       copy_size_override = (isa_->GetMinorVersion() == 0 && isa_->GetStepping() == 10) ?
@@ -2568,6 +2563,10 @@ hsa_status_t GpuAgent::PcSamplingIterateConfig(hsa_ven_amd_pcs_iterate_configura
   if (ret != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
 
   for (uint32_t i = 0; i < size; i++) {
+    if ((isa_->GetMajorVersion() == 12 && (isa_->GetMinorVersion() == 0)) &&
+        sampleInfoList[i].method == HSA_PC_SAMPLING_METHOD_KIND_STOCHASTIC_V1) {
+      continue;
+    }
     hsa_ven_amd_pcs_configuration_t hsaPcSampling;
     if (ConvertHsaKmtPcSamplingInfoToHsa(&sampleInfoList[i], &hsaPcSampling) == HSA_STATUS_SUCCESS
         && cb(&hsaPcSampling, cb_data) == HSA_STATUS_INFO_BREAK)
@@ -2614,6 +2613,10 @@ hsa_status_t GpuAgent::PcSamplingCreateFromId(HsaPcSamplingTraceId ioctlId,
   if (sampling_method == HSA_VEN_AMD_PCS_METHOD_HOSTTRAP_V1) {
     pcs_data = &pcs_hosttrap_data_;
   } else if (sampling_method == HSA_VEN_AMD_PCS_METHOD_STOCHASTIC_V1) {
+    if (isa_->GetMajorVersion() == 12 && (isa_->GetMinorVersion() == 0)) {
+      return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+    }
+
     pcs_data = &pcs_stochastic_data_;
   } else {
     // Unsupported sampling method
@@ -3093,6 +3096,7 @@ hsa_status_t GpuAgent::PcSamplingFlushDeviceBuffers(
 
   const uint32_t atomic_ex_cmd_sz = 9;
   const uint32_t wait_reg_mem_cmd_sz = 7;
+  const uint32_t acquire_mem_cmd_sz = 8;
   const uint32_t dma_data_cmd_sz = 7;
   const uint32_t copy_data_cmd_sz = 6;
   const uint32_t write_data_cmd_sz = 5;
@@ -3225,6 +3229,20 @@ hsa_status_t GpuAgent::PcSamplingFlushDeviceBuffers(
   cmd_data[i++] = PM4_WAIT_REG_MEM_DW6(PM4_WAIT_REG_MEM_POLL_INTERVAL(4) |
                                        PM4_WAIT_REG_MEM_OPTIMIZE_ACE_OFFLOAD_MODE);
 
+  // For GFX1200 and GFX1201 only - add an ACQUIRE_MEM packet to flush L2 cache before DMA.
+  // This ensures that any data written by the trap handler is visible to the DMA engine.
+  if ((isa_->GetMajorVersion() == 12) && (isa_->GetMinorVersion() == 0)) {
+    cmd_data[i++] =
+        PM4_HDR(PM4_HDR_IT_OPCODE_ACQUIRE_MEM, acquire_mem_cmd_sz, isa_->GetMajorVersion());
+    cmd_data[i++] = 0;                                // DW1: COHER_CNTL
+    cmd_data[i++] = 0;                                // DW2: COHER_SIZE
+    cmd_data[i++] = 0;                                // DW3: COHER_SIZE_HI
+    cmd_data[i++] = 0;                                // DW4: COHER_BASE_LO
+    cmd_data[i++] = 0;                                // DW5: COHER_BASE_HI
+    cmd_data[i++] = 4;                                // DW6: POLL_INTERVAL
+    cmd_data[i++] = PM4_ACQUIRE_MEM_GCR_CNTL_GL2_WB;  // DW7: GCR_CNTL (GL2_WB=1, RANGE=ALL)
+  }
+
   uint8_t* buffer_temp = buffer[which_buffer];
 
   for (copy_bytes = std::min(to_copy, (uint32_t)CP_DMA_DATA_TRANSFER_CNT_MAX); 0 < to_copy;

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_topology.cpp

@@ -68,6 +68,9 @@
 #include "core/inc/amd_memory_region.h"
 #include "core/inc/runtime.h"
 #include "core/util/utils.h"
+#ifdef HSAKMT_VIRTIO_ENABLED
+#include "core/inc/amd_virtio_driver.h"
+#endif
 
 extern r_debug _amdgpu_r_debug;
 
@@ -75,16 +78,21 @@ namespace rocr {
 namespace AMD {
 // Anonymous namespace.
 namespace {
-#if _WIN32
-constexpr size_t num_drivers = 0;
-#elif __linux__
-constexpr size_t num_drivers = 2;
-#endif
 
-const std::array<std::function<hsa_status_t(std::unique_ptr<core::Driver>&)>, num_drivers>
+const std::array<std::function<hsa_status_t(std::unique_ptr<core::Driver>&)>,
+#if _WIN32
+                 0
+#elif __linux__
+                 static_cast<size_t>(core::DriverType::NUM_DRIVER_TYPES)
+#endif
+                 >
     discover_driver_funcs = {
 #ifdef __linux__
-        KfdDriver::DiscoverDriver, XdnaDriver::DiscoverDriver
+        KfdDriver::DiscoverDriver,
+        XdnaDriver::DiscoverDriver,
+#ifdef HSAKMT_VIRTIO_ENABLED
+        KfdVirtioDriver::DiscoverDriver,
+#endif
 #endif
 };
 
@@ -110,14 +118,14 @@ bool InitializeDriver(std::unique_ptr<core::Driver>& driver) {
   return true;
 }
 
-void DiscoverCpu(HSAuint32 node_id, HsaNodeProperties& node_prop) {
-  CpuAgent* cpu = new CpuAgent(node_id, node_prop);
+void DiscoverCpu(HSAuint32 node_id, HsaNodeProperties& node_prop, core::DriverType driver_type) {
+  CpuAgent* cpu = new CpuAgent(node_id, node_prop, driver_type);
   cpu->Enable();
   core::Runtime::runtime_singleton_->RegisterAgent(cpu, true);
 }
 
 GpuAgent* DiscoverGpu(HSAuint32 node_id, HsaNodeProperties& node_prop, bool xnack_mode,
-                      bool enabled) {
+                      bool enabled, core::DriverType driver_type) {
   GpuAgent* gpu = nullptr;
   if (node_prop.NumFComputeCores == 0) {
       // Ignore non GPUs.
@@ -125,7 +133,7 @@ GpuAgent* DiscoverGpu(HSAuint32 node_id, HsaNodeProperties& node_prop, bool xnac
   }
   try {
     gpu = new GpuAgent(node_id, node_prop, xnack_mode,
-                       core::Runtime::runtime_singleton_->gpu_agents().size());
+                       core::Runtime::runtime_singleton_->gpu_agents().size(), driver_type);
 
     const HsaVersionInfo& kfd_version = core::Runtime::runtime_singleton_->KfdVersion().version;
 
@@ -152,7 +160,7 @@ GpuAgent* DiscoverGpu(HSAuint32 node_id, HsaNodeProperties& node_prop, bool xnac
         node_prop.Capability.ui32.SRAM_EDCSupport = 1;
         delete gpu;
         gpu = new GpuAgent(node_id, node_prop, xnack_mode,
-                           core::Runtime::runtime_singleton_->gpu_agents().size());
+                           core::Runtime::runtime_singleton_->gpu_agents().size(), driver_type);
       }
     }
   } catch (const hsa_exception& e) {
@@ -257,24 +265,29 @@ void SurfaceGpuList(std::vector<int32_t>& gpu_list, bool xnack_mode, bool enable
   const int32_t invalidIdx = -1;
   int32_t list_sz = gpu_list.size();
   HsaNodeProperties node_prop = {0};
-  const auto& gpu_driver = core::Runtime::runtime_singleton_->AgentDriver(core::DriverType::KFD);
-  for (int32_t idx = 0; idx < list_sz; idx++) {
-    if (gpu_list[idx] == invalidIdx) {
-      break;
+  for (const auto& gpu_driver : core::Runtime::runtime_singleton_->AgentDrivers()) {
+    if (!core::Runtime::IsGPUDriver(gpu_driver->kernel_driver_type_)) {
+      continue;
     }
 
-    // Obtain properties of the node
-    hsa_status_t ret = gpu_driver.GetNodeProperties(node_prop, gpu_list[idx]);
-    assert(ret == HSA_STATUS_SUCCESS && "Error in getting Node Properties");
+    for (int32_t idx = 0; idx < list_sz; idx++) {
+      if (gpu_list[idx] == invalidIdx) {
+        break;
+      }
 
-    // disable interrupt signal for DTIF platform
-    if (core::Runtime::runtime_singleton_->flag().enable_dtif())
-      core::g_use_interrupt_wait = false;
+      // Obtain properties of the node
+      hsa_status_t ret = gpu_driver->GetNodeProperties(node_prop, gpu_list[idx]);
+      assert(ret == HSA_STATUS_SUCCESS && "Error in getting Node Properties");
 
-    // Instantiate a Gpu device. The IO links
-    // of this node have already been registered
-    assert((node_prop.NumFComputeCores != 0) && "Improper node used for GPU device discovery.");
-    DiscoverGpu(gpu_list[idx], node_prop, xnack_mode, enabled);
+      // disable interrupt signal for DTIF platform
+      if (core::Runtime::runtime_singleton_->flag().enable_dtif())
+        core::g_use_interrupt_wait = false;
+
+      // Instantiate a Gpu device. The IO links
+      // of this node have already been registered
+      assert((node_prop.NumFComputeCores != 0) && "Improper node used for GPU device discovery.");
+      DiscoverGpu(gpu_list[idx], node_prop, xnack_mode, enabled, gpu_driver->kernel_driver_type_);
+    }
   }
 }
 
@@ -335,7 +348,7 @@ bool BuildTopology() {
     /// @todo: Add support for AIEs.
     // Query if env ROCR_VISIBLE_DEVICES is defined. If defined
     // determine number and order of GPU devices to be surfaced.
-    if (filter && driver->kernel_driver_type_ == core::DriverType::KFD) {
+    if (filter && (core::Runtime::IsGPUDriver(driver->kernel_driver_type_))) {
       rvdFilter.BuildRvdTokenList();
       rvdFilter.BuildDeviceUuidList(node_props_vec);
       visibleCnt = rvdFilter.BuildUsrDeviceList();
@@ -350,7 +363,7 @@ bool BuildTopology() {
     for (auto& node_props : node_props_vec) {
       if (node_props.NumCPUCores) {
         // Node has CPU cores so instantiate a CPU agent.
-        DiscoverCpu(node_id, node_props);
+        DiscoverCpu(node_id, node_props, driver->kernel_driver_type_);
       }
 
       if (node_props.NumNeuralCores) {

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/runtime.cpp

@@ -54,9 +54,7 @@
 #include <sys/mman.h>
 #include <sys/socket.h>
 #include <sys/un.h>
-#include <iostream>
 #include <thread>
-#include <chrono>
 
 #include "core/inc/runtime.h"
 #include "core/inc/hsa_table_interface.h"
@@ -3078,8 +3076,8 @@ Agent* Runtime::GetSVMPrefetchAgent(void* ptr, size_t size) {
   return agents_by_node_[prefetch_node][0];
 }
 
-hsa_status_t Runtime::DmaBufExport(const void* ptr, size_t size, int* dmabuf,
-                                           uint64_t* offset, uint64_t flags) {
+hsa_status_t Runtime::DmaBufExport(const void* ptr, size_t size, int* dmabuf, uint64_t* offset,
+                                   uint64_t flags) {
 #ifdef __linux__
   ScopedAcquire<KernelSharedMutex::Shared> lock(memory_lock_.shared());
   // Lookup containing allocation.
@@ -3090,18 +3088,23 @@ hsa_status_t Runtime::DmaBufExport(const void* ptr, size_t size, int* dmabuf,
         (ptr < reinterpret_cast<const uint8_t*>(mem->first) + mem->second.size)) {
       // Check size is in bounds.
       if (uintptr_t(ptr) - uintptr_t(mem->first) + size <= mem->second.size) {
-        // Check allocation is on GPU
-        if (mem->second.region->owner()->device_type() != Agent::kAmdGpuDevice)
-          return HSA_STATUS_ERROR_INVALID_AGENT;
+        switch (mem->second.region->owner()->device_type()) {
+          case Agent::kAmdGpuDevice: {
+            auto* owner = static_cast<AMD::GpuAgent*>(mem->second.region->owner());
 
-        rocr::AMD::GpuAgent* owner =
-                    static_cast<AMD::GpuAgent*>(mem->second.region->owner());
-
-        if (flags & HSA_AMD_DMABUF_MAPPING_TYPE_PCIE &&
-            !owner->is_xgmi_cpu_gpu() &&
-            !owner->LargeBarEnabled()) {
-            return (hsa_status_t)HSA_STATUS_ERROR_NOT_SUPPORTED;
+            if (flags & HSA_AMD_DMABUF_MAPPING_TYPE_PCIE && !owner->is_xgmi_cpu_gpu() &&
+                !owner->LargeBarEnabled()) {
+              return static_cast<hsa_status_t>(HSA_STATUS_ERROR_NOT_SUPPORTED);
+            }
+          } break;
+          case Agent::kAmdCpuDevice:
+            return HSA_STATUS_ERROR_INVALID_AGENT;
+          case Agent::kAmdAieDevice:
+            break;
+          case Agent::kUnknownDevice:
+            return HSA_STATUS_ERROR_INVALID_AGENT;
         }
+
         int fd;
         uint64_t off;
         hsa_status_t err = mem->second.region->owner()->driver().ExportDMABuf(
@@ -3194,10 +3197,11 @@ hsa_status_t Runtime::VMemoryAddressFree(void* va, size_t size) {
 
   if (it->second.use_count > 0) return HSA_STATUS_ERROR_RESOURCE_FREE;
 
-  if (it->second.registered)
+  if (it->second.registered) {
     if (HSAKMT_CALL(hsaKmtFreeMemory(it->second.os_addr, size)) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
-  else
+  } else {
     if (munmap(it->second.os_addr, size)) return HSA_STATUS_ERROR;
+  }
 
   reserved_address_map_.erase(it);
   return HSA_STATUS_SUCCESS;
@@ -3213,7 +3217,7 @@ hsa_status_t Runtime::VMemoryHandleCreate(const MemoryRegion* region, size_t siz
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
 
   ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
-  void *user_mode_driver_handle;
+  ThunkHandle user_mode_driver_handle;
   hsa_status_t status =
       region->Allocate(size, alloc_flags, &user_mode_driver_handle, 0);
   if (status == HSA_STATUS_SUCCESS) {
@@ -3230,7 +3234,7 @@ hsa_status_t Runtime::VMemoryHandleCreate(const MemoryRegion* region, size_t siz
 
 hsa_status_t Runtime::VMemoryHandleRelease(hsa_amd_vmem_alloc_handle_t memoryOnlyHandle) {
   ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
-  auto memoryHandleIt = memory_handle_map_.find(reinterpret_cast<void*>(memoryOnlyHandle.handle));
+  auto memoryHandleIt = memory_handle_map_.find(MemoryHandle::Convert(memoryOnlyHandle));
 
   if (memoryHandleIt == memory_handle_map_.end()) {
     debug_warning(false && "Can't find memory handle");
@@ -3285,7 +3289,7 @@ hsa_status_t Runtime::VMemoryHandleMap(void* va, size_t size, size_t in_offset,
     if (reinterpret_cast<uint8_t*>(va) + size > lowerMappedHandleIt->first) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   }
 
-  auto memoryHandleIt = memory_handle_map_.find(reinterpret_cast<void*>(memoryOnlyHandle.handle));
+  auto memoryHandleIt = memory_handle_map_.find(MemoryHandle::Convert(memoryOnlyHandle));
   if (memoryHandleIt == memory_handle_map_.end()) {
     debug_warning(false && "Can't find memory handle");
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
@@ -3648,7 +3652,7 @@ hsa_status_t Runtime::VMemoryExportShareableHandle(int* dmabuf_fd,
                                                    hsa_amd_vmem_alloc_handle_t handle,
                                                    uint64_t flags) {
   *dmabuf_fd = -1;
-  auto memoryHandle = memory_handle_map_.find((void*)handle.handle);
+  auto memoryHandle = memory_handle_map_.find(MemoryHandle::Convert(handle));
   if (memoryHandle == memory_handle_map_.end()) {
     debug_warning(false && "Can't find memory handle");
     return HSA_STATUS_ERROR_INVALID_ALLOCATION;
@@ -3744,7 +3748,7 @@ hsa_status_t Runtime::VMemoryRetainAllocHandle(hsa_amd_vmem_alloc_handle_t* mapp
 hsa_status_t Runtime::VMemoryGetAllocPropertiesFromHandle(hsa_amd_vmem_alloc_handle_t allocHandle,
                                                           const core::MemoryRegion** mem_region,
                                                           hsa_amd_memory_type_t* type) {
-  auto memoryHandleIt = memory_handle_map_.find(reinterpret_cast<void*>(allocHandle.handle));
+  auto memoryHandleIt = memory_handle_map_.find(MemoryHandle::Convert(allocHandle));
   if (memoryHandleIt == memory_handle_map_.end()) return HSA_STATUS_ERROR_INVALID_ALLOCATION;
 
   *mem_region = memoryHandleIt->second.region;

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/trap_handler/trap_handler_gfx12.s

@@ -56,6 +56,7 @@
 .set SQ_WAVE_EXCP_FLAG_PRIV_HT_SHIFT           , 7
 .set SQ_WAVE_EXCP_FLAG_PRIV_WAVE_START_SHIFT   , 8
 .set SQ_WAVE_EXCP_FLAG_PRIV_WAVE_END_SHIFT     , 9
+.set SQ_WAVE_EXCP_FLAG_PRIV_PERF_SNAPSHOT      , 10
 .set SQ_WAVE_EXCP_FLAG_PRIV_TRAP_AFTER_INST_SHIFT , 11
 .set SQ_WAVE_EXCP_FLAG_PRIV_XNACK_ERROR_SHIFT  , 12
 
@@ -74,6 +75,7 @@
 .set SQ_WAVE_STATE_PRIV_HALT_BFE               , (SQ_WAVE_STATE_PRIV_HALT_SHIFT | (1 << 16))
 .set SQ_WAVE_STATE_PRIV_HALT_SHIFT             , 14
 .set SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT , 2
+
 .set TRAP_ID_ABORT                             , 2
 .set TRAP_ID_DEBUGTRAP                         , 3
 .set TTMP6_SAVED_STATUS_HALT_MASK              , (1 << TTMP6_SAVED_STATUS_HALT_SHIFT)
@@ -87,140 +89,273 @@
 .set TTMP11_DEBUG_ENABLED_SHIFT                , 23
 .set TTMP_PC_HI_SHIFT                          , 7
 
-// ABI between first and second level trap handler:
-//   { ttmp1, ttmp0 } = TrapID[3:0], zeros, PC[47:0]
-//   ttmp11 = 0[7:0], DebugEnabled[0], 0[15:0], NoScratch[0], 0[5:0]
-//   ttmp12 = SQ_WAVE_STATE_PRIV
-//   ttmp14 = TMA[31:0]
-//   ttmp15 = TMA[63:32]
+.set TTMP13_HT_FLAG_BIT                        , 22           // TTMP13 bit for host‑trap
+.set TTMP13_STOCH_FLAG_BIT                     , 21           // TTMP13 bit for stochastic
+.set TTMP13_BUF_FULL_BIT                       , 31           // TTMP13 bit – buf full mark
+.set TTMP8_DISPATCH_ID_MASK                    , 0X1FFFFFF
+// Per-sample data layout within the device buffer. Each sample is 64 bytes.
+// These are offsets from the start of a specific sample slot in the device buffer.
 
-trap_entry:
-  // Clear ttmp3 as it will contain the exception code.
-  s_mov_b32            ttmp3, 0
+.set SAMPLE_OFF_BYTES_PER_SAMPLE               , 0x40         // bytes per sample slot
 
+.set SAMPLE_OFF_PC_HOST                        , 0x00         // original PC (host only)
+.set SAMPLE_OFF_EXEC_LOHI                      , 0x08         // saved EXEC low/high
+.set SAMPLE_OFF_WGID_XY                        , 0x10         // WG id X / Y
+.set SAMPLE_OFF_WGID_Z_WAVE                    , 0x18         // WG id Z
+.set SAMPLE_OFF_TIMESTAMP                      , 0x30         // 64 bit realtime counter
+.set SAMPLE_OFF_HW_ID                          , 0x20         // HW_ID (values combined from the HW_ID1 + HW_ID2)
+.set SAMPLE_OFF_SNAPSHOT_DATA                  , 0x24
+.set SAMPLE_OFF_CORRELATION                    , 0x38         // doorbell + dispatch id
+.set SAMPLE_OFF_BUF_WRITTEN_VAL                , 0x10         // Offset to buf_written_val0/1 in pcs_sampling_data_t
+.set SAMPLE_OFF_BUF_SIZE                       , 0x8          // Offset to buf_size in pcs_sampling_data_t
+.set SAMPLE_OFF_DONE_SIG0                      , 0x18         // Offset for done_sig0 (hsa_signal_t handle for buffer 0)
+.set SAMPLE_OFF_DONE_SIG1                      , 0x28         // Offset for done_sig1 (hsa_signal_t handle for buffer 1)
+.set SAMPLE_OFF_SIGNAL_VALUE                   , 0x8          // Offset within signal structure to value field
+.set SAMPLE_OFF_EVENT_MAILBOX0                 , 0x10         // Offset for event mailbox pointer for buffer 0
+.set SAMPLE_OFF_EVENT_MAILBOX1                 , 0x20         // Offset for event mailbox pointer for buffer 1
+
+.set WAVE_ID_MASK                              , 0x1f         // Mask to extract Wave ID from TTMP register.
+.set BUF_INDEX_MASK                            , 0x7fffffff   // strip bit31 from add_x2
+.set SAMPLE_OFF_BUF_WRITTEN_VAL                , 0x10         // Offset to buf_written_val0/1 in pcs_sampling_data_t
+.set SAMPLE_INDEX_WIDTH                        , 31           // The sample index is 63 bits; the high part is 31 bits.
+
+.set HW_REG_SHADER_HW_ID1                      , 0xf817
+.set HW_REG_SHADER_HW_ID2                      , 0xf818
+.set HW_REG_SQ_PERF_SNAPSHOT_PC_LO             , 0xf80b
+.set HW_REG_SQ_PERF_SNAPSHOT_PC_HI             , 0xf80c
+.set HW_REG_SQ_PERF_SNAPSHOT_DATA1             , 0xf80f
+.set HW_REG_SQ_PERF_SNAPSHOT_DATA2             , 0xf810
+.set HW_REG_SQ_PERF_SNAPSHOT_DATA              , 0xf81b
+
+  // Macro to store the Correlation ID (Dispatch ID and Doorbell ID) into the current sample slot
+  //
+  // Assumes the following registers are set before it is called:
+  //   v[0:1]:Must contain the 64-bit base address of the target sample slot
+  //   ttmp8 :Must contain the dispatch ID in bits [24:0]
+  //   exec  :Must be set to 0x1 to ensure operations apply only to lane 0
+  //
+  // Clobbers the following registers:
+  //   v[2:3]:Used for [dispatch_id, doorbell_id]
+  //   ttmp6 :Used as scratch register
+.macro STORE_CORRELATION_ID
+  s_sendmsg_rtn_b32 ttmp6, sendmsg(MSG_RTN_GET_DOORBELL)    // Gets current queue's doorbell ID into ttmp6.
+  s_wait_kmcnt      0
+  s_and_b32         ttmp6, ttmp6, DOORBELL_ID_MASK          // Mask to get actual doorbell ID.
+  v_writelane_b32   v3, ttmp6, 0                            // Store doorbell ID into high part of v[2:3] (via v3).
+  s_and_b32         ttmp6, ttmp8, TTMP8_DISPATCH_ID_MASK    // Get dispatch ID from ttmp8 into ttmp6
+  v_writelane_b32   v2, ttmp6, 0                            // Store dispatch ID into low part of v[2:3] (via v2)
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_CORRELATION, scope:SCOPE_SYS  // Store {dispatch_id, doorbell_id} into sample slot.
+                                                                       // v[0:1] = sample slot base address.
+                                                                       // v[2] = dispatch_id, v[3] = doorbell_id.
+.endm
+
+  // Macro to store the HW_ID registers into the current sample slot
+  //
+  // Assumes the following registers are set before it is called:
+  //   v[0:1]: Must contain the 64-bit base address of the target sample slot.
+  //   exec  : Must be set to 0x1 to ensure operations apply only to lane 0.
+  //
+  // Clobbers the following registers:
+  //   v[2:3]: Used to stage the data for the global store.
+  //   ttmp6 : Used as scratch registers.
+.macro STORE_HW_ID
+  // Current ROCr API determines single dword for HW_ID, while this information is scattered accross two
+  // dword registers HW_ID1 and HW_ID2 on GFX10+ architectures.
+  // Thus, we combine values from HW_ID1 and HW_ID2 into a single dword HW_ID with the following layout:
+  // WAVE_ID[4:0]
+  // QUEUE_ID[8:5]
+  // RESERVED [9]
+  // WGP_ID[13:10]
+  // SIMD_ID[15:14]
+  // SA_ID[16]
+  // ME_ID[17]
+  // SE_ID[19:18]
+  // PIPE_ID[21:20]
+  // RESERVED [22]
+  // WG_ID[27:23]
+  // VM_ID[31:28]
+
+  // Note: We don't show DP_RATE and STATE_ID that are useless for compute kernels
+  // Also, we reduced SE_ID to 2 bits as there's only a maximum of 4 SEs on existing gfx12.0 parts
+  // Finally, ME_ID is reduced to 1 bit as wavefronts are dispatched from either ME0 or ME1 in gfx12.
+  // Bits 9 and 22 are reserved for a future use.
+
+  s_getreg_b32      ttmp6, HW_REG_SHADER_HW_ID1             // Put HW_ID1 in ttmp6
+  v_and_b32         v2, ttmp6, 0x1feffcff                   // Mask DP_RATE, SE_ID[2] and SIMD_ID
+  v_and_b32         v3, ttmp6, 0x300                        // Put SIMD_ID into ttmp6[8:9]
+  v_lshl_or_b32     v2, v3, 6, v2                           // Put SIMD_ID into v2[15:14]
+  s_getreg_b32      ttmp6, HW_REG_SHADER_HW_ID2             // Put HW_ID2 in ttmp6
+  v_and_b32         v3, ttmp6, 0xf000000                    // v3 = VM_ID in bits 27:24
+  v_lshl_or_b32     v2, v3, 4, v2                           // Put VM_ID into v2[31:28]
+  v_and_b32         v3, ttmp6, 0x1f0000                     // v3 = WG_ID in bits 20:16
+  v_lshl_or_b32     v2, v3, 7, v2                           // Put WG_ID in v2[27:23]
+  v_and_b32         v3, ttmp6, 0x100                        // v3 = ME_ID[0] in bit 8
+  v_lshl_or_b32     v2, v3, 9, v2                           // Put ME_ID in v2[17]
+  v_and_b32         v3, ttmp6, 0x30                         // v3 = PIPE_ID in bits 5:4
+  v_lshl_or_b32     v2, v3, 16, v2                          // Put PIPE_ID in v2[21:20]
+  v_and_b32         v3, ttmp6, 0xf                          // v3 = QUEUE_ID in bits 3:0
+  v_lshl_or_b32     v2, v3, 5, v2                           // Put QUEUE_ID in v2[8:5]
+  global_store_b32  v[0:1], v2, off, offset:SAMPLE_OFF_HW_ID, scope:SCOPE_SYS  // store HW_ID
+.endm
+
+// ABI (Application Binary Interface) between first and second-level trap handler:
+//   ttmp0: PC_LO[31:0] (Program Counter Low)
+//   ttmp1: PC_HI[15:0] (Program Counter High, bits 0-15), TrapID[3:0] (in bits 28-31 of original PC_HI)
+//   ttmp11: 0[7:0], DebugEnabled[0], 0[15:0], NoScratch[0], 0[5:0]
+//   ttmp12: SQ_WAVE_STATE_PRIV (Private wave state register value).
+//   ttmp14: TMA[31:0] - TMA_LO (Trap Memory Argument Low - base address for trap handler data, low 32 bits).
+//   ttmp15: TTMA[63:32] - TMA_HI (Trap Memory Argument High - base address for trap handler data, high 32 bits).
+//   For PC Sampling, this points to pcs_hosttrap_data_ or pcs_stochastic_data_
+ trap_entry:
+
+  s_mov_b32         ttmp3, 0
+
+.check_hosttrap:
+
+  // ttmp[14:15] points to TMA.
+  // Available: ttmp[2:3], ttmp[4:5], ttmp6, ttmp[10:11]
+  s_getreg_b32      ttmp2, hwreg(HW_REG_EXCP_FLAG_PRIV)     // On gfx12, EXCP_FLAG_PRIV.b7
+  s_bitcmp1_b32     ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_HT_SHIFT
+  s_cbranch_scc0    .check_stochastic
+
+  // It's a Host Trap event.
+  s_load_b64        ttmp[14:15], ttmp[14:15], 0x0, scope:SCOPE_CU         // ttmp[14:15]=*host_trap_buffers
+  s_bitset1_b32     ttmp13, TTMP13_HT_FLAG_BIT              // set bit 22 in TTMP13
+
+  // Clear the Host Trap flag in the hardware register to acknowledge the event
+  s_setreg_imm32_b32 hwreg(HW_REG_EXCP_FLAG_PRIV, SQ_WAVE_EXCP_FLAG_PRIV_HT_SHIFT,1), 0
+  s_wait_kmcnt      0                                       // Ensure previous load is complete.
+  s_branch          .profile_trap_handlers
+
+.check_stochastic:
+  s_getreg_b32      ttmp2, hwreg(HW_REG_EXCP_FLAG_PRIV)     // EXCP_FLAG_PRIV.b10=stochastic_sample_trap
+  s_bitcmp1_b32     ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_PERF_SNAPSHOT // Test Performance Snapshot bit.
+
+  s_cbranch_scc0    .check_exceptions                       // If not Stochastic, check for other exceptions.
+
+  s_load_b64           ttmp[14:15], ttmp[14:15], 0x8, scope:SCOPE_CU         // ttmp[14:15]=*stoch_trap_buf
+  s_wait_kmcnt      0
+
+  s_bitset1_b32     ttmp13, TTMP13_STOCH_FLAG_BIT           // set bit 21 in TTMP13
+
+  s_setreg_imm32_b32 hwreg(HW_REG_EXCP_FLAG_PRIV, SQ_WAVE_EXCP_FLAG_PRIV_PERF_SNAPSHOT,1), 0 // Clear the perf_snapshot flag
+  s_branch          .profile_trap_handlers
+
+  // Check if this is a trap (s_trap instruction) or a hardware exception.
+  // Extract TrapID from ttmp1 (which contains PC_HI).
   // Branch if not a trap (an exception instead).
-  s_bfe_u32            ttmp2, ttmp1, SQ_WAVE_PC_HI_TRAP_ID_BFE
-  s_cbranch_scc0       .check_exceptions
+  s_bfe_u32         ttmp2, ttmp1, SQ_WAVE_PC_HI_TRAP_ID_BFE // ttmp2 = TrapID
+  s_cbranch_scc0       .check_exceptions			             // If TrapID is 0, it's an exception, so branch.
 
   // If caused by s_trap then advance PC, then figure out the trap ID:
   // - if trapID is DEBUGTRAP and debugger is attach, report WAVE_TRAP,
   // - if trapID is ABORTTRAP, report WAVE_ABORT,
   // - report WAVE_TRAP for any other trap ID.
-  s_add_u32            ttmp0, ttmp0, 0x4
-  s_addc_u32           ttmp1, ttmp1, 0x0
+  s_add_u32         ttmp0, ttmp0, 0x4                       // PC_LO += 4
+  s_addc_u32        ttmp1, ttmp1, 0x0                       // PC_HI += carry.
 
   // If llvm.debugtrap and debugger is not attached.
-  s_cmp_eq_u32         ttmp2, TRAP_ID_DEBUGTRAP
-  s_cbranch_scc0       .not_debug_trap
+  s_cmp_eq_u32      ttmp2, TRAP_ID_DEBUGTRAP
+  s_cbranch_scc0    .not_debug_trap
 
-  s_bitcmp1_b32        ttmp11, TTMP11_DEBUG_ENABLED_SHIFT
-  s_cbranch_scc0       .check_exceptions
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
+  s_bitcmp1_b32     ttmp11, TTMP11_DEBUG_ENABLED_SHIFT
+  s_cbranch_scc0    .check_exceptions
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
 
 .not_debug_trap:
-  s_cmp_eq_u32         ttmp2, TRAP_ID_ABORT
-  s_cbranch_scc0       .not_abort_trap
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_ABORT_M0
-  s_branch             .check_exceptions
+  s_cmp_eq_u32      ttmp2, TRAP_ID_ABORT
+  s_cbranch_scc0    .not_abort_trap
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_ABORT_M0
+  s_branch          .check_exceptions
 
 .not_abort_trap:
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
+
+  s_bitcmp1_b32     ttmp8, TTMP8_DEBUG_FLAG_SHIFT
+  s_cbranch_scc0    .check_exceptions
 
-  // We need to explititly look for all exceptions we want to report to the
-  // host:
-  // - EXCP_FLAG_PRIV.XNACK_ERROR (&& EXCP_FLAG_PRIV.MEMVIOL)
-  //                                                 -> WAVE_MEMORY_VIOLATION
-  // - EXCP_FLAG_PRIV.MEMVIOL (and !EXCP_FLAG_PRIV.XNACK_ERROR)
-  //                                                 -> WAVE_APERTURE_VIOLATION
-  // - EXCP_FLAG_PRIV.ILLEGAL_INST                   -> WAVE_ILLEGAL_INSTRUCTION
-  // - EXCP_FLAG_PRIV.WAVE_START                     -> WAVE_TRAP
-  // - EXCP_FLAG_PRIV.WAVE_END && TRAP_CTRL.WAVE_END -> WAVE_TRAP
-  // - TRAP_CTRL.TRAP_AFTER_INST                     -> WAVE_TRAP
-  // - EXCP_FLAG_PRIV.ADDR_WATCH && TRAP_CTL.WATCH   -> WAVE_TRAP
-  // - (EXCP_FLAG_USER[ALU] & TRAP_CTRL[ALU]) != 0   -> WAVE_MATH_ERROR
 .check_exceptions:
-  s_getreg_b32         ttmp2, hwreg(HW_REG_EXCP_FLAG_PRIV)
-  s_getreg_b32         ttmp13, hwreg(HW_REG_TRAP_CTRL)
+  s_getreg_b32      ttmp2, hwreg(HW_REG_EXCP_FLAG_PRIV)
+  s_getreg_b32      ttmp13, hwreg(HW_REG_TRAP_CTRL)
 
-  s_bitcmp1_b32        ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_XNACK_ERROR_SHIFT
-  s_cbranch_scc0       .not_memory_violation
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_MEMORY_VIOLATION_M0
+  s_bitcmp1_b32     ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_XNACK_ERROR_SHIFT
+  s_cbranch_scc0    .not_memory_violation
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_MEMORY_VIOLATION_M0
 
   // Aperture violation requires XNACK_ERROR == 0.
-  s_branch             .not_aperture_violation
+  s_branch          .not_aperture_violation
 
 .not_memory_violation:
-  s_bitcmp1_b32        ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_MEMVIOL_SHIFT
-  s_cbranch_scc0       .not_aperture_violation
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_APERTURE_VIOLATION_M0
+  s_bitcmp1_b32     ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_MEMVIOL_SHIFT
+  s_cbranch_scc0    .not_aperture_violation
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_APERTURE_VIOLATION_M0
 
 .not_aperture_violation:
-  s_bitcmp1_b32        ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_ILLEGAL_INST_SHIFT
-  s_cbranch_scc0       .not_illegal_instruction
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_ILLEGAL_INSTRUCTION_M0
+  s_bitcmp1_b32     ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_ILLEGAL_INST_SHIFT
+  s_cbranch_scc0    .not_illegal_instruction
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_ILLEGAL_INSTRUCTION_M0
 
 .not_illegal_instruction:
-  s_bitcmp1_b32        ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_WAVE_START_SHIFT
-  s_cbranch_scc0       .not_wave_end
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
+  s_bitcmp1_b32     ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_WAVE_START_SHIFT
+  s_cbranch_scc0    .not_wave_end
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
 
 .not_wave_start:
-  s_bitcmp1_b32        ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_WAVE_END_SHIFT
-  s_cbranch_scc0       .not_wave_end
-  s_bitcmp1_b32        ttmp13, SQ_WAVE_TRAP_CTRL_WAVE_END_SHIFT
-  s_cbranch_scc0       .not_wave_end
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
+  s_bitcmp1_b32     ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_WAVE_END_SHIFT
+  s_cbranch_scc0    .not_wave_end
+  s_bitcmp1_b32     ttmp13, SQ_WAVE_TRAP_CTRL_WAVE_END_SHIFT
+  s_cbranch_scc0    .not_wave_end
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
 
 .not_wave_end:
-  s_bitcmp1_b32        ttmp13, SQ_WAVE_TRAP_CTRL_TRAP_AFTER_INST
-  s_cbranch_scc0       .not_trap_after_inst
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
+  s_bitcmp1_b32     ttmp13, SQ_WAVE_TRAP_CTRL_TRAP_AFTER_INST
+  s_cbranch_scc0    .not_trap_after_inst
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
 
 .not_trap_after_inst:
-  s_and_b32            ttmp2, ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_ADDR_WATCH_MASK
-  s_cbranch_scc0       .not_addr_watch
-  s_bitcmp1_b32        ttmp13, SQ_WAVE_TRAP_CTRL_ADDR_WATCH_SHIFT
-  s_cbranch_scc0       .not_addr_watch
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
+  s_and_b32         ttmp2, ttmp2, SQ_WAVE_EXCP_FLAG_PRIV_ADDR_WATCH_MASK
+  s_cbranch_scc0    .not_addr_watch
+  s_bitcmp1_b32     ttmp13, SQ_WAVE_TRAP_CTRL_ADDR_WATCH_SHIFT
+  s_cbranch_scc0    .not_addr_watch
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_TRAP_M0
 
 .not_addr_watch:
-  s_getreg_b32         ttmp2, hwreg(HW_REG_EXCP_FLAG_USER, SQ_WAVE_EXCP_FLAG_USER_MATH_EXCP_SHIFT, SQ_WAVE_EXCP_FLAG_USER_MATH_EXCP_SIZE)
-  s_and_b32            ttmp13, ttmp13, SQ_WAVE_TRAP_CTRL_MATH_EXCP_MASK
-  s_and_b32            ttmp2, ttmp2, ttmp13
-  s_cbranch_scc0       .not_math_exception
-  s_or_b32             ttmp3, ttmp3, EC_QUEUE_WAVE_MATH_ERROR_M0
+  s_getreg_b32      ttmp2, hwreg(HW_REG_EXCP_FLAG_USER, SQ_WAVE_EXCP_FLAG_USER_MATH_EXCP_SHIFT, SQ_WAVE_EXCP_FLAG_USER_MATH_EXCP_SIZE)
+  s_and_b32         ttmp13, ttmp13, SQ_WAVE_TRAP_CTRL_MATH_EXCP_MASK
+  s_and_b32         ttmp2, ttmp2, ttmp13
+  s_cbranch_scc0    .not_math_exception
+  s_or_b32          ttmp3, ttmp3, EC_QUEUE_WAVE_MATH_ERROR_M0
 
 .not_math_exception:
-  s_cmp_eq_u32         ttmp3, 0
+  s_cmp_eq_u32      ttmp3, 0
   // This was not a s_trap we are interested in or an exception, return to
   // the user code.
-  s_cbranch_scc1       .exit_trap
+  s_cbranch_scc1    .exit_trap
 
 .send_interrupt:
   // Fetch doorbell id for our queue.
-  s_sendmsg_rtn_b32    ttmp2, sendmsg(MSG_RTN_GET_DOORBELL)
-  s_wait_kmcnt         0
-  s_and_b32            ttmp2, ttmp2, DOORBELL_ID_MASK
-  s_or_b32             ttmp3, ttmp2, ttmp3
+  s_sendmsg_rtn_b32 ttmp2, sendmsg(MSG_RTN_GET_DOORBELL)
+  s_wait_kmcnt      0
+  s_and_b32         ttmp2, ttmp2, DOORBELL_ID_MASK
+  s_or_b32          ttmp3, ttmp2, ttmp3
 
   // Save trap id and halt status in ttmp6.
-  s_andn2_b32          ttmp6, ttmp6, (TTMP6_SAVED_TRAP_ID_MASK | TTMP6_SAVED_STATUS_HALT_MASK)
-  s_bfe_u32            ttmp2, ttmp1, SQ_WAVE_PC_HI_TRAP_ID_BFE
-  s_min_u32            ttmp2, ttmp2, 0xF
-  s_lshl_b32           ttmp2, ttmp2, TTMP6_SAVED_TRAP_ID_SHIFT
-  s_or_b32             ttmp6, ttmp6, ttmp2
-  s_bfe_u32            ttmp2, ttmp12, SQ_WAVE_STATE_PRIV_HALT_BFE
-  s_lshl_b32           ttmp2, ttmp2, TTMP6_SAVED_STATUS_HALT_SHIFT
-  s_or_b32             ttmp6, ttmp6, ttmp2
+  s_andn2_b32       ttmp6, ttmp6, (TTMP6_SAVED_TRAP_ID_MASK | TTMP6_SAVED_STATUS_HALT_MASK)
+  s_bfe_u32         ttmp2, ttmp1, SQ_WAVE_PC_HI_TRAP_ID_BFE
+  s_min_u32         ttmp2, ttmp2, 0xF
+  s_lshl_b32        ttmp2, ttmp2, TTMP6_SAVED_TRAP_ID_SHIFT
+  s_or_b32          ttmp6, ttmp6, ttmp2
+  s_bfe_u32         ttmp2, ttmp12, SQ_WAVE_STATE_PRIV_HALT_BFE
+  s_lshl_b32        ttmp2, ttmp2, TTMP6_SAVED_STATUS_HALT_SHIFT
+  s_or_b32          ttmp6, ttmp6, ttmp2
 
   // m0 = interrupt data = (exception_code << DOORBELL_ID_SIZE) | doorbell_id
-  s_mov_b32            ttmp2, m0
-  s_mov_b32            m0, ttmp3
-  s_nop                0x0 // Manually inserted wait states
-  s_sendmsg            sendmsg(MSG_INTERRUPT)
+  s_mov_b32         ttmp2, m0
+  s_mov_b32         m0, ttmp3
+  s_sendmsg         sendmsg(MSG_INTERRUPT)
   // Wait for the message to go out.
-  s_wait_kmcnt         0
-  s_mov_b32            m0, ttmp2
+  s_wait_kmcnt      0
+  s_mov_b32         m0, ttmp2
 
   // Parking the wave requires saving the original pc in the preserved ttmps.
   // Register layout before parking the wave:
@@ -234,44 +369,488 @@ trap_entry:
   // ttmp11: 1st_level_ttmp11[31:23] pc_hi[15:0] 1st_level_ttmp11[6:0]
   //
   // Save the PC
-  s_mov_b32            ttmp10, ttmp0
-  s_and_b32            ttmp1, ttmp1, SQ_WAVE_PC_HI_ADDRESS_MASK
-  s_lshl_b32           ttmp1, ttmp1, TTMP_PC_HI_SHIFT
-  s_andn2_b32          ttmp11, ttmp11, (SQ_WAVE_PC_HI_ADDRESS_MASK << TTMP_PC_HI_SHIFT)
-  s_or_b32             ttmp11, ttmp11, ttmp1
+  s_mov_b32         ttmp10, ttmp0
+  s_and_b32         ttmp1, ttmp1, SQ_WAVE_PC_HI_ADDRESS_MASK
+  s_lshl_b32        ttmp1, ttmp1, TTMP_PC_HI_SHIFT
+  s_andn2_b32       ttmp11, ttmp11, (SQ_WAVE_PC_HI_ADDRESS_MASK << TTMP_PC_HI_SHIFT)
+  s_or_b32          ttmp11, ttmp11, ttmp1
 
   // Park the wave
-  s_getpc_b64          [ttmp0, ttmp1]
-  s_add_u32            ttmp0, ttmp0, .parked - .
-  s_addc_u32           ttmp1, ttmp1, 0x0
+  s_getpc_b64       [ttmp0, ttmp1]
+  s_add_u32         ttmp0, ttmp0, .parked - .
+  s_addc_u32        ttmp1, ttmp1, 0x0
 
 .halt_wave:
   // Halt the wavefront upon restoring STATUS below.
-  s_bitset1_b32        ttmp6, TTMP6_WAVE_STOPPED_SHIFT
-  s_bitset1_b32        ttmp12, SQ_WAVE_STATE_PRIV_HALT_SHIFT
+  s_bitset1_b32     ttmp6, TTMP6_WAVE_STOPPED_SHIFT
+  s_bitset1_b32     ttmp12, SQ_WAVE_STATE_PRIV_HALT_SHIFT
 
   // Initialize TTMP registers
-  s_bitcmp1_b32        ttmp8, TTMP8_DEBUG_FLAG_SHIFT
-  s_cbranch_scc1       .ttmps_initialized
-  s_mov_b32            ttmp4, 0
-  s_mov_b32            ttmp5, 0
-  s_bitset1_b32        ttmp8, TTMP8_DEBUG_FLAG_SHIFT
+  s_bitcmp1_b32     ttmp8, TTMP8_DEBUG_FLAG_SHIFT
+  s_cbranch_scc1    .ttmps_initialized
+  s_mov_b32         ttmp4, 0
+  s_mov_b32         ttmp5, 0
+  s_bitset1_b32     ttmp8, TTMP8_DEBUG_FLAG_SHIFT
 .ttmps_initialized:
+  s_branch          .exit_trap
+
+.profile_trap_handlers:
+  // Register state at the start of profile_trap_handlers:
+  //
+  // ttmp0:  PC_LO[31:0] - Contains program counter low bits
+  // ttmp1:  PC_HI[15:0] - Contains program counter high bits
+  // ttmp2:  Contains HW_REG_EXCP_FLAG_PRIV
+  // ttmp3:  Initialized to 0, available for use
+  // ttmp4:  Available - Can be freely used
+  // ttmp5:  Available - Can be freely used
+  // ttmp6:  Initially contains flags  - trap ID and halt status - reused after saving
+  // ttmp7:  Contains WGID_Y in high 16 bits, WGID_Z in low 16 bits
+  // ttmp8:  Contains dispatch ID in bits [24:0] and debug flag
+  // ttmp9:  Contains WGID_X
+  // ttmp10: Available - Used next to save exec_lo
+  // ttmp11: Contains debug flags - Used next to save exec_hi
+  // ttmp12: Contains SQ_WAVE_STATE_PRIV
+  // ttmp13: Contains flag bits for sampling type - HT_FLAG_BIT or STOCH_FLAG_BIT
+  // ttmp[14:15]: Contains HT or ST buffer base address
+  //
+  // v[0:3] contain user shader data that must be preserved/restored
+  // exec: Contains user's execution mask
+  s_mov_b64         ttmp[10:11], exec                       // save exec to ttmp[10:11]
+  s_mov_b64         exec, 0x1                               // turn on lane 0 only
+
+  v_readlane_b32    ttmp2, v0, 0
+  v_readlane_b32    ttmp3, v1, 0                            // Save out lane 0’s first 2 VGPRs
+
+  // At this point, ttmp[4:5], ttmp6 and v[0:1] are free
+  // Atomically get current sample slot index and select buffer
+  // pcs_sampling_data_t.buf_write_val (uint64_t) stores:
+  //   Bit 63: current_buffer_id (0 or 1)
+  //   Bits 62-0: current_sample_index_in_buffer
+  // v0 = 1 (value to add to the low part of buf_write_val)
+  // v1 = 0 (value to add to the high part of buf_write_val, bit 63 is buffer selector)
+
+  v_mov_b32         v0, 1
+  v_mov_b32         v1, 0
+
+  global_atomic_add_u64 v[0:1], v1, v[0:1], ttmp[14:15], scope:SCOPE_SYS th:TH_ATOMIC_RETURN
+  s_wait_loadcnt    0                                       // Wait for atomic operation to complete and return value
+
+  // At this point, ttmp[4:5] and ttmp6 are free
+  // v[0:1] (lane 0) now holds the previous value of buf_write_val.
+  // This previous value gives the slot index for the current sample.
+
+  v_readlane_b32    ttmp6, v1, 0x0                          // previous buf_write_val[63:32]
+  s_lshr_b32        ttmp6, ttmp6, TTMP13_BUF_FULL_BIT       // ttmp6 = previous_buffer_id (0 or 1, from bit 63 of original uint64_t)
+                                                            // This ttmp6 is used to select which buffer's metadata (size, watermark, signal) to use.
+                                                            // It's also used to calculate the base address of the sample buffer.
+  s_bitset0_b32     ttmp13, TTMP13_BUF_FULL_BIT             // Clear our local buffer full flag for now
+
+  s_cmp_eq_u32      ttmp6, 0                                // store off buf_to_use
+  s_cbranch_scc1    .skip_bufbit_set                        // into bit31 of ttmp13
+  s_bitset1_b32     ttmp13, TTMP13_BUF_FULL_BIT
+
+.skip_bufbit_set:
+  // ttmp[2:3]=v[0:1]-backup, ttmp[4:5]=free, ttmp6=buf_to_use (also in ttmp13.b31)
+  // ttmp[10:11]=EXEC backup. ttmp[14:15]=tma
+  // v[0:1].lane0=local_entry, v[2:3]=original, EXEC=0x1
+
+  v_bfe_u32         v1, v1, 0, SAMPLE_INDEX_WIDTH           // v[0:1] = new local_entry
+                                                            // removes bit 31 from v1, returning v1 & 0x7FFFFFFF.
+
+  v_readlane_b32    ttmp5, v1, 0                            // ttmp5 = high 31 bits of sample index (if index > 2^32-1).
+  s_cmp_lg_u32      ttmp5, 0                                // Check if sample index is very large (overflowed 32 bits).
+
+  s_cbranch_scc1    .lost_sample                            // If ttmp5 > 0, index is too large, treat as lost sample.
+
+  s_load_b32           ttmp5, ttmp[14:15], SAMPLE_OFF_BUF_SIZE, scope:SCOPE_CU // ttmp5 = pcs_sampling_data_t.buf_size
+  v_readlane_b32    ttmp4, v0, 0                            // ttmp4 = sample_index_for_current_sample (from v0)
+  s_wait_kmcnt      0                                       // Wait for buf_size load.
+
+  s_cmp_ge_u32      ttmp4, ttmp5                            // if local_entry >= buf_size
+  s_cbranch_scc1    .lost_sample                            // If index >= buf_size, buffer is full, sample is lost.
+                                                            // This also sets TTMP13_BUF_FULL_BIT implicitly by branching.
+
+  // Register state before calculating the sample buffer address:
+  // ttmp2 = backup of original shader's v0
+  // ttmp3 = backup of original shader's v1
+  // ttmp4 = sample_index_for_current_sample (from v0)
+  // ttmp5 = buf_size
+  // ttmp6 = buffer_id (0 or 1)
+  // ttmp[10:11] = original shader's [exec_lo, exec_hi]
+  // ttmp[14:15] = base_address_of_pcs_sampling_data_t (TMA)
+  // ttmp13.b31 = buffer_id (0 or 1, same as ttmp6)
+  // v[0:1].lane0 = sample index value from atomic
+  // v[2:3] = original user shader's v[2:3] values
+  // exec = backup of user shader's v[0:1]
+  s_mov_b64         exec, ttmp[2:3]                         // stash into EXEC to free up ttmp
+
+  // Calculate the base address of the correct sample buffer (buffer0 or buffer1).
+  // The buffers are located after the pcs_sampling_data_t struct header.
+  // Address = (TMA + SAMPLE_OFF_BYTES_PER_SAMPLE) + (buffer_id * buf_size * 64)
+  s_mul_i32         ttmp2, ttmp5, ttmp6                     // low 32 bits
+  s_mul_hi_u32      ttmp3, ttmp5, ttmp6                     // high 32 bits
+
+  // Multiply by 64 bytes per sample slot (shift left by 6 bits)
+  // This converts from units of samples to units of bytes
+  s_lshl_b64        ttmp[2:3], ttmp[2:3], 6
+  s_add_u32         ttmp2, ttmp2, SAMPLE_OFF_BYTES_PER_SAMPLE
+  s_addc_u32        ttmp3, ttmp3, 0
+  s_add_u32         ttmp4, ttmp14, ttmp2                    // ttmp4 = TMA_base_lo + total_offset_lo. This is low part of &bufferX
+  s_addc_u32        ttmp5, ttmp15, ttmp3                    // ttmp5 = TMA_base_hi + total_offset_hi + carry. This is high part of &bufferX
+                                                            // ttmp[4:5] now correctly points to the base of the selected sample buffer array
+
+  s_bitcmp1_b32     ttmp13, TTMP13_HT_FLAG_BIT              // if ttmp13.b22==1, this is hosttrap
+  s_cbranch_scc1    .fill_sample_ht
+  s_bitcmp1_b32     ttmp13, TTMP13_STOCH_FLAG_BIT
+  s_cbranch_scc1    .fill_sample_stoch
+
+  s_mov_b64         ttmp[2:3], exec                         // Restore user v[0:1] backup to ttmp[2:3]
+  v_readlane_b32    ttmp4, v2, 0                            // Backup user v[2:3] to ttmp[4:5] for restore.
+  v_readlane_b32    ttmp5, v3, 0
+  s_branch          .restore_vector_before_exit_trap
+
+.fill_sample_ht:
+  // At this point, v[0:1] is local_entry (but v1 is 0)
+  // v[2:3] is original user-data
+  // ttmp[2:3] is free
+  // ttmp[4:5] holds &buffer
+  // ttmp6 holds buf_to_use
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // [ttmp14:15]=‘tma’, ttmp13.b31 = buf_to_use
+  // EXEC holds holds backup of original shader’s v[0:1]
+
+  v_readlane_b32    ttmp6, v0, 0                              // ttmp6=local_entry
+  s_mul_i32         ttmp2, ttmp6, SAMPLE_OFF_BYTES_PER_SAMPLE // into buffer for 64B objects
+  s_mul_hi_u32      ttmp3, ttmp6, SAMPLE_OFF_BYTES_PER_SAMPLE // ttmp[2:3] now holds the offset
+  s_add_u32         ttmp2, ttmp2, ttmp4
+  s_addc_u32        ttmp3, ttmp3, ttmp5                     // ttmp[2:3]=&bufferX[local_entry]
+  v_readlane_b32    ttmp4, v2, 0x0                          // ttmp[4:5] now holds backup of
+  v_readlane_b32    ttmp5, v3, 0x0                          // user-data from v[2:3]
+  v_writelane_b32   v0, ttmp2, 0x0
+  v_writelane_b32   v1, ttmp3, 0x0                          // v[0:1]=&buffer[local_entry]
+
+  s_sendmsg_rtn_b64 ttmp[2:3], sendmsg(MSG_RTN_GET_REALTIME)
+  s_wait_kmcnt      0                                       // Wait for timestamp
+
+  // v[0:1] = &buffer[local_entry]
+  // v[2:3] = free
+  // ttmp[2:3] holds the thing we want to store
+  // ttmp[4:5] holds backup of original shaders v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shaders [exec_lo,exec_hi]
+  // ttmp[14:15]=tma, ttmp13.b31 = buf_to_use
+  // EXEC holds backup of original shaders v[0:1]
+
+  v_writelane_b32   v2, ttmp2, 0                            // bring output data to v[2:3]
+  v_writelane_b32   v3, ttmp3, 0
+
+  s_mov_b64         ttmp[2:3], exec                         // vector stores need EXEC set
+  s_mov_b64         exec, 1                                 // so ttmp[2:3] holds it for now
+
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_TIMESTAMP, scope:SCOPE_SYS // store out timestamp
+
+  // v[0:1] = &buffer[local_entry]
+  // v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=‘tma’, ttmp13.b31 = buf_to_use
+  // EXEC is 0x1
+
+  s_and_b32         ttmp1, ttmp1, SQ_WAVE_PC_HI_ADDRESS_MASK // Clear out extra data from PC_HI
+  v_writelane_b32   v2, ttmp0, 0
+  v_writelane_b32   v3, ttmp1, 0
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_PC_HOST, scope:SCOPE_SYS  // store out PC
+
+  v_writelane_b32   v2, ttmp10, 0
+  v_writelane_b32   v3, ttmp11, 0
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_EXEC_LOHI, scope:SCOPE_SYS  // store out original EXEC
+
+  // Store Workgroup ID X and Y at offset SAMPLE_OFF_WGID_XY (0x10).
+  // ttmp9 = WGID_X (from first-level handler).
+  // ttmp7 contains WGID_Y in high 16 bits.
+  v_writelane_b32   v2, ttmp9, 0                            // wg_id_x
+  s_bfe_u32         ttmp6, ttmp7, (16<<16)                  // extract bits 15:0, wg_id_y
+  v_writelane_b32   v3, ttmp6, 0
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_WGID_XY, scope:SCOPE_SYS  // store wg_id_x and wg_id_y
+
+  // Store Workgroup ID Z and Wave ID at offset SAMPLE_OFF_WGID_Z_WAVE (0x18).
+  // ttmp7 contains WGID_Z in low 16 bits.
+  // ttmp11 contains Wave ID in low 6 bits (from EXEC_hi).
+  s_bfe_u32         ttmp6, ttmp7, (16|16<<16)               // extract bits 31:16, wg_id_z
+  v_writelane_b32   v2, ttmp6, 0
+  v_writelane_b32   v3, ttmp8, 0x0                          // wave_in_wg is bits 29:25
+  v_lshrrev_b32     v3, 25, v3                              // Shift wave_in_wg to 4:0
+  v_and_b32         v3, v3, WAVE_ID_MASK                    // put (ttmp8>>25)&0x1f into v3
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_WGID_Z_WAVE, scope:SCOPE_SYS  // store wg_id_z and wave_id
+
+  // v[0:1] = &buffer[local_entry]
+  // v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=‘tma’, ttmp13.b31 = buf_to_use
+  // EXEC is 0x1
+  // Get HW_ID1 & 2 with S_GETREG_B32 with size=32 (F8 in upper bits), offset=0, and:
+  // HW_ID1 = 23 (0x17), HW_ID2 = 24 (0x18)
+
+  STORE_HW_ID
+
+  // The following is still true as we get ready to jump to correlation ID check
+  // v[0:1] = &buffer[local_entry]
+  // v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15=‘tma’, ttmp13.b31 = buf_to_use
+  // EXEC is 0x1
+
+  STORE_CORRELATION_ID
+  // Ensure all stores have completed before returning and incrementing written_val
+  s_wait_storecnt   0
+
+  // Still true after returning back from correlation ID check
+  // v[0:1] = &buffer[local_entry], but we no longer need it
+  // v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=‘tma’, ttmp13.b31 = buf_to_use
+  // EXEC is 0x1
+  //
+  s_branch          .ret_from_fill_sample
+
+.fill_sample_stoch:
+  // v0 contains local_entry, v1 is free
+  // v[2:3] is original user-data
+  // ttmp[2:3] is free
+  // ttmp[4:5] holds &buffer
+  // ttmp6 holds buf_to_use
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // [ttmp14:15]=‘tma’, ttmp13.b31 = buf_to_use
+  // EXEC holds holds backup of original shader’s v[0:1]
+
+  v_readlane_b32    ttmp6, v0, 0x0                            // ttmp2=local_entry
+  s_mul_i32         ttmp2, ttmp6, SAMPLE_OFF_BYTES_PER_SAMPLE // into buffer for 64B objects
+  s_mul_hi_u32      ttmp3, ttmp6, SAMPLE_OFF_BYTES_PER_SAMPLE // ttmp[2:3] now holds the offset
+  s_add_u32         ttmp2, ttmp2, ttmp4
+  s_addc_u32        ttmp3, ttmp3, ttmp5                       // ttmp[2:3]=&bufferX[local_entry]
+  v_readlane_b32    ttmp4, v2, 0x0                            // ttmp[4:5] now holds backup of
+  v_readlane_b32    ttmp5, v3, 0x0                            // user-data from v[2:3]
+  v_writelane_b32   v0, ttmp2, 0x0
+  v_writelane_b32   v1, ttmp3, 0x0                            // v[0:1]=&buffer[local_entry]
+  s_sendmsg_rtn_b64 ttmp[2:3], sendmsg(MSG_RTN_GET_REALTIME)
+  s_wait_kmcnt      0                                         // Wait for timestamp
+
+  // v[0:1] = &buffer[local_entry]
+  // v[2:3] = free
+  // ttmp[2:3] holds the thing we want to store
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=‘tma’, ttmp13.b31 = buf_to_use
+  // EXEC holds backup of original shader’s v[0:1]
+
+  v_writelane_b32   v2, ttmp2, 0                            // bring output data to v[2:3]
+  v_writelane_b32   v3, ttmp3, 0
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_TIMESTAMP, scope:SCOPE_SYS  // store out timestamp
+
+  // v[0:1] = &buffer[local_entry]
+  // v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=‘tma’, ttmp13.b31 = buf_to_use
+  // EXEC is 0x1
+  v_writelane_b32   v2, ttmp10, 0
+  v_writelane_b32   v3, ttmp11, 0
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_EXEC_LOHI, scope:SCOPE_SYS  // store out original EXEC
+  v_writelane_b32   v2, ttmp9, 0                            // wg_id_x
+  s_bfe_u32         ttmp6, ttmp7, (0 | (16 << 16))          // extract bits 15:0, wg_id_y
+  v_writelane_b32   v3, ttmp6, 0
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_WGID_XY, scope:SCOPE_SYS  // store wg_id_x and wg_id_y
+  s_bfe_u32         ttmp6, ttmp7, (16|16<<16)               // extract bits 31:16, wg_id_z
+  v_writelane_b32   v2, ttmp6, 0                            // put wg_id_z in v2
+  v_writelane_b32   v3, ttmp8, 0x0                          // wave_in_wg is bits 29:25
+
+  v_lshrrev_b32     v3, 25, v3                              // Shift wave_in_wg to 4:0
+
+  v_and_b32         v3, v3, WAVE_ID_MASK                    // put (ttmp8>>25)&0x1f into v3
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_WGID_Z_WAVE, scope:SCOPE_SYS  // store wg_id_z and wave_id
+
+  STORE_HW_ID
+
+  //Read SNAPSHOT Data
+  s_getreg_b32      ttmp6, HW_REG_SQ_PERF_SNAPSHOT_DATA1
+  v_writelane_b32   v2, ttmp6, 0x0
+  s_getreg_b32      ttmp6, HW_REG_SQ_PERF_SNAPSHOT_DATA2
+  v_writelane_b32   v3, ttmp6, 0x0
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_SNAPSHOT_DATA + 4, scope:SCOPE_SYS  // store snapshot DATA1 and DATA2
+
+  s_getreg_b32      ttmp2, HW_REG_SQ_PERF_SNAPSHOT_DATA
+  v_writelane_b32   v2, ttmp2, 0
+  global_store_b32  v[0:1], v2, off, offset:SAMPLE_OFF_SNAPSHOT_DATA, scope:SCOPE_SYS  // store perf snapshot DATA
+
+  s_getreg_b32      ttmp6, HW_REG_SQ_PERF_SNAPSHOT_PC_LO
+  v_writelane_b32   v2, ttmp6, 0x0
+  s_getreg_b32      ttmp6, HW_REG_SQ_PERF_SNAPSHOT_PC_HI
+  v_writelane_b32   v3, ttmp6, 0x0
+  global_store_b64  v[0:1], v[2:3], off, offset:SAMPLE_OFF_PC_HOST, scope:SCOPE_SYS  // store PC_HI:PC_LO
+
+  // The following is still true as we get ready to jump to correlation ID check
+  // v[0:1] = &buffer[local_entry]
+  // v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=tma, ttmp13.b31 tells us buf_to_use
+  // EXEC is 0x1
+
+  STORE_CORRELATION_ID
+  // Ensure all stores have completed before returning and incrementing written_val
+  s_wait_storecnt   0
+
+.ret_from_fill_sample:
+  // v[0:1] = free
+  // v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=‘tma’, ttmp13.b31 tells us buf_to_use
+  // EXEC is 0x1
+
+  // Sample data has been written to the device buffer.
+  // Now, atomically increment the count of written samples for the current buffer.
+  // This is pcs_sampling_data_t.buf_written_val0 or buf_written_val1.
+  s_lshr_b32        ttmp6, ttmp13, 31                       // ttmp6 is buf_to_use
+  s_mulk_i32        ttmp6, 0x10                             // ttmp6=offset from
+                                                            // written_val0 to written_val_X
+  s_add_u32         ttmp14, ttmp14, ttmp6                   // now ttmp[14:15] points to base for
+  s_addc_u32        ttmp15, ttmp15, 0                       // buf_written_valX atomic operation
+
+  // Atomically increment the chosen buf_written_val.
+  // v0 = 0 (value to add - low part), v1 = 1 (value to add - high part, effectively just adding 1 to uint32_t)
+
+  v_mov_b32         v0, 0                                   // want to atomic increment
+  v_mov_b32         v1, 1                                   // buf_written_valX
+  global_atomic_add_u32 v0, v0, v1, ttmp[14:15], offset:SAMPLE_OFF_BUF_WRITTEN_VAL, scope:SCOPE_SYS th:TH_ATOMIC_RETURN
+  s_wait_loadcnt    0
+
+  // v0 = done, v1 = free, v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=buf_written_valX-0x10, EXEC=0x1
+  // Check Watermark and Signal Host
+
+  s_mov_b64         exec, ttmp[4:5]                         // stash user’s v[2:3] in EXEC
+  s_load_b32        ttmp5, ttmp[14:15], 0x14, scope:SCOPE_CU // load watermark into ttmp5
+  v_readlane_b32    ttmp4, v0, 0                            // put done into ttmp4
+  s_wait_kmcnt      0                                       // wait for watermark to load
+  s_cmp_lg_u32      ttmp4, ttmp5                            // if done != watermark, exit
+  s_add_u32         ttmp4, ttmp4, 1                         // ttmp4 is now current_sample_count (count_before_inc + 1)
+  s_cmp_lt_u32      ttmp4, ttmp5                            // if (current_sample_count < watermark), don't signal
+  s_mov_b64         ttmp[4:5], exec                         // restore user’s v[2:3]
+  s_mov_b64         exec, 1
+  s_cbranch_scc1    .restore_vector_before_exit_trap
+
+.send_signal:
+  // v[0:3] = free, ttmp[2:5] = backups of original v[0:3], ttmp6=free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=buf_written_valX-0x10, EXEC=old copy of original shader v[2:3]
+  // write done-signal and optional interrupt
+
+  // Watermark reached or exceeded. Signal the host.
+  // Load the hsa_signal_t handle for the current buffer.
+  // done_sig0 is at offset 0x18. done_sig1 is at 0x28.
+  // addr = ttmp[14:15] + 0x18 + (buffer_id * 0x10).
+  // ttmp0 still holds buffer_id * 0x10.
+
+  s_load_b64           ttmp[14:15], ttmp[14:15], SAMPLE_OFF_DONE_SIG0, scope:SCOPE_CU // load done_sig into ttmp[14:15]
+  s_mov_b64         exec, 1
+  s_wait_kmcnt      0
+
+  v_mov_b32         v0, 0
+  v_mov_b32         v1, 0                                   // value to store into v[0:1]
+  v_writelane_b32   v2, ttmp14, 0
+  v_writelane_b32   v3, ttmp15, 0                           // Put signal address into v[2:3]
+  global_store_b64  v[2:3], v[0:1], off, offset:SAMPLE_OFF_SIGNAL_VALUE, scope:SCOPE_SYS // zero out signal value
+
+  s_load_b32           ttmp6, ttmp[14:15], 0x18, scope:SCOPE_CU           // load event_id into ttmp6
+  s_load_b64           ttmp[14:15], ttmp[14:15], SAMPLE_OFF_EVENT_MAILBOX0, scope:SCOPE_CU     // load event mailbox ptr into 14:15
+  s_wait_kmcnt      0
+
+  s_cmp_eq_u64      ttmp[14:15], 0                          // null mailbox means no interrupt
+  s_cbranch_scc1    .restore_vector_before_exit_trap
+  s_cmp_eq_u32      ttmp6, 0                                // event_id zero means no interrupt
+  s_cbranch_scc1    .restore_vector_before_exit_trap
+  v_writelane_b32   v2, ttmp14, 0
+  v_writelane_b32   v3, ttmp15, 0                           // Put mailbox address into v[2:3]
+
+  s_wait_storecnt   0
+  v_writelane_b32   v0, ttmp6, 0x0                          // put event_id into v0
+  global_store_b32  v[2:3], v0, off, offset:0x0, scope:SCOPE_SYS // Send event ID to the mailbox
+  s_wait_storecnt   0
+  s_mov_b32         ttmp14, m0                              // save off m0
+  v_readlane_b32    ttmp15, v0, 0                           // Put ID into message payload
+  s_mov_b32         m0, ttmp15
+  s_sendmsg         sendmsg(MSG_INTERRUPT)                  // send interrupt message
+  s_wait_kmcnt      0
+  s_mov_b32         m0, ttmp14                              // restore m0
+
+  // v[0:1] = free
+  // v[2:3] = free
+  // ttmp[2:3] holds backup of original shader’s v[0:1]
+  // ttmp[4:5] holds backup of original shader’s v[2:3]
+  // ttmp6 = free
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=somewhere in tma region, EXEC is junk
+
+.restore_vector_before_exit_trap:
+  v_writelane_b32   v2, ttmp4, 0
+  v_writelane_b32   v3, ttmp5, 0
+
+.lost_sample:
+  // v0 contains local_entry, v1 is free
+  // v[2:3] is original user-data
+  // ttmp[2:3] [local_entry, buf_size]
+  // ttmp[4:5] = free
+  // ttmp6=buf_to_use (also in ttmp13.b31)
+  // ttmp[10:11] holds original shader’s [exec_lo,exec_hi]
+  // ttmp[14:15]=tma
+  // EXEC=0x1
+  // Restore vector registers before exiting
+
+  s_bitcmp1_b32     ttmp13, TTMP13_STOCH_FLAG_BIT           // Check if stochastic sampling
+  s_cbranch_scc0    .lost_sample_restore                    // If not, just restore and exit
+  s_getreg_b32      ttmp6, HW_REG_SQ_PERF_SNAPSHOT_PC_HI    // Read PC_HI to release lock
+
+.lost_sample_restore:
+  v_writelane_b32   v0, ttmp2, 0                            // restore v[0:1] to user data
+  v_writelane_b32   v1, ttmp3, 0
+  s_mov_b64         exec, ttmp[10:11]                       // restore exec mask
 
 .exit_trap:
   // Restore SQ_WAVE_STATUS.
-  s_and_b64            exec, exec, exec // Restore STATUS.EXECZ, not writable by s_setreg_b32
-  s_and_b64            vcc, vcc, vcc    // Restore STATUS.VCCZ, not writable by s_setreg_b32
-  s_setreg_b32         hwreg(HW_REG_STATE_PRIV, 0, SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT), ttmp12
-  s_lshr_b32           ttmp12, ttmp12, (SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT + 1)
-  s_setreg_b32         hwreg(HW_REG_STATE_PRIV, SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT + 1, 32 - SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT - 1), ttmp12
+  s_and_b64         exec, exec, exec                        // Restore STATUS.EXECZ, not writable by s_setreg_b32
+  s_and_b64         vcc, vcc, vcc                           // Restore STATUS.VCCZ, not writable by s_setreg_b32
+  s_setreg_b32      hwreg(HW_REG_STATE_PRIV, 0, SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT), ttmp12
+  s_lshr_b32        ttmp12, ttmp12, (SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT + 1)
+  s_setreg_b32      hwreg(HW_REG_STATE_PRIV, SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT + 1, 32 - SQ_WAVE_STATE_PRIV_BARRIER_COMPLETE_SHIFT - 1), ttmp12
 
-  // Return to original (possibly modified) PC.
-  s_rfe_b64            [ttmp0, ttmp1]
+  s_rfe_b64         [ttmp0, ttmp1]
 
 .parked:
-  s_trap               0x2
-  s_branch             .parked
+  s_trap            0x2
+  s_branch          .parked
 
 // Add s_code_end padding so instruction prefetch always has something to read.
 .rept (256 - ((. - trap_entry) % 64)) / 4