Fix concurrent XNNPACK allocator initialization

onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.cc

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <mutex>
 #include <string_view>
 #include <unordered_map>
 #include <unordered_set>
@@ -183,20 +184,24 @@ XnnpackExecutionProvider::XnnpackExecutionProvider(const XnnpackExecutionProvide
 
 std::vector<AllocatorPtr> XnnpackExecutionProvider::CreatePreferredAllocators() {
   const auto& [stored_allocator, xnn_allocator] = GetStoredAllocator();
-  if (!stored_allocator) {
-    const AllocatorCreationInfo allocator_info(
-        [](int) {
-          // lazy create the allocator
-          return std::make_unique<CPUAllocator>(OrtMemoryInfo(kXnnpackExecutionProvider,
-                                                              OrtAllocatorType::OrtDeviceAllocator));
-        });
-    stored_allocator = CreateAllocator(allocator_info);
-  }
-  xnn_allocator->context = stored_allocator.get();
-  const xnn_status st = xnn_initialize(xnn_allocator);
-  if (st != xnn_status_success) {
-    ORT_THROW("XNNPACK initialization failed with status ", st);
-  }
+  // XNNPACK retains the allocator context globally, so publish it once and keep
+  // the owning allocator alive across concurrent EP creation and destruction.
+  static std::once_flag init_once;
+  std::call_once(init_once, [&]() {
+    if (!stored_allocator) {
+      const AllocatorCreationInfo allocator_info(
+          [](int) {
+            return std::make_unique<CPUAllocator>(OrtMemoryInfo(kXnnpackExecutionProvider,
+                                                                OrtAllocatorType::OrtDeviceAllocator));
+          });
+      stored_allocator = CreateAllocator(allocator_info);
+    }
+    xnn_allocator->context = stored_allocator.get();
+    const xnn_status st = xnn_initialize(xnn_allocator);
+    if (st != xnn_status_success) {
+      ORT_THROW("XNNPACK initialization failed with status ", st);
+    }
+  });
   return std::vector<AllocatorPtr>{stored_allocator};
 }
 

onnxruntime/test/providers/xnnpack/xnnpack_basic_test.cc

@@ -1,8 +1,14 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <array>
+#include <condition_variable>
+#include <cstdlib>
+#include <exception>
+#include <mutex>
 #include <random>
 #include <string>
+#include <thread>
 
 #include "core/common/logging/logging.h"
 #include "core/common/span_utils.h"
@@ -29,6 +35,7 @@
 #endif
 
 #include "gtest/gtest.h"
+#include "xnnpack.h"
 
 using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::logging;
@@ -41,6 +48,88 @@ extern std::unique_ptr<Ort::Env> ort_env;
 namespace onnxruntime {
 namespace test {
 
+#if GTEST_HAS_DEATH_TEST
+namespace {
+void TestConcurrentAllocatorInitialization() {
+  constexpr size_t num_threads = 32;
+  std::array<std::unique_ptr<IExecutionProvider>, num_threads> providers;
+  std::array<AllocatorPtr, num_threads> allocators;
+  std::array<std::exception_ptr, num_threads> errors;
+  for (auto& provider : providers) {
+    provider = DefaultXnnpackExecutionProvider();
+  }
+
+  std::mutex mutex;
+  std::condition_variable cv;
+  size_t ready = 0;
+  bool start = false;
+  InlinedVector<std::thread> threads;
+  threads.reserve(num_threads);
+  for (size_t i = 0; i < num_threads; ++i) {
+    threads.emplace_back([&, i]() {
+      {
+        std::unique_lock<std::mutex> lock(mutex);
+        ++ready;
+        cv.notify_all();
+        cv.wait(lock, [&]() { return start; });
+      }
+      ORT_TRY {
+        auto preferred_allocators = providers[i]->CreatePreferredAllocators();
+        ORT_ENFORCE(preferred_allocators.size() == 1);
+        allocators[i] = preferred_allocators[0];
+      }
+      ORT_CATCH(...) {
+        errors[i] = std::current_exception();
+      }
+    });
+  }
+  {
+    std::unique_lock<std::mutex> lock(mutex);
+    cv.wait(lock, [&]() { return ready == num_threads; });
+    start = true;
+  }
+  cv.notify_all();
+  for (auto& thread : threads) {
+    thread.join();
+  }
+
+  ASSERT_NE(allocators[0], nullptr);
+  for (size_t i = 0; i < num_threads; ++i) {
+    ASSERT_EQ(errors[i], nullptr);
+    ASSERT_EQ(allocators[i], allocators[0]);
+    providers[i].reset();
+  }
+
+  auto* allocator = allocators[0].get();
+  for (auto& retained_allocator : allocators) {
+    retained_allocator.reset();
+  }
+  // Exercise XNNPACK's retained allocator context after all EP owners are gone.
+  xnn_workspace_t workspace = nullptr;
+  ASSERT_EQ(xnn_create_workspace(&workspace), xnn_status_success);
+  ASSERT_EQ(xnn_release_workspace(workspace), xnn_status_success);
+
+  auto preferred_allocators = DefaultXnnpackExecutionProvider()->CreatePreferredAllocators();
+  ASSERT_EQ(preferred_allocators.size(), 1U);
+  ASSERT_EQ(preferred_allocators[0].get(), allocator);
+  void* buffer = allocator->Alloc(64);
+  ASSERT_NE(buffer, nullptr);
+  allocator->Free(buffer);
+}
+}  // namespace
+
+TEST(XnnpackEPDeathTest, ConcurrentAllocatorInitialization) {
+  // Re-exec so earlier XNNPACK tests cannot hide a race in first-time initialization.
+  GTEST_FLAG_SET(death_test_style, "threadsafe");
+  ASSERT_EXIT(
+      {
+        TestConcurrentAllocatorInitialization();
+        std::_Exit(::testing::Test::HasFailure() ? EXIT_FAILURE : EXIT_SUCCESS);
+      },
+      ::testing::ExitedWithCode(EXIT_SUCCESS), "");
+}
+#endif
+
 // test uses ONNX model so can't be run in a minimal build.
 // TODO: When we need XNNPACK in a minimal build we should add an ORT format version of the model
 #if !defined(ORT_MINIMAL_BUILD)