From 1c5b11fd045f0cbe656fcb8b330e29d125d67734 Mon Sep 17 00:00:00 2001
From: Christophe Paquot <christophe.paquot@amd.com>
Date: Mon, 18 Apr 2022 14:28:34 -0700
Subject: [PATCH] SWDEV-322620 - Virtual Memory Management

Implement hipMemAddressFree and hipMemAddressReserve

Change-Id: I1b09e433db5c5c6d6b9fe2a7a15fe2c09f1e1874


[ROCm/clr commit: fd17b53ea464f78a19f3b085865a754f36169cb4]
---
 projects/clr/hipamd/src/hip_vm.cpp | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/projects/clr/hipamd/src/hip_vm.cpp b/projects/clr/hipamd/src/hip_vm.cpp
index 0f21c2a121..0f9c55a3fc 100644
--- a/projects/clr/hipamd/src/hip_vm.cpp
+++ b/projects/clr/hipamd/src/hip_vm.cpp
@@ -29,6 +29,10 @@ hipError_t hipMemAddressFree(void* devPtr, size_t size) {
     HIP_RETURN(hipErrorInvalidValue);
   }
 
+  for (auto& dev: g_devices) {
+    dev->devices()[0]->virtualFree(devPtr);
+  }
+
   HIP_RETURN(hipSuccess);
 }
 
@@ -40,6 +44,25 @@ hipError_t hipMemAddressReserve(void** ptr, size_t size, size_t alignment, void*
     HIP_RETURN(hipErrorInvalidValue);
   }
 
+  *ptr = nullptr;
+
+  void* startAddress = addr;
+
+  for (auto& dev : g_devices) {
+    *ptr = dev->devices()[0]->virtualAlloc(startAddress, size, alignment);
+
+    // if addr==0 we generate the va and use it for other devices
+    if (startAddress == nullptr) {
+      startAddress = *ptr;
+    } else if (*ptr != startAddress) {
+      // if we cannot reserve the same VA on other devices, just fail
+      for (auto& d : g_devices) {
+        if (d == dev) HIP_RETURN(hipErrorOutOfMemory);
+        d->devices()[0]->virtualFree(startAddress);
+      }
+    }
+  }
+
   HIP_RETURN(hipSuccess);
 }