Function apply_algo

exir/memory_planning.py:1411–1544

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 neve

(
    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,
)

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.