Disable im2row transform by default, add a flag to enable it (pytorch#18693)

Differential Revision: D97356070

Pull Request resolved: pytorch#18693

Author: mcremon-meta

backends/cadence/aot/compiler.py

@@ -43,6 +43,7 @@
 from torch.export.exported_program import ExportedProgram
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e
 
+from .pass_utils import EdgePassesConfig
 from .passes import apply_exir_ops_passes, apply_torch_ops_passes
 from .utils import print_ops_info
 
@@ -355,12 +356,15 @@ def _lower_ep_to_cadence(
     program: ExportedProgram,
     dump_graphs: bool = False,
     opt_level: int = 1,
+    edge_passes_config: Optional[EdgePassesConfig] = 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)
+    cadence_prog_manager = apply_exir_ops_passes(
+        opt_level, edge_prog_manager, edge_passes_config
+    )
     return cadence_prog_manager
 
 
@@ -369,9 +373,12 @@ def export_to_cadence(
     inputs: tuple[object, ...],
     dump_graphs: bool = False,
     opt_level: int = 1,
+    edge_passes_config: Optional[EdgePassesConfig] = 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)
+    cadence_prog_manager = apply_exir_ops_passes(
+        opt_level, edge_prog_manager, edge_passes_config
+    )
     return cadence_prog_manager
 
 
@@ -380,6 +387,7 @@ def quantize_and_export_to_cadence(
     inputs: tuple[object, ...],
     dump_graphs: bool = False,
     opt_level: int = 1,
+    edge_passes_config: Optional[EdgePassesConfig] = None,
 ) -> EdgeProgramManager:
     """
     Trace, quantize, lower a model/inputs pair to edge IR and apply frontend
@@ -391,6 +399,7 @@ def quantize_and_export_to_cadence(
         quantized_model,
         opt_level=opt_level,
         dump_graphs=dump_graphs,
+        edge_passes_config=edge_passes_config,
     )
 
 

backends/cadence/aot/pass_utils.py

@@ -6,6 +6,7 @@
 
 # pyre-strict
 
+import dataclasses
 from abc import abstractmethod
 from dataclasses import dataclass
 from typing import Callable, List, Optional, override, Set, Type, TypeVar, Union
@@ -28,11 +29,20 @@ def allow_lifetime_and_storage_overlap(opt_level: int) -> bool:
     return opt_level >= 2
 
 
+# A dataclass that bundles feature flags for edge passes.
+# When adding a new flag, add a matching bool field to both this class and
+# CadencePassAttribute; the pass filter will pick it up automatically.
+@dataclass(frozen=True)
+class EdgePassesConfig:
+    use_im2row_transform: bool = False
+
+
 # A dataclass that stores the attributes of an ExportPass.
 @dataclass(frozen=True)
 class CadencePassAttribute:
     opt_level: Optional[int] = None
     debug_pass: bool = False
+    use_im2row_transform: bool = False
 
 
 # A dictionary that maps an ExportPass to its attributes.
@@ -58,17 +68,38 @@ def get_all_available_cadence_passes() -> Set[Type[PassBase]]:
     return set(ALL_CADENCE_PASSES.keys())
 
 
+def _check_feature_flags(
+    pass_attribute: CadencePassAttribute,
+    config: EdgePassesConfig,
+) -> bool:
+    """Check all feature flags: a pass is included only if every feature it
+    requires is enabled in the config. Iterates over EdgePassesConfig fields
+    so new flags are handled automatically."""
+    for field in dataclasses.fields(EdgePassesConfig):
+        if getattr(pass_attribute, field.name, False) and not getattr(
+            config, field.name
+        ):
+            return False
+    return True
+
+
 # Create a new filter to filter out relevant passes from all passes.
 def create_cadence_pass_filter(
-    opt_level: int, debug: bool = False
+    opt_level: int,
+    debug: bool = False,
+    edge_passes_config: Optional[EdgePassesConfig] = None,
 ) -> Callable[[Type[PassBase]], bool]:
+    if edge_passes_config is None:
+        edge_passes_config = EdgePassesConfig()
+
     def _filter(p: Type[PassBase]) -> bool:
         pass_attribute = get_cadence_pass_attribute(p)
         return (
             pass_attribute is not None
             and pass_attribute.opt_level is not None
             and pass_attribute.opt_level <= opt_level
             and (not pass_attribute.debug_pass or debug)
+            and _check_feature_flags(pass_attribute, edge_passes_config)
         )
 
     return _filter

backends/cadence/aot/passes.py

@@ -19,6 +19,7 @@
 from executorch.backends.cadence.aot.pass_utils import (
     CadencePassAttribute,
     create_cadence_pass_filter,
+    EdgePassesConfig,
     register_cadence_pass,
 )
 
@@ -101,9 +102,12 @@ def get_passes_in_default_order() -> list[Type[ExportPass]]:
 def apply_exir_ops_passes(
     opt_level: int,
     edge_prog_manager: EdgeProgramManager,
+    edge_passes_config: Optional[EdgePassesConfig] = None,
 ) -> EdgeProgramManager:
     passes = get_passes_in_default_order()
-    pass_filter = create_cadence_pass_filter(opt_level)
+    pass_filter = create_cadence_pass_filter(
+        opt_level, edge_passes_config=edge_passes_config
+    )
     cadence_passes = [
         (
             lambda graph_module, filtered_pass=filtered_pass: filtered_pass()(

backends/cadence/aot/replace_ops.py

@@ -1327,7 +1327,7 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         return True
 
 
-@register_cadence_pass(CadencePassAttribute(opt_level=2))
+@register_cadence_pass(CadencePassAttribute(opt_level=2, use_im2row_transform=True))
 class ReplaceConvWithIm2RowAndLinear(RemoveOrReplacePassInterface):
     """
     Replace convolution where groups=1 with im2row followed by a linear op.