Traverse pNext chain for memory allocations extensions

Related-To: LOCI-4036

Signed-off-by: Raiyan Latif <raiyan.latif@intel.com>

level_zero/core/source/context/context_imp.cpp

@@ -272,8 +272,11 @@ 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 = static_cast<const ze_base_desc_t *>(deviceDesc->pNext);
-        const ze_base_desc_t *extendedDesc = reinterpret_cast<const ze_base_desc_t *>(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<const ze_relaxed_allocation_limits_exp_desc_t *>(extendedDesc);
@@ -284,6 +287,8 @@ ze_result_t ContextImp::allocSharedMem(ze_device_handle_t hDevice,
             } else if (extendedDesc->stype == ZE_STRUCTURE_TYPE_RAYTRACING_MEM_ALLOC_EXT_DESC) {
                 rayTracingAllocation = true;
             }
+            extendedDesc = static_cast<const ze_base_desc_t *>(extendedDesc->pNext);
+        }
     }
 
     if (relaxedSizeAllowed == false &&

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<MockProductHelper>();
@@ -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()) {