memory_planning.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

import functools
import itertools
import logging
import operator
from collections import defaultdict
from dataclasses import dataclass, field
from typing import (
    Any,
    Callable,
    cast,
    Dict,
    Iterable,
    List,
    Optional,
    Set,
    Tuple,
    Union,
)

import torch
from executorch.exir import memory
from executorch.exir.control_flow import while_loop as exir_while
from executorch.exir.delegate import executorch_call_delegate
from executorch.exir.error import internal_assert, InternalError
from executorch.exir.operator.convert import is_inplace_variant, is_out_variant
from executorch.exir.schema import TensorShapeDynamism
from executorch.exir.tensor import TensorSpec
from torch import fx
from torch.export.exported_program import (
    ConstantArgument,
    ExportGraphSignature,
    InputKind,
)
from torch.fx import Node
from torch.utils._pytree import tree_flatten

REGISTERED_ALGOS: Dict[str, Callable[..., List[int]]] = {}


class Verifier:
    """
    Verify if the outcome of a memory planning algorithm makes sense.
    E.g., make sure tensors having overlapping lifetime does not have overlapping
    storage/buffer.
    """

    def __init__(
        self,
        graph_module: torch.fx.GraphModule,
        alloc_graph_input: bool,
        alloc_graph_output: bool,
        alloc_mutable_buffers: bool,
        graph_signature: Optional[ExportGraphSignature] = None,
    ) -> None:
        self.graph_module = graph_module
        self.graph_signature = graph_signature
        self.alloc_graph_input = alloc_graph_input
        self.alloc_graph_output = alloc_graph_output
        self.alloc_mutable_buffers = alloc_mutable_buffers

    @classmethod
    def mem_obj_id_match(
        cls, lhs_spec: TensorSpec, rhs_spec: TensorSpec, accept_both_none: bool = True
    ) -> bool:
        """
        Given two `TensorSpec`, return if their `mem_obj_id` are the same. Note that if
        both are None, this function will return True if `accept_both_none` is True and
        False otherwise.
        """
        if lhs_spec.mem_id != rhs_spec.mem_id:
            return False

        # both are None
        if lhs_spec.mem_obj_id is None and rhs_spec.mem_obj_id is None:
            return accept_both_none

        return lhs_spec.mem_obj_id == rhs_spec.mem_obj_id

    @classmethod
    def has_overlap(cls, lhs_ivl: List[int], rhs_ivl: List[int]) -> bool:
        r"""
        The passed in intervals are inclusive in both sides. Return if they have
        overlapping.
        """
        # empty interval
        if lhs_ivl[0] > lhs_ivl[1] or rhs_ivl[0] > rhs_ivl[1]:
            return False

        return (lhs_ivl[0] >= rhs_ivl[0] and lhs_ivl[0] <= rhs_ivl[1]) or (
            rhs_ivl[0] >= lhs_ivl[0] and rhs_ivl[0] <= lhs_ivl[1]
        )

    @classmethod
    def lifetime_overlap(cls, lhs_spec: TensorSpec, rhs_spec: TensorSpec) -> bool:
        lhs_lifetime = lhs_spec.lifetime
        rhs_lifetime = rhs_spec.lifetime
        internal_assert(
            lhs_lifetime[0] is not None and lhs_lifetime[1] is not None,
            f"{lhs_spec} should have valid start and end",
        )
        internal_assert(
            rhs_lifetime[0] is not None and rhs_lifetime[1] is not None,
            f"{rhs_spec} should have valid start and end",
        )
        return cls.has_overlap(lhs_lifetime, rhs_lifetime)

    @classmethod
    def storage_overlap(cls, lhs_spec: TensorSpec, rhs_spec: TensorSpec) -> bool:
        intervals = []
        if lhs_spec.mem_id != rhs_spec.mem_id:
            return False
        for spec in [lhs_spec, rhs_spec]:
            internal_assert(
                spec.allocated_memory >= 0,
                f"{spec} should have non-zero allocated memory",
            )
            internal_assert(
                isinstance(spec.mem_offset, int) and spec.mem_offset >= 0,
                f"{spec} should have specified memory offset",
            )
            intervals.append(
                [spec.mem_offset, spec.mem_offset + spec.allocated_memory - 1]
            )
        has_overlap = cls.has_overlap(*intervals)

        return has_overlap

    @classmethod
    def _debug_message_from_specs(
        cls, lhs_spec: TensorSpec, rhs_spec: TensorSpec
    ) -> str:
        message = (
            f"lhs life time: {lhs_spec.lifetime}, rhs lifetime: {rhs_spec.lifetime} "
        )
        message += f"lhs: mem_id {lhs_spec.mem_id} storage: {lhs_spec.mem_offset}, {lhs_spec.allocated_memory} "
        message += f"rhs: mem_id {rhs_spec.mem_id} storage: {rhs_spec.mem_offset}, {rhs_spec.allocated_memory}"
        return message

    def verify_storage_reuse(
        self, allow_lifetime_and_storage_overlap: bool = False
    ) -> int:
        """
        'allow_lifetime_and_storage_overlap' allows tensors to overlap in both
        lifetime and storage. If is it False, and two tensors have both overlapping
        lifetime and storage, throw an exception.
        Returns:
            Number of pairs of tenors that have overlapping storage.
        """
        num_reuse_pairs = 0

        # unique tensors specs
        all_specs = list(
            collect_specs_from_nodes(
                self.graph_module.graph.nodes,
                self.graph_signature,
                ignore_const=True,
                ignore_graph_input=not self.alloc_graph_input,
                ignore_graph_output=not self.alloc_graph_output,
                ignore_mutable_buffers=not self.alloc_mutable_buffers,
                do_assertion=False,
                ignore_out_var_node=False,
                dedup=True,
            )
        )

        for lhs_spec_idx, lhs_spec in enumerate(all_specs):
            for rhs_spec in all_specs[lhs_spec_idx + 1 :]:
                # Check that both specs are consistent about whether mem_obj_id is defined
                if (lhs_spec.mem_obj_id is None) != (rhs_spec.mem_obj_id is None):
                    raise InternalError(
                        "Specs do not agree on whether mem_obj_id is defined."
                    )

                has_storage_overlap = Verifier.storage_overlap(lhs_spec, rhs_spec)
                if not has_storage_overlap:
                    continue

                if not allow_lifetime_and_storage_overlap and self.lifetime_overlap(
                    lhs_spec, rhs_spec
                ):
                    raise InternalError(
                        f"Unexpected storage overlap: {Verifier._debug_message_from_specs(lhs_spec, rhs_spec)}"
                    )

                # Check that each mem_obj_id is consistent with whether the tensors have
                # storage overlap
                if not Verifier.mem_obj_id_match(lhs_spec, rhs_spec):
                    raise InternalError(
                        f"Unexpected mem_obj_id mismatch: lhs {lhs_spec}, rhs {rhs_spec}"
                    )

                num_reuse_pairs += 1

        return num_reuse_pairs

    def verify_graph_input_output(self) -> None:
        r"""
        alloc_graph_input / alloc_graph_output indicates if memory for graph
        input/output is allocated by the compiler. If not, the runtime will
        set them using buffers provided by users.
        """
        graph_module = self.graph_module
        # There is one tricky case here. If the graph input and graph output
        # tensors have overlap, but alloc_graph_input != alloc_graph_output,
        # then the overlapped tensor will cause assertion failure below.
        # The current behavior is if either alloc_graph_input or alloc_graph_output
        # is false, those overlapped tensor will not have memory allocated.
        #
        # Ignore the check in this case for now.
        overlap = get_graph_input_tensors(
            graph_module.graph.nodes, self.graph_signature
        ) & get_graph_output_tensors(graph_module.graph.nodes)
        if overlap and (self.alloc_graph_input != self.alloc_graph_output):
            logging.debug(
                "Having overlapping graph input/output tensors while the allocation decision for graph input/output mismatch."
            )
            return

        graph_input_allocated = None
        graph_output_allocated = None

        has_dynamic_unbound_input = False
        has_dynamic_unbound_output = False

        check_list = {"placeholder", "output"} & {
            node.op for node in graph_module.graph.nodes
        }
        assert "output" in check_list, f"graph module has no output: {graph_module}"

        # Collect mutable buffer specs so we can filter them when they appear on
        # non-placeholder nodes (e.g., aliased on the output node via SpecPropPass).
        mutable_buffer_specs = _get_mutable_buffer_specs(
            graph_module.graph.nodes, self.graph_signature
        )
        for nd in graph_module.graph.nodes:
            if nd.op in check_list:
                if not (specs := get_node_tensor_specs(nd)):
                    continue
                if _is_mutable_buffer(nd, self.graph_signature):
                    continue
                specs = list(
                    filter(
                        lambda spec: not spec.const
                        and spec not in mutable_buffer_specs,
                        specs,
                    )
                )
                if len(specs) == 0:
                    # all outputs are const so no need to allocate memory just say we succeeded
                    graph_output_allocated = self.alloc_graph_output
                    continue
                allocated = any(
                    spec is None or spec.mem_offset is not None for spec in specs
                )
                has_dynamic_unbound_tensor = any(
                    spec is None
                    or spec.shape_dynamism == TensorShapeDynamism.DYNAMIC_UNBOUND
                    for spec in specs
                )
                assert (
                    all(spec is None or spec.mem_offset is not None for spec in specs)
                    == allocated
                ), "Either all or non of the tensors should be allocated memory"
                if nd.op == "placeholder":
                    graph_input_allocated = allocated
                    has_dynamic_unbound_input |= has_dynamic_unbound_tensor
                else:
                    graph_output_allocated = allocated
                    has_dynamic_unbound_output |= has_dynamic_unbound_tensor

        # only check if inputs are allocated if there are user inputs:
        user_inputs_exist = _do_user_inputs_exist(graph_signature=self.graph_signature)

        if "placeholder" in check_list and user_inputs_exist:
            assert graph_input_allocated is not None, "graph_input_allocated not set"
            if not has_dynamic_unbound_input:
                assert (
                    graph_input_allocated == self.alloc_graph_input
                ), f"Misallocate graph input: {graph_input_allocated} v.s. {self.alloc_graph_input}"

        assert graph_output_allocated is not None, "graph_output_allocated not set"
        if not has_dynamic_unbound_output:
            assert (
                graph_output_allocated == self.alloc_graph_output
            ), f"Misallocate graph output {graph_output_allocated} v.s. {self.alloc_graph_output}"


def _is_out_var_node(node: torch.fx.Node) -> bool:
    return (
        node.op == "call_function"
        and isinstance(node.target, torch._ops.OpOverload)
        and is_out_variant(node.target._schema.name, node.target._schema.overload_name)
    )


def _is_inplace_node(node: torch.fx.Node) -> bool:
    return (
        node.op == "call_function"
        and isinstance(node.target, torch._ops.OpOverload)
        and is_inplace_variant(
            node.target._schema.name, node.target._schema.overload_name
        )
    )


def update_tensor_lifetime(
    node: torch.fx.Node,
    spec: TensorSpec,
    node_idx: int,
    max_node_idx: int,
    gs: Optional[ExportGraphSignature] = None,
) -> None:
    r"""
    Update the lifetime of the tensor to cover node_idx. A tensor's lifetime
    are represented by the index of the first and last node referring
    that tensor in its inputs/outputs.

    Arguments:
        spec: the TensorSpec for the tensor
        node_idx: extend the tensor's lifetime to cover node_idx
    """
    start, end = spec.lifetime
    if node.op == "placeholder":
        start = 0
    else:
        start = node_idx if start is None or start > node_idx else start

    if node.op == "placeholder" and _is_mutable_buffer(node, gs):
        # mutable buffers are never freed
        end = max_node_idx
    else:
        end = node_idx if end is None or end < node_idx else end
    spec.lifetime = [start, end]


# pyre-ignore
def filter_nodes(inputs: Iterable[Any]) -> Iterable[Node]:
    """
    This method need return Node object embedded inside List/Dict as well.
    """
    return [nd for nd in tree_flatten(list(inputs))[0] if isinstance(nd, Node)]


def _is_mutable_buffer(
    node: Node, graph_signature: Optional[ExportGraphSignature] = None
) -> bool:
    """
    Check if the node is mutable buffer according to the provided graph signature.
    """
    # graph signature is None for memory planning passes not called from EdgeProgramManager, these paths are deprecated so mutable buffers are not supported on them.
    if graph_signature is None:
        return False
    if node.op == "placeholder":
        if isinstance(node.target, str):
            if node.target in graph_signature.inputs_to_buffers:
                fqn = graph_signature.inputs_to_buffers[node.target]
                # if the buffer is mutated then record that
                if fqn in graph_signature.buffers_to_mutate.values():
                    return True
    return False


def _get_mutable_buffer_specs(
    nodes: Iterable[Node], graph_signature: Optional[ExportGraphSignature]
) -> Set[TensorSpec]:
    mutable_buffer_specs: Set[TensorSpec] = set()
    for node in nodes:
        if _is_mutable_buffer(node, graph_signature):
            for spec in get_node_tensor_specs(node):
                mutable_buffer_specs.add(spec)
    return mutable_buffer_specs


def _do_user_inputs_exist(graph_signature: Optional[ExportGraphSignature]) -> bool:
    if graph_signature is None:
        return False

    user_inputs = list(
        filter(
            lambda input: input.kind == InputKind.USER_INPUT,
            graph_signature.input_specs,
        )
    )

    # Return false if:
    # - there are no inputs.
    # - if user inputs are all prims (as this currently
    #   causes the memory planning verifier to blow up).
    # Otherwise, return true.
    return any(
        not isinstance(input.arg, ConstantArgument)
        or not isinstance(input.arg.value, (int, float, bool, str))
        for input in user_inputs
    )


def get_graph_input_tensors(
    nodes: Iterable[Node], graph_signature: Optional[ExportGraphSignature] = None
) -> Set[TensorSpec]:
    graph_input_tensors = set()
    for node in nodes:
        if node.op == "placeholder" and not _is_mutable_buffer(node, graph_signature):
            for spec in get_node_tensor_specs(node):
                graph_input_tensors.add(spec)

    return graph_input_tensors


def get_graph_output_tensors(nodes: Iterable[Node]) -> Set[TensorSpec]:
    graph_output_tensors = set()
    for node in nodes:
        if node.op == "output":
            for spec in get_node_tensor_specs(node):
                graph_output_tensors.add(spec)

    return graph_output_tensors


def collect_specs_from_nodes(  # noqa: C901
    nodes: Iterable[Node],
    graph_signature: Optional[ExportGraphSignature] = None,
    ignore_graph_input: bool = False,
    ignore_graph_output: bool = False,
    ignore_mutable_buffers: bool = False,
    ignore_const: bool = True,
    ignore_out_var_node: bool = True,
    dedup: bool = True,
    do_assertion: bool = True,
    ignore_dynamic_unbound_tensor: bool = True,
) -> Iterable[TensorSpec]:
    r"""
    Collect specs from the passed in nodes. Do filtering as controlled by
    arguments.
    Arguments:
        ignore_graph_input: ignore graph input tensors from placeholder nodes
        ignore_const: whether to ignore the const
        ignore_out_var_node: whether to ignore out variant node
        dedup: whether do dedup
        do_assertion: whether to assert the filtered nodes belong to a resticted set like alloc, getitem
    """
    unique_spec = set()
    graph_input_tensors: Set[TensorSpec] = (
        get_graph_input_tensors(nodes, graph_signature) if ignore_graph_input else set()
    )
    graph_output_tensors: Set[TensorSpec] = (
        get_graph_output_tensors(nodes) if ignore_graph_output else set()
    )

    # Collect mutable buffer specs so we can filter them when they appear on
    # non-placeholder nodes (e.g., aliased on the output node via SpecPropPass).
    mutable_buffer_specs: Set[TensorSpec] = set()
    if ignore_mutable_buffers:
        mutable_buffer_specs = _get_mutable_buffer_specs(nodes, graph_signature)

    for node in nodes:
        # ignore the specs from unrelevant Fx ops for now.
        if node.op in ["get_attr"]:
            continue

        # don't reallocate memory for out-variant op's output tensors,
        # since they are just input tenors.
        if ignore_out_var_node and _is_out_var_node(node):
            continue

        if not (specs := get_node_tensor_specs(node)):
            continue

        if _is_inplace_node(node):
            continue

        if _is_mutable_buffer(node, graph_signature) and ignore_mutable_buffers:
            continue

        if do_assertion:
            internal_assert(
                node.op in ("placeholder", "output")
                or node.target
                in [
                    memory.alloc,
                    memory.view,
                    operator.getitem,
                    torch.ops.higher_order.cond,
                    exir_while,
                    torch.ops.higher_order.map_impl,
                    executorch_call_delegate,
                ],
                f"Unexpected op {node.op}, target {node.target}",
            )
        for spec in specs:
            if spec is None:
                continue
            # Dynamic unbound tensors' memory will be allocated by the runtime.
            # Memory planning should ignore them.
            if (
                ignore_dynamic_unbound_tensor
                and spec.shape_dynamism == TensorShapeDynamism.DYNAMIC_UNBOUND
            ):
                continue

            # Note: graph input may be the output of other ops (e.g. the return op)
            # If ignore_graph_input is true, we should ignore those Tensor so
            # we skip planning memory for graph input.
            if ignore_graph_input and spec in graph_input_tensors:
                continue
            if ignore_graph_output and spec in graph_output_tensors:
                continue
            if ignore_mutable_buffers and spec in mutable_buffer_specs:
                continue
            if (
                ignore_const
                and spec.const
                and not node.meta.get("weight_has_gradient", False)
            ):
                continue
            if dedup:
                if spec in unique_spec:
                    continue
                else:
                    unique_spec.add(spec)
            yield spec


def update_all_tensors_lifetime(
    graph_module: torch.fx.GraphModule,
    graph_signature: Optional[ExportGraphSignature] = None,
) -> Set[TensorSpec]:
    r"""
    Set the lifetime for all the tensors encountered in the Fx graph.
    """
    specs = set()
    max_node_idx = len(graph_module.graph.nodes) - 1
    for node_idx, node in enumerate(graph_module.graph.nodes):
        for spec in collect_specs_from_nodes(
            filter_nodes(itertools.chain([node], node.args, node.kwargs.values())),
            graph_signature,
            ignore_graph_input=False,
            ignore_const=False,
            ignore_out_var_node=False,
            dedup=False,
            do_assertion=False,
            ignore_dynamic_unbound_tensor=False,
        ):
            update_tensor_lifetime(node, spec, node_idx, max_node_idx, graph_signature)
            specs.add(spec)
    return specs


@dataclass
class AllocationSpec:
    """
    AllocationSpec is used to represent the allocation of a tensor.
    """

    # The offset of the tensor in the shared object/pool.
    offset: int
    # TensorSpec
    spec: TensorSpec


@dataclass
class SharedObject:
    r"""
    We define the concept of shared object, which represents a segment
    in the memory buffer that can be shared by multiple tensors. In order to
    check if a shared object is available for a tensor, we maintain the
    last_used_index attribute. The shared object will be available for nodes
    with index greater than last_used_index.
    """

    # index of the shared object in the list of shared objects, used as a unique id
    idx: int
    # offset in the memory buffer
    offset: int
    # size of this shared object in bytes
    size: int
    # When the object is first created
    first_used_index: int
    # the object will be available for index (last_used_index + 1)
    last_used_index: int
    # list of allocations belong to this shared object
    allocations: List[AllocationSpec] = field(default_factory=list)

    def __repr__(self) -> str:
        return f"SharedObject(idx={self.idx}, offset={self.offset}, size={self.size}, lifetime=[{self.first_used_index, self.last_used_index}])"


@dataclass
class SpecAllocResult:
    """These are the values that a memory plannig algorithm assigns to a spec.
    These are not directly written back into the spec object, but are used to
    track the allocation decisions and assigned back to the spec object in the
    end, based on which algorithm is picked as the best performing one.
    """

    mem_id: int
    mem_obj_id: int
    mem_offset: int


@dataclass
class MemoryAlgoResult:
    """This is the result returned by a memory planning algorithm that is
    invoked by memory_planning_algorithm_suite. It contains the allocation
    decisions of that algorithm for all the specs, and the size of the buffer
    that was used for different memory hierarchies.
    """

    spec_dict: Dict[TensorSpec, SpecAllocResult]
    bufsizes: List[int]


def materialize_buffer(
    shared_objects: List[SharedObject], input_total_size: int = 0
) -> int:
    r"""
    Assign concrete location in the buffer for each SharedObject.offset.

    Assuming all the passed in shared objects belong to the same memory buffer.
    """
    total_size = input_total_size
    for sobj in shared_objects:
        sobj.offset = total_size
        total_size += sobj.size
    return total_size


def _does_not_overlap(sobj: SharedObject, spec: TensorSpec) -> bool:
    r"""
    Check if a shared object and a tensor do not overlap.
    """
    for alloc in sobj.allocations:
        if not (
            spec.lifetime[1] < alloc.spec.lifetime[0]
            or spec.lifetime[0] > alloc.spec.lifetime[1]
        ):
            return False
    return True


def _find_max_overlapping_allocations_offset(
    sobj: SharedObject, spec: TensorSpec
) -> int:
    max_offset = 0
    for alloc in sobj.allocations:
        if (
            spec.lifetime[1] < alloc.spec.lifetime[0]
            or spec.lifetime[0] > alloc.spec.lifetime[1]
        ):
            continue
        max_offset = max(alloc.offset + alloc.spec.allocated_memory, max_offset)
    return max_offset


def pick_shared_obj(
    shared_objects: List[SharedObject],
    spec: TensorSpec,
    allow_overlapping_allocations: bool = True,
) -> SharedObject:
    r"""
    Pick the available shared object to which to assign this spec,
    or create a new one
    Algorithm details
    Previous: Look at every spec in chronological order. Find if previously allocated object
    allows it to fit in. If not, allocate a new object.
    New:
    - Sort all the specs by allocation size
    - Process the specs in order
    - If the spec's size in smaller than previously allocated buckets:
        - Conditions under which previously allocated bucket can be used:
          - Lifetime of the spec does not overlap with lifetime of the bucket.
              - In this case allocate spec to that bucket and expand its lifetime.
              - Spec is allocated at offset = 0 in this bucket.
              - Add this spec to allocated object's list of specs.
          - Lifetime of the spec overlaps with lifetime of the bucket,
            partially or fully (e.g. spec's lifetime subset of bucket's lifetime)
              - If none of the specs in the bucket overlaps with spec's lifetime.
                - Allocate spec to the bucket at offset = 0.
                - Add this spec to the bucket's list of specs.
                - Expand bucket's lifetime accounting for added spec's lifetime.
              - If one or more specs in the bucket overlaps with spec's lifetime.
                - Collect offsets (at which the given overlapping spec is allocated in the bucket).
                  of all the overlapping specs, and find the max offset.
                - Allocate spec to the bucket at offset = max_offset + max_offset_spec_size.
                - Add this spec to the bucket's list of specs.
                - Expand bucket's lifetime accounting for added spec's lifetime.
        - If none of these conditions are met, allocate a new bucket.
            - Add spec to this bucket.
            - Update bucket's lifetime to that of the spec.
    - If the spec's size is larger than previously allocated buckets, allocate a new bucket.
        - Size and lifetime of this bucket is that of the spec

    Proof of correctness:
    - If allocating a new bucket, it is correct.
    - If allocating spec to an existing bucket, whose lifetime does not overlap with any
      of the previously allocated specs' lifetime, then the allocation is correct.
    Proof of correctness by induction when adding spec to an existing bucket:
    - If all previous allocations in the given bucket are correct:
        - Then the new one being added must be correct because when the requested allocation
          overlaps with one or more previous allocations, we find the largest offset among
          all the overlapping allocations, and allocate the new spec at that offset. Hence,
          the allocation at such an offset, will not overlap with any previous allocations.
    Base case: A newly added allocation within a bucket with single allocation is correct:
    because a) it must fit and b) its lifetime must not overlap with object's lifetime.
    This holds true because of the following invariants:
    - Once a bucket is created, it is never resized.
    - All the allocations within a bucket follow this:
      - Span, defined by allocation's offset + size, of two allocations can only overlap,
        if their timelines do not overlap.
    """
    picked = None
    for sobj in shared_objects:
        if _does_not_overlap(sobj, spec):
            assert sobj.size >= spec.allocated_memory, "Allocation specs are not sorted"
            picked = sobj
            sobj.first_used_index = min(sobj.first_used_index, spec.lifetime[0])
            sobj.last_used_index = max(sobj.last_used_index, spec.lifetime[1])
            allocation_spec = AllocationSpec(0, spec)
            picked.allocations.append(allocation_spec)
            break

    if picked is None and allow_overlapping_allocations:
        for sobj in shared_objects:
            max_offset = _find_max_overlapping_allocations_offset(sobj, spec)
            if max_offset > 0:
                if max_offset + spec.allocated_memory <= sobj.size:
                    picked = sobj
                    sobj.first_used_index = min(sobj.first_used_index, spec.lifetime[0])
                    sobj.last_used_index = max(sobj.last_used_index, spec.lifetime[1])
                    allocation_spec = AllocationSpec(max_offset, spec)
                    picked.allocations.append(allocation_spec)
                    break

    if picked is None:
        picked = SharedObject(
            len(shared_objects),
            -1,
            spec.allocated_memory,
            spec.lifetime[0],
            spec.lifetime[1],
        )
        allocation_spec = AllocationSpec(0, spec)
        picked.allocations.append(allocation_spec)
        picked.first_used_index = spec.lifetime[0]
        picked.last_used_index = spec.lifetime[1]
        shared_objects.append(picked)

    return picked


def get_node_tensor_specs(
    node: torch.fx.Node,
) -> Union[List[TensorSpec], Tuple[TensorSpec]]:
    r"""
    Return the list of the tensor specs for the node or empty list if the node
    has no tensor specs.
    """
    # get tensor specs
    if node.target == memory.view:
        base = node.args[0]
        assert isinstance(base, torch.fx.Node)
        specs = base.meta.get("spec")
    else:
        specs = node.meta.get("spec")

    if isinstance(specs, TensorSpec):
        specs = [specs]
    if not isinstance(specs, (list, tuple)):
        return []
    else:
        return [
            spec
            for spec in specs
            if not isinstance(spec, (int, float, bool, str, type(None)))
        ]


# Little bit hacky to check if the graph contains
# XNNPACK delegate
# Why?


def _contains_xnnpack_delegate(graph_module: torch.fx.GraphModule) -> bool:
    for node in graph_module.graph.nodes:
        if node.target == executorch_call_delegate:
            lowered_module = getattr(
                graph_module.graph.owning_module, node.args[0].target
            )
            if "xnnpack" in lowered_module.backend_id.lower():
                return True
    return False


def greedy(
    alignment: int,
    specs: Set[TensorSpec],
    graph_module: torch.fx.GraphModule,
    graph_signature: ExportGraphSignature,
    extra_padding: int = 0,
    *,
    allow_overlapping_allocations: bool = True,
) -> MemoryAlgoResult:
    r"""Greedy algorithm to allocate memory for tensors in the graph.

    Args:
        alignment: Memory alignment requirement
        specs: Set of TensorSpec objects with updated lifetimes
        graph_module: Graph module
        graph_signature: Graph signature
        extra_padding: Additional padding to add to each memory buffer (in bytes)
        allow_overlapping_allocations: If set to true, allows for allocations that overlap
            in their lifetime but are at different offsets in the storage. By default true.
            This flag is added to allow for Vulkan to use MemoryPlanningPass with overlapping
            allocations disabled

    Returns:
        MemoryAlgoResult containing the allocation decisions
    """
    greedy_result = MemoryAlgoResult({}, [])
    spec2obj = {}
    shared_objects = defaultdict(list)

    # For each tensor, pick the available shared object with closest size to
    # the tensor. If there are no available shared object left, create a new
    # one.
    import bisect

    sorted_specs = []
    for spec in specs:
        bisect.insort(sorted_specs, spec, key=lambda x: x.allocated_memory)

    sorted_specs.reverse()

    for spec in sorted_specs:
        # Create an entry for this TensorSpec in the result object that we'll be
        # returning from this algorithm.
        spec_alloc_result = greedy_result.spec_dict.get(spec, SpecAllocResult(0, 0, 0))
        if spec.mem_id is None:
            spec_alloc_result.mem_id = 1
        else:
            spec_alloc_result.mem_id = spec.mem_id
        greedy_result.spec_dict[spec] = spec_alloc_result
        spec.realign(alignment)
        spec2obj[spec] = pick_shared_obj(
            shared_objects[spec_alloc_result.mem_id],
            spec,
            allow_overlapping_allocations,
        )

    if len(shared_objects) == 0:
        # Cannot find any tensor in the graph that needs to be allocated.
        # Return [0, 0] to be consistent with default behavior of naive.
        total_sizes = [0, 0]
    else:
        total_sizes = [0] * (max(shared_objects.keys()) + 1)
        num_specs_processed = 0
        for mem_id in shared_objects:
            input_total_size = 0
            if bufsizes := getattr(graph_module, "input_mem_buffer_sizes", None):
                assert isinstance(bufsizes, list)
                if len(bufsizes) > mem_id:
                    input_total_size = bufsizes[mem_id]
            total_sizes[mem_id] = materialize_buffer(
                shared_objects[mem_id], input_total_size
            )
            total_sizes[mem_id] += extra_padding

            # Since we now know the number of shared objects we need and the size of
            # each shared object, we can assign offset in the memory buffer for each
            # shared object.
            for sobj in shared_objects[mem_id]:
                for alloc in sobj.allocations:
                    spec = alloc.spec
                    # Get the spec_alloc_result for this spec and update it with the
                    # mem_obj_id and mem_offset generated by this algorithm.
                    spec_alloc_result = greedy_result.spec_dict.get(spec, None)
                    assert spec_alloc_result is not None, f"Spec {spec} not found."
                    spec_alloc_result.mem_obj_id = sobj.idx
                    spec_alloc_result.mem_offset = sobj.offset + alloc.offset
                    num_specs_processed += 1
        assert (
            len(spec2obj) == num_specs_processed
        ), f"All specs should be processed but there were {len(spec2obj)} specs and processed {num_specs_processed} specs"

    logging.debug(f"greedy algorithm returns bufsizes: {total_sizes}")
    greedy_result.bufsizes = total_sizes
    return greedy_result


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()
                }
            )
            == 1
        ), "Different memory planning algorithms should have the same number of buffers allocated."

        # Find the algorithm that minimizes the total memory usage.
        best_algo = min(
            mem_algo_results, key=lambda k: sum(mem_algo_results[k].bufsizes)
        )
        logging.debug(f"Best memory planning algo for this model is {best_algo}")
        bufsizes = mem_algo_results[best_algo].bufsizes

        # Update the mem_id and mem_offset for each spec in the graph module based on the
        # values provided by the best memory planning algorithm.
        for spec in mem_algo_results[best_algo].spec_dict:
            spec_alloc_result = mem_algo_results[best_algo].spec_dict[spec]
            spec.mem_id = spec_alloc_result.mem_id
            spec.mem_offset = spec_alloc_result.mem_offset
            spec.mem_obj_id = spec_alloc_result.mem_obj_id

        return bufsizes


def naive(
    alignment: int,
    specs: Set[TensorSpec],
    graph_module: torch.fx.GraphModule,
    graph_signature: ExportGraphSignature,
    extra_padding: int,
) -> MemoryAlgoResult:
    """Naive algorithm to allocate memory for tensors in the graph.

    This algorithm simply allocates memory for each tensor sequentially without reusing memory.

    Args:
        alignment: Memory alignment requirement
        specs: Set of TensorSpec objects with updated lifetimes
        graph_module: Graph module
        graph_signature: Graph signature
        extra_padding: Additional padding to add to each memory buffer (in bytes)

    Returns:
        MemoryAlgoResult containing the allocation decisions
    """
    naive_result = MemoryAlgoResult({}, [])

    # allocate 'allocated' bytes from buffer with id mem_id.
    # return the starting offset of the allocated buffer.
    def _allocate_buf(bufsizes: List[int], mem_id: int, allocated: int) -> int:
        if mem_id >= len(bufsizes):
            bufsizes.extend([0] * (mem_id - len(bufsizes) + 1))
        ret = bufsizes[mem_id]
        bufsizes[mem_id] += allocated
        return ret

    bufsizes = getattr(graph_module, "input_mem_buffer_sizes", None)
    if bufsizes is None:
        bufsizes = [0, 0]
    bufsizes = cast(List[int], bufsizes)

    for spec in specs:
        spec_alloc_result = naive_result.spec_dict.get(spec, SpecAllocResult(0, 0, 0))
        # assume a single memory layer which has mem_id 1
        if spec.mem_id is None:
            spec_alloc_result.mem_id = 1
        else:
            spec_alloc_result.mem_id = spec.mem_id
        naive_result.spec_dict[spec] = spec_alloc_result

        # allocate spec.allocated_memory bytes in the buffer
        # with the corresponding mem_id
        spec.realign(alignment)
        spec_alloc_result.mem_offset = _allocate_buf(
            bufsizes, spec_alloc_result.mem_id, spec.allocated_memory
        )

    logging.debug(f"naive algorithm returns bufsizes: {bufsizes}")
    naive_result.bufsizes = bufsizes
    return naive_result


def get_cond_nodes(graph_module: torch.fx.GraphModule) -> Iterable[Node]:
    for nd in graph_module.graph.nodes:
        if nd.target is torch.ops.higher_order.cond:
            yield nd


def get_while_nodes(graph_module: torch.fx.GraphModule) -> Iterable[Node]:
    for nd in graph_module.graph.nodes:
        if nd.target is exir_while:
            yield nd


def get_map_nodes(graph_module: torch.fx.GraphModule) -> Iterable[Node]:
    for nd in graph_module.graph.nodes:
        if nd.target is torch.ops.higher_order.map_impl:
            yield nd


def get_scan_nodes(graph_module: torch.fx.GraphModule) -> Iterable[Node]:
    for nd in graph_module.graph.nodes:
        if nd.target is torch.ops.higher_order.scan:
            yield nd


def get_return_specs(graph_module: fx.GraphModule) -> Set[TensorSpec]:
    return_specs = set()
    nodes = graph_module.graph.nodes
    if len(nodes) > 0:
        last_node = next(iter(reversed(nodes)))
        for spec in tree_flatten(last_node.meta["spec"])[0]:
            return_specs.add(spec)
    return return_specs


def get_input_specs(graph_module: fx.GraphModule) -> Set[TensorSpec]:
    input_specs = set()
    nodes = graph_module.graph.nodes
    for node in nodes:
        if node.op == "placeholder":
            for spec in tree_flatten(node.meta["spec"])[0]:
                input_specs.add(spec)
    return input_specs


def insert_calls_to_free(
    graph_module: fx.GraphModule, allspecs: Set[TensorSpec]
) -> None:
    """
    Insert calls to free for dynamic unbound tensors that goes out of lifetime.

    Only handle the module itself. Submodule is handles in separate calls of
    this function.

    NOTE: this method will invalidate lifetime recorded in TensorSpec because
    of extra free node added to the graph.
    """
    # Note: we should never free a output tensor
    return_specs = get_return_specs(graph_module)
    # Note: we should never free a input tensor since buffer for input tensor
    # may be passed in from user.
    input_specs = get_input_specs(graph_module)
    idx_to_dead_specs = defaultdict(list)
    for spec in allspecs:
        if (
            spec.shape_dynamism == TensorShapeDynamism.DYNAMIC_UNBOUND
            and spec not in return_specs
            and spec not in input_specs
        ):
            idx_to_dead_specs[spec.lifetime[1]].append(spec)

    num_nodes = len(graph_module.graph.nodes)
    # iterate in reverse order so inserted node does not disturbe node
    # numbering.
    for node, node_idx in zip(
        reversed(graph_module.graph.nodes), range(num_nodes - 1, -1, -1)
    ):
        dead_specs = idx_to_dead_specs.get(node_idx, [])
        if not dead_specs:
            continue
        with graph_module.graph.inserting_after(node):
            for spec in dead_specs:
                graph_module.graph.call_function(memory.free, (spec,))
    graph_module.recompile()


def _merge_bufsizes(bufsizes: list[int], new_bufsizes: list[int]) -> list[int]:
    """Combine two buffer size lists."""
    if len(bufsizes) < len(new_bufsizes):
        bufsizes.extend([0] * (len(new_bufsizes) - len(bufsizes)))
    for i in range(len(new_bufsizes)):
        bufsizes[i] = max(bufsizes[i], new_bufsizes[i])
    return bufsizes


def _handle_submodule(
    algo: Callable[..., list[int]],
    parent_graph_module: torch.fx.GraphModule,
    alignment: int,
    submodule_node: torch.fx.Node,
    graph_signature: Optional[ExportGraphSignature] = None,
    alloc_graph_input: bool = False,
) -> list[int]:
    """Apply algo to nodes in a submodule of the graph module."""
    assert submodule_node.op == "get_attr"
    submodule = getattr(parent_graph_module, submodule_node.target)

    logging.debug(f"Planning memory for submodule {submodule_node.name}...")
    bufsizes = apply_algo(
        algo,
        submodule,
        alignment,
        graph_signature,
        alloc_graph_input=alloc_graph_input,
        alloc_graph_output=True,
    )
    submodule.meta.update({"non_const_buffer_sizes": bufsizes})
    logging.debug(f"Buffer sizes for submodule {submodule_node.name}: {bufsizes}")
    return bufsizes


def _apply_algo_to_submodules(
    algo: Callable[..., list[int]],
    graph_module: torch.fx.GraphModule,
    alignment: int,
    graph_signature: Optional[ExportGraphSignature] = None,
) -> list[int]:
    """Apply algo to map/cond/while/scan nodes in the graph module.

    This method will popuate graph_module.meta["non_const_buffer_sizes"] for
    all submodules and return a bufsizes list that is the maximum size of all
    buffers.
    """

    # Bufsizes for submodules.
    bufsizes: list[int] = []

    def _handle(
        submodule_node: torch.fx.Node,
        alloc_graph_input: bool = False,
    ) -> None:
        current_bufsizes = _handle_submodule(
            algo,
            graph_module,
            alignment,
            submodule_node,
            graph_signature,
            alloc_graph_input=alloc_graph_input,
        )
        nonlocal bufsizes
        _merge_bufsizes(bufsizes, current_bufsizes)

    for cond_node in get_cond_nodes(graph_module):
        _handle(cast(torch.fx.Node, cond_node.args[1]))
        _handle(cast(torch.fx.Node, cond_node.args[2]))

    for while_node in get_while_nodes(graph_module):
        _handle(cast(torch.fx.Node, while_node.args[0]))
        _handle(cast(torch.fx.Node, while_node.args[1]))

    for map_node in get_map_nodes(graph_module):
        _handle(cast(torch.fx.Node, map_node.args[0]), alloc_graph_input=True)

    for scan_node in get_scan_nodes(graph_module):
        _handle(cast(torch.fx.Node, scan_node.args[0]), alloc_graph_input=True)

    # TODO: We can handle delegates the same way as map/cond/while.
    # Maybe populate the graph_module.meta["non_const_buffer_sizes"] for delegates.

    return bufsizes


def apply_algo(
    algo: Callable[..., list[int]],
    graph_module: torch.fx.GraphModule,
    alignment: int,
    graph_signature: Optional[ExportGraphSignature] = None,
    alloc_graph_input: bool = True,
    alloc_graph_output: bool = True,
    alloc_mutable_buffers: bool = True,
) -> list[int]:
    """
    Recursively apply algo to graph_module and its submodules for control flow.

    Algo implementation should handle one of two meta entries for submodules:
    1. input_mem_buffer_sizes: List of int offset bytes. Memory allocated by
       `algo` should start at the offset specified by this list;
    OR
    2. non_const_buffer_sizes: List of bufsizes for planned memory in submodule.
       `algo` should reserve the space specified by this list for the lifetime
       of the submodule node (e.g. cond, while, map).

    TODO: Missing optimizations:
    1. To handle maps, we set `alloc_graph_input=True`, which allocates
    appropriate space for mapped arg but ends up allocating extra space for
    `operand` arg. The memory for operands is unused.
    """
    # Extract the nodes and their lifespans from the graph_module
    # Difficult to just filter the list of specs returned by this due to
    # how we flag trainable weights.
    _ = update_all_tensors_lifetime(graph_module, graph_signature)

    # Filter specs based on alloc_graph_input and alloc_graph_output
    specs = collect_specs_from_nodes(
        graph_module.graph.nodes,
        graph_signature,
        do_assertion=False,
        ignore_graph_input=not alloc_graph_input,
        ignore_graph_output=not alloc_graph_output,
        ignore_mutable_buffers=not alloc_mutable_buffers,
    )

    # Get temporary specs for submodules to set aside space during execution
    # of submodules.
    submodule_bufsizes = _apply_algo_to_submodules(
        algo, graph_module, alignment, graph_signature
    )

    # Update `input_mem_buffer_sizes` in graph_module. This will allow existing
    # algos to work using `input_mem_buffer_sizes` or use
    # `non_const_buffer_sizes` directly.
    # pyre-ignore[16]: `torch.fx.GraphModule` has no attribute `input_mem_buffer_sizes`.
    graph_module.input_mem_buffer_sizes = submodule_bufsizes

    # Get extra padding for XNNPACK if needed
    extra_padding = 0
    if _contains_xnnpack_delegate(graph_module):
        extra_padding = 64

    # Pass the filtered specs to the algorithm
    bufsizes: list[int] = algo(
        alignment,
        specs,
        graph_module,
        graph_signature,
        extra_padding,
    )

    # pyre-ignore[6]: Incompatible parameter type [6]
    # In call `insert_calls_to_free`, for 2nd positional argument, expected `Set[TensorSpec]` but got `Iterable[TensorSpec]`
    insert_calls_to_free(graph_module, specs)

    graph_module.meta.update({"non_const_buffer_sizes": bufsizes})
    return bufsizes