[Executorch][QNN Recipes] Introduce QNN fp16 recipe

backends/qualcomm/_passes/TARGETS

@@ -12,6 +12,7 @@ runtime.python_library(
     ],
     deps = [
         "//executorch/backends/transforms:addmm_mm_to_linear",
+        "//executorch/backends/transforms:decompose_sdpa",
         "//executorch/exir/backend:backend_details",
         "//executorch/exir/backend:compile_spec_schema",
     ],

backends/qualcomm/recipes/TARGETS

@@ -0,0 +1,53 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+runtime.python_library(
+    name = "qnn_recipes",
+    srcs = [
+        "__init__.py",
+    ],
+    visibility = [
+        "//executorch/...",
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        "//executorch/export:recipe_registry",
+        ":qnn_recipe_provider",
+        ":qnn_recipe_types",
+    ],
+)
+
+runtime.python_library(
+    name = "qnn_recipe_provider",
+    srcs = [
+        "qnn_recipe_provider.py",
+    ],
+    visibility = [
+        "//executorch/...",
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/export:lib",
+        "//executorch/backends/qualcomm/partition:partition",
+        "//executorch/backends/qualcomm/serialization:serialization",
+        "//executorch/backends/qualcomm/utils:utils",
+        "//executorch/backends/qualcomm/_passes:passes",
+        ":qnn_recipe_types",
+    ],
+)
+
+runtime.python_library(
+    name = "qnn_recipe_types",
+    srcs = [
+        "qnn_recipe_types.py",
+    ],
+    visibility = [
+        "//executorch/...",
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        "//executorch/export:lib",
+    ],
+)

backends/qualcomm/recipes/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# 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.
+
+"""QNN Recipe module for ExecuTorch"""
+from executorch.export import recipe_registry
+
+from .qnn_recipe_provider import QNNRecipeProvider
+from .qnn_recipe_types import QNNRecipeType
+
+# Auto-register XNNPACK recipe provider
+recipe_registry.register_backend_recipe_provider(QNNRecipeProvider())
+
+__all__ = ["QNNRecipeProvider", "QNNRecipeType"]

backends/qualcomm/recipes/qnn_recipe_provider.py

@@ -0,0 +1,179 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# 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.
+
+# pyre-strict
+
+import logging
+from typing import Any, Optional, Sequence
+
+from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager
+from executorch.backends.qualcomm.partition.qnn_partitioner import QnnPartitioner
+from executorch.backends.qualcomm.recipes.qnn_recipe_types import (
+    QNN_BACKEND,
+    QNNRecipeType,
+)
+from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset
+from executorch.backends.qualcomm.utils.utils import (
+    generate_htp_compiler_spec,
+    generate_qnn_executorch_compiler_spec,
+    get_soc_to_chipset_map,
+    qnn_edge_config,
+)
+from executorch.export import (
+    BackendRecipeProvider,
+    ExportRecipe,
+    LoweringRecipe,
+    RecipeType,
+)
+
+
+class QNNRecipeProvider(BackendRecipeProvider):
+    @property
+    def backend_name(self) -> str:
+        return QNN_BACKEND
+
+    def get_supported_recipes(self) -> Sequence[RecipeType]:
+        return list(QNNRecipeType)
+
+    def create_recipe(
+        self, recipe_type: RecipeType, **kwargs: Any
+    ) -> Optional[ExportRecipe]:
+        """Create QNN recipe for different precisions and SoC targets"""
+
+        if recipe_type not in self.get_supported_recipes():
+            return None
+
+        self._validate_recipe_kwargs(recipe_type, kwargs)
+
+        if recipe_type == QNNRecipeType.FP16:
+            return self._build_fp16_recipe(recipe_type, kwargs)
+
+        return None
+
+    def _validate_recipe_kwargs(self, recipe_type: RecipeType, kwargs: Any) -> None:
+        """Validate kwargs for each recipe type"""
+        expected_keys = self._get_expected_keys(recipe_type)
+
+        unexpected = set(kwargs.keys()) - expected_keys
+        if unexpected:
+            logging.warning(
+                f"QNN Recipe '{recipe_type.value}' received unexpected parameters: {list(unexpected)}, ignoring them"
+            )
+
+        self._validate_soc_parameter(kwargs)
+        self._validate_partitioner_parameters(kwargs)
+
+    def _get_expected_keys(self, recipe_type: RecipeType) -> set:
+        """Get expected parameter keys for a recipe type"""
+        _ = recipe_type
+        common_keys = {
+            "soc_model",
+            "skip_node_id_set",
+            "skip_node_op_set",
+            "skip_mutable_buffer",
+        }
+        return common_keys
+
+    def _validate_soc_parameter(self, kwargs: Any) -> None:
+        """Validate soc_model parameter"""
+        if "soc_model" in kwargs:
+            soc_model = kwargs["soc_model"]
+            if isinstance(soc_model, str):
+                try:
+                    soc_model = get_soc_to_chipset_map()[soc_model]
+                    kwargs["soc_model"] = soc_model
+                except KeyError:
+                    raise ValueError(
+                        f"Invalid SoC model '{soc_model}'. Supported models: {[e.name for e in get_soc_to_chipset_map()]}"
+                    )
+            elif not isinstance(soc_model, QcomChipset):
+                raise ValueError(
+                    f"Parameter 'soc_model' must be a QcomChipset enum or string, got {type(soc_model)}"
+                )
+        else:
+            raise ValueError("Parameter 'soc_model' is required")
+
+    def _validate_partitioner_parameters(self, kwargs: Any) -> None:
+        """Validate partitioner parameters"""
+        if "skip_node_id_set" in kwargs:
+            skip_node_id_set = kwargs["skip_node_id_set"]
+            if skip_node_id_set is not None and not isinstance(skip_node_id_set, set):
+                raise ValueError(
+                    f"Parameter 'skip_node_id_set' must be a set or None, got {type(skip_node_id_set)}"
+                )
+
+        if "skip_node_op_set" in kwargs:
+            skip_node_op_set = kwargs["skip_node_op_set"]
+            if skip_node_op_set is not None and not isinstance(skip_node_op_set, set):
+                raise ValueError(
+                    f"Parameter 'skip_node_op_set' must be a set or None, got {type(skip_node_op_set)}"
+                )
+
+        if "skip_mutable_buffer" in kwargs:
+            skip_mutable_buffer = kwargs["skip_mutable_buffer"]
+            if not isinstance(skip_mutable_buffer, bool):
+                raise ValueError(
+                    f"Parameter 'skip_mutable_buffer' must be a boolean, got {type(skip_mutable_buffer)}"
+                )
+
+    def _build_fp16_recipe(
+        self,
+        recipe_type: RecipeType,
+        kwargs: Any,
+    ) -> ExportRecipe:
+        soc_model = kwargs["soc_model"]
+        skip_node_id_set = kwargs.get("skip_node_id_set", None)
+        skip_node_op_set = kwargs.get("skip_node_op_set", None)
+        skip_mutable_buffer = kwargs.get("skip_mutable_buffer", False)
+
+        lowering_recipe = self._get_qnn_lowering_recipe(
+            use_fp16=True,
+            soc_model=soc_model,
+            skip_node_id_set=skip_node_id_set,
+            skip_node_op_set=skip_node_op_set,
+            skip_mutable_buffer=skip_mutable_buffer,
+        )
+
+        return ExportRecipe(
+            name=recipe_type.value,
+            aten_transform_passes=[
+                lambda method_, ep: QnnPassManager().transform_for_export_pipeline(ep)
+            ],
+            lowering_recipe=lowering_recipe,
+        )
+
+    def _get_qnn_lowering_recipe(
+        self,
+        use_fp16: bool,
+        soc_model: QcomChipset,
+        skip_node_id_set: Optional[set] = None,
+        skip_node_op_set: Optional[set] = None,
+        skip_mutable_buffer: bool = False,
+    ) -> LoweringRecipe:
+        """Get QNN lowering recipe with optional precision and SoC target"""
+        backend_options = generate_htp_compiler_spec(use_fp16=use_fp16)
+
+        compile_specs = generate_qnn_executorch_compiler_spec(
+            soc_model=soc_model,
+            backend_options=backend_options,
+        )
+
+        partitioner = QnnPartitioner(
+            compiler_specs=compile_specs,
+            skip_node_id_set=skip_node_id_set,
+            skip_node_op_set=skip_node_op_set,
+            skip_mutable_buffer=skip_mutable_buffer,
+        )
+
+        edge_compile_config = qnn_edge_config()
+
+        return LoweringRecipe(
+            partitioners=[partitioner],
+            edge_transform_passes=[
+                lambda method_, ep: QnnPassManager().get_to_edge_transform_passes(ep)
+            ],
+            edge_compile_config=edge_compile_config,
+        )

backends/qualcomm/recipes/qnn_recipe_types.py

@@ -0,0 +1,28 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# 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.
+
+# pyre-strict
+
+from executorch.export import RecipeType
+
+
+QNN_BACKEND: str = "qnn"
+
+
+class QNNRecipeType(RecipeType):
+    """QNN-specific recipe types"""
+
+    # FP16 precision recipe, accepts kwargs:
+    # 1. soc_model
+    # 2. skip_node_id_set
+    # 3. skip_node_op_set
+    # 4. skip_mutable_buffer
+
+    FP16 = "qnn_fp16"
+
+    @classmethod
+    def get_backend_name(cls) -> str:
+        return QNN_BACKEND

export/export.py

@@ -195,12 +195,12 @@ def _build_stages(self, stages: List[StageType]) -> Dict[StageType, Stage]:
             elif stage_type == StageType.QUANTIZE:
                 stage = QuantizeStage(self._quant_recipe)
             elif stage_type == StageType.TORCH_EXPORT:
-                pre_edge_passes = None
-                if self._export_recipe.pre_edge_transform_passes is not None:
-                    pre_edge_passes = list(
-                        self._export_recipe.pre_edge_transform_passes
+                aten_transform_passes = None
+                if self._export_recipe.aten_transform_passes is not None:
+                    aten_transform_passes = list(
+                        self._export_recipe.aten_transform_passes
                     )
-                stage = TorchExportStage(pre_edge_passes)
+                stage = TorchExportStage(aten_transform_passes)
             elif stage_type == StageType.TO_EDGE_TRANSFORM_AND_LOWER:
                 stage = EdgeTransformAndLowerStage.from_recipe(self._lowering_recipe)
             elif stage_type == StageType.TO_EDGE:

export/recipe.py

@@ -7,9 +7,10 @@
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
 from enum import Enum, EnumMeta
-from typing import Callable, List, Optional, Sequence
+from typing import Callable, List, Optional
 
 import torch
+from executorch.exir import ExportedProgram
 
 from executorch.exir._warnings import experimental
 
@@ -117,12 +118,15 @@ class LoweringRecipe:
 
     Attributes:
         partitioners: Optional list of partitioners for model partitioning
-        edge_transform_passes: Optional sequence of transformation passes to apply
+        edge_transform_passes: Optional list of callables that take (method_name: str, exported_program: ExportedProgram) as arguments
+                               and return a list of passes (PassType) to be executed during lowering stages.
         edge_compile_config: Optional edge compilation configuration
     """
 
     partitioners: Optional[List[Partitioner]] = None
-    edge_transform_passes: Optional[Sequence[PassType]] = None
+    edge_transform_passes: (
+        None | List[Callable[[str, ExportedProgram], List[PassType]]]
+    ) = None
     # pyre-ignore[11]: Type not defined
     edge_compile_config: Optional[EdgeCompileConfig] = None
 
@@ -141,8 +145,8 @@ class ExportRecipe:
     Attributes:
         name: Optional name for the recipe
         quantization_recipe: Optional quantization recipe for model quantization
-        pre_edge_transform_passes: Optional function to apply transformation passes
-                                  before edge lowering
+        aten_transform_passes: Optional list of functions to apply transformation passes to the program before edge lowering.
+                               These callables are invoked to modify and return the transformed program.
         lowering_recipe: Optional lowering recipe for model lowering and partitioning
         executorch_backend_config: Optional backend configuration for ExecuTorch
         pipeline_stages: Optional list of stages to execute, defaults to a standard pipeline.
@@ -151,7 +155,9 @@ class ExportRecipe:
 
     name: Optional[str] = None
     quantization_recipe: Optional[QuantizationRecipe] = None
-    pre_edge_transform_passes: Optional[Sequence[PassType]] = None
+    aten_transform_passes: Optional[
+        List[Callable[[str, ExportedProgram], ExportedProgram]]
+    ] = None
     lowering_recipe: Optional[LoweringRecipe] = None
     # pyre-ignore[11]: Type not defined
     executorch_backend_config: Optional[ExecutorchBackendConfig] = None
@@ -240,8 +246,8 @@ def _combine_recipes(  # noqa: C901
 
         for recipe in backend_recipes:
             # Collect pre-edge transform passes
-            if recipe.pre_edge_transform_passes:
-                all_pre_edge_passes.extend(recipe.pre_edge_transform_passes)
+            if recipe.aten_transform_passes:
+                all_pre_edge_passes.extend(recipe.aten_transform_passes)
 
             # Collect partitioners from lowering recipes
             if recipe.lowering_recipe and recipe.lowering_recipe.partitioners:
@@ -307,7 +313,7 @@ def _combine_recipes(  # noqa: C901
         return cls(
             name=recipe_name,
             quantization_recipe=combined_quantization_recipe,
-            pre_edge_transform_passes=all_pre_edge_passes,
+            aten_transform_passes=all_pre_edge_passes,
             lowering_recipe=combined_lowering_recipe,
             executorch_backend_config=combined_backend_config,
         )