Fix control flow output memory aliasing

exir/emit/test/test_emit.py

@@ -681,8 +681,12 @@ def false_fn(y: torch.Tensor) -> torch.Tensor:
 
         num_mm = 0
         num_add = 0
+        num_move = 0
         num_other = 0
         for inst in program.execution_plan[0].chains[0].instructions:
+            if isinstance(inst.instr_args, MoveCall):
+                num_move += 1
+                continue
             if not isinstance(inst.instr_args, KernelCall):
                 continue
 
@@ -697,8 +701,63 @@ def false_fn(y: torch.Tensor) -> torch.Tensor:
 
         self.assertEqual(num_mm, 2)
         self.assertEqual(num_add, 1)
+        self.assertEqual(num_move, 2)
         self.assertEqual(num_other, 0)
 
+    def test_emit_cond_output_lifetime(self) -> None:
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.register_buffer(
+                    "state", torch.tensor([torch.nan], dtype=torch.float64)
+                )
+
+            def forward(self, data: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+                value = data[0].unsqueeze(0)
+                pred = value >= 0.0
+                first = torch_cond(
+                    pred,
+                    lambda x: x.clone(),
+                    lambda x: self.state.clone(),
+                    [value],
+                )
+                second = torch_cond(
+                    pred,
+                    lambda x: torch_cond(
+                        torch.isnan(self.state),
+                        lambda y: y.clone(),
+                        lambda y: y + 1.0,
+                        [x],
+                    ),
+                    lambda x: self.state.clone(),
+                    [first],
+                )
+                self.state.copy_(first)
+                return first, second
+
+        inputs = (
+            torch.tensor([5.0], dtype=torch.float64),
+            torch.tensor([6.0], dtype=torch.float64),
+        )
+        model = M().eval()
+        program = to_edge(export(model, (inputs[0],), strict=True)).to_executorch(
+            config=ExecutorchBackendConfig(
+                passes=[InitializedMutableBufferPass(["state"])],
+                memory_planning_pass=MemoryPlanningPass(
+                    alloc_graph_input=False,
+                    alloc_graph_output=True,
+                ),
+            )
+        )
+        method = Runtime.get().load_program(program.buffer).load_method("forward")
+        eager_model = M().eval()
+
+        for value in inputs:
+            actual = method.execute([value])
+            expected = eager_model(value)
+            for actual_output, expected_output in zip(actual, expected):
+                torch.testing.assert_close(actual_output, expected_output)
+
     def test_emit_map(self) -> None:
         class Foo(torch.nn.Module):
             def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

exir/memory_planning.py

@@ -1326,6 +1326,46 @@ def _merge_bufsizes(bufsizes: list[int], new_bufsizes: list[int]) -> list[int]:
     return bufsizes
 
 
+def _allocate_submodule_buffer(
+    size: int,
+    lifetime: List[int],
+    allocations: List[Tuple[List[int], int, int]],
+) -> int:
+    overlapping = sorted(
+        (offset, allocated_size)
+        for allocated_lifetime, offset, allocated_size in allocations
+        if Verifier.has_overlap(lifetime, allocated_lifetime)
+    )
+    offset = 0
+    for allocated_offset, allocated_size in overlapping:
+        if offset + size <= allocated_offset:
+            break
+        offset = max(offset, allocated_offset + allocated_size)
+    allocations.append((lifetime, offset, size))
+    return offset
+
+
+def _shift_submodule_allocations(
+    submodule: torch.fx.GraphModule, offsets: List[int]
+) -> None:
+    shifted_specs: Set[TensorSpec] = set()
+    for module in submodule.modules():
+        if not isinstance(module, torch.fx.GraphModule):
+            continue
+        for node in module.graph.nodes:
+            for spec in get_node_tensor_specs(node):
+                if spec in shifted_specs:
+                    continue
+                shifted_specs.add(spec)
+                if spec.mem_id is None or spec.mem_offset is None:
+                    continue
+                internal_assert(
+                    spec.mem_id < len(offsets),
+                    f"Missing submodule buffer offset for mem_id {spec.mem_id}",
+                )
+                spec.mem_offset += offsets[spec.mem_id]
+
+
 def _handle_submodule(
     algo: Callable[..., list[int]],
     parent_graph_module: torch.fx.GraphModule,
@@ -1365,37 +1405,71 @@ def _apply_algo_to_submodules(
     buffers.
     """
 
-    # Bufsizes for submodules.
     bufsizes: list[int] = []
+    allocations: Dict[int, List[Tuple[List[int], int, int]]] = defaultdict(list)
 
-    def _handle(
-        submodule_node: torch.fx.Node,
-        alloc_graph_input: bool = False,
-    ) -> None:
-        current_bufsizes = _handle_submodule(
-            algo,
-            graph_module,
-            alignment,
-            submodule_node,
-            graph_signature,
-            alloc_graph_input=alloc_graph_input,
-        )
-        nonlocal bufsizes
-        _merge_bufsizes(bufsizes, current_bufsizes)
-
-    for cond_node in get_cond_nodes(graph_module):
-        _handle(cast(torch.fx.Node, cond_node.args[1]))
-        _handle(cast(torch.fx.Node, cond_node.args[2]))
+    for node_idx, node in enumerate(graph_module.graph.nodes):
+        submodule_nodes: List[torch.fx.Node] = []
+        alloc_graph_input = False
+        lifetime = [node_idx, node_idx]
+
+        if node.target is torch.ops.higher_order.cond:
+            submodule_nodes = [
+                cast(torch.fx.Node, node.args[1]),
+                cast(torch.fx.Node, node.args[2]),
+            ]
+            # MoveCall lets branch outputs escape their submodule, so the arena
+            # cannot be reused until the corresponding cond outputs are dead.
+            output_lifetimes = [
+                cast(int, spec.lifetime[1])
+                for spec in get_node_tensor_specs(node)
+                if spec.lifetime[1] is not None
+            ]
+            if output_lifetimes:
+                lifetime[1] = max(output_lifetimes)
+        elif node.target is exir_while:
+            submodule_nodes = [
+                cast(torch.fx.Node, node.args[0]),
+                cast(torch.fx.Node, node.args[1]),
+            ]
+        elif node.target is torch.ops.higher_order.map_impl:
+            submodule_nodes = [cast(torch.fx.Node, node.args[0])]
+            alloc_graph_input = True
+        elif node.target is torch.ops.higher_order.scan:
+            submodule_nodes = [cast(torch.fx.Node, node.args[0])]
+            alloc_graph_input = True
+        else:
+            continue
 
-    for while_node in get_while_nodes(graph_module):
-        _handle(cast(torch.fx.Node, while_node.args[0]))
-        _handle(cast(torch.fx.Node, while_node.args[1]))
+        current_bufsizes: list[int] = []
+        for submodule_node in submodule_nodes:
+            _merge_bufsizes(
+                current_bufsizes,
+                _handle_submodule(
+                    algo,
+                    graph_module,
+                    alignment,
+                    submodule_node,
+                    graph_signature,
+                    alloc_graph_input=alloc_graph_input,
+                ),
+            )
 
-    for map_node in get_map_nodes(graph_module):
-        _handle(cast(torch.fx.Node, map_node.args[0]), alloc_graph_input=True)
+        offsets = [0] * len(current_bufsizes)
+        for mem_id, size in enumerate(current_bufsizes):
+            if mem_id == 0 or size == 0:
+                continue
+            offsets[mem_id] = _allocate_submodule_buffer(
+                size, lifetime, allocations[mem_id]
+            )
+            if len(bufsizes) <= mem_id:
+                bufsizes.extend([0] * (mem_id + 1 - len(bufsizes)))
+            bufsizes[mem_id] = max(bufsizes[mem_id], offsets[mem_id] + size)
 
-    for scan_node in get_scan_nodes(graph_module):
-        _handle(cast(torch.fx.Node, scan_node.args[0]), alloc_graph_input=True)
+        for submodule_node in submodule_nodes:
+            _shift_submodule_allocations(
+                getattr(graph_module, submodule_node.target), offsets
+            )
 
     # TODO: We can handle delegates the same way as map/cond/while.
     # Maybe populate the graph_module.meta["non_const_buffer_sizes"] for delegates.

exir/tests/test_memory_planning.py

@@ -33,6 +33,7 @@
     apply_algo,
     collect_specs_from_nodes,
     filter_nodes,
+    get_cond_nodes,
     get_node_tensor_specs,
     greedy,
     MemoryAlgoResult,
@@ -54,6 +55,7 @@
 from functorch.experimental.control_flow import map as torch_map
 from parameterized import parameterized
 from torch import nn
+from torch._higher_order_ops import cond as torch_cond
 from torch.ao.quantization import (  # @manual=//caffe2:torch
     float_qparams_weight_only_qconfig,
 )
@@ -1122,6 +1124,70 @@ def _get_specs(gm: torch.fx.GraphModule) -> set[TensorSpec]:
     )
 
 
+class ConsecutiveCondModel(torch.nn.Module):
+    def forward(
+        self, pred: torch.Tensor, data: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        first = torch_cond(
+            pred,
+            lambda x: x + 1,
+            lambda x: x - 1,
+            [data],
+        )
+        second = torch_cond(
+            pred,
+            lambda x: x + 2,
+            lambda x: x - 2,
+            [first],
+        )
+        return first, second
+
+
+class TestCond(unittest.TestCase):
+    def test_consecutive_cond_output_lifetime(self) -> None:
+        graph_module = (
+            to_edge(
+                export(
+                    ConsecutiveCondModel(),
+                    (torch.tensor(True), torch.ones(2)),
+                    strict=True,
+                )
+            )
+            .exported_program()
+            .graph_module
+        )
+        graph_module = PassManager(
+            passes=[
+                SpecPropPass(),
+                ToOutVarPass(),
+            ],
+        )(graph_module).graph_module
+        graph_module = MemoryPlanningPass().run(graph_module).graph_module
+
+        cond_nodes = list(get_cond_nodes(graph_module))
+        self.assertEqual(len(cond_nodes), 2)
+
+        def branch_output_specs(
+            cond_node: torch.fx.Node,
+        ) -> List[List[TensorSpec]]:
+            outputs = []
+            for branch_node in cond_node.args[1:3]:
+                self.assertIsInstance(branch_node, torch.fx.Node)
+                branch = getattr(graph_module, branch_node.target)
+                outputs.append(get_node_tensor_specs(branch.graph.output_node()))
+            return outputs
+
+        first_outputs = branch_output_specs(cond_nodes[0])
+        second_outputs = branch_output_specs(cond_nodes[1])
+        for first_branch_outputs in first_outputs:
+            for second_branch_outputs in second_outputs:
+                for first_spec in first_branch_outputs:
+                    for second_spec in second_branch_outputs:
+                        self.assertFalse(
+                            Verifier.storage_overlap(first_spec, second_spec)
+                        )
+
+
 class MapModel(torch.nn.Module):
     def __init__(self) -> None:
         super().__init__()