Allow allocations bigger than maxMemAllocSize in L0 provider

by using ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE flag
when appropriate.

Author: igchor

src/provider/provider_level_zero.c

@@ -30,6 +30,8 @@ typedef struct ze_memory_provider_t {
 
     ze_device_handle_t *resident_device_handles;
     uint32_t resident_device_count;
+
+    ze_device_properties_t device_properties;
 } ze_memory_provider_t;
 
 typedef struct ze_ops_t {
@@ -54,6 +56,8 @@ typedef struct ze_ops_t {
     ze_result_t (*zeContextMakeMemoryResident)(ze_context_handle_t,
                                                ze_device_handle_t, void *,
                                                size_t);
+    ze_result_t (*zeDeviceGetProperties)(ze_device_handle_t,
+                                         ze_device_properties_t *);
 } ze_ops_t;
 
 static ze_ops_t g_ze_ops;
@@ -86,12 +90,15 @@ static void init_ze_global_state(void) {
         util_get_symbol_addr(0, "zeMemCloseIpcHandle", lib_name);
     *(void **)&g_ze_ops.zeContextMakeMemoryResident =
         util_get_symbol_addr(0, "zeContextMakeMemoryResident", lib_name);
+    *(void **)&g_ze_ops.zeDeviceGetProperties =
+        util_get_symbol_addr(0, "zeDeviceGetProperties", lib_name);
 
     if (!g_ze_ops.zeMemAllocHost || !g_ze_ops.zeMemAllocDevice ||
         !g_ze_ops.zeMemAllocShared || !g_ze_ops.zeMemFree ||
         !g_ze_ops.zeMemGetIpcHandle || !g_ze_ops.zeMemOpenIpcHandle ||
         !g_ze_ops.zeMemCloseIpcHandle ||
-        !g_ze_ops.zeContextMakeMemoryResident) {
+        !g_ze_ops.zeContextMakeMemoryResident ||
+        !g_ze_ops.zeDeviceGetProperties) {
         // g_ze_ops.zeMemPutIpcHandle can be NULL because it was introduced
         // starting from Level Zero 1.6
         LOG_ERR("Required Level Zero symbols not found.");
@@ -123,6 +130,13 @@ umf_result_t ze_memory_provider_initialize(void *params, void **provider) {
     ze_provider->device = ze_params->level_zero_device_handle;
     ze_provider->memory_type = (ze_memory_type_t)ze_params->memory_type;
 
+    ze_result_t ret = g_ze_ops.zeDeviceGetProperties(
+        ze_provider->device, &ze_provider->device_properties);
+    if (ret != ZE_RESULT_SUCCESS) {
+        umf_ba_global_free(ze_provider);
+        return UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC;
+    }
+
     *provider = ze_provider;
 
     return UMF_RESULT_SUCCESS;
@@ -147,6 +161,16 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
 
     ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
 
+    bool useRelaxedAllocationFlag =
+        size > ze_provider->device_properties.maxMemAllocSize;
+    ze_relaxed_allocation_limits_exp_desc_t relaxed_desc = {
+        .stype = ZE_STRUCTURE_TYPE_RELAXED_ALLOCATION_LIMITS_EXP_DESC,
+        .pNext = NULL,
+        .flags = 0};
+    if (useRelaxedAllocationFlag) {
+        relaxed_desc.flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE;
+    }
+
     ze_result_t ze_result = ZE_RESULT_SUCCESS;
     switch (ze_provider->memory_type) {
     case UMF_MEMORY_TYPE_HOST: {
@@ -161,7 +185,7 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
     case UMF_MEMORY_TYPE_DEVICE: {
         ze_device_mem_alloc_desc_t dev_desc = {
             .stype = ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC,
-            .pNext = NULL,
+            .pNext = useRelaxedAllocationFlag ? &relaxed_desc : NULL,
             .flags = 0,
             .ordinal = 0 // TODO
         };
@@ -177,8 +201,8 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
             .flags = 0};
         ze_device_mem_alloc_desc_t dev_desc = {
             .stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
-            .pNext = NULL,
-            .flags = 0,
+            .flags = NULL,
+            .pNext = useRelaxedAllocationFlag ? &relaxed_desc : NULL,
             .ordinal = 0 // TODO
         };
         ze_result = g_ze_ops.zeMemAllocShared(ze_provider->context, &dev_desc,