From e81fb20505083419015a3ab88efcd7d67c128a35 Mon Sep 17 00:00:00 2001 From: Raiyan Latif Date: Thu, 23 Mar 2023 12:39:38 +0000 Subject: [PATCH] Traverse pNext chain for memory allocations extensions Related-To: LOCI-4036 Signed-off-by: Raiyan Latif --- .../core/source/context/context_imp.cpp | 25 +++++--- .../unit_tests/sources/memory/test_memory.cpp | 64 +++++++++++++++++++ 2 files changed, 79 insertions(+), 10 deletions(-) diff --git a/level_zero/core/source/context/context_imp.cpp b/level_zero/core/source/context/context_imp.cpp index bd2ec8a1f7..60fb7fe83b 100644 --- a/level_zero/core/source/context/context_imp.cpp +++ b/level_zero/core/source/context/context_imp.cpp @@ -272,17 +272,22 @@ ze_result_t ContextImp::allocSharedMem(ze_device_handle_t hDevice, bool relaxedSizeAllowed = NEO::DebugManager.flags.AllowUnrestrictedSize.get(); bool rayTracingAllocation = false; - if (deviceDesc->pNext) { - const ze_base_desc_t *extendedDesc = reinterpret_cast(deviceDesc->pNext); - if (extendedDesc->stype == ZE_STRUCTURE_TYPE_RELAXED_ALLOCATION_LIMITS_EXP_DESC) { - const ze_relaxed_allocation_limits_exp_desc_t *relaxedLimitsDesc = - reinterpret_cast(extendedDesc); - if (!(relaxedLimitsDesc->flags & ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE)) { - return ZE_RESULT_ERROR_INVALID_ARGUMENT; + const ze_base_desc_t *extendedDesc = static_cast(deviceDesc->pNext); + bool isDeviceMemAllocDescStypeValid = (deviceDesc->stype == ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC) ? true : false; + + if (isDeviceMemAllocDescStypeValid) { + while (extendedDesc) { + if (extendedDesc->stype == ZE_STRUCTURE_TYPE_RELAXED_ALLOCATION_LIMITS_EXP_DESC) { + const ze_relaxed_allocation_limits_exp_desc_t *relaxedLimitsDesc = + reinterpret_cast(extendedDesc); + if (!(relaxedLimitsDesc->flags & ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE)) { + return ZE_RESULT_ERROR_INVALID_ARGUMENT; + } + relaxedSizeAllowed = true; + } else if (extendedDesc->stype == ZE_STRUCTURE_TYPE_RAYTRACING_MEM_ALLOC_EXT_DESC) { + rayTracingAllocation = true; } - relaxedSizeAllowed = true; - } else if (extendedDesc->stype == ZE_STRUCTURE_TYPE_RAYTRACING_MEM_ALLOC_EXT_DESC) { - rayTracingAllocation = true; + extendedDesc = static_cast(extendedDesc->pNext); } } 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 fee1d4b1ed..c909197dd3 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 @@ -901,6 +901,7 @@ TEST_F(MemoryTest, givenProductWith48bForRTWhenAllocatingSharedMemoryAsRayTracin ze_raytracing_mem_alloc_ext_desc_t rtDesc = {}; rtDesc.stype = ZE_STRUCTURE_TYPE_RAYTRACING_MEM_ALLOC_EXT_DESC; + deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC; deviceDesc.pNext = &rtDesc; auto mockProductHelper = std::make_unique(); @@ -939,6 +940,7 @@ TEST_F(MemoryTest, givenProductWithNon48bForRTWhenAllocatingSharedMemoryAsRayTra ze_host_mem_alloc_desc_t hostDesc = {}; ze_raytracing_mem_alloc_ext_desc_t rtDesc = {}; + deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC; rtDesc.stype = ZE_STRUCTURE_TYPE_RAYTRACING_MEM_ALLOC_EXT_DESC; deviceDesc.pNext = &rtDesc; @@ -1970,6 +1972,36 @@ TEST_F(MemoryRelaxedSizeTests, EXPECT_EQ(ZE_RESULT_SUCCESS, result); } +TEST_F(MemoryRelaxedSizeTests, + givenCallToDeviceAllocAsRayTracingAllocationWithLargerThanAllowedSizeAndRelaxedFlagThenAllocationIsMade) { + if (device->getNEODevice()->areSharedSystemAllocationsAllowed()) { + GTEST_SKIP(); + } + size_t size = device->getNEODevice()->getDeviceInfo().maxMemAllocSize + 1; + size_t alignment = 1u; + void *ptr = nullptr; + + ze_device_mem_alloc_desc_t deviceDesc = {}; + deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC; + + ze_raytracing_mem_alloc_ext_desc_t rtDesc = {}; + rtDesc.stype = ZE_STRUCTURE_TYPE_RAYTRACING_MEM_ALLOC_EXT_DESC; + deviceDesc.pNext = &rtDesc; + + ze_relaxed_allocation_limits_exp_desc_t relaxedSizeDesc = {}; + relaxedSizeDesc.stype = ZE_STRUCTURE_TYPE_RELAXED_ALLOCATION_LIMITS_EXP_DESC; + relaxedSizeDesc.flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE; + rtDesc.pNext = &relaxedSizeDesc; + ze_result_t result = context->allocDeviceMem(device->toHandle(), + &deviceDesc, + size, alignment, &ptr); + EXPECT_EQ(ZE_RESULT_SUCCESS, result); + EXPECT_NE(nullptr, ptr); + + result = context->freeMem(ptr); + EXPECT_EQ(ZE_RESULT_SUCCESS, result); +} + TEST_F(MemoryRelaxedSizeTests, givenCallToDeviceAllocWithLargerThanAllowedSizeAndDebugFlagThenAllocationIsMade) { if (device->getNEODevice()->areSharedSystemAllocationsAllowed()) { @@ -2112,6 +2144,38 @@ TEST_F(MemoryRelaxedSizeTests, EXPECT_EQ(ZE_RESULT_SUCCESS, result); } +TEST_F(MemoryRelaxedSizeTests, + givenCallToSharedAllocAsRayTracingAllocationWithLargerThanAllowedSizeAndRelaxedFlagThenAllocationIsMade) { + if (device->getNEODevice()->areSharedSystemAllocationsAllowed()) { + GTEST_SKIP(); + } + size_t size = device->getNEODevice()->getDeviceInfo().maxMemAllocSize + 1; + size_t alignment = 1u; + void *ptr = nullptr; + + ze_device_mem_alloc_desc_t deviceDesc = {}; + deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC; + + ze_raytracing_mem_alloc_ext_desc_t rtDesc = {}; + rtDesc.stype = ZE_STRUCTURE_TYPE_RAYTRACING_MEM_ALLOC_EXT_DESC; + deviceDesc.pNext = &rtDesc; + + ze_relaxed_allocation_limits_exp_desc_t relaxedSizeDesc = {}; + relaxedSizeDesc.stype = ZE_STRUCTURE_TYPE_RELAXED_ALLOCATION_LIMITS_EXP_DESC; + relaxedSizeDesc.flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE; + rtDesc.pNext = &relaxedSizeDesc; + ze_host_mem_alloc_desc_t hostDesc = {}; + ze_result_t result = context->allocSharedMem(device->toHandle(), + &deviceDesc, + &hostDesc, + size, alignment, &ptr); + EXPECT_EQ(ZE_RESULT_SUCCESS, result); + EXPECT_NE(nullptr, ptr); + + result = context->freeMem(ptr); + EXPECT_EQ(ZE_RESULT_SUCCESS, result); +} + TEST_F(MemoryRelaxedSizeTests, givenCallToSharedAllocWithLargerThanAllowedSizeAndDebugFlagThenAllocationIsMade) { if (device->getNEODevice()->areSharedSystemAllocationsAllowed()) {