Limit allocations to available global mem size

Signed-off-by: Jaime Arteaga <jaime.a.arteaga.molina@intel.com>

Author: Jaime Arteaga

level_zero/core/source/context/context_imp.cpp

@@ -142,6 +142,12 @@ ze_result_t ContextImp::allocDeviceMem(ze_device_handle_t hDevice,
         return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
     }
 
+    if (relaxedSizeAllowed &&
+        (size > this->driverHandle->devices[0]->getNEODevice()->getDeviceInfo().globalMemSize)) {
+        *ptr = nullptr;
+        return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
+    }
+
     auto neoDevice = Device::fromHandle(hDevice)->getNEODevice();
     auto rootDeviceIndex = neoDevice->getRootDeviceIndex();
     auto deviceBitfields = this->driverHandle->deviceBitfields;
@@ -191,6 +197,12 @@ ze_result_t ContextImp::allocSharedMem(ze_device_handle_t hDevice,
         return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
     }
 
+    if (relaxedSizeAllowed &&
+        (size > this->driverHandle->devices[0]->getNEODevice()->getDeviceInfo().globalMemSize)) {
+        *ptr = nullptr;
+        return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
+    }
+
     auto neoDevice = this->devices.begin()->second->getNEODevice();
 
     auto deviceBitfields = this->deviceBitfields;

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

@@ -362,6 +362,25 @@ TEST_F(MemoryRelaxedSizeTests,
     EXPECT_EQ(ZE_RESULT_SUCCESS, result);
 }
 
+TEST_F(MemoryRelaxedSizeTests,
+       givenCallToDeviceAllocWithLargerThanGlobalMemSizeAndRelaxedFlagThenAllocationIsNotMade) {
+    size_t size = device->getNEODevice()->getDeviceInfo().globalMemSize + 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_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;
+    deviceDesc.pNext = &relaxedSizeDesc;
+    ze_result_t result = context->allocDeviceMem(device->toHandle(),
+                                                 &deviceDesc,
+                                                 size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, result);
+    EXPECT_EQ(nullptr, ptr);
+}
+
 TEST_F(MemoryRelaxedSizeTests,
        givenCallToDeviceAllocWithLargerThanAllowedSizeAndRelaxedFlagWithIncorrectFlagThenAllocationIsNotMade) {
     size_t size = device->getNEODevice()->getHardwareCapabilities().maxMemAllocSize + 1;
@@ -459,6 +478,27 @@ TEST_F(MemoryRelaxedSizeTests,
     EXPECT_EQ(ZE_RESULT_SUCCESS, result);
 }
 
+TEST_F(MemoryRelaxedSizeTests,
+       givenCallToSharedAllocWithLargerThanGlobalMemSizeAndRelaxedFlagThenAllocationIsNotMade) {
+    size_t size = device->getNEODevice()->getDeviceInfo().globalMemSize + 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_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;
+    deviceDesc.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_ERROR_UNSUPPORTED_SIZE, result);
+    EXPECT_EQ(nullptr, ptr);
+}
+
 TEST_F(MemoryRelaxedSizeTests,
        givenCallToSharedAllocWithLargerThanAllowedSizeAndRelaxedFlagWithIncorrectFlagThenAllocationIsNotMade) {
     size_t size = device->getNEODevice()->getHardwareCapabilities().maxMemAllocSize + 1;