use std::call_once

Author: uditagarwal97

sycl/source/detail/device_binary_image.cpp

@@ -710,15 +710,10 @@ CompressedRTDeviceBinaryImage::CompressedRTDeviceBinaryImage(
       static_cast<size_t>(Bin->BinaryEnd - Bin->BinaryStart));
 }
 
-// Decompress the device binary image if it is compressed. This function is
-// thread-safe and will only decompress once even if called from multiple
-// threads.
+// std::call_once ensures that this function is thread_safe and prevents
+// race during image decompression.
 void CompressedRTDeviceBinaryImage::Decompress() {
-  ImageState expected = ImageState::Compressed;
-  ImageState desired = ImageState::DecompressionInProgress;
-
-  // Decompress if not already done by another thread.
-  if (DecompState.compare_exchange_strong(expected, desired)) {
+  auto DecompressFunc = [&]() {
     size_t CompressedDataSize =
         static_cast<size_t>(Bin->BinaryEnd - Bin->BinaryStart);
 
@@ -733,17 +728,9 @@ void CompressedRTDeviceBinaryImage::Decompress() {
 
     Bin->Format = ur::getBinaryImageFormat(Bin->BinaryStart, getSize());
     Format = static_cast<ur::DeviceBinaryType>(Bin->Format);
+  };
 
-    DecompState.store(ImageState::Decompressed);
-  } else {
-    // Wait until the decompression is done by another thread.
-    while (DecompState.load() == ImageState::DecompressionInProgress) {
-      // Just spin.
-    }
-  }
-
-  assert(DecompState.load() == ImageState::Decompressed &&
-         "Image should be decompressed by now");
+  std::call_once(InitFlag, DecompressFunc);
 }
 
 CompressedRTDeviceBinaryImage::~CompressedRTDeviceBinaryImage() {

sycl/source/detail/device_binary_image.hpp

@@ -19,6 +19,7 @@
 #include <atomic>
 #include <cstring>
 #include <memory>
+#include <mutex>
 
 namespace sycl {
 inline namespace _V1 {
@@ -309,14 +310,6 @@ class DynRTDeviceBinaryImage : public RTDeviceBinaryImage {
 // actually used to build a program.
 // Also, frees the decompressed data in destructor.
 class CompressedRTDeviceBinaryImage : public RTDeviceBinaryImage {
-
-  // CompressedRTDeviceBinaryImage is in one of the following state.
-  enum ImageState {
-    Compressed = 0,
-    DecompressionInProgress = 1,
-    Decompressed = 2
-  };
-
 public:
   CompressedRTDeviceBinaryImage(sycl_device_binary Bin);
   ~CompressedRTDeviceBinaryImage() override;
@@ -339,8 +332,8 @@ class CompressedRTDeviceBinaryImage : public RTDeviceBinaryImage {
   std::unique_ptr<char[]> m_DecompressedData;
   size_t m_ImageSize = 0;
 
-  // Atomic variable used to prevent race during image decompression.
-  std::atomic<ImageState> DecompState{ImageState::Compressed};
+  // Flag to ensure decompression happens only once.
+  std::once_flag InitFlag;
 };
 #endif // SYCL_RT_ZSTD_AVAILABLE