From 40c8da9b9ca637afbd887b1e5a3deb7f7581febc Mon Sep 17 00:00:00 2001
From: "Spruit, Neil R" <neil.r.spruit@intel.com>
Date: Mon, 2 Oct 2023 08:32:07 -0700
Subject: [PATCH] [UR][L0] Check Global Mem Size as Limit for Free Memory

Signed-off-by: Spruit, Neil R <neil.r.spruit@intel.com>
---
 source/adapters/level_zero/device.cpp | 35 ++++++++++++++++++---------
 source/adapters/level_zero/device.hpp |  5 ++++
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/source/adapters/level_zero/device.cpp b/source/adapters/level_zero/device.cpp
index f5b00d80cc..ec6a294c21 100644
--- a/source/adapters/level_zero/device.cpp
+++ b/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.
diff --git a/source/adapters/level_zero/device.hpp b/source/adapters/level_zero/device.hpp
index 35404c6525..bdae64beba 100644
--- a/source/adapters/level_zero/device.hpp
+++ b/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;
 };