Add memory planning pass for fully delegated path (#4176)

Summary: Pull Request resolved: #4176

Reviewed By: tarun292

Differential Revision: D59507586

Pulled By: cccclai

fbshipit-source-id: 8b4c196cfbb7eacaef37bc382dabdd290f0fa89a

Author: haowhsu-quic

exir/lowered_backend_module.py

@@ -139,14 +139,15 @@ def buffer(
         segment_alignment: int = 4096,
         constant_tensor_alignment: Optional[int] = None,
         delegate_alignment: Optional[int] = None,
+        memory_planning: MemoryPlanningPass = None,
     ) -> bytes:
         """
         Returns a buffer containing the serialized ExecuTorch binary.
         """
         # TODO(T181463742): avoid calling bytes(..) which incurs large copies.
         out = bytes(
             _serialize_pte_binary(
-                program=self.program(),
+                program=self.program(memory_planning=memory_planning),
                 extract_delegate_segments=extract_delegate_segments,
                 segment_alignment=segment_alignment,
                 constant_tensor_alignment=constant_tensor_alignment,
@@ -157,7 +158,11 @@ def buffer(
 
     # TODO(chenlai): re-consider recapture instead of manually constructing the program because
     # the meta data construction is done manually.
-    def program(self, emit_stacktrace: bool = False) -> Program:
+    def program(
+        self,
+        emit_stacktrace: bool = False,
+        memory_planning: MemoryPlanningPass = None,
+    ) -> Program:
         # Fix autodpes introuces cyclic dependencies:
         # program -> verifier -> lowered_backend_module -> program
         # @manual
@@ -319,9 +324,9 @@ def program(self, emit_stacktrace: bool = False) -> Program:
             example_inputs=None,
             verifier=lowered_exported_program.verifier,
         )
-        exported_program = _transform(
-            exported_program, SpecPropPass(), MemoryPlanningPass("greedy")
-        )
+        if memory_planning is None:
+            memory_planning = MemoryPlanningPass("greedy")
+        exported_program = _transform(exported_program, SpecPropPass(), memory_planning)
         emitted_program = emit_program(
             exported_program, emit_stacktrace=emit_stacktrace
         ).program