SWDEV-558848 - Update thunk interface signature for vmm enablement (#2259)

Co-authored-by: Rahul Manocha <rmanocha@amd.com>
2025-12-11 08:43:28 -08:00
@@ -553,6 +553,7 @@ hsaKmtExportDMABufHandle(
 HSAKMT_STATUS
 HSAKMTAPI
 hsaKmtGetMemoryHandle(
+    void* va,                     // IN
    void* MemoryAddress,          // IN
    HSAuint64 SizeInBytes,        // IN
    uint64_t* SharedMemoryHandle  // OUT
@@ -931,8 +931,8 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtGetAMDGPUDeviceHandle(HSAuint32 NodeId,
 }

 HSAKMT_STATUS HSAKMTAPI
-hsaKmtGetMemoryHandle(void *MemoryAddress, HSAuint64 SizeInBytes,
-                      uint64_t *SharedMemoryHandle) {
+hsaKmtGetMemoryHandle(void* va, void* MemoryAddress, HSAuint64 SizeInBytes,
+                      uint64_t* SharedMemoryHandle) {
 	CHECK_KFD_OPEN();

 	return HSAKMT_STATUS_NOT_SUPPORTED;
@@ -415,9 +415,10 @@ hsa_status_t KfdDriver::AllocQueueGWS(HSA_QUEUEID queue_id, uint32_t num_gws,
  return HSA_STATUS_SUCCESS;
 }

-hsa_status_t KfdDriver::GetShareableHandle(void* mem, size_t size, core::ShareableHandle* handle) {
+hsa_status_t KfdDriver::GetShareableHandle(void* va, void* mem, size_t size,
+                                           core::ShareableHandle* handle) {
  uint64_t mem_handle;
-  HSAKMT_STATUS status = HSAKMT_CALL(hsaKmtGetMemoryHandle(mem, size, &mem_handle));
+  HSAKMT_STATUS status = HSAKMT_CALL(hsaKmtGetMemoryHandle(va, mem, size, &mem_handle));
  if (status != HSAKMT_STATUS_SUCCESS) {
    return HSA_STATUS_ERROR;
  }
@@ -1029,7 +1029,7 @@ hsa_status_t XdnaDriver::MakeMemoryResident(const void* mem, size_t size, uint64

 hsa_status_t XdnaDriver::MakeMemoryUnresident(const void* mem) const { return HSA_STATUS_ERROR; }

-hsa_status_t XdnaDriver::GetShareableHandle(void* mem, size_t size, core::ShareableHandle* handle) {
+hsa_status_t XdnaDriver::GetShareableHandle(void* va, void* mem, size_t size, core::ShareableHandle* handle) {
  return HSA_STATUS_ERROR;
 }

@@ -116,7 +116,7 @@ public:
  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 GetShareableHandle(void* mem, size_t size, core::ShareableHandle* handle) override;
+  hsa_status_t GetShareableHandle(void* va, void* mem, size_t size, core::ShareableHandle* handle) override;
  hsa_status_t SPMAcquire(uint32_t preferred_node_id) const override;
  hsa_status_t SPMRelease(uint32_t preferred_node_id) const override;
  hsa_status_t SPMSetDestBuffer(uint32_t preferred_node_id, uint32_t size_bytes, uint32_t* timeout,
@@ -225,7 +225,7 @@ public:
                   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 GetShareableHandle(void* mem, size_t size, core::ShareableHandle* handle) override;
+  hsa_status_t GetShareableHandle(void* va, void* mem, size_t size, core::ShareableHandle* handle) override;
  hsa_status_t ReleaseShareableHandle(core::ShareableHandle &handle) override;

  /// @brief Submits @p num_pkts packets in a command chain.
@@ -233,11 +233,11 @@ public:
                             size_t offset, size_t size) = 0;

  /// @brief Get Shareable Memory Handle for physical memory
-  ///
+  /// @param[in] va virtual address
  /// @param[in] mem  physical memory handle
  /// @param[in] size size of memory allocated in bytes
  /// @param[out] handle handle of the memory object
-  virtual hsa_status_t GetShareableHandle(void* mem, size_t size,
+  virtual hsa_status_t GetShareableHandle(void* va, void* mem, size_t size,
                                          core::ShareableHandle* handle) = 0;

  /// @brief Releases the object associated with the handle.
@@ -332,9 +332,10 @@ class ThunkLoader {
                                      HsaAisFlags flags, \
                                      HSAuint64 *SizeCopiedInBytes, \
                                      HSAint32 *status);
-    typedef HSAKMT_STATUS (HSAKMT_DEF(hsaKmtGetMemoryHandle))(void *MemoryAddress, \
+    typedef HSAKMT_STATUS (HSAKMT_DEF(hsaKmtGetMemoryHandle))(void* va, \
+                                      void* MemoryAddress, \
                                      HSAuint64 SizeInBytes, \
-                                      uint64_t  *SharedMemoryHandle);
+                                      uint64_t* SharedMemoryHandle);
    /* drm API */
    typedef int (DRM_DEF(amdgpu_device_initialize))(int fd, \
                                      uint32_t *major_version, \
@@ -3681,7 +3681,7 @@ hsa_status_t Runtime::VMemoryHandleMap(void* va, size_t size, size_t in_offset,
    ret = GetAmdgpuDeviceArgs(agent, shareable_handle, &drm_fd, &drm_cpu_addr);
    if (ret) return HSA_STATUS_ERROR;
  } else {
-    hsa_status_t status = agent_driver.GetShareableHandle(memoryHandleIt->first, size, &shareable_handle);
+    hsa_status_t status = agent_driver.GetShareableHandle(va, memoryHandleIt->first, size, &shareable_handle);
    if (status != HSA_STATUS_SUCCESS) {
      return status;
    }
@@ -459,10 +459,13 @@ uint64_t HostTotalPhysicalMemory() {
  return totalPhys;
 }

-bool UnmapMemory(void* addr, size_t size) { return VirtualFree(addr, size, MEM_RELEASE) != 0; }
+bool UnmapMemory(void* addr, size_t size) { return UncommitMemory(addr, size); }

 bool MapMemory(void* addr, size_t size, MemProt perms, int fd [[maybe_unused]],
               uint64_t cpu_addr [[maybe_unused]]) {
+  if (perms == MEM_PROT_NONE) {
+    return true;
+  }
  DWORD OldProtect;
  return VirtualProtect(addr, size, memProtToOsProt(perms), &OldProtect) != 0;
 }