Add preferred agent info to pointer info struct.

Lookup blit agent via pointer info in memory_fill.

Change-Id: I02feaf68bb9726858e8cb0ede6bc5f2b3707f5af

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

@@ -453,8 +453,42 @@ hsa_status_t Runtime::CopyMemory(void* dst, core::Agent& dst_agent,
 }
 
 hsa_status_t Runtime::FillMemory(void* ptr, uint32_t value, size_t count) {
-  assert(blit_agent_ != NULL);
-  return blit_agent_->DmaFill(ptr, value, count);
+  // Choose blit agent from pointer info
+  hsa_amd_pointer_info_t info;
+  uint32_t agent_count;
+  hsa_agent_t* accessible = nullptr;
+  info.size = sizeof(info);
+  MAKE_SCOPE_GUARD([&]() { free(accessible); });
+  hsa_status_t err = PtrInfo(ptr, &info, malloc, &agent_count, &accessible);
+  if (err != HSA_STATUS_SUCCESS) return err;
+
+  ptrdiff_t endPtr = (ptrdiff_t)ptr + count * sizeof(uint32_t);
+
+  // Check for GPU fill
+  // Selects GPU fill for SVM and Locked allocations if a GPU address is given and is mapped.
+  if (info.agentBaseAddress <= ptr &&
+      endPtr <= (ptrdiff_t)info.agentBaseAddress + info.sizeInBytes) {
+    core::Agent* blit_agent = core::Agent::Convert(info.agentOwner);
+    if (blit_agent->device_type() != core::Agent::DeviceType::kAmdGpuDevice) {
+      blit_agent = nullptr;
+      for (int i = 0; i < agent_count; i++) {
+        if (core::Agent::Convert(accessible[i])->device_type() ==
+            core::Agent::DeviceType::kAmdGpuDevice) {
+          blit_agent = core::Agent::Convert(accessible[i]);
+          break;
+        }
+      }
+    }
+    if (blit_agent) return blit_agent->DmaFill(ptr, value, count);
+  }
+
+  // Host and unmapped SVM addresses copy via host.
+  if (info.hostBaseAddress <= ptr && endPtr <= (ptrdiff_t)info.hostBaseAddress + info.sizeInBytes) {
+    memset(ptr, value, count * sizeof(uint32_t));
+    return HSA_STATUS_SUCCESS;
+  }
+
+  return HSA_STATUS_ERROR_INVALID_ALLOCATION;
 }
 
 hsa_status_t Runtime::AllowAccess(uint32_t num_agents,
@@ -646,6 +680,8 @@ hsa_status_t Runtime::PtrInfo(void* ptr, hsa_amd_pointer_info_t* info, void* (*a
   HsaPointerInfo thunkInfo;
   uint32_t* mappedNodes;
 
+  hsa_amd_pointer_info_t retInfo;
+
   // check output struct is at least as large as the first info revision.
   if (info->size < sizeof(struct hsa_amd_pointer_info_v1_s)) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
 
@@ -674,12 +710,15 @@ hsa_status_t Runtime::PtrInfo(void* ptr, hsa_amd_pointer_info_t* info, void* (*a
   static_assert((int)HSA_POINTER_REGISTERED_GRAPHICS == (int)HSA_EXT_POINTER_TYPE_GRAPHICS,
                 "Thunk pointer info mismatch");
 
-  info->size = Min(info->size, sizeof(struct hsa_amd_pointer_info_v1_s));
-  info->type = (hsa_amd_pointer_type_t)thunkInfo.Type;
-  info->agentBaseAddress = (void*)thunkInfo.GPUAddress;
-  info->hostBaseAddress = thunkInfo.CPUAddress;
-  info->sizeInBytes = thunkInfo.SizeInBytes;
-  info->userData = thunkInfo.UserData;
+  retInfo.size = Min(info->size, sizeof(hsa_amd_pointer_info_t));
+  retInfo.type = (hsa_amd_pointer_type_t)thunkInfo.Type;
+  retInfo.agentBaseAddress = reinterpret_cast<void*>(thunkInfo.GPUAddress);
+  retInfo.hostBaseAddress = thunkInfo.CPUAddress;
+  retInfo.sizeInBytes = thunkInfo.SizeInBytes;
+  retInfo.userData = thunkInfo.UserData;
+  retInfo.agentOwner = agents_by_node_[thunkInfo.Node][0]->public_handle();
+
+  memcpy(info, &retInfo, retInfo.size);
 
   if (returnListData) {
     uint32_t count = 0;

runtime/hsa-runtime/inc/hsa_ext_amd.h

@@ -1164,14 +1164,17 @@ hsa_status_t HSA_API hsa_amd_memory_unlock(void* host_ptr);
  *
  * @param[in] count Number of uint32_t element to be set to the value.
  *
- * @retval ::HSA_STATUS_SUCCESS The function has been executed successfully.
+ * @retval HSA_STATUS_SUCCESS The function has been executed successfully.
  *
- * @retval ::HSA_STATUS_ERROR_NOT_INITIALIZED The HSA runtime has not been
+ * @retval HSA_STATUS_ERROR_NOT_INITIALIZED The HSA runtime has not been
  * initialized.
  *
- * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT @p ptr is NULL or
+ * @retval HSA_STATUS_ERROR_INVALID_ARGUMENT @p ptr is NULL or
  * not 4 bytes aligned
  *
+ * @retval HSA_STATUS_ERROR_INVALID_ALLOCATION if the given memory
+ * region was not allocated with HSA runtime APIs.
+ *
  */
 hsa_status_t HSA_API
     hsa_amd_memory_fill(void* ptr, uint32_t value, size_t count);
@@ -1346,6 +1349,23 @@ typedef struct hsa_amd_pointer_info_v1_s {
   Application provided value.
   */
   void* userData;
+} hsa_amd_pointer_info_v1_t;
+
+/**
+ * @brief Minor version updates to pointer info.
+ */
+#ifdef __cplusplus
+typedef struct hsa_amd_pointer_info_v2_s : hsa_amd_pointer_info_v1_t {
+#else
+typedef struct hsa_amd_pointer_info_v2_t {
+  struct hsa_amd_pointer_info_v1_t;
+#endif
+  /*
+  Reports an agent which "owns" (ie has preferred access to) the pool in which the allocation was
+  made.  When multiple agents share equal access to a pool (ex: multiple CPU agents, or multi-die
+  GPU boards) any such agent may be returned.
+  */
+  hsa_agent_t agentOwner;
 } hsa_amd_pointer_info_t;
 
 /**