AidanBeltonS · Oct 26, 2026 · Oct 26, 2026 · Oct 27, 2026 · Oct 27, 2026 · Oct 27, 2026
@@ -2758,60 +2758,76 @@ bool Device::virtualFree(void* addr) {
 }
 
 // ================================================================================================
-// Direct synchronous map/unmap path bypassing VirtualMapCommand. Used by the
-// non-async hipMemMap/hipMemUnmap entry points (wired up in subsequent
-// commits). No execution() lock and no barrier packet: the HIP layer is
-// responsible for draining access-device queues from the CPU side before
-// calling, and hsa_amd_vmem_{map,unmap} are synchronous w.r.t. the CPU.
-cl_int Device::virtualMap(void* va, size_t size, amd::Memory* phys) {
-  if (phys == nullptr) {
-    LogError("virtualMap: phys is nullptr");
-    return CL_INVALID_VALUE;
-  }
+// Single map implementation shared by the synchronous hipMemMap path (Device::virtualMap) and
+// the command path (VirtualGPU::submitVirtualMap). track_mapping is false for the mempool /
+// graph callers that manage the VA bookkeeping themselves.
+cl_int Device::MapPhysicalToVirtual(void* va, size_t size, amd::Memory* phys,
+                                    bool track_mapping) const {
+  assert(phys != nullptr && "MapPhysicalToVirtual requires a physical memory object!");
 
-  amd::Memory* vaddr_sub_obj = MapMemObjBookkeeping(phys, va, size);
-  if (vaddr_sub_obj == nullptr) {
-    LogError("virtualMap: MapMemObjBookkeeping failed");
-    return CL_INVALID_VALUE;
+  amd::Memory* vaddr_sub_obj = nullptr;
+  if (track_mapping) {
+    vaddr_sub_obj = MapMemObjBookkeeping(phys, va, size);
+    if (vaddr_sub_obj == nullptr) {
+      LogError("MapPhysicalToVirtual: MapMemObjBookkeeping failed");
+      return CL_INVALID_VALUE;
+    }
   }
 
-  hsa_amd_vmem_alloc_handle_t opaque_hsa_handle;
-  opaque_hsa_handle.handle = phys->getUserData().hsa_handle;
-  hsa_status_t hsa_status = Hsa::vmem_map(vaddr_sub_obj->getSvmPtr(), size,
-                                          vaddr_sub_obj->getOffset(), opaque_hsa_handle, 0);
+  hsa_amd_vmem_alloc_handle_t phys_handle{phys->getUserData().hsa_handle};
+  const size_t phys_offset = (vaddr_sub_obj != nullptr) ? vaddr_sub_obj->getOffset() : 0;
+  hsa_status_t hsa_status = Hsa::vmem_map(va, size, phys_offset, phys_handle, 0);
   if (hsa_status != HSA_STATUS_SUCCESS) {
-    LogPrintfError("virtualMap: hsa_amd_vmem_map failed with status: %d", hsa_status);
-    // Roll back the bookkeeping so the sub_obj/MemObjMap doesn't leak.
-    // FinalizeMapMemObjBookkeeping was never called, so MemObjMap doesn't
-    // contain va and the cross-links are not wired. Just tear down the
-    // sub-buffer view directly.
-    vaddr_sub_obj->getContext().devices()[0]->DestroyVirtualBuffer(vaddr_sub_obj);
-    vaddr_sub_obj->release();
+    LogPrintfError("MapPhysicalToVirtual: hsa_amd_vmem_map failed with status: %d", hsa_status);
+    if (vaddr_sub_obj != nullptr) {
+      // Roll back the bookkeeping so the sub_obj doesn't leak. FinalizeMapMemObjBookkeeping was
+      // never called, so MemObjMap doesn't contain va and the cross-links are not wired. Just
+      // tear down the sub-buffer view directly.
+      vaddr_sub_obj->getContext().devices()[0]->DestroyVirtualBuffer(vaddr_sub_obj);
+      vaddr_sub_obj->release();
+    }
     return CL_OUT_OF_HOST_MEMORY;
   }
 
-  constexpr bool kImportVmmForInterprocess = true;
-  FinalizeMapMemObjBookkeeping(vaddr_sub_obj, phys, va, kImportVmmForInterprocess);
-  // Resolve owningAgent now that the handle is reachable, refresh every device view.
-  for (auto* device : vaddr_sub_obj->getContext().devices()) {
-    if (auto* devMem = static_cast<Memory*>(vaddr_sub_obj->getDeviceMemory(*device))) {
-      devMem->refreshOwningAgentFromPointerInfo();
+  if (vaddr_sub_obj != nullptr) {
+    constexpr bool kImportVmmForInterprocess = true;
+    FinalizeMapMemObjBookkeeping(vaddr_sub_obj, phys, va, kImportVmmForInterprocess);
+
+    // Views created before the map inherited the VA reservation's owner. Views created later
+    // resolve it in computeAndSetOwningAgent().
+    constexpr bool kNoAlloc = false;
+    for (const auto* ctx_device : vaddr_sub_obj->getContext().devices()) {
+      if (auto* dev_mem =
+              static_cast<Memory*>(vaddr_sub_obj->getDeviceMemory(*ctx_device, kNoAlloc))) {
+        dev_mem->refreshOwningAgent();
+      }
     }
   }
   return CL_SUCCESS;
 }
 
 // ================================================================================================
-cl_int Device::virtualUnmap(void* va, size_t size) {
+// Single unmap implementation shared by Device::virtualUnmap and VirtualGPU::submitVirtualMap.
+cl_int Device::UnmapVirtual(void* va, size_t size, bool track_mapping) const {
+  if (!track_mapping) {
+    hsa_status_t hsa_status = Hsa::vmem_unmap(va, size);
+    if (hsa_status != HSA_STATUS_SUCCESS) {
+      LogPrintfError("UnmapVirtual: hsa_amd_vmem_unmap failed with status: %d", hsa_status);
+      return CL_INVALID_VALUE;
+    }
+    return CL_SUCCESS;
+  }
+
   amd::Memory* vaddr_sub_obj = amd::MemObjMap::FindMemObj(va);
+  assert(vaddr_sub_obj != nullptr && "Tracked VA has no sub_obj in MemObjMap!");
   if (vaddr_sub_obj == nullptr) {
-    LogPrintfError("virtualUnmap: no sub_obj for va: %p", va);
+    LogPrintfError("UnmapVirtual: no sub_obj for va: %p", va);
     return CL_INVALID_VALUE;
   }
 
   hsa_status_t hsa_status = Hsa::vmem_unmap(vaddr_sub_obj->getSvmPtr(), size);
   if (hsa_status != HSA_STATUS_SUCCESS) {
-    LogPrintfError("virtualUnmap: hsa_amd_vmem_unmap failed with status: %d", hsa_status);
+    LogPrintfError("UnmapVirtual: hsa_amd_vmem_unmap failed with status: %d", hsa_status);
     return CL_INVALID_VALUE;
   }
 
@@ -2821,6 +2837,28 @@ cl_int Device::virtualUnmap(void* va, size_t size) {
   return CL_SUCCESS;
 }
 
+// ================================================================================================
+// Direct synchronous map/unmap path bypassing VirtualMapCommand. Used by the
+// non-async hipMemMap/hipMemUnmap entry points. No execution() lock and no
+// barrier packet: the HIP layer is responsible for draining access-device
+// queues from the CPU side before calling, and hsa_amd_vmem_{map,unmap} are
+// synchronous w.r.t. the CPU.
+cl_int Device::virtualMap(void* va, size_t size, amd::Memory* phys) {
+  if (phys == nullptr) {
+    LogError("virtualMap: phys is nullptr");
+    return CL_INVALID_VALUE;
+  }
+
+  constexpr bool kTrackMapping = true;
+  return MapPhysicalToVirtual(va, size, phys, kTrackMapping);
+}
+
+// ================================================================================================
+cl_int Device::virtualUnmap(void* va, size_t size) {
+  constexpr bool kTrackMapping = true;
+  return UnmapVirtual(va, size, kTrackMapping);
+}
+
 // ================================================================================================
 bool Device::SetMemAccess(void* va_addr, size_t va_size, VmmAccess access_flags,
                           VmmLocationType access_location, int numaNode) {

@@ -517,6 +517,16 @@ class Device : public NullDevice {
   virtual cl_int virtualMap(void* va, size_t size, amd::Memory* phys) override;
   virtual cl_int virtualUnmap(void* va, size_t size) override;
 
+  //! Maps phys at va. With track_mapping, creates the VA sub-object, wires MemObjMap, resolves
+  //! the owning agent on its existing device views, and rolls the sub-object back if the map
+  //! fails. phys must not be nullptr.
+  [[nodiscard]] cl_int MapPhysicalToVirtual(void* va, size_t size, amd::Memory* phys,
+                                            bool track_mapping) const;
+
+  //! Unmaps va. With track_mapping, also tears down the VA sub-object bookkeeping, and va must
+  //! name a range this runtime previously mapped and bookkept.
+  [[nodiscard]] cl_int UnmapVirtual(void* va, size_t size, bool track_mapping) const;
+
   virtual bool SetMemAccess(void* va_addr, size_t va_size, VmmAccess access_flags,
                             VmmLocationType = VmmLocationType::kDevice,
                             int numaNode = -1) override;

@@ -1197,51 +1197,50 @@ bool Buffer::create(bool alloc_local, amd::Memory::ErrorCode* errorCode) {
   return (success = (deviceMemory_ != nullptr));
 }
 
-// Resolve an imported VMM buffer's owning agent from its physical memory handle, since
-// pointer_info reports a null owner for imports. Returns false for non-VMM memory.
-static bool vmmOwningAgent(const Memory* self, hsa_agent_t* agent) {
-  amd::Memory* owner = self->owner();
-  if (owner == nullptr || !owner->vmmImported() || owner->getUserData().phys_mem_obj == nullptr) {
-    return false;
-  }
-  hsa_amd_vmem_alloc_handle_t handle{owner->getUserData().phys_mem_obj->getUserData().hsa_handle};
-  hsa_amd_vmem_handle_info_t info = {};
-  info.size = sizeof(info);
-  if (Hsa::vmem_get_vmem_info(handle, &info) != HSA_STATUS_SUCCESS || info.agent.handle == 0) {
-    return false;
-  }
-  *agent = info.agent;
-  return true;
-}
-
-// For VMM-imported memory, resolve the owner from the allocation handle (vmem_get_vmem_info):
-// pointer_info is not sufficient there, as it reports a null owner for imports. All other memory
-// uses pointer_info.
-void Memory::refreshOwningAgentFromPointerInfo() {
-  if (deviceMemory_ == nullptr) {
-    return;
+// Resolve the agent that owns the backing pages, which for imported memory is a peer device.
+// A mapped VMM range asks its physical handle, since pointer_info reports no owner for a handle
+// imported from another process. Other memory (IPC-opened) asks pointer_info, keyed off a valid
+// agentOwner because ROCr does not always tag IPC memory with HSA_EXT_POINTER_TYPE_IPC. A fabric
+// import reports no owner either way and falls back to this device.
+void Memory::refreshOwningAgent() {
+  hsa_agent_t agent = {};
+  amd::Memory* phys = (owner() != nullptr) ? owner()->getUserData().phys_mem_obj : nullptr;
+
+  if (phys != nullptr) {
+    hsa_amd_vmem_alloc_handle_t handle{phys->getUserData().hsa_handle};
+    hsa_amd_vmem_handle_info_t handle_info = {};
+    handle_info.size = sizeof(handle_info);
+    if (Hsa::vmem_get_vmem_info(handle, &handle_info) == HSA_STATUS_SUCCESS) {
+      agent = handle_info.agent;
+    }
+  } else {
+    assert(deviceMemory_ != nullptr && "refreshOwningAgent called before memory was backed!");
+    hsa_amd_pointer_info_t ptr_info = {};
+    ptr_info.size = sizeof(ptr_info);
+    if (hsa_amd_pointer_info(reinterpret_cast<address>(deviceMemory_), &ptr_info, nullptr, nullptr,
+                             nullptr) == HSA_STATUS_SUCCESS) {
+      agent = ptr_info.agentOwner;
+    }
   }
 
-  hsa_agent_t agent;
-  if (!vmmOwningAgent(this, &agent)) {
-    hsa_amd_pointer_info_t info = {};
-    info.size = sizeof(info);
-    hsa_status_t err = hsa_amd_pointer_info(reinterpret_cast<address>(deviceMemory_), &info,
-                                            nullptr, nullptr, nullptr);
-
-    // info.agentOwner is the agent that actually owns the backing pages (a peer device for
-    // IPC memory). Note ROCr does not always tag IPC-opened memory with
-    // HSA_EXT_POINTER_TYPE_IPC, so key off a valid agentOwner rather than the pointer type.
-    agent = (err == HSA_STATUS_SUCCESS && info.agentOwner.handle != 0) ? info.agentOwner
-                                                                       : dev().getBackendDevice();
-  }
-  setOwningAgent(agent);
+  // Always set an agent: computeAndSetOwningAgent() calls this with nothing cached, and
+  // getOwningAgent() feeds the copy agents in rocblit.
+  setOwningAgent(agent.handle != 0 ? agent : dev().getBackendDevice());
 }
 
 // Helper function to compute and cache the owning agent
 void Buffer::computeAndSetOwningAgent() {
   hsa_agent_t agent;
 
+  // Memory backed by another device's pages is owned by that device. A mapped VMM range belongs
+  // to its physical handle's agent rather than the VA reservation it is carved from, so this
+  // must precede the sub-buffer inheritance below.
+  if (owner() != nullptr &&
+      (owner()->getUserData().phys_mem_obj != nullptr || owner()->ipcShared())) {
+    refreshOwningAgent();
+    return;
+  }
+
   // Sub-buffers must inherit agent from parent, not recompute it
   if (owner() != nullptr && owner()->parent() != nullptr) {
     const Memory* parentMemory = static_cast<const Memory*>(
@@ -1255,11 +1254,6 @@ void Buffer::computeAndSetOwningAgent() {
     LogWarning("Sub-buffer parent not available for agent inheritance");
   }
 
-  if (owner() != nullptr && (owner()->ipcShared() || owner()->vmmImported())) {
-    refreshOwningAgentFromPointerInfo();
-    return;
-  }
-
   if (kind_ == MEMORY_KIND_ARENA || kind_ == MEMORY_KIND_HOST) {
     // Arena and host memory use CPU agent
     agent = dev().getCpuAgent();

@@ -120,11 +120,12 @@ class Memory : public device::Memory {
     return owningAgentIndex_.load(std::memory_order_acquire);
   }
 
-  //! Recompute and cache the owning agent by querying pointer_info on the (already-backed)
-  //! device memory. Needed for VMM-mapped / VMM-imported (interprocess) memory: the true
-  //! owner is only known after hsa_amd_vmem_map backs the virtual address, which happens
-  //! after create() runs, so the agent cannot be resolved at create() time.
-  void refreshOwningAgentFromPointerInfo();
+  //! Recompute and cache the owning agent: from the physical handle for a mapped VMM range,
+  //! whose backing pages may live on a peer device, and from pointer_info for IPC memory.
+  //! Called after hsa_amd_vmem_map backs the virtual address, since the owner of an imported
+  //! range is not knowable at create() time, and from computeAndSetOwningAgent() for views
+  //! created after the map.
+  void refreshOwningAgent();
 
   //! Validates allocated memory for possible workarounds
   virtual bool ValidateMemory() { return true; }

@@ -4495,63 +4495,20 @@ void VirtualGPU::submitVirtualMap(amd::VirtualMapCommand& vcmd) {
     return;
   }
 
-  // Get the amd::Memory object for the physical address
-  amd::Memory* phys_mem_obj = vcmd.memory();
-  hsa_status_t hsa_status = HSA_STATUS_SUCCESS;
-
-  // If Physical address is not set, then it is map command. If set, it is unmap command.
-  if (phys_mem_obj != nullptr) {
-    amd::Memory* vaddr_sub_obj = nullptr;
-    if (vcmd.trackMapping()) {
-      vaddr_sub_obj =
-          dev().MapMemObjBookkeeping(phys_mem_obj, const_cast<void*>(vcmd.ptr()), vcmd.size());
-      if (vaddr_sub_obj == nullptr) {
-        LogError("HSA Command: MapMemObjBookkeeping failed!");
-        profilingEnd();
-        return;
-      }
-    }
-    // Map the physical to virtual address the hsa api
-    hsa_amd_vmem_alloc_handle_t opaque_hsa_handle;
-    opaque_hsa_handle.handle = phys_mem_obj->getUserData().hsa_handle;
-    size_t phys_offset = (vaddr_sub_obj != nullptr) ? vaddr_sub_obj->getOffset() : 0;
-    if ((hsa_status = Hsa::vmem_map(const_cast<void*>(vcmd.ptr()), vcmd.size(), phys_offset,
-                                    opaque_hsa_handle, 0)) != HSA_STATUS_SUCCESS) {
-      LogError("HSA Command: hsa_amd_vmem_map failed!");
-    } else if (vaddr_sub_obj != nullptr) {
-      constexpr bool kImportVmmForInterprocess = true;
-      dev().FinalizeMapMemObjBookkeeping(vaddr_sub_obj, phys_mem_obj, const_cast<void*>(vcmd.ptr()),
-                                         kImportVmmForInterprocess);
-      // The VA is now backed, so the true owning agent (a peer device for imported memory)
-      // can finally be resolved. create() ran before the mapping existed, so cache it now.
-      if (auto* devMem = static_cast<Memory*>(vaddr_sub_obj->getDeviceMemory(dev()))) {
-        devMem->refreshOwningAgentFromPointerInfo();
-      }
-    }
+  void* va = const_cast<void*>(vcmd.ptr());
+  cl_int status = CL_SUCCESS;
+  // If Physical address is set, then it is map command. If not set, it is unmap command.
+  if (vcmd.memory() != nullptr) {
+    status = dev().MapPhysicalToVirtual(va, vcmd.size(), vcmd.memory(), vcmd.trackMapping());
   } else {
     dispatchBarrierPacket(kBarrierPacketHeader, false);
     Barriers().WaitCurrent();
 
-    if (!vcmd.trackMapping()) {
-      if ((hsa_status = Hsa::vmem_unmap(const_cast<void*>(vcmd.ptr()), vcmd.size())) !=
-          HSA_STATUS_SUCCESS) {
-        LogError("HSA Command: hsa_amd_vmem_unmap failed");
-      }
-    } else {
-      amd::Memory* vaddr_sub_obj = amd::MemObjMap::FindMemObj(vcmd.ptr());
-      assert(vaddr_sub_obj != nullptr);
-
-      // Unmap the object, since the physical addr is set.
-      if ((hsa_status = Hsa::vmem_unmap(vaddr_sub_obj->getSvmPtr(), vcmd.size())) ==
-          HSA_STATUS_SUCCESS) {
-        constexpr bool kDestroyVirtualBuffer = true;
-        constexpr bool kReleaseSubObj = true;
-        dev().UnmapMemObjBookkeeping(vaddr_sub_obj, const_cast<void*>(vcmd.ptr()),
-                                     kDestroyVirtualBuffer, kReleaseSubObj);
-      } else {
-        LogError("HSA Command: hsa_amd_vmem_unmap failed");
-      }
-    }
+    status = dev().UnmapVirtual(va, vcmd.size(), vcmd.trackMapping());
+  }
+
+  if (status != CL_SUCCESS) {
+    vcmd.setStatus(status);
   }
 
   // Since this is a memory operation, the HW event set for barrier packet