From 57d55bea14035bb71c069e6217ced14d91352cfc Mon Sep 17 00:00:00 2001
From: Gregory Comer <gjcomer@meta.com>
Date: Thu, 21 May 2026 18:01:29 -0700
Subject: [PATCH] Always lock workspace in buck builds (#19734)

Summary:

In older versions of the XNNPACK backend, we always locked when accessing the XNNPACK workspace. This gives guarantees that we don't re-size the workspace, even when callers violate thread safety.

In D95475030, an OSS contributor disabled locking when workspace sharing is off. This is safe as long as the caller observes proper thread safety. However, we're seeing some crashes that look like they might be caused by the workspace resizing or freeing during inference.

This diff restores the locking for buck builds in the workspace sharing disabled case as a hardening measure. I've left it enabled in CMake.

Differential Revision: D106026285
---
 backends/xnnpack/runtime/XNNExecutor.cpp      | 22 ++++++++++++++-----
 backends/xnnpack/runtime/XNNExecutor.h        | 10 +++++++++
 .../xnnpack/runtime/XNNWorkspaceManager.cpp   |  2 ++
 backends/xnnpack/targets.bzl                  |  2 ++
 4 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/backends/xnnpack/runtime/XNNExecutor.cpp b/backends/xnnpack/runtime/XNNExecutor.cpp
index 103a8812931..39f067f9ec9 100644
--- a/backends/xnnpack/runtime/XNNExecutor.cpp
+++ b/backends/xnnpack/runtime/XNNExecutor.cpp
@@ -71,6 +71,13 @@ ET_NODISCARD Error XNNExecutor::initialize(
  * delegate->execute()
  */
 ET_NODISCARD Error XNNExecutor::prepare_args(Span<EValue*> args) {
+  ET_DCHECK_MSG(
+      !destroyed_.load(std::memory_order_acquire),
+      "XNNExecutor::prepare_args called after destroy");
+  ET_DCHECK_MSG(
+      !in_use_.exchange(true, std::memory_order_acquire),
+      "XNNExecutor::prepare_args called concurrently");
+
   ET_CHECK_OR_RETURN_ERROR(
       runtime_ != nullptr,
       Internal,
@@ -160,11 +167,14 @@ ET_NODISCARD Error XNNExecutor::forward(BackendExecutionContext& context) {
   xnn_status status = xnn_setup_runtime_v2(
       runtime_.get(), externals_.size(), externals_.data());
 
-  ET_CHECK_OR_RETURN_ERROR(
-      status == xnn_status_success,
-      Internal,
-      "Internal Error: Setting up the runtime failed with code: %s",
-      xnn_status_to_string(status));
+  if (status != xnn_status_success) {
+    in_use_.store(false, std::memory_order_release);
+    ET_LOG(
+        Error,
+        "Internal Error: Setting up the runtime failed with code: %s",
+        xnn_status_to_string(status));
+    return Error::Internal;
+  }
 
   auto error = profiler_.start(context.event_tracer());
   if (error != Error::Ok) {
@@ -184,6 +194,8 @@ ET_NODISCARD Error XNNExecutor::forward(BackendExecutionContext& context) {
         static_cast<unsigned int>(error));
   }
 
+  in_use_.store(false, std::memory_order_release);
+
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,
       Internal,
diff --git a/backends/xnnpack/runtime/XNNExecutor.h b/backends/xnnpack/runtime/XNNExecutor.h
index fa7c8360be4..2d709678c1c 100644
--- a/backends/xnnpack/runtime/XNNExecutor.h
+++ b/backends/xnnpack/runtime/XNNExecutor.h
@@ -16,6 +16,7 @@
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 
 #include <xnnpack.h>
+#include <atomic>
 #include <memory>
 #include <vector>
 
@@ -36,11 +37,20 @@ class XNNExecutor {
   std::vector<xnn_external_value> externals_;
   std::vector<std::string> packed_data_names_;
   std::shared_ptr<XNNWorkspace> workspace_;
+  std::atomic<bool> in_use_{false};
+  std::atomic<bool> destroyed_{false};
 
  public:
   XNNExecutor(std::shared_ptr<XNNWorkspace> workspace)
       : workspace_(workspace) {}
 
+  ~XNNExecutor() {
+    ET_DCHECK_MSG(
+        !in_use_.load(std::memory_order_acquire),
+        "XNNExecutor destroyed while in use");
+    destroyed_.store(true, std::memory_order_release);
+  }
+
   inline size_t getNumInputs() {
     return input_ids_.size();
   }
diff --git a/backends/xnnpack/runtime/XNNWorkspaceManager.cpp b/backends/xnnpack/runtime/XNNWorkspaceManager.cpp
index d3550da5cc7..e115074a108 100644
--- a/backends/xnnpack/runtime/XNNWorkspaceManager.cpp
+++ b/backends/xnnpack/runtime/XNNWorkspaceManager.cpp
@@ -61,7 +61,9 @@ XNNWorkspaceManager::get_or_create_workspace(
       return create_result.error();
     }
 
+#ifndef XNNPACK_WORKSPACE_ALWAYS_LOCK
     create_result.get()->disable_locking();
+#endif
     return create_result.get();
   } else if (mode == WorkspaceSharingMode::PerModel) {
     return get_or_create_model_workspace(program_id);
diff --git a/backends/xnnpack/targets.bzl b/backends/xnnpack/targets.bzl
index 868e68e5b8c..b3af589df10 100644
--- a/backends/xnnpack/targets.bzl
+++ b/backends/xnnpack/targets.bzl
@@ -14,6 +14,8 @@ def _get_preprocessor_flags():
     if native.read_config("executorch", "xnnpack_weights_cache", "0") != "0":
         preprocessor_flags.append("-DENABLE_XNNPACK_WEIGHTS_CACHE")
 
+    preprocessor_flags.append("-DXNNPACK_WORKSPACE_ALWAYS_LOCK")
+
     # Enable if not disabled through config
     return preprocessor_flags