Allow zero agent handle in AsyncCopy APIs.

IPC use cases with RVD set can't convey proper agent handles. Runtime discovery is required to properly route the copy in this case. Change-Id: I4c97e132fb4b6ac1040de1cb17fe5a3e36d6be48 [ROCm/ROCR-Runtime commit: c289a43e88]
2022-04-19 22:51:53 -05:00
commit 289a86785b
@@ -115,7 +115,8 @@ class Agent : public Checked<0xF6BC25EB17E6F917> {
  explicit Agent(uint32_t node_id, DeviceType type)
      : node_id_(node_id),
        device_type_(uint32_t(type)),
-        profiling_enabled_(false) {
+        profiling_enabled_(false),
+        enabled_(false) {
    public_handle_ = Convert(this);
  }

@@ -267,6 +268,10 @@ class Agent : public Checked<0xF6BC25EB17E6F917> {
    return stat;
  }

+  __forceinline bool Enabled() const { return enabled_; }
+
+  __forceinline void Enable() { enabled_ = true; }
+
  virtual void Trim() {
    for (auto region : regions()) region->Trim();
  }
@@ -306,6 +311,8 @@ class Agent : public Checked<0xF6BC25EB17E6F917> {

  bool profiling_enabled_;

+  bool enabled_;
+
  // Used by an Agent's MemoryRegions to ensure serial memory operation on the device.
  // Serial memory operations are needed to ensure, among other things, that allocation failures are
  // due to true OOM conditions and per region caching (Trim and Allocate must be serial and
@@ -104,6 +104,11 @@ class CpuAgent : public core::Agent {
                           uint32_t group_segment_size,
                           core::Queue** queue) override;

+  // @brief Override from core::Agent.
+  hsa_status_t DmaCopy(void* dst, core::Agent& dst_agent, const void* src, core::Agent& src_agent,
+                       size_t size, std::vector<core::Signal*>& dep_signals,
+                       core::Signal& out_signal) override;
+
  // @brief Returns number of data caches.
  __forceinline size_t num_cache() const { return cache_props_.size(); }

@@ -132,7 +132,7 @@ class Runtime {

  /// @brief Insert agent into agent list ::agents_.
  /// @param [in] agent Pointer to the agent object.
-  void RegisterAgent(Agent* agent);
+  void RegisterAgent(Agent* agent, bool Enabled);

  /// @brief Delete all agent objects from ::agents_.
  void DestroyAgents();
@@ -221,10 +221,9 @@ class Runtime {
  ///
  /// @retval ::HSA_STATUS_SUCCESS if copy command has been submitted
  /// successfully to the agent DMA queue.
-  hsa_status_t CopyMemory(void* dst, core::Agent& dst_agent, const void* src,
-                          core::Agent& src_agent, size_t size,
-                          std::vector<core::Signal*>& dep_signals,
-                          core::Signal& completion_signal);
+  hsa_status_t CopyMemory(void* dst, core::Agent* dst_agent, const void* src,
+                          core::Agent* src_agent, size_t size,
+                          std::vector<core::Signal*>& dep_signals, core::Signal& completion_signal);

  /// @brief Fill the first @p count of uint32_t in ptr with value.
  ///
@@ -310,6 +309,8 @@ class Runtime {

  const std::vector<Agent*>& gpu_agents() { return gpu_agents_; }

+  const std::vector<Agent*>& disabled_gpu_agents() { return disabled_gpu_agents_; }
+
  const std::vector<uint32_t>& gpu_ids() { return gpu_ids_; }

  Agent* region_gpu() { return region_gpu_; }
@@ -505,6 +506,9 @@ class Runtime {
  // Agent list containing all compatible GPU agents in the platform.
  std::vector<Agent*> gpu_agents_;

+  // Agent list containing incompletely initialized GPU agents not to be used by the process.
+  std::vector<Agent*> disabled_gpu_agents_;
+
  // Agent map containing all agents indexed by their KFD node IDs.
  std::map<uint32_t, std::vector<Agent*> > agents_by_node_;

@@ -44,6 +44,7 @@

 #include <algorithm>
 #include <cstring>
+#include <thread>

 #include "core/inc/amd_memory_region.h"
 #include "core/inc/host_queue.h"
@@ -374,5 +375,37 @@ hsa_status_t CpuAgent::QueueCreate(size_t size, hsa_queue_type32_t queue_type,
  return HSA_STATUS_ERROR;
 }

+hsa_status_t CpuAgent::DmaCopy(void* dst, core::Agent& dst_agent, const void* src,
+                               core::Agent& src_agent, size_t size,
+                               std::vector<core::Signal*>& dep_signals, core::Signal& out_signal) {
+  // For cpu to cpu, fire and forget a copy thread.
+  const bool profiling_enabled = (dst_agent.profiling_enabled() || src_agent.profiling_enabled());
+  if (profiling_enabled) out_signal.async_copy_agent(this);
+  std::thread(
+      [](void* dst, const void* src, size_t size, std::vector<core::Signal*> dep_signals,
+         core::Signal* completion_signal, bool profiling_enabled) {
+        for (core::Signal* dep : dep_signals) {
+          dep->WaitRelaxed(HSA_SIGNAL_CONDITION_EQ, 0, UINT64_MAX, HSA_WAIT_STATE_BLOCKED);
+        }
+
+        if (profiling_enabled) {
+          core::Runtime::runtime_singleton_->GetSystemInfo(HSA_SYSTEM_INFO_TIMESTAMP,
+                                                           &completion_signal->signal_.start_ts);
+        }
+
+        memcpy(dst, src, size);
+
+        if (profiling_enabled) {
+          core::Runtime::runtime_singleton_->GetSystemInfo(HSA_SYSTEM_INFO_TIMESTAMP,
+                                                           &completion_signal->signal_.end_ts);
+        }
+
+        completion_signal->SubRelease(1);
+      },
+      dst, src, size, dep_signals, &out_signal, profiling_enabled)
+      .detach();
+  return HSA_STATUS_SUCCESS;
+}
+
 }  // namespace amd
 }  // namespace rocr
@@ -109,12 +109,14 @@ CpuAgent* DiscoverCpu(HSAuint32 node_id, HsaNodeProperties& node_prop) {
  }

  CpuAgent* cpu = new CpuAgent(node_id, node_prop);
-  core::Runtime::runtime_singleton_->RegisterAgent(cpu);
+  cpu->Enable();
+  core::Runtime::runtime_singleton_->RegisterAgent(cpu, true);

  return cpu;
 }

-GpuAgent* DiscoverGpu(HSAuint32 node_id, HsaNodeProperties& node_prop, bool xnack_mode) {
+GpuAgent* DiscoverGpu(HSAuint32 node_id, HsaNodeProperties& node_prop, bool xnack_mode,
+                      bool enabled) {
  GpuAgent* gpu = nullptr;
  if (node_prop.NumFComputeCores == 0) {
      // Ignore non GPUs.
@@ -163,7 +165,8 @@ GpuAgent* DiscoverGpu(HSAuint32 node_id, HsaNodeProperties& node_prop, bool xnac
      throw;
    }
  }
-  core::Runtime::runtime_singleton_->RegisterAgent(gpu);
+  if (enabled) gpu->Enable();
+  core::Runtime::runtime_singleton_->RegisterAgent(gpu, enabled);
  return gpu;
 }

@@ -242,7 +245,7 @@ void RegisterLinkInfo(uint32_t node_id, uint32_t num_link) {
 /**
 * Process the list of Gpus that are surfaced to user
 */
-static void SurfaceGpuList(std::vector<int32_t>& gpu_list, bool xnack_mode) {
+static void SurfaceGpuList(std::vector<int32_t>& gpu_list, bool xnack_mode, bool enabled) {
  // Process user visible Gpu devices
  int32_t invalidIdx = -1;
  int32_t list_sz = gpu_list.size();
@@ -259,7 +262,7 @@ static void SurfaceGpuList(std::vector<int32_t>& gpu_list, bool xnack_mode) {
    // Instantiate a Gpu device. The IO links
    // of this node have already been registered
    assert((node_prop.NumFComputeCores != 0) && "Improper node used for GPU device discovery.");
-    DiscoverGpu(gpu_list[idx], node_prop, xnack_mode);
+    DiscoverGpu(gpu_list[idx], node_prop, xnack_mode, enabled);
  }
 }

@@ -299,6 +302,7 @@ void BuildTopology() {
  int32_t invalidIdx = -1;
  uint32_t visibleCnt = 0;
  std::vector<int32_t> gpu_usr_list;
+  std::vector<int32_t> gpu_disabled;
  bool filter = RvdFilter::FilterDevices();
  if (filter) {
    rvdFilter.BuildRvdTokenList();
@@ -329,6 +333,8 @@ void BuildTopology() {
        int32_t devRank = rvdFilter.GetUsrDeviceRank(kfdIdx);
        if (devRank != (-1)) {
          gpu_usr_list[devRank] = node_id;
+        } else {
+          gpu_disabled.push_back(node_id);
        }
      } else {
        gpu_usr_list.push_back(node_id);
@@ -347,7 +353,8 @@ void BuildTopology() {
  bool xnack_mode = BindXnackMode();

  // Instantiate ROCr objects to encapsulate Gpu devices
-  SurfaceGpuList(gpu_usr_list, xnack_mode);
+  SurfaceGpuList(gpu_usr_list, xnack_mode, true);
+  SurfaceGpuList(gpu_disabled, xnack_mode, false);

  // Parse HSA_CU_MASK with GPU and CU count limits.
  uint32_t maxGpu = core::Runtime::runtime_singleton_->gpu_agents().size();
@@ -106,6 +106,12 @@ struct ValidityError<const T*> {
      return hsa_status_t(ValidityError<decltype(ptr)>::value); \
  } while (false)

+#define IS_NULL_OR_VALID(ptr)                                                                      \
+  do {                                                                                             \
+    if ((ptr) != NULL && !(ptr)->IsValid())                                                        \
+      return hsa_status_t(ValidityError<decltype(ptr)>::value);                                    \
+  } while (false)
+
 #define CHECK_ALLOC(ptr)                                         \
  do {                                                           \
    if ((ptr) == NULL) return HSA_STATUS_ERROR_OUT_OF_RESOURCES; \
@@ -240,10 +246,10 @@ hsa_status_t hsa_amd_memory_async_copy(void* dst, hsa_agent_t dst_agent_handle,
  }

  core::Agent* dst_agent = core::Agent::Convert(dst_agent_handle);
-  IS_VALID(dst_agent);
+  IS_NULL_OR_VALID(dst_agent);

  core::Agent* src_agent = core::Agent::Convert(src_agent_handle);
-  IS_VALID(src_agent);
+  IS_NULL_OR_VALID(src_agent);

  std::vector<core::Signal*> dep_signal_list(num_dep_signals);
  if (num_dep_signals > 0) {
@@ -260,8 +266,8 @@ hsa_status_t hsa_amd_memory_async_copy(void* dst, hsa_agent_t dst_agent_handle,
  bool rev_copy_dir = core::Runtime::runtime_singleton_->flag().rev_copy_dir();
  if (size > 0) {
    return core::Runtime::runtime_singleton_->CopyMemory(
-        dst, (rev_copy_dir ? *src_agent  : *dst_agent),
-        src, (rev_copy_dir ? *dst_agent  : *src_agent),
+        dst, (rev_copy_dir ? src_agent : dst_agent),
+        src, (rev_copy_dir ? dst_agent : src_agent),
        size, dep_signal_list, *out_signal_obj);
  }

@@ -46,7 +46,6 @@
 #include <atomic>
 #include <cstring>
 #include <string>
-#include <thread>
 #include <vector>

 #include "core/common/shared.h"
@@ -150,7 +149,7 @@ bool Runtime::IsOpen() {

 // Register agent information only.  Must not call anything that may use the registered information
 // since those tables are incomplete.
-void Runtime::RegisterAgent(Agent* agent) {
+void Runtime::RegisterAgent(Agent* agent, bool Enabled) {
  // Record the agent in the node-to-agent reverse lookup table.
  agents_by_node_[agent->node_id()].push_back(agent);

@@ -199,12 +198,13 @@ void Runtime::RegisterAgent(Agent* agent) {
      }
    }
  } else if (agent->device_type() == Agent::DeviceType::kAmdGpuDevice) {
-    gpu_agents_.push_back(agent);
-
-    gpu_ids_.push_back(agent->node_id());
-
-    // Assign the first discovered gpu agent as region gpu.
-    if (region_gpu_ == NULL) region_gpu_ = agent;
+    if (Enabled) {
+      gpu_agents_.push_back(agent);
+      gpu_ids_.push_back(agent->node_id());
+      // Assign the first discovered gpu agent as region gpu.
+      if (region_gpu_ == NULL) region_gpu_ = agent;
+    } else
+      disabled_gpu_agents_.push_back(agent);
  }
 }

@@ -214,6 +214,9 @@ void Runtime::DestroyAgents() {
  std::for_each(gpu_agents_.begin(), gpu_agents_.end(), DeleteObject());
  gpu_agents_.clear();

+  std::for_each(disabled_gpu_agents_.begin(), disabled_gpu_agents_.end(), DeleteObject());
+  gpu_agents_.clear();
+
  gpu_ids_.clear();

  std::for_each(cpu_agents_.begin(), cpu_agents_.end(), DeleteObject());
@@ -462,53 +465,37 @@ hsa_status_t Runtime::CopyMemory(void* dst, const void* src, size_t size) {
  return err;
 }

-hsa_status_t Runtime::CopyMemory(void* dst, core::Agent& dst_agent,
-                                 const void* src, core::Agent& src_agent,
-                                 size_t size,
+hsa_status_t Runtime::CopyMemory(void* dst, core::Agent* dst_agent, const void* src,
+                                 core::Agent* src_agent, size_t size,
                                 std::vector<core::Signal*>& dep_signals,
                                 core::Signal& completion_signal) {
-  const bool dst_gpu =
-      (dst_agent.device_type() == core::Agent::DeviceType::kAmdGpuDevice);
-  const bool src_gpu =
-      (src_agent.device_type() == core::Agent::DeviceType::kAmdGpuDevice);
-  if (dst_gpu || src_gpu) {
-    core::Agent* copy_agent = (src_gpu) ? &src_agent : &dst_agent;
-    return copy_agent->DmaCopy(dst, dst_agent, src, src_agent, size, dep_signals,
-                               completion_signal);
-  }
+  auto lookupAgent = [this](core::Agent* agent, const void* ptr) {
+    if (agent == nullptr) {
+      hsa_amd_pointer_info_t info;
+      PtrInfoBlockData block;
+      info.size = sizeof(info);
+      PtrInfo(ptr, &info, nullptr, nullptr, nullptr, &block);
+      // Limit to IPC and GFX types for now.  These are the only types for which the application may
+      // not posess a proper agent handle.
+      if ((info.type != HSA_EXT_POINTER_TYPE_IPC) && (info.type != HSA_EXT_POINTER_TYPE_GRAPHICS))
+        return;
+      agent = Agent::Convert(info.agentOwner);
+    }
+  };

-  // For cpu to cpu, fire and forget a copy thread.
-  const bool profiling_enabled =
-      (dst_agent.profiling_enabled() || src_agent.profiling_enabled());
-  if (profiling_enabled) completion_signal.async_copy_agent(&dst_agent);
-  std::thread(
-      [](void* dst, const void* src, size_t size,
-         std::vector<core::Signal*> dep_signals,
-         core::Signal* completion_signal, bool profiling_enabled) {
+  lookupAgent(dst_agent, dst);
+  lookupAgent(src_agent, src);
+  if (dst_agent == nullptr || src_agent == nullptr) return HSA_STATUS_ERROR_INVALID_AGENT;

-        for (core::Signal* dep : dep_signals) {
-          dep->WaitRelaxed(HSA_SIGNAL_CONDITION_EQ, 0, UINT64_MAX,
-                           HSA_WAIT_STATE_BLOCKED);
-        }
+  // At least one agent must be available for operation in the current process.
+  if (!dst_agent->Enabled() && !src_agent->Enabled()) return HSA_STATUS_ERROR_INVALID_AGENT;

-        if (profiling_enabled) {
-          core::Runtime::runtime_singleton_->GetSystemInfo(HSA_SYSTEM_INFO_TIMESTAMP,
-                                                           &completion_signal->signal_.start_ts);
-        }
-
-        memcpy(dst, src, size);
-
-        if (profiling_enabled) {
-          core::Runtime::runtime_singleton_->GetSystemInfo(HSA_SYSTEM_INFO_TIMESTAMP,
-                                                           &completion_signal->signal_.end_ts);
-        }
-
-        completion_signal->SubRelease(1);
-      },
-      dst, src, size, dep_signals, &completion_signal,
-      profiling_enabled).detach();
-
-  return HSA_STATUS_SUCCESS;
+  const bool dst_gpu = (dst_agent->device_type() == core::Agent::DeviceType::kAmdGpuDevice);
+  const bool src_gpu = (src_agent->device_type() == core::Agent::DeviceType::kAmdGpuDevice);
+  core::Agent* copy_agent = (src_gpu) ? src_agent : dst_agent;
+  if (!copy_agent->Enabled()) copy_agent = (copy_agent == src_agent) ? dst_agent : src_agent;
+  return copy_agent->DmaCopy(dst, *dst_agent, src, *src_agent, size, dep_signals,
+                             completion_signal);
 }

 hsa_status_t Runtime::FillMemory(void* ptr, uint32_t value, size_t count) {
@@ -1110,6 +1110,8 @@ hsa_status_t HSA_API hsa_amd_memory_pool_free(void* ptr);
 *
 * @param[in] dst_agent Agent associated with the @p dst. The agent must be able to directly
 * access both the source and destination buffers in their current locations.
+ * May be zero in which case the runtime will attempt to discover the destination agent.
+ * Discovery may have variable and/or high latency.
 *
 * @param[in] src A valid pointer to the source of data to be copied. The source
 * buffer must not overlap with the destination buffer, otherwise the copy will succeed
@@ -1117,6 +1119,8 @@ hsa_status_t HSA_API hsa_amd_memory_pool_free(void* ptr);
 *
 * @param[in] src_agent Agent associated with the @p src. The agent must be able to directly
 * access both the source and destination buffers in their current locations.
+ * May be zero in which case the runtime will attempt to discover the destination agent.
+ * Discovery may have variable and/or high latency.
 *
 * @param[in] size Number of bytes to copy. If @p size is 0, no copy is
 * performed and the function returns success. Copying a number of bytes larger
@@ -1127,9 +1131,9 @@ hsa_status_t HSA_API hsa_amd_memory_pool_free(void* ptr);
 *
 * @param[in] dep_signals List of signals that must be waited on before the copy
 * operation starts. The copy will start after every signal has been observed with
- * the value 0. The dependent signal should not include completion signal from hsa_amd_memory_async_copy
- * operation to be issued in future as that can result in a deadlock. If @p num_dep_signals is 0, this
- * argument is ignored.
+ * the value 0. The dependent signal should not include completion signal from
+ * hsa_amd_memory_async_copy operation to be issued in future as that can result
+ * in a deadlock. If @p num_dep_signals is 0, this argument is ignored.
 *
 * @param[in] completion_signal Signal used to indicate completion of the copy
 * operation. When the copy operation is finished, the value of the signal is
@@ -1144,7 +1148,7 @@ hsa_status_t HSA_API hsa_amd_memory_pool_free(void* ptr);
 * @retval ::HSA_STATUS_ERROR_NOT_INITIALIZED The HSA runtime has not been
 * initialized.
 *
- * @retval ::HSA_STATUS_ERROR_INVALID_AGENT The agent is invalid.
+ * @retval ::HSA_STATUS_ERROR_INVALID_AGENT An agent is invalid or no discovered agent has access.
 *
 * @retval ::HSA_STATUS_ERROR_INVALID_SIGNAL @p completion_signal is invalid.
 *