memory_planning algos take the specs as inputs instead of calculating them themselves (#9952)

Summary:

Refactor the algos to not calculate the active set of nodes themselves. Future work should refactor apply_algo more to not be recursive and actually smartly handle lifespans on control flow.

Made the memory planning suite a class so that the algo list is easily configureable without having to make a wrapper function.

Reviewed By: skrtskrtfb

Differential Revision: D72600295

Author: JacobSzwejbka

backends/cadence/aot/memory_planning.py

@@ -12,7 +12,7 @@
 import math
 import typing
 from functools import partial
-from typing import Iterable, List, Optional, Tuple
+from typing import Iterable, List, Optional, Set, Tuple
 
 import torch
 from executorch.backends.cadence.aot.memory_constraints import (
@@ -73,11 +73,11 @@ def collect_specs_from_graph_module(
 # the fastest memory available
 # flake8: noqa 'position_based_greedy_with_hierarchy' is too complex (13)
 def position_based_greedy_with_hierarchy(
-    graph_module: torch.fx.GraphModule,
     alignment: int,
+    specs: Set[TensorSpec],
+    graph_module: torch.fx.GraphModule,
     graph_signature: ExportGraphSignature,
-    alloc_graph_input: bool,
-    alloc_graph_output: bool,
+    extra_padding: int = 0,
     *,
     memory_config: MemoryConfig,
     mem_constraints: MemConstraints,
@@ -119,9 +119,7 @@ def memory_available(spec: TensorSpec) -> bool:
 
     # Iterate over all the specs in sorted order
     for spec in sorted(
-        collect_specs_from_graph_module(
-            graph_module, graph_signature, alloc_graph_input, alloc_graph_output
-        ),
+        specs,
         key=lambda spec: spec.allocated_memory,
         reverse=True,
     ):
@@ -167,11 +165,11 @@ def memory_available(spec: TensorSpec) -> bool:
 
 # Greedy tensor placement with the heuristics from arxiv.org/pdf/2001.03288.pdf
 def greedy_by_size_for_offset_calculation_with_hierarchy(
-    graph_module: torch.fx.GraphModule,
     alignment: int,
+    specs: Set[TensorSpec],
+    graph_module: torch.fx.GraphModule,
     graph_signature: ExportGraphSignature,
-    alloc_graph_input: bool,
-    alloc_graph_output: bool,
+    extra_padding: int = 0,
     *,
     memory_config: MemoryConfig,
     mem_constraints: MemConstraints,
@@ -199,9 +197,7 @@ def greedy_by_size_for_offset_calculation_with_hierarchy(
 
     # Iterate over all the specs in sorted order
     for spec in sorted(
-        collect_specs_from_graph_module(
-            graph_module, graph_signature, alloc_graph_input, alloc_graph_output
-        ),
+        specs,
         key=lambda spec: spec.allocated_memory,
         reverse=True,
     ):

backends/vulkan/vulkan_preprocess.py

@@ -47,7 +47,7 @@
 )
 from executorch.exir.backend.utils import DelegateMappingBuilder
 
-from executorch.exir.memory_planning import greedy, memory_planning_algorithm_suite
+from executorch.exir.memory_planning import greedy, MemoryPlanningAlgorithmSuite
 from executorch.exir.pass_base import ExportPass, PassBase
 
 from executorch.exir.passes import MemoryPlanningPass, SpecPropPass
@@ -199,9 +199,7 @@ def preprocess(  # noqa: C901
         # Finally, apply dynamic shape passes and memory planning pass. These passes
         # must be applied only when the graph structure is finalized.
         greedy_memory_planning = partial(greedy, allow_overlapping_allocations=False)
-        mem_planning_suite = partial(
-            memory_planning_algorithm_suite, algo_list=[greedy_memory_planning]
-        )
+        mem_planning_suite = MemoryPlanningAlgorithmSuite(algo_list=[greedy_memory_planning])
         program = apply_passes(
             program,
             [