From 51670abc297dc41da70215a2f7dc4915c9e8a303 Mon Sep 17 00:00:00 2001
From: "Spruit, Neil R" <neil.r.spruit@intel.com>
Date: Wed, 2 Aug 2023 01:45:07 +0000
Subject: [PATCH] feature: Add Support for Making Memory Resident on P2P
 capable devices

Related-To: LOCI-4583

- Provide support for allowing a user to make memory resident on a peer
to peer connected device.

Signed-off-by: Spruit, Neil R <neil.r.spruit@intel.com>
---
 .../core/source/context/context_imp.cpp       | 11 +++-
 .../sources/context/test_context.cpp          | 60 +++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/level_zero/core/source/context/context_imp.cpp b/level_zero/core/source/context/context_imp.cpp
index aed9ec0719..657bd7206b 100644
--- a/level_zero/core/source/context/context_imp.cpp
+++ b/level_zero/core/source/context/context_imp.cpp
@@ -448,7 +448,16 @@ ze_result_t ContextImp::makeMemoryResident(ze_device_handle_t hDevice, void *ptr
         neoDevice->getRootDeviceIndex(),
         nullptr);
     if (allocation == nullptr) {
-        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
+        NEO::SvmAllocationData *allocData = nullptr;
+        DriverHandleImp *driverHandleImp = static_cast<DriverHandleImp *>(this->driverHandle);
+        bool foundBuffer = driverHandleImp->findAllocationDataForRange(ptr, size, &allocData);
+        if (foundBuffer) {
+            uintptr_t alignedPtr = reinterpret_cast<uintptr_t>(ptr);
+            allocation = driverHandleImp->getPeerAllocation(device, allocData, ptr, &alignedPtr, nullptr);
+        }
+        if (allocation == nullptr) {
+            return ZE_RESULT_ERROR_INVALID_ARGUMENT;
+        }
     }
 
     NEO::MemoryOperationsHandler *memoryOperationsIface = neoDevice->getRootDeviceEnvironment().memoryOperationsInterface.get();
diff --git a/level_zero/core/test/unit_tests/sources/context/test_context.cpp b/level_zero/core/test/unit_tests/sources/context/test_context.cpp
index 96894c748c..8f32668cb1 100644
--- a/level_zero/core/test/unit_tests/sources/context/test_context.cpp
+++ b/level_zero/core/test/unit_tests/sources/context/test_context.cpp
@@ -195,6 +195,66 @@ TEST_F(MultiDeviceContextTests,
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
 }
 
+TEST_F(MultiDeviceContextTests,
+       GivenDeviceMemoryWhenMakeResidentCalledOnPeerDeviceThenSuccessReturned) {
+    ze_context_handle_t hContext;
+    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
+
+    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    driverHandle->devices[0]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
+        std::make_unique<NEO::MockMemoryOperations>();
+    driverHandle->devices[1]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[1]->memoryOperationsInterface =
+        std::make_unique<NEO::MockMemoryOperations>();
+
+    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
+
+    const size_t size = 4096;
+    void *ptr = nullptr;
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    res = context->allocDeviceMem(driverHandle->devices[0],
+                                  &deviceDesc,
+                                  size,
+                                  0,
+                                  &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    res = contextImp->makeMemoryResident(driverHandle->devices[1], ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    res = contextImp->freeMem(ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    res = contextImp->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+}
+
+TEST_F(MultiDeviceContextTests,
+       GivenInvalidDeviceMemoryWhenMakeResidentCalledOnPeerDeviceThenSuccessReturned) {
+    ze_context_handle_t hContext;
+    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
+
+    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    driverHandle->devices[0]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
+        std::make_unique<NEO::MockMemoryOperations>();
+    driverHandle->devices[1]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[1]->memoryOperationsInterface =
+        std::make_unique<NEO::MockMemoryOperations>();
+
+    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
+
+    const size_t size = 4096;
+    void *ptr = reinterpret_cast<void *>(0x1234);
+
+    res = contextImp->makeMemoryResident(driverHandle->devices[1], ptr, size);
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);
+
+    res = contextImp->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+}
+
 TEST_F(MultiDeviceContextTests,
        whenMappingReservedMemoryOnPhysicalMemoryOnMultiDeviceThenSuccessReturned) {
     ze_context_handle_t hContext;