Introduce CompileMode compatibility API (pytorch#21721)

Differential Revision: D115453091

Pull Request resolved: pytorch#21721

Author: mcremon-meta

backends/cadence/aot/compiler.py

@@ -45,7 +45,12 @@
 from torch.export.exported_program import ExportedProgram
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e
 
-from .pass_utils import EdgePassesConfig
+from .pass_utils import (
+    CompileMode,
+    EdgePassesConfig,
+    normalize_compile_mode,
+    resolve_opt_level,
+)
 from .passes import apply_exir_ops_passes, apply_torch_ops_passes
 from .utils import print_ops_info
 
@@ -364,15 +369,16 @@ def quantize_and_export_to_edge(
 def _lower_ep_to_cadence(
     program: ExportedProgram,
     dump_graphs: bool = False,
-    opt_level: int = 1,
+    mode: CompileMode | int | None = None,
     edge_passes_config: Optional[EdgePassesConfig] = None,
+    opt_level: Optional[int] = None,
 ) -> EdgeProgramManager:
     """
     Lower an existing ExportedProgram to edge IR and apply frontend optimization passes.
     """
     edge_prog_manager = _lower_ep_to_edge(program, dump_graphs=dump_graphs)
     cadence_prog_manager = apply_exir_ops_passes(
-        opt_level, edge_prog_manager, edge_passes_config
+        resolve_opt_level(mode, opt_level), edge_prog_manager, edge_passes_config
     )
     return cadence_prog_manager
 
@@ -381,12 +387,13 @@ def export_to_cadence(
     model: torch.nn.Module,
     inputs: tuple[object, ...],
     dump_graphs: bool = False,
-    opt_level: int = 1,
+    mode: CompileMode | int | None = None,
     edge_passes_config: Optional[EdgePassesConfig] = None,
+    opt_level: Optional[int] = None,
 ) -> EdgeProgramManager:
     edge_prog_manager = export_to_edge(model, inputs, dump_graphs=dump_graphs)
     cadence_prog_manager = apply_exir_ops_passes(
-        opt_level, edge_prog_manager, edge_passes_config
+        resolve_opt_level(mode, opt_level), edge_prog_manager, edge_passes_config
     )
     return cadence_prog_manager
 
@@ -395,8 +402,9 @@ def quantize_and_export_to_cadence(
     model: torch.nn.Module,
     inputs: tuple[object, ...],
     dump_graphs: bool = False,
-    opt_level: int = 1,
+    mode: CompileMode | int | None = None,
     edge_passes_config: Optional[EdgePassesConfig] = None,
+    opt_level: Optional[int] = None,
 ) -> EdgeProgramManager:
     """
     Trace, quantize, lower a model/inputs pair to edge IR and apply frontend
@@ -406,6 +414,7 @@ def quantize_and_export_to_cadence(
 
     return _lower_ep_to_cadence(
         quantized_model,
+        mode=mode,
         opt_level=opt_level,
         dump_graphs=dump_graphs,
         edge_passes_config=edge_passes_config,
@@ -416,17 +425,19 @@ def export_to_executorch_gen_etrecord(
     model: torch.nn.Module,
     inputs: tuple[object, ...],
     output_dir: Optional[str] = None,
-    opt_level: int = 1,
+    mode: CompileMode | int | None = None,
     mem_algo: int = 0,
     alloc_graph_input: bool = True,
     alloc_graph_output: bool = True,
     memory_config: Optional[MemoryConfig] = None,
     dump_graphs: bool = False,
+    opt_level: Optional[int] = None,
 ) -> ExecutorchProgramManager:
     ep = torch.export.export(model, inputs, strict=True)
     return _lower_ep_to_cadence_gen_etrecord(
         ep,
         output_dir=output_dir,
+        mode=mode,
         opt_level=opt_level,
         mem_algo=mem_algo,
         alloc_graph_input=alloc_graph_input,
@@ -442,15 +453,18 @@ def export_to_executorch_gen_etrecord(
 def _lower_ep_to_cadence_gen_etrecord(
     ep: ExportedProgram,
     output_dir: Optional[str] = None,
-    opt_level: int = 1,
+    mode: CompileMode | int | None = None,
     mem_algo: int = 0,
     alloc_graph_input: bool = True,
     alloc_graph_output: bool = True,
     memory_config: Optional[MemoryConfig] = None,
     dump_graphs: bool = False,
+    opt_level: Optional[int] = None,
 ) -> ExecutorchProgramManager:
+    normalized_mode = normalize_compile_mode(mode, opt_level)
+    resolved_opt_level = resolve_opt_level(normalized_mode)
     edge_prog_manager = _lower_ep_to_edge(ep, dump_graphs)
-    cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager)
+    cadence_prog_manager = apply_exir_ops_passes(resolved_opt_level, edge_prog_manager)
 
     # Print some information to terminal
     print_ops_info(
@@ -463,7 +477,7 @@ def _lower_ep_to_cadence_gen_etrecord(
 
     memory_planning_pass = CadenceMemoryPlanning(
         memory_config,
-        opt_level=opt_level,
+        mode=normalized_mode,
         mem_algo=mem_algo,
         alloc_graph_input=alloc_graph_input,
         alloc_graph_output=alloc_graph_output,
@@ -483,7 +497,7 @@ def _lower_ep_to_cadence_gen_etrecord(
     print_memory_planning_info(
         exec_prog,
         memory_config,
-        opt_level,
+        normalized_mode,
         alloc_graph_input,
         alloc_graph_output,
     )

backends/cadence/aot/memory_constraints.py

@@ -16,8 +16,11 @@
 import torch.fx
 from executorch.backends.cadence.aot.pass_utils import (
     CadencePassAttribute,
+    CompileMode,
     create_cadence_pass_filter,
+    normalize_compile_mode,
     register_cadence_pass,
+    resolve_opt_level,
 )
 from executorch.backends.cadence.aot.utils import get_shape, is_node_in_flattened_output
 from executorch.exir import memory
@@ -63,11 +66,16 @@ class MemConstraints:
 
     def __init__(
         self,
-        opt_level: int = 1,
+        mode: CompileMode | int | None = None,
         alloc_graph_input: bool = True,
         alloc_graph_output: bool = True,
+        opt_level: Optional[int] = None,
     ) -> None:
-        self.opt_level = opt_level
+        normalized_mode = normalize_compile_mode(mode, opt_level)
+        self.opt_level = resolve_opt_level(normalized_mode)
+        self.mode: Optional[CompileMode] = (
+            normalized_mode if isinstance(normalized_mode, CompileMode) else None
+        )
         self.alloc_graph_input = alloc_graph_input
         self.alloc_graph_output = alloc_graph_output
 

backends/cadence/aot/memory_planning.py

@@ -19,6 +19,11 @@
     MemoryPlanningAlgo,
     MemoryPlanningState,
 )
+from executorch.backends.cadence.aot.pass_utils import (
+    CompileMode,
+    normalize_compile_mode,
+    resolve_opt_level,
+)
 from executorch.backends.cadence.aot.utils import (
     MemoryConfig,
     MemoryPlanningAlgoFailure,
@@ -287,13 +292,15 @@ def find_peak_memory_usage(
 def print_memory_planning_info(
     executorch_prog: ExecutorchProgramManager,
     memory_config: MemoryConfig,
-    opt_level: int,
-    alloc_graph_input: bool,
-    alloc_graph_output: bool,
+    mode: CompileMode | int | None = None,
+    alloc_graph_input: bool = True,
+    alloc_graph_output: bool = True,
     log_level=logging.INFO,
+    opt_level: Optional[int] = None,
 ) -> None:
     # Get the peak memory usages per memory space
     mem_constraints = MemConstraints(
+        mode=mode,
         opt_level=opt_level,
         alloc_graph_input=alloc_graph_input,
         alloc_graph_output=alloc_graph_output,
@@ -397,20 +404,25 @@ class CadenceMemoryPlanning:
     def __init__(
         self,
         memory_config: MemoryConfig,
-        opt_level: int,
-        mem_algo: int,
+        mode: CompileMode | int | None = None,
+        mem_algo: int = 0,
         alloc_graph_input: bool = True,
         alloc_graph_output: bool = True,
         additional_constraint_gen_passes: Optional[Sequence[ConstraintsGenPass]] = None,
+        opt_level: Optional[int] = None,
     ) -> None:
         self.memory_config = memory_config
-        self.opt_level = opt_level
+        normalized_mode = normalize_compile_mode(mode, opt_level)
+        self.opt_level = resolve_opt_level(normalized_mode)
+        self.mode: Optional[CompileMode] = (
+            normalized_mode if isinstance(normalized_mode, CompileMode) else None
+        )
         self.alloc_graph_input = alloc_graph_input
         self.alloc_graph_output = alloc_graph_output
 
         self.algo: MemoryPlanningAlgo = self.get_mem_algos(
             memory_config,
-            opt_level,
+            normalized_mode,
             alloc_graph_input,
             alloc_graph_output,
             additional_constraint_gen_passes,
@@ -419,22 +431,24 @@ def __init__(
     @staticmethod
     def get_mem_algos(
         memory_config: MemoryConfig,
-        opt_level: int,
-        alloc_graph_input: bool,
-        alloc_graph_output: bool,
-        additional_constraint_gen_passes: Optional[Sequence[ConstraintsGenPass]],
+        mode: CompileMode | int | None = None,
+        alloc_graph_input: bool = True,
+        alloc_graph_output: bool = True,
+        additional_constraint_gen_passes: Optional[Sequence[ConstraintsGenPass]] = None,
+        opt_level: Optional[int] = None,
     ) -> list[MemoryPlanningAlgo]:
+        normalized_mode = normalize_compile_mode(mode, opt_level)
         return [
             PositionBasedGreedyWithHierarchy(
                 memory_config=memory_config,
-                opt_level=opt_level,
+                mode=normalized_mode,
                 alloc_graph_input=alloc_graph_input,
                 alloc_graph_output=alloc_graph_output,
                 additional_constraint_gen_passes=additional_constraint_gen_passes,
             ),
             GreedyWithHeuristic(
                 memory_config=memory_config,
-                opt_level=opt_level,
+                mode=normalized_mode,
                 alloc_graph_input=alloc_graph_input,
                 alloc_graph_output=alloc_graph_output,
                 additional_constraint_gen_passes=additional_constraint_gen_passes,

backends/cadence/aot/memory_planning_algo.py

@@ -15,6 +15,11 @@
     GenerateMemConstraints,
     MemConstraints,
 )
+from executorch.backends.cadence.aot.pass_utils import (
+    CompileMode,
+    normalize_compile_mode,
+    resolve_opt_level,
+)
 from executorch.backends.cadence.aot.utils import MemoryConfig
 from executorch.exir.memory_planning import Verifier
 from executorch.exir.tensor import TensorSpec
@@ -85,16 +90,21 @@ class MemoryPlanningAlgo(ABC):
     def __init__(
         self,
         memory_config: MemoryConfig,
-        opt_level: int = 1,
+        mode: CompileMode | int | None = None,
         alloc_graph_input: bool = True,
         alloc_graph_output: bool = True,
         additional_constraint_gen_passes: Optional[Sequence[ConstraintsGenPass]] = None,
+        opt_level: Optional[int] = None,
     ) -> None:
         self.memory_config: MemoryConfig = memory_config
         self.additional_constraint_gen_passes: Optional[
             Sequence[ConstraintsGenPass]
         ] = additional_constraint_gen_passes
-        self.opt_level: int = opt_level
+        normalized_mode = normalize_compile_mode(mode, opt_level)
+        self.opt_level: int = resolve_opt_level(normalized_mode)
+        self.mode: Optional[CompileMode] = (
+            normalized_mode if isinstance(normalized_mode, CompileMode) else None
+        )
         self.alloc_graph_input: bool = alloc_graph_input
         self.alloc_graph_output: bool = alloc_graph_output
         self.memory_id_is_valid: list[bool] = [True] * self.get_num_memories()

backends/cadence/aot/pass_utils.py

@@ -7,7 +7,9 @@
 # pyre-strict
 
 import dataclasses
+import warnings
 from dataclasses import dataclass
+from enum import Enum
 from typing import Callable, List, Optional, Set, Type, Union
 
 import torch
@@ -25,6 +27,59 @@
 from torch._ops import OpOverloadPacket
 
 
+class CompileMode(Enum):
+    """Selects which pass pipeline the Cadence backend runs.
+
+    During migration, omitting ``mode`` preserves the legacy ``opt_level=1``
+    behavior, while explicitly passing ``DEFAULT`` selects the former O3
+    pipeline with all graph optimizations.
+    """
+
+    MINIMAL = "minimal"
+    DEFAULT = "default"
+    SIZE = "size"
+
+
+def opt_level_from_compile_mode(mode: CompileMode) -> int:
+    """Translate the new API to the legacy pass levels during migration."""
+    if mode is CompileMode.MINIMAL:
+        return 0
+    if mode is CompileMode.SIZE:
+        return 4
+    return 3
+
+
+def resolve_opt_level(
+    mode: CompileMode | int | None = None,
+    opt_level: Optional[int] = None,
+) -> int:
+    if opt_level is not None:
+        if mode is not None:
+            warnings.warn(
+                "Both mode and opt_level were provided; deprecated opt_level takes "
+                "precedence.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        return opt_level
+    if mode is None:
+        return 1
+    if isinstance(mode, int):
+        return mode
+    return opt_level_from_compile_mode(mode)
+
+
+def normalize_compile_mode(
+    mode: CompileMode | int | None = None,
+    opt_level: Optional[int] = None,
+) -> CompileMode | int:
+    """Resolve compatibility arguments while preserving an explicit enum mode."""
+    resolved_opt_level = resolve_opt_level(mode, opt_level)
+    if opt_level is None and isinstance(mode, CompileMode):
+        return mode
+    return resolved_opt_level
+
+
 # Is an overlap in tensor lifetime and storage allowed at the current opt level?
 # We allow overlap at opt level >= 2.
 def allow_lifetime_and_storage_overlap(opt_level: int) -> bool: