[UR][L0] Check Global Mem Size as Limit for Free Memory

Signed-off-by: Spruit, Neil R <[email protected]>

Author: nrspruit

source/adapters/level_zero/device.cpp

@@ -88,6 +88,24 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGet(
   return UR_RESULT_SUCCESS;
 }
 
+uint64_t calculateGlobalMemSize(ur_device_handle_t Device) {
+  // Cache GlobalMemSize
+  Device->ZeGlobalMemSize.Compute =
+      [Device](struct ze_global_memsize &GlobalMemSize) {
+        for (const auto &ZeDeviceMemoryExtProperty :
+             Device->ZeDeviceMemoryProperties->second) {
+          GlobalMemSize.value += ZeDeviceMemoryExtProperty.physicalSize;
+        }
+        if (GlobalMemSize.value == 0) {
+          for (const auto &ZeDeviceMemoryProperty :
+               Device->ZeDeviceMemoryProperties->first) {
+            GlobalMemSize.value += ZeDeviceMemoryProperty.totalSize;
+          }
+        }
+      };
+  return Device->ZeGlobalMemSize.operator->()->value;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
     ur_device_handle_t Device,  ///< [in] handle of the device instance
     ur_device_info_t ParamName, ///< [in] type of the info to retrieve
@@ -251,20 +269,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
   case UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE:
     return ReturnValue(uint64_t{Device->ZeDeviceProperties->maxMemAllocSize});
   case UR_DEVICE_INFO_GLOBAL_MEM_SIZE: {
-    uint64_t GlobalMemSize = 0;
     // Support to read physicalSize depends on kernel,
     // so fallback into reading totalSize if physicalSize
     // is not available.
-    for (const auto &ZeDeviceMemoryExtProperty :
-         Device->ZeDeviceMemoryProperties->second) {
-      GlobalMemSize += ZeDeviceMemoryExtProperty.physicalSize;
-    }
-    if (GlobalMemSize == 0) {
-      for (const auto &ZeDeviceMemoryProperty :
-           Device->ZeDeviceMemoryProperties->first) {
-        GlobalMemSize += ZeDeviceMemoryProperty.totalSize;
-      }
-    }
+    uint64_t GlobalMemSize = calculateGlobalMemSize(Device);
     return ReturnValue(uint64_t{GlobalMemSize});
   }
   case UR_DEVICE_INFO_LOCAL_MEM_SIZE:
@@ -637,6 +645,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
                       static_cast<int32_t>(ZE_RESULT_ERROR_UNINITIALIZED));
       return UR_RESULT_ERROR_ADAPTER_SPECIFIC;
     }
+    // Calculate the global memory size as the max limit that can be reported as
+    // "free" memory for the user to allocate.
+    uint64_t GlobalMemSize = calculateGlobalMemSize(Device);
     // Only report device memory which zeMemAllocDevice can allocate from.
     // Currently this is only the one enumerated with ordinal 0.
     uint64_t FreeMemory = 0;
@@ -661,7 +672,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
         }
       }
     }
-    return ReturnValue(FreeMemory);
+    return ReturnValue(std::min(GlobalMemSize, FreeMemory));
   }
   case UR_DEVICE_INFO_MEMORY_CLOCK_RATE: {
     // If there are not any memory modules then return 0.

source/adapters/level_zero/device.hpp

@@ -39,6 +39,10 @@ enum EventsScope {
   LastCommandInBatchHostVisible
 };
 
+struct ze_global_memsize {
+  uint64_t value;
+};
+
 struct ur_device_handle_t_ : _ur_object {
   ur_device_handle_t_(ze_device_handle_t Device, ur_platform_handle_t Plt,
                       ur_device_handle_t ParentDevice = nullptr)
@@ -170,4 +174,5 @@ struct ur_device_handle_t_ : _ur_object {
       ZeDeviceMemoryAccessProperties;
   ZeCache<ZeStruct<ze_device_cache_properties_t>> ZeDeviceCacheProperties;
   ZeCache<ZeStruct<ze_device_ip_version_ext_t>> ZeDeviceIpVersionExt;
+  ZeCache<struct ze_global_memsize> ZeGlobalMemSize;
 };