From bb7e32a47d764e31ab41e0e8bc27fc7cf4a510a6 Mon Sep 17 00:00:00 2001
From: "Mateusz P. Nowak" <mateusz.p.nowak@intel.com>
Date: Mon, 28 Apr 2025 16:23:16 +0200
Subject: [PATCH] [UR][L0 adapter] Move init code to platform

Signed-off-by: Mateusz P. Nowak <mateusz.p.nowak@intel.com>
---
 .../adapters/level_zero/image_common.cpp      | 24 +++++--------------
 .../source/adapters/level_zero/platform.cpp   |  8 +++++++
 .../source/adapters/level_zero/platform.hpp   |  5 ++++
 3 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/unified-runtime/source/adapters/level_zero/image_common.cpp b/unified-runtime/source/adapters/level_zero/image_common.cpp
index d80ca7d9b8522..3a067927ea0bb 100644
--- a/unified-runtime/source/adapters/level_zero/image_common.cpp
+++ b/unified-runtime/source/adapters/level_zero/image_common.cpp
@@ -19,6 +19,7 @@
 #include "helpers/memory_helpers.hpp"
 #include "image_common.hpp"
 #include "logger/ur_logger.hpp"
+#include "platform.hpp"
 #include "sampler.hpp"
 #include "ur_interface_loader.hpp"
 
@@ -26,11 +27,7 @@ typedef ze_result_t(ZE_APICALL *zeMemGetPitchFor2dImage_pfn)(
     ze_context_handle_t hContext, ze_device_handle_t hDevice, size_t imageWidth,
     size_t imageHeight, unsigned int elementSizeInBytes, size_t *rowPitch);
 
-typedef ze_result_t(ZE_APICALL *zeImageGetDeviceOffsetExp_pfn)(
-    ze_image_handle_t hImage, uint64_t *pDeviceOffset);
-
 zeMemGetPitchFor2dImage_pfn zeMemGetPitchFor2dImageFunctionPtr = nullptr;
-zeImageGetDeviceOffsetExp_pfn zeImageGetDeviceOffsetExpFunctionPtr = nullptr;
 
 namespace {
 
@@ -351,25 +348,16 @@ ur_result_t bindlessImagesCreateImpl(ur_context_handle_t hContext,
     return UR_RESULT_ERROR_INVALID_VALUE;
   }
 
-  static std::once_flag InitFlag;
-  std::call_once(InitFlag, [&]() {
-    ze_driver_handle_t DriverHandle = hContext->getPlatform()->ZeDriver;
-    auto Result = zeDriverGetExtensionFunctionAddress(
-        DriverHandle, "zeImageGetDeviceOffsetExp",
-        (void **)&zeImageGetDeviceOffsetExpFunctionPtr);
-    if (Result != ZE_RESULT_SUCCESS)
-      logger::error("zeDriverGetExtensionFunctionAddress "
-                    "zeImageGetDeviceOffsetExpv failed, err = {}",
-                    Result);
-  });
-  if (!zeImageGetDeviceOffsetExpFunctionPtr)
+  if (!hDevice->Platform->ZeImageGetDeviceOffsetExt.Supported)
     return UR_RESULT_ERROR_INVALID_OPERATION;
+
   uint64_t DeviceOffset{};
   ze_image_handle_t ZeImageTranslated;
   ZE2UR_CALL(zelLoaderTranslateHandle,
              (ZEL_HANDLE_IMAGE, ZeImage, (void **)&ZeImageTranslated));
-  ZE2UR_CALL(zeImageGetDeviceOffsetExpFunctionPtr,
-             (ZeImageTranslated, &DeviceOffset));
+  ZE2UR_CALL(
+      hDevice->Platform->ZeImageGetDeviceOffsetExt.zeImageGetDeviceOffsetExp,
+      (ZeImageTranslated, &DeviceOffset));
   *phImage = DeviceOffset;
 
   std::shared_lock<ur_shared_mutex> Lock(hDevice->Mutex);
diff --git a/unified-runtime/source/adapters/level_zero/platform.cpp b/unified-runtime/source/adapters/level_zero/platform.cpp
index 2cd7b63effcf8..d45582e53037d 100644
--- a/unified-runtime/source/adapters/level_zero/platform.cpp
+++ b/unified-runtime/source/adapters/level_zero/platform.cpp
@@ -543,6 +543,14 @@ ur_result_t ur_platform_handle_t_::initialize() {
             .zeCommandListImmediateAppendCommandListsExp != nullptr;
   }
 
+  ZE_CALL_NOCHECK(zeDriverGetExtensionFunctionAddress,
+                  (ZeDriver, "zeImageGetDeviceOffsetExp",
+                   reinterpret_cast<void **>(
+                       &ZeImageGetDeviceOffsetExt.zeImageGetDeviceOffsetExp)));
+
+  ZeImageGetDeviceOffsetExt.Supported =
+      ZeImageGetDeviceOffsetExt.zeImageGetDeviceOffsetExp != nullptr;
+
   return UR_RESULT_SUCCESS;
 }
 
diff --git a/unified-runtime/source/adapters/level_zero/platform.hpp b/unified-runtime/source/adapters/level_zero/platform.hpp
index 76a5bb72f60d8..efaa09841053c 100644
--- a/unified-runtime/source/adapters/level_zero/platform.hpp
+++ b/unified-runtime/source/adapters/level_zero/platform.hpp
@@ -169,4 +169,9 @@ struct ur_platform_handle_t_ : public ur_platform {
         ze_command_list_handle_t, uint32_t, ze_command_list_handle_t *,
         ze_event_handle_t, uint32_t, ze_event_handle_t *);
   } ZeCommandListImmediateAppendExt;
+
+  struct ZeImageGetDeviceOffsetExtension {
+    bool Supported = false;
+    ze_result_t (*zeImageGetDeviceOffsetExp)(ze_image_handle_t, uint64_t *);
+  } ZeImageGetDeviceOffsetExt;
 };