try to fix k80 CI

Author: xhuohai

modules/Nncase.Modules.NTT/Evaluator/Distributed/Boxing.cs

@@ -47,6 +47,14 @@ IRType VisitD2D(DistributedType inv, DistributedType outv)
             {
                 if (inv.Partial is not null && outv.AxisPolicies[i] is SBPSplit s)
                 {
+                    if (inv.AxisPolicies[i] is SBPSplit splitIn)
+                    {
+                        if (splitIn.Axes.Except(s.Axes).Any())
+                        {
+                            return new InvalidType("Not Supported Split-> Split.");
+                        }
+                    }
+
                     if (s.Axes.Except(inv.Partial.Axes).ToArray() != s.Axes)
                     {
                         if (s.Axes.Except(inv.Partial.Axes).Any())

src/Nncase.Passes/Distributed/AutoDistributed.cs

@@ -557,6 +557,11 @@ string DescribeSbp(IRType? type)
                     }
                 }
 
+                if (expr.Target is not Boxing && ((Call)newExpr).Arguments.AsValueEnumerable().Any(a => a.CheckedType is DistributedType dt && dt.Partial is not null))
+                {
+                    continue;
+                }
+
                 if (!newExpr.InferenceType(_inferencer_cache) || newExpr.CheckedType is InvalidType)
                 {
                     continue;
@@ -969,7 +974,7 @@ private IRType CheckBoxingType(IRType inType, IRType outType, bool isReshape = f
     {
         IRType VisitD2D(DistributedType inv, DistributedType outv)
         {
-            if (DistributedUtility.AreSamePolicies(inv.AxisPolicies, outv.AxisPolicies))
+            if (inv.Partial == outv.Partial && DistributedUtility.AreSamePolicies(inv.AxisPolicies, outv.AxisPolicies))
             {
                 return new InvalidType("Same DistributedType");
             }
@@ -984,6 +989,14 @@ IRType VisitD2D(DistributedType inv, DistributedType outv)
             {
                 if (inv.Partial is not null && outv.AxisPolicies[i] is SBPSplit s)
                 {
+                    if (inv.AxisPolicies[i] is SBPSplit splitIn)
+                    {
+                        if (splitIn.Axes.Except(s.Axes).Any())
+                        {
+                            return new InvalidType("Not Supported Split-> Split.");
+                        }
+                    }
+
                     if (s.Axes.Except(inv.Partial.Axes).ToArray() != s.Axes)
                     {
                         if (s.Axes.Except(inv.Partial.Axes).Any())

src/Nncase.Passes/Rules/Distributed/UpdateBoxingTensorType.cs

@@ -0,0 +1,43 @@
+// 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.Text;
+using System.Threading.Tasks;
+using Nncase.IR;
+using Nncase.IR.Distributed;
+using Nncase.PatternMatch;
+using static Nncase.IR.F.NN;
+
+using static Nncase.IR.TypePatternUtility;
+using static Nncase.PatternMatch.F.Distributed;
+using static Nncase.PatternMatch.Utility;
+
+namespace Nncase.Passes.Rules;
+
+[RuleGenerator]
+public partial class UpdateBoxingTensorType : RewriteRule<Pattern>
+{
+    /// <inheritdoc/>
+    public override Pattern Pattern { get; } = IsBoxing(
+      target_name: "boxing",
+      _ => true,
+      IsWildcard("input"));
+
+    private Expr? GetReplace(Boxing boxing, Expr input, RunPassContext context)
+    {
+        var type = input.CheckedType;
+        if (type is DistributedType dt1 && boxing.NewType is DistributedType dt2)
+        {
+            var ttype = dt1.TensorType;
+            var dtype = dt2 with { TensorType = ttype };
+            var newBoxing = new Call(new IR.Distributed.Boxing(dtype), input);
+            context.MatchOptions.SuppressPattern(newBoxing, Pattern);
+            return newBoxing;
+        }
+
+        return null;
+    }
+}

src/Nncase.Schedule/Transforms/AutoTilePass.cs

@@ -11,6 +11,7 @@
 using Nncase.IR;
 using Nncase.IR.Affine;
 using Nncase.Passes.GraphPartition;
+using Nncase.Passes.Rules;
 using Nncase.Schedule;
 using QuikGraph;
 using QuikGraph.Algorithms;
@@ -116,6 +117,7 @@ protected override Task<BaseFunction> RunCoreAsync(BaseFunction input, RunPassCo
 
         var constructor = new AutoTileReconstructor(tiler, ModuleKind, CompileOptions, condenseAlgo, dimVars.ToArray());
         var post = constructor.Construct();
+        post = CompilerServices.Rewrite(post, [new UpdateBoxingTensorType()], new());
         return Task.FromResult((BaseFunction)func.With(body: post));
     }
 }