fix graph partition

kendryte · Feb 12, 2025 · daf2556 · daf2556
1 parent f1daad0
commit daf2556
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/modules/Nncase.Modules.CPU/Evaluator/TIR/CPU/Binary.cs b/modules/Nncase.Modules.CPU/Evaluator/TIR/CPU/Binary.cs
@@ -2,8 +2,8 @@
 // Licensed under the Apache license. See LICENSE file in the project root for full license information.
 
 using Google.OrTools.ConstraintSolver;
-using Nncase.IR;
 using Nncase.Diagnostics;
+using Nncase.IR;
 using Nncase.IR.Affine;
 using Nncase.Schedule;
 using Nncase.TIR.CPU;
@@ -35,7 +35,6 @@ public string Visit(IPrintOpContext context, Binary target)
         return $"Binary({target.DisplayProperty()}, {context.GetArgument(target, Binary.Lhs)}, {context.GetArgument(target, Binary.Rhs)}, {context.GetArgument(target, Binary.Output)})";
     }
 
-
     private static IntExpr GetComputeCycle(IntExpr[][] bufferShapes, Solver solver, MicroKernelContext context)
     {
         var factora = System.Math.Min(context.BufferShapes[0][^1], 32);

diff --git a/modules/Nncase.Modules.CPU/Passes/CPUFunctionPartitionPass.cs b/modules/Nncase.Modules.CPU/Passes/CPUFunctionPartitionPass.cs
@@ -167,7 +167,7 @@ protected override Expr OnComplexCluster(ClusteredBidirectionalGraph<ExprVertex,
         }
 
         var cloner = new ExprClusterCloner(extractDict);
-        var outVertices = cluster.OutVertices().ToArray();
+        var outVertices = cluster.OutVertices(Algo.ClusteredGraph).ToArray();
         var clones = new List<Expr>();
         foreach (var outVertex in outVertices)
         {

diff --git a/src/Nncase.Graph/Graphs/GraphExtensions.cs b/src/Nncase.Graph/Graphs/GraphExtensions.cs
@@ -16,14 +16,23 @@ public static IEnumerable<TEdge> InEdges<TVertex, TEdge>(this IBidirectionalGrap
     public static IEnumerable<TEdge> OutEdges<TVertex, TEdge>(this IBidirectionalGraph<TVertex, TEdge> subGraph, IBidirectionalGraph<TVertex, TEdge> parentGraph)
          where TEdge : IEdge<TVertex> => subGraph.Vertices.Select(v => parentGraph.OutEdges(v)).SelectMany(e => e).Where(e => !subGraph.ContainsVertex(e.Target));
 
+    /// <summary>
+    /// Get the in vertices which in degree is 0 in this graph.
+    /// </summary>
     public static IEnumerable<TVertex> InVertices<TVertex, TEdge>(this IBidirectionalGraph<TVertex, TEdge> graph)
          where TEdge : IEdge<TVertex>
         => graph.Vertices.Where(v => graph.InDegree(v) == 0);
 
+    /// <summary>
+    /// Get the out vertices which out degree is 0 in this graph.
+    /// </summary>
     public static IEnumerable<TVertex> OutVertices<TVertex, TEdge>(this IBidirectionalGraph<TVertex, TEdge> graph)
          where TEdge : IEdge<TVertex>
         => graph.Vertices.Where(v => graph.OutDegree(v) == 0);
 
+    /// <summary>
+    /// Get the out vertices which used by external in this graph.
+    /// </summary>
     public static IEnumerable<TVertex> OutVertices<TVertex, TEdge>(this IBidirectionalGraph<TVertex, TEdge> subGraph, IBidirectionalGraph<TVertex, TEdge> parentGraph)
          where TEdge : IEdge<TVertex>
     {