Function naive

exir/memory_planning.py:1111–1166

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_modul

(
    alignment: int,
    specs: Set[TensorSpec],
    graph_module: torch.fx.GraphModule,
    graph_signature: ExportGraphSignature,
    extra_padding: int,
)

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