fix: trim device usm pools when alloc fails

When normal allocation fails trim pools before deferred frees.

Related-To: NEO-6893

Signed-off-by: Dominik Dabek <dominik.dabek@intel.com>
This commit is contained in:
Dominik Dabek
2024-09-17 09:25:27 +00:00
committed by Compute-Runtime-Automation
parent 135c0b43ae
commit 731cebd721
2 changed files with 30 additions and 0 deletions

View File

@ -281,6 +281,9 @@ ze_result_t ContextImp::allocDeviceMem(ze_device_handle_t hDevice,
void *usmPtr =
this->driverHandle->svmAllocsManager->createUnifiedMemoryAllocation(size, unifiedMemoryProperties);
if (usmPtr == nullptr) {
if (neoDevice->getUsmMemAllocPoolsManager()) {
neoDevice->getUsmMemAllocPoolsManager()->trim();
}
if (driverHandle->svmAllocsManager->getNumDeferFreeAllocs() > 0) {
this->driverHandle->svmAllocsManager->freeSVMAllocDeferImpl();
usmPtr =

View File

@ -291,5 +291,32 @@ TEST_F(AllocUsmDeviceEnabledMemoryNewVersionTest, givenContextWhenAllocatingAndF
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
}
TEST_F(AllocUsmDeviceEnabledMemoryNewVersionTest, givenContextWhenNormalAllocFailsThenPoolsAreTrimmed) {
auto usmMemAllocPoolsManager = driverHandle->devices[0]->getNEODevice()->getUsmMemAllocPoolsManager();
ASSERT_NE(nullptr, usmMemAllocPoolsManager);
auto mockUsmMemAllocPoolsManager = reinterpret_cast<MockUsmMemAllocPoolsManager *>(usmMemAllocPoolsManager);
auto deviceHandle = driverHandle->devices[0]->toHandle();
ze_device_mem_alloc_desc_t deviceAllocDesc{};
void *allocation = nullptr;
const auto startingPoolSize = 20 * MemoryConstants::megaByte;
ze_result_t result = context->allocDeviceMem(deviceHandle, &deviceAllocDesc, 2 * MemoryConstants::megaByte + 1, 0u, &allocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, allocation);
EXPECT_TRUE(usmMemAllocPoolsManager->isInitialized());
result = context->freeMem(allocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_EQ(2 * MemoryConstants::megaByte + 1 + startingPoolSize, mockUsmMemAllocPoolsManager->totalSize);
auto mockMemoryManager = reinterpret_cast<MockMemoryManager *>(driverHandle->getMemoryManager());
mockMemoryManager->failInDevicePoolWithError = true;
void *failAllocation = nullptr;
result = context->allocDeviceMem(deviceHandle, &deviceAllocDesc, 2 * MemoryConstants::megaByte + 1, 0u, &failAllocation);
EXPECT_EQ(nullptr, failAllocation);
EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, result);
EXPECT_EQ(startingPoolSize, mockUsmMemAllocPoolsManager->totalSize);
}
} // namespace ult
} // namespace L0