rocr/aie: Handle sideband

* Adding a sideband interface to get a handle from a virtual address. This is just temporary and should be removed when we move to the vmem API

* Removing unused queue structure

* Removing the queue from the arguments of hsa_amd_get_handle_from_vaddr

* Free XDNA BOs

---------

Co-authored-by: Yiannis Papadopoulos <[email protected]>

Author: 2 people

runtime/hsa-runtime/core/common/hsa_table_interface.cpp

@@ -930,6 +930,11 @@ hsa_amd_queue_hw_ctx_config(const hsa_queue_t *queue,
   return amdExtTable->hsa_amd_queue_hw_ctx_config_fn(queue, config_type, args);
 }
 
+// Mirrors AMD Extension APIs.
+hsa_status_t hsa_amd_get_handle_from_vaddr(void* ptr, uint32_t* handle) {
+  return amdExtTable->hsa_amd_get_handle_from_vaddr_fn(ptr, handle);
+}
+
 // Mirrors Amd Extension Apis
 hsa_status_t HSA_API hsa_amd_queue_cu_set_mask(const hsa_queue_t* queue,
                                                uint32_t num_cu_mask_count,

runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp

@@ -249,7 +249,12 @@ KfdDriver::ConfigHwCtx(core::Queue &queue,
                        hsa_amd_queue_hw_ctx_config_param_t config_type,
                        void *args) {
   // Only AIE queues support this for now.
-  return HSA_STATUS_ERROR_INVALID_QUEUE;
+  return HSA_STATUS_ERROR_INVALID_AGENT;
+}
+
+hsa_status_t KfdDriver::GetHandleFromVaddr(void* ptr, uint32_t* handle) {
+  // Only AIE queues support this for now.
+  return HSA_STATUS_ERROR_INVALID_AGENT;
 }
 
 void *KfdDriver::AllocateKfdMemory(const HsaMemFlags &flags, uint32_t node_id,

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

@@ -134,17 +134,17 @@ hsa_status_t
 XdnaDriver::AllocateMemory(const core::MemoryRegion &mem_region,
                            core::MemoryRegion::AllocateFlags alloc_flags,
                            void **mem, size_t size, uint32_t node_id) {
-  const MemoryRegion &m_region(static_cast<const MemoryRegion &>(mem_region));
+  const auto &region = static_cast<const MemoryRegion &>(mem_region);
   amdxdna_drm_create_bo create_bo_args{.size = size};
   amdxdna_drm_get_bo_info get_bo_info_args{0};
   drm_gem_close close_bo_args{0};
   void *mapped_mem(nullptr);
 
-  if (!m_region.IsSystem()) {
+  if (!region.IsSystem()) {
     return HSA_STATUS_ERROR_INVALID_REGION;
   }
 
-  if (m_region.kernarg()) {
+  if (region.kernarg()) {
     create_bo_args.type = AMDXDNA_BO_CMD;
   } else {
     create_bo_args.type = AMDXDNA_BO_DEV;
@@ -187,11 +187,30 @@ XdnaDriver::AllocateMemory(const core::MemoryRegion &mem_region,
   }
 
   vmem_handle_mappings.emplace(create_bo_args.handle, mapped_mem);
+  vmem_handle_mappings_reverse.emplace(mapped_mem, create_bo_args.handle);
 
   return HSA_STATUS_SUCCESS;
 }
 
-hsa_status_t XdnaDriver::FreeMemory(void *mem, size_t size) {
+hsa_status_t XdnaDriver::FreeMemory(void* ptr, size_t size) {
+  auto it = vmem_handle_mappings_reverse.find(ptr);
+  if (it == vmem_handle_mappings_reverse.end())
+    return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+  // TODO:ypapadop-amd: need to unmap memory, but we don't know if it's mapped or not as we don't have
+  // region information
+
+  auto handle = it->second;
+
+  drm_gem_close close_args = {};
+  close_args.handle = handle;
+  if (ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_args) < 0) {
+    return HSA_STATUS_ERROR;
+  }
+
+  vmem_handle_mappings.erase(handle);
+  vmem_handle_mappings_reverse.erase(it);
+
   return HSA_STATUS_SUCCESS;
 }
 
@@ -257,6 +276,14 @@ XdnaDriver::ConfigHwCtx(core::Queue &queue,
   }
 }
 
+hsa_status_t XdnaDriver::GetHandleFromVaddr(void* ptr, uint32_t* handle) {
+  auto it = vmem_handle_mappings_reverse.find(ptr);
+  if (it == vmem_handle_mappings_reverse.end())
+    return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+  *handle = it->second;
+  return HSA_STATUS_SUCCESS;
+}
+
 hsa_status_t XdnaDriver::QueryDriverVersion() {
   amdxdna_drm_query_aie_version aie_version{0, 0};
   amdxdna_drm_get_info args{DRM_AMDXDNA_QUERY_AIE_VERSION, sizeof(aie_version),

runtime/hsa-runtime/core/inc/amd_kfd_driver.h

@@ -82,6 +82,7 @@ class KfdDriver : public core::Driver {
   hsa_status_t ConfigHwCtx(core::Queue &queue,
                            hsa_amd_queue_hw_ctx_config_param_t config_type,
                            void *args) override;
+  hsa_status_t GetHandleFromVaddr(void* ptr, uint32_t* handle) override;
 
 private:
   /// @brief Allocate agent accessible memory (system / local memory).

runtime/hsa-runtime/core/inc/amd_xdna_driver.h

@@ -89,6 +89,8 @@ class XdnaDriver : public core::Driver {
                            hsa_amd_queue_hw_ctx_config_param_t config_type,
                            void *args) override;
 
+  hsa_status_t GetHandleFromVaddr(void* ptr, uint32_t* handle) override;
+
 private:
   hsa_status_t QueryDriverVersion();
   /// @brief Allocate device accesible heap space.
@@ -110,6 +112,9 @@ class XdnaDriver : public core::Driver {
   /// to manage some of this for now.
   std::unordered_map<uint32_t, void *> vmem_handle_mappings;
 
+  // TODO: Remove this once we move to the vmem API
+  std::unordered_map<void*, uint32_t> vmem_handle_mappings_reverse;
+
   /// @brief Virtual address range allocated for the device heap.
   ///
   /// Allocate a large enough space so we can carve out the device heap in

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

@@ -138,6 +138,8 @@ class Driver {
   ConfigHwCtx(Queue &queue, hsa_amd_queue_hw_ctx_config_param_t config_type,
               void *args) = 0;
 
+  virtual hsa_status_t GetHandleFromVaddr(void* ptr, uint32_t* handle) = 0;
+
   /// Unique identifier for supported kernel-mode drivers.
   const DriverType kernel_driver_type_;
 

runtime/hsa-runtime/core/inc/hsa_ext_amd_impl.h

@@ -113,6 +113,10 @@ hsa_status_t
 hsa_amd_queue_hw_ctx_config(const hsa_queue_t *queue,
                             hsa_amd_queue_hw_ctx_config_param_t config_type,
                             void *args);
+
+// Mirrors AMD Extension APIs.
+hsa_status_t hsa_amd_get_handle_from_vaddr(void* ptr, uint32_t* handle);
+
 // Mirrors Amd Extension Apis
 hsa_status_t hsa_amd_queue_cu_set_mask(const hsa_queue_t* queue,
                                                uint32_t num_cu_mask_count,

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

@@ -403,6 +403,8 @@ class Runtime {
                                                    const core::MemoryRegion** mem_region,
                                                    hsa_amd_memory_type_t* type);
 
+  hsa_status_t GetHandleFromVaddr(void* ptr, uint32_t* handle);
+
   hsa_status_t EnableLogging(uint8_t* flags, void* file);
 
   const std::vector<Agent*>& cpu_agents() { return cpu_agents_; }

runtime/hsa-runtime/core/runtime/hsa_api_trace.cpp

@@ -80,7 +80,7 @@ void HsaApiTable::Init() {
   // they can add preprocessor macros on the new functions
 
   constexpr size_t expected_core_api_table_size = 1016;
-  constexpr size_t expected_amd_ext_table_size = 592;
+  constexpr size_t expected_amd_ext_table_size = 600;//592;
   constexpr size_t expected_image_ext_table_size = 120;
   constexpr size_t expected_finalizer_ext_table_size = 64;
   constexpr size_t expected_tools_table_size = 64;
@@ -407,6 +407,7 @@ void HsaApiTable::UpdateAmdExts() {
   amd_ext_api.hsa_amd_async_function_fn = AMD::hsa_amd_async_function;
   amd_ext_api.hsa_amd_signal_wait_any_fn = AMD::hsa_amd_signal_wait_any;
   amd_ext_api.hsa_amd_queue_hw_ctx_config_fn = AMD::hsa_amd_queue_hw_ctx_config;
+  amd_ext_api.hsa_amd_get_handle_from_vaddr_fn = AMD::hsa_amd_get_handle_from_vaddr;
   amd_ext_api.hsa_amd_queue_cu_set_mask_fn = AMD::hsa_amd_queue_cu_set_mask;
   amd_ext_api.hsa_amd_queue_cu_get_mask_fn = AMD::hsa_amd_queue_cu_get_mask;
   amd_ext_api.hsa_amd_memory_pool_get_info_fn = AMD::hsa_amd_memory_pool_get_info;

runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp

@@ -634,6 +634,19 @@ hsa_amd_queue_hw_ctx_config(const hsa_queue_t *queue,
   CATCH;
 }
 
+
+hsa_status_t hsa_amd_get_handle_from_vaddr(void* ptr, uint32_t* handle) {
+  TRY;
+  IS_OPEN();
+
+  IS_BAD_PTR(ptr);
+  IS_BAD_PTR(handle);
+
+  return core::Runtime::runtime_singleton_->GetHandleFromVaddr(ptr, handle);
+
+  CATCH;
+}
+
 hsa_status_t hsa_amd_queue_cu_set_mask(const hsa_queue_t* queue, uint32_t num_cu_mask_count,
                                        const uint32_t* cu_mask) {
   TRY;

runtime/hsa-runtime/core/runtime/runtime.cpp

@@ -3605,6 +3605,17 @@ hsa_status_t Runtime::VMemoryGetAllocPropertiesFromHandle(hsa_amd_vmem_alloc_han
   return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t Runtime::GetHandleFromVaddr(void* ptr, uint32_t* handle) {
+  auto it = allocation_map_.find(ptr);
+  if (it == allocation_map_.end()) {
+    return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+  }
+
+  auto* agent = it->second.region->owner();
+  auto& driver = AgentDriver(agent->driver_type);
+  return driver.GetHandleFromVaddr(ptr, handle);
+}
+
 hsa_status_t Runtime::EnableLogging(uint8_t* flags, void* file) {
   memcpy(log_flags, flags, sizeof(log_flags));
 

runtime/hsa-runtime/hsacore.so.def

@@ -179,6 +179,7 @@ global:
 	hsa_amd_async_function;
 	hsa_amd_image_get_info_max_dim;
 	hsa_amd_queue_hw_ctx_config;
+	hsa_amd_get_handle_from_vaddr;
 	hsa_amd_queue_cu_set_mask;
 	hsa_amd_queue_cu_get_mask;
 	hsa_amd_memory_fill;

runtime/hsa-runtime/inc/hsa_api_trace.h

@@ -205,6 +205,7 @@ struct AmdExtTable {
   decltype(hsa_amd_async_function)* hsa_amd_async_function_fn;
   decltype(hsa_amd_signal_wait_any)* hsa_amd_signal_wait_any_fn;
   decltype(hsa_amd_queue_hw_ctx_config) *hsa_amd_queue_hw_ctx_config_fn;
+  decltype(hsa_amd_get_handle_from_vaddr)* hsa_amd_get_handle_from_vaddr_fn;
   decltype(hsa_amd_queue_cu_set_mask)* hsa_amd_queue_cu_set_mask_fn;
   decltype(hsa_amd_memory_pool_get_info)* hsa_amd_memory_pool_get_info_fn;
   decltype(hsa_amd_agent_iterate_memory_pools)* hsa_amd_agent_iterate_memory_pools_fn;

runtime/hsa-runtime/inc/hsa_ext_amd.h

@@ -1375,6 +1375,9 @@ hsa_status_t HSA_API hsa_amd_queue_hw_ctx_config(
     const hsa_queue_t *queue, hsa_amd_queue_hw_ctx_config_param_t config_type,
     void *args);
 
+
+hsa_status_t HSA_API hsa_amd_get_handle_from_vaddr(void* ptr, uint32_t* handle);
+
 /**
  * @brief Set a queue's CU affinity mask.
  *