Support hsa_amd_agents_allow_access on page fragments.

Since access may only be manipulated on whole pages, suballocator fragments must cooperate to set the page's access.
Since the KFD does not migrate memory on access changes this implementation makes agent access sticky across the requests in a fragmented page.

Change-Id: I88479ed45fb40e9782b704526a7b8ffb22e7bd76

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

@@ -50,6 +50,7 @@
 #include "core/inc/agent.h"
 #include "core/inc/memory_region.h"
 #include "core/util/simple_heap.h"
+#include "core/util/locks.h"
 
 #include "inc/hsa_ext_amd.h"
 
@@ -98,8 +99,7 @@ class MemoryRegion : public core::MemoryRegion {
 
   ~MemoryRegion();
 
-  hsa_status_t Allocate(size_t size, AllocateFlags alloc_flags,
-                        void** address) const;
+  hsa_status_t Allocate(size_t& size, AllocateFlags alloc_flags, void** address) const;
 
   hsa_status_t Free(void* address, size_t size) const;
 
@@ -181,6 +181,8 @@ class MemoryRegion : public core::MemoryRegion {
 
   HSAuint64 virtual_size_;
 
+  mutable KernelMutex access_lock_;
+
   static const size_t kPageSize_ = 4096;
 
   class BlockAllocator {

runtime/hsa-runtime/core/inc/memory_region.h

@@ -90,8 +90,7 @@ class MemoryRegion : public Checked<0x9C961F19EE175BB3> {
 
   typedef uint32_t AllocateFlags;
 
-  virtual hsa_status_t Allocate(size_t size, AllocateFlags alloc_flags,
-                                void** address) const = 0;
+  virtual hsa_status_t Allocate(size_t& size, AllocateFlags alloc_flags, void** address) const = 0;
 
   virtual hsa_status_t Free(void* address, size_t size) const = 0;
 

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

@@ -148,8 +148,7 @@ MemoryRegion::MemoryRegion(bool fine_grain, bool full_profile, core::Agent* owne
 
 MemoryRegion::~MemoryRegion() {}
 
-hsa_status_t MemoryRegion::Allocate(size_t size, AllocateFlags alloc_flags,
-                                    void** address) const {
+hsa_status_t MemoryRegion::Allocate(size_t& size, AllocateFlags alloc_flags, void** address) const {
   if (address == NULL) {
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   }
@@ -171,13 +170,16 @@ hsa_status_t MemoryRegion::Allocate(size_t size, AllocateFlags alloc_flags,
       (alloc_flags & AllocateDoubleMap ? 1 : 0);
 
   // Only allow using the suballocator for ordinary VRAM.
-  bool useSubAlloc = !core::Runtime::runtime_singleton_->flag().disable_fragment_alloc();
-  useSubAlloc &= IsLocalMemory();
-  useSubAlloc &= (alloc_flags == AllocateRestrict);
-  useSubAlloc &= (size <= fragment_allocator_.max_alloc());
-  if (useSubAlloc) {
-    *address = fragment_allocator_.alloc(size);
-    return HSA_STATUS_SUCCESS;
+  if (IsLocalMemory()) {
+    bool useSubAlloc = !core::Runtime::runtime_singleton_->flag().disable_fragment_alloc();
+    useSubAlloc &= (alloc_flags == AllocateRestrict);
+    useSubAlloc &= (size <= fragment_allocator_.max_alloc());
+    if (useSubAlloc) {
+      *address = fragment_allocator_.alloc(size);
+      return HSA_STATUS_SUCCESS;
+    }
+    // Pad up larger VRAM allocations.
+    size = AlignUp(size, fragment_allocator_.max_alloc());
   }
 
   *address = AllocateKfdMemory(kmt_alloc_flags, owner()->node_id(), size);
@@ -444,6 +446,34 @@ hsa_status_t MemoryRegion::AllowAccess(uint32_t num_agents,
     return HSA_STATUS_ERROR;
   }
 
+  // Adjust for fragments.  Make accessibility sticky for fragments since this will satisfy the
+  // union of accessible agents between the fragments in the block.
+  hsa_amd_pointer_info_t info;
+  uint32_t agent_count = 0;
+  hsa_agent_t* accessible = nullptr;
+  MAKE_SCOPE_GUARD([&]() { free(accessible); });
+  core::Runtime::PtrInfoBlockData blockInfo;
+  std::vector<uint64_t> union_agents;
+  info.size = sizeof(info);
+
+  ScopedAcquire<KernelMutex> lock(&access_lock_);
+  if (core::Runtime::runtime_singleton_->PtrInfo(const_cast<void*>(ptr), &info, malloc,
+                                                 &agent_count, &accessible,
+                                                 &blockInfo) == HSA_STATUS_SUCCESS) {
+    if (blockInfo.length != size || info.sizeInBytes != size) {
+      for (int i = 0; i < num_agents; i++) union_agents.push_back(agents[i].handle);
+      for (int i = 0; i < agent_count; i++) union_agents.push_back(accessible[i].handle);
+      std::sort(union_agents.begin(), union_agents.end());
+      const auto& last = std::unique(union_agents.begin(), union_agents.end());
+      union_agents.erase(last, union_agents.end());
+
+      agents = reinterpret_cast<hsa_agent_t*>(&union_agents[0]);
+      num_agents = union_agents.size();
+      size = blockInfo.length;
+      ptr = blockInfo.base;
+    }
+  }
+
   bool cpu_in_list = false;
 
   std::set<GpuAgentInt*> whitelist_gpus;
@@ -492,6 +522,8 @@ hsa_status_t MemoryRegion::AllowAccess(uint32_t num_agents,
     }
   }
 
+  lock.Release();
+
   for (GpuAgentInt* gpu : whitelist_gpus) {
     gpu->InitDma();
   }
@@ -603,9 +635,9 @@ void* MemoryRegion::BlockAllocator::alloc(size_t request_size, size_t& allocated
   assert(request_size <= block_size() && "BlockAllocator alloc request exceeds block size.");
 
   void* ret;
+  size_t bsize = block_size();
   hsa_status_t err = region_.Allocate(
-      block_size(), core::MemoryRegion::AllocateRestrict | core::MemoryRegion::AllocateDirect,
-      &ret);
+      bsize, core::MemoryRegion::AllocateRestrict | core::MemoryRegion::AllocateDirect, &ret);
   if (err != HSA_STATUS_SUCCESS)
     throw new ::AMD::hsa_exception(err, "MemoryRegion::BlockAllocator::alloc failed.");
   assert(ret != nullptr && "Region returned nullptr on success.");

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

@@ -339,7 +339,10 @@ hsa_status_t Runtime::FreeMemory(void* ptr) {
   size = it->second.size;
 
   // Imported fragments can't be released with FreeMemory.
-  if (region == nullptr) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  if (region == nullptr) {
+    assert(false && "Can't release imported memory with free.");
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
 
   allocation_map_.erase(it);