From daf25563efca09460746c24b6b2d50d4925ac4ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=83=91=E5=90=AF=E8=88=AA?= <597323109@qq.com>
Date: Wed, 12 Feb 2025 17:54:25 +0800
Subject: [PATCH] fix graph partition

---
 modules/Nncase.Modules.CPU/Evaluator/TIR/CPU/Binary.cs   | 3 +--
 .../Passes/CPUFunctionPartitionPass.cs                   | 2 +-
 src/Nncase.Graph/Graphs/GraphExtensions.cs               | 9 +++++++++
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/modules/Nncase.Modules.CPU/Evaluator/TIR/CPU/Binary.cs b/modules/Nncase.Modules.CPU/Evaluator/TIR/CPU/Binary.cs
index c1d9614a71..89a4249c3f 100644
--- 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
index fb3efd348d..9ba79080a3 100644
--- 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
index c55eb909b8..82db08c90d 100644
--- 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>
     {