fix: getMemAddressRange for pooled usm

Use pool methods for getting base ptr and size if passed ptr is in a usm
pool.

Related-To: NEO-6893

Signed-off-by: Dominik Dabek <dominik.dabek@intel.com>
This commit is contained in:
Dominik Dabek 2025-05-08 10:29:39 +00:00 committed by Compute-Runtime-Automation
parent 6f4a397cfc
commit 501aebfc45
2 changed files with 71 additions and 7 deletions

View File

@ -618,15 +618,31 @@ ze_result_t ContextImp::getMemAddressRange(const void *ptr,
size_t *pSize) {
NEO::SvmAllocationData *allocData = this->driverHandle->svmAllocsManager->getSVMAlloc(ptr);
if (allocData) {
NEO::GraphicsAllocation *alloc;
alloc = allocData->gpuAllocations.getDefaultGraphicsAllocation();
if (pBase) {
uint64_t *allocBase = reinterpret_cast<uint64_t *>(pBase);
*allocBase = alloc->getGpuAddress();
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) {
if (pBase) {
*pBase = pool->getPooledAllocationBasePtr(ptr);
}
if (pSize) {
*pSize = allocData->size;
if (pSize) {
*pSize = pool->getPooledAllocationSize(ptr);
}
} else {
NEO::GraphicsAllocation *alloc;
alloc = allocData->gpuAllocations.getDefaultGraphicsAllocation();
if (pBase) {
uint64_t *allocBase = reinterpret_cast<uint64_t *>(pBase);
*allocBase = alloc->getGpuAddress();
}
if (pSize) {
*pSize = allocData->size;
}
}
return ZE_RESULT_SUCCESS;

View File

@ -177,6 +177,30 @@ TEST_F(AllocUsmHostEnabledMemoryTest, givenDriverHandleWhenCallingAllocHostMemWi
EXPECT_EQ(0u, mockHostMemAllocPool->allocations.getNumAllocs());
}
TEST_F(AllocUsmHostEnabledMemoryTest, givenPooledAllocationWhenCallingGetMemAddressRangeThenCorrectValuesAreReturned) {
auto pool = &driverHandle->usmHostMemAllocPool;
void *pooledAllocation = nullptr;
ze_host_mem_alloc_desc_t hostDesc = {};
ze_result_t result = context->allocHostMem(&hostDesc, 1u, 0u, &pooledAllocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, pooledAllocation);
EXPECT_TRUE(pool->isInPool(pooledAllocation));
size_t size = 0u;
context->getMemAddressRange(pooledAllocation, nullptr, &size);
EXPECT_GE(size, 1u);
EXPECT_EQ(pool->getPooledAllocationSize(pooledAllocation), size);
void *basePtr = nullptr;
context->getMemAddressRange(pooledAllocation, &basePtr, nullptr);
EXPECT_NE(nullptr, basePtr);
EXPECT_EQ(pool->getPooledAllocationBasePtr(pooledAllocation), basePtr);
result = context->freeMem(pooledAllocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
}
TEST_F(AllocUsmHostEnabledMemoryTest, givenDrmDriverModelWhenOpeningIpcHandleFromPooledAllocationThenOffsetIsApplied) {
auto mockHostMemAllocPool = reinterpret_cast<MockUsmMemAllocPool *>(&driverHandle->usmHostMemAllocPool);
EXPECT_TRUE(driverHandle->usmHostMemAllocPool.isInitialized());
@ -376,6 +400,30 @@ TEST_F(AllocUsmDeviceEnabledMemoryTest, givenDeviceWhenCallingAllocDeviceMemWith
}
}
TEST_F(AllocUsmDeviceEnabledMemoryTest, givenPooledAllocationWhenCallingGetMemAddressRangeThenCorrectValuesAreReturned) {
auto pool = l0Devices[0]->getNEODevice()->getUsmMemAllocPool();
void *pooledAllocation = nullptr;
ze_device_mem_alloc_desc_t deviceDesc = {};
ze_result_t result = context->allocDeviceMem(l0Devices[0], &deviceDesc, poolAllocationThreshold, 0u, &pooledAllocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
EXPECT_NE(nullptr, pooledAllocation);
EXPECT_TRUE(pool->isInPool(pooledAllocation));
size_t size = 0u;
context->getMemAddressRange(pooledAllocation, nullptr, &size);
EXPECT_GE(size, poolAllocationThreshold);
EXPECT_EQ(pool->getPooledAllocationSize(pooledAllocation), size);
void *basePtr = nullptr;
context->getMemAddressRange(pooledAllocation, &basePtr, nullptr);
EXPECT_NE(nullptr, basePtr);
EXPECT_EQ(pool->getPooledAllocationBasePtr(pooledAllocation), basePtr);
result = context->freeMem(pooledAllocation);
EXPECT_EQ(ZE_RESULT_SUCCESS, result);
}
TEST_F(AllocUsmDeviceEnabledMemoryTest, givenDrmDriverModelWhenOpeningIpcHandleFromPooledAllocationThenOffsetIsApplied) {
auto mockDeviceMemAllocPool = reinterpret_cast<MockUsmMemAllocPool *>(l0Devices[0]->getNEODevice()->getUsmMemAllocPool());
ASSERT_NE(nullptr, mockDeviceMemAllocPool);