TODO: Deprecated capture_program with to_edge_transform_and_lower_to_qnn Captures and transforms a PyTorch module into an Exir exported program. Args: module (Union[torch.nn.Module, torch.fx.GraphModule]): The PyTorch module or fx.GraphModule to be captured. inputs (Tu
(
module: Union[torch.nn.Module, torch.fx.GraphModule],
inputs: Tuple[torch.Tensor],
dep_table: Optional[Dict] = None,
passes_job: OrderedDict = None,
dynamic_shapes: Dict = None,
)
| 479 | |
| 480 | |
| 481 | def capture_program( |
| 482 | module: Union[torch.nn.Module, torch.fx.GraphModule], |
| 483 | inputs: Tuple[torch.Tensor], |
| 484 | dep_table: Optional[Dict] = None, |
| 485 | passes_job: OrderedDict = None, |
| 486 | dynamic_shapes: Dict = None, |
| 487 | ) -> exir.ExirExportedProgram: |
| 488 | """ |
| 489 | TODO: Deprecated capture_program with to_edge_transform_and_lower_to_qnn |
| 490 | |
| 491 | Captures and transforms a PyTorch module into an Exir exported program. |
| 492 | |
| 493 | Args: |
| 494 | module (Union[torch.nn.Module, torch.fx.GraphModule]): The PyTorch module or fx.GraphModule to be captured. |
| 495 | inputs (Tuple[torch.Tensor]): The input tensors for the module. |
| 496 | dep_table (Optional[Dict]): Dependency table for the transformation passes. |
| 497 | passes_job (OrderedDict, optional): Ordered dictionary of transformation passes. |
| 498 | dynamic_shapes (Dict, optional): Information about dynamic shapes. |
| 499 | |
| 500 | Returns: |
| 501 | exir.ExirExportedProgram: The transformed Exir exported program ready for lowering to QNN backend. |
| 502 | """ |
| 503 | warnings.warn( |
| 504 | "capture_program is deprecated. Use to_edge_transform_and_lower_to_qnn instead.", |
| 505 | DeprecationWarning, |
| 506 | stacklevel=1, |
| 507 | ) |
| 508 | ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes, strict=True) |
| 509 | ep = QnnPassManager().transform_for_export_pipeline(ep) |
| 510 | # TODO: Handle stack op. If we want to run annotate_decomposed pass for stack op, |
| 511 | # we need to make stack op decompose, which means we need to find a method to |
| 512 | # remove it from skip_decomp table |
| 513 | decomposed_ep = ep.run_decompositions(get_decomp_table(passes_job)) |
| 514 | core_ep = ExirExportedProgram(decomposed_ep, False) |
| 515 | edge_ep = core_ep.to_edge(qnn_edge_config()) |
| 516 | transform_passes = QnnPassManager().get_to_edge_transform_passes( |
| 517 | edge_ep.exported_program, |
| 518 | passes_job=passes_job, |
| 519 | dep_table=dep_table, |
| 520 | ) |
| 521 | edge_ep.transform(*transform_passes) |
| 522 | return edge_ep |
| 523 | |
| 524 | |
| 525 | def _partition_graph_into_submodules(gm, subgm_tag, subgm_cb, ptn): |