Speedup toposort

Author: ricardoV94

pytensor/graph/basic.py

@@ -1522,13 +1522,38 @@ def _compute_deps_cache_(io):
     return rlist
 
 
+def node_toposort(
+    output_nodes: Iterable[Apply] = (),
+    input_nodes: Iterable[Apply] = (),
+) -> Generator[Apply]:
+    computed = {None, *input_nodes}
+    todo = list(output_nodes)
+    while todo:
+        cur = todo[-1]
+        if cur in computed:
+            todo.pop()
+            continue
+        # Since computed includes None we don't need to filter it in this check
+        if all(i.owner in computed for i in cur.inputs):
+            computed.add(cur)
+            yield todo.pop()
+        else:
+            todo.extend(i.owner for i in cur.inputs if i.owner is not None)
+
+
 def io_toposort(
     inputs: Iterable[Variable],
     outputs: Reversible[Variable],
     orderings: dict[Apply, list[Apply]] | None = None,
     clients: dict[Variable, list[Variable]] | None = None,
 ) -> list[Apply]:
-    """Perform topological sort from input and output nodes.
+    """Perform topological of nodes between input and output variables.
+
+    Notes
+    -----
+    If sorting from root or single-output node variables, without orderings or clients,
+    it's slightly faster to use `list(node_toposort((o.owner for o in outputs)))` instead,
+    as the individual variables can be ignored
 
     Parameters
     ----------
@@ -1543,32 +1568,35 @@ def io_toposort(
         each node in the subgraph that is sorted.
 
     """
-    if not orderings and clients is None:  # ordering can be None or empty dict
-        # Specialized function that is faster when more then ~10 nodes
-        # when no ordering.
-
-        # Do a new stack implementation with the vm algo.
-        # This will change the order returned.
+    if not orderings and clients is None:
+        # Specialized function that is faster when no special orderings are required and clients need not be computed.
         computed = set(inputs)
-        todo = [o.owner for o in reversed(outputs) if o.owner]
+        todo = [o.owner for o in outputs if o.owner is not None]
         order = []
         while todo:
-            cur = todo.pop()
-            if all(out in computed for out in cur.outputs):
+            cur = todo[-1]
+            # It's faster to short circuit on the first output, as most nodes will have all edges non-computed
+            # Starting the `all` iterator has a non-negligeable cost
+            if cur.outputs[0] in computed and all(
+                out in computed for out in cur.outputs[1:]
+            ):
+                todo.pop()
                 continue
             if all(i in computed or i.owner is None for i in cur.inputs):
                 computed.update(cur.outputs)
-                order.append(cur)
+                order.append(todo.pop())
             else:
-                todo.append(cur)
                 todo.extend(
-                    i.owner for i in cur.inputs if (i.owner and i not in computed)
+                    i.owner
+                    for i in cur.inputs
+                    if ((i.owner is not None) and (i not in computed))
                 )
         return order
 
     iset = set(inputs)
 
-    if not orderings:  # ordering can be None or empty dict
+    if not orderings:
+        # ordering can be None or empty dict
         # Specialized function that is faster when no ordering.
         # Also include the cache in the function itself for speed up.
 

pytensor/graph/fg.py

@@ -15,13 +15,13 @@
     clone_get_equiv,
     graph_inputs,
     io_toposort,
+    node_toposort,
     vars_between,
 )
 from pytensor.graph.basic import as_string as graph_as_string
 from pytensor.graph.features import AlreadyThere, Feature, ReplaceValidate
 from pytensor.graph.op import Op
 from pytensor.graph.utils import MetaObject, MissingInputError, TestValueError
-from pytensor.misc.ordered_set import OrderedSet
 
 
 ClientType = tuple[Apply, int]
@@ -353,21 +353,16 @@ def import_node(
         apply_node : Apply
             The node to be imported.
         check : bool
-            Check that the inputs for the imported nodes are also present in
-            the `FunctionGraph`.
+            Check that the inputs for the imported nodes are also present in the `FunctionGraph`.
         reason : str
             The name of the optimization or operation in progress.
         import_missing : bool
             Add missing inputs instead of raising an exception.
         """
         # We import the nodes in topological order. We only are interested in
-        # new nodes, so we use all variables we know of as if they were the
-        # input set.  (The functions in the graph module only use the input set
-        # to know where to stop going down.)
-        new_nodes = io_toposort(self.variables, apply_node.outputs)
-
-        if check:
-            for node in new_nodes:
+        # new nodes, so we use all nodes we know of as inputs to interrupt the toposort
+        for node in node_toposort([apply_node], node_inputs=self.apply_nodes):
+            if check:
                 for var in node.inputs:
                     if (
                         var.owner is None
@@ -387,8 +382,6 @@ def import_node(
                             )
                             raise MissingInputError(error_msg, variable=var)
 
-        for node in new_nodes:
-            assert node not in self.apply_nodes
             self.apply_nodes.add(node)
             if not hasattr(node.tag, "imported_by"):
                 node.tag.imported_by = []
@@ -753,11 +746,12 @@ def toposort(self) -> list[Apply]:
           :meth:`FunctionGraph.orderings`.
 
         """
-        if len(self.apply_nodes) < 2:
-            # No sorting is necessary
-            return list(self.apply_nodes)
-
-        return io_toposort(self.inputs, self.outputs, self.orderings())
+        orderings = self.orderings()
+        if orderings:
+            return io_toposort(self.inputs, self.outputs, orderings)
+        else:
+            # Faster implementation when no orderings are needed
+            return node_toposort(o.owner for o in self.outputs)
 
     def orderings(self) -> dict[Apply, list[Apply]]:
         """Return a map of node to node evaluation dependencies.
@@ -776,28 +770,14 @@ def orderings(self) -> dict[Apply, list[Apply]]:
         take care of computing the dependencies by itself.
 
         """
-        assert isinstance(self._features, list)
         all_orderings: list[dict] = []
 
         for feature in self._features:
             if hasattr(feature, "orderings"):
                 orderings = feature.orderings(self)
-                if not isinstance(orderings, dict):
-                    raise TypeError(
-                        "Non-deterministic return value from "
-                        + str(feature.orderings)
-                        + ". Nondeterministic object is "
-                        + str(orderings)
-                    )
-                if len(orderings) > 0:
+                if orderings:
                     all_orderings.append(orderings)
-                    for node, prereqs in orderings.items():
-                        if not isinstance(prereqs, list | OrderedSet):
-                            raise TypeError(
-                                "prereqs must be a type with a "
-                                "deterministic iteration order, or toposort "
-                                " will be non-deterministic."
-                            )
+
         if len(all_orderings) == 1:
             # If there is only 1 ordering, we reuse it directly.
             return all_orderings[0].copy()

pytensor/graph/rewriting/basic.py

@@ -25,7 +25,7 @@
     Constant,
     Variable,
     applys_between,
-    io_toposort,
+    node_toposort,
     vars_between,
 )
 from pytensor.graph.features import AlreadyThere, Feature, NodeFinder
@@ -2027,7 +2027,7 @@ def apply(self, fgraph, start_from=None):
         callback_before = fgraph.execute_callbacks_time
         nb_nodes_start = len(fgraph.apply_nodes)
         t0 = time.perf_counter()
-        q = deque(io_toposort(fgraph.inputs, start_from))
+        q = deque(node_toposort(output_nodes=(o.owner for o in start_from)))
         io_t = time.perf_counter() - t0
 
         def importer(node):
@@ -2320,11 +2320,6 @@ def add_requirements(self, fgraph):
     def apply(self, fgraph, start_from=None):
         change_tracker = ChangeTracker()
         fgraph.attach_feature(change_tracker)
-        if start_from is None:
-            start_from = fgraph.outputs
-        else:
-            for node in start_from:
-                assert node in fgraph.outputs
 
         changed = True
         max_use_abort = False
@@ -2403,7 +2398,7 @@ def apply_cleanup(profs_dict):
             changed |= apply_cleanup(iter_cleanup_sub_profs)
 
             topo_t0 = time.perf_counter()
-            q = deque(io_toposort(fgraph.inputs, start_from))
+            q = deque(node_toposort(o.owner for o in fgraph.outputs))
             io_toposort_timing.append(time.perf_counter() - topo_t0)
 
             nb_nodes.append(len(q))

pytensor/scan/rewriting.py

@@ -23,6 +23,7 @@
     equal_computations,
     graph_inputs,
     io_toposort,
+    node_toposort,
 )
 from pytensor.graph.destroyhandler import DestroyHandler
 from pytensor.graph.features import ReplaceValidate
@@ -225,7 +226,7 @@ def scan_push_out_non_seq(fgraph, node):
 
     node_inputs, node_outputs = node.op.inner_inputs, node.op.inner_outputs
 
-    local_fgraph_topo = io_toposort(node_inputs, node_outputs)
+    local_fgraph_topo = node_toposort(node_outputs)
     local_fgraph_outs_set = set(node_outputs)
     local_fgraph_outs_map = {v: k for k, v in enumerate(node_outputs)}
 
@@ -435,7 +436,7 @@ def scan_push_out_seq(fgraph, node):
 
     node_inputs, node_outputs = node.op.inner_inputs, node.op.inner_outputs
 
-    local_fgraph_topo = io_toposort(node_inputs, node_outputs)
+    local_fgraph_topo = node_toposort(node_outputs)
     local_fgraph_outs_set = set(node_outputs)
     local_fgraph_outs_map = {v: k for k, v in enumerate(node_outputs)}
 

pytensor/tensor/rewriting/blas.py

@@ -59,6 +59,7 @@
 
 import numpy as np
 
+from pytensor.graph.basic import node_toposort
 from pytensor.tensor.rewriting.basic import register_specialize
 
 
@@ -459,6 +460,8 @@ def apply(self, fgraph):
             callbacks_before = fgraph.execute_callbacks_times.copy()
             callback_before = fgraph.execute_callbacks_time
 
+        nodelist = node_toposort(o.owner for o in fgraph.outputs)
+
         def on_import(new_node):
             if new_node is not node:
                 nodelist.append(new_node)
@@ -470,7 +473,6 @@ def on_import(new_node):
         while did_something:
             nb_iter += 1
             t0 = time.perf_counter()
-            nodelist = pytensor.graph.basic.io_toposort(fgraph.inputs, fgraph.outputs)
             time_toposort += time.perf_counter() - t0
             did_something = False
             nodelist.reverse()

tests/graph/test_basic.py

@@ -23,6 +23,7 @@
     get_var_by_name,
     graph_inputs,
     io_toposort,
+    node_toposort,
     orphans_between,
     truncated_graph_inputs,
     variable_depends_on,
@@ -329,6 +330,23 @@ def test_multi_output_nodes(self):
             out.owner,
         }
 
+    @pytest.mark.parametrize(
+        "toposort_func",
+        [
+            lambda x: io_toposort([], [x]),
+            lambda x: list(node_toposort([x.owner])),
+        ],
+        ids=["io_toposort", "node_toposort"],
+    )
+    def test_benchmark(self, toposort_func, benchmark):
+        r1 = MyVariable(1)
+        out = r1
+        for i in range(50):
+            out = MyOp(out, out)
+
+        assert toposort_func(out) == io_toposort([r1], [out])
+        benchmark(toposort_func, out)
+
 
 class TestEval:
     def setup_method(self):