diff --git a/exir/memory_planning.py b/exir/memory_planning.py index 012cf8dd144..6790b84ed4e 100644 --- a/exir/memory_planning.py +++ b/exir/memory_planning.py @@ -1232,28 +1232,27 @@ def _allocate_buf(bufsizes: List[int], mem_id: int, allocated: int) -> int: return naive_result -def get_cond_nodes(graph_module: torch.fx.GraphModule) -> Iterable[Node]: - for nd in graph_module.graph.nodes: - if nd.target is torch.ops.higher_order.cond: - yield nd - - -def get_while_nodes(graph_module: torch.fx.GraphModule) -> Iterable[Node]: - for nd in graph_module.graph.nodes: - if nd.target is exir_while: - yield nd - - -def get_map_nodes(graph_module: torch.fx.GraphModule) -> Iterable[Node]: - for nd in graph_module.graph.nodes: - if nd.target is torch.ops.higher_order.map_impl: - yield nd - +def _hop_output_last_use( + hop_node: torch.fx.Node, hop_idx: int, node_index: dict[torch.fx.Node, int] +) -> int: + """Return the max node index at which any output of a HOP node is last consumed. -def get_scan_nodes(graph_module: torch.fx.GraphModule) -> Iterable[Node]: - for nd in graph_module.graph.nodes: - if nd.target is torch.ops.higher_order.scan: - yield nd + HOP outputs are accessed via getitem nodes. We find the last consumer of + any getitem user of the HOP node. + """ + last_use = hop_idx + for user in hop_node.users: + spec = user.meta.get("spec") + if isinstance(spec, TensorSpec) and spec.lifetime[1] is not None: + last_use = max(last_use, spec.lifetime[1]) + elif isinstance(spec, (list, tuple)): + for s in spec: + if isinstance(s, TensorSpec) and s.lifetime[1] is not None: + last_use = max(last_use, s.lifetime[1]) + else: + for consumer in user.users: + last_use = max(last_use, node_index[consumer]) + return last_use def get_return_specs(graph_module: fx.GraphModule) -> Set[TensorSpec]: @@ -1333,11 +1332,15 @@ def _handle_submodule( submodule_node: torch.fx.Node, graph_signature: Optional[ExportGraphSignature] = None, alloc_graph_input: bool = False, + input_mem_buffer_sizes: Optional[list[int]] = None, ) -> list[int]: """Apply algo to nodes in a submodule of the graph module.""" assert submodule_node.op == "get_attr" submodule = getattr(parent_graph_module, submodule_node.target) + if input_mem_buffer_sizes is not None: + submodule.input_mem_buffer_sizes = input_mem_buffer_sizes + logging.debug(f"Planning memory for submodule {submodule_node.name}...") bufsizes = apply_algo( algo, @@ -1360,42 +1363,86 @@ def _apply_algo_to_submodules( ) -> list[int]: """Apply algo to map/cond/while/scan nodes in the graph module. - This method will popuate graph_module.meta["non_const_buffer_sizes"] for - all submodules and return a bufsizes list that is the maximum size of all - buffers. - """ + Processes HOP nodes in graph order, tracking output lifetimes. When a later + HOP's execution start falls within an earlier HOP's output lifetime, the + later HOP's submodules allocate at offsets AFTER the earlier HOP's memory + to prevent overwriting still-live outputs. - # Bufsizes for submodules. + Branches within the SAME HOP (e.g. true/false of one cond) are mutually + exclusive and safely share memory (merged via max). + """ bufsizes: list[int] = [] - 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) + parent_offset = getattr(graph_module, "input_mem_buffer_sizes", None) - 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])) + # (submodule_arg_indices, alloc_graph_input, outputs_in_arena) + # outputs_in_arena: cond/while outputs live in the submodule arena (via + # MoveCall), so a later HOP sharing that arena would corrupt them. + # map/scan outputs are allocated in the parent graph — no conflict. + hop_targets: dict = { + torch.ops.higher_order.cond: ((1, 2), False, True), + exir_while: ((0, 1), False, True), + torch.ops.higher_order.map_impl: ((0,), True, False), + torch.ops.higher_order.scan: ((0,), True, False), + } - 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])) + # Track HOPs whose outputs live in the arena: + # (bufsizes_for_this_hop, hop_start_idx, hop_output_last_use) + hop_allocs: list[tuple[list[int], int, int]] = [] - for map_node in get_map_nodes(graph_module): - _handle(cast(torch.fx.Node, map_node.args[0]), alloc_graph_input=True) + node_index: Optional[Dict[torch.fx.Node, int]] = None - for scan_node in get_scan_nodes(graph_module): - _handle(cast(torch.fx.Node, scan_node.args[0]), alloc_graph_input=True) + for node in graph_module.graph.nodes: + if node.op != "call_function" or node.target not in hop_targets: + continue + + if node_index is None: + node_index = {n: i for i, n in enumerate(graph_module.graph.nodes)} + + submodule_arg_indices, alloc_graph_input, outputs_in_arena = hop_targets[ + node.target + ] + hop_start_idx = node_index[node] + hop_end_idx = _hop_output_last_use(node, hop_start_idx, node_index) + + # Check if this HOP overlaps with any prior HOP's live arena outputs. + overlap_bufsizes: list[int] = [] + for prior_bufsizes, _prior_start, prior_end in hop_allocs: + if hop_start_idx <= prior_end: + _merge_bufsizes(overlap_bufsizes, prior_bufsizes) + + if parent_offset: + if overlap_bufsizes: + _merge_bufsizes(overlap_bufsizes, parent_offset) + else: + overlap_bufsizes = list(parent_offset) + + input_mem_buffer_sizes = overlap_bufsizes if overlap_bufsizes else None + + # Plan all branches of this HOP, taking the max (they're mutually exclusive) + this_hop_bufsizes: list[int] = [] + for arg_idx in submodule_arg_indices: + submodule_node = cast(torch.fx.Node, node.args[arg_idx]) + branch_bufsizes = _handle_submodule( + algo, + graph_module, + alignment, + submodule_node, + graph_signature, + alloc_graph_input=alloc_graph_input, + input_mem_buffer_sizes=input_mem_buffer_sizes, + ) + _merge_bufsizes(this_hop_bufsizes, branch_bufsizes) + + # Record this HOP's allocation for future overlap checks, but only + # if its outputs live in the submodule arena. + effective_bufsizes = list(overlap_bufsizes) if overlap_bufsizes else [] + _merge_bufsizes(effective_bufsizes, this_hop_bufsizes) + if outputs_in_arena: + hop_allocs.append((effective_bufsizes, hop_start_idx, hop_end_idx)) + + # Merge into the global submodule bufsizes (the max across all HOPs) + _merge_bufsizes(bufsizes, effective_bufsizes) # TODO: We can handle delegates the same way as map/cond/while. # Maybe populate the graph_module.meta["non_const_buffer_sizes"] for delegates. @@ -1563,10 +1610,13 @@ def apply_algo( # Only apply submodule pre-allocation for CPU specs; device buffers # do not share memory space with CPU submodule arenas. + effective_bufsizes = submodule_bufsizes if device_key == _CPU_KEY else [] + pre_existing = getattr(graph_module, "input_mem_buffer_sizes", None) + if pre_existing: + effective_bufsizes = list(effective_bufsizes) if effective_bufsizes else [] + _merge_bufsizes(effective_bufsizes, pre_existing) # pyre-ignore[16]: `torch.fx.GraphModule` has no attribute `input_mem_buffer_sizes`. - graph_module.input_mem_buffer_sizes = ( - submodule_bufsizes if device_key == _CPU_KEY else [] - ) + graph_module.input_mem_buffer_sizes = effective_bufsizes # Run algorithm independently on this device's specs device_bufsizes = algo( diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py index 31f3b1844c2..159c611fa62 100644 --- a/exir/tests/test_memory_planning.py +++ b/exir/tests/test_memory_planning.py @@ -1163,6 +1163,104 @@ def get_random_inputs(self) -> tuple[torch.Tensor, ...]: return self.map_model.get_random_inputs() +class OverlappingCondModel(torch.nn.Module): + """Two cond nodes where cond_0's outputs are still live when cond_1 executes.""" + + IV = torch.zeros(1, dtype=torch.float32) + INC1 = torch.tensor(1, dtype=torch.int32) + INC2 = torch.tensor(2, dtype=torch.int32) + + def __init__(self) -> None: + super().__init__() + self.register_buffer("aggr_v", torch.tensor(0, dtype=torch.int32)) + self.register_buffer("prev_v", self.IV.clone()) + self.register_buffer("curr_v", self.IV.clone()) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]: + aggr_v = self.aggr_v.clone() + v = x[0].unsqueeze(0) + is_valid = (v >= 0) & (x[1] >= 0) + + prev_v, curr_v = torch.cond( + is_valid, + lambda v: (self.curr_v.clone(), v.clone()), + lambda v: (self.prev_v.clone(), self.curr_v.clone()), + [v], + ) + self.prev_v.copy_(prev_v) + self.curr_v.copy_(curr_v) + + aggr_v += torch.cond( + is_valid, + lambda: self.INC1.clone(), + lambda: self.INC2.clone(), + ) + self.aggr_v.copy_(aggr_v) + + return self.aggr_v, self.prev_v, self.curr_v + + +class NonOverlappingCondModel(torch.nn.Module): + """Two cond nodes where cond_0's outputs are consumed before cond_1 executes.""" + + INC1 = torch.tensor(1, dtype=torch.int32) + INC2 = torch.tensor(2, dtype=torch.int32) + + def __init__(self) -> None: + super().__init__() + self.register_buffer("counter", torch.tensor(0, dtype=torch.int32)) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]: + is_positive = x[0] > 0 + + inc1 = torch.cond( + is_positive, + lambda: self.INC1.clone(), + lambda: self.INC2.clone(), + ) + self.counter.copy_(self.counter + inc1) + + inc2 = torch.cond( + is_positive, + lambda: self.INC2.clone(), + lambda: self.INC1.clone(), + ) + self.counter.copy_(self.counter + inc2) + + return (self.counter,) + + +class NestedCondModel(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, ...]: + pred = data >= 0.0 + + result1 = torch.cond( + pred, + lambda v: v.clone(), + lambda v: self.state.clone(), + [data], + ) + + result2 = torch.cond( + pred, + lambda r: torch.cond( + torch.isnan(self.state), + lambda r2: r2.clone(), + lambda r2: r2 + 1.0, + [r], + ), + lambda r: self.state.clone(), + [result1], + ) + + self.state.copy_(result1) + return result1, result2 + + class TestMap(unittest.TestCase): def test_map(self) -> None: @@ -1765,3 +1863,145 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: if node.op == "call_function" ) self.assertFalse(has_inplace) + + +class TestCondOverlap(unittest.TestCase): + """Tests for memory planning with multiple cond nodes whose outputs overlap.""" + + def test_overlapping_cond_outputs_no_corruption(self) -> None: + """Two cond nodes where cond_0's outputs are still live when cond_1 executes.""" + + module = OverlappingCondModel().eval() + INPUTS = [ + torch.tensor(d, dtype=torch.float32) + for d in [[5, 1], [6, 1], [7, -1], [8, 1]] + ] + # State: aggr_v=0, prev_v=[0.], curr_v=[0.] + # [5,1]: valid=T, prev_v←curr_v=[0.], curr_v←v=[5.], aggr_v=0+1=1 + # [6,1]: valid=T, prev_v←[5.], curr_v←[6.], aggr_v=1+1=2 + # [7,-1]: valid=F, prev_v←[5.], curr_v←[6.], aggr_v=2+2=4 + # [8,1]: valid=T, prev_v←[6.], curr_v←[8.], aggr_v=4+1=5 + EXPECTED = [ + ( + torch.tensor(1, dtype=torch.int32), + torch.tensor([0.0]), + torch.tensor([5.0]), + ), + ( + torch.tensor(2, dtype=torch.int32), + torch.tensor([5.0]), + torch.tensor([6.0]), + ), + ( + torch.tensor(4, dtype=torch.int32), + torch.tensor([5.0]), + torch.tensor([6.0]), + ), + ( + torch.tensor(5, dtype=torch.int32), + torch.tensor([6.0]), + torch.tensor([8.0]), + ), + ] + + forward_ep = export(module, (INPUTS[0],)) + edge = to_edge(forward_ep) + et = edge.to_executorch() + + import executorch.runtime + + rt_prog = executorch.runtime.Runtime.get().load_program(et.buffer) + rt_forward = rt_prog.load_method("forward") + + for inp, expected in zip(INPUTS, EXPECTED): + result = tuple(rt_forward.execute([inp])) + for r, e in zip(result, expected): + self.assertTrue( + torch.equal(r, e), + f"Mismatch for input {inp}: got {r}, expected {e}", + ) + + def test_non_overlapping_conds_share_memory(self) -> None: + """When cond outputs are consumed before the next cond, memory should be shared.""" + + module = NonOverlappingCondModel().eval() + inp = torch.tensor([5.0, 1.0]) + + forward_ep = export(module, (inp,)) + edge = to_edge(forward_ep) + + gm = edge.exported_program().graph_module + + mem_algo = MemoryPlanningAlgorithmSuite(algo_list=[greedy]) + gm = PassManager( + passes=[SpecPropPass(), ToOutVarPass()], + )(gm).graph_module + mem_planning_pass = MemoryPlanningPass( + mem_algo, + alloc_graph_input=True, + alloc_graph_output=True, + alloc_mutable_buffers=True, + ) + gm = mem_planning_pass.run(gm).graph_module + + cond_nodes = list( + gm.graph.find_nodes(op="call_function", target=torch.ops.higher_order.cond) + ) + self.assertEqual(len(cond_nodes), 2) + + sub1_true = getattr(gm, cond_nodes[0].args[1].target) + sub2_true = getattr(gm, cond_nodes[1].args[1].target) + bufsizes1 = sub1_true.meta["non_const_buffer_sizes"] + bufsizes2 = sub2_true.meta["non_const_buffer_sizes"] + + self.assertEqual(bufsizes1, bufsizes2) + + def test_nested_cond_no_corruption(self) -> None: + """Nested cond branches must not overwrite still-live outer cond outputs.""" + from executorch.exir.passes.init_mutable_pass import ( + InitializedMutableBufferPass, + ) + + INPUTS = [ + torch.tensor(5.0, dtype=torch.float64), + torch.tensor(6.0, dtype=torch.float64), + ] + # state starts as NaN. + # Call 1 (5.0): pred=True, result1=5.0, isnan(state)=True so result2=5.0, state←5.0 + # Call 2 (6.0): pred=True, result1=6.0, isnan(state)=False so result2=6.0+1.0=7.0 + EXPECTED = [ + ( + torch.tensor(5.0, dtype=torch.float64), + torch.tensor(5.0, dtype=torch.float64), + ), + ( + torch.tensor(6.0, dtype=torch.float64), + torch.tensor(7.0, dtype=torch.float64), + ), + ] + + module = NestedCondModel().eval() + forward_ep = export(module, (INPUTS[0],)) + edge = to_edge(forward_ep) + et = edge.to_executorch( + config=ExecutorchBackendConfig( + passes=[InitializedMutableBufferPass(list(module._buffers.keys()))], + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, + alloc_graph_output=True, + ), + ) + ) + + import executorch.runtime + + rt_prog = executorch.runtime.Runtime.get().load_program(et.buffer) + rt_forward = rt_prog.load_method("forward") + + for inp, expected in zip(INPUTS, EXPECTED): + result = tuple(rt_forward.execute([inp])) + for r, e in zip(result, expected): + self.assertTrue( + torch.equal(r, e), + f"Mismatch for input {inp}: got {r}, expected {e}", + )