Refactor symbolic shape: Make shape as expr & Visit attributes

kendryte · Feb 7, 2025 · b9e5684 · b9e5684
1 parent 2fc1a28
commit b9e5684
Show file tree

Hide file tree

Showing 69 changed files with 1,358 additions and 1,314 deletions.
diff --git a/modules/Nncase.Modules.CPU/Passes/Distributed/AutoDistributed.cs b/modules/Nncase.Modules.CPU/Passes/Distributed/AutoDistributed.cs
@@ -496,12 +496,6 @@ private Dictionary<IRType, List<Expr>> VisitLeafArgument(ParameterKind parameter
                 case (ParameterKind.Attribute, TensorConst e):
                     updateBuckets(buckets, new[] { e.With() }); // remove all old users.
                     break;
-                case (ParameterKind.Attribute, ShapeConst e):
-                    updateBuckets(buckets, new[] { e.With() }); // remove all old users.
-                    break;
-                case (ParameterKind.Attribute, DimensionConst e):
-                    updateBuckets(buckets, new[] { e.With() }); // remove all old users.
-                    break;
                 case (ParameterKind.Attribute, None e):
                     updateBuckets(buckets, new[] { e.With() });
                     break;

diff --git a/src/Nncase.Compiler/Compiler.cs b/src/Nncase.Compiler/Compiler.cs
@@ -144,6 +144,7 @@ public void TargetIndependentPass(IPassManager passManager)
             p.Add<Passes.Rules.ShapeExpr.GatherToGetItem>();
             p.Add<Passes.Rules.ShapeExpr.FoldGetItemShapeOf>();
             p.Add<Passes.Rules.Neutral.FoldGetItemConcat>();
+            p.Add<Passes.Rules.Neutral.FoldGetItemReshape>();
             p.Add<Passes.Rules.Neutral.FoldIf>();
             p.Add<Passes.Rules.Neutral.FoldNopReduce>();
             p.Add<Passes.Rules.Neutral.SliceToGetItem>();

diff --git a/src/Nncase.Core/CostModel/Cost.cs b/src/Nncase.Core/CostModel/Cost.cs
@@ -202,7 +202,7 @@ public static UInt128 GetMemoryAccess(IRType type)
     {
         return type switch
         {
-            TensorType t => (UInt128)(t.Shape.Aggregate(1D, (acc, x) => acc * (x.IsFixed ? x.FixedValue : 1)) * t.DType.SizeInBytes),
+            TensorType t => (UInt128)t.Shape.ProdWithDynamicAsOne() * (UInt128)t.DType.SizeInBytes,
             TupleType t => t.Fields.Sum(GetMemoryAccess),
             DistributedType t => GetMemoryAccess(Utilities.DistributedUtility.GetDividedTensorType(t)),
             _ => 0,
@@ -218,7 +218,7 @@ public static UInt128 GetFakeMemoryAccess(IRType type, uint bits)
     {
         return type switch
         {
-            TensorType t => (UInt128)Math.Ceiling((float)t.Shape.Aggregate(1D, (acc, x) => acc * (x.IsFixed ? x.FixedValue : 1)) * t.DType.SizeInBytes * bits / 8),
+            TensorType t => (UInt128)Math.Ceiling((float)t.Shape.ProdWithDynamicAsOne() * t.DType.SizeInBytes * bits / 8),
             TupleType t => t.Fields.Sum(x => GetFakeMemoryAccess(x, bits)),
             _ => 0,
         };
@@ -228,7 +228,7 @@ public static UInt128 GetCPUCycles(IRType type, double cyclesPerElement = 1)
     {
         return type switch
         {
-            TensorType t => (UInt128)(t.Shape.Aggregate(1D, (acc, x) => acc * (x.IsFixed ? x.FixedValue : 1)) * cyclesPerElement),
+            TensorType t => (UInt128)(t.Shape.ProdWithDynamicAsOne() * cyclesPerElement),
             TupleType t => t.Fields.Sum(GetMemoryAccess),
             DistributedType t => GetCPUCycles(Utilities.DistributedUtility.GetDividedTensorType(t)),
             _ => 0,

diff --git a/src/Nncase.Core/Evaluator/Metric.cs b/src/Nncase.Core/Evaluator/Metric.cs
@@ -189,7 +189,7 @@ public static UInt128 GetFLOPs(IRType type, long scale = 1)
     {
         return type switch
         {
-            TensorType t => (UInt128)t.Shape.Aggregate(scale, (acc, x) => acc * (x.IsFixed ? x.FixedValue : 1)),
+            TensorType t => (UInt128)t.Shape.ProdWithDynamicAsOne(scale),
             TupleType t => t.Fields.Sum(f => GetFLOPs(f, scale)),
             _ => 0,
         };

diff --git a/src/Nncase.Core/IR/Const.cs b/src/Nncase.Core/IR/Const.cs
@@ -156,10 +156,6 @@ public static Const FromValue(IValue value)
                     : new TensorConst(tv.AsTensor());
             case TupleValue tpv:
                 return new TupleConst(tpv);
-            case ShapeValue sv:
-                return new ShapeConst(sv.Dimensions.ToArray());
-            case DimensionValue dv:
-                return new DimensionConst(dv.Dimension);
             default:
                 throw new ArgumentOutOfRangeException(nameof(value));
         }