Support for unified recipe registry and user interfaces (#12248)

Summary:
Implements the RFC: #12248


1. `BackendRecipeProvider` -> Abstract interface that all backends must implement while providing recipes.
1. `recipe_registry` -> Singleton registry that maintains `BackendRecipeProviders`
1. `RecipeType` -> Abstract enum, backends extend this to provide support for specific recipes.
1. `RecipeBuilder` -> Helps building the recipes
1. `ExportRecipe` will have two class methods
    a. `get_recipe` -> Queries registry to get recipe w/ specific backend
    b. `recipe_builder` -> To get a recipe builder, starting with a base recipe to further customize

Differential Revision: D78034047

Author: abhinaykukkadapu

export/TARGETS

@@ -24,6 +24,7 @@ python_library(
     deps = [
         ":recipe",
         "//executorch/runtime:runtime",
+        ":recipe_registry"
     ]
 )
 
@@ -35,5 +36,30 @@ python_library(
     deps = [
         ":export",
         ":recipe",
+        ":recipe_registry",
+        ":recipe_provider"
     ],
 )
+
+
+python_library(
+    name = "recipe_registry",
+    srcs = [
+        "recipe_registry.py",
+    ],
+    deps = [
+        ":recipe",
+        ":recipe_provider"
+    ],
+)
+
+
+python_library(
+    name = "recipe_provider",
+    srcs = [
+        "recipe_provider.py",
+    ],
+    deps = [
+        ":recipe",
+    ]
+)

export/__init__.py

@@ -4,6 +4,8 @@
 # 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
+
 """
 ExecuTorch export module.
 
@@ -12,13 +14,19 @@
 export management.
 """
 
-# pyre-strict
-
 from .export import export, ExportSession
-from .recipe import ExportRecipe
+from .recipe import ExportRecipe, QuantizationRecipe, RecipeBuilder, RecipeType
+from .recipe_provider import BackendRecipeProvider
+from .recipe_registry import recipe_registry
+
 
 __all__ = [
     "ExportRecipe",
+    "QuantizationRecipe",
     "ExportSession",
     "export",
+    "BackendRecipeProvider",
+    "recipe_registry",
+    "RecipeBuilder",
+    "RecipeType",
 ]

export/export.py

@@ -1,3 +1,9 @@
+# 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.
+
 from abc import ABC, abstractmethod
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
@@ -10,11 +16,14 @@
     ExecutorchProgramManager,
     to_edge_transform_and_lower,
 )
+from executorch.exir.program._program import _transform
 from executorch.exir.schema import Program
 from executorch.extension.export_util.utils import save_pte_program
 from executorch.runtime import Runtime, Verification
 from tabulate import tabulate
 from torch import nn
+
+from torch._export.pass_base import PassType
 from torch.export import ExportedProgram
 from torchao.quantization import quantize_
 from torchao.quantization.pt2e import allow_exported_model_train_eval
@@ -95,9 +104,7 @@ class ExportStage(Stage):
 
     def __init__(
         self,
-        pre_edge_transform_passes: Optional[
-            Callable[[ExportedProgram], ExportedProgram]
-        ] = None,
+        pre_edge_transform_passes: Optional[List[PassType]] = None,
     ) -> None:
         self._exported_program: Dict[str, ExportedProgram] = {}
         self._pre_edge_transform_passes = pre_edge_transform_passes
@@ -153,10 +160,10 @@ def run(
                 )
 
                 # Apply pre-edge transform passes if available
-                if self._pre_edge_transform_passes is not None:
-                    for pre_edge_transform_pass in self._pre_edge_transform_passes:
-                        self._exported_program[method_name] = pre_edge_transform_pass(
-                            self._exported_program[method_name]
+                if pre_edge_transform_passes := self._pre_edge_transform_passes or []:
+                    for pass_ in pre_edge_transform_passes:
+                        self._exported_program[method_name] = _transform(
+                            self._exported_program[method_name], pass_
                         )
 
     def get_artifacts(self) -> Dict[str, ExportedProgram]:

export/recipe.py

@@ -13,18 +13,30 @@
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import Callable, List, Optional, Sequence
+from typing import Any, List, Optional, Sequence
 
 from executorch.exir._warnings import experimental
 
 from executorch.exir.backend.partitioner import Partitioner
 from executorch.exir.capture import EdgeCompileConfig, ExecutorchBackendConfig
 from executorch.exir.pass_manager import PassType
-from torch.export import ExportedProgram
 from torchao.core.config import AOBaseConfig
 from torchao.quantization.pt2e.quantizer import Quantizer
 
 
+class RecipeType(Enum):
+    """
+    Base recipe type class that backends can extend to define their own recipe types.
+    Backends should create their own enum classes that inherit from RecipeType:
+    Example:
+        class MyBackendRecipeType(RecipeType):
+            FP32 = "fp32"
+            INT8_DYNAMIC = "int8_dynamic"
+    """
+
+    pass
+
+
 class Mode(str, Enum):
     """
     Export mode enumeration.
@@ -52,7 +64,7 @@ class QuantizationRecipe:
     quantizers: Optional[List[Quantizer]] = None
     ao_base_config: Optional[List[AOBaseConfig]] = None
 
-    def get_quantizers(self) -> Optional[Quantizer]:
+    def get_quantizers(self) -> Optional[List[Quantizer]]:
         """
         Get the quantizer associated with this recipe.
 
@@ -89,17 +101,164 @@ class ExportRecipe:
 
     name: Optional[str] = None
     quantization_recipe: Optional[QuantizationRecipe] = None
-    edge_compile_config: Optional[EdgeCompileConfig] = (
-        None  # pyre-ignore[11]: Type not defined
-    )
-    pre_edge_transform_passes: Optional[
-        Callable[[ExportedProgram], ExportedProgram]
-        | List[Callable[[ExportedProgram], ExportedProgram]]
-    ] = None
+    # pyre-ignore[11]: Type not defined
+    edge_compile_config: Optional[EdgeCompileConfig] = None
+    pre_edge_transform_passes: Optional[Sequence[PassType]] = None
     edge_transform_passes: Optional[Sequence[PassType]] = None
     transform_check_ir_validity: bool = True
     partitioners: Optional[List[Partitioner]] = None
-    executorch_backend_config: Optional[ExecutorchBackendConfig] = (
-        None  # pyre-ignore[11]: Type not defined
-    )
+    # pyre-ignore[11]: Type not defined
+    executorch_backend_config: Optional[ExecutorchBackendConfig] = None
     mode: Mode = Mode.RELEASE
+
+    @classmethod
+    def get_recipe(
+        cls, recipe_type: "RecipeType", backend: str, **kwargs
+    ) -> "ExportRecipe":
+        """Get a clean base recipe from backend"""
+        from .recipe_registry import recipe_registry
+
+        recipe = recipe_registry.create_recipe(recipe_type, backend, **kwargs)
+        if recipe is None:
+            supported = recipe_registry.get_supported_recipes(backend)
+            raise ValueError(
+                f"Recipe '{recipe_type.value}' not supported by '{backend}'. "
+                f"Supported: {[r.value for r in supported]}"
+            )
+        return recipe
+
+    @classmethod
+    def recipe_builder(
+        cls, recipe_type: "RecipeType", backend: str, **kwargs
+    ) -> "RecipeBuilder":
+        """Create a recipe builder from a base recipe"""
+        from .recipe_registry import recipe_registry
+
+        base_recipe = recipe_registry.create_recipe(recipe_type, backend, **kwargs)
+        if base_recipe is None:
+            supported = recipe_registry.get_supported_recipes(backend)
+            raise ValueError(
+                f"Recipe '{recipe_type.value}' not supported by '{backend}'. "
+                f"Supported: {[r.value for r in supported]}"
+            )
+        return RecipeBuilder(base_recipe)
+
+
+class RecipeBuilder:
+    def __init__(
+        self, base_recipe: Optional[ExportRecipe] = None, name: str = "custom_recipe"
+    ):
+        """
+        Initialize builder with optional base recipe.
+
+        Args:
+            base_recipe: Existing recipe to customize or None to start fresh
+            name: Name for the recipe
+        """
+        if base_recipe is not None:
+            # Initialize from existing recipe
+            self._name = base_recipe.name
+            self._quantization_recipe = base_recipe.quantization_recipe
+            self._partitioners = base_recipe.partitioners
+            self._edge_compile_config = base_recipe.edge_compile_config
+            self._executorch_backend_config = base_recipe.executorch_backend_config
+            self._pre_edge_transform_passes = base_recipe.pre_edge_transform_passes
+            self._edge_transform_passes = base_recipe.edge_transform_passes
+            self._transform_check_ir_validity = getattr(
+                base_recipe, "transform_check_ir_validity", True
+            )
+            self._mode = getattr(base_recipe, "mode", None)
+        else:
+            self._name = name
+            self._quantization_recipe: Optional[QuantizationRecipe] = None
+            self._partitioners: Optional[List[Any]] = None
+            self._edge_compile_config: Optional[Any] = None
+            self._executorch_backend_config: Optional[Any] = None
+            self._pre_edge_transform_passes: Optional[Sequence[PassType]] = None
+            self._edge_transform_passes: Optional[Sequence[PassType]] = None
+            self._transform_check_ir_validity: bool = True
+            self._mode = None
+
+    def with_name(self, name: str) -> "RecipeBuilder":
+        """Set recipe name"""
+        self._name = name
+        return self
+
+    def with_quantization(
+        self, quantization_recipe: QuantizationRecipe
+    ) -> "RecipeBuilder":
+        """Set quantization recipe"""
+        self._quantization_recipe = quantization_recipe
+        return self
+
+    def with_partitioners(self, partitioners: List[Any]) -> "RecipeBuilder":
+        """Set partitioners"""
+        self._partitioners = partitioners
+        return self
+
+    def with_pre_edge_passes(self, passes: List[PassType]) -> "RecipeBuilder":
+        """Set pre-edge transform passes"""
+        self._pre_edge_transform_passes = passes
+        return self
+
+    def with_edge_passes(self, passes: Sequence[PassType]) -> "RecipeBuilder":
+        """Set edge transform passes"""
+        self._edge_transform_passes = passes
+        return self
+
+    def with_edge_compile_config(self, config: Any) -> "RecipeBuilder":
+        """Set edge compile config"""
+        self._edge_compile_config = config
+        return self
+
+    def with_backend_config(self, config: Any) -> "RecipeBuilder":
+        """Set backend config"""
+        self._executorch_backend_config = config
+        return self
+
+    def with_ir_validity_check(self, check: bool) -> "RecipeBuilder":
+        """Set IR validity check"""
+        self._transform_check_ir_validity = check
+        return self
+
+    def with_mode(self, mode: Any) -> "RecipeBuilder":
+        """Set export mode"""
+        self._mode = mode
+        return self
+
+    def add_pre_edge_pass(self, pass_: PassType) -> "RecipeBuilder":
+        """Add a single pre-edge transform pass to existing passes"""
+        if self._pre_edge_transform_passes is None:
+            self._pre_edge_transform_passes = []
+        self._pre_edge_transform_passes = list(self._pre_edge_transform_passes) + [
+            pass_
+        ]
+        return self
+
+    def add_edge_pass(self, pass_: PassType) -> "RecipeBuilder":
+        """Add a single edge transform pass to existing passes"""
+        if self._edge_transform_passes is None:
+            self._edge_transform_passes = []
+        self._edge_transform_passes = list(self._edge_transform_passes) + [pass_]
+        return self
+
+    def add_partitioner(self, partitioner: Any) -> "RecipeBuilder":
+        """Add a partitioner to existing partitioners"""
+        if self._partitioners is None:
+            self._partitioners = []
+        self._partitioners = list(self._partitioners) + [partitioner]
+        return self
+
+    def build(self) -> ExportRecipe:
+        """Build the final recipe"""
+        return ExportRecipe(
+            name=self._name,
+            quantization_recipe=self._quantization_recipe,
+            partitioners=self._partitioners,
+            edge_compile_config=self._edge_compile_config,
+            executorch_backend_config=self._executorch_backend_config,
+            pre_edge_transform_passes=self._pre_edge_transform_passes,
+            edge_transform_passes=self._edge_transform_passes,
+            transform_check_ir_validity=self._transform_check_ir_validity,
+            mode=self._mode,
+        )

export/recipe_provider.py

@@ -0,0 +1,60 @@
+# 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.
+
+# pyre-strict
+
+"""
+Recipe registry for managing backend recipe providers.
+
+This module provides the registry system for backend recipe providers and
+the abstract interface that all backends must implement.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Sequence
+
+from .recipe import ExportRecipe, RecipeType
+
+
+class BackendRecipeProvider(ABC):
+    """
+    Abstract recipe provider that all backends must implement
+    """
+
+    @property
+    @abstractmethod
+    def backend_name(self) -> str:
+        """
+        Name of the backend (ex: 'xnnpack', 'qnn' etc)
+        """
+        pass
+
+    @abstractmethod
+    def get_supported_recipes(self) -> Sequence[RecipeType]:
+        """
+        Get list of supported recipes.
+        """
+        pass
+
+    @abstractmethod
+    def create_recipe(
+        self, recipe_type: RecipeType, **kwargs: Any
+    ) -> Optional[ExportRecipe]:
+        """
+        Create a recipe for the given type.
+        Returns None if the recipe is not supported by this backend.
+
+        Args:
+            recipe_type: The type of recipe to create
+            **kwargs: Recipe-specific parameters (ex: group_size)
+
+        Returns:
+            ExportRecipe if supported, None otherwise
+        """
+        pass
+
+    def supports_recipe(self, recipe_type: RecipeType) -> bool:
+        return recipe_type in self.get_supported_recipes()