From 09c775f34782c5b7d7b70283b3fa362988cf9f53 Mon Sep 17 00:00:00 2001
From: Mateusz Jablonski <mateusz.jablonski@intel.com>
Date: Tue, 15 Oct 2019 13:58:22 +0200
Subject: [PATCH] Add method to get total number of devices in context

Related-To: NEO-3691

Change-Id: I6e468118f99967a40dd8dce30013ada2142ca727
Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>
---
 runtime/context/context.cpp          |  8 ++++++++
 runtime/context/context.h            |  1 +
 unit_tests/context/context_tests.cpp | 23 +++++++++++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/runtime/context/context.cpp b/runtime/context/context.cpp
index 8749573908..20df1d2978 100644
--- a/runtime/context/context.cpp
+++ b/runtime/context/context.cpp
@@ -249,6 +249,14 @@ size_t Context::getNumDevices() const {
     return devices.size();
 }
 
+size_t Context::getTotalNumDevices() const {
+    size_t numAvailableDevices = 0u;
+    for (auto &device : devices) {
+        numAvailableDevices += device->getNumAvailableDevices();
+    }
+    return numAvailableDevices;
+}
+
 Device *Context::getDevice(size_t deviceOrdinal) {
     return (Device *)devices[deviceOrdinal];
 }
diff --git a/runtime/context/context.h b/runtime/context/context.h
index d3ce432a37..7ab078f25a 100644
--- a/runtime/context/context.h
+++ b/runtime/context/context.h
@@ -74,6 +74,7 @@ class Context : public BaseObject<_cl_context> {
                                     cl_image_format *imageFormats, cl_uint *numImageFormats);
 
     size_t getNumDevices() const;
+    size_t getTotalNumDevices() const;
     Device *getDevice(size_t deviceOrdinal);
 
     MemoryManager *getMemoryManager() {
diff --git a/unit_tests/context/context_tests.cpp b/unit_tests/context/context_tests.cpp
index 0d9f048c98..6bebb6e99b 100644
--- a/unit_tests/context/context_tests.cpp
+++ b/unit_tests/context/context_tests.cpp
@@ -13,6 +13,7 @@
 #include "runtime/helpers/options.h"
 #include "runtime/sharings/sharing.h"
 #include "unit_tests/fixtures/platform_fixture.h"
+#include "unit_tests/helpers/variable_backup.h"
 #include "unit_tests/mocks/mock_context.h"
 #include "unit_tests/mocks/mock_deferred_deleter.h"
 #include "unit_tests/mocks/mock_device.h"
@@ -314,6 +315,28 @@ TEST(Context, givenFtrSvmFalseWhenContextIsCreatedThenSVMAllocsManagerIsNotCreat
     EXPECT_EQ(nullptr, svmManager);
 }
 
+TEST(MultiDeviceContextTest, givenContextWithMultipleDevicesWhenGettingTotalNumberOfDevicesThenNumberOfAllAvailableDevicesIsReturned) {
+    DebugManagerStateRestore restorer;
+    const uint32_t numDevices = 2u;
+    const uint32_t numSubDevices = 3u;
+    VariableBackup<size_t> numDevicesBackup(&numPlatformDevices);
+    numDevicesBackup = numDevices;
+    DebugManager.flags.CreateMultipleSubDevices.set(numSubDevices);
+    platform()->initialize();
+    auto device0 = platform()->getDevice(0);
+    auto device1 = platform()->getDevice(1);
+    cl_device_id clDevices[2]{device0, device1};
+
+    DeviceVector deviceVector(clDevices, 2);
+    cl_int retVal = CL_OUT_OF_HOST_MEMORY;
+    auto context = std::unique_ptr<Context>(Context::create<Context>(nullptr, deviceVector, nullptr, nullptr, retVal));
+    EXPECT_EQ(CL_SUCCESS, retVal);
+    EXPECT_EQ(numSubDevices, device0->getNumAvailableDevices());
+    EXPECT_EQ(numSubDevices, device1->getNumAvailableDevices());
+    EXPECT_EQ(numDevices, context->getNumDevices());
+    EXPECT_EQ(numDevices * numSubDevices, context->getTotalNumDevices());
+}
+
 class ContextWithAsyncDeleterTest : public ::testing::WithParamInterface<bool>,
                                     public ::testing::Test {
   public: