Add pointer info support.

Change-Id: I3edcc0bfddbf12465065c9bc3b6565288faff1b8 [ROCm/ROCR-Runtime commit: 9dd76dbeda]
2016-11-09 22:21:24 -06:00
parent b5d60f56dd
commit c9e78a451b
10 changed files with 279 additions and 26 deletions
@@ -469,11 +469,15 @@ hsa_status_t MemoryRegion::AllowAccess(uint32_t num_agents,
  HsaMemMapFlags map_flag = map_flag_;
  map_flag.ui32.HostAccess |= (cpu_in_list) ? 1 : 0;

-  uint64_t alternate_va = 0;
-  if (!amd::MemoryRegion::MakeKfdMemoryResident(
-          whitelist_nodes.size(), &whitelist_nodes[0], const_cast<void*>(ptr),
-          size, &alternate_va, map_flag)) {
-    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  
+  {
+    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),
+      size, &alternate_va, map_flag)) {
+        return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+    }
  }

  for (GpuAgentInt* gpu : whitelist_gpus) {
@@ -358,6 +358,8 @@ void HsaApiTable::UpdateAmdExts() {
  amd_ext_api.hsa_amd_memory_fill_fn = AMD::hsa_amd_memory_fill;
  amd_ext_api.hsa_amd_interop_map_buffer_fn = AMD::hsa_amd_interop_map_buffer;
  amd_ext_api.hsa_amd_interop_unmap_buffer_fn = AMD::hsa_amd_interop_unmap_buffer;
+  amd_ext_api.hsa_amd_pointer_info_fn = AMD::hsa_amd_pointer_info;
+  amd_ext_api.hsa_amd_pointer_info_set_userdata_fn = AMD::hsa_amd_pointer_info_set_userdata;
 }

 class Init {
@@ -597,4 +597,18 @@ hsa_status_t hsa_amd_interop_unmap_buffer(void* ptr) {
  return HSA_STATUS_SUCCESS;
 }

+hsa_status_t hsa_amd_pointer_info(void* ptr, hsa_amd_pointer_info_t* info, void* (*alloc)(size_t),
+                              uint32_t* num_accessible, hsa_agent_t** accessible) {
+  IS_OPEN();
+  IS_BAD_PTR(ptr);
+  IS_BAD_PTR(info);
+  return core::Runtime::runtime_singleton_->PtrInfo(ptr, info, alloc, num_accessible, accessible);
+}
+
+hsa_status_t hsa_amd_pointer_info_set_userdata(void* ptr, void* userdata) {
+  IS_OPEN();
+  IS_BAD_PTR(ptr);
+  return core::Runtime::runtime_singleton_->SetPtrInfoData(ptr, userdata);
+}
+
 } // end of AMD namespace
@@ -142,6 +142,10 @@ bool Runtime::IsOpen() {
 }

 void Runtime::RegisterAgent(Agent* agent) {
+  // Record the agent in the node-to-agent reverse lookup table.
+  agents_by_node_[agent->node_id()].push_back(agent);
+
+  // Process agent as a cpu or gpu device.
  if (agent->device_type() == Agent::DeviceType::kAmdCpuDevice) {
    cpu_agents_.push_back(agent);

@@ -230,6 +234,8 @@ void Runtime::RegisterAgent(Agent* agent) {
 }

 void Runtime::DestroyAgents() {
+  agents_by_node_.clear();
+
  std::for_each(gpu_agents_.begin(), gpu_agents_.end(), DeleteObject());
  gpu_agents_.clear();

@@ -302,11 +308,11 @@ hsa_status_t Runtime::IterateAgent(hsa_status_t (*callback)(hsa_agent_t agent,
 hsa_status_t Runtime::AllocateMemory(const MemoryRegion* region, size_t size,
                                     MemoryRegion::AllocateFlags alloc_flags,
                                     void** address) {
+  ScopedAcquire<KernelMutex> lock(&memory_lock_);
  hsa_status_t status = region->Allocate(size, alloc_flags, address);

  // Track the allocation result so that it could be freed properly.
  if (status == HSA_STATUS_SUCCESS) {
-    ScopedAcquire<KernelMutex> lock(&memory_lock_);
    allocation_map_[*address] = AllocationRegion(region, size);
  }

@@ -320,22 +326,18 @@ hsa_status_t Runtime::FreeMemory(void* ptr) {

  const MemoryRegion* region = NULL;
  size_t size = 0;
-  {
-    ScopedAcquire<KernelMutex> lock(&memory_lock_);
+  ScopedAcquire<KernelMutex> lock(&memory_lock_);

-    std::map<const void*, AllocationRegion>::const_iterator it =
-        allocation_map_.find(ptr);
+  std::map<const void*, AllocationRegion>::const_iterator it = allocation_map_.find(ptr);

-    if (it == allocation_map_.end()) {
-      assert(false && "Can't find address in allocation map");
-      return HSA_STATUS_ERROR;
-    }
-
-    region = it->second.region;
-    size = it->second.size;
-
-    allocation_map_.erase(it);
+  if (it == allocation_map_.end()) {
+    assert(false && "Can't find address in allocation map");
+    return HSA_STATUS_ERROR;
  }
+  region = it->second.region;
+  size = it->second.size;
+
+  allocation_map_.erase(it);

  return region->Free(ptr, size);
 }
@@ -463,8 +465,7 @@ hsa_status_t Runtime::AllowAccess(uint32_t num_agents,
  {
    ScopedAcquire<KernelMutex> lock(&memory_lock_);

-    std::map<const void*, AllocationRegion>::const_iterator it =
-        allocation_map_.find(ptr);
+    std::map<const void*, AllocationRegion>::const_iterator it = allocation_map_.find(ptr);

    if (it == allocation_map_.end()) {
      return HSA_STATUS_ERROR;
@@ -634,8 +635,7 @@ hsa_status_t Runtime::InteropMap(uint32_t num_agents, Agent** agents,
  return HSA_STATUS_SUCCESS;
 }

-hsa_status_t Runtime::InteropUnmap(void* ptr)
-{
+hsa_status_t Runtime::InteropUnmap(void* ptr) {
  if(hsaKmtUnmapMemoryToGPU(ptr)!=HSAKMT_STATUS_SUCCESS)
    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
  if(hsaKmtDeregisterMemory(ptr)!=HSAKMT_STATUS_SUCCESS)
@@ -643,6 +643,78 @@ hsa_status_t Runtime::InteropUnmap(void* ptr)
  return HSA_STATUS_SUCCESS;
 }

+hsa_status_t Runtime::PtrInfo(void* ptr, hsa_amd_pointer_info_t* info, void* (*alloc)(size_t),
+                              uint32_t* num_agents_accessible, hsa_agent_t** accessible) {
+  HsaPointerInfo thunkInfo;
+  uint32_t* mappedNodes;
+
+  // 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;
+
+  bool returnListData =
+      ((alloc != nullptr) && (num_agents_accessible != nullptr) && (accessible != nullptr));
+  if (returnListData) {
+    size_t max_agents = cpu_agents_.size() + gpu_agents_.size();
+    mappedNodes = (uint32_t*)alloca(max_agents * sizeof(uint32_t));
+    // memory_lock protects access to the NMappedNodes array since this changes with calls to memory
+    // APIs.
+    ScopedAcquire<KernelMutex> lock(&memory_lock_);
+    hsaKmtQueryPointerInfo(ptr, &thunkInfo);
+    assert(thunkInfo.NMappedNodes <= max_agents &&
+           "PointerInfo: Thunk returned more than all agents in NMappedNodes.");
+    memcpy(mappedNodes, thunkInfo.MappedNodes, thunkInfo.NMappedNodes * sizeof(uint32_t));
+  } else {
+    hsaKmtQueryPointerInfo(ptr, &thunkInfo);
+  }
+
+  static_assert((int)HSA_POINTER_UNKNOWN == (int)HSA_EXT_POINTER_TYPE_UNKNOWN,
+                "Thunk pointer info mismatch");
+  static_assert((int)HSA_POINTER_ALLOCATED == (int)HSA_EXT_POINTER_TYPE_HSA,
+                "Thunk pointer info mismatch");
+  static_assert((int)HSA_POINTER_REGISTERED_USER == (int)HSA_EXT_POINTER_TYPE_LOCKED,
+                "Thunk pointer info mismatch");
+  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;
+
+  if (returnListData) {
+    uint32_t count = 0;
+    for (int i = 0; i < thunkInfo.NMappedNodes; i++) {
+      assert(mappedNodes[i] < agents_by_node_.size() &&
+             "PointerInfo: Invalid node ID returned from thunk.");
+      count += agents_by_node_[mappedNodes[i]].size();
+    }
+
+    *accessible = (hsa_agent_t*)alloc(sizeof(hsa_agent_t) * count);
+    if ((*accessible) == nullptr) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+    *num_agents_accessible = count;
+
+    uint32_t index = 0;
+    for (int i = 0; i < thunkInfo.NMappedNodes; i++) {
+      auto& list = agents_by_node_[mappedNodes[i]];
+      for (int j = 0; j < list.size(); j++) {
+        (*accessible)[index] = list[j]->public_handle();
+        index++;
+      }
+    }
+  }
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t Runtime::SetPtrInfoData(void* ptr, void* userptr) {
+  if (hsaKmtSetMemoryUserData(ptr, userptr) == HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_SUCCESS;
+  else
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+}
+
 void Runtime::AsyncEventsLoop(void*) {
  auto& async_events_control_ = runtime_singleton_->async_events_control_;
  auto& async_events_ = runtime_singleton_->async_events_;