fix: create graphicsAllocation per each rootDevice in unified sharing

Resolved: NEO-12585
Signed-off-by: Jaroslaw Warchulski <jaroslaw.warchulski@intel.com>

opencl/source/sharings/unified/unified_buffer.cpp

@@ -1,13 +1,11 @@
 /*
- * Copyright (C) 2019-2023 Intel Corporation
+ * Copyright (C) 2019-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
  */
 
-#include "unified_buffer.h"
-
-#include "shared/source/helpers/get_info.h"
+#include "opencl/source/sharings/unified/unified_buffer.h"
 
 #include "opencl/source/context/context.h"
 #include "opencl/source/mem_obj/buffer.h"
@@ -15,19 +13,12 @@
 using namespace NEO;
 
 Buffer *UnifiedBuffer::createSharedUnifiedBuffer(Context *context, cl_mem_flags flags, UnifiedSharingMemoryDescription extMem, cl_int *errcodeRet) {
-    ErrorCodeHelper errorCode(errcodeRet, CL_SUCCESS);
-
-    auto graphicsAllocation = UnifiedBuffer::createGraphicsAllocation(context, extMem, AllocationType::sharedBuffer);
-    if (!graphicsAllocation) {
-        errorCode.set(CL_INVALID_MEM_OBJECT);
+    auto multiGraphicsAllocation = UnifiedBuffer::createMultiGraphicsAllocation(context, extMem, nullptr, AllocationType::sharedBuffer, errcodeRet);
+    if (!multiGraphicsAllocation) {
         return nullptr;
     }
 
-    UnifiedSharingFunctions *sharingFunctions = context->getSharing<UnifiedSharingFunctions>();
-    auto sharingHandler = new UnifiedBuffer(sharingFunctions, extMem.type);
-    auto rootDeviceIndex = graphicsAllocation->getRootDeviceIndex();
-    auto multiGraphicsAllocation = MultiGraphicsAllocation(rootDeviceIndex);
-    multiGraphicsAllocation.addAllocation(graphicsAllocation);
+    auto sharingHandler = new UnifiedBuffer(context->getSharing<UnifiedSharingFunctions>(), extMem.type);
 
-    return Buffer::createSharedBuffer(context, flags, sharingHandler, std::move(multiGraphicsAllocation));
+    return Buffer::createSharedBuffer(context, flags, sharingHandler, std::move(*multiGraphicsAllocation));
 }

opencl/source/sharings/unified/unified_image.cpp

@@ -23,37 +23,32 @@ namespace NEO {
 
 Image *UnifiedImage::createSharedUnifiedImage(Context *context, cl_mem_flags flags, UnifiedSharingMemoryDescription description,
                                               const cl_image_format *imageFormat, const cl_image_desc *imageDesc, cl_int *errcodeRet) {
-    ErrorCodeHelper errorCode(errcodeRet, CL_SUCCESS);
-    UnifiedSharingFunctions *sharingFunctions = context->getSharing<UnifiedSharingFunctions>();
 
     auto *clSurfaceFormat = Image::getSurfaceFormatFromTable(flags, imageFormat, context->getDevice(0)->getHardwareInfo().capabilityTable.supportsOcl21Features);
     ImageInfo imgInfo = {};
     imgInfo.imgDesc = Image::convertDescriptor(*imageDesc);
     imgInfo.surfaceFormat = &clSurfaceFormat->surfaceFormat;
 
-    GraphicsAllocation *graphicsAllocation = createGraphicsAllocation(context, description, &imgInfo, AllocationType::sharedImage);
-    if (!graphicsAllocation) {
-        errorCode.set(CL_INVALID_MEM_OBJECT);
+    auto multiGraphicsAllocation = createMultiGraphicsAllocation(context, description, &imgInfo, AllocationType::sharedImage, errcodeRet);
+    if (!multiGraphicsAllocation) {
         return nullptr;
     }
 
-    swapGmm(graphicsAllocation, context, &imgInfo);
+    for (auto graphicsAllocation : multiGraphicsAllocation->getGraphicsAllocations()) {
+        swapGmm(graphicsAllocation, context, &imgInfo);
 
-    auto &memoryManager = *context->getMemoryManager();
-    if (graphicsAllocation->getDefaultGmm()->unifiedAuxTranslationCapable()) {
-        const auto &hwInfo = context->getDevice(0)->getHardwareInfo();
-        const auto &productHelper = context->getDevice(0)->getProductHelper();
-        graphicsAllocation->getDefaultGmm()->setCompressionEnabled(productHelper.isPageTableManagerSupported(hwInfo) ? memoryManager.mapAuxGpuVA(graphicsAllocation) : true);
+        auto &memoryManager = *context->getMemoryManager();
+        if (graphicsAllocation->getDefaultGmm()->unifiedAuxTranslationCapable()) {
+            const auto &hwInfo = context->getDevice(0)->getHardwareInfo();
+            const auto &productHelper = context->getDevice(0)->getProductHelper();
+            graphicsAllocation->getDefaultGmm()->setCompressionEnabled(productHelper.isPageTableManagerSupported(hwInfo) ? memoryManager.mapAuxGpuVA(graphicsAllocation) : true);
+        }
     }
 
-    const uint32_t baseMipmapIndex = 0u;
-    const uint32_t sharedMipmapsCount = imageDesc->num_mip_levels;
-    auto sharingHandler = new UnifiedImage(sharingFunctions, description.type);
-    auto multiGraphicsAllocation = MultiGraphicsAllocation(context->getDevice(0)->getRootDeviceIndex());
-    multiGraphicsAllocation.addAllocation(graphicsAllocation);
+    auto sharingHandler = new UnifiedImage(context->getSharing<UnifiedSharingFunctions>(), description.type);
 
-    return Image::createSharedImage(context, sharingHandler, McsSurfaceInfo{}, std::move(multiGraphicsAllocation), nullptr,
-                                    flags, 0, clSurfaceFormat, imgInfo, __GMM_NO_CUBE_MAP, baseMipmapIndex, sharedMipmapsCount, false);
+    return Image::createSharedImage(context, sharingHandler, McsSurfaceInfo{}, std::move(*multiGraphicsAllocation), nullptr,
+                                    flags, 0, clSurfaceFormat, imgInfo, __GMM_NO_CUBE_MAP, 0u, imageDesc->num_mip_levels, false);
 }
 
 } // namespace NEO

opencl/source/sharings/unified/unified_sharing.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2024 Intel Corporation
+ * Copyright (C) 2019-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -7,6 +7,7 @@
 
 #include "opencl/source/sharings/unified/unified_sharing.h"
 
+#include "shared/source/helpers/get_info.h"
 #include "shared/source/helpers/string.h"
 #include "shared/source/helpers/surface_format_info.h"
 #include "shared/source/helpers/timestamp_packet.h"
@@ -34,23 +35,36 @@ void UnifiedSharing::synchronizeObject(UpdateData &updateData) {
 void UnifiedSharing::releaseResource(MemObj *memObject, uint32_t rootDeviceIndex) {
 }
 
-GraphicsAllocation *UnifiedSharing::createGraphicsAllocation(Context *context, UnifiedSharingMemoryDescription description, AllocationType allocationType) {
-    return createGraphicsAllocation(context, description, nullptr, allocationType);
-}
-
-GraphicsAllocation *UnifiedSharing::createGraphicsAllocation(Context *context, UnifiedSharingMemoryDescription description, ImageInfo *imgInfo, AllocationType allocationType) {
+std::unique_ptr<MultiGraphicsAllocation> UnifiedSharing::createMultiGraphicsAllocation(Context *context, UnifiedSharingMemoryDescription description, ImageInfo *imgInfo, AllocationType allocationType, cl_int *errcodeRet) {
+    ErrorCodeHelper errorCode(errcodeRet, CL_SUCCESS);
     auto memoryManager = context->getMemoryManager();
 
-    if (description.type != UnifiedSharingHandleType::win32Nt && description.type != UnifiedSharingHandleType::linuxFd && description.type != UnifiedSharingHandleType::win32Shared)
+    if (description.type != UnifiedSharingHandleType::win32Nt && description.type != UnifiedSharingHandleType::linuxFd && description.type != UnifiedSharingHandleType::win32Shared) {
+        errorCode.set(CL_INVALID_MEM_OBJECT);
         return nullptr;
+    }
 
-    const AllocationProperties properties{context->getDevice(0)->getRootDeviceIndex(),
-                                          false, // allocateMemory
-                                          imgInfo,
-                                          allocationType,
-                                          context->getDeviceBitfieldForAllocation(context->getDevice(0)->getRootDeviceIndex())};
+    auto pRootDeviceIndices = &context->getRootDeviceIndices();
+    auto multiGraphicsAllocation = std::make_unique<MultiGraphicsAllocation>(context->getMaxRootDeviceIndex());
     MemoryManager::OsHandleData osHandleData{description.handle};
-    return memoryManager->createGraphicsAllocationFromSharedHandle(osHandleData, properties, false, false, true, nullptr);
+
+    for (auto &rootDeviceIndex : *pRootDeviceIndices) {
+        const AllocationProperties properties{rootDeviceIndex,
+                                              false, // allocateMemory
+                                              imgInfo,
+                                              allocationType,
+                                              context->getDeviceBitfieldForAllocation(rootDeviceIndex)};
+
+        auto graphicsAllocation = memoryManager->createGraphicsAllocationFromSharedHandle(osHandleData, properties, false, false, true, nullptr);
+        if (!graphicsAllocation) {
+            errorCode.set(CL_INVALID_MEM_OBJECT);
+            return nullptr;
+        }
+
+        multiGraphicsAllocation->addAllocation(graphicsAllocation);
+    }
+
+    return multiGraphicsAllocation;
 }
 
 template <>

opencl/source/sharings/unified/unified_sharing.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2023 Intel Corporation
+ * Copyright (C) 2019-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -10,8 +10,11 @@
 #include "opencl/source/sharings/sharing.h"
 #include "opencl/source/sharings/unified/unified_sharing_types.h"
 
+#include "CL/cl.h"
+
 namespace NEO {
 struct ImageInfo;
+class MultiGraphicsAllocation;
 
 class UnifiedSharingFunctions : public SharingFunctions {
   public:
@@ -32,8 +35,7 @@ class UnifiedSharing : public SharingHandler {
     void synchronizeObject(UpdateData &updateData) override;
     void releaseResource(MemObj *memObject, uint32_t rootDeviceIndex) override;
 
-    static GraphicsAllocation *createGraphicsAllocation(Context *context, UnifiedSharingMemoryDescription description, AllocationType allocationType);
-    static GraphicsAllocation *createGraphicsAllocation(Context *context, UnifiedSharingMemoryDescription description, ImageInfo *imgInfo, AllocationType allocationType);
+    static std::unique_ptr<MultiGraphicsAllocation> createMultiGraphicsAllocation(Context *context, UnifiedSharingMemoryDescription description, ImageInfo *imgInfo, AllocationType allocationType, cl_int *errcodeRet);
 
   private:
     UnifiedSharingFunctions *sharingFunctions;

opencl/test/unit_test/sharings/unified/unified_sharing_buffer_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2023 Intel Corporation
+ * Copyright (C) 2019-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -7,6 +7,7 @@
 
 #include "opencl/source/mem_obj/buffer.h"
 #include "opencl/source/sharings/unified/unified_buffer.h"
+#include "opencl/test/unit_test/fixtures/multi_root_device_fixture.h"
 #include "opencl/test/unit_test/sharings/unified/unified_sharing_fixtures.h"
 #include "opencl/test/unit_test/sharings/unified/unified_sharing_mocks.h"
 
@@ -72,3 +73,17 @@ TEST_F(UnifiedSharingBufferTestsWithMemoryManager, givenValidContextAndMemoryMan
     auto buffer = std::unique_ptr<Buffer>(UnifiedBuffer::createSharedUnifiedBuffer(context.get(), flags, desc, &retVal));
     ASSERT_EQ(CL_SUCCESS, retVal);
 }
+
+using UnifiedSharingBufferTestsWithMultiRootDevice = MultiRootDeviceFixture;
+
+TEST_F(UnifiedSharingBufferTestsWithMultiRootDevice, givenValidContextWithMultiRootDeviceWhenCreatingBufferFromSharedHandleThenGraphicsAllocationsAreCorrectlySet) {
+    cl_mem_flags flags{};
+    UnifiedSharingMemoryDescription desc{};
+    desc.handle = reinterpret_cast<void *>(0x1234);
+    desc.type = UnifiedSharingHandleType::win32Nt;
+    cl_int retVal{};
+    auto buffer = std::unique_ptr<Buffer>(UnifiedBuffer::createSharedUnifiedBuffer(context.get(), flags, desc, &retVal));
+    EXPECT_NE(nullptr, buffer->getGraphicsAllocation(device1->getRootDeviceIndex()));
+    EXPECT_NE(nullptr, buffer->getGraphicsAllocation(device2->getRootDeviceIndex()));
+    ASSERT_EQ(CL_SUCCESS, retVal);
+}

opencl/test/unit_test/sharings/unified/unified_sharing_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2024 Intel Corporation
+ * Copyright (C) 2019-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -177,7 +177,7 @@ struct UnifiedSharingCreateAllocationTests : UnifiedSharingTestsWithMemoryManage
     };
 
     struct MockSharingHandler : UnifiedSharing {
-        using UnifiedSharing::createGraphicsAllocation;
+        using UnifiedSharing::createMultiGraphicsAllocation;
     };
 
     void SetUp() override {
@@ -195,7 +195,7 @@ TEST_F(UnifiedSharingCreateAllocationTests, givenWindowsNtHandleWhenCreateGraphi
     desc.handle = reinterpret_cast<void *>(0x1234);
     desc.type = UnifiedSharingHandleType::win32Nt;
     AllocationType allocationType = AllocationType::sharedImage;
-    MockSharingHandler::createGraphicsAllocation(this->context.get(), desc, allocationType);
+    MockSharingHandler::createMultiGraphicsAllocation(this->context.get(), desc, nullptr, allocationType, nullptr);
 
     EXPECT_TRUE(memoryManager->createFromSharedHandleCalled);
     EXPECT_EQ(toOsHandle(desc.handle), memoryManager->handle);
@@ -206,7 +206,7 @@ TEST_F(UnifiedSharingCreateAllocationTests, givenWindowsSharedHandleWhenCreateGr
     desc.handle = reinterpret_cast<void *>(0x1234);
     desc.type = UnifiedSharingHandleType::win32Shared;
     AllocationType allocationType = AllocationType::sharedImage;
-    MockSharingHandler::createGraphicsAllocation(this->context.get(), desc, allocationType);
+    MockSharingHandler::createMultiGraphicsAllocation(this->context.get(), desc, nullptr, allocationType, nullptr);
 
     EXPECT_TRUE(memoryManager->createFromSharedHandleCalled);
     EXPECT_EQ(toOsHandle(desc.handle), memoryManager->handle);
@@ -219,7 +219,7 @@ TEST_F(UnifiedSharingCreateAllocationTests, givenLinuxSharedHandleWhenCreateGrap
     desc.handle = reinterpret_cast<void *>(0x1234);
     desc.type = UnifiedSharingHandleType::linuxFd;
     AllocationType allocationType = AllocationType::sharedImage;
-    MockSharingHandler::createGraphicsAllocation(this->context.get(), desc, allocationType);
+    MockSharingHandler::createMultiGraphicsAllocation(this->context.get(), desc, nullptr, allocationType, nullptr);
 
     EXPECT_TRUE(memoryManager->createFromSharedHandleCalled);
     EXPECT_EQ(toOsHandle(desc.handle), memoryManager->handle);