Return user size on zeMemGetAddressRange

Signed-off-by: Jaime Arteaga <[email protected]>

Author: Jaime Arteaga

level_zero/core/source/context/context_imp.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2021 Intel Corporation
+ * Copyright (C) 2020-2022 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -393,7 +393,7 @@ ze_result_t ContextImp::getMemAddressRange(const void *ptr,
         }
 
         if (pSize) {
-            *pSize = alloc->getUnderlyingBufferSize();
+            *pSize = allocData->size;
         }
 
         return ZE_RESULT_SUCCESS;

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

@@ -3111,6 +3111,30 @@ TEST_F(ContextMemoryTest, whenRetrievingAddressRangeForDeviceAllocationThenRange
     EXPECT_EQ(result, ZE_RESULT_SUCCESS);
 }
 
+TEST_F(ContextMemoryTest, whenRetrievingSizeForDeviceAllocationThenUserSizeIsReturned) {
+    size_t allocSize = 100;
+    size_t alignment = 1u;
+    void *allocPtr = nullptr;
+
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    ze_result_t result = context->allocDeviceMem(device->toHandle(),
+                                                 &deviceDesc,
+                                                 allocSize, alignment, &allocPtr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_NE(nullptr, allocPtr);
+
+    void *base = nullptr;
+    size_t size = 0u;
+    void *pPtr = reinterpret_cast<void *>(reinterpret_cast<uint64_t>(allocPtr) + 77);
+    result = context->getMemAddressRange(pPtr, &base, &size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_EQ(base, allocPtr);
+    EXPECT_EQ(size, allocSize);
+
+    result = context->freeMem(allocPtr);
+    EXPECT_EQ(result, ZE_RESULT_SUCCESS);
+}
+
 TEST_F(ContextMemoryTest, whenRetrievingAddressRangeForDeviceAllocationWithNoBaseArgumentThenSizeIsCorrectAndSuccessIsReturned) {
     size_t allocSize = 4096u;
     size_t alignment = 1u;