Class MemoryPlanningAlgorithmSuite

exir/memory_planning.py:1031–1108

class MemoryPlanningAlgorithmSuite:
    def __init__(
        self,
        algo_list: Optional[List[Callable[..., MemoryAlgoResult]]] = None,
    ) -> None:
        if algo_list is None:
            algo_list = [greedy]
        self.algo_list: List[Callable[..., MemoryAlgoResult]] = algo_list

    def __call__(
        self,
        alignment: int,
        specs: Set[TensorSpec],
        graph_module: torch.fx.GraphModule,
        graph_signature: ExportGraphSignature,
        extra_padding: int,
    ) -> List[int]:
        r"""
        Memory planning algorithm suite that runs a list of memory planning algorithms
        and returns the result of the algorithm that minimizes the total memory usage.

        Args:
            graph_module: The graph module to allocate memory for
            alignment: Memory alignment requirement
            graph_signature: Optional graph signature
            alloc_graph_input: Whether to allocate memory for graph input
            alloc_graph_output: Whether to allocate memory for graph output
            allow_overlapping_allocations: Whether to allow overlapping allocations
            algo_list: List of memory planning algorithms to run
            specs: Optional set of TensorSpec objects with updated lifetimes. If None, they will be
                calculated from the graph_module.

        Returns:
            List of buffer sizes for each memory hierarchy
        """

        mem_algo_results = {}
        for algo in self.algo_list:
            if isinstance(algo, functools.partial):
                name = algo.func.__name__
            else:
                name = getattr(algo, "__name__", None)

            mem_algo_results[name] = algo(
                alignment,
                specs,
                graph_module,
                graph_signature,
                extra_padding,
            )

        # All the algorithms should have the same number of buffers allocated.
        assert (
            len(
                {
                    len(mem_algo_result.bufsizes)
                    for mem_algo_result in mem_algo_results.values()
                }