Function update_tensor_lifetime

exir/memory_planning.py:443–470

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

(
    node: torch.fx.Node,
    spec: TensorSpec,
    node_idx: int,
    max_node_idx: int,
    gs: Optional[ExportGraphSignature] = None,
)

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