Creates a memory timlines, where each step in the timeline is a list of active allocations at that timestep.
(graph: torch.fx.Graph)
| 49 | |
| 50 | |
| 51 | def create_tensor_allocation_info(graph: torch.fx.Graph) -> List[MemoryTimeline]: |
| 52 | """ |
| 53 | Creates a memory timlines, where each step in the timeline is a list of active |
| 54 | allocations at that timestep. |
| 55 | """ |
| 56 | nodes = graph.nodes |
| 57 | memory_timeline: List[Optional[MemoryTimeline]] = [None for _ in range(len(nodes))] |
| 58 | unique_specs: set[TensorSpec] = set() |
| 59 | for _, node in enumerate(nodes): |
| 60 | if node.op == "output": |
| 61 | continue |
| 62 | if node.target == memory.alloc or node.target == memory.view: |
| 63 | continue |
| 64 | tensor_specs = get_node_tensor_specs(node) |
| 65 | if tensor_specs is None: |
| 66 | continue |
| 67 | for tensor_spec in tensor_specs: |
| 68 | # TODO: Make use of mem_id in the allocation info |
| 69 | if tensor_spec is None or tensor_spec.mem_id is None or tensor_spec.const: |
| 70 | continue |
| 71 | if tensor_spec in unique_specs: |
| 72 | continue |
| 73 | unique_specs.add(tensor_spec) |
| 74 | start, end = tensor_spec.lifetime |
| 75 | size = num_bytes_from_shape_and_dtype( |
| 76 | typing.cast(torch.Size, tensor_spec.shape), tensor_spec.dtype |
| 77 | ) |
| 78 | stack_trace = node.meta.get("stack_trace") |
| 79 | fqn = _get_module_hierarchy(node) |
| 80 | for j in range(start, end + 1): |
| 81 | memory_timeline_j = memory_timeline[j] |
| 82 | if memory_timeline_j is None: |
| 83 | memory_timeline_j = MemoryTimeline() |
| 84 | memory_timeline[j] = memory_timeline_j |
| 85 | assert memory_timeline_j |
| 86 | memory_timeline_j.allocations.append( |
| 87 | Allocation( |
| 88 | node.name, |
| 89 | node.target, |
| 90 | tensor_spec.mem_id, |
| 91 | tensor_spec.mem_offset, |
| 92 | size, |
| 93 | fqn, |
| 94 | stack_trace, |
| 95 | ) |
| 96 | ) |
| 97 | return memory_timeline # type: ignore[return-value] |
| 98 | |
| 99 | |
| 100 | def _validate_memory_planning_is_done(exported_program: ExportedProgram): |
no test coverage detected