intel/compute-runtime
1
0
Fork 0
mirror of https://github.com/intel/compute-runtime.git synced 2026-01-04 15:53:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: l0 usm pooling, ipc handling

Browse Source 
Move getting usm pool for ptr to its own method.
Move trying to free via pooling to its own method.

Use base ptr of pool for tracking in IPCHandleMap.
Track ipc handle refcount in freeMem.
Return error when trying to use a not allocated ptr within pool.

Related-To: NEO-6893

Signed-off-by: Dominik Dabek <dominik.dabek@intel.com>
This commit is contained in:
Dominik Dabek
2025-05-23 13:47:41 +00:00
committed by Compute-Runtime-Automation
parent 19361031fe
commit ad4d7a2ce8
6 changed files with 228 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

151
level_zero/core/source/context/context_imp.cpp
View File

@@ -451,6 +451,42 @@ void ContextImp::freePeerAllocations(const void *ptr, bool blocking, Device *dev
}
}
NEO::UsmMemAllocPool *ContextImp::getUsmPoolOwningPtr(const void *ptr, NEO::SvmAllocationData *svmData) {
DEBUG_BREAK_IF(nullptr == svmData);
NEO::UsmMemAllocPool *usmPool = nullptr;
if (InternalMemoryType::hostUnifiedMemory == svmData->memoryType &&
driverHandle->usmHostMemAllocPool.isInPool(ptr)) {
usmPool = &driverHandle->usmHostMemAllocPool;
} else if (InternalMemoryType::deviceUnifiedMemory == svmData->memoryType) {
if (svmData->device->getUsmMemAllocPool() &&
svmData->device->getUsmMemAllocPool()->isInPool(ptr)) {
usmPool = svmData->device->getUsmMemAllocPool();
} else if (svmData->device->getUsmMemAllocPoolsManager()) {
usmPool = svmData->device->getUsmMemAllocPoolsManager()->getPoolContainingAlloc(ptr);
}
}
return usmPool;
}
bool ContextImp::tryFreeViaPooling(const void *ptr, NEO::SvmAllocationData *svmData, NEO::UsmMemAllocPool *usmPool) {
if (usmPool) {
[[maybe_unused]] auto status = usmPool->freeSVMAlloc(ptr, false);
DEBUG_BREAK_IF(false == status);
return true;
} else if (InternalMemoryType::deviceUnifiedMemory == svmData->memoryType) {
if (auto deviceUsmPoolsManager = svmData->device->getUsmMemAllocPoolsManager()) {
DEBUG_BREAK_IF(false == deviceUsmPoolsManager->isInitialized());
if (deviceUsmPoolsManager->recycleSVMAlloc(const_cast<void *>(ptr),
false)) {
return true;
}
}
}
return false;
}
ze_result_t ContextImp::freeMem(const void *ptr) {
return this->freeMem(ptr, false);
}
@@ -461,15 +497,29 @@ ze_result_t ContextImp::freeMem(const void *ptr, bool blocking) {
return ZE_RESULT_ERROR_INVALID_ARGUMENT;
}
uint64_t addressForIpc = reinterpret_cast<uint64_t>(ptr);
auto *usmPool = getUsmPoolOwningPtr(ptr, allocation);
if (usmPool) {
if (nullptr == usmPool->getPooledAllocationBasePtr(ptr)) {
// ptr is within usm pool address space but is not allocated
return ZE_RESULT_ERROR_INVALID_ARGUMENT;
} else {
addressForIpc = usmPool->getPoolAddress();
}
}
std::map<uint64_t, IpcHandleTracking *>::iterator ipcHandleIterator;
auto lockIPC = this->driverHandle->lockIPCHandleMap();
ipcHandleIterator = this->driverHandle->getIPCHandleMap().begin();
while (ipcHandleIterator != this->driverHandle->getIPCHandleMap().end()) {
if (ipcHandleIterator->second->ptr == reinterpret_cast<uint64_t>(ptr)) {
auto *memoryManager = driverHandle->getMemoryManager();
memoryManager->closeInternalHandle(ipcHandleIterator->second->ipcData.handle, ipcHandleIterator->second->handleId, nullptr);
delete ipcHandleIterator->second;
this->driverHandle->getIPCHandleMap().erase(ipcHandleIterator->first);
if (ipcHandleIterator->second->ptr == addressForIpc) {
ipcHandleIterator->second->refcnt -= 1;
if (ipcHandleIterator->second->refcnt == 0 || nullptr == usmPool) {
auto *memoryManager = driverHandle->getMemoryManager();
memoryManager->closeInternalHandle(ipcHandleIterator->second->ipcData.handle, ipcHandleIterator->second->handleId, ipcHandleIterator->second->alloc);
delete ipcHandleIterator->second;
this->driverHandle->getIPCHandleMap().erase(ipcHandleIterator->first);
}
break;
}
ipcHandleIterator++;
@@ -478,26 +528,11 @@ ze_result_t ContextImp::freeMem(const void *ptr, bool blocking) {
for (auto &pairDevice : this->devices) {
this->freePeerAllocations(ptr, blocking, Device::fromHandle(pairDevice.second));
}
if (InternalMemoryType::hostUnifiedMemory == allocation->memoryType) {
if (this->driverHandle->usmHostMemAllocPool.freeSVMAlloc(ptr, blocking)) {
return ZE_RESULT_SUCCESS;
}
} else if (InternalMemoryType::deviceUnifiedMemory == allocation->memoryType) {
if (auto deviceUsmPoolsManager = allocation->device->getUsmMemAllocPoolsManager()) {
DEBUG_BREAK_IF(false == deviceUsmPoolsManager->isInitialized());
if (deviceUsmPoolsManager->freeSVMAlloc(ptr, blocking)) {
return ZE_RESULT_SUCCESS;
}
if (deviceUsmPoolsManager->recycleSVMAlloc(const_cast<void *>(ptr),
blocking)) {
return ZE_RESULT_SUCCESS;
}
} else if (auto deviceUsmPool = allocation->device->getUsmMemAllocPool()) {
if (deviceUsmPool->freeSVMAlloc(ptr, blocking)) {
return ZE_RESULT_SUCCESS;
}
}
if (this->tryFreeViaPooling(ptr, allocation, usmPool)) {
return ZE_RESULT_SUCCESS;
}
this->driverHandle->svmAllocsManager->freeSVMAlloc(const_cast<void *>(ptr), blocking);
return ZE_RESULT_SUCCESS;
@@ -514,21 +549,18 @@ ze_result_t ContextImp::freeMemExt(const ze_memory_free_ext_desc_t *pMemFreeDesc
if (allocation == nullptr) {
return ZE_RESULT_ERROR_INVALID_ARGUMENT;
}
auto *usmPool = getUsmPoolOwningPtr(ptr, allocation);
if (usmPool && nullptr == usmPool->getPooledAllocationBasePtr(ptr)) {
// ptr is within usm pool address space but is not allocated
return ZE_RESULT_ERROR_INVALID_ARGUMENT;
}
for (auto &pairDevice : this->devices) {
this->freePeerAllocations(ptr, false, Device::fromHandle(pairDevice.second));
}
if (InternalMemoryType::hostUnifiedMemory == allocation->memoryType) {
if (this->driverHandle->usmHostMemAllocPool.freeSVMAlloc(ptr, false)) {
return ZE_RESULT_SUCCESS;
}
} else if (InternalMemoryType::deviceUnifiedMemory == allocation->memoryType) {
if (auto deviceUsmPool = allocation->device->getUsmMemAllocPool()) {
if (deviceUsmPool->freeSVMAlloc(ptr, false)) {
return ZE_RESULT_SUCCESS;
}
}
if (this->tryFreeViaPooling(ptr, allocation, usmPool)) {
return ZE_RESULT_SUCCESS;
}
this->driverHandle->svmAllocsManager->freeSVMAllocDefer(const_cast<void *>(ptr));
@@ -630,19 +662,18 @@ ze_result_t ContextImp::getMemAddressRange(const void *ptr,
size_t *pSize) {
NEO::SvmAllocationData *allocData = this->driverHandle->svmAllocsManager->getSVMAlloc(ptr);
if (allocData) {
NEO::UsmMemAllocPool *pool = nullptr;
if (driverHandle->usmHostMemAllocPool.isInPool(ptr)) {
pool = &driverHandle->usmHostMemAllocPool;
} else if (allocData->device && allocData->device->getUsmMemAllocPool() && allocData->device->getUsmMemAllocPool()->isInPool(ptr)) {
pool = allocData->device->getUsmMemAllocPool();
}
if (pool) {
auto usmPool = getUsmPoolOwningPtr(ptr, allocData);
if (usmPool) {
if (nullptr == usmPool->getPooledAllocationBasePtr(ptr)) {
// ptr is within usm pool address space but is not allocated
return ZE_RESULT_ERROR_UNKNOWN;
}
if (pBase) {
*pBase = pool->getPooledAllocationBasePtr(ptr);
*pBase = usmPool->getPooledAllocationBasePtr(ptr);
}
if (pSize) {
*pSize = pool->getPooledAllocationSize(ptr);
*pSize = usmPool->getPooledAllocationSize(ptr);
}
} else {
NEO::GraphicsAllocation *alloc;
@@ -683,31 +714,16 @@ ze_result_t ContextImp::putIpcMemHandle(ze_ipc_mem_handle_t ipcHandle) {
return ZE_RESULT_SUCCESS;
}
void ContextImp::setIPCHandleData(NEO::GraphicsAllocation *graphicsAllocation, uint64_t handle, IpcMemoryData &ipcData, uint64_t ptrAddress, uint8_t type) {
void ContextImp::setIPCHandleData(NEO::GraphicsAllocation *graphicsAllocation, uint64_t handle, IpcMemoryData &ipcData, uint64_t ptrAddress, uint8_t type, NEO::UsmMemAllocPool *usmPool) {
std::map<uint64_t, IpcHandleTracking *>::iterator ipcHandleIterator;
ipcData = {};
ipcData.handle = handle;
ipcData.type = type;
if (this->driverHandle->usmHostMemAllocPool.isInPool(addrToPtr(ptrAddress))) {
ipcData.poolOffset = this->driverHandle->usmHostMemAllocPool.getOffsetInPool(addrToPtr(ptrAddress));
} else {
for (auto const &devicePair : this->getDevices()) {
auto device = Device::fromHandle(devicePair.second);
auto neoDevice = device->getNEODevice();
if (auto deviceUsmMemAllocPoolsManager = neoDevice->getUsmMemAllocPoolsManager()) {
if (auto poolOffset = deviceUsmMemAllocPoolsManager->getOffsetInPool(addrToPtr(ptrAddress))) {
ipcData.poolOffset = poolOffset;
break;
}
} else if (auto deviceUsmMemAllocPool = neoDevice->getUsmMemAllocPool()) {
if (auto poolOffset = deviceUsmMemAllocPool->getOffsetInPool(addrToPtr(ptrAddress))) {
ipcData.poolOffset = poolOffset;
break;
}
}
}
if (usmPool) {
ipcData.poolOffset = usmPool->getOffsetInPool(addrToPtr(ptrAddress));
ptrAddress = usmPool->getPoolAddress();
}
auto lock = this->driverHandle->lockIPCHandleMap();
@@ -757,6 +773,14 @@ ze_result_t ContextImp::getIpcMemHandlesImpl(const void *ptr,
if (!allocData) {
return ZE_RESULT_ERROR_INVALID_ARGUMENT;
}
auto type = allocData->memoryType;
auto *usmPool = getUsmPoolOwningPtr(ptr, allocData);
if (usmPool && nullptr == usmPool->getPooledAllocationBasePtr(ptr)) {
// ptr is within usm pool address space but is not allocated
return ZE_RESULT_ERROR_INVALID_ARGUMENT;
}
auto memoryManager = this->driverHandle->getMemoryManager();
auto alloc = allocData->gpuAllocations.getDefaultGraphicsAllocation();
@@ -772,7 +796,6 @@ ze_result_t ContextImp::getIpcMemHandlesImpl(const void *ptr,
return ZE_RESULT_SUCCESS;
}
auto type = allocData->memoryType;
auto ipcType = InternalIpcMemoryType::deviceUnifiedMemory;
if (type == InternalMemoryType::hostUnifiedMemory) {
ipcType = InternalIpcMemoryType::hostUnifiedMemory;
@@ -789,7 +812,7 @@ ze_result_t ContextImp::getIpcMemHandlesImpl(const void *ptr,
memoryManager->registerIpcExportedAllocation(alloc);
IpcMemoryData &ipcData = *reinterpret_cast<IpcMemoryData *>(pIpcHandles[i].data);
setIPCHandleData(alloc, handle, ipcData, reinterpret_cast<uint64_t>(ptr), static_cast<uint8_t>(ipcType));
setIPCHandleData(alloc, handle, ipcData, reinterpret_cast<uint64_t>(ptr), static_cast<uint8_t>(ipcType), usmPool);
}
return ZE_RESULT_SUCCESS;

4
level_zero/core/source/context/context_imp.h
View File

@@ -197,9 +197,11 @@ struct ContextImp : Context, NEO::NonCopyableAndNonMovableClass {
protected:
ze_result_t getIpcMemHandlesImpl(const void *ptr, uint32_t *numIpcHandles, ze_ipc_mem_handle_t *pIpcHandles);
void setIPCHandleData(NEO::GraphicsAllocation *graphicsAllocation, uint64_t handle, IpcMemoryData &ipcData, uint64_t ptrAddress, uint8_t type);
void setIPCHandleData(NEO::GraphicsAllocation *graphicsAllocation, uint64_t handle, IpcMemoryData &ipcData, uint64_t ptrAddress, uint8_t type, NEO::UsmMemAllocPool *usmPool);
bool isAllocationSuitableForCompression(const StructuresLookupTable &structuresLookupTable, Device &device, size_t allocSize);
size_t getPageAlignedSizeRequired(size_t size, NEO::HeapIndex *heapRequired, size_t *pageSizeRequired);
NEO::UsmMemAllocPool *getUsmPoolOwningPtr(const void *ptr, NEO::SvmAllocationData *svmData);
bool tryFreeViaPooling(const void *ptr, NEO::SvmAllocationData *svmData, NEO::UsmMemAllocPool *usmPool);
std::map<uint32_t, ze_device_handle_t> devices;
std::vector<ze_device_handle_t> deviceHandles;

129
level_zero/core/test/unit_tests/sources/memory/test_memory_pooling.cpp
View File

@@ -68,6 +68,13 @@ struct AllocUsmPoolMemoryTest : public ::testing::Test {
context->destroy();
}
void setMaxMemoryUsed(NEO::Device &device) {
const auto isIntegrated = device.getHardwareInfo().capabilityTable.isIntegratedDevice;
const uint64_t deviceMemory = isIntegrated ? device.getDeviceInfo().globalMemSize : device.getDeviceInfo().localMemSize;
auto mockMemoryManager = reinterpret_cast<MockMemoryManager *>(driverHandle->getMemoryManager());
mockMemoryManager->localMemAllocsSize[device.getRootDeviceIndex()] = deviceMemory;
}
DebugManagerStateRestore restorer;
std::unique_ptr<Mock<L0::DriverHandleImp>> driverHandle;
const uint32_t numRootDevices = 2u;
@@ -175,11 +182,23 @@ TEST_F(AllocUsmHostEnabledMemoryTest, givenDriverHandleWhenCallingAllocHostMemWi
result = context->freeMem(ptrExportMemory);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_EQ(0u, mockHostMemAllocPool->allocations.getNumAllocs());
void *notAllocatedPoolPtr = mockHostMemAllocPool->pool;
EXPECT_EQ(nullptr, mockHostMemAllocPool->getPooledAllocationBasePtr(notAllocatedPoolPtr));
result = context->freeMemExt(&memFreeDesc, notAllocatedPoolPtr);
EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, result);
}
TEST_F(AllocUsmHostEnabledMemoryTest, givenPooledAllocationWhenCallingGetMemAddressRangeThenCorrectValuesAreReturned) {
auto pool = &driverHandle->usmHostMemAllocPool;
{
auto mockDeviceMemAllocPool = reinterpret_cast<MockUsmMemAllocPool *>(pool);
void *notAllocatedPoolPtr = mockDeviceMemAllocPool->pool;
ze_result_t result = context->getMemAddressRange(notAllocatedPoolPtr, nullptr, nullptr);
EXPECT_EQ(ZE_RESULT_ERROR_UNKNOWN, result);
}
void *pooledAllocation = nullptr;
ze_host_mem_alloc_desc_t hostDesc = {};
ze_result_t result = context->allocHostMem(&hostDesc, 1u, 0u, &pooledAllocation);
@@ -282,6 +301,12 @@ TEST_F(AllocUsmDeviceEnabledMemoryTest, givenDeviceWhenCallingAllocDeviceMemWith
EXPECT_FALSE(poolIsCompressed);
}
{
void *notAllocatedPoolPtr = mockDeviceMemAllocPool->pool;
ze_result_t result = context->freeMem(notAllocatedPoolPtr);
EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, result);
}
{
void *ptr1Byte = nullptr;
ze_device_mem_alloc_desc_t deviceDesc = {};
@@ -403,6 +428,13 @@ TEST_F(AllocUsmDeviceEnabledMemoryTest, givenDeviceWhenCallingAllocDeviceMemWith
TEST_F(AllocUsmDeviceEnabledMemoryTest, givenPooledAllocationWhenCallingGetMemAddressRangeThenCorrectValuesAreReturned) {
auto pool = l0Devices[0]->getNEODevice()->getUsmMemAllocPool();
{
auto mockDeviceMemAllocPool = reinterpret_cast<MockUsmMemAllocPool *>(pool);
void *notAllocatedPoolPtr = mockDeviceMemAllocPool->pool;
ze_result_t result = context->getMemAddressRange(notAllocatedPoolPtr, nullptr, nullptr);
EXPECT_EQ(ZE_RESULT_ERROR_UNKNOWN, result);
}
void *pooledAllocation = nullptr;
ze_device_mem_alloc_desc_t deviceDesc = {};
ze_result_t result = context->allocDeviceMem(l0Devices[0], &deviceDesc, poolAllocationThreshold, 0u, &pooledAllocation);
@@ -467,6 +499,13 @@ TEST_F(AllocUsmDeviceEnabledMemoryTest, givenDrmDriverModelWhenOpeningIpcHandleF
executionEnvironment->rootDeviceEnvironments[0]->osInterface.reset(new NEO::OSInterface());
executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::make_unique<NEO::MockDriverModelDRM>());
{
void *notAllocatedPoolPtr = mockDeviceMemAllocPool->pool;
ze_ipc_mem_handle_t ipcHandle{};
ze_result_t result = context->getIpcMemHandle(notAllocatedPoolPtr, &ipcHandle);
EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, result);
}
void *allocation = nullptr;
ze_device_mem_alloc_desc_t deviceDesc = {};
ze_result_t result = context->allocDeviceMem(l0Devices[0], &deviceDesc, 1u, 1u, &allocation);
@@ -494,6 +533,56 @@ TEST_F(AllocUsmDeviceEnabledMemoryTest, givenDrmDriverModelWhenOpeningIpcHandleF
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
}
TEST_F(AllocUsmDeviceEnabledMemoryTest, givenMultiplePooledAllocationsWhenOpeningIpcHandlesAndFreeingMemoryThenTrackRefCountCorrectly) {
auto mockDeviceMemAllocPool = reinterpret_cast<MockUsmMemAllocPool *>(l0Devices[0]->getNEODevice()->getUsmMemAllocPool());
ASSERT_NE(nullptr, mockDeviceMemAllocPool);
EXPECT_TRUE(mockDeviceMemAllocPool->isInitialized());
void *allocation1 = nullptr;
ze_device_mem_alloc_desc_t deviceDesc = {};
ze_result_t result = context->allocDeviceMem(l0Devices[0], &deviceDesc, 1u, 0u, &allocation1);
auto poolAllocationData = driverHandle->svmAllocsManager->getSVMAlloc(mockDeviceMemAllocPool->pool);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, allocation1);
EXPECT_TRUE(mockDeviceMemAllocPool->isInPool(allocation1));
EXPECT_EQ(poolAllocationData, driverHandle->svmAllocsManager->getSVMAlloc(allocation1));
void *allocation2 = nullptr;
result = context->allocDeviceMem(l0Devices[0], &deviceDesc, 1u, 0u, &allocation2);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, allocation2);
EXPECT_TRUE(mockDeviceMemAllocPool->isInPool(allocation2));
EXPECT_EQ(poolAllocationData, driverHandle->svmAllocsManager->getSVMAlloc(allocation2));
ze_ipc_mem_handle_t ipcHandle1{};
result = context->getIpcMemHandle(allocation1, &ipcHandle1);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
IpcMemoryData &ipcData1 = *reinterpret_cast<IpcMemoryData *>(ipcHandle1.data);
EXPECT_EQ(mockDeviceMemAllocPool->getOffsetInPool(allocation1), ipcData1.poolOffset);
ze_ipc_mem_handle_t ipcHandle2{};
result = context->getIpcMemHandle(allocation1, &ipcHandle2);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
IpcMemoryData &ipcData2 = *reinterpret_cast<IpcMemoryData *>(ipcHandle2.data);
EXPECT_EQ(mockDeviceMemAllocPool->getOffsetInPool(allocation1), ipcData2.poolOffset);
EXPECT_TRUE(0u != ipcData1.poolOffset || 0u != ipcData2.poolOffset); // at most one pooled allocation can have offset == 0
auto &ipcHandleMap = driverHandle->getIPCHandleMap();
ASSERT_EQ(1u, ipcHandleMap.size());
auto ipcHandleTracking = ipcHandleMap.begin()->second;
EXPECT_EQ(2u, ipcHandleTracking->refcnt);
EXPECT_EQ(mockDeviceMemAllocPool->getPoolAddress(), ipcHandleTracking->ptr);
result = context->freeMem(allocation1);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
ASSERT_EQ(1u, ipcHandleMap.size());
EXPECT_EQ(1u, ipcHandleTracking->refcnt);
result = context->freeMem(allocation2);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_EQ(0u, ipcHandleMap.size());
}
using AllocUsmDeviceEnabledMemoryNewVersionTest = AllocUsmPoolMemoryTest<-1, 1, 2>;
TEST_F(AllocUsmDeviceEnabledMemoryNewVersionTest, givenContextWhenAllocatingAndFreeingDeviceUsmThenPoolingIsUsed) {
@@ -569,8 +658,7 @@ TEST_F(AllocUsmDeviceEnabledMemoryNewVersionTest, givenContextWhenAllocatingAndF
result = context->allocDeviceMem(deviceHandle, &deviceAllocDesc, 2 * MemoryConstants::megaByte + 1, 0u, &allocation2MB1B);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, allocation2MB1B);
auto mockMemoryManager = reinterpret_cast<MockMemoryManager *>(driverHandle->getMemoryManager());
mockMemoryManager->localMemAllocsSize[mockRootDeviceIndex] = driverHandle->getMemoryManager()->getLocalMemorySize(mockRootDeviceIndex, static_cast<uint32_t>(driverHandle->devices[0]->getNEODevice()->getDeviceBitfield().to_ulong()));
this->setMaxMemoryUsed(*driverHandle->devices[0]->getNEODevice());
result = context->freeMem(allocation2MB1B); // should not be recycled, because too much device memory is used
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
void *allocation2MB1BNotRecycled = nullptr;
@@ -598,6 +686,43 @@ TEST_F(AllocUsmDeviceEnabledMemoryNewVersionTest, givenContextWhenAllocatingAndF
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
}
TEST_F(AllocUsmDeviceEnabledMemoryNewVersionTest, givenContextWhenCallingFreeMemExtThenPoolingIsUsedCorrectly) {
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{};
const size_t allocationSize = 4 * MemoryConstants::megaByte;
void *allocation = nullptr;
ze_result_t result = context->allocDeviceMem(deviceHandle, &deviceAllocDesc, allocationSize, 0u, &allocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, allocation);
ze_memory_free_ext_desc_t memFreeDesc = {};
memFreeDesc.freePolicy = ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_DEFER_FREE;
result = context->freeMemExt(&memFreeDesc, allocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
const auto &poolInfo = usmMemAllocPoolsManager->poolInfos[UsmMemAllocPoolsManager::firstNonPreallocatedIndex];
auto &poolsList = mockUsmMemAllocPoolsManager->pools[poolInfo];
ASSERT_EQ(1u, poolsList.size());
auto usmPool = reinterpret_cast<MockUsmMemAllocPool *>(poolsList[0].get());
EXPECT_EQ(allocation, usmPool->pool);
EXPECT_EQ(allocationSize, usmPool->poolSize);
usmMemAllocPoolsManager->trim();
EXPECT_EQ(0u, poolsList.size());
result = context->allocDeviceMem(deviceHandle, &deviceAllocDesc, allocationSize, 0u, &allocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, allocation);
this->setMaxMemoryUsed(*driverHandle->devices[0]->getNEODevice());
result = context->freeMemExt(&memFreeDesc, allocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_EQ(0u, poolsList.size()); // not recycled
}
TEST_F(AllocUsmDeviceEnabledMemoryNewVersionTest, givenContextWhenNormalAllocFailsThenPoolsAreTrimmed) {
auto usmMemAllocPoolsManager = driverHandle->devices[0]->getNEODevice()->getUsmMemAllocPoolsManager();
ASSERT_NE(nullptr, usmMemAllocPoolsManager);