Support Get and Set access for memory mappings

Mapping memory handles to virtual memory addresses do not make them accessible. The set access function is needed to make the memory mappings accessible to specific agents. The get access function returns current access properties for individual agents. This is part of patch series for Virtual Memory API. Change-Id: I152ba0557fd2a802eb9d840568b68cdd1911b72c [ROCm/ROCR-Runtime commit: 13fbd8a232]
2022-12-19 00:04:08 +00:00
commit 933aac4cda
@@ -1259,6 +1259,17 @@ hsa_status_t HSA_API hsa_amd_vmem_unmap(void* va, size_t size) {
  return amdExtTable->hsa_amd_vmem_unmap_fn(va, size);
 }

+hsa_status_t HSA_API hsa_amd_vmem_set_access(void* va, size_t size,
+                                             const hsa_amd_memory_access_desc_t* desc,
+                                             const size_t desc_cnt) {
+  return amdExtTable->hsa_amd_vmem_set_access_fn(va, size, desc, desc_cnt);
+}
+
+hsa_status_t HSA_API hsa_amd_vmem_get_access(void* va, hsa_access_permission_t* perms,
+                                             const hsa_agent_t agent_handle) {
+  return amdExtTable->hsa_amd_vmem_get_access_fn(va, perms, agent_handle);
+}
+
 // Tools only table interfaces.
 namespace rocr {

@@ -320,6 +320,15 @@ hsa_status_t hsa_amd_vmem_map(void* va, size_t size, size_t in_offset,
 // Mirrors Amd Extension Apis
 hsa_status_t hsa_amd_vmem_unmap(void* va, size_t size);

+// Mirrors Amd Extension Apis
+hsa_status_t hsa_amd_vmem_set_access(void* va, size_t size,
+                                     const hsa_amd_memory_access_desc_t* desc,
+                                     const size_t desc_cnt);
+
+// Mirrors Amd Extension Apis
+hsa_status_t hsa_amd_vmem_get_access(void* va, hsa_access_permission_t* flags,
+                                     const hsa_agent_t agent_handle);
+
 }  // namespace amd
 }  // namespace rocr

@@ -369,6 +369,12 @@ class Runtime {

  hsa_status_t VMemoryHandleUnmap(void* va, size_t size);

+  hsa_status_t VMemorySetAccess(void* va, size_t size, const hsa_amd_memory_access_desc_t* desc,
+                                size_t desc_cnt);
+
+  hsa_status_t VMemoryGetAccess(const void* va, hsa_access_permission_t* perms,
+                                hsa_agent_t agent_handle);
+
  const std::vector<Agent*>& cpu_agents() { return cpu_agents_; }

  const std::vector<Agent*>& gpu_agents() { return gpu_agents_; }
@@ -729,6 +735,30 @@ class Runtime {
  };
  std::map<ThunkHandle, MemoryHandle> memory_handle_map_;

+  struct MappedHandle;
+  struct MappedHandleAllowedAgent {
+    MappedHandleAllowedAgent()
+        : va(NULL), permissions(HSA_ACCESS_PERMISSION_NONE), mappedHandle(NULL), ldrm_bo(0) {}
+    MappedHandleAllowedAgent(MappedHandle* _mappedHandle, Agent* targetAgent, void* va, size_t size,
+                             hsa_access_permission_t perms)
+        : va(va),
+          size(size),
+          targetAgent(targetAgent),
+          permissions(perms),
+          mappedHandle(_mappedHandle),
+          ldrm_bo(0) {}
+
+    hsa_status_t RemoveAccess();
+    hsa_status_t EnableAccess(hsa_access_permission_t perms);
+
+    void* va;
+    size_t size;
+    Agent* targetAgent;
+    hsa_access_permission_t permissions;
+    MappedHandle* mappedHandle;
+    amdgpu_bo_handle ldrm_bo;
+  };
+
  struct MappedHandle {
    MappedHandle()
        : mem_handle(NULL),
@@ -762,6 +792,7 @@ class Runtime {
    int drm_fd;
    void* drm_cpu_addr;  // CPU Buffer address
    amdgpu_bo_handle ldrm_bo;
+    std::map<Agent*, MappedHandleAllowedAgent> allowed_agents;
  };
  std::map<const void*, MappedHandle> mapped_handle_map_;  // Indexed by VA

@@ -408,6 +408,8 @@ void HsaApiTable::UpdateAmdExts() {
  amd_ext_api.hsa_amd_vmem_handle_release_fn = AMD::hsa_amd_vmem_handle_release;
  amd_ext_api.hsa_amd_vmem_map_fn = AMD::hsa_amd_vmem_map;
  amd_ext_api.hsa_amd_vmem_unmap_fn = AMD::hsa_amd_vmem_unmap;
+  amd_ext_api.hsa_amd_vmem_set_access_fn = AMD::hsa_amd_vmem_set_access;
+  amd_ext_api.hsa_amd_vmem_get_access_fn = AMD::hsa_amd_vmem_get_access;
 }

 void LoadInitialHsaApiTable() {
@@ -1276,5 +1276,29 @@ hsa_status_t hsa_amd_vmem_unmap(void* va, size_t size) {
  CATCH;
 }

+hsa_status_t hsa_amd_vmem_set_access(void* va, size_t size,
+                                     const hsa_amd_memory_access_desc_t* desc,
+                                     size_t desc_cnt) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(va);
+  IS_BAD_PTR(desc);
+  IS_ZERO(desc_cnt);
+
+  return core::Runtime::runtime_singleton_->VMemorySetAccess(va, size, desc, desc_cnt);
+  CATCH;
+}
+
+hsa_status_t hsa_amd_vmem_get_access(void* va, hsa_access_permission_t* perms,
+                                     hsa_agent_t agent_handle) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(va);
+  IS_BAD_PTR(perms);
+
+  return core::Runtime::runtime_singleton_->VMemoryGetAccess(va, perms, agent_handle);
+  CATCH;
+}
+
 }   //  namespace amd
 }   //  namespace rocr
@@ -2620,7 +2620,22 @@ hsa_status_t Runtime::VMemoryHandleUnmap(void* va, size_t size) {

  if (mappedHandleIt->second.size != size) return HSA_STATUS_ERROR_INVALID_ARGUMENT;

-  if (mappedHandleIt->second.ldrm_bo) ret = amdgpu_bo_free(mappedHandleIt->second.ldrm_bo);
+  for (auto agentPermsIt = mappedHandleIt->second.allowed_agents.begin();
+       agentPermsIt != mappedHandleIt->second.allowed_agents.end();) {
+    assert(va == agentPermsIt->second.va);
+    if (agentPermsIt->second.ldrm_bo)
+      ret = amdgpu_bo_va_op(agentPermsIt->second.ldrm_bo, mappedHandleIt->second.offset, size,
+                            reinterpret_cast<uint64_t>(va), 0, AMDGPU_VA_OP_UNMAP);
+    else
+      ret = munmap(va, size);
+    if (ret) return HSA_STATUS_ERROR;
+    agentPermsIt = mappedHandleIt->second.allowed_agents.erase(agentPermsIt);
+  }
+
+  if (mappedHandleIt->second.ldrm_bo)
+    ret = amdgpu_bo_free(mappedHandleIt->second.ldrm_bo);
+  else
+    ret = munmap(va, size);

  if (ret) return HSA_STATUS_ERROR;

@@ -2641,5 +2656,135 @@ hsa_status_t Runtime::VMemoryHandleUnmap(void* va, size_t size) {
  mapped_handle_map_.erase(mappedHandleIt);
  return HSA_STATUS_SUCCESS;
 }
+
+hsa_status_t Runtime::MappedHandleAllowedAgent::EnableAccess(hsa_access_permission_t perms) {
+  if (targetAgent->device_type() == core::Agent::DeviceType::kAmdCpuDevice) {
+    void* ret_cpu_addr =
+        mmap(va, size, mmap_perm(perms), MAP_SHARED | MAP_FIXED, mappedHandle->drm_fd,
+             reinterpret_cast<uint64_t>(mappedHandle->drm_cpu_addr));
+    assert(ret_cpu_addr == va);
+  } else {
+    int ret;
+    ret = amdgpu_bo_va_op(mappedHandle->ldrm_bo, mappedHandle->offset, mappedHandle->size,
+                          reinterpret_cast<uint64_t>(va), drm_perm(perms), AMDGPU_VA_OP_MAP);
+    if (ret) return HSA_STATUS_ERROR;
+
+    ldrm_bo = mappedHandle->ldrm_bo;
+  }
+  permissions = perms;
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t Runtime::MappedHandleAllowedAgent::RemoveAccess() {
+  int ret;
+
+  if (!ldrm_bo)  // Mapped to host
+    ret = munmap(va, mappedHandle->size);
+  else  // Mapped to device
+    ret = amdgpu_bo_va_op(ldrm_bo, mappedHandle->offset, mappedHandle->size,
+                          reinterpret_cast<uint64_t>(va), 0, AMDGPU_VA_OP_UNMAP);
+
+  return (ret) ? HSA_STATUS_ERROR : HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t Runtime::VMemorySetAccess(void* va, size_t size,
+                                       const hsa_amd_memory_access_desc_t* desc,
+                                       const size_t desc_cnt) {
+  int nodesCnt = 0;
+  ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
+
+  auto mappedHandleIt = mapped_handle_map_.find(va);
+  if (mappedHandleIt == mapped_handle_map_.end() || mappedHandleIt->second.size != size)
+    return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+  static const int tinyArraySize = 8;
+
+  HSAuint32 short_nodes[tinyArraySize];
+  HSAuint32* nodes = short_nodes;
+  if (desc_cnt > tinyArraySize) {
+    nodes = new HSAuint32[desc_cnt];
+
+    if (nodes == NULL) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+  MAKE_SCOPE_GUARD([&]() {
+    if (desc_cnt > tinyArraySize) delete[] nodes;
+  });
+
+  for (int i = 0; i < desc_cnt; i++) {
+    Agent* targetAgent = Agent::Convert(desc[i].agent_handle);
+
+    if (targetAgent == NULL || !targetAgent->IsValid()) return HSA_STATUS_ERROR_INVALID_AGENT;
+
+    if (targetAgent->device_type() == core::Agent::DeviceType::kAmdGpuDevice)
+      nodes[nodesCnt++] = targetAgent->node_id();
+
+    auto agentPermsIt = mappedHandleIt->second.allowed_agents.find(targetAgent);
+    if (agentPermsIt == mappedHandleIt->second.allowed_agents.end()) {
+      /* Agent not previously allowed, we need a new entry */
+      MappedHandleAllowedAgent newAgentPerms(&mappedHandleIt->second, targetAgent, va, size,
+                                             desc[i].permissions);
+      if (newAgentPerms.EnableAccess(desc[i].permissions) != HSA_STATUS_SUCCESS)
+        return HSA_STATUS_ERROR;
+
+      mappedHandleIt->second.allowed_agents[targetAgent] = newAgentPerms;
+    } else {
+      /* Previous permissions are same as current permission */
+      if (agentPermsIt->second.permissions == desc[i].permissions) continue;
+
+      if (agentPermsIt->second.RemoveAccess() != HSA_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+      if (agentPermsIt->second.EnableAccess(desc[i].permissions) != HSA_STATUS_SUCCESS) {
+        mappedHandleIt->second.allowed_agents.erase(agentPermsIt);
+        return HSA_STATUS_ERROR;
+      }
+    }
+  }
+
+  // Remove agents that were previously allowed but not included in current list
+  for (auto agentPermsIt = mappedHandleIt->second.allowed_agents.begin();
+       agentPermsIt != mappedHandleIt->second.allowed_agents.end();) {
+    bool agent_removed = true;
+    for (int i = 0; i < desc_cnt; i++) {
+      if (agentPermsIt->first == Agent::Convert(desc[i].agent_handle)) {
+        agent_removed = false;
+        continue;
+      }
+    }
+    if (agent_removed) {
+      assert(agentPermsIt->second.va == va);
+
+      if (agentPermsIt->second.RemoveAccess() != HSA_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+
+      agentPermsIt = mappedHandleIt->second.allowed_agents.erase(agentPermsIt);
+    } else {
+      ++agentPermsIt;
+    }
+  }
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t Runtime::VMemoryGetAccess(const void* va, hsa_access_permission_t* perms,
+                                       hsa_agent_t agent_handle) {
+  *perms = HSA_ACCESS_PERMISSION_NONE;
+
+  ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
+
+  auto mappedHandleIt = mapped_handle_map_.find(va);
+  if (mappedHandleIt == mapped_handle_map_.end()) {
+    return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+  }
+
+  Agent* agent = Agent::Convert(agent_handle);
+  if (agent == NULL || !agent->IsValid() || agent->device_type() != core::Agent::kAmdGpuDevice)
+    return HSA_STATUS_ERROR_INVALID_AGENT;
+
+  auto agentPermsIt = mappedHandleIt->second.allowed_agents.find(agent);
+  if (agentPermsIt != mappedHandleIt->second.allowed_agents.end()) {
+    *perms = agentPermsIt->second.permissions;
+    return HSA_STATUS_SUCCESS;
+  }
+  return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+}
+
 }  // namespace core
 }  // namespace rocr
@@ -239,6 +239,8 @@ global:
 	hsa_amd_vmem_handle_release;
 	hsa_amd_vmem_map;
 	hsa_amd_vmem_unmap;
+	hsa_amd_vmem_set_access;
+	hsa_amd_vmem_get_access;

 local:
    *;
@@ -200,6 +200,8 @@ struct AmdExtTable {
  decltype(hsa_amd_vmem_handle_release)* hsa_amd_vmem_handle_release_fn;
  decltype(hsa_amd_vmem_map)* hsa_amd_vmem_map_fn;
  decltype(hsa_amd_vmem_unmap)* hsa_amd_vmem_unmap_fn;
+  decltype(hsa_amd_vmem_set_access)* hsa_amd_vmem_set_access_fn;
+  decltype(hsa_amd_vmem_get_access)* hsa_amd_vmem_get_access_fn;
 };

 // Table to export HSA Core Runtime Apis
@@ -2844,6 +2844,60 @@ hsa_status_t hsa_amd_vmem_map(void* va, size_t size, size_t in_offset,
 */
 hsa_status_t hsa_amd_vmem_unmap(void* va, size_t size);

+typedef struct hsa_amd_memory_access_desc_s {
+  hsa_access_permission_t permissions;
+  hsa_agent_t agent_handle;
+} hsa_amd_memory_access_desc_t;
+
+/*
+ * @brief Make a memory mapping accessible
+ *
+ * Make previously mapped virtual address accessible to specific agents. @p size must be equal to
+ * size of previously mapped virtual memory handle.
+ * Calling hsa_amd_vmem_set_access multiple times on the same @p va will overwrite previous
+ * permissions for all agents
+ *
+ * @param[in] va previously mapped virtual address
+ * @param[in] size of memory mapping
+ * @param[in] desc list of access permissions for each agent
+ * @param[in] desc_cnt number of elements in desc
+ *
+ * @retval ::HSA_STATUS_SUCCESS
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT va, size or memory_handle are invalid
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ALLOCATION memory_handle is invalid
+ *
+ * @retval ::HSA_STATUS_ERROR_OUT_OF_RESOURCES Insufficient resources
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_AGENT Invalid agent in desc
+ *
+ * @retval ::HSA_STATUS_ERROR Unexpected internal error
+ */
+hsa_status_t hsa_amd_vmem_set_access(void* va, size_t size,
+                                     const hsa_amd_memory_access_desc_t* desc,
+                                     size_t desc_cnt);
+
+/*
+ * @brief Get current access permissions for memory mapping
+ *
+ * Get access permissions for memory mapping for specific agent.
+ *
+ * @param[in] va previously mapped virtual address
+ * @param[in] perms current permissions
+ * @param[in] agent_handle agent
+ *
+ * @retval ::HSA_STATUS_SUCCESS
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_AGENT Invalid agent
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ALLOCATION va is not mapped or permissions never set for this
+ * agent
+ *
+ * @retval ::HSA_STATUS_ERROR Unexpected internal error
+ */
+hsa_status_t hsa_amd_vmem_get_access(void* va, hsa_access_permission_t* perms,
+                                     hsa_agent_t agent_handle);
 #ifdef __cplusplus
 }  // end extern "C" block
 #endif