From eed69f45ed377edd20d665e9fe6dffbbd5024e28 Mon Sep 17 00:00:00 2001
From: Bartosz Dunajski <bartosz.dunajski@intel.com>
Date: Tue, 22 Oct 2024 14:17:18 +0000
Subject: [PATCH] fix: Enable 64k pages for TSB allocation

Related-To: HSD-18040274716

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>
---
 shared/source/memory_manager/memory_manager.cpp                 | 1 +
 .../memory_manager_allocate_in_preferred_pool_tests.cpp         | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/shared/source/memory_manager/memory_manager.cpp b/shared/source/memory_manager/memory_manager.cpp
index 2ea171128f..fb4d7bde04 100644
--- a/shared/source/memory_manager/memory_manager.cpp
+++ b/shared/source/memory_manager/memory_manager.cpp
@@ -464,6 +464,7 @@ bool MemoryManager::getAllocationData(AllocationData &allocationData, const Allo
     case AllocationType::gpuTimestampDeviceBuffer:
     case AllocationType::preemption:
     case AllocationType::syncDispatchToken:
+    case AllocationType::timestampPacketTagBuffer:
         allow64KbPages = true;
     default:
         break;
diff --git a/shared/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.cpp b/shared/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.cpp
index 8c23640df6..ef4ef246e3 100644
--- a/shared/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.cpp
+++ b/shared/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.cpp
@@ -243,7 +243,6 @@ INSTANTIATE_TEST_SUITE_P(Allow32BitAnd64kbPagesTypes,
                          ::testing::ValuesIn(allocationTypesWith32BitAnd64KbPagesAllowed));
 
 static const AllocationType allocationTypesWith32BitAnd64KbPagesNotAllowed[] = {AllocationType::commandBuffer,
-                                                                                AllocationType::timestampPacketTagBuffer,
                                                                                 AllocationType::profilingTagBuffer,
                                                                                 AllocationType::image,
                                                                                 AllocationType::instructionHeap,
@@ -612,6 +611,7 @@ HWTEST_F(GetAllocationDataTestHw, givenTimestampPacketTagBufferTypeWhenGetAlloca
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
     EXPECT_FALSE(allocData.flags.useSystemMemory);
     EXPECT_TRUE(allocData.flags.requiresCpuAccess);
+    EXPECT_TRUE(allocData.flags.allow64kbPages);
 }
 
 TEST(MemoryManagerTest, givenProfilingTagBufferTypeWhenGetAllocationDataIsCalledThenSystemMemoryIsRequested) {