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>

level_zero/core/source/context/context_imp.cpp

@@ -618,6 +618,21 @@ 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) {
+            if (pBase) {
+                *pBase = pool->getPooledAllocationBasePtr(ptr);
+            }
+
+            if (pSize) {
+                *pSize = pool->getPooledAllocationSize(ptr);
+            }
+        } else {
             NEO::GraphicsAllocation *alloc;
             alloc = allocData->gpuAllocations.getDefaultGraphicsAllocation();
             if (pBase) {
@@ -628,6 +643,7 @@ ze_result_t ContextImp::getMemAddressRange(const void *ptr,
             if (pSize) {
                 *pSize = allocData->size;
             }
+        }
 
         return ZE_RESULT_SUCCESS;
     }

level_zero/core/test/unit_tests/sources/memory/test_memory_pooling.cpp

@@ -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);