From e096bf881e3fdaf5dad37ec876e92a5550c3f68d Mon Sep 17 00:00:00 2001
From: Bartosz Dunajski <bartosz.dunajski@intel.com>
Date: Tue, 5 Oct 2021 13:13:42 +0000
Subject: [PATCH] Fix assigning SubDevices to Program

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>
---
 opencl/source/program/program.cpp             | 10 ++---
 .../unit_test/device/sub_device_tests.cpp     | 41 +++++++++++++++++++
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/opencl/source/program/program.cpp b/opencl/source/program/program.cpp
index 8d79a83378..ee900e7b72 100644
--- a/opencl/source/program/program.cpp
+++ b/opencl/source/program/program.cpp
@@ -54,12 +54,10 @@ Program::Program(Context *context, bool isBuiltIn, const ClDeviceVector &clDevic
             maxRootDeviceIndex = device->getRootDeviceIndex();
         }
         deviceBuildInfos[device] = {};
-        if (device->getNumGenericSubDevices() > 1) {
-            for (auto i = 0u; i < device->getNumGenericSubDevices(); i++) {
-                auto subDevice = device->getNearestGenericSubDevice(i);
-                if (isDeviceAssociated(*subDevice)) {
-                    deviceBuildInfos[device].associatedSubDevices.push_back(subDevice);
-                }
+        for (auto i = 0u; i < device->getNumSubDevices(); i++) {
+            auto subDevice = device->getSubDevice(i);
+            if (isDeviceAssociated(*subDevice)) {
+                deviceBuildInfos[device].associatedSubDevices.push_back(subDevice);
             }
         }
     }
diff --git a/opencl/test/unit_test/device/sub_device_tests.cpp b/opencl/test/unit_test/device/sub_device_tests.cpp
index cb6d972c49..be0302a12d 100644
--- a/opencl/test/unit_test/device/sub_device_tests.cpp
+++ b/opencl/test/unit_test/device/sub_device_tests.cpp
@@ -19,7 +19,9 @@
 #include "opencl/source/cl_device/cl_device.h"
 #include "opencl/test/unit_test/libult/ult_command_stream_receiver.h"
 #include "opencl/test/unit_test/mocks/mock_cl_device.h"
+#include "opencl/test/unit_test/mocks/mock_context.h"
 #include "opencl/test/unit_test/mocks/mock_platform.h"
+#include "opencl/test/unit_test/mocks/mock_program.h"
 #include "test.h"
 
 using namespace NEO;
@@ -876,6 +878,45 @@ HWTEST2_F(EngineInstancedDeviceTests, givenEngineInstancedDeviceWhenProgrammingC
     EXPECT_TRUE(cfeState->getSingleSliceDispatchCcsMode());
 }
 
+HWTEST_F(EngineInstancedDeviceTests, givenEngineInstancedDeviceWhenCreatingProgramThenAssignAllSubDevices) {
+    constexpr uint32_t genericDevicesCount = 2;
+    constexpr uint32_t ccsCount = 2;
+
+    if (!createDevices(genericDevicesCount, ccsCount)) {
+        GTEST_SKIP();
+    }
+
+    const char *source = "text";
+    size_t sourceSize = strlen(source);
+
+    auto clRootDevice = std::make_unique<ClDevice>(*rootDevice, nullptr);
+    auto clSubDevice = clRootDevice->getSubDevice(0);
+    auto clSubSubDevice0 = clSubDevice->getSubDevice(0);
+    auto clSubSubDevice1 = clSubDevice->getSubDevice(1);
+
+    cl_device_id device_ids[] = {clSubDevice, clSubSubDevice0, clSubSubDevice1};
+    ClDeviceVector deviceVector{device_ids, 3};
+    MockContext context(deviceVector);
+
+    cl_int retVal = CL_INVALID_PROGRAM;
+    auto program = std::unique_ptr<MockProgram>(Program::create<MockProgram>(
+        &context,
+        1,
+        &source,
+        &sourceSize,
+        retVal));
+
+    ASSERT_NE(nullptr, program.get());
+    ASSERT_EQ(CL_SUCCESS, retVal);
+
+    ASSERT_TRUE(program->deviceBuildInfos.find(clSubDevice) != program->deviceBuildInfos.end());
+
+    auto &associatedSubDevices = program->deviceBuildInfos[clSubDevice].associatedSubDevices;
+    ASSERT_EQ(2u, associatedSubDevices.size());
+    EXPECT_EQ(clSubSubDevice0, associatedSubDevices[0]);
+    EXPECT_EQ(clSubSubDevice1, associatedSubDevices[1]);
+}
+
 TEST(SubDevicesTest, whenInitializeRootCsrThenDirectSubmissionIsNotInitialized) {
     auto device = std::make_unique<MockDevice>();
     device->initializeRootCommandStreamReceiver();