Unmap GPUs when allow_access removes them from system pools.

Change-Id: Ib9eb88622fded43ebd9eddbf78ad6771a5b91e77


[ROCm/ROCR-Runtime commit: e38ff18990]

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

@@ -86,13 +86,11 @@ class MemoryRegion : public core::MemoryRegion {
   static void DeregisterMemory(void* ptr);
 
   /// @brief Pin memory.
-  static bool MakeKfdMemoryResident(size_t num_node, const uint32_t* nodes,
-                                    void* ptr, size_t size,
-                                    uint64_t* alternate_va,
-                                    HsaMemMapFlags map_flag);
+  static bool MakeKfdMemoryResident(size_t num_node, const uint32_t* nodes, const void* ptr,
+                                    size_t size, uint64_t* alternate_va, HsaMemMapFlags map_flag);
 
   /// @brief Unpin memory.
-  static void MakeKfdMemoryUnresident(void* ptr);
+  static void MakeKfdMemoryUnresident(const void* ptr);
 
   MemoryRegion(bool fine_grain, bool full_profile, core::Agent* owner,
                const HsaMemoryProperties& mem_props);

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

@@ -315,12 +315,11 @@ class Runtime {
   static void AsyncEventsLoop(void*);
 
   struct AllocationRegion {
-    AllocationRegion() : region(NULL), assigned_agent_(NULL), size(0) {}
+    AllocationRegion() : region(NULL), size(0) {}
     AllocationRegion(const MemoryRegion* region_arg, size_t size_arg)
-        : region(region_arg), assigned_agent_(NULL), size(size_arg) {}
+        : region(region_arg), size(size_arg) {}
 
     const MemoryRegion* region;
-    const Agent* assigned_agent_;
     size_t size;
   };
 

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

@@ -81,23 +81,21 @@ bool MemoryRegion::RegisterMemory(void* ptr, size_t size, size_t num_nodes,
 
 void MemoryRegion::DeregisterMemory(void* ptr) { hsaKmtDeregisterMemory(ptr); }
 
-bool MemoryRegion::MakeKfdMemoryResident(size_t num_node, const uint32_t* nodes,
-                                         void* ptr, size_t size,
-                                         uint64_t* alternate_va,
+bool MemoryRegion::MakeKfdMemoryResident(size_t num_node, const uint32_t* nodes, const void* ptr,
+                                         size_t size, uint64_t* alternate_va,
                                          HsaMemMapFlags map_flag) {
   assert(num_node > 0);
   assert(nodes != NULL);
 
   *alternate_va = 0;
-  const HSAKMT_STATUS status =
-      hsaKmtMapMemoryToGPUNodes(ptr, size, alternate_va, map_flag, num_node,
-                                const_cast<uint32_t*>(nodes));
+  const HSAKMT_STATUS status = hsaKmtMapMemoryToGPUNodes(
+      const_cast<void*>(ptr), size, alternate_va, map_flag, num_node, const_cast<uint32_t*>(nodes));
 
   return (status == HSAKMT_STATUS_SUCCESS);
 }
 
-void MemoryRegion::MakeKfdMemoryUnresident(void* ptr) {
-  hsaKmtUnmapMemoryToGPU(ptr);
+void MemoryRegion::MakeKfdMemoryUnresident(const void* ptr) {
+  hsaKmtUnmapMemoryToGPU(const_cast<void*>(ptr));
 }
 
 MemoryRegion::MemoryRegion(bool fine_grain, bool full_profile,
@@ -454,15 +452,16 @@ hsa_status_t MemoryRegion::AllowAccess(uint32_t num_agents,
   if (whitelist_nodes.size() == 0 && IsSystem()) {
     assert(cpu_in_list);
     // This is a system region and only CPU agents in the whitelist.
-    // No need to call map.
+    // Remove old mappings.
+    amd::MemoryRegion::MakeKfdMemoryUnresident(ptr);
     return HSA_STATUS_SUCCESS;
   }
 
   // If this is a local memory region, the owning gpu always needs to be in
   // the whitelist.
   if (IsPublic() &&
-      std::find(whitelist_nodes.begin(), whitelist_nodes.end(),
-                owner()->node_id()) == whitelist_nodes.end()) {
+      std::find(whitelist_nodes.begin(), whitelist_nodes.end(), owner()->node_id()) ==
+          whitelist_nodes.end()) {
     whitelist_nodes.push_back(owner()->node_id());
     whitelist_gpus.insert(reinterpret_cast<GpuAgentInt*>(owner()));
   }
@@ -470,12 +469,11 @@ hsa_status_t MemoryRegion::AllowAccess(uint32_t num_agents,
   HsaMemMapFlags map_flag = map_flag_;
   map_flag.ui32.HostAccess |= (cpu_in_list) ? 1 : 0;
 
-  
   {
     ScopedAcquire<KernelMutex> lock(&core::Runtime::runtime_singleton_->memory_lock_);
     uint64_t alternate_va = 0;
     if (!amd::MemoryRegion::MakeKfdMemoryResident(
-      whitelist_nodes.size(), &whitelist_nodes[0], const_cast<void*>(ptr),
+      whitelist_nodes.size(), &whitelist_nodes[0], ptr,
       size, &alternate_va, map_flag)) {
         return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
     }