Add hsa_amd_memory_lock_to_pool.

Makes malloc memory accessible to GPUs so that the memory has the
capabilities of the pool it is locked to.
This admits fine grained locked memory and reserves API space for any future
special CPU pools.

Change-Id: If8c3dd8582a43f19d3d36b3763c1a688cc419ef0


[ROCm/ROCR-Runtime commit: a535e18cc1]

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

@@ -1015,6 +1015,14 @@ hsa_status_t HSA_API hsa_amd_memory_lock(void* host_ptr, size_t size,
                                      host_ptr, size, agents, num_agent, agent_ptr);
 }
 
+// Mirrors Amd Extension Apis
+hsa_status_t HSA_API hsa_amd_memory_lock_to_pool(void* host_ptr, size_t size, hsa_agent_t* agents,
+                                                 int num_agent, hsa_amd_memory_pool_t pool,
+                                                 uint32_t flags, void** agent_ptr) {
+  return amdExtTable->hsa_amd_memory_lock_to_pool_fn(host_ptr, size, agents, num_agent, pool, flags,
+                                                     agent_ptr);
+}
+
 // Mirrors Amd Extension Apis
 hsa_status_t HSA_API hsa_amd_memory_unlock(void* host_ptr) {
   return amdExtTable->hsa_amd_memory_unlock_fn(host_ptr);

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

@@ -82,8 +82,7 @@ class MemoryRegion : public core::MemoryRegion {
   /// @brief Free agent accessible memory (system / local memory).
   static void FreeKfdMemory(void* ptr, size_t size);
 
-  static bool RegisterMemory(void* ptr, size_t size, size_t num_nodes,
-                             const uint32_t* nodes);
+  static bool RegisterMemory(void* ptr, size_t size, const HsaMemFlags& MemFlags);
 
   static void DeregisterMemory(void* ptr);
 

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

@@ -173,6 +173,11 @@ hsa_status_t HSA_API hsa_amd_memory_lock(void* host_ptr, size_t size,
                                          hsa_agent_t* agents, int num_agent,
                                          void** agent_ptr);
 
+// Mirrors Amd Extension Apis
+hsa_status_t HSA_API hsa_amd_memory_lock_to_pool(void* host_ptr, size_t size, hsa_agent_t* agents,
+                                                 int num_agent, hsa_amd_memory_pool_t pool,
+                                                 uint32_t flags, void** agent_ptr);
+
 // Mirrors Amd Extension Apis
 hsa_status_t HSA_API hsa_amd_memory_unlock(void* host_ptr);
 

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

@@ -68,15 +68,11 @@ void MemoryRegion::FreeKfdMemory(void* ptr, size_t size) {
   assert(status == HSAKMT_STATUS_SUCCESS);
 }
 
-bool MemoryRegion::RegisterMemory(void* ptr, size_t size, size_t num_nodes,
-                                  const uint32_t* nodes) {
+bool MemoryRegion::RegisterMemory(void* ptr, size_t size, const HsaMemFlags& MemFlags) {
   assert(ptr != NULL);
   assert(size != 0);
-  assert(num_nodes != 0);
-  assert(nodes != NULL);
 
-  const HSAKMT_STATUS status = hsaKmtRegisterMemoryToNodes(
-      ptr, size, num_nodes, const_cast<uint32_t*>(nodes));
+  const HSAKMT_STATUS status = hsaKmtRegisterMemoryWithFlags(ptr, size, MemFlags);
   return (status == HSAKMT_STATUS_SUCCESS);
 }
 
@@ -120,8 +116,6 @@ MemoryRegion::MemoryRegion(bool fine_grain, bool full_profile, core::Agent* owne
         (mem_props_.HeapType == HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE) ? 0 : 1;
     mem_flag_.ui32.NonPaged = 1;
 
-    map_flag_.ui32.PageSize = HSA_PAGE_SIZE_4KB;
-
     virtual_size_ = kGpuVmSize;
   } else if (IsSystem()) {
     mem_flag_.ui32.PageSize = HSA_PAGE_SIZE_4KB;
@@ -129,9 +123,6 @@ MemoryRegion::MemoryRegion(bool fine_grain, bool full_profile, core::Agent* owne
     mem_flag_.ui32.HostAccess = 1;
     mem_flag_.ui32.CachePolicy = HSA_CACHING_CACHED;
 
-    map_flag_.ui32.HostAccess = 1;
-    map_flag_.ui32.PageSize = HSA_PAGE_SIZE_4KB;
-
     virtual_size_ =
         (full_profile) ? os::GetUserModeVirtualMemorySize() : kGpuVmSize;
   }
@@ -584,8 +575,7 @@ hsa_status_t MemoryRegion::Lock(uint32_t num_agents, const hsa_agent_t* agents,
   }
 
   // Call kernel driver to register and pin the memory.
-  if (RegisterMemory(host_ptr, size, whitelist_nodes.size(),
-                     &whitelist_nodes[0])) {
+  if (RegisterMemory(host_ptr, size, mem_flag_)) {
     uint64_t alternate_va = 0;
     if (MakeKfdMemoryResident(whitelist_nodes.size(), &whitelist_nodes[0],
                               host_ptr, size, &alternate_va, map_flag_)) {

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

@@ -386,6 +386,7 @@ void HsaApiTable::UpdateAmdExts() {
   amd_ext_api.hsa_amd_queue_set_priority_fn = AMD::hsa_amd_queue_set_priority;
   amd_ext_api.hsa_amd_memory_async_copy_rect_fn = AMD::hsa_amd_memory_async_copy_rect;
   amd_ext_api.hsa_amd_runtime_queue_create_register_fn = AMD::hsa_amd_runtime_queue_create_register;
+  amd_ext_api.hsa_amd_memory_lock_to_pool_fn = AMD::hsa_amd_memory_lock_to_pool;
 }
 
 class Init {

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

@@ -549,14 +549,48 @@ hsa_status_t hsa_amd_memory_lock(void* host_ptr, size_t size,
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   }
 
-  const amd::MemoryRegion* system_region =
-      reinterpret_cast<const amd::MemoryRegion*>(
-          core::Runtime::runtime_singleton_->system_regions_fine()[0]);
+  // Check for APU
+  if (core::Runtime::runtime_singleton_->system_regions_coarse().size() == 0) {
+    assert(core::Runtime::runtime_singleton_->system_regions_fine()[0]->full_profile() &&
+           "Missing coarse grain host memory on dGPU system.");
+    *agent_ptr = host_ptr;
+    return HSA_STATUS_SUCCESS;
+  }
+
+  const amd::MemoryRegion* system_region = static_cast<const amd::MemoryRegion*>(
+      core::Runtime::runtime_singleton_->system_regions_coarse()[0]);
 
   return system_region->Lock(num_agent, agents, host_ptr, size, agent_ptr);
   CATCH;
 }
 
+hsa_status_t hsa_amd_memory_lock_to_pool(void* host_ptr, size_t size, hsa_agent_t* agents,
+                                         int num_agent, hsa_amd_memory_pool_t pool, uint32_t flags,
+                                         void** agent_ptr) {
+  TRY;
+  IS_OPEN();
+  *agent_ptr = NULL;
+
+  if (size == 0 || host_ptr == NULL || agent_ptr == NULL || flags != 0) {
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  if ((agents != NULL && num_agent == 0) || (agents == NULL && num_agent != 0)) {
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  hsa_region_t region = {pool.handle};
+  const amd::MemoryRegion* mem_region = amd::MemoryRegion::Convert(region);
+  if (mem_region == nullptr) {
+    return (hsa_status_t)HSA_STATUS_ERROR_INVALID_MEMORY_POOL;
+  }
+  if (mem_region->owner()->device_type() != core::Agent::kAmdCpuDevice)
+    return (hsa_status_t)HSA_STATUS_ERROR_INVALID_MEMORY_POOL;
+
+  return mem_region->Lock(num_agent, agents, host_ptr, size, agent_ptr);
+  CATCH;
+}
+
 hsa_status_t hsa_amd_memory_unlock(void* host_ptr) {
   TRY;
   IS_OPEN();
@@ -615,7 +649,7 @@ hsa_status_t hsa_amd_memory_pool_allocate(hsa_amd_memory_pool_t memory_pool, siz
   TRY;
   IS_OPEN();
 
-  if (size == 0 || ptr == NULL) {
+  if (size == 0 || ptr == NULL || flags != 0) {
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   }