Function greedy

exir/memory_planning.py:935–1028

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 p

(
    alignment: int,
    specs: Set[TensorSpec],
    graph_module: torch.fx.GraphModule,
    graph_signature: ExportGraphSignature,
    extra_padding: int = 0,
    *,
    allow_overlapping_allocations: bool = True,
)

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.