Transforms the program according to the provided passes. Args: passes: This param can be one of: 1) a list of passes - all methods in the given EdgeProgramManager will be transformed with the provided passes.
(
self,
passes: Union[Sequence[PassType], Dict[str, Sequence[PassType]], PassManager],
compile_config: Optional[EdgeCompileConfig] = None,
)
| 1603 | |
| 1604 | @et_logger("transform") |
| 1605 | def transform( |
| 1606 | self, |
| 1607 | passes: Union[Sequence[PassType], Dict[str, Sequence[PassType]], PassManager], |
| 1608 | compile_config: Optional[EdgeCompileConfig] = None, |
| 1609 | ) -> "EdgeProgramManager": |
| 1610 | """ |
| 1611 | Transforms the program according to the provided passes. |
| 1612 | |
| 1613 | Args: |
| 1614 | passes: This param can be one of: |
| 1615 | 1) a list of passes - |
| 1616 | all methods in the given EdgeProgramManager |
| 1617 | will be transformed with the provided passes. |
| 1618 | 2) a dictionary mapping method names to lists of passes - |
| 1619 | only method names specified in the dictionary will be |
| 1620 | transformed with their corresponding passes. |
| 1621 | 3) a PassManager instance - |
| 1622 | all methods in the given EdgeProgramManager will be |
| 1623 | transformed with the given PassManager instance. |
| 1624 | compile_config: Compile config to use for veriy the correctness of model |
| 1625 | graph after each pass. If not specified, the compile config of the |
| 1626 | calling EdgeProgramManager will be used. It will be used in as compile |
| 1627 | config of returned EdgeProgramManager. |
| 1628 | |
| 1629 | Returns: |
| 1630 | EdgeProgramManager: A copy of the calling EdgeProgramManager with the |
| 1631 | transformations applied. |
| 1632 | """ |
| 1633 | |
| 1634 | compile_config = compile_config or self.compile_config |
| 1635 | new_programs: Dict[str, ExportedProgram] = {} |
| 1636 | |
| 1637 | # Cast passes parameter upfront. |
| 1638 | passes_seq: Optional[Sequence[PassType]] = None |
| 1639 | passes_dict: Optional[Dict[str, Sequence[PassType]]] = None |
| 1640 | pass_manager: Optional[PassManager] = None |
| 1641 | |
| 1642 | if isinstance(passes, Sequence): |
| 1643 | passes_seq = passes |
| 1644 | if isinstance(passes, dict): |
| 1645 | passes_dict = passes |
| 1646 | if isinstance(passes, PassManager): |
| 1647 | pass_manager = passes |
| 1648 | |
| 1649 | for name, program in self._edge_programs.items(): |
| 1650 | # If the method name is enforced, but not matched, we skip transformation. |
| 1651 | if ( |
| 1652 | isinstance(passes, dict) |
| 1653 | and passes_dict |
| 1654 | and name not in passes_dict.keys() |
| 1655 | ): |
| 1656 | new_programs[name] = copy.deepcopy(program) |
| 1657 | continue |
| 1658 | |
| 1659 | # Depending on the passes parameter, call the corresponding transform function. |
| 1660 | if passes_seq is not None: |
| 1661 | new_programs[name] = _transform(program, *passes_seq) |
| 1662 | elif passes_dict is not None: |