fix: usm reuse, check for in use before returning

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

Author: ddabek-i

shared/source/memory_manager/unified_memory_manager.cpp

@@ -86,6 +86,18 @@ bool SVMAllocsManager::SvmAllocationCache::allocUtilizationAllows(size_t request
     return true;
 }
 
+bool SVMAllocsManager::SvmAllocationCache::isInUse(SvmAllocationData *svmData) {
+    if (svmData->cpuAllocation && memoryManager->allocInUse(*svmData->cpuAllocation)) {
+        return true;
+    }
+    for (auto &gpuAllocation : svmData->gpuAllocations.getGraphicsAllocations()) {
+        if (gpuAllocation && memoryManager->allocInUse(*gpuAllocation)) {
+            return true;
+        }
+    }
+    return false;
+}
+
 void *SVMAllocsManager::SvmAllocationCache::get(size_t size, const UnifiedMemoryProperties &unifiedMemoryProperties) {
     if (false == sizeAllowed(size)) {
         return nullptr;
@@ -102,7 +114,8 @@ void *SVMAllocsManager::SvmAllocationCache::get(size_t size, const UnifiedMemory
         UNRECOVERABLE_IF(!svmAllocData);
         if (svmAllocData->device == unifiedMemoryProperties.device &&
             svmAllocData->allocationFlagsProperty.allFlags == unifiedMemoryProperties.allocationFlags.allFlags &&
-            svmAllocData->allocationFlagsProperty.allAllocFlags == unifiedMemoryProperties.allocationFlags.allAllocFlags) {
+            svmAllocData->allocationFlagsProperty.allAllocFlags == unifiedMemoryProperties.allocationFlags.allAllocFlags &&
+            false == isInUse(svmAllocData)) {
             if (svmAllocData->device) {
                 auto lock = svmAllocData->device->obtainAllocationsReuseLock();
                 svmAllocData->device->recordAllocationGetFromReuse(allocationIter->allocationSize);
@@ -759,6 +772,7 @@ void SVMAllocsManager::initUsmDeviceAllocationsCache(Device &device) {
         this->usmDeviceAllocationsCache.allocations.reserve(128u);
     }
     this->usmDeviceAllocationsCache.svmAllocsManager = this;
+    this->usmDeviceAllocationsCache.memoryManager = memoryManager;
 }
 
 void SVMAllocsManager::initUsmHostAllocationsCache() {

shared/source/memory_manager/unified_memory_manager.h

@@ -167,6 +167,7 @@ class SVMAllocsManager {
         static bool sizeAllowed(size_t size) { return size <= SvmAllocationCache::maxServicedSize; }
         bool insert(size_t size, void *ptr, SvmAllocationData *svmData);
         static bool allocUtilizationAllows(size_t requestedSize, size_t reuseCandidateSize);
+        bool isInUse(SvmAllocationData *svmData);
         void *get(size_t size, const UnifiedMemoryProperties &unifiedMemoryProperties);
         void trim();
 

shared/test/unit_test/memory_manager/unified_memory_manager_cache_tests.cpp

@@ -105,6 +105,35 @@ TEST(SvmAllocationCacheSimpleTest, givenDifferentSizesWhenCheckingIfSizeAllowsTh
     EXPECT_FALSE(SVMAllocsManager::SvmAllocationCache::sizeAllowed(256 * MemoryConstants::megaByte + 1));
 }
 
+TEST(SvmAllocationCacheSimpleTest, givenAllocationsWhenCheckingIsInUseThenReturnCorrectValue) {
+    SVMAllocsManager::SvmAllocationCache allocationCache;
+    MockMemoryManager memoryManager;
+    MockSVMAllocsManager svmAllocsManager(&memoryManager, false);
+
+    allocationCache.memoryManager = &memoryManager;
+    allocationCache.svmAllocsManager = &svmAllocsManager;
+
+    {
+        memoryManager.deferAllocInUse = false;
+        MockGraphicsAllocation gpuGfxAllocation;
+        SvmAllocationData svmAllocData(mockRootDeviceIndex);
+        EXPECT_FALSE(allocationCache.isInUse(&svmAllocData));
+        svmAllocData.gpuAllocations.addAllocation(&gpuGfxAllocation);
+        EXPECT_FALSE(allocationCache.isInUse(&svmAllocData));
+        memoryManager.deferAllocInUse = true;
+        EXPECT_TRUE(allocationCache.isInUse(&svmAllocData));
+    }
+    {
+        memoryManager.deferAllocInUse = false;
+        MockGraphicsAllocation cpuGfxAllocation;
+        SvmAllocationData svmAllocData(mockRootDeviceIndex);
+        svmAllocData.cpuAllocation = &cpuGfxAllocation;
+        EXPECT_FALSE(allocationCache.isInUse(&svmAllocData));
+        memoryManager.deferAllocInUse = true;
+        EXPECT_TRUE(allocationCache.isInUse(&svmAllocData));
+    }
+}
+
 struct SvmAllocationCacheTestFixture {
     SvmAllocationCacheTestFixture() : executionEnvironment(defaultHwInfo.get()) {}
     void setUp() {
@@ -717,6 +746,38 @@ TEST_F(SvmDeviceAllocationCacheTest, givenAllocationWithIsInternalAllocationSetW
     svmManager->trimUSMDeviceAllocCache();
 }
 
+TEST_F(SvmDeviceAllocationCacheTest, givenAllocationInUsageWhenAllocatingAfterFreeThenDoNotReuseAllocation) {
+    std::unique_ptr<UltDeviceFactory> deviceFactory(new UltDeviceFactory(1, 1));
+    RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
+    std::map<uint32_t, DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};
+    DebugManagerStateRestore restore;
+    debugManager.flags.ExperimentalEnableDeviceAllocationCache.set(1);
+    auto device = deviceFactory->rootDevices[0];
+    auto svmManager = std::make_unique<MockSVMAllocsManager>(device->getMemoryManager(), false);
+    svmManager->initUsmAllocationsCaches(*device);
+    EXPECT_TRUE(svmManager->usmDeviceAllocationsCacheEnabled);
+    svmManager->usmDeviceAllocationsCache.maxSize = 1 * MemoryConstants::gigaByte;
+
+    SVMAllocsManager::UnifiedMemoryProperties unifiedMemoryProperties(InternalMemoryType::deviceUnifiedMemory, 1, rootDeviceIndices, deviceBitfields);
+    unifiedMemoryProperties.device = device;
+    auto allocation = svmManager->createUnifiedMemoryAllocation(10u, unifiedMemoryProperties);
+    EXPECT_NE(allocation, nullptr);
+    svmManager->freeSVMAlloc(allocation);
+    EXPECT_EQ(svmManager->usmDeviceAllocationsCache.allocations.size(), 1u);
+
+    MockMemoryManager *mockMemoryManager = reinterpret_cast<MockMemoryManager *>(device->getMemoryManager());
+    mockMemoryManager->deferAllocInUse = true;
+    auto testedAllocation = svmManager->createUnifiedMemoryAllocation(10u, unifiedMemoryProperties);
+    EXPECT_EQ(svmManager->usmDeviceAllocationsCache.allocations.size(), 1u);
+    auto svmData = svmManager->getSVMAlloc(testedAllocation);
+    EXPECT_NE(nullptr, svmData);
+
+    svmManager->freeSVMAlloc(testedAllocation);
+    EXPECT_EQ(svmManager->usmDeviceAllocationsCache.allocations.size(), 2u);
+
+    svmManager->trimUSMDeviceAllocCache();
+}
+
 using SvmHostAllocationCacheTest = Test<SvmAllocationCacheTestFixture>;
 
 TEST_F(SvmHostAllocationCacheTest, givenAllocationCacheDisabledWhenCheckingIfEnabledThenItIsDisabled) {
@@ -1278,4 +1339,35 @@ TEST_F(SvmHostAllocationCacheTest, givenHostOutOfMemoryWhenAllocatingThenCacheIs
     svmManager->trimUSMHostAllocCache();
     ASSERT_EQ(svmManager->usmHostAllocationsCache.allocations.size(), 0u);
 }
+
+TEST_F(SvmHostAllocationCacheTest, givenAllocationInUsageWhenAllocatingAfterFreeThenDoNotReuseAllocation) {
+    std::unique_ptr<UltDeviceFactory> deviceFactory(new UltDeviceFactory(1, 1));
+    RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
+    std::map<uint32_t, DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};
+    DebugManagerStateRestore restore;
+    debugManager.flags.ExperimentalEnableHostAllocationCache.set(1);
+    auto device = deviceFactory->rootDevices[0];
+    auto svmManager = std::make_unique<MockSVMAllocsManager>(device->getMemoryManager(), false);
+    svmManager->initUsmAllocationsCaches(*device);
+    EXPECT_TRUE(svmManager->usmHostAllocationsCacheEnabled);
+    svmManager->usmHostAllocationsCache.maxSize = 1 * MemoryConstants::gigaByte;
+
+    SVMAllocsManager::UnifiedMemoryProperties unifiedMemoryProperties(InternalMemoryType::hostUnifiedMemory, 1, rootDeviceIndices, deviceBitfields);
+    auto allocation = svmManager->createUnifiedMemoryAllocation(10u, unifiedMemoryProperties);
+    EXPECT_NE(allocation, nullptr);
+    svmManager->freeSVMAlloc(allocation);
+    EXPECT_EQ(svmManager->usmHostAllocationsCache.allocations.size(), 1u);
+
+    MockMemoryManager *mockMemoryManager = reinterpret_cast<MockMemoryManager *>(device->getMemoryManager());
+    mockMemoryManager->deferAllocInUse = true;
+    auto testedAllocation = svmManager->createUnifiedMemoryAllocation(10u, unifiedMemoryProperties);
+    EXPECT_EQ(svmManager->usmHostAllocationsCache.allocations.size(), 1u);
+    auto svmData = svmManager->getSVMAlloc(testedAllocation);
+    EXPECT_NE(nullptr, svmData);
+
+    svmManager->freeSVMAlloc(testedAllocation);
+    EXPECT_EQ(svmManager->usmHostAllocationsCache.allocations.size(), 2u);
+
+    svmManager->trimUSMHostAllocCache();
+}
 } // namespace NEO