From 0650b96999aef4cf4a2e0a2973250d10dc5cb062 Mon Sep 17 00:00:00 2001
From: Jaroslaw Warchulski <jaroslaw.warchulski@intel.com>
Date: Mon, 24 Mar 2025 11:42:27 +0000
Subject: [PATCH] fix: set proper allocation in MemObj::getMemObjectInfo

Related-To: NEO-12585
Signed-off-by: Jaroslaw Warchulski <jaroslaw.warchulski@intel.com>
---
 opencl/source/mem_obj/mem_obj.cpp             |  7 ++--
 .../mem_obj/get_mem_object_info_tests.cpp     | 36 ++++++++++++++++++-
 .../os_interface/linux/device_os_tests.cpp    |  7 ++--
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/opencl/source/mem_obj/mem_obj.cpp b/opencl/source/mem_obj/mem_obj.cpp
index e3b6830a23..75bbc5676b 100644
--- a/opencl/source/mem_obj/mem_obj.cpp
+++ b/opencl/source/mem_obj/mem_obj.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -134,7 +134,8 @@ cl_int MemObj::getMemObjectInfo(cl_mem_info paramName,
     cl_mem clAssociatedMemObject = static_cast<cl_mem>(this->associatedMemObject);
     cl_context ctx = nullptr;
     uint64_t internalHandle = 0llu;
-    auto allocation = getMultiGraphicsAllocation().getDefaultGraphicsAllocation();
+    auto rootDeviceIndex = context->getDevice(0)->getRootDeviceIndex();
+    auto allocation = multiGraphicsAllocation.getGraphicsAllocation(rootDeviceIndex);
     cl_bool usesCompression;
 
     switch (paramName) {
@@ -204,7 +205,7 @@ cl_int MemObj::getMemObjectInfo(cl_mem_info paramName,
         break;
 
     case CL_MEM_ALLOCATION_HANDLE_INTEL: {
-        auto retVal = multiGraphicsAllocation.getDefaultGraphicsAllocation()->peekInternalHandle(this->memoryManager, internalHandle);
+        auto retVal = allocation->peekInternalHandle(this->memoryManager, internalHandle);
         if (retVal != 0) {
             return CL_OUT_OF_RESOURCES;
         }
diff --git a/opencl/test/unit_test/mem_obj/get_mem_object_info_tests.cpp b/opencl/test/unit_test/mem_obj/get_mem_object_info_tests.cpp
index fd3b246f59..21ac5e7705 100644
--- a/opencl/test/unit_test/mem_obj/get_mem_object_info_tests.cpp
+++ b/opencl/test/unit_test/mem_obj/get_mem_object_info_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -479,3 +479,37 @@ TEST_F(GetMemObjectInfo, GivenFailureOnGettingInternalHandleWhenGettingMemObject
         &sizeReturned);
     EXPECT_EQ(CL_OUT_OF_RESOURCES, retVal);
 }
+
+TEST_F(GetMemObjectInfo, GivenTwoRootDevicesInReverseOrderWhenGettingMemObjectInfoThenCorrectHandleIsReturned) {
+    auto device1 = std::make_unique<MockClDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr, 1));
+    auto device2 = std::make_unique<MockClDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr, 0));
+
+    ClDeviceVector devices;
+    devices.push_back(device1.get());
+    devices.push_back(device2.get());
+    auto context = MockContext(devices);
+
+    auto graphicsAllocation1 = new MockGraphicsAllocation();
+    auto graphicsAllocation2 = new MockGraphicsAllocation(1, nullptr, 0);
+    graphicsAllocation2->internalHandle = 0xf00d;
+    graphicsAllocation2->peekInternalHandleResult = 0;
+
+    MultiGraphicsAllocation multiGraphicsAllocation(1);
+    multiGraphicsAllocation.addAllocation(graphicsAllocation1);
+    multiGraphicsAllocation.addAllocation(graphicsAllocation2);
+
+    MemObj buffer(&context, CL_MEM_OBJECT_BUFFER, {}, CL_MEM_COPY_HOST_PTR, 0, 64, nullptr, nullptr, std::move(multiGraphicsAllocation), true, false, false);
+
+    size_t sizeReturned = 0;
+    uint64_t internalHandle = std::numeric_limits<uint64_t>::max();
+
+    auto retVal = buffer.getMemObjectInfo(
+        CL_MEM_ALLOCATION_HANDLE_INTEL,
+        sizeof(internalHandle),
+        &internalHandle,
+        &sizeReturned);
+
+    EXPECT_EQ(CL_SUCCESS, retVal);
+    EXPECT_EQ(sizeof(internalHandle), sizeReturned);
+    EXPECT_EQ(internalHandle, graphicsAllocation2->internalHandle);
+}
diff --git a/opencl/test/unit_test/os_interface/linux/device_os_tests.cpp b/opencl/test/unit_test/os_interface/linux/device_os_tests.cpp
index 0e6b3375fe..894a713f5a 100644
--- a/opencl/test/unit_test/os_interface/linux/device_os_tests.cpp
+++ b/opencl/test/unit_test/os_interface/linux/device_os_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -69,7 +69,7 @@ TEST(DeviceOsTest, WhenDeviceIsCreatedThenSimultaneousInteropsIsSupported) {
     EXPECT_TRUE(pDevice->simultaneousInterops == expected);
 }
 
-TEST(ApiOsTest, GivenUnupportedApiTokensWhenGettingInfoThenInvalidValueErrorIsReturned) {
+TEST(ApiOsTest, GivenUnsupportedApiTokensWhenGettingInfoThenInvalidValueErrorIsReturned) {
     MockContext context;
     MockBuffer buffer;
 
@@ -79,8 +79,11 @@ TEST(ApiOsTest, GivenUnupportedApiTokensWhenGettingInfoThenInvalidValueErrorIsRe
     EXPECT_EQ(CL_INVALID_VALUE, retVal);
 
     void *paramVal = nullptr;
+    buffer.context = &context;
     retVal = buffer.getMemObjectInfo(CL_MEM_D3D10_RESOURCE_KHR, sizeof(void *), paramVal, &size);
     EXPECT_EQ(CL_INVALID_VALUE, retVal);
+
+    buffer.context = nullptr;
 }
 
 TEST(ApiOsTest, GivenUnsupportedApiWhenGettingDispatchThenNullIsReturned) {