@@ -457,6 +457,51 @@ TEST_F(SvmDeviceAllocationCacheTest, givenAllocationsWithDifferentSizesWhenAlloc
457457 EXPECT_EQ (svmManager->usmDeviceAllocationsCache .allocations .size (), 0u );
458458}
459459
460+ TEST_F (SvmDeviceAllocationCacheTest, givenAllocationWithDifferentSizeWhenAllocatingAfterFreeThenCorrectSizeIsSet) {
461+ std::unique_ptr<UltDeviceFactory> deviceFactory (new UltDeviceFactory (1 , 1 ));
462+ RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
463+ std::map<uint32_t , DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};
464+ DebugManagerStateRestore restore;
465+ debugManager.flags .ExperimentalEnableDeviceAllocationCache .set (1 );
466+ auto device = deviceFactory->rootDevices [0 ];
467+ auto svmManager = std::make_unique<MockSVMAllocsManager>(device->getMemoryManager (), false );
468+ device->maxAllocationsSavedForReuseSize = 1 * MemoryConstants::gigaByte;
469+ svmManager->initUsmAllocationsCaches (*device);
470+ EXPECT_TRUE (svmManager->usmDeviceAllocationsCacheEnabled );
471+
472+ SVMAllocsManager::UnifiedMemoryProperties unifiedMemoryProperties (InternalMemoryType::deviceUnifiedMemory, 1 , rootDeviceIndices, deviceBitfields);
473+ unifiedMemoryProperties.device = device;
474+
475+ size_t firstAllocationSize = allocationSizeBasis - 2 ;
476+ size_t secondAllocationSize = allocationSizeBasis - 1 ;
477+
478+ auto allocation = svmManager->createUnifiedMemoryAllocation (firstAllocationSize, unifiedMemoryProperties);
479+ EXPECT_NE (allocation, nullptr );
480+
481+ size_t expectedCacheSize = 0u ;
482+ EXPECT_EQ (svmManager->usmDeviceAllocationsCache .allocations .size (), expectedCacheSize);
483+
484+ svmManager->freeSVMAlloc (allocation);
485+
486+ EXPECT_EQ (svmManager->usmDeviceAllocationsCache .allocations .size (), 1u );
487+
488+ SvmAllocationData *svmData = svmManager->getSVMAlloc (allocation);
489+ EXPECT_EQ (svmData->size , firstAllocationSize);
490+
491+ auto secondAllocation = svmManager->createUnifiedMemoryAllocation (secondAllocationSize, unifiedMemoryProperties);
492+ EXPECT_EQ (svmManager->usmDeviceAllocationsCache .allocations .size (), 1u );
493+ EXPECT_NE (secondAllocation, allocation);
494+
495+ svmManager->freeSVMAlloc (secondAllocation);
496+ EXPECT_EQ (svmManager->usmDeviceAllocationsCache .allocations .size (), 2u );
497+
498+ svmData = svmManager->getSVMAlloc (secondAllocation);
499+ EXPECT_EQ (svmData->size , secondAllocationSize);
500+
501+ svmManager->cleanupUSMAllocCaches ();
502+ EXPECT_EQ (svmManager->usmDeviceAllocationsCache .allocations .size (), 0u );
503+ }
504+
460505TEST_F (SvmDeviceAllocationCacheTest, givenAllocationsWithDifferentSizesWhenAllocatingAfterFreeThenLimitMemoryWastage) {
461506 std::unique_ptr<UltDeviceFactory> deviceFactory (new UltDeviceFactory (1 , 1 ));
462507 RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
0 commit comments