From 015595a3f81d1e114f5f20db43d2773a81c24968 Mon Sep 17 00:00:00 2001
From: Bartosz Dunajski <bartosz.dunajski@intel.com>
Date: Tue, 11 May 2021 18:43:28 +0000
Subject: [PATCH] Remove root device index duplicates from Event pool
 allocation

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>
---
 level_zero/core/source/event/event.cpp        | 12 +++----
 .../unit_tests/sources/event/test_event.cpp   | 34 +++++++++++++++++++
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/level_zero/core/source/event/event.cpp b/level_zero/core/source/event/event.cpp
index 1cf908c342..e0bcbd29bd 100644
--- a/level_zero/core/source/event/event.cpp
+++ b/level_zero/core/source/event/event.cpp
@@ -26,13 +26,12 @@
 #include "level_zero/core/source/driver/driver_handle_imp.h"
 #include "level_zero/tools/source/metrics/metric.h"
 
-#include <queue>
-#include <unordered_map>
+#include <set>
 
 namespace L0 {
 
 ze_result_t EventPoolImp::initialize(DriverHandle *driver, Context *context, uint32_t numDevices, ze_device_handle_t *phDevices, uint32_t numEvents) {
-    std::vector<uint32_t> rootDeviceIndices;
+    std::set<uint32_t> rootDeviceIndices;
     uint32_t maxRootDeviceIndex = 0u;
 
     void *eventPoolPtr = nullptr;
@@ -63,7 +62,7 @@ ze_result_t EventPoolImp::initialize(DriverHandle *driver, Context *context, uin
         }
 
         devices.push_back(eventDevice);
-        rootDeviceIndices.push_back(eventDevice->getNEODevice()->getRootDeviceIndex());
+        rootDeviceIndices.insert(eventDevice->getNEODevice()->getRootDeviceIndex());
         if (maxRootDeviceIndex < eventDevice->getNEODevice()->getRootDeviceIndex()) {
             maxRootDeviceIndex = eventDevice->getNEODevice()->getRootDeviceIndex();
         }
@@ -71,10 +70,11 @@ ze_result_t EventPoolImp::initialize(DriverHandle *driver, Context *context, uin
 
     eventPoolAllocations = std::make_unique<NEO::MultiGraphicsAllocation>(maxRootDeviceIndex);
 
-    NEO::AllocationProperties allocationProperties{rootDeviceIndices.at(0), alignedSize, allocationType, systemMemoryBitfield};
+    NEO::AllocationProperties allocationProperties{*rootDeviceIndices.begin(), alignedSize, allocationType, systemMemoryBitfield};
     allocationProperties.alignment = eventAlignment;
 
-    eventPoolPtr = driver->getMemoryManager()->createMultiGraphicsAllocationInSystemMemoryPool(rootDeviceIndices,
+    std::vector<uint32_t> rootDeviceIndicesVector = {rootDeviceIndices.begin(), rootDeviceIndices.end()};
+    eventPoolPtr = driver->getMemoryManager()->createMultiGraphicsAllocationInSystemMemoryPool(rootDeviceIndicesVector,
                                                                                                allocationProperties,
                                                                                                *eventPoolAllocations);
 
diff --git a/level_zero/core/test/unit_tests/sources/event/test_event.cpp b/level_zero/core/test/unit_tests/sources/event/test_event.cpp
index a3aac8dd7a..eb34fa5fce 100644
--- a/level_zero/core/test/unit_tests/sources/event/test_event.cpp
+++ b/level_zero/core/test/unit_tests/sources/event/test_event.cpp
@@ -598,6 +598,40 @@ TEST_F(EventPoolCreateMultiDevice, whenCreatingEventPoolWithMultipleDevicesThenE
     delete[] devices;
 }
 
+TEST_F(EventPoolCreateMultiDevice, whenCreatingEventPoolWithMultipleDevicesThenDontDuplicateRootDeviceIndices) {
+    ze_event_pool_desc_t eventPoolDesc = {};
+    eventPoolDesc.stype = ZE_STRUCTURE_TYPE_EVENT_POOL_DESC;
+    eventPoolDesc.flags = ZE_EVENT_POOL_FLAG_HOST_VISIBLE;
+    eventPoolDesc.count = 32;
+
+    uint32_t deviceCount = 1;
+    ze_device_handle_t rootDeviceHandle;
+
+    ze_result_t result = zeDeviceGet(driverHandle.get(), &deviceCount, &rootDeviceHandle);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    deviceCount = 0;
+    result = zeDeviceGetSubDevices(rootDeviceHandle, &deviceCount, nullptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_TRUE(deviceCount >= 2);
+
+    auto subDeviceHandle = std::make_unique<ze_device_handle_t[]>(deviceCount);
+    result = zeDeviceGetSubDevices(rootDeviceHandle, &deviceCount, subDeviceHandle.get());
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    std::unique_ptr<L0::EventPool> eventPool(EventPool::create(driverHandle.get(),
+                                                               context,
+                                                               deviceCount,
+                                                               subDeviceHandle.get(),
+                                                               &eventPoolDesc));
+    EXPECT_NE(nullptr, eventPool);
+
+    auto allocation = &eventPool->getAllocation();
+    EXPECT_NE(nullptr, allocation);
+
+    EXPECT_EQ(allocation->getGraphicsAllocations().size(), 1u);
+}
+
 TEST_F(EventPoolCreateMultiDevice, whenCreatingEventPoolWithNoDevicesThenEventPoolCreateSucceedsAndAllDeviceAreUsed) {
     ze_event_pool_desc_t eventPoolDesc = {};
     eventPoolDesc.stype = ZE_STRUCTURE_TYPE_EVENT_POOL_DESC;