From 0e6729062a35eefcf68815c01e094b812d551af1 Mon Sep 17 00:00:00 2001
From: Szymon Morek <szymon.morek@intel.com>
Date: Thu, 18 Jul 2024 14:27:20 +0000
Subject: [PATCH] performance: enable compression on shared USM

Signed-off-by: Szymon Morek <szymon.morek@intel.com>
---
 .../os_agnostic_memory_manager.cpp            | 15 ++++++++++++
 .../memory_manager/unified_memory_manager.cpp |  7 ++++++
 .../unified_memory_manager_tests.cpp          | 24 +++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/shared/source/memory_manager/os_agnostic_memory_manager.cpp b/shared/source/memory_manager/os_agnostic_memory_manager.cpp
index c750ddaa6c..dc2eef99f2 100644
--- a/shared/source/memory_manager/os_agnostic_memory_manager.cpp
+++ b/shared/source/memory_manager/os_agnostic_memory_manager.cpp
@@ -636,6 +636,21 @@ GraphicsAllocation *OsAgnosticMemoryManager::allocateGraphicsMemoryInDevicePool(
         allocation = new MemoryAllocation(allocationData.rootDeviceIndex, numHandles, allocationData.type, storage, storage, canonizedGpuAddress,
                                           allocationData.size, counter, MemoryPool::localMemory, false, allocationData.flags.flushL3, maxOsContextCount);
         counter++;
+        if (allocationData.flags.preferCompressed) {
+            auto &productHelper = executionEnvironment.rootDeviceEnvironments[allocationData.rootDeviceIndex]->getHelper<ProductHelper>();
+            GmmRequirements gmmRequirements{};
+            gmmRequirements.allowLargePages = true;
+            gmmRequirements.preferCompressed = true;
+
+            auto gmm = std::make_unique<Gmm>(executionEnvironment.rootDeviceEnvironments[allocationData.rootDeviceIndex]->getGmmHelper(),
+                                             allocationData.hostPtr,
+                                             allocationData.size,
+                                             MemoryConstants::pageSize2M,
+                                             CacheSettingsHelper::getGmmUsageType(allocationData.type, allocationData.flags.uncacheable, productHelper),
+                                             allocationData.storageInfo,
+                                             gmmRequirements);
+            allocation->setDefaultGmm(gmm.release());
+        }
     } else {
         std::unique_ptr<Gmm> gmm;
         size_t sizeAligned64k = 0;
diff --git a/shared/source/memory_manager/unified_memory_manager.cpp b/shared/source/memory_manager/unified_memory_manager.cpp
index 715388bba6..02d48140fe 100644
--- a/shared/source/memory_manager/unified_memory_manager.cpp
+++ b/shared/source/memory_manager/unified_memory_manager.cpp
@@ -648,6 +648,13 @@ void *SVMAllocsManager::createUnifiedAllocationWithDeviceStorage(size_t size, co
                                        subDevices};
 
     gpuProperties.alignment = alignment;
+    auto compressionSupported = false;
+    if (unifiedMemoryProperties.device) {
+        compressionSupported = memoryManager->usmCompressionSupported(unifiedMemoryProperties.device);
+        compressionSupported &= memoryManager->isCompressionSupportedForShareable(unifiedMemoryProperties.allocationFlags.flags.shareable);
+    }
+    gpuProperties.flags.preferCompressed = compressionSupported;
+
     MemoryPropertiesHelper::fillCachePolicyInProperties(gpuProperties, false, svmProperties.readOnly, false, cacheRegion);
     GraphicsAllocation *allocationGpu = memoryManager->allocateGraphicsMemoryWithProperties(gpuProperties, svmPtr);
     if (!allocationGpu) {
diff --git a/shared/test/unit_test/memory_manager/unified_memory_manager_tests.cpp b/shared/test/unit_test/memory_manager/unified_memory_manager_tests.cpp
index 5a4fc92935..39976ab6e5 100644
--- a/shared/test/unit_test/memory_manager/unified_memory_manager_tests.cpp
+++ b/shared/test/unit_test/memory_manager/unified_memory_manager_tests.cpp
@@ -6,6 +6,7 @@
  */
 
 #include "shared/source/helpers/aligned_memory.h"
+#include "shared/source/helpers/gfx_core_helper.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/mocks/mock_command_stream_receiver.h"
 #include "shared/test/common/mocks/mock_device.h"
@@ -587,3 +588,26 @@ TEST_F(SVMLocalMemoryAllocatorTest, givenInternalAllocationWhenNewAllocationIsCr
     svmManager->freeSVMAlloc(ptr);
     svmManager->freeSVMAlloc(ptr2);
 }
+
+TEST_F(SVMLocalMemoryAllocatorTest, givenLocalMemoryEnabledAndCompressionEnabledThenDeviceSideSharedUsmIsCompressed) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.RenderCompressedBuffersEnabled.set(1);
+    std::unique_ptr<UltDeviceFactory> deviceFactory(new UltDeviceFactory(1, 2));
+    auto device = deviceFactory->rootDevices[0];
+
+    auto &hwInfo = device->getHardwareInfo();
+    auto &gfxCoreHelper = device->getGfxCoreHelper();
+    if (!gfxCoreHelper.usmCompressionSupported(hwInfo)) {
+        GTEST_SKIP();
+    }
+
+    void *cmdQ = reinterpret_cast<void *>(0x12345);
+    SVMAllocsManager::UnifiedMemoryProperties unifiedMemoryProperties(InternalMemoryType::sharedUnifiedMemory, 1, rootDeviceIndices, deviceBitfields);
+    unifiedMemoryProperties.device = device;
+    auto ptr = svmManager->createSharedUnifiedMemoryAllocation(4096, unifiedMemoryProperties, &cmdQ);
+    EXPECT_NE(nullptr, ptr);
+    auto svmData = svmManager->getSVMAlloc(ptr);
+    EXPECT_TRUE(svmData->gpuAllocations.getDefaultGraphicsAllocation()->isCompressionEnabled());
+
+    svmManager->freeSVMAlloc(ptr);
+}
\ No newline at end of file