[SYCL] Fix race during image decompression

Author: uditagarwal97

sycl/source/detail/device_binary_image.cpp

@@ -710,22 +710,40 @@ 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.
 void CompressedRTDeviceBinaryImage::Decompress() {
+  ImageState expected = ImageState::Compressed;
+  ImageState desired = ImageState::DecompressionInProgress;
 
-  size_t CompressedDataSize =
-      static_cast<size_t>(Bin->BinaryEnd - Bin->BinaryStart);
+  // Decompress if not already done by another thread.
+  if (DecompState.compare_exchange_strong(expected, desired)) {
+    size_t CompressedDataSize =
+        static_cast<size_t>(Bin->BinaryEnd - Bin->BinaryStart);
 
-  size_t DecompressedSize = 0;
-  m_DecompressedData = ZSTDCompressor::DecompressBlob(
-      reinterpret_cast<const char *>(Bin->BinaryStart), CompressedDataSize,
-      DecompressedSize);
+    size_t DecompressedSize = 0;
+    m_DecompressedData = ZSTDCompressor::DecompressBlob(
+        reinterpret_cast<const char *>(Bin->BinaryStart), CompressedDataSize,
+        DecompressedSize);
 
-  Bin->BinaryStart =
-      reinterpret_cast<const unsigned char *>(m_DecompressedData.get());
-  Bin->BinaryEnd = Bin->BinaryStart + DecompressedSize;
+    Bin->BinaryStart =
+        reinterpret_cast<const unsigned char *>(m_DecompressedData.get());
+    Bin->BinaryEnd = Bin->BinaryStart + DecompressedSize;
 
-  Bin->Format = ur::getBinaryImageFormat(Bin->BinaryStart, getSize());
-  Format = static_cast<ur::DeviceBinaryType>(Bin->Format);
+    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");
 }
 
 CompressedRTDeviceBinaryImage::~CompressedRTDeviceBinaryImage() {

sycl/source/detail/device_binary_image.hpp

@@ -309,6 +309,14 @@ 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;
@@ -330,6 +338,9 @@ class CompressedRTDeviceBinaryImage : public RTDeviceBinaryImage {
 private:
   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};
 };
 #endif // SYCL_RT_ZSTD_AVAILABLE
 

sycl/unittests/compression/CompressionTests.cpp

@@ -8,6 +8,7 @@
 
 #include "../thread_safety/ThreadUtils.h"
 #include <detail/compression.hpp>
+#include <detail/device_binary_image.hpp>
 #include <sycl/sycl.hpp>
 
 #include <string>
@@ -113,3 +114,61 @@ TEST(CompressionTest, ConcurrentCompressionDecompression) {
     ::ThreadPool MPool(ThreadCount, testCompressDecompress);
   }
 }
+
+// Test to decompress CompressedRTDeviceImage using multiple threads.
+// The idea behind this test is to ensure that a device image is
+// decompressed only once even if multiple threads try to decompress
+// it at the same time.
+TEST(CompressionTest, ConcurrentDecompressionOfDeviceImage) {
+  // Data to compress.
+  std::string data = "Hello World! I'm about to get compressed :P";
+
+  // Compress this data.
+  size_t compressedSize = 0;
+  auto compressedData = ZSTDCompressor::CompressBlob(data.c_str(), data.size(),
+                                                     compressedSize, 1);
+
+  unsigned char *compressedDataPtr =
+      reinterpret_cast<unsigned char *>(compressedData.get());
+
+  const char *EntryName = "Entry";
+  _sycl_offload_entry_struct EntryStruct = {
+      /*addr*/ nullptr, const_cast<char *>(EntryName), strlen(EntryName),
+      /*flags*/ 0, /*reserved*/ 0};
+  sycl_device_binary_struct BinStruct{/*Version*/ 1,
+                                      /*Kind*/ 4,
+                                      /*Format*/ SYCL_DEVICE_BINARY_TYPE_SPIRV,
+                                      /*DeviceTargetSpec*/ nullptr,
+                                      /*CompileOptions*/ nullptr,
+                                      /*LinkOptions*/ nullptr,
+                                      /*ManifestStart*/ nullptr,
+                                      /*ManifestEnd*/ nullptr,
+                                      /*BinaryStart*/ compressedDataPtr,
+                                      /*BinaryEnd*/ compressedDataPtr +
+                                          compressedSize,
+                                      /*EntriesBegin*/ &EntryStruct,
+                                      /*EntriesEnd*/ &EntryStruct + 1,
+                                      /*PropertySetsBegin*/ nullptr,
+                                      /*PropertySetsEnd*/ nullptr};
+  sycl_device_binary Bin = &BinStruct;
+  CompressedRTDeviceBinaryImage Img{Bin};
+
+  // Decompress the image with multiple threads.
+  constexpr size_t ThreadCount = 20;
+  Barrier b(ThreadCount);
+  {
+    auto testDecompress = [&](size_t threadId) {
+      b.wait();
+      Img.Decompress();
+
+      // Check if decompressed data is same as original data.
+      // Img.getRawData will change if there's a race in image decompression
+      // and the check will fail.
+      for (size_t i = 0; i < Img.getSize(); ++i) {
+        ASSERT_EQ(data[i], Img.getRawData().BinaryStart[i]);
+      }
+    };
+
+    ::ThreadPool MPool(ThreadCount, testDecompress);
+  }
+}