Class PassManager

exir/pass_manager.py:21–78

Class to run multiple passes on a given graph module. The PassManager is callable so to run it, we can just call the PassManager instance. Private Attributes: * **passes**: A list of callable passes * **params**: An instance of PassManagerParams containing the result of

class PassManager(fx.PassManager):
    """
    Class to run multiple passes on a given graph module. The PassManager is
    callable so to run it, we can just call the PassManager instance.

    Private Attributes:
        * **passes**: A list of callable passes
        * **params**: An instance of PassManagerParams containing the result of the
            flags set in the constructor.
    """

    def __init__(
        self,
        passes: Optional[Union[List[PassType], List[List[PassType]]]] = None,
        run_checks_after_each_pass: bool = False,
        suppress_check_failures: bool = False,
    ) -> None:
        r"""
        Args:
            passes: A list of passes
            enable_debug_pass: set to true to enable the debug passes
            run_checks_after_each_pass: whether to run checks and linting after each pass
        """

        # Flatten the passes to a list of callables
        passes = passes if passes else []
        flattened_passes = [
            fx.pass_result_wrapper(fn) for fn in pytree.tree_flatten(passes)[0]
        ]

        super().__init__(
            flattened_passes,
            run_checks_after_each_pass=run_checks_after_each_pass,
            suppress_check_failures=suppress_check_failures,
        )

    def check(self, module: torch.nn.Module) -> None:
        """
        Runs various checks on the given graph module to make sure it contains
        the needed data for passes.

        Some checks that need to be run:
            - Ensure that types of operator node match the types specified in
              the node's spec field (ex. if the op returns a tuple then the
              node's spec field is a tuple)
            - Ensure that the graph module has type torch.fx.GraphModule
        """
        assert isinstance(module, fx.GraphModule)
        module.recompile()
        module.graph.lint()
        # TODO(qihan): use verifier.check_is_exir

        for node in module.graph.nodes:
            if node.op == "call_method":
                raise ExportError(
                    ExportErrorType.NOT_SUPPORTED,
                    f"call_method `{node}` is not supported except for backend delegate.",
                )