Generate the memory timeline from the given ExecuTorch program. Args: executorch_program The ExecuTorch program to be analyzed. Returns: Chrome trace in JSON format: Format: Each thread represents a unit of time. Thus to navigate timeline scroll up and do
(
executorch_program_manager: ExecutorchProgramManager,
chrome_trace_filename: str,
enable_memory_offsets: bool = False,
method_name: str = "forward",
ommit_metadata: bool = False,
)
| 109 | |
| 110 | |
| 111 | def generate_memory_trace( |
| 112 | executorch_program_manager: ExecutorchProgramManager, |
| 113 | chrome_trace_filename: str, |
| 114 | enable_memory_offsets: bool = False, |
| 115 | method_name: str = "forward", |
| 116 | ommit_metadata: bool = False, |
| 117 | ): |
| 118 | """ |
| 119 | Generate the memory timeline from the given ExecuTorch program. |
| 120 | Args: |
| 121 | executorch_program The ExecuTorch program to be analyzed. |
| 122 | Returns: |
| 123 | Chrome trace in JSON format: |
| 124 | Format: |
| 125 | Each thread represents a unit of time. Thus to navigate timeline scroll up and down. |
| 126 | For each thread, the x axis represents live tensor objects that are normalized according the allocation size. |
| 127 | """ |
| 128 | if not isinstance(executorch_program_manager, ExecutorchProgramManager): |
| 129 | raise ValueError( |
| 130 | f"generate_memory_trace expects ExecutorchProgramManager instance but got {type(executorch_program_manager)}" |
| 131 | ) |
| 132 | |
| 133 | exported_program = executorch_program_manager.exported_program(method_name) |
| 134 | if not _validate_memory_planning_is_done(exported_program): |
| 135 | raise ValueError("Executorch program does not have memory planning.") |
| 136 | |
| 137 | memory_timeline = create_tensor_allocation_info(exported_program.graph) |
| 138 | root = {} |
| 139 | trace_events: List[Dict[str, Any]] = [] |
| 140 | root["traceEvents"] = trace_events |
| 141 | |
| 142 | tid = 0 |
| 143 | for memory_timeline_event in memory_timeline: |
| 144 | start_time = 0 |
| 145 | if memory_timeline_event is None: |
| 146 | continue |
| 147 | for allocation in memory_timeline_event.allocations: |
| 148 | e: Dict[str, Any] = {} |
| 149 | e["name"] = allocation.name |
| 150 | e["cat"] = "memory_allocation" |
| 151 | e["ph"] = "X" |
| 152 | e["ts"] = ( |
| 153 | int(allocation.memory_offset) |
| 154 | if enable_memory_offsets |
| 155 | else int(start_time) |
| 156 | ) |
| 157 | allocation_size_kb = allocation.size_bytes |
| 158 | e["dur"] = int(allocation_size_kb) |
| 159 | e["pid"] = int(allocation.memory_id) |
| 160 | e["tid"] = tid |
| 161 | e["args"] = {} |
| 162 | if not ommit_metadata: |
| 163 | e["args"]["op_name"] = f"{allocation.op_name}" |
| 164 | # ID refers to memory space, typically from 1 to N. |
| 165 | # For CPU, everything is allocated on one "space", other backends may have multiple. |
| 166 | e["args"]["Memory ID"] = allocation.memory_id |
| 167 | e["args"]["fqn"] = f"{allocation.fqn}" |
| 168 | e["args"]["source"] = f"{allocation.file_and_line_num}" |
no test coverage detected