Metal: fix race condition when using a progress monitor callback

.github/workflows/internal.ci.yml

@@ -948,6 +948,16 @@ jobs:
       artifact-path: build install deps
       cmd: scripts/build.py install --full -D OIDN_SANITIZER=Address
 
+  build-macos-arm-tsan:
+    secrets: inherit
+    uses: intel-innersource/libraries.devops.renderkit.workflows/.github/workflows/macos.yml@main
+    with:
+      project: oidn
+      runs-on: '[ "macOS", "ARM64", "build" ]'
+      artifact-out: build-macos-arm-tsan
+      artifact-path: build install deps
+      cmd: scripts/build.py install --full -D OIDN_SANITIZER=Thread
+
   build-ios-arm:
     secrets: inherit
     uses: intel-innersource/libraries.devops.renderkit.workflows/.github/workflows/macos.yml@main
@@ -1051,6 +1061,21 @@ jobs:
         export TBB_ENABLE_SANITIZERS=1
         scripts/test.py --device cpu --minimal --log test.log
 
+  test-macos-arm-tsan:
+    needs: build-macos-arm-tsan
+    secrets: inherit
+    uses: intel-innersource/libraries.devops.renderkit.workflows/.github/workflows/macos.yml@main
+    with:
+      project: oidn
+      runs-on: '[ "macOS", "ARM64", "build" ]' # match the build job due to TSAN
+      artifact-in: build-macos-arm-tsan
+      artifact-out: test-macos-arm-tsan
+      artifact-path: test.log
+      artifact-on-failure: true
+      cmd: |
+        export TBB_ENABLE_SANITIZERS=1
+        scripts/test.py --device cpu --minimal --log test.log
+
   test-macos-arm-metal:
     needs: build-macos-arm
     secrets: inherit
@@ -1118,6 +1143,21 @@ jobs:
         export TBB_ENABLE_SANITIZERS=1
         scripts/test.py --device metal --minimal --log test.log
 
+  test-macos-arm-tsan-metal:
+    needs: build-macos-arm-tsan
+    secrets: inherit
+    uses: intel-innersource/libraries.devops.renderkit.workflows/.github/workflows/macos.yml@main
+    with:
+      project: oidn
+      runs-on: '[ "macOS", "ARM64", "build" ]' # match the build job due to TSAN
+      artifact-in: build-macos-arm-tsan
+      artifact-out: test-macos-arm-tsan-metal
+      artifact-path: test.log
+      artifact-on-failure: true
+      cmd: |
+        export TBB_ENABLE_SANITIZERS=1
+        scripts/test.py --device metal --minimal --log test.log
+
   build-macos:
     secrets: inherit
     uses: intel-innersource/libraries.devops.renderkit.workflows/.github/workflows/macos.yml@main

CHANGELOG.md

@@ -13,6 +13,8 @@ Version History
     combined with the handle type flag when importing external memory with
     dedicated allocation
 -   Fixed corrupted output on Apple M5 Pro/Max GPUs using Metal
+-   Fixed a race condition in the Metal device that could cause a crash when
+    using a progress monitor callback
 -   Fixed device detection failure or crash on Windows if some old Intel
     integrated GPU drivers are installed (fix requires building with oneAPI
     DPC++ Compiler 6.1.0)

devices/metal/metal_engine.h

@@ -7,6 +7,9 @@
 #include "core/engine.h"
 #include "metal_device.h"
 
+#include <mutex>
+#include <condition_variable>
+
 OIDN_NAMESPACE_BEGIN
 
   struct MetalPipeline : public RefCount
@@ -132,10 +135,20 @@ OIDN_NAMESPACE_BEGIN
     void wait() override;
 
   private:
+    // Synchronization state for host functions submitted as command buffer completion handlers,
+    // which run on background (GCD) threads
+    struct HostFuncState
+    {
+      std::mutex mutex;
+      std::condition_variable cond;
+      int numOutstanding = 0; // number of submitted host functions that have not finished yet
+    };
+
     MetalDevice* device;
     id<MTLCommandQueue> commandQueue;
     MPSCommandBuffer* commandBuffer = nullptr;
     id<MTLLibrary> library;
+    HostFuncState hostFuncState;
   };
 
 OIDN_NAMESPACE_END

devices/metal/metal_engine.mm

@@ -161,12 +161,36 @@ OIDN_NAMESPACE_BEGIN
   void MetalEngine::submitHostFunc(std::function<void()>&& f, const Ref<CancellationToken>& ct)
   {
     auto fPtr = new std::function<void()>(std::move(f));
+    HostFuncState* state = &hostFuncState;
 
+    // Register the host function as outstanding
+    {
+      std::lock_guard<std::mutex> lock(state->mutex);
+      ++state->numOutstanding;
+    }
+
+    // Submit a command buffer with a completion handler that will run the host function on a background thread
     auto commandBuffer = getMTLCommandBuffer();
     [commandBuffer addCompletedHandler: ^(id<MTLCommandBuffer> commandBuffer)
     {
       std::unique_ptr<std::function<void()>> fSmartPtr(fPtr);
+
+      // Acquire the mutex once to synchronize with the submitting thread, which released it after
+      // constructing the function. This publishes the function to this thread so it can be run
+      // safely below, without holding the lock across the call.
+      {
+        std::lock_guard<std::mutex> lock(state->mutex);
+      }
+
+      // Call the host function
       (*fSmartPtr)();
+
+      // Mark the host function as finished and wake up any thread waiting in wait()
+      {
+        std::lock_guard<std::mutex> lock(state->mutex);
+        if (--state->numOutstanding == 0)
+          state->cond.notify_all();
+      }
     }];
 
     flush();
@@ -203,6 +227,13 @@ OIDN_NAMESPACE_BEGIN
       [commandBuffer release];
       commandBuffer = nullptr;
     }
+
+    // Wait until all submitted host functions (completion handlers) have finished running on their
+    // background threads, so their effects are complete and synchronized before returning
+    {
+      std::unique_lock<std::mutex> lock(hostFuncState.mutex);
+      hostFuncState.cond.wait(lock, [&] { return hostFuncState.numOutstanding == 0; });
+    }
   }
 
 OIDN_NAMESPACE_END