diff --git a/level_zero/core/source/context/context_imp.cpp b/level_zero/core/source/context/context_imp.cpp index 1552338081..9abe840abb 100644 --- a/level_zero/core/source/context/context_imp.cpp +++ b/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; diff --git a/level_zero/core/test/unit_tests/sources/memory/test_memory.cpp b/level_zero/core/test/unit_tests/sources/memory/test_memory.cpp index 7cc8f26b76..c61dcb785b 100644 --- a/level_zero/core/test/unit_tests/sources/memory/test_memory.cpp +++ b/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(reinterpret_cast(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;