Remove unused vars

src/Nncase.Compiler/Compiler.cs

@@ -260,6 +260,19 @@ public void RegisterShapeBucket(IPassManager p)
         }
 
         p.AddWithName<AddFunctionToModule>("AddIfToModule");
+
+        // 1. Optimize if
+        // 2. Optimize prefill/decode
+        for (int i = 0; i < 2; i++)
+        {
+            p.AddWithName<DataflowPass>("RemoveUnusedVars").Configure(c =>
+            {
+                c.Add<RemoveUnusedVarsByCall>();
+                c.Add<RemoveUnusedVarsByIf>();
+            });
+            p.AddWithName<AddFunctionToModule>("AddNewFunctionToModule");
+            p.AddWithName<RemoveUnusedFunctions>("RemoveUnusedFunctions");
+        }
     }
 
     public void ClearFixShape(IPassManager p)

src/Nncase.Core/PatternMatch/IfPattern.cs

@@ -0,0 +1,99 @@
+// Copyright (c) Canaan Inc. All rights reserved.
+// Licensed under the Apache license. See LICENSE file in the project root for full license information.
+
+using System;
+using Nncase.IR;
+using Nncase.IR.Tensors;
+
+namespace Nncase.PatternMatch;
+
+/// <summary>
+/// Pattern for <see cref="If"/>.
+/// </summary>
+/// <param name="Then">Then pattern.</param>
+/// <param name="Else">Else pattern.</param>
+/// <param name="Arguments">Arguments pattern.</param>
+/// <param name="Name"> name. </param>
+public sealed record IfPattern(Pattern Then, Pattern Else, VArgsPattern Arguments, string? Name) : Pattern<If>(Name)
+{
+    /// <summary>
+    /// Initializes a new instance of the <see cref="IfPattern"/> class.
+    /// </summary>
+    /// <param name="if"><see cref="If"/> expression.</param>
+    /// <param name="name">name.</param>
+    public IfPattern(If @if, string? name)
+        : this(@if.Then, @if.Else, new VArgsPattern(@if.Arguments.ToArray(), null), name)
+    {
+    }
+
+    /// <summary>
+    /// Get parameter pattern.
+    /// </summary>
+    /// <param name="parameter">Parameter info.</param>
+    /// <returns>Parameter pattern.</returns>
+    public Pattern this[ParameterInfo parameter]
+    {
+        get => Arguments.Fields[parameter.Index];
+    }
+}
+
+/// <summary>
+/// PatternMatch Utility.
+/// </summary>
+public static partial class Utility
+{
+    /// <summary>
+    /// is call .
+    /// </summary>
+    /// <param name="name">name.</param>
+    /// <param name="then">then.</param>
+    /// <param name="else">else.</param>
+    /// <param name="parameters">params.</param>
+    /// <returns>call pattern.</returns>
+    public static IfPattern IsIf(string? name, Pattern then, Pattern @else, VArgsPattern parameters) => new IfPattern(then, @else, parameters, name);
+
+    /// <summary>
+    /// is call .
+    /// </summary>
+    /// <param name="then">then.</param>
+    /// <param name="else">else.</param>
+    /// <param name="parameters">params.</param>
+    /// <returns>call pattern.</returns>
+    public static IfPattern IsIf(Pattern then, Pattern @else, VArgsPattern parameters) => IsIf(null, then, @else, parameters);
+
+    /// <summary>
+    /// is call .
+    /// </summary>
+    /// <param name="name">name.</param>
+    /// <param name="then">then.</param>
+    /// <param name="else">else.</param>
+    /// <param name="parameters">params.</param>
+    /// <returns>call pattern.</returns>
+    public static IfPattern IsIf(string? name, Pattern then, Pattern @else, params Pattern[] parameters) => new IfPattern(then, @else, new VArgsPattern(parameters, null), name);
+
+    /// <summary>
+    /// is call .
+    /// </summary>
+    /// <param name="name">name.</param>
+    /// <param name="then">then.</param>
+    /// <param name="else">else.</param>
+    /// <param name="parameters">params.</param>
+    public static IfPattern IsIf(string? name, FunctionPattern then, FunctionPattern @else, params Pattern[] parameters) => new IfPattern(then, @else, new VArgsPattern(parameters, null), name);
+
+    /// <summary>
+    /// is call .
+    /// </summary>
+    /// <param name="name">name.</param>
+    /// <param name="then">then.</param>
+    /// <param name="else">else.</param>
+    /// <param name="parameters">params.</param>
+    public static IfPattern IsIf(string? name, FunctionPattern then, FunctionPattern @else, VArgsPattern parameters) => new IfPattern(then, @else, parameters, name);
+
+    /// <summary>
+    /// is call .
+    /// </summary>
+    /// <param name="then">then.</param>
+    /// <param name="else">else.</param>
+    /// <param name="parameters">params.</param>
+    public static IfPattern IsIf(Pattern then, Pattern @else, params Pattern[] parameters) => IsIf(null, then, @else, parameters);
+}

src/Nncase.Core/Utilities/ArrayUtility.cs

@@ -86,4 +86,27 @@ public static Expr[] ToExprArray(ReadOnlySpan<int> values)
 
         return array;
     }
+
+    public static T[] SelectByIndices<T>(ReadOnlySpan<T> source, ReadOnlySpan<int> indices)
+    {
+        var result = new T[indices.Length];
+        for (int i = 0; i < indices.Length; i++)
+        {
+            result[i] = source[indices[i]];
+        }
+
+        return result;
+    }
+
+    public static T[] SelectByIndices<T>(ReadOnlySpan<T> source, IEnumerable<int> indices)
+    {
+        var result = new T[indices.Count()];
+        int i = 0;
+        foreach (var index in indices)
+        {
+            result[i++] = source[index];
+        }
+
+        return result;
+    }
 }

src/Nncase.EGraph/PatternMatch/EGraphMatcher.cs

@@ -66,6 +66,7 @@ private IReadOnlyList<MatchScope> Visit(IReadOnlyList<MatchScope> matchScopes, I
             (FusionPattern fusionPattern, Fusion fusion) => VisitLeaf(matchScopes, fusionPattern, enode, fusion),
             (FunctionPattern functionPat, Function func) => Visit(matchScopes, functionPat, enode, func),
             (CallPattern callPat, Call call) => Visit(matchScopes, callPat, enode, call),
+            (IfPattern ifPat, If @if) => Visit(matchScopes, ifPat, enode, @if),
             (MarkerPattern mkPat, Marker mk) => Visit(matchScopes, mkPat, enode, mk),
             (TuplePattern tuplePat, IR.Tuple tuple) => Visit(matchScopes, tuplePat, enode, tuple),
             (IOpPattern opPat, Op op) => VisitLeaf(matchScopes, opPat, enode, op),
@@ -165,6 +166,35 @@ private IReadOnlyList<MatchScope> Visit(IReadOnlyList<MatchScope> matchScopes, C
         return context.NewScopes;
     }
 
+    private IReadOnlyList<MatchScope> Visit(IReadOnlyList<MatchScope> matchScopes, IfPattern pattern, ENode enode, If expr)
+    {
+        var context = new MatchContext(matchScopes, pattern, expr);
+
+        if (context.HasCandidates
+            && pattern.MatchLeaf(expr)
+            && pattern.Then.MatchLeaf(expr.Then)
+            && pattern.Else.MatchLeaf(expr.Else)
+            && pattern.Arguments.MatchLeaf(expr.Arguments))
+        {
+            var newScopes = Visit(context.Candidates, pattern.Then, enode.Children[0]);
+            if (newScopes.Count > 0)
+            {
+                newScopes = Visit(newScopes, pattern.Else, enode.Children[1]);
+                if (newScopes.Count > 0)
+                {
+                    newScopes = Visit(newScopes, pattern.Arguments, enode.Children.Skip(2));
+                    if (newScopes.Count > 0)
+                    {
+                        context.NewScopes.AddRange(newScopes);
+                        context.MatchCandidates(pattern, expr);
+                    }
+                }
+            }
+        }
+
+        return context.NewScopes;
+    }
+
     private IReadOnlyList<MatchScope> Visit(IReadOnlyList<MatchScope> matchScopes, MarkerPattern pattern, ENode enode, Marker expr)
     {
         var context = new MatchContext(matchScopes, pattern, expr);

src/Nncase.Graph/PatternMatch/Matcher.g.cs

@@ -29,6 +29,21 @@ protected override bool VisitCall(Call expr, IPattern pattern)
         return DefaultVisit(expr, pattern);
     }
 
+    /// <inheritdoc/>
+    protected override bool VisitIf(If expr, IPattern pattern)
+    {
+        if (pattern is IfPattern exprPattern)
+        {
+            return exprPattern.MatchLeaf(expr)
+                && Visit(expr.Then, exprPattern.Then)
+                && Visit(expr.Else, exprPattern.Else)
+                && VisitVArgsPattern(expr.Arguments, exprPattern.Arguments)
+            ;
+        }
+
+        return DefaultVisit(expr, pattern);
+    }
+
     /// <inheritdoc/>
     protected override bool VisitFunction(Function expr, IPattern pattern)
     {

src/Nncase.Passes/Rules/Neutral/RemoveUnusedFunctions.cs

@@ -0,0 +1,58 @@
+// Copyright (c) Canaan Inc. All rights reserved.
+// Licensed under the Apache license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Reactive;
+using System.Threading.Tasks;
+using System.Xml;
+using NetFabric.Hyperlinq;
+using Nncase.IR;
+using Nncase.IR.Math;
+using Nncase.IR.Tensors;
+using Nncase.PatternMatch;
+using Nncase.Targets;
+using static Nncase.PatternMatch.F.Math;
+using static Nncase.PatternMatch.Utility;
+using static Nncase.Utilities.ReplaceUtility;
+
+namespace Nncase.Passes.Rules.ShapeBucket;
+
+public sealed class RemoveUnusedFunctions : ModulePass
+{
+    public RemoveUnusedFunctions(CompileOptions compileOptions)
+    {
+        CompileOptions = compileOptions;
+    }
+
+    public CompileOptions CompileOptions { get; }
+
+    protected override Task<IRModule> RunCoreAsync(IRModule input, RunPassContext context)
+    {
+        while (true)
+        {
+            var funcsToRemove = new HashSet<BaseFunction>(ReferenceEqualityComparer.Instance);
+            foreach (var func in input.Functions)
+            {
+                if (!ReferenceEquals(func, input.Entry)
+                    && func.Users.Count == 1)
+                {
+                    funcsToRemove.Add(func);
+                }
+            }
+
+            if (funcsToRemove.Count == 0)
+            {
+                break;
+            }
+
+            foreach (var func in funcsToRemove)
+            {
+                input.Remove(func);
+            }
+        }
+
+        return Task.FromResult(input);
+    }
+}