Change the image-sampler mapping to track the UR sampler object and free the memory accordingly

Author: GeorgeWeb

unified-runtime/source/adapters/level_zero/device.hpp

@@ -237,9 +237,9 @@ struct ur_device_handle_t_ : ur_object {
   std::unordered_map<ur_exp_image_native_handle_t, ze_image_handle_t>
       ZeOffsetToImageHandleMap;
 
-  // Map (sampled) bindless image handle to its corresponding sampler.
-  std::unordered_map<ur_exp_image_native_handle_t, ze_sampler_handle_t>
-      ZeImageToSamplerMap;
+  // Map (sampled) bindless image handle to its corresponding sampler handle.
+  std::unordered_map<ur_exp_image_native_handle_t, ur_sampler_handle_t>
+      UrImageToSamplerMap;
 
   // unique ephemeral identifer of the device in the adapter
   std::optional<DeviceId> Id;

unified-runtime/source/adapters/level_zero/image_common.cpp

@@ -349,6 +349,8 @@ ur_result_t bindlessImagesCreateImpl(ur_context_handle_t hContext,
                (zeCtx, hDevice->ZeDevice, &ZeImageDesc, &ZeImage));
     ZE2UR_CALL(zeContextMakeImageResident, (zeCtx, hDevice->ZeDevice, ZeImage));
   } else {
+    // Release sampler handle memory to avoid memory leak.
+    UR_CALL(releaseSamplerHandle(hSampler));
     return UR_RESULT_ERROR_INVALID_VALUE;
   }
 
@@ -363,8 +365,11 @@ ur_result_t bindlessImagesCreateImpl(ur_context_handle_t hContext,
                     "zeImageGetDeviceOffsetExpv failed, err = {}",
                     Result);
   });
-  if (!zeImageGetDeviceOffsetExpFunctionPtr)
+  if (!zeImageGetDeviceOffsetExpFunctionPtr) {
+    // Release sampler handle memory to avoid memory leak.
+    UR_CALL(releaseSamplerHandle(hSampler));
     return UR_RESULT_ERROR_INVALID_OPERATION;
+  }
   uint64_t DeviceOffset{};
   ze_image_handle_t ZeImageTranslated;
   ZE2UR_CALL(zelLoaderTranslateHandle,
@@ -380,7 +385,7 @@ ur_result_t bindlessImagesCreateImpl(ur_context_handle_t hContext,
     std::scoped_lock LockAll(Lock, SLock);
     // Associate the bindless image with host object handle and sampler handle.
     hDevice->ZeOffsetToImageHandleMap[*phImage] = ZeImage;
-    hDevice->ZeImageToSamplerMap[*phImage] = hSampler->ZeSampler;
+    hDevice->UrImageToSamplerMap[*phImage] = hSampler;
   } else {
     Lock.lock();
     hDevice->ZeOffsetToImageHandleMap[*phImage] = ZeImage;
@@ -1022,11 +1027,11 @@ ur_result_t urBindlessImagesSampledImageHandleDestroyExp(
       hContext, hDevice, hImage));
 
   std::shared_lock<ur_shared_mutex> Lock(hDevice->Mutex);
-  auto Item = hDevice->ZeImageToSamplerMap.find(hImage);
-  if (Item != hDevice->ZeImageToSamplerMap.end()) {
-    hDevice->ZeImageToSamplerMap.erase(Item);
+  auto Item = hDevice->UrImageToSamplerMap.find(hImage);
+  if (Item != hDevice->UrImageToSamplerMap.end()) {
+    hDevice->UrImageToSamplerMap.erase(Item);
     Lock.release();
-    ZE2UR_CALL(zeSamplerDestroy, (Item->second));
+    UR_CALL(releaseSamplerHandle(Item->second));
   } else {
     Lock.release();
     return UR_RESULT_ERROR_INVALID_NULL_HANDLE;

unified-runtime/source/adapters/level_zero/sampler.cpp

@@ -17,6 +17,37 @@
 #include "context.hpp"
 #endif
 
+ur_result_t releaseSamplerHandle(
+    /// [in] handle of the sampler object to release
+    ur_sampler_handle_t Sampler) {
+  if (!Sampler->RefCount.decrementAndTest())
+    return UR_RESULT_SUCCESS;
+
+  if (checkL0LoaderTeardown()) {
+#ifndef NDEBUG
+    // Verify that the ZeSampler memory has not been destroyed already,
+    // otherwise this will return ZE_RESULT_ERROR_INVALID_NULL_HANDLE. This is
+    // logically impossible due to the way we have managed the lifetime of the
+    // sampler handle, but it may still be a good sanity check in debug.
+    std::shared_lock<ur_shared_mutex> Lock(Sampler->Mutex);
+    assert(Sampler->ZeSampler &&
+           "Invalid ZeSampler handle in releaseSamplerHandle");
+    Lock.release();
+#endif
+    auto ZeResult = ZE_CALL_NOCHECK(zeSamplerDestroy, (Sampler->ZeSampler));
+    // Gracefully handle the case that L0 was already unloaded.
+    if (ZeResult && (ZeResult != ZE_RESULT_ERROR_UNINITIALIZED &&
+                     ZeResult != ZE_RESULT_ERROR_UNKNOWN))
+      return ze2urResult(ZeResult);
+    if (ZeResult == ZE_RESULT_ERROR_UNKNOWN) {
+      ZeResult = ZE_RESULT_ERROR_UNINITIALIZED;
+    }
+  }
+  delete Sampler;
+
+  return UR_RESULT_SUCCESS;
+}
+
 namespace ur::level_zero {
 
 ur_result_t urSamplerCreate(
@@ -126,22 +157,7 @@ ur_result_t urSamplerRetain(
 ur_result_t urSamplerRelease(
     /// [in] handle of the sampler object to release
     ur_sampler_handle_t Sampler) {
-  if (!Sampler->RefCount.decrementAndTest())
-    return UR_RESULT_SUCCESS;
-
-  if (checkL0LoaderTeardown()) {
-    auto ZeResult = ZE_CALL_NOCHECK(zeSamplerDestroy, (Sampler->ZeSampler));
-    // Gracefully handle the case that L0 was already unloaded.
-    if (ZeResult && (ZeResult != ZE_RESULT_ERROR_UNINITIALIZED &&
-                     ZeResult != ZE_RESULT_ERROR_UNKNOWN))
-      return ze2urResult(ZeResult);
-    if (ZeResult == ZE_RESULT_ERROR_UNKNOWN) {
-      ZeResult = ZE_RESULT_ERROR_UNINITIALIZED;
-    }
-  }
-  delete Sampler;
-
-  return UR_RESULT_SUCCESS;
+  return releaseSamplerHandle(Sampler);
 }
 
 ur_result_t urSamplerGetInfo(

unified-runtime/source/adapters/level_zero/sampler.hpp

@@ -19,3 +19,7 @@ struct ur_sampler_handle_t_ : ur_object {
 
   ZeStruct<ze_sampler_desc_t> ZeSamplerDesc;
 };
+
+ur_result_t releaseSamplerHandle(
+    /// [in] handle of the sampler object to release
+    ur_sampler_handle_t Sampler);