rocr: support reserving non-registered VA

Extend hsa_amd_vmem_address_reserve/hsa_amd_vmem_address_reserve_align
to support HSA_AMD_VMEM_ADDRESS_NO_REGISTER flag. This allocation can be
used to reserve virtual address ranges that can later be used by
hsa_amd_svm_attributes_set for SVM based memory allocations.


[ROCm/ROCR-Runtime commit: b3c48cc68c]

projects/rocr-runtime/runtime/hsa-runtime/core/inc/runtime.h

@@ -802,11 +802,14 @@ class Runtime {
   typedef void* ThunkHandle;
 
   struct AddressHandle {
-    AddressHandle() : size(0), use_count(0) {}
-    AddressHandle(size_t size) : size(size), use_count(0) {}
+    AddressHandle() : os_addr(nullptr), size(0), use_count(0), registered(false) {}
+    AddressHandle(void* addr, size_t _size, bool _registered) : os_addr(addr), size(_size), use_count(0), registered(_registered) {}
 
+    // Address returned by OS. May be different from user address when adjusted for alignment
+    void *os_addr;
     size_t size;
     int use_count;
+    bool registered;
   };
   std::map<const void*, AddressHandle> reserved_address_map_;  // Indexed by VA
 

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp

@@ -1343,7 +1343,10 @@ hsa_status_t hsa_amd_vmem_address_reserve(void** va, size_t size, uint64_t addre
   TRY;
   IS_OPEN();
   IS_ZERO(size);
-  IS_TRUE(core::Runtime::runtime_singleton_->VirtualMemApiSupported());
+
+  if (!(flags & HSA_AMD_VMEM_ADDRESS_NO_REGISTER))
+    IS_TRUE(core::Runtime::runtime_singleton_->VirtualMemApiSupported());
+
   return core::Runtime::runtime_singleton_->VMemoryAddressReserve(va, size, address, 0, flags);
   CATCH;
 }

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/runtime.cpp

@@ -3144,6 +3144,26 @@ hsa_status_t Runtime::VMemoryAddressReserve(void** va, size_t size, uint64_t add
 
   ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
 
+  if (flags & HSA_AMD_VMEM_ADDRESS_NO_REGISTER) {
+    size_t requested = size + alignment - sysconf(_SC_PAGE_SIZE);
+    auto mem = mmap(addr, requested, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE | MAP_NORESERVE, -1, 0);
+    if (mem == MAP_FAILED)
+      return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+
+    auto aligned = AlignUp(mem, alignment);
+
+    // Hint to enable THP for large host allocations which can help in performance gain
+    constexpr size_t kLargePageSize = 2*1024*1024;
+    if (size >= kLargePageSize) {
+      if (madvise(aligned, size, MADV_HUGEPAGE))
+        debug_warning(false && "madvise with MADV_HUGEPAGE failed");
+    }
+
+    reserved_address_map_[aligned] = AddressHandle(mem, size, false);
+    *va = aligned;
+    return HSA_STATUS_SUCCESS;
+  }
+
   memFlags.ui32.OnlyAddress = 1;
   memFlags.ui32.FixedAddress = 1;
 
@@ -3155,7 +3175,7 @@ hsa_status_t Runtime::VMemoryAddressReserve(void** va, size_t size, uint64_t add
       return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
   }
 
-  reserved_address_map_[addr] = AddressHandle(size);
+  reserved_address_map_[addr] = AddressHandle(addr, size, true);
   *va = addr;
   return HSA_STATUS_SUCCESS;
 }
@@ -3173,7 +3193,10 @@ hsa_status_t Runtime::VMemoryAddressFree(void* va, size_t size) {
 
   if (it->second.use_count > 0) return HSA_STATUS_ERROR_RESOURCE_FREE;
 
-  if (HSAKMT_CALL(hsaKmtFreeMemory(va, size)) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+  if (it->second.registered)
+    if (HSAKMT_CALL(hsaKmtFreeMemory(it->second.os_addr, size)) != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+  else
+    if (munmap(it->second.os_addr, size)) return HSA_STATUS_ERROR;
 
   reserved_address_map_.erase(it);
   return HSA_STATUS_SUCCESS;