Refactor XNN workspace sharing to allow runtime gating

Author: GregoryComer

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -2237,24 +2237,22 @@ ET_NODISCARD Error XNNCompiler::compileModel(
   xnn_weights_cache_t weights_cache_ptr = nullptr;
 #endif
 
-#ifdef ENABLE_XNNPACK_SHARED_WORKSPACE
-  ET_CHECK_OR_RETURN_ERROR(
-      workspace != nullptr, Internal, "Failed to initialize XNNPACK workspace");
+  if (workspace != nullptr) {
   status = xnn_create_runtime_v4(
       subgraph.get(),
       weights_cache_ptr,
       workspace,
       ::executorch::extension::threadpool::get_pthreadpool(),
       runtime_flags,
       &runtime_ptr);
-#else
-  status = xnn_create_runtime_v3(
-      subgraph.get(),
-      weights_cache_ptr,
-      ::executorch::extension::threadpool::get_pthreadpool(),
-      runtime_flags,
-      &runtime_ptr);
-#endif
+  } else {
+    status = xnn_create_runtime_v3(
+        subgraph.get(),
+        weights_cache_ptr,
+        ::executorch::extension::threadpool::get_pthreadpool(),
+        runtime_flags,
+        &runtime_ptr);
+  }
 
   ET_CHECK_OR_RETURN_ERROR(
       xnn_status_success == status,

backends/xnnpack/runtime/XNNPACKBackend.cpp

@@ -51,21 +51,7 @@ class XnnpackBackend final
     }
 
 #ifdef ENABLE_XNNPACK_SHARED_WORKSPACE
-    // Create a workspace for the XNNExecutor to use. This workspace will be
-    // shared across all delegate instances.
-    ET_LOG(Debug, "Creating XNN workspace");
-    xnn_workspace_t workspace = nullptr;
-    status = xnn_create_workspace(&workspace);
-    if (status != xnn_status_success) {
-      ET_LOG(
-          Error,
-          "Failed to create XNN workspace, XNNPACK status: 0x%x",
-          (unsigned int)status);
-      workspace = nullptr;
-      return;
-    }
-    workspace_.reset(workspace);
-    ET_LOG(Debug, "Created XNN workspace: %p", workspace_.get());
+    enable_shared_workspace_ = true;
 #endif // ENABLE_XNNPACK_SHARED_WORKSPACE
   }
 
@@ -86,9 +72,28 @@ class XnnpackBackend final
     const NamedDataMap* named_data_map = context.get_named_data_map();
     // thread safe. This can heppen when multiple threads call init() on
     // the same backend instance.
-#ifdef ENABLE_XNNPACK_SHARED_WORKSPACE
-    const std::lock_guard<std::mutex> lock(workspace_mutex_);
-#endif
+
+    std::unique_lock<std::mutex> lock(workspace_mutex_, std::defer_lock);
+    if (enable_shared_workspace_) {
+      lock.lock();
+      if (!workspace_) {
+        // Create a workspace for the XNNExecutor to use. This workspace will be
+        // shared across all delegate instances.
+        ET_LOG(Debug, "Creating XNN workspace");
+        xnn_workspace_t workspace = nullptr;
+        auto status = xnn_create_workspace(&workspace);
+        if (status != xnn_status_success) {
+          ET_LOG(
+              Error,
+              "Failed to create XNN workspace, XNNPACK status: 0x%x",
+              (unsigned int)status);
+          workspace = nullptr;
+          return Error::Internal;
+        }
+        workspace_.reset(workspace);
+        ET_LOG(Debug, "Created XNN workspace: %p", workspace_.get());
+      }
+    }
 
 #ifdef ENABLE_XNNPACK_WEIGHTS_CACHE
     const std::lock_guard<std::mutex> lock_weight_cache(weights_cache_mutex_);
@@ -129,9 +134,10 @@ class XnnpackBackend final
       EValue** args) const override {
     auto executor = static_cast<xnnpack::delegate::XNNExecutor*>(handle);
 
-#ifdef ENABLE_XNNPACK_SHARED_WORKSPACE
-    const std::lock_guard<std::mutex> lock(workspace_mutex_);
-#endif
+    std::unique_lock<std::mutex> lock(workspace_mutex_, std::defer_lock);
+    if (enable_shared_workspace_) {
+      lock.lock();
+    }
 
 #ifdef ENABLE_XNNPACK_WEIGHTS_CACHE
     const std::lock_guard<std::mutex> lock_weights_cache(weights_cache_mutex_);
@@ -160,9 +166,10 @@ class XnnpackBackend final
       // This is needed to serialize access to xnn_delete_runtime which is not
       // thread safe. This can heppen when multiple threads call destroy() on
       // the same backend instance.
-#ifdef ENABLE_XNNPACK_SHARED_WORKSPACE
-      const std::lock_guard<std::mutex> lock(workspace_mutex_);
-#endif
+      std::unique_lock<std::mutex> lock(workspace_mutex_, std::defer_lock);
+      if (enable_shared_workspace_) {
+        lock.lock();
+      }
 
       auto executor = static_cast<xnnpack::delegate::XNNExecutor*>(handle);
 
@@ -181,10 +188,15 @@ class XnnpackBackend final
     }
   }
 
+  void set_workspace_sharing_enabled(bool enable) {
+    this->enable_shared_workspace_ = enable;
+  }
+
  private:
+  bool enable_shared_workspace_;
   // This is a global workspace for all delegate instances.
   mutable std::mutex workspace_mutex_;
-  std::unique_ptr<xnn_workspace, decltype(&xnn_release_workspace)> workspace_{
+  mutable std::unique_ptr<xnn_workspace, decltype(&xnn_release_workspace)> workspace_{
       nullptr,
       &xnn_release_workspace};
 
@@ -204,5 +216,11 @@ Backend backend{"XnnpackBackend", &cls};
 static auto success_with_compiler = register_backend(backend);
 } // namespace
 
+namespace executorch::backend::xnnpack {
+    ET_EXPERIMENTAL void set_workspace_sharing_enabled(bool enable) {
+      
+    }
+}
+
 } // namespace backends
 } // namespace executorch

backends/xnnpack/runtime/XNNPACKBackend.h

@@ -0,0 +1,5 @@
+#include <executorch/runtime/compiler.h>
+
+namespace executorch::backend::xnnpack {
+    ET_EXPERIMENTAL void set_workspace_sharing_enabled(bool enable);
+}