1# Copyright © 2024 Apple Inc. All rights reserved. 2# 3# Please refer to the license found in the LICENSE file in the root directory of the source tree. 4 5import copy 6import errno 7import json 8import os 9 10import subprocess 11from dataclasses import dataclass 12from pathlib import Path 13 14from typing import Any, Dict, Final, List, Optional, Tuple, Union 15 16import executorch.exir as exir 17 18import pandas as pd 19import torch 20from executorch.backends.apple.coreml.compiler import CoreMLBackend 21from executorch.backends.apple.coreml.partition import CoreMLPartitioner 22 23from executorch.devtools import BundledProgram, generate_etrecord, Inspector 24from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite 25from executorch.devtools.bundled_program.serialize import ( 26 serialize_from_bundled_program_to_flatbuffer, 27) 28from executorch.devtools.inspector import Event 29 30from executorch.exir import ( 31 EdgeProgramManager, 32 ExecutorchBackendConfig, 33 ExecutorchProgramManager, 34 ExirExportedProgram, 35 to_edge, 36) 37from executorch.exir.backend.compile_spec_schema import CompileSpec 38from executorch.exir.tracer import Value 39 40from torch.export import export, ExportedProgram 41 42COREML_METADATA_KEYS: Final[List[Tuple[str, str]]] = [ 43 ("operatorName", "coreml_operator"), 44 ("estimatedCost", "coreml_estimated_cost"), 45 ("preferredComputeUnit", "coreml_preferred_device"), 46 ("supportedComputeUnits", "coreml_supported_devices"), 47] 48 49 50def build_devtools_runner_including_coreml( 51 root_dir_path: Path, 52 conda_env_name: str, 53 force: bool = False, 54): 55 if not force: 56 devtools_executable_path = ( 57 root_dir_path / "cmake-out" / "examples" / "devtools" / "example_runner" 58 ) 59 print(devtools_executable_path) 60 if devtools_executable_path.is_file(): 61 return 62 63 cd_root_command: str = f"cd {root_dir_path.resolve()}" 64 conda_activate_env_command: str = f"source conda activate {conda_env_name}" 65 build_devtools_runner_command: str = ( 66 "./examples/devtools/build_example_runner.sh --coreml" 67 ) 68 build_command: str = ( 69 f"{cd_root_command} && {conda_activate_env_command} && {build_devtools_runner_command}" 70 ) 71 subprocess.run( 72 f'bash -c "{build_command}"', shell=True, check=True 73 ).check_returncode() 74 75 76_EDGE_COMPILE_CONFIG = exir.EdgeCompileConfig( 77 _check_ir_validity=False, 78 _skip_dim_order=True, 79) 80 81_EDGE_BACKEND_CONFIG = exir.ExecutorchBackendConfig( 82 extract_delegate_segments=True, 83) 84 85 86def to_core_aten( 87 module: torch.nn.Module, 88 example_inputs: Tuple[Value, ...], 89) -> ExportedProgram: 90 core_aten_program = export( 91 mod=module, 92 args=example_inputs, 93 ) 94 return core_aten_program 95 96 97def core_aten_to_edge( 98 core_aten_program: ExportedProgram, 99 edge_compile_config: exir.EdgeCompileConfig, 100) -> EdgeProgramManager: 101 edge_manager = to_edge( 102 programs=core_aten_program, 103 compile_config=edge_compile_config, 104 ) 105 return edge_manager 106 107 108def module_to_edge( 109 module: torch.nn.Module, 110 example_inputs: Tuple[Value, ...], 111 edge_compile_config: exir.EdgeCompileConfig = _EDGE_COMPILE_CONFIG, 112) -> EdgeProgramManager: 113 module.eval() 114 core_aten_program = to_core_aten( 115 module=module, 116 example_inputs=example_inputs, 117 ) 118 return core_aten_to_edge( 119 core_aten_program=core_aten_program, 120 edge_compile_config=edge_compile_config, 121 ) 122 123 124def lower_and_export_edge_to_coreml( 125 edge_program: EdgeProgramManager, 126 compile_specs: List[CompileSpec], 127 config: ExecutorchBackendConfig, 128 skip_ops_for_coreml_delegation: Optional[List[str]] = None, 129) -> ExirExportedProgram: 130 partitioner = CoreMLPartitioner( 131 skip_ops_for_coreml_delegation=skip_ops_for_coreml_delegation, 132 compile_specs=compile_specs, 133 ) 134 delegated_program_manager = edge_program.to_backend( 135 partitioner, 136 ) 137 executorch_program = delegated_program_manager.to_executorch( 138 config=config, 139 ) 140 return executorch_program 141 142 143def write_to_file(buffer: bytes, file_path: Path): 144 with open(file_path.resolve(), "wb") as file: 145 file.write(buffer) 146 147 148def generate_bundled_program( 149 executorch_program: ExecutorchProgramManager, 150 example_inputs: Tuple[Value, ...], 151 method_name: str, 152 bundled_program_path: Path, 153): 154 method_test_suites = [ 155 MethodTestSuite( 156 method_name=method_name, 157 test_cases=[MethodTestCase(inputs=example_inputs)], 158 ) 159 ] 160 161 bundled_program = BundledProgram(executorch_program, method_test_suites) 162 bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer( 163 bundled_program 164 ) 165 166 write_to_file(buffer=bundled_program_buffer, file_path=bundled_program_path) 167 168 169def generate_etdump_with_intermediate_values( 170 root_dir_path: Path, 171 bundled_program_path: Path, 172 et_dump_path: Path, 173 debug_buffer_path: Path, 174 debug_buffer_size: int, 175): 176 devtools_executable_path = ( 177 root_dir_path / "cmake-out" / "examples" / "devtools" / "example_runner" 178 ) 179 if not devtools_executable_path.is_file(): 180 raise FileNotFoundError( 181 errno.ENOENT, 182 os.strerror(errno.ENOENT), 183 str(devtools_executable_path.resolve()), 184 ) 185 186 devtools_runner_command: str = f""" 187 {devtools_executable_path.resolve()} -dump_intermediate_outputs\ 188 -bundled_program_path {bundled_program_path.resolve()}\ 189 -etdump_path {et_dump_path.resolve()}\ 190 -debug_output_path {debug_buffer_path.resolve()}\ 191 -debug_buffer_size {debug_buffer_size}""" 192 subprocess.run( 193 f'bash -c "{devtools_runner_command}"', shell=True, check=True 194 ).check_returncode() 195 196 197def create_inspector( 198 edge_program: EdgeProgramManager, 199 executorch_program: ExecutorchProgramManager, 200 example_inputs: Tuple[Value, ...], 201 model_name: str, 202 root_dir_path: Path, 203 working_dir_path: Path, 204 method_name: str = "forward", 205 debug_buffer_size: int = 1 * 1024 * 1024 * 1024, 206 delegate_metadata_parser=None, 207 delegate_time_scale_converter=None, 208) -> Inspector: 209 et_record_path = working_dir_path / f"{model_name}_etrecord.bin" 210 generate_etrecord( 211 et_record=et_record_path.resolve(), 212 edge_dialect_program=edge_program, 213 executorch_program=executorch_program, 214 ) 215 216 bundled_program_path = working_dir_path / f"{model_name}.bpte" 217 generate_bundled_program( 218 executorch_program=executorch_program, 219 example_inputs=example_inputs, 220 method_name=method_name, 221 bundled_program_path=bundled_program_path, 222 ) 223 224 et_dump_path: Path = working_dir_path / f"{model_name}_etdump.etdp" 225 debug_buffer_path: Path = working_dir_path / f"{model_name}_debug_output.bin" 226 generate_etdump_with_intermediate_values( 227 root_dir_path=root_dir_path, 228 bundled_program_path=bundled_program_path, 229 et_dump_path=et_dump_path, 230 debug_buffer_path=debug_buffer_path, 231 debug_buffer_size=debug_buffer_size, 232 ) 233 234 return Inspector( 235 etdump_path=str(et_dump_path.resolve()), 236 etrecord=str(et_record_path.resolve()), 237 debug_buffer_path=str(debug_buffer_path.resolve()), 238 enable_module_hierarchy=True, 239 delegate_metadata_parser=delegate_metadata_parser, 240 delegate_time_scale_converter=delegate_time_scale_converter, 241 ) 242 243 244def parse_coreml_delegate_metadata(delegate_metadatas: List[str]) -> Dict[str, Any]: 245 if len(delegate_metadatas) == 0: 246 return 247 try: 248 coreml_metadata: Dict[str, Any] = json.loads(delegate_metadatas[0]) 249 result: Dict[str, str] = {} 250 for col_key, col_name in COREML_METADATA_KEYS: 251 value = coreml_metadata.get(col_key, None) 252 if value is not None: 253 result[col_name] = value 254 return result 255 256 except ValueError: 257 return {} 258 259 260def convert_coreml_delegate_time( 261 event_name: Union[str, int], input_time: Union[int, float] 262) -> Union[int, float]: 263 return input_time / (1000 * 1000) 264 265 266def create_inspector_coreml( 267 edge_program: EdgeProgramManager, 268 compile_specs: List[CompileSpec], 269 example_inputs: Tuple[Value, ...], 270 model_name: str, 271 root_dir_path: Path, 272 working_dir_path: Path, 273 method_name: str = "forward", 274 debug_buffer_size: int = 1 * 1024 * 1024 * 1024, 275) -> Inspector: 276 edge_program_copy = copy.deepcopy(edge_program) 277 executorch_program = lower_and_export_edge_to_coreml( 278 edge_program=edge_program_copy, 279 compile_specs=compile_specs, 280 config=_EDGE_BACKEND_CONFIG, 281 ) 282 return create_inspector( 283 edge_program=edge_program, 284 executorch_program=executorch_program, 285 example_inputs=example_inputs, 286 root_dir_path=root_dir_path, 287 model_name=f"{model_name}_coreml", 288 working_dir_path=working_dir_path, 289 method_name=method_name, 290 debug_buffer_size=debug_buffer_size, 291 delegate_metadata_parser=parse_coreml_delegate_metadata, 292 delegate_time_scale_converter=convert_coreml_delegate_time, 293 ) 294 295 296def create_inspector_reference( 297 edge_program: EdgeProgramManager, 298 example_inputs: Tuple[Value, ...], 299 model_name: str, 300 root_dir_path: Path, 301 working_dir_path: Path, 302 method_name: str = "forward", 303 debug_buffer_size: int = 1 * 1024 * 1024 * 1024, 304) -> Inspector: 305 edge_program_copy = copy.deepcopy(edge_program) 306 return create_inspector( 307 edge_program=edge_program, 308 executorch_program=edge_program_copy.to_executorch(), 309 example_inputs=example_inputs, 310 root_dir_path=root_dir_path, 311 model_name=f"{model_name}_default", 312 working_dir_path=working_dir_path, 313 method_name=method_name, 314 debug_buffer_size=debug_buffer_size, 315 ) 316 317 318def get_debug_handle_to_event_map( 319 inspector: Inspector, 320 event_block_name: str = "Execute", 321) -> Dict[int, Event]: 322 result = {} 323 324 def is_not_blank(s): 325 return bool(s and not s.isspace()) 326 327 event_names_to_ignore = {"DELEGATE_CALL", "OPERATOR_CALL"} 328 for event_block in inspector.event_blocks: 329 if event_block.name == event_block_name: 330 for event in event_block.events: 331 if is_not_blank(event.name) and event.name not in event_names_to_ignore: 332 debug_handles = [] 333 if isinstance(event.debug_handles, int): 334 debug_handles.append(event.debug_handles) 335 elif isinstance(event.debug_handles, list): 336 debug_handles.extend(event.debug_handles) 337 debug_handles.sort() 338 for debug_handle in debug_handles: 339 if len(event.debug_data) > 0: 340 result[debug_handle] = event 341 return result 342 343 344@dataclass 345class EventData: 346 tag: str 347 event: Event 348 349 350@dataclass 351class ComparisonResult: 352 datas: List[tuple[EventData, EventData]] 353 354 def to_dataframe( 355 self, 356 atol: float = 1e-3, 357 rtol: float = 1e-3, 358 ) -> pd.DataFrame: 359 def get_compute_device(event: Event) -> str: 360 if event.delegate_backend_name == CoreMLBackend.__name__: 361 return event.delegate_debug_metadatas.get( 362 "coreml_preferred_device", "CPU" 363 ) 364 365 return "CPU" 366 367 if len(self.datas) == 0: 368 return 369 370 (data1, data2) = self.datas[0] 371 dict = { 372 data1.tag: [], 373 f"{data1.tag}_compute_unit": [], 374 data2.tag: [], 375 f"{data2.tag}_compute_unit": [], 376 "max_diff": [], 377 } 378 379 for data1, data2 in self.datas: 380 event1 = data1.event 381 event2 = data2.event 382 debug_data1 = event1.debug_data[0] 383 debug_data2 = event2.debug_data[0] 384 385 if debug_data1.size() != debug_data2.size(): 386 continue 387 388 max_diff = 0.0 389 indices = torch.isclose( 390 debug_data1, debug_data2, atol=atol, rtol=rtol 391 ).logical_not() 392 393 # Find the maximum difference 394 if torch.count_nonzero(indices) > 0: 395 values1 = torch.masked_select(debug_data1, indices) 396 values2 = torch.masked_select(debug_data2, indices) 397 diff = torch.abs(values1 - values2) 398 max_diff = torch.max(diff).item() 399 400 dict[f"{data1.tag}_compute_unit"].append(get_compute_device(event1)) 401 dict[f"{data2.tag}_compute_unit"].append(get_compute_device(event2)) 402 dict["max_diff"].append(max_diff) 403 dict[data1.tag].append(event1.name) 404 dict[data2.tag].append(event2.name) 405 406 return pd.DataFrame(dict) 407 408 409def get_comparison_result( 410 inspector1: Inspector, 411 tag1: str, 412 inspector2: Inspector, 413 tag2: str, 414) -> ComparisonResult: 415 debug_handle_event_map_1 = get_debug_handle_to_event_map(inspector1) 416 debug_handle_event_map_2 = get_debug_handle_to_event_map(inspector2) 417 418 event_datas = [] 419 for handle, event1 in debug_handle_event_map_1.items(): 420 event2 = debug_handle_event_map_2.get(handle, None) 421 if event2 is None: 422 continue 423 424 event_data1 = EventData(tag=tag1, event=event1) 425 event_data2 = EventData(tag=tag2, event=event2) 426 event_datas.append((event_data1, event_data2)) 427 428 return ComparisonResult(datas=event_datas) 429