From 235e167565eafa33605520cbdd30b0d47794a901 Mon Sep 17 00:00:00 2001
From: Jaime Arteaga <jaime.a.arteaga.molina@intel.com>
Date: Thu, 7 Jul 2022 07:37:56 +0000
Subject: [PATCH] Correctly return number of IPC handles in zexMemGetIpcHandles

Signed-off-by: Jaime Arteaga <jaime.a.arteaga.molina@intel.com>
---
 .../core/source/context/context_imp.cpp       |  9 ++++--
 .../unit_tests/sources/memory/test_memory.cpp | 28 +++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/level_zero/core/source/context/context_imp.cpp b/level_zero/core/source/context/context_imp.cpp
index 6fdef04100..4d14165d27 100644
--- a/level_zero/core/source/context/context_imp.cpp
+++ b/level_zero/core/source/context/context_imp.cpp
@@ -436,12 +436,17 @@ ze_result_t ContextImp::getIpcMemHandles(const void *ptr,
     if (allocData) {
         auto alloc = allocData->gpuAllocations.getDefaultGraphicsAllocation();
         uint32_t numHandles = alloc->getNumHandles();
-        if (pIpcHandles == nullptr) {
+        UNRECOVERABLE_IF(numIpcHandles == nullptr);
+
+        if (*numIpcHandles == 0 || *numIpcHandles > numHandles) {
             *numIpcHandles = numHandles;
+        }
+
+        if (pIpcHandles == nullptr) {
             return ZE_RESULT_SUCCESS;
         }
 
-        for (uint32_t i = 0; i < numHandles; i++) {
+        for (uint32_t i = 0; i < *numIpcHandles; i++) {
             int handle = static_cast<int>(allocData->gpuAllocations.getDefaultGraphicsAllocation()->peekInternalHandle(this->driverHandle->getMemoryManager(), i));
             memcpy_s(reinterpret_cast<void *>(pIpcHandles[i].data),
                      sizeof(ze_ipc_mem_handle_t),
diff --git a/level_zero/core/test/unit_tests/sources/memory/test_memory.cpp b/level_zero/core/test/unit_tests/sources/memory/test_memory.cpp
index 809a51d684..4eb6c7670e 100644
--- a/level_zero/core/test/unit_tests/sources/memory/test_memory.cpp
+++ b/level_zero/core/test/unit_tests/sources/memory/test_memory.cpp
@@ -74,6 +74,34 @@ TEST_F(MemoryExportImportImplicitScalingTest,
     EXPECT_EQ(ZE_RESULT_SUCCESS, result);
 }
 
+TEST_F(MemoryExportImportImplicitScalingTest,
+       givenCallToGetIpcHandleWithDeviceAllocationThenNumIpcHandlesIsUpdatedAlways) {
+    size_t size = 10;
+    size_t alignment = 1u;
+    void *ptr = nullptr;
+
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    ze_result_t result = context->allocDeviceMem(device->toHandle(),
+                                                 &deviceDesc,
+                                                 size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_NE(nullptr, ptr);
+
+    uint32_t numIpcHandles = 0;
+    result = context->getIpcMemHandles(ptr, &numIpcHandles, nullptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_EQ(numIpcHandles, 2u);
+
+    numIpcHandles *= 4;
+    std::vector<ze_ipc_mem_handle_t> ipcHandles(numIpcHandles);
+    result = context->getIpcMemHandles(ptr, &numIpcHandles, ipcHandles.data());
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_EQ(numIpcHandles, 2u);
+
+    result = context->freeMem(ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+}
+
 TEST_F(MemoryExportImportImplicitScalingTest,
        whenCallingOpenIpcHandlesWithIpcHandleThenDeviceAllocationIsReturned) {
     size_t size = 10;