Class FusePatternsPass

backends/vulkan/_passes/fuse_patterns.py:17–34

class FusePatternsPass(ExportPass):
    def __init__(self) -> None:
        super().__init__()
        self._exported_program: Optional[ExportedProgram] = None

    def call(self, graph_module: torch.fx.GraphModule):
        assert self._exported_program is not None

        total_replaced = vk_patterns.replace_all_fusable_subgraphs(
            self._exported_program, graph_module
        )

        if total_replaced > 0:
            graph_module.recompile()
            # Re-trace the graph
            graph_module = super().call(graph_module).graph_module

        return PassResult(graph_module, total_replaced > 0)