executorch 1.6.0__cp315-cp315-win_amd64.whl

executorch/exir/memory_planning.py

@@ -0,0 +1,1567 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# 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.inplace_aliasing import is_inplace_node
+from executorch.exir.operator.convert import is_out_variant
+from executorch.exir.schema import DeviceType, NonConstBufferDevice, 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 _storage_root_chain(cls, spec: TensorSpec) -> set[TensorSpec]:
+        """Return spec and all TensorSpecs backing its storage."""
+        seen: Set[TensorSpec] = set()
+        root = spec
+        while True:
+            internal_assert(
+                root not in seen,
+                "Circular storage_base relationship is not supported.",
+            )
+            seen.add(root)
+            if root.storage_base is None:
+                break
+            root = root.storage_base
+        return seen
+
+    @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
+                ):
+                    # Some ops, such as in-place ops, intentionally place one
+                    # TensorSpec inside another TensorSpec's storage despite
+                    # overlapping lifetimes.
+                    # This is OK if one spec is storage-backed by the other.
+                    # Specs that merely share a root, such as siblings, should
+                    # still be checked as normal allocations.
+                    lhs_chain = Verifier._storage_root_chain(lhs_spec)
+                    rhs_chain = Verifier._storage_root_chain(rhs_spec)
+                    is_common_base_pair = lhs_spec in rhs_chain or rhs_spec in lhs_chain
+                    if not is_common_base_pair:
+                        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 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) or _is_device_resident_constant(node)
+    ):
+        # Mutable buffers and device resident constants keep their data across
+        # calls, so they 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 _is_device_resident_constant(node: Node) -> bool:
+    """Whether node is a constant that PropagateDevicePass placed in device
+    memory. The runtime copies its data from the program or its data file
+    into its planned memory once, at load, so it is planned like a mutable
+    buffer instead of being used in place."""
+    return node.op == "placeholder" and node.meta.get("device_resident_constant", 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)
+            and not _is_device_resident_constant(node)
+        ):
+            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()
+    )
+
+    # Mutation outputs alias persistent state, whose allocation is controlled
+    # by ignore_mutable_buffers independently of graph output allocation.
+    mutable_buffer_specs: Set[TensorSpec] = set()
+    if ignore_mutable_buffers or ignore_graph_output:
+        mutable_buffer_specs = _get_mutable_buffer_specs(nodes, graph_signature)
+    graph_output_tensors -= mutable_buffer_specs
+
+    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_mutable_buffer(node, graph_signature) and ignore_mutable_buffers:
+            continue
+
+        if do_assertion:
+            internal_assert(
+                node.op in ("placeholder", "output")
+                or is_inplace_node(node)
+                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)
+                and not _is_device_resident_constant(node)
+            ):
+                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)
+    # Scratch is live only while the delegate is executing, but it is reached
+    # through the delegate node, so the walk above extends it every time a
+    # consumer of the call's outputs mentions that node. Recompute it from the
+    # delegate calls alone rather than letting those consumers stretch it.
+    #
+    # Accumulated across calls rather than assigned per call: fx copies
+    # node.meta shallowly, so a pass that clones a delegate call leaves two
+    # nodes sharing one spec. Covering both is conservative; assigning would
+    # leave the buffer looking dead while the earlier call still writes to it.
+    scratch_lifetimes: Dict[TensorSpec, Tuple[int, int]] = {}
+    for node_idx, node in enumerate(graph_module.graph.nodes):
+        for spec in memory.delegate_scratch_specs(node):
+            start, end = scratch_lifetimes.get(spec, (node_idx, node_idx))
+            scratch_lifetimes[spec] = (min(start, node_idx), max(end, node_idx))
+    for spec, (start, end) in scratch_lifetimes.items():
+        spec.lifetime = [start, end]
+    _extend_storage_base_lifetimes(specs)
+    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 _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.
+
+    A delegate call's scratch buffers are included. They are not values the
+    node produces, but every caller that walks specs to place, verify or
+    attribute memory needs them.
+    """
+    # 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)):
+        specs = []
+    else:
+        specs = [
+            spec
+            for spec in specs
+            if not isinstance(spec, (int, float, bool, str, type(None)))
+        ]
+
+    if scratch := memory.delegate_scratch_specs(node):
+        return list(specs) + scratch
+    return specs
+
+
+# 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 _extend_storage_base_lifetimes(specs: Set[TensorSpec]) -> None:
+    for spec in specs:
+        if spec.storage_base is None:
+            continue
+        internal_assert(
+            spec.lifetime[0] is not None and spec.lifetime[1] is not None,
+            "Storage-backed TensorSpec must have a lifetime.",
+        )
+        start = cast(int, spec.lifetime[0])
+        end = cast(int, spec.lifetime[1])
+        seen: Set[TensorSpec] = {spec}
+        base = spec.storage_base
+        while base is not None:
+            internal_assert(
+                base not in seen,
+                "Circular storage_base relationship is not supported.",
+            )
+            seen.add(base)
+            internal_assert(
+                base.lifetime[0] is not None and base.lifetime[1] is not None,
+                "storage_base TensorSpec must have a lifetime.",
+            )
+            base.lifetime[0] = min(cast(int, base.lifetime[0]), start)
+            base.lifetime[1] = max(cast(int, base.lifetime[1]), end)
+            base = base.storage_base
+
+
+def _resolve_storage_base_specs(
+    deferred_storage_base: List[TensorSpec],
+    spec2obj: Dict[TensorSpec, SharedObject],
+    greedy_result: MemoryAlgoResult,
+) -> None:
+    remaining = list(deferred_storage_base)
+    while remaining:
+        progress = False
+        next_remaining = []
+        for spec in remaining:
+            base = spec.storage_base
+            internal_assert(
+                base is not None,
+                "Deferred storage-backed TensorSpec should have a storage_base.",
+            )
+            if base not in spec2obj:
+                next_remaining.append(spec)
+                continue
+            progress = True
+            sobj = spec2obj[base]
+
+            base_alloc_result = greedy_result.spec_dict[base]
+            spec_alloc_result = greedy_result.spec_dict[spec]
+            spec_alloc_result.mem_id = base_alloc_result.mem_id
+
+            allocated_memory = spec.allocated_memory
+            storage_base_offset = spec.storage_base_offset
+            internal_assert(
+                storage_base_offset >= 0,
+                "storage_base_offset must be non-negative.",
+            )
+            base_alloc_offset = None
+            base_allocated_memory = None
+            for alloc_entry in sobj.allocations:
+                if alloc_entry.spec is base:
+                    base_alloc_offset = alloc_entry.offset
+                    base_allocated_memory = alloc_entry.spec.allocated_memory
+                    break
+            assert base_alloc_offset is not None, (
+                f"Base allocation entry not found in shared object for spec "
+                f"with allocated_memory={spec.allocated_memory}"
+            )
+            assert base_allocated_memory is not None, (
+                f"Base allocation entry not found in shared object for spec "
+                f"with allocated_memory={allocated_memory}"
+            )
+            internal_assert(
+                (base_alloc_offset + storage_base_offset) % spec.alignment == 0,
+                f"Storage-backed TensorSpec allocation must respect alignment, got offset {storage_base_offset} inside parent with offset {base_alloc_offset} for alignment {spec.alignment}.",
+            )
+            internal_assert(
+                storage_base_offset + allocated_memory <= base_allocated_memory,
+                "Storage-backed TensorSpec allocation must fit within storage_base.",
+            )
+            sobj.first_used_index = min(sobj.first_used_index, spec.lifetime[0])
+            sobj.last_used_index = max(sobj.last_used_index, spec.lifetime[1])
+            sobj.allocations.append(
+                AllocationSpec(base_alloc_offset + storage_base_offset, spec)
+            )
+            spec2obj[spec] = sobj
+        if not progress:
+            unresolved = ", ".join(
+                f"allocated_memory={s.allocated_memory}" for s in next_remaining
+            )
+            raise InternalError(
+                "Circular or unresolvable storage_base dependency chain: " + unresolved
+            )
+        remaining = next_remaining
+
+
+def _compute_total_sizes(
+    shared_objects: Dict[int, List[SharedObject]],
+    graph_module: torch.fx.GraphModule,
+    extra_padding: int,
+    greedy_result: MemoryAlgoResult,
+    num_specs_expected: int,
+) -> List[int]:
+    if len(shared_objects) == 0:
+        return [0, 0]
+
+    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_size = input_total_size
+        for sobj in shared_objects[mem_id]:
+            sobj.offset = total_size
+            total_size += sobj.size
+        total_sizes[mem_id] = total_size
+        total_sizes[mem_id] += extra_padding
+
+        for sobj in shared_objects[mem_id]:
+            for alloc in sobj.allocations:
+                spec_alloc_result = greedy_result.spec_dict.get(alloc.spec, None)
+                assert spec_alloc_result is not None, f"Spec {alloc.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 (
+        num_specs_expected == num_specs_processed
+    ), f"All specs should be processed but there were {num_specs_expected} specs and processed {num_specs_processed} specs"
+    return total_sizes
+
+
+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: Dict[TensorSpec, SharedObject] = {}
+    shared_objects: Dict[int, List[SharedObject]] = defaultdict(list)
+
+    import bisect
+
+    sorted_specs = []
+    for spec in specs:
+        bisect.insort(sorted_specs, spec, key=lambda x: x.allocated_memory)
+
+    sorted_specs.reverse()
+
+    deferred_storage_base: List[TensorSpec] = []
+
+    for spec in sorted_specs:
+        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)
+
+        if spec.storage_base is not None:
+            deferred_storage_base.append(spec)
+            continue
+
+        spec2obj[spec] = pick_shared_obj(
+            shared_objects[spec_alloc_result.mem_id],
+            spec,
+            allow_overlapping_allocations,
+        )
+
+    _resolve_storage_base_specs(deferred_storage_base, spec2obj, greedy_result)
+
+    total_sizes = _compute_total_sizes(
+        shared_objects, graph_module, extra_padding, greedy_result, len(spec2obj)
+    )
+
+    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:
+        if spec.storage_base is not None:
+            raise InternalError(
+                "The naive memory planning algorithm does not support storage-backed "
+                "TensorSpecs. Use the greedy algorithm instead."
+            )
+
+        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
+
+
+_CPU_KEY: tuple[DeviceType, int] = (DeviceType.CPU, 0)
+
+
+def _partition_specs_by_device(
+    all_specs: list[TensorSpec],
+    enable_non_cpu_memory_planning: bool,
+) -> dict[tuple[DeviceType, int], list[TensorSpec]]:
+    """Partition specs by (device_type, device_index).
+
+    Different device indices on the same device type (e.g. CUDA:0 vs CUDA:1)
+    get separate memory buffers.
+
+    When ``enable_non_cpu_memory_planning`` is False (legacy), all specs are
+    placed into a single CPU:0 bucket regardless of their device attribute.
+
+    Insertion order is preserved within each partition because order-sensitive
+    algorithms (e.g. greedy with bisect.insort) rely on it for stable tie-breaking.
+    """
+    specs_by_device: dict[tuple[DeviceType, int], list[TensorSpec]] = defaultdict(list)
+    if not enable_non_cpu_memory_planning:
+        specs_by_device[_CPU_KEY] = all_specs
+        return specs_by_device
+
+    has_non_cpu_specs = False
+    has_pre_assigned_mem_id = False
+    for spec in all_specs:
+        device_key = (spec.device, spec.device_index)
+        specs_by_device[device_key].append(spec)
+        if spec.device != DeviceType.CPU:
+            has_non_cpu_specs = True
+        if spec.mem_id is not None:
+            has_pre_assigned_mem_id = True
+
+    # Custom pool passes pre-assign mem_ids (e.g. mem_id=2, 3, …) to place
+    # tensors into specific memory arenas.  Per-device partitioning appends
+    # device buffers after the CPU buffers, and the remap formula
+    #   global_mem_id = (local_mem_id - 1) + base_mem_id
+    # assumes the algo-local numbering starts at 1.  If a custom pass has
+    # already set mem_ids > 1 on the CPU side, the device-buffer slots may
+    # collide with those custom pool slots.
+    # TODO(gasoonjia): support custom pools + per-device planning by reserving
+    # device slots after the highest custom pool id.
+    if has_non_cpu_specs and has_pre_assigned_mem_id:
+        raise NotImplementedError(
+            "enable_non_cpu_memory_planning is not yet compatible with "
+            "custom memory pool passes that pre-assign spec.mem_id. "
+            "The per-device buffer slots may collide with custom pool "
+            "mem_ids. Please disable enable_non_cpu_memory_planning or "
+            "remove the custom mem_id assignments."
+        )
+
+    return specs_by_device
+
+
+def _build_non_const_buffer_device(
+    buffer_devices: list[tuple[DeviceType, int]],
+) -> Optional[list[NonConstBufferDevice]]:
+    """Build the non-CPU buffer device list for serialization.
+
+    Returns ``None`` when all buffers are CPU (the default), so that no
+    redundant device metadata is emitted.
+    """
+    if not any(dk[0] != DeviceType.CPU for dk in buffer_devices):
+        return None
+    return [
+        NonConstBufferDevice(buffer_idx=i, device_type=dt, device_index=di)
+        for i, (dt, di) in enumerate(buffer_devices)
+        if (dt, di) != _CPU_KEY
+    ]
+
+
+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,
+    enable_non_cpu_memory_planning: bool = False,
+) -> list[int]:
+    """
+    Recursively apply algo to graph_module and its submodules for control flow.
+
+    Partitions specs by device type and device idx, and runs the memory planning
+    algorithm independently per device, then merges results into separate buffers.
+    This ensures device memory and CPU memory are never mixed.
+
+    When enable_non_cpu_memory_planning is False (default), all specs are planned
+    into a single CPU memory pool regardless of their device attribute. This
+    preserves the legacy behavior. Set to True to enable per-device partitioning.
+
+    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
+    _ = update_all_tensors_lifetime(graph_module, graph_signature)
+
+    # Collect specs into an ordered list so we can iterate multiple times and
+    # partition by device.  Order matters: order-sensitive algorithms (e.g.
+    # greedy with bisect.insort) rely on insertion order for stable tie-breaking,
+    # and `collect_specs_from_nodes` already deduplicates via its `dedup` flag.
+    all_specs: list[TensorSpec] = list(
+        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.
+    # NOTE: submodule_bufsizes are currently applied only to the CPU partition.
+    # This assumes all control-flow submodule tensors (cond/while/map) live in
+    # CPU memory.  Today this is safe because on-device tensors only appear as
+    # delegate blob I/O, which never lives inside control-flow submodules.
+    # If device tensors ever appear in submodules, _apply_algo_to_submodules
+    # will need per-device partitioning as well.
+    submodule_bufsizes = _apply_algo_to_submodules(
+        algo, graph_module, alignment, graph_signature
+    )
+
+    # Get extra padding for XNNPACK if needed
+    extra_padding = 0
+    if _contains_xnnpack_delegate(graph_module):
+        extra_padding = 64
+
+    specs_by_device = _partition_specs_by_device(
+        all_specs, enable_non_cpu_memory_planning
+    )
+
+    # Plan each device independently
+    global_bufsizes: list[int] = [0]  # index 0 reserved for constants
+    buffer_devices: list[tuple[DeviceType, int]] = [_CPU_KEY]
+
+    # Process CPU:0 first (if present), then other devices sorted by
+    # (type.value, index) so the ordering is deterministic.
+    device_order = sorted(
+        specs_by_device.keys(),
+        key=lambda dk: (dk != _CPU_KEY, dk[0].value, dk[1]),
+    )
+
+    for device_key in device_order:
+        device_specs = specs_by_device[device_key]
+
+        # Only apply submodule pre-allocation for CPU specs; device buffers
+        # do not share memory space with CPU submodule arenas.
+        # pyre-ignore[16]: `torch.fx.GraphModule` has no attribute `input_mem_buffer_sizes`.
+        graph_module.input_mem_buffer_sizes = (
+            submodule_bufsizes if device_key == _CPU_KEY else []
+        )
+
+        # Run algorithm independently on this device's specs
+        device_bufsizes = algo(
+            alignment, device_specs, graph_module, graph_signature, extra_padding
+        )
+
+        # Calculate base mem_id in global space
+        base_mem_id = len(global_bufsizes)
+
+        # Append buffer sizes (skip index 0 which is constants placeholder)
+        global_bufsizes.extend(device_bufsizes[1:])
+
+        # Track device key for each new buffer slot
+        for _ in device_bufsizes[1:]:
+            buffer_devices.append(device_key)
+
+        # Remap spec mem_ids from algo-local to global.
+        # At this point spec.mem_id has been set by MemoryPlanningAlgorithmSuite:
+        # the suite runs each algorithm (e.g. greedy), picks the best result,
+        # and writes the winning mem_id/mem_offset/mem_obj_id back onto each
+        # spec.  For specs with no pre-assigned mem_id the algorithm defaults
+        # to mem_id=1; custom-pool passes may pre-assign other values (e.g. 3).
+        # We remap from the algo-local numbering (1-based) to the global
+        # position: global_mem_id = (local_mem_id - 1) + base_mem_id.
+        for spec in device_specs:
+            if spec.mem_id is not None:
+                spec.mem_id = (spec.mem_id - 1) + base_mem_id
+
+    # Ensure backward compatibility: at least [0, 0] when no specs exist
+    if len(global_bufsizes) < 2:
+        global_bufsizes.append(0)
+        buffer_devices.append(_CPU_KEY)
+
+    # Insert free calls and build device buffer mapping
+    insert_calls_to_free(graph_module, all_specs)
+
+    non_const_buffer_device = _build_non_const_buffer_device(buffer_devices)
+    graph_module.meta["non_const_buffer_sizes"] = global_bufsizes
+    if non_const_buffer_device is not None:
+        graph_module.meta["non_const_buffer_device"] = non_const_buffer_device
+    return global_bufsizes