refactor: remove topsort because expr chain is already sorted as it's built

Author: polvalente

nx/lib/nx/defn/graph_splitter.ex

@@ -138,10 +138,7 @@ defmodule Nx.Defn.GraphSplitter do
         end
       )
 
-    # Apply topological sort to the expr_chain
-    sorted_expr_chain = topological_sort(expr_chain)
-
-    {sorted_expr_chain, cache, Map.delete(state, :expression_chain)}
+    {expr_chain, cache, Map.delete(state, :expression_chain)}
   end
 
   defp composite_eval(expr, state, cache) do
@@ -289,42 +286,4 @@ defmodule Nx.Defn.GraphSplitter do
   end
 
   defp rewrite_subtree(other, _, acc), do: {other, acc}
-
-  defp topological_sort(expr_chain) do
-    # Create a new directed graph
-    graph = :digraph.new()
-
-    # Add vertices for each stage output
-    Enum.each(expr_chain, fn %Stage{id: id, arguments: arguments} ->
-      Enum.with_index(arguments, fn _, idx ->
-        :digraph.add_vertex(graph, {id, idx})
-      end)
-    end)
-
-    # Add edges based on argument sources
-    Enum.each(expr_chain, fn %Stage{id: id, argument_sources: sources} ->
-      Enum.each(sources, fn {_arg_id, {source_stage_id, source_index}} ->
-        if source_stage_id != nil do
-          Enum.each(0..(map_size(sources) - 1)//1, fn idx ->
-            :digraph.add_edge(graph, {source_stage_id, source_index}, {id, idx})
-          end)
-        end
-      end)
-    end)
-
-    # Perform topological sort
-    sorted_ids =
-      graph
-      |> :digraph_utils.topsort()
-      |> Enum.map(fn {stage_id, _} -> stage_id end)
-      |> Enum.uniq()
-
-    # Clean up the graph
-    :digraph.delete(graph)
-
-    # Return the sorted stages based on sorted_ids
-    Enum.map(sorted_ids, fn stage_id ->
-      Enum.find(expr_chain, fn %Stage{id: id} -> id == stage_id end)
-    end)
-  end
 end

nx/test/nx/defn/graph_splitter_test.exs

@@ -426,18 +426,12 @@ defmodule Nx.Defn.GraphSplitterTest do
 
       chain = GraphSplitter.traverse(expr, split_fn)
 
-      assert [root, side1, side2, merge] = chain
+      assert [root, right, left, merge] = chain
 
       assert {%T{data: %Expr{op: :multiply, args: [arg0, arg1]}}} = root.expr
       assert %T{data: %Expr{op: :parameter, args: [0]}} = arg0
       assert %T{data: %Expr{op: :parameter, args: [1]}} = arg1
 
-      # because things are balanced, we don't know which of side1 and side2 are left and right
-      # in our expr, so we should disambiguate:
-
-      {[%Stage{} = left], [%Stage{} = right]} =
-        Enum.split_with([side1, side2], fn %Stage{expr: {expr}} -> expr.data.op == :multiply end)
-
       # left should depend on exactly the same parameters as the root, as it's pulling from
       # the global scope
       assert {%T{data: %Expr{op: :multiply, args: [x, arg1_left]}}} = left.expr