insert_write_back_for_buffers_pass should inject copy_ nodes at the earliest possible spot. #7345

Activity

1. JacobSzwejbka

   added

   module: exirIssues related to Export IR and the code under exir/Issues related to Export IR and the code under exir/

   bug

   on Dec 17, 2024
2. JacobSzwejbka

   added

   triagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate moduleThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

   on Feb 4, 2025
3. jackzhxng

   removed

   bug

   on Feb 21, 2025
4. github-project-automation

   added this to ExecuTorch Coreon Feb 21, 2025
5. github-project-automation

   moved this to To triage in ExecuTorch Coreon Feb 21, 2025
6. JacobSzwejbka

   added

   good first issueGood for newcomersGood for newcomers

   on Mar 10, 2025
7. iseeyuan

   moved this from To triage to Backlog in ExecuTorch Coreon Mar 11, 2025
8. byjlw

   added this to New Contributors Projects and Issueson Mar 20, 2025
9. prmaynard commented on Apr 2, 2025

   

   prmaynard

   on Apr 2, 2025

   More actions

   Looks interesting as a first issue, can I give this one a try?
10. JacobSzwejbka commented on Apr 7, 2025

    

    JacobSzwejbka

    on Apr 7, 2025

    ContributorAuthor

    More actions

    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.

    👍React with 👍1Reacted by prmaynard
11. metascroy commented on Apr 22, 2025

    

    metascroy

    on Apr 22, 2025

    Contributor

    More actions

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

    

    prmaynard

    on Apr 23, 2025

    More actions

    @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?
13. RihaanBH-1810 commented on Aug 21, 2025

    

    RihaanBH-1810

    on Aug 21, 2025

    More actions

    @JacobSzwejbka @metascroy can I please take this up as my first issue, since there hasn't been much activity by the previous assignee.
14. prmaynard commented on Aug 21, 2025

    

    prmaynard

    on Aug 21, 2025

    Last edited by prmaynard

    More actions

    @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.
15. 1 remaining item

    Load more

16. JPL11 commented on Jul 6, 2026

    

    JPL11

    on Jul 6, 2026

    More actions

    I picked this up and put up #20744. It inserts each write-back copy_ at the earliest point that preserves the current end-of-graph semantics, with a forward alias walk (schema alias_info, conservative for getitem/submodules/unknown schemas) to handle the last-use-of-aliases problem discussed above — a reader of the buffer's old value keeps the copy_ after it. There's a regression test for exactly that case. Happy to iterate on review feedback.
17. JPL11

    added 14 commits that reference this issue on Oct 10, 2026

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit 1bfe8cd is merged into the 'main' branch.

    1bfe8cd

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit fc6ef41 is merged into the 'main' branch.

    fc6ef41

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit 518bbb9 is merged into the 'main' branch.

    518bbb9

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit d46a942 is merged into the 'main' branch.

    d46a942

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit b13fcf9 is merged into the 'main' branch.

    b13fcf9

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit b7ca939 is merged into the 'main' branch.

    b7ca939

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit d6b3823 is merged into the 'main' branch.

    d6b3823

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit fae6eb4 is merged into the 'main' branch.

    fae6eb4

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit a10d28b is merged into the 'main' branch.

    a10d28b

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit 435110b is merged into the 'main' branch.

    435110b

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit 87b0905 is merged into the 'main' branch.

    87b0905

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit 3cb4fc6 is merged into the 'main' branch.

    3cb4fc6

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit e5e451e is merged into the 'main' branch.

    e5e451e

    Insert write-back copy_ nodes at the earliest safe point

    ...

    This issue will close once commit db676b4 is merged into the 'main' branch.

    db676b4