From 3909bd0fea3e9c205535faab3830e8e29946b598 Mon Sep 17 00:00:00 2001
From: Zbigniew Zdanowicz <zbigniew.zdanowicz@intel.com>
Date: Fri, 9 Apr 2021 17:52:17 +0000
Subject: [PATCH] Add setter function to event class

Related-To: NEO-5244

Signed-off-by: Zbigniew Zdanowicz <zbigniew.zdanowicz@intel.com>
---
 level_zero/core/source/event/event.h                 |  1 +
 .../core/test/unit_tests/fixtures/device_fixture.h   |  4 ++--
 .../core/test/unit_tests/fixtures/module_fixture.h   | 10 ++++++++++
 .../test/unit_tests/sources/event/test_event.cpp     | 12 ++++++++++++
 shared/source/device/root_device.cpp                 |  7 ++++---
 5 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/level_zero/core/source/event/event.h b/level_zero/core/source/event/event.h
index d3877837d1..36769a3b53 100644
--- a/level_zero/core/source/event/event.h
+++ b/level_zero/core/source/event/event.h
@@ -46,6 +46,7 @@ struct Event : _ze_event_handle_t {
     virtual NEO::GraphicsAllocation &getAllocation();
 
     void increasePacketsInUse() { packetsInUse++; }
+    void setPacketsInUse(uint32_t newPacketsInUse) { packetsInUse = newPacketsInUse; }
     void resetPackets() { packetsInUse = 0; }
     uint64_t getGpuAddress() { return gpuAddress; }
     void *getHostAddress() { return hostAddress; }
diff --git a/level_zero/core/test/unit_tests/fixtures/device_fixture.h b/level_zero/core/test/unit_tests/fixtures/device_fixture.h
index 0e45af7b8e..715a0a1359 100644
--- a/level_zero/core/test/unit_tests/fixtures/device_fixture.h
+++ b/level_zero/core/test/unit_tests/fixtures/device_fixture.h
@@ -42,8 +42,8 @@ struct MultiDeviceFixture {
     DebugManagerStateRestore restorer;
     std::unique_ptr<Mock<L0::DriverHandleImp>> driverHandle;
     std::vector<NEO::Device *> devices;
-    const uint32_t numRootDevices = 2u;
-    const uint32_t numSubDevices = 2u;
+    uint32_t numRootDevices = 2u;
+    uint32_t numSubDevices = 2u;
     L0::ContextImp *context = nullptr;
 };
 
diff --git a/level_zero/core/test/unit_tests/fixtures/module_fixture.h b/level_zero/core/test/unit_tests/fixtures/module_fixture.h
index 6a8c5e51e4..683fea99c4 100644
--- a/level_zero/core/test/unit_tests/fixtures/module_fixture.h
+++ b/level_zero/core/test/unit_tests/fixtures/module_fixture.h
@@ -249,6 +249,15 @@ struct MultiDeviceModuleFixture : public MultiDeviceFixture {
                                                       moduleBuildLog, ModuleType::User));
     }
 
+    void createKernel(uint32_t rootDeviceIndex) {
+        ze_kernel_desc_t desc = {};
+        desc.pKernelName = kernelName.c_str();
+
+        kernel = std::make_unique<WhiteBox<::L0::Kernel>>();
+        kernel->module = modules[rootDeviceIndex].get();
+        kernel->initialize(&desc);
+    }
+
     void TearDown() override {
         MultiDeviceFixture::TearDown();
     }
@@ -257,6 +266,7 @@ struct MultiDeviceModuleFixture : public MultiDeviceFixture {
     const std::string kernelName = "test";
     const uint32_t numKernelArguments = 6;
     std::vector<std::unique_ptr<L0::Module>> modules;
+    std::unique_ptr<WhiteBox<::L0::Kernel>> kernel;
 };
 
 struct ImportHostPointerModuleFixture : public ModuleFixture {
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 d719dd35f5..aedae3ef2e 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
@@ -357,6 +357,18 @@ TEST_F(TimestampEventCreate, givenTimestampEventThenAllocationsIsOfPacketTagBuff
     EXPECT_EQ(NEO::GraphicsAllocation::AllocationType::TIMESTAMP_PACKET_TAG_BUFFER, allocation->getAllocationType());
 }
 
+TEST_F(TimestampEventCreate, givenEventTimestampWhenPacketCountIsSetThenCorrectOffsetIsReturned) {
+    EXPECT_EQ(0u, event->getPacketsInUse());
+    auto gpuAddr = event->getGpuAddress();
+    EXPECT_EQ(gpuAddr, event->getTimestampPacketAddress());
+
+    event->setPacketsInUse(4u);
+    EXPECT_EQ(4u, event->getPacketsInUse());
+
+    gpuAddr += (4u * NEO::TimestampPackets<uint32_t>::getSinglePacketSize());
+    EXPECT_EQ(gpuAddr, event->getTimestampPacketAddress());
+}
+
 TEST_F(TimestampEventCreate, givenEventTimestampWhenPacketCountIsIncreasedThenCorrectOffsetIsReturned) {
     EXPECT_EQ(0u, event->getPacketsInUse());
     auto gpuAddr = event->getGpuAddress();
diff --git a/shared/source/device/root_device.cpp b/shared/source/device/root_device.cpp
index d3c2657929..9b7e1c1e7d 100644
--- a/shared/source/device/root_device.cpp
+++ b/shared/source/device/root_device.cpp
@@ -44,7 +44,8 @@ SubDevice *RootDevice::createSubDevice(uint32_t subDeviceIndex) {
 
 bool RootDevice::createDeviceImpl() {
     auto deviceMask = executionEnvironment->rootDeviceEnvironments[this->rootDeviceIndex]->deviceAffinityMask;
-    deviceBitfield = maxNBitValue(HwHelper::getSubDevicesCount(&getHardwareInfo()));
+    uint32_t subDeviceCount = HwHelper::getSubDevicesCount(&getHardwareInfo());
+    deviceBitfield = maxNBitValue(subDeviceCount);
     deviceBitfield &= deviceMask;
     numSubDevices = static_cast<uint32_t>(deviceBitfield.count());
     if (numSubDevices == 1) {
@@ -52,8 +53,8 @@ bool RootDevice::createDeviceImpl() {
     }
     UNRECOVERABLE_IF(!subdevices.empty());
     if (numSubDevices) {
-        subdevices.resize(HwHelper::getSubDevicesCount(&getHardwareInfo()), nullptr);
-        for (auto i = 0u; i < HwHelper::getSubDevicesCount(&getHardwareInfo()); i++) {
+        subdevices.resize(subDeviceCount, nullptr);
+        for (auto i = 0u; i < subDeviceCount; i++) {
             if (!deviceBitfield.test(i)) {
                 continue;
             }