use partitioner instance directly in to_backend (#2513)

Summary:
Pull Request resolved: #2513

to_backend either takes partitioner or a dict of partitioner `key: method_name, value: partitioner`. We shouldn't do key as the backend name and value as the partitioner.

Reviewed By: mergennachin

Differential Revision: D55078939

fbshipit-source-id: 7c2b8587095d2227d17a198c3b4e975910f4df00

Author: cccclai

examples/models/llama2/builder.py

@@ -12,7 +12,7 @@
 import logging
 from enum import Enum
 from json import JSONDecodeError
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional
 
 import torch
 from executorch.backends.transforms.duplicate_dynamic_quant_chain import (
@@ -288,30 +288,18 @@ def export_to_edge(
             )
         return self
 
-    def to_backend(
-        self, partitioner: Union[Partitioner, Dict[str, Partitioner]]
-    ) -> "LlamaEdgeManager":
+    def to_backend(self, partitioner: Optional[Partitioner]) -> "LlamaEdgeManager":
         """
         Partition the model and lower to different backends. The signature is
         aligned with the signature of `to_backend` method of EdgeManager.
         Args:
-            partitioner (Union[Partitioner, Dict[str, Partitioner]]): One or more
+            partitioner (Optional[Partitioner]): One or more
                 partitioner to be sent to EdgeManager.to_backend().
         """
         assert self.edge_manager is not None, "Need to run export_to_edge() first"
-        if isinstance(partitioner, dict):
-            for key, p in partitioner.items():
-                assert self.edge_manager is not None
-                self.edge_manager = self.edge_manager.to_backend(p)
-                if self.verbose:
-                    logging.info(
-                        print_delegated_graph(
-                            self.edge_manager.exported_program().graph_module
-                        )
-                    )
-                    logging.info(f"Applied partitioners: {key}")
-        elif isinstance(partitioner, Partitioner):
-            assert self.edge_manager is not None
+        if partitioner is None:
+            logging.info("No partitioner provided, passing...")
+        else:
             self.edge_manager = self.edge_manager.to_backend(partitioner)
             if self.verbose:
                 logging.info(
@@ -320,8 +308,6 @@ def to_backend(
                     )
                 )
                 logging.info(f"Applied partitioners: {partitioner}")
-        else:
-            logging.warning("Invalid partitioner, skipping...")
         return self
 
     def to_executorch(self) -> "LlamaEdgeManager":

examples/models/llama2/export_llama_lib.py

@@ -539,22 +539,18 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
     ).export_to_edge(quantizers)
 
     # to_backend
-    partitioners = {}
+    partitioner = None
     if pt2e_quant_params is not None and pt2e_quant_params.quantize_linear is not None:
-        partitioners[XnnpackDynamicallyQuantizedPartitioner.__name__] = (
-            XnnpackDynamicallyQuantizedPartitioner()
-        )
+        partitioner = XnnpackDynamicallyQuantizedPartitioner()
         modelname = f"xnnpack_dq_{modelname}"
 
     if args.xnnpack:
         # Following changes due to.
         # 1. We need dynamically quantized partitioner for both pt2e_quantize options
         #    as well as "qmode 8da4w" which is also dynamic quantizes linear layers.
         # 2. XNNPACK partitioner seems to result in seg fault for non dqlinear ops.
-        partitioners[XnnpackDynamicallyQuantizedPartitioner.__name__] = (
-            XnnpackDynamicallyQuantizedPartitioner()
-        )
-        # partitioners[XnnpackPartitioner.__name__] = XnnpackPartitioner()
+        partitioner = XnnpackDynamicallyQuantizedPartitioner()
+        # partitioner = XnnpackPartitioner()
         modelname = f"xnnpack_{modelname}"
 
     if args.vulkan:
@@ -565,7 +561,7 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
             args.quantization_mode is None
         ), "Vulkan backend does not support quantization at the moment"
 
-        partitioners[VulkanPartitioner.__name__] = VulkanPartitioner()
+        partitioner = VulkanPartitioner()
         modelname = f"vulkan_{modelname}"
 
     if args.mps:
@@ -584,7 +580,7 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
 
         compile_specs = [CompileSpec("use_fp16", bytes([True]))]
         # pyre-ignore: Undefined attribute [16]: Module `executorch.backends` has no attribute `apple`.
-        partitioners[MPSPartitioner.__name__] = MPSPartitioner(compile_specs)
+        partitioner = MPSPartitioner(compile_specs)
         modelname = f"mps_{modelname}"
 
     if args.coreml:
@@ -615,7 +611,7 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
             model_type=CoreMLBackend.MODEL_TYPE.MODEL,
         )
         # pyre-ignore: Undefined attribute [16]: Module `executorch.backends` has no attribute `apple`
-        partitioners[CoreMLPartitioner.__name__] = CoreMLPartitioner(
+        partitioner = CoreMLPartitioner(
             skip_ops_for_coreml_delegation=None, compile_specs=compile_specs
         )
         modelname = f"coreml_{modelname}"
@@ -627,7 +623,7 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
         logging.info("Generating etrecord")
         # Copy the edge manager which will be serialized into etrecord. This is memory-wise expensive.
         edge_manager_copy = copy.deepcopy(builder_exported_to_edge.edge_manager)
-        builder = builder_exported_to_edge.to_backend(partitioners).to_executorch()
+        builder = builder_exported_to_edge.to_backend(partitioner).to_executorch()
 
         # Generate ETRecord
         if edge_manager_copy:
@@ -638,7 +634,7 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
             )
             logging.info("Generated etrecord.bin")
     else:
-        builder = builder_exported_to_edge.to_backend(partitioners).to_executorch()
+        builder = builder_exported_to_edge.to_backend(partitioner).to_executorch()
 
     if args.profile_memory:
         generate_memory_trace(builder.export_program, "memory_profile.json")