From 74d12b7f2742ef86cab49fd5dbe3b99f7eaf8426 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabian=20Zwoli=C5=84ski?= <fabian.zwolinski@intel.com>
Date: Fri, 25 Jul 2025 13:57:04 +0000
Subject: [PATCH] fix: enable TimestampPoolAllocator on hw mode
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Related-To: NEO-12287
Signed-off-by: Fabian Zwoliński <fabian.zwolinski@intel.com>
---
 .../utilities/timestamp_pool_allocator.cpp    |  3 +-
 .../timestamp_pool_allocator_tests.cpp        | 41 +++++++++++++++----
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/shared/source/utilities/timestamp_pool_allocator.cpp b/shared/source/utilities/timestamp_pool_allocator.cpp
index 656731493f..318ac7afa4 100644
--- a/shared/source/utilities/timestamp_pool_allocator.cpp
+++ b/shared/source/utilities/timestamp_pool_allocator.cpp
@@ -11,6 +11,7 @@
 #include "shared/source/helpers/aligned_memory.h"
 #include "shared/source/memory_manager/allocation_properties.h"
 #include "shared/source/memory_manager/memory_manager.h"
+#include "shared/source/os_interface/device_factory.h"
 #include "shared/source/utilities/buffer_pool_allocator.inl"
 
 namespace NEO {
@@ -62,7 +63,7 @@ bool TimestampPoolAllocator::isEnabled() const {
         return NEO::debugManager.flags.EnableTimestampPoolAllocator.get();
     }
 
-    return false;
+    return NEO::DeviceFactory::isHwModeSelected() && device->getProductHelper().is2MBLocalMemAlignmentEnabled();
 }
 
 SharedTimestampAllocation *TimestampPoolAllocator::requestGraphicsAllocationForTimestamp(size_t size) {
diff --git a/shared/test/unit_test/utilities/timestamp_pool_allocator_tests.cpp b/shared/test/unit_test/utilities/timestamp_pool_allocator_tests.cpp
index e2b615e098..e183720b0b 100644
--- a/shared/test/unit_test/utilities/timestamp_pool_allocator_tests.cpp
+++ b/shared/test/unit_test/utilities/timestamp_pool_allocator_tests.cpp
@@ -106,34 +106,59 @@ TEST_F(TimestampPoolAllocatorTest, givenTimestampPoolAllocatorWhenRequestExceeds
 }
 
 TEST_F(TimestampPoolAllocatorTest, whenCheckingIsEnabledWithDifferentSettingsThenReturnsExpectedValue) {
+    DebugManagerStateRestore restorer;
+
     auto mockProductHelper = new MockProductHelper;
     pDevice->getRootDeviceEnvironmentRef().productHelper.reset(mockProductHelper);
     auto &timestampAllocator = pDevice->getDeviceTimestampPoolAllocator();
 
+    constexpr int32_t csrHwMode = static_cast<int32_t>(CommandStreamReceiverType::hardware);
+    constexpr int32_t csrNonHwMode = static_cast<int32_t>(CommandStreamReceiverType::tbx);
+    auto setHwMode = [&](bool hwMode) {
+        debugManager.flags.SetCommandStreamReceiver.set(hwMode ? csrHwMode : csrNonHwMode);
+    };
+
     {
         debugManager.flags.EnableTimestampPoolAllocator.set(0);
         mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
-
-        EXPECT_FALSE(timestampAllocator.isEnabled());
-    }
-    {
-        debugManager.flags.EnableTimestampPoolAllocator.set(-1);
-        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = false;
+        setHwMode(true);
 
         EXPECT_FALSE(timestampAllocator.isEnabled());
     }
     {
         debugManager.flags.EnableTimestampPoolAllocator.set(1);
         mockProductHelper->is2MBLocalMemAlignmentEnabledResult = false;
+        setHwMode(false);
 
         EXPECT_TRUE(timestampAllocator.isEnabled());
     }
+
+    debugManager.flags.EnableTimestampPoolAllocator.set(-1);
+
     {
-        debugManager.flags.EnableTimestampPoolAllocator.set(-1);
-        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
+        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = false;
+        setHwMode(false);
 
         EXPECT_FALSE(timestampAllocator.isEnabled());
     }
+    {
+        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = false;
+        setHwMode(true);
+
+        EXPECT_FALSE(timestampAllocator.isEnabled());
+    }
+    {
+        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
+        setHwMode(false);
+
+        EXPECT_FALSE(timestampAllocator.isEnabled());
+    }
+    {
+        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
+        setHwMode(true);
+
+        EXPECT_TRUE(timestampAllocator.isEnabled());
+    }
 }
 
 TEST_F(TimestampPoolAllocatorTest, givenTimestampPoolAllocatorWhenPoolSizeAlignmentRequestedThenReturnsAlignedSize) {