Fix bug with mixed weight cache + workspace sharing

Differential Revision: D106412035

Pull Request resolved: #19777

Author: GregoryComer

backends/xnnpack/runtime/XNNExecutor.cpp

@@ -93,7 +93,7 @@ ET_NODISCARD Error XNNExecutor::initialize(
  * delegate->execute()
  */
 ET_NODISCARD Error XNNExecutor::prepare_args(Span<EValue*> args) {
-  ET_CHECK_MSG(
+  ET_DCHECK_MSG(
       !destroyed_.load(std::memory_order_acquire),
       "XNNExecutor::prepare_args called after destroy");
 

backends/xnnpack/runtime/XNNExecutor.h

@@ -45,7 +45,7 @@ class XNNExecutor {
       : workspace_(workspace) {}
 
   ~XNNExecutor() {
-    ET_CHECK_MSG(
+    ET_DCHECK_MSG(
         !in_use_.load(std::memory_order_acquire),
         "XNNExecutor destroyed while in use");
     destroyed_.store(true, std::memory_order_release);

backends/xnnpack/runtime/XNNPACKBackend.cpp

@@ -16,7 +16,6 @@
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/executor/pte_data_map.h>
 
-#include <cinttypes>
 #include <memory>
 #include <mutex>
 
@@ -101,6 +100,7 @@ class XnnpackBackend final
       lock_weights_cache.lock();
       weights_cache_->initialize_for_runtime(
           context.get_runtime_allocator(), named_data_map);
+      workspace->set_uses_weight_cache();
     }
 
     auto [workspace_lock, workspace_ptr] = workspace->acquire();
@@ -131,16 +131,6 @@ class XnnpackBackend final
       return err;
     }
 
-    ET_LOG(
-        Info,
-        "XnnpackBackend::init delegate=%p workspace_id=%" PRIu64
-        " workspace_ptr=%p program_id=0x%" PRIxPTR " weight_cache=%s",
-        (void*)executor,
-        workspace->id(),
-        (void*)workspace_ptr,
-        program_id,
-        use_weight_cache ? "true" : "false");
-
     return executor;
   }
 
@@ -151,18 +141,10 @@ class XnnpackBackend final
     auto executor = static_cast<xnnpack::delegate::XNNExecutor*>(handle);
 
     auto workspace = executor->get_workspace();
-    ET_LOG(
-        Info,
-        "XnnpackBackend::execute begin delegate=%p workspace_id=%" PRIu64
-        " num_args=%zu weight_cache=%s",
-        (void*)executor,
-        workspace->id(),
-        (size_t)args.size(),
-        executor->uses_weight_cache() ? "true" : "false");
 
     std::unique_lock<std::mutex> lock_weights_cache(
         weights_cache_mutex_, std::defer_lock);
-    if (executor->uses_weight_cache()) {
+    if (executor->uses_weight_cache() || workspace->uses_weight_cache()) {
       lock_weights_cache.lock();
     }
 
@@ -183,14 +165,6 @@ class XnnpackBackend final
     // Convert output data types if necessary (e.g., int32 -> int64 for Long)
     err = executor->convert_outputs(args);
 
-    ET_LOG(
-        Info,
-        "XnnpackBackend::execute end delegate=%p workspace_id=%" PRIu64
-        " err=0x%x",
-        (void*)executor,
-        workspace->id(),
-        (unsigned int)err);
-
     return err;
   }
 
@@ -199,12 +173,6 @@ class XnnpackBackend final
       auto executor = static_cast<xnnpack::delegate::XNNExecutor*>(handle);
       auto workspace = executor->get_workspace();
 
-      ET_LOG(
-          Info,
-          "XnnpackBackend::destroy delegate=%p workspace_id=%" PRIu64,
-          (void*)executor,
-          workspace->id());
-
       const std::lock_guard<std::mutex> lock_weights_cache(
           weights_cache_mutex_);
 

backends/xnnpack/runtime/XNNWorkspace.h

@@ -59,6 +59,14 @@ class XNNWorkspace {
     lock_required_ = false;
   }
 
+  void set_uses_weight_cache() {
+    uses_weight_cache_.store(true, std::memory_order_release);
+  }
+
+  bool uses_weight_cache() const {
+    return uses_weight_cache_.load(std::memory_order_acquire);
+  }
+
   static runtime::Result<std::shared_ptr<XNNWorkspace>> create() {
     // Because this class can't be moved, we need to construct it in-place.
     xnn_workspace_t workspace = nullptr;
@@ -80,6 +88,7 @@ class XNNWorkspace {
   std::mutex mutex_;
   uint64_t id_;
   bool lock_required_ = true;
+  std::atomic<bool> uses_weight_cache_{false};
   WorkspacePtr workspace_;
 };