diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py index d6273b01d90..e458ba93302 100644 --- a/exir/emit/test/test_emit.py +++ b/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: diff --git a/exir/memory_planning.py b/exir/memory_planning.py index 012cf8dd144..709c164eb0a 100644 --- a/exir/memory_planning.py +++ b/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. diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py index 31f3b1844c2..ec3982207e2 100644 --- a/exir/tests/test_memory_planning.py +++ b/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__()