[SYCL][L0] Implement pi_device and pi_platform cache

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -139,6 +139,10 @@ class ReturnHelper {
   size_t *param_value_size_ret;
 };
 
+// save discovered pi_devices & pi_platforms for quick return
+static std::vector<pi_device> piDevicesCache;
+static std::vector<pi_platform> piPlatformsCache;
+
 } // anonymous namespace
 
 // TODO:: In the following 4 methods we may want to distinguish read access vs.
@@ -448,6 +452,8 @@ pi_result piPlatformsGet(pi_uint32 NumEntries, pi_platform *Platforms,
     return PI_INVALID_VALUE;
   }
 
+  static std::vector<pi_platform> piPlatformsCache;
+
   // This is a good time to initialize Level Zero.
   // TODO: We can still safely recover if something goes wrong during the init.
   // Implement handling segfault using sigaction.
@@ -482,6 +488,16 @@ pi_result piPlatformsGet(pi_uint32 NumEntries, pi_platform *Platforms,
     assert(ZeDriverCount == 1);
     ZE_CALL(zeDriverGet(&ZeDriverCount, &ZeDriver));
 
+    for (uint32_t i = 0; i < piPlatformsCache.size(); i++) {
+      if (piPlatformsCache[i]->ZeDriver == ZeDriver) {
+        Platforms[0] = piPlatformsCache[i];
+        if (NumPlatforms)
+          *NumPlatforms = 1;
+
+        return PI_SUCCESS;
+      }
+    }
+
     try {
       // TODO: figure out how/when to release this memory
       *Platforms = new _pi_platform(ZeDriver);
@@ -507,6 +523,9 @@ pi_result piPlatformsGet(pi_uint32 NumEntries, pi_platform *Platforms,
       Platforms[0]->ZeDriverApiVersion =
           std::to_string(ZE_MAJOR_VERSION(ZeApiVersion)) + std::string(".") +
           std::to_string(ZE_MINOR_VERSION(ZeApiVersion));
+
+      // save a copy in the cache for future uses
+      piPlatformsCache.push_back(Platforms[0]);
     } catch (const std::bad_alloc &) {
       return PI_OUT_OF_HOST_MEMORY;
     } catch (...) {
@@ -602,7 +621,18 @@ pi_result piDevicesGet(pi_platform Platform, pi_device_type DeviceType,
   uint32_t ZeDeviceCount = 0;
   const bool AskingForGPU = (DeviceType & PI_DEVICE_TYPE_GPU);
   const bool AskingForDefault = (DeviceType == PI_DEVICE_TYPE_DEFAULT);
-  ZE_CALL(zeDeviceGet(ZeDriver, &ZeDeviceCount, nullptr));
+
+  if (piDevicesCache.size() != 0) {
+    for (uint32_t i = 0; i < piDevicesCache.size(); i++) {
+      if (piDevicesCache[i]->Platform == Platform) {
+        ZeDeviceCount++;
+      }
+    }
+  }
+  if (ZeDeviceCount == 0) {
+    ZE_CALL(zeDeviceGet(ZeDriver, &ZeDeviceCount, nullptr));
+  }
+
   if (ZeDeviceCount == 0 || !(AskingForGPU || AskingForDefault)) {
     if (NumDevices)
       *NumDevices = 0;
@@ -618,6 +648,17 @@ pi_result piDevicesGet(pi_platform Platform, pi_device_type DeviceType,
     return PI_SUCCESS;
   }
 
+  // if devices are already captured in cache, return them from the cache.
+  uint32_t count = 0;
+  for (uint32_t i = 0; i < piDevicesCache.size(); i++) {
+    if (piDevicesCache[i]->Platform == Platform) {
+      Devices[count++] = piDevicesCache[i];
+    }
+  }
+  if (count == ZeDeviceCount) {
+    return PI_SUCCESS;
+  }
+
   try {
     std::vector<ze_device_handle_t> ZeDevices(ZeDeviceCount);
     ZE_CALL(zeDeviceGet(ZeDriver, &ZeDeviceCount, ZeDevices.data()));
@@ -629,6 +670,8 @@ pi_result piDevicesGet(pi_platform Platform, pi_device_type DeviceType,
         if (Result != PI_SUCCESS) {
           return Result;
         }
+        // save a copy in the cache for future uses.
+        piDevicesCache.push_back(Devices[I]);
       }
     }
   } catch (const std::bad_alloc &) {
@@ -657,6 +700,13 @@ pi_result piDeviceRelease(pi_device Device) {
   if (--(Device->RefCount) == 0) {
     // Destroy the command list used for initializations
     ZE_CALL(zeCommandListDestroy(Device->ZeCommandListInit));
+    // invalidate piDeviceCache entry
+    for (uint32_t i = 0; i < piDevicesCache.size(); i++) {
+      if (Device == piDevicesCache[i]) {
+        piDevicesCache.erase(piDevicesCache.begin() + i);
+        break;
+      }
+    }
     delete Device;
   }