QNN: construct partitioners from options and allocate LLM I/O in rpcmem

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/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/tests/CMakeLists.txt

@@ -0,0 +1,38 @@
+# 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.
+
+# Standalone Linux host tests: cmake -S backends/qualcomm/tests -B <build>
+cmake_minimum_required(VERSION 3.24)
+project(qnn_prerequisite_tests LANGUAGES C CXX)
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+get_filename_component(
+  EXECUTORCH_ROOT "${CMAKE_CURRENT_LIST_DIR}/../../.." ABSOLUTE
+)
+set(EXECUTORCH_BUILD_QNN ON CACHE BOOL "" FORCE)
+set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON CACHE BOOL "" FORCE)
+set(EXECUTORCH_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+set(EXECUTORCH_BUILD_EXECUTOR_RUNNER OFF CACHE BOOL "" FORCE)
+add_subdirectory("${EXECUTORCH_ROOT}" executorch)
+add_subdirectory("${EXECUTORCH_ROOT}/third-party/googletest" googletest)
+enable_testing()
+
+add_executable(
+  qnn_io_buffer_allocator_test test_io_buffer_allocator.cpp
+  ../llm/io_buffer_allocator.cpp
+)
+target_include_directories(
+  qnn_io_buffer_allocator_test PRIVATE
+  "${EXECUTORCH_ROOT}/.."
+  "${EXECUTORCH_ROOT}/runtime/core/portable_type/c10"
+  "${EXECUTORCH_ROOT}/third-party/flatbuffers/include"
+  "${CMAKE_BINARY_DIR}/schema/include"
+)
+target_link_libraries(
+  qnn_io_buffer_allocator_test PRIVATE executorch_core GTest::gtest_main
+)
+include(GoogleTest)
+gtest_discover_tests(qnn_io_buffer_allocator_test)

backends/qualcomm/tests/test_io_buffer_allocator.cpp

@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+#include <executorch/backends/qualcomm/runtime/QnnExecuTorch.h>
+#include <executorch/extension/llm/runner/io_manager/io_manager.h>
+#include <gtest/gtest.h>
+
+namespace {
+void* allocated_block;
+size_t requested_bytes;
+size_t requested_alignment;
+void* freed_block;
+void* bound_tensor;
+void* bound_block;
+} // namespace
+
+void* QnnExecuTorchAllocCustomMem(size_t bytes, size_t alignment) {
+  requested_bytes = bytes;
+  requested_alignment = alignment;
+  return allocated_block;
+}
+
+void QnnExecuTorchFreeCustomMem(void* block) {
+  freed_block = block;
+}
+
+void QnnExecuTorchAddCustomMemTensorAddr(void* tensor, void* block) {
+  bound_tensor = tensor;
+  bound_block = block;
+}
+
+TEST(QnnIOBufferAllocatorTest, RegistersAndForwardsAllocationBindingAndFree) {
+  auto& factory = executorch::extension::llm::io_buffer_allocator_factory();
+  ASSERT_TRUE(factory);
+  auto allocator = factory();
+  ASSERT_NE(allocator, nullptr);
+  alignas(64) char memory[128];
+  allocated_block = memory;
+  EXPECT_EQ(allocator->allocate(sizeof(memory), 64), memory);
+  EXPECT_EQ(requested_bytes, sizeof(memory));
+  EXPECT_EQ(requested_alignment, 64);
+  allocator->bind(memory, memory + 32);
+  EXPECT_EQ(bound_tensor, memory + 32);
+  EXPECT_EQ(bound_block, memory);
+  allocator->deallocate(memory);
+  EXPECT_EQ(freed_block, memory);
+}
+
+TEST(QnnIOBufferAllocatorTest, PreservesFailedAllocationForHeapFallback) {
+  auto& factory = executorch::extension::llm::io_buffer_allocator_factory();
+  ASSERT_TRUE(factory);
+  allocated_block = nullptr;
+  auto allocator = factory();
+  EXPECT_EQ(allocator->allocate(1024, 64), nullptr);
+  EXPECT_EQ(requested_bytes, 1024);
+  EXPECT_EQ(requested_alignment, 64);
+}