fix: add infrastructure to limit device usm reuse max memory used

Related-To: NEO-12924

Signed-off-by: Dominik Dabek <[email protected]>

Author: ddabek-i

shared/source/ail/ail_configuration.h

@@ -82,6 +82,8 @@ class AILConfiguration {
 
     virtual bool disableBindlessAddressing() = 0;
 
+    virtual bool limitAmountOfDeviceMemoryForRecycling() = 0;
+
   protected:
     virtual void applyExt(RuntimeCapabilityTable &runtimeCapabilityTable) = 0;
     std::string processName;
@@ -96,6 +98,7 @@ extern const std::set<std::string_view> applicationsContextSyncFlag;
 extern const std::set<std::string_view> applicationsForceRcsDg2;
 extern const std::set<std::string_view> applicationsBufferPoolDisabled;
 extern const std::set<std::string_view> applicationsOverfetchDisabled;
+extern const std::set<std::string_view> applicationsDeviceUSMRecyclingLimited;
 
 template <PRODUCT_FAMILY product>
 class AILConfigurationHw : public AILConfiguration {
@@ -115,6 +118,7 @@ class AILConfigurationHw : public AILConfiguration {
     bool forceRcs() override;
     bool handleDivergentBarriers() override;
     bool disableBindlessAddressing() override;
+    bool limitAmountOfDeviceMemoryForRecycling() override;
 
     bool shouldForceRcs = false;
     bool shouldHandleDivergentBarriers = false;

shared/source/ail/ail_configuration_base.inl

@@ -61,4 +61,9 @@ inline void AILConfigurationHw<product>::setDisableBindlessAddressing(bool val)
     shouldDisableBindlessAddressing = val;
 }
 
+template <PRODUCT_FAMILY product>
+inline bool AILConfigurationHw<product>::limitAmountOfDeviceMemoryForRecycling() {
+    return false;
+}
+
 } // namespace NEO

shared/source/ail/ail_configuration_extra.cpp

@@ -36,6 +36,8 @@ const std::set<std::string_view> applicationsBufferPoolDisabled = {};
 
 const std::set<std::string_view> applicationsOverfetchDisabled = {};
 
+const std::set<std::string_view> applicationsDeviceUSMRecyclingLimited = {};
+
 AILConfigurationCreateFunctionType ailConfigurationFactory[IGFX_MAX_PRODUCT];
 
 void AILConfiguration::apply(RuntimeCapabilityTable &runtimeCapabilityTable) {

shared/source/ail/xe2_hpg_core/lnl/ail_configuration_lnl.cpp

@@ -18,6 +18,12 @@ bool AILConfigurationHw<IGFX_LUNARLAKE>::isBufferPoolEnabled() {
     return iterator == applicationsBufferPoolDisabled.end();
 }
 
+template <>
+bool AILConfigurationHw<IGFX_LUNARLAKE>::limitAmountOfDeviceMemoryForRecycling() {
+    auto iterator = applicationsDeviceUSMRecyclingLimited.find(processName);
+    return iterator != applicationsDeviceUSMRecyclingLimited.end();
+}
+
 template <>
 bool AILConfigurationHw<IGFX_LUNARLAKE>::is256BPrefetchDisableRequired() {
     auto iterator = applicationsOverfetchDisabled.find(processName);

shared/source/ail/xe_hpg_core/mtl/ail_configuration_mtl.cpp

@@ -58,6 +58,12 @@ bool AILConfigurationHw<IGFX_METEORLAKE>::isBufferPoolEnabled() {
     return iterator == applicationsBufferPoolDisabled.end();
 }
 
+template <>
+bool AILConfigurationHw<IGFX_METEORLAKE>::limitAmountOfDeviceMemoryForRecycling() {
+    auto iterator = applicationsDeviceUSMRecyclingLimited.find(processName);
+    return iterator != applicationsDeviceUSMRecyclingLimited.end();
+}
+
 template class AILConfigurationHw<IGFX_METEORLAKE>;
 
 } // namespace NEO

shared/source/memory_manager/unified_memory_manager.cpp

@@ -7,6 +7,7 @@
 
 #include "shared/source/memory_manager/unified_memory_manager.h"
 
+#include "shared/source/ail/ail_configuration.h"
 #include "shared/source/command_stream/command_stream_receiver.h"
 #include "shared/source/device/sub_device.h"
 #include "shared/source/execution_environment/execution_environment.h"
@@ -713,7 +714,9 @@ void SVMAllocsManager::freeZeroCopySvmAllocation(SvmAllocationData *svmData) {
 void SVMAllocsManager::initUsmDeviceAllocationsCache(Device &device) {
     this->usmDeviceAllocationsCache.allocations.reserve(128u);
     const auto totalDeviceMemory = device.getGlobalMemorySize(static_cast<uint32_t>(device.getDeviceBitfield().to_ulong()));
-    auto fractionOfTotalMemoryForRecycling = 0.08;
+    auto ailConfiguration = device.getAilConfigurationHelper();
+    const bool limitDeviceMemoryForReuse = ailConfiguration && ailConfiguration->limitAmountOfDeviceMemoryForRecycling();
+    auto fractionOfTotalMemoryForRecycling = limitDeviceMemoryForReuse ? 0.02 : 0.08;
     if (debugManager.flags.ExperimentalEnableDeviceAllocationCache.get() != -1) {
         fractionOfTotalMemoryForRecycling = 0.01 * std::min(100, debugManager.flags.ExperimentalEnableDeviceAllocationCache.get());
     }

shared/test/common/mocks/mock_ail_configuration.h

@@ -33,6 +33,11 @@ class MockAILConfiguration : public AILConfiguration {
         return isBufferPoolEnabledReturn;
     }
 
+    bool limitAmountOfDeviceMemoryForRecyclingReturn = false;
+    bool limitAmountOfDeviceMemoryForRecycling() override {
+        return limitAmountOfDeviceMemoryForRecyclingReturn;
+    }
+
     bool fallbackToLegacyValidationLogic = false;
     bool useLegacyValidationLogic() override {
         return fallbackToLegacyValidationLogic;

shared/test/unit_test/memory_manager/unified_memory_manager_cache_tests.cpp

@@ -8,6 +8,7 @@
 #include "shared/source/helpers/api_specific_config.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/helpers/raii_product_helper.h"
+#include "shared/test/common/mocks/mock_ail_configuration.h"
 #include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
 #include "shared/test/common/mocks/mock_memory_manager.h"
@@ -136,6 +137,8 @@ HWTEST_F(SvmDeviceAllocationCacheTest, givenOclApiSpecificConfigWhenCheckingIfEn
     std::unique_ptr<UltDeviceFactory> deviceFactory(new UltDeviceFactory(1, 1));
     auto device = deviceFactory->rootDevices[0];
     RAIIProductHelperFactory<MockProductHelper> raii(*device->getExecutionEnvironment()->rootDeviceEnvironments[0]);
+    MockAILConfiguration mockAilConfigurationHelper;
+    device->mockAilConfigurationHelper = &mockAilConfigurationHelper;
     {
         raii.mockProductHelper->isDeviceUsmAllocationReuseSupportedResult = false;
         auto svmManager = std::make_unique<MockSVMAllocsManager>(device->getMemoryManager(), false);
@@ -152,6 +155,16 @@ HWTEST_F(SvmDeviceAllocationCacheTest, givenOclApiSpecificConfigWhenCheckingIfEn
         const auto expectedMaxSize = static_cast<size_t>(0.08 * device->getGlobalMemorySize(static_cast<uint32_t>(device->getDeviceBitfield().to_ullong())));
         EXPECT_EQ(expectedMaxSize, svmManager->usmDeviceAllocationsCache.maxSize);
     }
+    {
+        raii.mockProductHelper->isDeviceUsmAllocationReuseSupportedResult = true;
+        mockAilConfigurationHelper.limitAmountOfDeviceMemoryForRecyclingReturn = true;
+        auto svmManager = std::make_unique<MockSVMAllocsManager>(device->getMemoryManager(), false);
+        EXPECT_FALSE(svmManager->usmDeviceAllocationsCacheEnabled);
+        svmManager->initUsmAllocationsCaches(*device);
+        EXPECT_TRUE(svmManager->usmDeviceAllocationsCacheEnabled);
+        const auto expectedMaxSize = static_cast<size_t>(0.02 * device->getGlobalMemorySize(static_cast<uint32_t>(device->getDeviceBitfield().to_ullong())));
+        EXPECT_EQ(expectedMaxSize, svmManager->usmDeviceAllocationsCache.maxSize);
+    }
 }
 
 struct SvmDeviceAllocationCacheSimpleTestDataType {