honor per-model packed-LHS disable in XNNPACK (#24245)

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -2131,7 +2131,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
     XNNWeightsCache* weights_cache,
     xnn_workspace_t workspace,
     const NamedDataMap* named_data_map,
-    bool use_weight_cache) {
+    const XNNCompilerOptions& options) {
   Result<XNNHeader> header = XNNHeader::Parse(buffer_pointer, num_bytes);
   const uint8_t* flatbuffer_data = nullptr;
   const uint8_t* constant_data = nullptr;
@@ -2239,7 +2239,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
     // With the weights cache the buffers land in the cache rather than in
     // unpacked_buffers, so track its list too.
     size_t prev_cached_buffers =
-        use_weight_cache ? weights_cache->get_num_unpacked_data() : 0;
+        options.use_weight_cache ? weights_cache->get_num_unpacked_data() : 0;
     err = defineTensor(
         subgraph.get(),
         remapped_ids,
@@ -2253,7 +2253,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
         named_data_map,
         unpacked_buffers,
         weights_cache,
-        use_weight_cache);
+        options.use_weight_cache);
 
     if (err != Error::Ok) {
       return err;
@@ -2268,7 +2268,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
         executor->unpacked_buffers_.push_back(std::move(unpacked_buffers[i]));
       }
       unpacked_buffers.resize(prev_buffers);
-      if (use_weight_cache) {
+      if (options.use_weight_cache) {
         weights_cache->take_unpacked_data_from(
             prev_cached_buffers, executor->unpacked_buffers_);
       }
@@ -2288,10 +2288,20 @@ ET_NODISCARD Error XNNCompiler::compileModel(
   runtime_flags |= XNN_FLAG_BASIC_PROFILING;
 #endif
 
+  if (options.disable_packed_lhs) {
+#if defined(XNN_FLAG_NO_PACKED_LHS)
+    runtime_flags |= XNN_FLAG_NO_PACKED_LHS;
+#else
+    ET_LOG(
+        Error,
+        "Ignoring disable_packed_lhs: this XNNPACK build does not support it");
+#endif
+  }
+
   xnn_runtime_t runtime_ptr = nullptr;
 
   xnn_weights_cache_t weights_cache_ptr = nullptr;
-  if (use_weight_cache) {
+  if (options.use_weight_cache) {
     ET_CHECK_OR_RETURN_ERROR(
         unpacked_buffers.size() == 0,
         Internal,
@@ -2319,7 +2329,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
       xnn_status_to_string(status));
 
   std::vector<std::string> packed_weights_names;
-  if (use_weight_cache) {
+  if (options.use_weight_cache) {
     // Constants XNNPACK does not pack were already moved to the executor in
     // the value loop above, so everything left here is safe to free.
     auto packed_weights_names_result = weights_cache->finalize_for_runtime();

backends/xnnpack/runtime/XNNCompiler.h

@@ -18,6 +18,11 @@ namespace backends {
 namespace xnnpack {
 namespace delegate {
 
+struct XNNCompilerOptions {
+  bool use_weight_cache{false};
+  bool disable_packed_lhs{false};
+};
+
 class XNNCompiler {
  public:
   // Takes Flatbuffer Serialized XNNPACK Model and rebuilds the xnn-subgraph
@@ -30,7 +35,7 @@ class XNNCompiler {
       XNNWeightsCache* weights_cache,
       xnn_workspace_t workspace,
       const NamedDataMap* named_data_map,
-      bool use_weight_cache);
+      const XNNCompilerOptions& options);
 };
 
 } // namespace delegate

backends/xnnpack/runtime/XNNExecutor.h

@@ -146,6 +146,7 @@ class XNNExecutor {
       executorch::runtime::Span<executorch::runtime::EValue*> args) const;
 
   friend class XNNCompiler;
+  friend class XNNPackedLhsTest;
 };
 
 } // namespace delegate

backends/xnnpack/runtime/XNNPACKBackend.cpp

@@ -90,12 +90,15 @@ class XnnpackBackend final
     }
     auto workspace = workspace_result.get();
 
-    bool use_weight_cache = options_.resolve_weight_cache(context);
+    const xnnpack::delegate::XNNCompilerOptions compiler_options{
+        .use_weight_cache = options_.resolve_weight_cache(context),
+        .disable_packed_lhs = options_.resolve_disable_packed_lhs(context),
+    };
     // Per-path weights cache: same-path PTEs share an instance and
     // serialize on its mutex; different paths run in parallel.
     std::shared_ptr<xnnpack::delegate::XNNWeightsCache> weights_cache;
     std::unique_lock<std::mutex> lock_weights_cache;
-    if (use_weight_cache) {
+    if (compiler_options.use_weight_cache) {
       // Only honor a path coming through runtime_spec (per-PTE opt-in).
       // Reading the backend-singleton global would let a non-opt-in PTE
       // inherit another model's cache file.
@@ -131,7 +134,7 @@ class XnnpackBackend final
         weights_cache.get(),
         workspace_ptr,
         named_data_map,
-        use_weight_cache);
+        compiler_options);
     // This backend does not need its processed data after compiling the model.
     processed->Free();
 
@@ -147,7 +150,7 @@ class XnnpackBackend final
 
     // Hand the cache to the executor (held by shared_ptr so it
     // outlives any sibling executors that share it).
-    if (use_weight_cache) {
+    if (compiler_options.use_weight_cache) {
       executor->set_weights_cache(std::move(weights_cache));
     }
 

backends/xnnpack/runtime/XNNPACKBackend.h

@@ -19,6 +19,13 @@ const char workspace_sharing_mode_option_key[] = "workspace_sharing_mode";
 // across delegate instances. Changes only affect subsequently loaded models.
 const char weight_cache_option_key[] = "weight_cache_enabled";
 
+/// The key for disabling packed-LHS graph rewrites. When enabled, XNNPACK
+/// preserves explicit dynamic-quantization nodes instead of folding them into
+/// packed-LHS consumers.
+// Must remain a C array for the BackendOptions template overloads.
+// @lint-ignore CLANGTIDY facebook-hte-CArray
+const char disable_packed_lhs_option_key[] = "disable_packed_lhs";
+
 /// Path for the packed weight file. When set, reserve_space() allocates from
 /// a MAP_SHARED file instead of heap; msync makes pages clean on iOS.
 // Must remain a C array (not const char*) so it can bind to the

backends/xnnpack/runtime/XnnpackBackendOptions.cpp

@@ -37,6 +37,8 @@ Error XnnpackBackendOptions::get_option(BackendOption& option) const {
     option.value = static_cast<int>(sharing_mode_.load());
   } else if (strcmp(option.key, weight_cache_option_key) == 0) {
     option.value = weight_cache_enabled_.load();
+  } else if (strcmp(option.key, disable_packed_lhs_option_key) == 0) {
+    option.value = disable_packed_lhs_.load();
   } else if (strcmp(option.key, packed_cache_path_option_key) == 0) {
     std::array<char, runtime::kMaxOptionValueLength> arr{};
     const std::lock_guard<std::mutex> lock(path_mutex_);
@@ -73,6 +75,14 @@ Error XnnpackBackendOptions::set_option(const BackendOption& option) {
     }
     ET_LOG(Debug, "Setting XNNPACK weight cache enabled to %d.", *val);
     weight_cache_enabled_.store(*val);
+  } else if (strcmp(option.key, disable_packed_lhs_option_key) == 0) {
+    auto* val = std::get_if<bool>(&option.value);
+    if (!val) {
+      ET_LOG(Error, "XNNPACK disable packed LHS must be a bool.");
+      return Error::InvalidArgument;
+    }
+    ET_LOG(Debug, "Setting XNNPACK disable packed LHS to %d.", *val);
+    disable_packed_lhs_.store(*val);
   } else if (strcmp(option.key, packed_cache_path_option_key) == 0) {
     auto* val = std::get_if<std::array<char, runtime::kMaxOptionValueLength>>(
         &option.value);
@@ -126,6 +136,12 @@ bool XnnpackBackendOptions::resolve_weight_cache(
       context, weight_cache_option_key, weight_cache_enabled_.load());
 }
 
+bool XnnpackBackendOptions::resolve_disable_packed_lhs(
+    const ET_RUNTIME_NAMESPACE::BackendInitContext& context) const {
+  return resolve_option<bool>(
+      context, disable_packed_lhs_option_key, disable_packed_lhs_.load());
+}
+
 runtime::Result<WorkspaceSharingMode>
 XnnpackBackendOptions::resolve_sharing_mode(
     const ET_RUNTIME_NAMESPACE::BackendInitContext& context) const {

backends/xnnpack/runtime/XnnpackBackendOptions.h

@@ -37,6 +37,11 @@ class XnnpackBackendOptions {
   bool resolve_weight_cache(
       const ET_RUNTIME_NAMESPACE::BackendInitContext& context) const;
 
+  // Resolve whether packed-LHS graph rewrites are disabled for a delegate,
+  // applying any runtime spec override.
+  bool resolve_disable_packed_lhs(
+      const ET_RUNTIME_NAMESPACE::BackendInitContext& context) const;
+
   // Resolve the effective workspace sharing mode for a delegate, applying
   // any runtime spec override. Returns InvalidArgument for out-of-range values.
   runtime::Result<WorkspaceSharingMode> resolve_sharing_mode(
@@ -95,6 +100,8 @@ class XnnpackBackendOptions {
   std::atomic<bool> weight_cache_enabled_{false};
 #endif
 
+  std::atomic<bool> disable_packed_lhs_{false};
+
   // The most-recently-set packed cache path. Today's callers (Cria
   // runner) set this via set_option(packed_cache_path_option_key) per
   // PTE, then invoke init() which reads it and calls

backends/xnnpack/test/CMakeLists.txt

@@ -36,6 +36,7 @@ et_cxx_test(
 target_include_directories(
   backends_xnnpack_test
   PRIVATE ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/XNNPACK/include
+          ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/XNNPACK
           ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/XNNPACK/src
           ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/cpuinfo/include
           ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/pthreadpool/include

backends/xnnpack/test/runtime/test_weight_cache.cpp

@@ -20,6 +20,7 @@
 
 using namespace ::testing;
 
+using executorch::backends::xnnpack::disable_packed_lhs_option_key;
 using executorch::backends::xnnpack::get_packed_cache_report;
 using executorch::backends::xnnpack::packed_cache_path_option_key;
 using executorch::backends::xnnpack::PackedCacheHeapReason;
@@ -53,6 +54,24 @@ static void set_and_check_weight_cache_enabled(bool enabled) {
   ASSERT_EQ(std::get<bool>(read_option.value), enabled);
 }
 
+static void set_and_check_disable_packed_lhs(bool disabled) {
+  executorch::runtime::runtime_init();
+
+  BackendOptions<1> backend_options;
+  backend_options.set_option(disable_packed_lhs_option_key, disabled);
+
+  auto status = executorch::runtime::set_option(
+      xnnpack_backend_key, backend_options.view());
+  ASSERT_EQ(status, Error::Ok);
+
+  BackendOption read_option;
+  strcpy(read_option.key, disable_packed_lhs_option_key);
+  read_option.value = !disabled;
+  status = get_option(xnnpack_backend_key, read_option);
+
+  ASSERT_EQ(std::get<bool>(read_option.value), disabled);
+}
+
 static TensorPtr create_input_tensor(float val) {
   std::vector<float> data(1000, val);
   return executorch::extension::make_tensor_ptr({10, 10, 10}, std::move(data));
@@ -101,6 +120,22 @@ TEST(WeightCache, SetInvalidType) {
   ASSERT_EQ(status, Error::InvalidArgument);
 }
 
+TEST(PackedLhs, SetDisabled) {
+  set_and_check_disable_packed_lhs(true);
+  set_and_check_disable_packed_lhs(false);
+}
+
+TEST(PackedLhs, SetInvalidType) {
+  executorch::runtime::runtime_init();
+
+  BackendOptions<1> backend_options;
+  backend_options.set_option(disable_packed_lhs_option_key, 1);
+
+  auto status = executorch::runtime::set_option(
+      xnnpack_backend_key, backend_options.view());
+  ASSERT_EQ(status, Error::InvalidArgument);
+}
+
 TEST(WeightCache, SetMultipleOptions) {
   executorch::runtime::runtime_init();