SWDEV-483895 - Use ROCr to reserve memory for HMM (#590)

rocclr/device/rocm/rocdevice.cpp

@@ -2199,6 +2199,27 @@ void Device::deviceVmemRelease(uint64_t mem_handle) const {
   }
 }
 
+void* Device::reserveMemory(size_t size, size_t alignment) const {
+  void* ptr = nullptr;
+  // Reserves non registered VA memory using HSA APIs.
+  hsa_status_t status = hsa_amd_vmem_address_reserve_align(&ptr, size, 0, alignment,
+                                                           HSA_AMD_VMEM_ADDRESS_NO_REGISTER);
+  ClPrint(amd::LOG_DEBUG, amd::LOG_MEM, "Reserve hsa device memory %p, size 0x%zx", ptr, size);
+  if (status != HSA_STATUS_SUCCESS) {
+    LogError("Fail to reserve memory");
+    return nullptr;
+  }
+  return ptr;
+}
+
+void Device::releaseMemory(void* ptr, size_t size) const {
+  hsa_status_t hsa_status = hsa_amd_vmem_address_free(ptr, size);
+  ClPrint(amd::LOG_DEBUG, amd::LOG_MEM, "Free hsa reserved memory %p", ptr);
+  if (hsa_status != HSA_STATUS_SUCCESS) {
+    LogError("hsa_amd_vmem_address_free failed \n");
+  }
+}
+
 void* Device::deviceLocalAlloc(size_t size, bool atomics, bool pseudo_fine_grain,
                                bool contiguous) const {
 

rocclr/device/rocm/rocdevice.hpp

@@ -415,7 +415,8 @@ class Device : public NullDevice {
   uint64_t deviceVmemAlloc(size_t size, uint64_t flags) const;
   void* deviceLocalAlloc(size_t size, bool atomics = false, bool pseudo_fine_grain=false,
                          bool contiguous = false) const;
-
+  void* reserveMemory(size_t size, size_t alignment) const;
+  void releaseMemory(void* ptr, size_t size) const;
   void memFree(void* ptr, size_t size) const;
 
   virtual void* svmAlloc(amd::Context& context, size_t size, size_t alignment,

rocclr/device/rocm/rocmemory.cpp

@@ -658,10 +658,8 @@ void Buffer::destroy() {
       if (isFineGrain) {
         if (memFlags & CL_MEM_ALLOC_HOST_PTR) {
           if (dev().info().hmmSupported_) {
-            // AMD HMM path. Destroy system memory
-            if (!(amd::Os::releaseMemory(deviceMemory_, size()))) {
-              ClPrint(amd::LOG_DEBUG, amd::LOG_MEM, "[ROCClr] munmap failed \n");
-            }
+            // AMD HMM path. Release reserved system memory
+            dev().releaseMemory(deviceMemory_, size());
           } else {
             dev().memFree(deviceMemory_, size());
           }
@@ -809,9 +807,8 @@ bool Buffer::create(bool alloc_local) {
       if (isFineGrain) {
         if (memFlags & CL_MEM_ALLOC_HOST_PTR) {
           if (dev().info().hmmSupported_) {
-            // AMD HMM path. Just allocate system memory and KFD will manage it
-            deviceMemory_ =  amd::Os::reserveMemory(
-                0, size(), amd::Os::pageSize(), amd::Os::MEM_PROT_RW);
+            // AMD HMM path. ROCr allocates system memory and KFD will manage it
+            deviceMemory_ = dev().reserveMemory(size(), amd::Os::pageSize());
             if (deviceMemory_ == NULL) {
               return false;
             }
@@ -1069,7 +1066,6 @@ bool Buffer::GetFDHandleForMem(void* dev_ptr, size_t size, bool vmm, void* handl
       return false;
     }
   }
-  
   if (dmabuffd <= 0) {
     LogPrintfError("Invalid file descriptor handle: %d returned", dmabuffd);
     return false;