feature: usm pool residency tracking

Initially under debug flag.
Track residency of pool and chunks.
If pool is already resident or already evicted, we can skip memory
operation on chunk from pool.

Return error on using not allocated chunk in pool.

Related-To: NEO-16303

Signed-off-by: Dominik Dabek <dominik.dabek@intel.com>
This commit is contained in:
Dominik Dabek
2025-10-23 14:48:00 +00:00
committed by Compute-Runtime-Automation
parent 145f03c294
commit ff48090c11
9 changed files with 324 additions and 0 deletions

View File

@@ -620,6 +620,10 @@ ze_result_t ContextImp::registerMemoryFreeCallback(zex_memory_free_callback_ext_
ze_result_t ContextImp::makeMemoryResident(ze_device_handle_t hDevice, void *ptr, size_t size) {
Device *device = L0::Device::fromHandle(hDevice);
NEO::Device *neoDevice = device->getNEODevice();
if (auto usmPool = neoDevice->getUsmPoolOwningPtr(ptr); usmPool && usmPool->isTrackingResidency()) {
auto result = usmPool->residencyOperation<NEO::UsmMemAllocPool::ResidencyOperationType::makeResident>(ptr);
return changeMemoryOperationStatusToL0ResultType(result);
}
auto allocation = device->getDriverHandle()->getDriverSystemMemoryAllocation(
ptr,
size,
@@ -660,6 +664,10 @@ ze_result_t ContextImp::makeMemoryResident(ze_device_handle_t hDevice, void *ptr
ze_result_t ContextImp::evictMemory(ze_device_handle_t hDevice, void *ptr, size_t size) {
Device *device = L0::Device::fromHandle(hDevice);
NEO::Device *neoDevice = device->getNEODevice();
if (auto usmPool = neoDevice->getUsmPoolOwningPtr(ptr); usmPool && usmPool->isTrackingResidency()) {
auto result = usmPool->residencyOperation<NEO::UsmMemAllocPool::ResidencyOperationType::evict>(ptr);
return changeMemoryOperationStatusToL0ResultType(result);
}
auto allocation = device->getDriverHandle()->getDriverSystemMemoryAllocation(
ptr,
size,

View File

@@ -417,12 +417,23 @@ void DriverHandleImp::initDeviceUsmAllocPool(NEO::Device &device, bool multiDevi
poolParams.poolSize = NEO::debugManager.flags.EnableDeviceUsmAllocationPool.get() * MemoryConstants::megaByte;
}
bool trackResidency = false;
if (NEO::debugManager.flags.EnableUsmPoolResidencyTracking.get() != -1) {
trackResidency = NEO::debugManager.flags.EnableUsmPoolResidencyTracking.get() != 0;
}
if (enabled) {
if (useUsmPoolManager) {
device.resetUsmAllocationPoolManager(new NEO::UsmMemAllocPoolsManager(InternalMemoryType::deviceUnifiedMemory, rootDeviceIndices, deviceBitfields, &device));
if (trackResidency) {
device.getUsmMemAllocPoolsManager()->enableResidencyTracking();
}
device.getUsmMemAllocPoolsManager()->initialize(this->svmAllocsManager);
} else {
device.resetUsmAllocationPool(new NEO::UsmMemAllocPool);
if (trackResidency) {
device.getUsmMemAllocPool()->enableResidencyTracking();
}
device.getUsmMemAllocPool()->initialize(this->svmAllocsManager, poolMemoryProperties, poolParams.poolSize, poolParams.minServicedSize, poolParams.maxServicedSize);
}
}

View File

@@ -672,6 +672,70 @@ TEST_F(AllocUsmDeviceEnabledSinglePoolMemoryTest, givenMultiplePooledAllocations
EXPECT_EQ(0u, ipcHandleMap.size());
}
TEST_F(AllocUsmDeviceEnabledSinglePoolMemoryTest, givenPooledAllocationWhenCallingResidencyOperationsThenSkipIfAllowed) {
DebugManagerStateRestore restorer;
NEO::debugManager.flags.EnableUsmPoolResidencyTracking.set(1);
auto mockDeviceMemAllocPool = reinterpret_cast<MockUsmMemAllocPool *>(l0Devices[0]->getNEODevice()->getUsmMemAllocPool());
ASSERT_NE(nullptr, mockDeviceMemAllocPool);
EXPECT_TRUE(mockDeviceMemAllocPool->isInitialized());
mockDeviceMemAllocPool->trackResidency = true;
void *allocation = nullptr;
ze_device_mem_alloc_desc_t deviceDesc = {};
ze_result_t result = context->allocDeviceMem(l0Devices[0], &deviceDesc, 1u, 0u, &allocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, allocation);
EXPECT_TRUE(mockDeviceMemAllocPool->isInPool(allocation));
void *secondAlloc = nullptr;
result = context->allocDeviceMem(l0Devices[0], &deviceDesc, 1u, 0u, &secondAlloc);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, secondAlloc);
EXPECT_TRUE(mockDeviceMemAllocPool->isInPool(secondAlloc));
auto mockMemoryOperationsHandler = static_cast<MockMemoryOperations *>(l0Devices[0]->getNEODevice()->getRootDeviceEnvironment().memoryOperationsInterface.get());
auto expectedMakeResidentCount = mockMemoryOperationsHandler->makeResidentCalledCount;
EXPECT_EQ(ZE_RESULT_SUCCESS, context->makeMemoryResident(l0Devices[0], allocation, 1u));
EXPECT_EQ(++expectedMakeResidentCount, mockMemoryOperationsHandler->makeResidentCalledCount);
EXPECT_EQ(ZE_RESULT_SUCCESS, context->makeMemoryResident(l0Devices[0], secondAlloc, 1u));
EXPECT_EQ(expectedMakeResidentCount, mockMemoryOperationsHandler->makeResidentCalledCount);
std::unique_ptr<UsmMemAllocPool> tempPoolSwap;
auto mockNeoDevice = reinterpret_cast<MockDevice *>(l0Devices[0]->getNEODevice());
std::swap(tempPoolSwap, mockNeoDevice->usmMemAllocPool);
void *nonPooledPtr = nullptr;
result = context->allocDeviceMem(l0Devices[0], &deviceDesc, 1u, 0u, &nonPooledPtr);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, nonPooledPtr);
EXPECT_FALSE(mockDeviceMemAllocPool->isInPool(nonPooledPtr));
std::swap(tempPoolSwap, mockNeoDevice->usmMemAllocPool);
EXPECT_EQ(ZE_RESULT_SUCCESS, context->makeMemoryResident(l0Devices[0], nonPooledPtr, 1u));
EXPECT_EQ(++expectedMakeResidentCount, mockMemoryOperationsHandler->makeResidentCalledCount);
auto expectedEvictMemoryCallCount = mockMemoryOperationsHandler->evictCalledCount;
EXPECT_EQ(ZE_RESULT_SUCCESS, context->evictMemory(l0Devices[0], nonPooledPtr, 1u));
EXPECT_EQ(++expectedEvictMemoryCallCount, mockMemoryOperationsHandler->evictCalledCount);
EXPECT_EQ(ZE_RESULT_SUCCESS, context->evictMemory(l0Devices[0], secondAlloc, 1u));
EXPECT_EQ(expectedEvictMemoryCallCount, mockMemoryOperationsHandler->evictCalledCount);
EXPECT_EQ(ZE_RESULT_SUCCESS, context->evictMemory(l0Devices[0], allocation, 1u));
EXPECT_EQ(++expectedEvictMemoryCallCount, mockMemoryOperationsHandler->evictCalledCount);
// already evicted
EXPECT_EQ(ZE_RESULT_SUCCESS, context->evictMemory(l0Devices[0], allocation, 1u));
EXPECT_EQ(expectedEvictMemoryCallCount, mockMemoryOperationsHandler->evictCalledCount);
EXPECT_EQ(ZE_RESULT_SUCCESS, context->freeMem(allocation));
// not allocated pool ptr
EXPECT_TRUE(mockDeviceMemAllocPool->isInPool(allocation));
EXPECT_EQ(nullptr, mockDeviceMemAllocPool->getPooledAllocationBasePtr(allocation));
EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, context->evictMemory(l0Devices[0], allocation, 1u));
EXPECT_EQ(expectedEvictMemoryCallCount, mockMemoryOperationsHandler->evictCalledCount);
EXPECT_EQ(ZE_RESULT_SUCCESS, context->freeMem(nonPooledPtr));
}
using AllocUsmDeviceEnabledMemoryNewVersionTest = AllocUsmPoolMemoryTest<-1, 1, -1>;
TEST_F(AllocUsmDeviceEnabledMemoryNewVersionTest, givenContextWhenAllocatingAndFreeingDeviceUsmThenPoolingIsUsed) {