Update

[ghstack-poisoned]

Author: mcr229

exir/backend/canonical_partitioners/all_node_partitioner.py

@@ -33,8 +33,8 @@ def __init__(
         compile_specs: List[CompileSpec],
     ):
         """
-        Partitioner that lowers every single node in the graph module to the
-        specified backend_id
+        Partitioner that lowers every single node in the graph module unconditionally
+        to the specified backend_id
         """
         super().__init__()
         self.delegation_spec = DelegationSpec(backend_id, compile_specs)

exir/program/_program.py

@@ -24,10 +24,7 @@
 from executorch.exir._serialize.data_serializer import DataSerializer
 from executorch.exir._warnings import experimental
 from executorch.exir.backend.backend_api import to_backend
-from executorch.exir.backend.canonical_partitioners.all_node_partitioner import (
-    AllNodePartitioner,
-)
-from executorch.exir.backend.partitioner import DelegationSpec, Partitioner
+from executorch.exir.backend.partitioner import Partitioner
 from executorch.exir.capture._config import EdgeCompileConfig, ExecutorchBackendConfig
 from executorch.exir.delegate import executorch_call_delegate, is_lowered_module
 from executorch.exir.emit import emit_program, EmitterOutput
@@ -1442,29 +1439,20 @@ def transform(
 
     @et_logger("to_backend")
     def to_backend(
-        self,
-        partitioner: Union[
-            DelegationSpec,
-            Dict[str, DelegationSpec],
-            Partitioner,
-            Dict[str, Partitioner],
-        ],
+        self, partitioner: Union[Partitioner, Dict[str, Partitioner]]
     ) -> "EdgeProgramManager":
         """
         Returns a semantically-equivalent program to the one given as input,
         but with portions of each program in the EdgeProgramManager targeted
         for delegation as determined by the partitioner.
 
         Args:
-            partitioner: The partitioner can be:
-                - Partitioner Subclass Instance; all programs in the EdgeProgramManager are lowered with
-                  this partitioner
-                - Dictionary mapping method name to partitioner subclass instance; Only method names specified
-                  in the dictionary will be lowered by the given partitioner.
-                - DelegationSpec; All programs are completely lowered to the backend_id specified in the
-                  DelegationSpec
-                - Dictionary mapping method name to DelegationSpec; Only method names specified in the dictionary
-                  will be lowered to the backend_id specified in the DelegationSpec
+            partitioner: The partitioner can either be a Partitioner subclass instance, or a
+                dictionary mapping method names to Partitioner subclass instance. If it is a
+                Partitioner subclass, all programs in the given EdgeProgramManager
+                will be lowered using the given partitioner. If it is a
+                dictionary, only method names specified in the dictionary will be
+                lowered with the given partitioner.
 
                 The Partitioner subclass instance is in charge with tagging portions of the
                 input program for delegation. A valid partitioner must return PartitionerResult including valid
@@ -1480,19 +1468,13 @@ def to_backend(
         if isinstance(partitioner, dict):
             for name, program in self._edge_programs.items():
                 if name in partitioner.keys():
-                    partitioner_to_use = partitioner[name]
-                    if isinstance(partitioner_to_use, DelegationSpec):
-                        partitioner_to_use = AllNodePartitioner(partitioner_to_use)
-                    new_edge_programs[name] = to_backend(program, partitioner_to_use)
+                    new_edge_programs[name] = to_backend(program, partitioner[name])
                 else:
                     new_edge_programs[name] = program
 
         else:  # apply partitioner to every method
             for name, program in self._edge_programs.items():
-                partitioner_to_use = partitioner
-                if isinstance(partitioner, DelegationSpec):
-                    partitioner_to_use = AllNodePartitioner(partitioner)
-                new_edge_programs[name] = to_backend(program, partitioner_to_use)
+                new_edge_programs[name] = to_backend(program, partitioner)
 
         config = EdgeCompileConfig(_check_ir_validity=False)
         return EdgeProgramManager(