QNN: bind delegate outputs in exported program order

backends/qualcomm/CMakeLists.txt

@@ -369,6 +369,24 @@ install(
   RUNTIME DESTINATION ${_qnn_install_dir}
 )
 
+# rpcmem I/O buffers for the LLM runner's IOManagers. It registers itself, so
+# runners link it whole-archive.
+if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER)
+  add_library(
+    qnn_llm_io_buffer_allocator STATIC
+    ${QNN_EXECUTORCH_ROOT_DIR}/llm/io_buffer_allocator.cpp
+  )
+  target_link_libraries(
+    qnn_llm_io_buffer_allocator PUBLIC qnn_executorch_backend executorch_core
+                                       extension_module_static extension_tensor
+  )
+  install(
+    TARGETS qnn_llm_io_buffer_allocator
+    EXPORT ExecuTorchTargets
+    ARCHIVE DESTINATION ${_qnn_install_dir}
+  )
+endif()
+
 if(${CMAKE_SYSTEM_PROCESSOR} MATCHES Hexagon)
   add_subdirectory(
     ${QNN_EXECUTORCH_ROOT_DIR}/runtime/backends/direct_mode

backends/qualcomm/builders/node_visitor.py

@@ -37,6 +37,7 @@
 from .utils import (
     deduce_dtype,
     get_parameter,
+    graph_output_position,
     is_graph_input,
     is_graph_output,
     is_mutable_buffer_input,
@@ -465,6 +466,8 @@ def get_tensor_name(
         # the input order between QNN and the original graph’s forward function may differ.
         # The `mutbuf_{id}` is utilized for mapping I/O of mutable buffer at runtime.
         # The `output_` is identified as the graph’s output at runtime to prevent confusion with per_tensor_dump.
+        # The `output_{id}` is the position among the graph's outputs, sorted on at runtime like `input_{id}`:
+        # the same passes can also reorder the ops that produce the outputs.
 
         # For more info about self.is_qnn_partitioner, check constructor for explanation.
         if self.is_qnn_partitioner:
@@ -482,7 +485,9 @@ def get_tensor_name(
                 ).index(node.name)
                 tensor_name = f"output_mutbuf_{position_index}_{tensor_name}"
             elif is_graph_output(node, wrapper_idx):
-                tensor_name = f"output_{tensor_name}"
+                position = graph_output_position(node, wrapper_idx)
+                prefix = "output_" if position is None else f"output_{position}_"
+                tensor_name = f"{prefix}{tensor_name}"
 
         # Only add qcom_tensor_name when enable tensor dump.
         # Only runs in qnn_preprocess (not op validation) since that's when

backends/qualcomm/builders/utils.py

@@ -141,6 +141,27 @@ def is_graph_output(node: torch.fx.Node, output_index: Optional[int] = None) ->
     return False
 
 
+def graph_output_position(
+    node: torch.fx.Node, output_index: Optional[int] = None
+) -> Optional[int]:
+    """The position of ``node`` (its ``output_index``-th output, for a
+    multi-output node) among the graph's outputs, or None."""
+    output = next(n for n in node.graph.nodes if n.op == "output")
+    outputs = output.args[0]
+    outputs = outputs if isinstance(outputs, (list, tuple)) else [outputs]
+    for position, value in enumerate(outputs):
+        if value is node:
+            return position
+        if (
+            isinstance(value, torch.fx.Node)
+            and getattr(value.target, "__name__", "") == "getitem"
+            and value.args[0] is node
+            and (output_index is None or value.args[1] == output_index)
+        ):
+            return position
+    return None
+
+
 def is_mutable_buffer_output(
     tensor: torch.fx.Node, edge_program: torch.export.ExportedProgram
 ) -> bool:

backends/qualcomm/llm/BUCK

@@ -0,0 +1,16 @@
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target")
+oncall("executorch")
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain xplat-only targets.
+
+load(":targets.bzl", "define_common_targets")
+
+
+non_fbcode_target(_kind = define_common_targets,)
+
+# !!!! fbcode/executorch/backends/qualcomm/runtime/TARGETS was merged into this file, see https://fburl.com/workplace/xl8l9yuo for more info !!!!
+
+load(":targets.bzl", "define_common_targets")
+
+
+fbcode_target(_kind = define_common_targets,)

backends/qualcomm/llm/io_buffer_allocator.cpp

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Puts LLM runner I/O buffers (KV cache, mask, logits) in rpcmem, so HTP reads
+// and writes them in place. Linking this library is enough: it registers
+// itself as the runner's I/O buffer allocator. Programs must be lowered with
+// shared buffers enabled for QNN to use the memory without copying.
+
+#include <executorch/backends/qualcomm/runtime/QnnExecuTorch.h>
+#include <executorch/extension/llm/runner/io_manager/io_manager.h>
+
+namespace executorch {
+namespace backends {
+namespace qnn {
+namespace {
+
+using ::executorch::extension::llm::IOBufferAllocator;
+
+class SharedBufferAllocator final : public IOBufferAllocator {
+ public:
+  void* allocate(size_t nbytes, size_t alignment) override {
+    return QnnExecuTorchAllocCustomMem(nbytes, alignment);
+  }
+  void deallocate(void* block) override {
+    QnnExecuTorchFreeCustomMem(block);
+  }
+  void bind(void* block, void* tensor_data) override {
+    QnnExecuTorchAddCustomMemTensorAddr(tensor_data, block);
+  }
+};
+
+const bool registered = ::executorch::extension::llm::set_io_buffer_allocator(
+    [] { return std::make_unique<SharedBufferAllocator>(); });
+
+} // namespace
+} // namespace qnn
+} // namespace backends
+} // namespace executorch

backends/qualcomm/llm/targets.bzl

@@ -0,0 +1,26 @@
+load(
+    "@fbsource//tools/build_defs:default_platform_defs.bzl",
+    "ANDROID",
+)
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+def define_common_targets():
+    """Defines targets that should be shared between fbcode and xplat.
+
+    The directory containing this targets.bzl file should also contain both
+    TARGETS and BUCK files that call this function.
+    """
+
+    # rpcmem I/O buffers for the LLM runner's IOManagers. Registers itself on
+    # load, so link it whole.
+    runtime.cxx_library(
+        name = "io_buffer_allocator",
+        srcs = ["io_buffer_allocator.cpp"],
+        link_whole = True,
+        platforms = [ANDROID],
+        deps = [
+            "//executorch/backends/qualcomm/runtime:runtime_android_build",
+            "//executorch/extension/llm/runner/io_manager:io_manager",
+        ],
+        visibility = ["PUBLIC"],
+    )

backends/qualcomm/partition/qnn_partitioner.py

@@ -4,11 +4,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 import copy
+import inspect
 import logging
-from collections import defaultdict
+from collections import defaultdict, OrderedDict
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
+from executorch.backends.qualcomm._passes.qnn_pass_manager import (
+    get_qnn_pass_manager_cls,
+)
 from executorch.backends.qualcomm.builders import node_visitor_manager
 from executorch.backends.qualcomm.builders.node_visitor import (
     PER_CHANNEL_GROUP_ENCODING,
@@ -41,6 +45,7 @@
     PartitionResult,
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
+from executorch.exir.pass_manager import PassType
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
@@ -199,14 +204,25 @@ def __init__(
         skip_node_id_set: set = None,
         skip_node_op_set: set = None,
         skip_mutable_buffer: bool = False,
+        dep_table: Optional[Dict] = None,
+        passes_job: Optional[OrderedDict] = None,
+        convert_linear_to_conv2d: bool = False,
     ):
         """
         Args:
             compiler_specs (List[CompileSpec]): Backend compiler specifications.
             skip_node_id_set (set, optional): Set of node IDs to exclude from partitioning.
             skip_node_op_set (set, optional): Set of OpOverload to exclude from partitioning.
             skip_mutable_buffer (bool, optional): If True, mutable buffers are not delegated to QNN.
+            dep_table (Dict, optional): Dependency table for the transformation passes.
+            passes_job (OrderedDict, optional): Ordered dictionary of transformation passes.
+            convert_linear_to_conv2d (bool, optional): Convert linear to conv2d where that
+                improves HTP performance.
         """
+        self.compiler_specs = compiler_specs
+        self.dep_table = dep_table
+        self.passes_job = passes_job
+        self.convert_linear_to_conv2d = convert_linear_to_conv2d
         self.compiler_specs_snapshot = copy.deepcopy(compiler_specs)
         self.backend = flatbuffer_to_option(
             generate_qnn_executorch_option(self.compiler_specs_snapshot)
@@ -220,6 +236,82 @@ def __init__(
         self.skip_node_op_set = set() if skip_node_op_set is None else skip_node_op_set
         self.skip_mutable_buffer = skip_mutable_buffer
 
+    @classmethod
+    def from_options(
+        cls,
+        soc_model: str,
+        backend: str = "htp",
+        **options: Any,
+    ) -> "QnnPartitioner":
+        """Build a QnnPartitioner from plain options, generating its compiler specs.
+
+        Lets a caller that only knows option names and values (for example an
+        exporter forwarding a ``dict`` of vendor settings) get a working
+        partitioner without importing the QNN compiler-spec builders.
+
+        Args:
+            soc_model (str): Target SoC name, e.g. ``"SM8750"``.
+            backend (str): QNN backend, one of ``"htp"``, ``"gpu"`` or ``"lpai"``.
+            **options: Each option is routed by name to the first of these that
+                accepts it: ``QnnPartitioner.__init__`` (``skip_node_id_set``,
+                ``passes_job``, ``convert_linear_to_conv2d``, ...),
+                ``generate_qnn_executorch_compiler_spec`` (``use_mha2sha``,
+                ``shared_buffer``, ...), then ``generate_<backend>_compiler_spec``
+                (``use_fp16``, ``htp_performance_mode``, ...). An option none of
+                them accepts raises ``TypeError``.
+        """
+        from executorch.backends.qualcomm.utils import utils as qnn_utils
+
+        backend_spec_fn = getattr(qnn_utils, f"generate_{backend}_compiler_spec", None)
+        if backend_spec_fn is None:
+            raise ValueError(f"Unsupported QNN backend {backend!r}.")
+
+        def accepted(fn):
+            return set(inspect.signature(fn).parameters) - {"self"}
+
+        partitioner_kwargs, compiler_spec_kwargs, backend_kwargs = {}, {}, {}
+        for name, value in options.items():
+            if name in accepted(cls.__init__) - {"compiler_specs"}:
+                partitioner_kwargs[name] = value
+            elif name in accepted(qnn_utils.generate_qnn_executorch_compiler_spec) - {
+                "soc_model",
+                "backend_options",
+            }:
+                compiler_spec_kwargs[name] = value
+            elif name in accepted(backend_spec_fn):
+                backend_kwargs[name] = value
+            else:
+                raise TypeError(
+                    f"Unknown QnnPartitioner option {name!r} for backend {backend!r}."
+                )
+
+        compiler_specs = qnn_utils.generate_qnn_executorch_compiler_spec(
+            soc_model=qnn_utils.get_soc_to_chipset_map()[soc_model],
+            backend_options=backend_spec_fn(**backend_kwargs),
+            **compiler_spec_kwargs,
+        )
+        return cls(compiler_specs, **partitioner_kwargs)
+
+    def transform_for_pre_decomposition(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgram:
+        """Run the QNN ATen-level passes before default decompositions."""
+        return get_qnn_pass_manager_cls(self.backend)().transform_for_export_pipeline(
+            exported_program
+        )
+
+    def transform_passes(self, exported_program: ExportedProgram) -> List[PassType]:
+        """Run the QNN edge-level passes before partitioning."""
+        return get_qnn_pass_manager_cls(self.backend)().get_to_edge_transform_passes(
+            exported_program,
+            passes_job=self.passes_job,
+            dep_table=self.dep_table,
+            compiler_specs=self.compiler_specs,
+            skip_node_id_set=self.skip_node_id_set,
+            skip_node_op_set=self.skip_node_op_set,
+            convert_linear_to_conv2d=self.convert_linear_to_conv2d,
+        )
+
     def generate_partitions(
         self, edge_program: torch.export.ExportedProgram
     ) -> List[Any]:

backends/qualcomm/runtime/QnnManager.cpp

@@ -14,6 +14,7 @@
 #include <executorch/backends/qualcomm/runtime/backends/QnnCustomProtocol.h>
 #include <executorch/backends/qualcomm/runtime/backends/QnnImplementation.h>
 #include <algorithm>
+#include <charconv>
 #include <cstdlib>
 #include <cstring>
 #include <fstream>
@@ -36,6 +37,35 @@ bool CompareExportedInput(
   return numA < numB;
 }
 
+// The graph output position AOT encodes as `output_{id}_...`, or -1 for an
+// output named without one (mutable buffers, older programs).
+int ExportedOutputPosition(const std::string& name) {
+  const std::string prefix = "output_";
+  if (name.rfind(prefix, 0) != 0) {
+    return -1;
+  }
+  const size_t end = name.find('_', prefix.size());
+  if (end == std::string::npos || end == prefix.size() ||
+      end + 1 == name.size()) {
+    return -1;
+  }
+  int position = -1;
+  const auto result =
+      std::from_chars(name.data() + prefix.size(), name.data() + end, position);
+  return result.ec == std::errc{} && result.ptr == name.data() + end &&
+          position >= 0
+      ? position
+      : -1;
+}
+
+bool CompareExportedOutput(
+    const std::shared_ptr<TensorWrapper>& a,
+    const std::shared_ptr<TensorWrapper>& b) {
+  const int pa = ExportedOutputPosition(a->GetName());
+  const int pb = ExportedOutputPosition(b->GetName());
+  return pa >= 0 && (pb < 0 || pa < pb);
+}
+
 int ExtractMutableBufferNumber(const std::string& name) {
   std::string prefix = "mutbuf_";
   size_t startPos = name.find(prefix);
@@ -389,6 +419,14 @@ Error QnnManager::AllocateTensor(const std::string& graph_name) {
     }
     output_tensors_[graph_name].emplace_back(std::move(tensor_wrapper));
   }
+  if (!options_->is_from_context_binary()) {
+    // Outputs with a position go in the order the program's outputs bind;
+    // the rest keep their place after them.
+    std::stable_sort(
+        output_tensors_[graph_name].begin(),
+        output_tensors_[graph_name].end(),
+        CompareExportedOutput);
+  }
   return Error::Ok;
 }
 
@@ -406,6 +444,12 @@ Error QnnManager::AllocateTensor(
         CompareExportedInput);
   }
   output_tensors_[graph_name] = std::move(outputs);
+  if (!options_->is_from_context_binary()) {
+    std::stable_sort(
+        output_tensors_[graph_name].begin(),
+        output_tensors_[graph_name].end(),
+        CompareExportedOutput);
+  }
   return Error::Ok;
 }