insert_write_back_for_buffers_pass should inject copy_ nodes at the earliest possible spot.

[JacobSzwejbka]

🐛 Describe the bug

Currently the pass injects all the 'copy_' nodes at the end of the graph. It should instead insert them immediately following the last usage of the thing being copied into the original mutable buffer. By waiting till the end we are arbitrarily extending the lifetime of those tensors potentially wasting a non trivial amount of space in the memory plan.

Versions

.

cc @angelayi

[prmaynard]

Looks interesting as a first issue, can I give this one a try?

[JacobSzwejbka]

Looks interesting as a first issue, can I give this one a try?

Yeah feel free to take a crack at it and I'll review any PR you put up. If you need some more guidance let me know.

the pass is called in _program.py.

[metascroy]

@prmaynard are you still working on this? If not, we'd like to return it to the pool.

[prmaynard]

@metascroy @JacobSzwejbka Yes, I have looked into this and have a proposed solution. I will try to submit a PR within a day or so. Should I add a test case to test_passes.py to verify we are inserting the copy nodes directly after their last use?

[RihaanBH-1810]

@JacobSzwejbka @metascroy can I please take this up as my first issue, since there hasn't been much activity by the previous assignee.

[prmaynard]

@RihaanBH-1810 please feel welcome to take a crack at this! Here is where I was able to get (forgive me it has been a while since I looked at this, life got busy)

def _insert_copy(
    gm: torch.fx.GraphModule,
    mutated_outputs: List[Optional[str]],
    input_name_to_node: Dict[str, torch.fx.Node],
):

    if all(name is None for name in mutated_outputs): # if all mutated outputs are none, do not execute
        return []    

    output_node = None
    for node in gm.graph.nodes:
        if node.op == "output":
            output_node = node
            break
    assert output_node is not None

    outputs = pytree.tree_flatten(output_node.args)[0]
    assert len(outputs) == len(mutated_outputs)

    user_returns = []
    buffer_copies = []

    for return_node, mutated_node_name in zip(outputs, mutated_outputs):

        if mutated_node_name is None:
            user_returns.append(return_node)
            continue

        # Mutable buffer grab the node
        if mutated_node_name in input_name_to_node:
            mutated_node = input_name_to_node[mutated_node_name]
        else:
            raise RuntimeError(
                f"Could not find {mutated_node_name} in either buffer or input nodes"
            )
    
        # needed to rebuild node index pairs after graph modification
        nodes = list(gm.graph.nodes)
        node_index = {node: i for i, node in enumerate(nodes)}

        # last reader
        last_read_idx = node_index[mutated_node]

        for user in mutated_node.users:
            if user.op != "output":
                last_read_idx = max(last_read_idx, node_index[user])


        # last consumer
        last_use_idx = node_index[return_node]

        for user in return_node.users:
            if user.op != "output":
                last_use_idx = max(last_use_idx, node_index[user])

        insert_after = max(last_read_idx, last_use_idx)

        if insert_after + 1 < len(nodes):
            insert_point = (nodes[insert_after + 1])
        else:
            insert_point = (output_node)


        with gm.graph.inserting_before(insert_point):
            buffer_output = gm.graph.call_function(
                torch.ops.aten.copy_.default, (mutated_node, return_node)
            )
            buffer_copies.append(buffer_output)

    # Re‑wire graph output so that the copy results precede user returns
    with gm.graph.inserting_before(output_node):
        buffer_copies.extend(user_returns)
        new_out = gm.graph.output(tuple(buffer_copies))
        output_node.replace_all_uses_with(new_out)
        gm.graph.erase_node(output_node)

    return buffer_copies

The logic for early node insertion is not too difficult, but the main issue I ran into was failing later unit tests. If I recall correctly, the problem was that you needed to be able to detect last use AND the last use of any aliases. I was unable to find a safe/encompassing way of detecting this.