Fix UT

Author: sunnycase

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

@@ -101,6 +101,7 @@ public Cost Visit(ICostEvaluateContext context, Boxing target)
                         cost = new Cost()
                         {
                             [CostFactorNames.MemoryLoad] = CostUtility.GetMemoryAccess(distributedType),
+                            [CostFactorNames.MemoryStore] = CostUtility.GetMemoryAccess(distributedType),
                         };
                         break;
                 }
@@ -112,6 +113,7 @@ public Cost Visit(ICostEvaluateContext context, Boxing target)
                     default:
                         cost = new Cost()
                         {
+                            [CostFactorNames.MemoryLoad] = CostUtility.GetMemoryAccess(distributedType),
                             [CostFactorNames.MemoryStore] = CostUtility.GetMemoryAccess(distributedType),
                             [CostFactorNames.Synchronization] = synchronizeCost,
                         };

ntt/include/nncase/ntt/kernels/pack.h

@@ -16,6 +16,7 @@
 #include "../apply.h"
 #include "../tensor_traits.h"
 #include "../ukernels/u_pack.h"
+#include "../vector.h"
 
 namespace nncase::ntt {
 namespace detail {
@@ -157,7 +158,8 @@ template <Tensor TIn, Tensor TOut> class pack_impl<TIn, TOut, 1> {
 template <Tensor TIn, class TOut, FixedDimensions TAxes>
 void pack(const TIn &input, TOut &&output, const TAxes &axes) noexcept {
     using TVec = typename std::decay_t<TOut>::element_type;
-    static_assert(TVec::rank() == TAxes::rank(),
+    static_assert(TVec::rank() ==
+                      vector_rank_v<typename TIn::value_type> + TAxes::rank(),
                   "Output vector rank must match axes rank");
     detail::pack_impl<TIn, std::decay_t<TOut>, TAxes::rank()> impl;
     impl(input, output, axes);

ntt/include/nncase/ntt/ukernels/u_pack.h

@@ -26,7 +26,7 @@ template <class T1, class T2, bool Arch> struct u_pack_policy {
     static constexpr size_t unroll = 4;
 };
 
-template <bool Arch, Scalar TIn, Vector TOut> class u_pack {
+template <bool Arch, ScalarOrVector TIn, Vector TOut> class u_pack {
   public:
     template <Dimension TM, Dimension TN, Dimension TMStrides>
     constexpr void operator()(const TIn *input, const TM &M, const TN &N,
@@ -82,7 +82,7 @@ class u_pack2d {
 };
 } // namespace ukernels
 
-template <Scalar TIn, Dimension TM, Dimension TN, Dimension TMStrides,
+template <ScalarOrVector TIn, Dimension TM, Dimension TN, Dimension TMStrides,
           Vector TOut>
 constexpr void u_pack(const TIn *input, const TM &M, const TN &N,
                       const TMStrides &m_strides, TOut *output) noexcept {

ntt/include/nncase/ntt/vector.h

@@ -15,6 +15,7 @@
 #pragma once
 #include "detail/shape_storage.h"
 #include "detail/vector_storage.h"
+#include "nncase/ntt/dimension.h"
 #include "tensor_traits.h"
 #include <type_traits>
 
@@ -45,8 +46,11 @@ class basic_vector
         return Lanes{}.template at<Index>();
     }
 
-    static basic_vector<T, Lanes> from_scalar(T value) noexcept;
-    static basic_vector<T, Lanes> unaligned_load_from(const T *ptr) noexcept;
+    template <ScalarOrVector U>
+    static basic_vector<T, Lanes> from_scalar(U value) noexcept;
+
+    template <ScalarOrVector U>
+    static basic_vector<T, Lanes> unaligned_load_from(const U *ptr) noexcept;
 
     constexpr basic_vector() noexcept = default;
     constexpr basic_vector(const buffer_type &buffer) noexcept
@@ -118,4 +122,14 @@ template <Vector T, size_t... Lanes> struct replace_lanes_type {
 
 template <Vector T, size_t... Lanes>
 using replace_lanes_t = typename replace_lanes_type<T, Lanes...>::type;
+
+template <class T> struct vector_rank {
+    static constexpr auto value = dim_zero;
+};
+
+template <Vector T> struct vector_rank<T> {
+    static constexpr auto value = fixed_dim_v<T::rank()>;
+};
+
+template <class T> constexpr inline auto vector_rank_v = vector_rank<T>::value;
 } // namespace nncase::ntt

ntt/include/nncase/ntt/vector_ops.h

@@ -547,19 +547,29 @@ namespace nncase::ntt::vector_ops {
 template <Vector TVector> struct vload_scalar {
     using T = typename TVector::element_type;
 
-    constexpr TVector operator()(const T &value) const noexcept {
+    template <ScalarOrVector U>
+    constexpr TVector operator()(const U &value) const noexcept {
+        const auto domain =
+            TVector::shape()
+                .template slice<0, TVector::rank() - vector_rank_v<U>>();
+
         TVector vec{};
-        ntt::apply(vec.shape(), [&](auto index) { vec(index) = value; });
+        ntt::apply(domain, [&](auto index) { vec(index) = value; });
         return vec;
     }
 };
 
 template <Vector TVector> struct vunaligned_load {
     using T = typename TVector::element_type;
 
-    constexpr TVector operator()(const T *ptr) const noexcept {
+    template <ScalarOrVector U>
+    constexpr TVector operator()(const U *ptr) const noexcept {
+        const auto domain =
+            TVector::shape()
+                .template slice<0, TVector::rank() - vector_rank_v<U>>();
+
         TVector vec{};
-        ntt::apply(vec.shape(), [&](auto index) { vec(index) = *ptr++; });
+        ntt::apply(domain, [&](auto index) { vec(index) = *ptr++; });
         return vec;
     }
 };
@@ -600,13 +610,15 @@ struct vmma {
 
 namespace nncase::ntt {
 template <Scalar T, FixedShape Lanes>
-basic_vector<T, Lanes> basic_vector<T, Lanes>::from_scalar(T value) noexcept {
+template <ScalarOrVector U>
+basic_vector<T, Lanes> basic_vector<T, Lanes>::from_scalar(U value) noexcept {
     return vector_ops::vload_scalar<basic_vector<T, Lanes>>()(value);
 }
 
 template <Scalar T, FixedShape Lanes>
+template <ScalarOrVector U>
 basic_vector<T, Lanes>
-basic_vector<T, Lanes>::unaligned_load_from(const T *ptr) noexcept {
+basic_vector<T, Lanes>::unaligned_load_from(const U *ptr) noexcept {
     return vector_ops::vunaligned_load<basic_vector<T, Lanes>>()(ptr);
 }
 

src/Nncase.Core/Tensor.cs

@@ -554,6 +554,13 @@ public Tensor CastTo(DataType type, CastMode castMode, long[] dimensions)
     /// <returns>Memory handle.</returns>
     public abstract MemoryHandle PinBuffer();
 
+    /// <summary>
+    /// Reshapes the current tensor to new dimensions, using the same backing storage.
+    /// </summary>
+    /// <param name="dimensions">An span of integers that represent the size of each dimension of the DenseTensor to create.</param>
+    /// <returns>A new tensor that reinterprets backing Buffer of this tensor with different dimensions.</returns>
+    public abstract Tensor Reshape(ReadOnlySpan<long> dimensions);
+
     /// <inheritdoc/>
     public IEnumerator GetEnumerator()
     {

src/Nncase.Core/TensorOfT.cs

@@ -180,7 +180,7 @@ public T GetValue(long index)
     /// </summary>
     /// <param name="dimensions">An span of integers that represent the size of each dimension of the DenseTensor to create.</param>
     /// <returns>A new tensor that reinterprets backing Buffer of this tensor with different dimensions.</returns>
-    public Tensor<T> Reshape(ReadOnlySpan<long> dimensions)
+    public override Tensor<T> Reshape(ReadOnlySpan<long> dimensions)
     {
         if (Length != TensorUtilities.GetProduct(dimensions))
         {

src/Nncase.Evaluator/Tensors/Unpack.cs

@@ -21,7 +21,12 @@ public IValue Visit(IEvaluateContext context, Unpack target)
     {
         var dt = context.CurrentCall.Arguments[Unpack.Input.Index].CheckedDataType;
         var elementType = dt is VectorType vt ? vt.ElemType : dt;
-        var oldLanesCount = ((VectorType)dt).Lanes.Count;
+        var oldLanesCount = dt switch
+        {
+            VectorType vt2 => vt2.Lanes.Count,
+            MaskVectorType => 1,
+            _ => throw new InvalidOperationException($"Unsupported input type: {dt}"),
+        };
         if (elementType == DataTypes.Float8E4M3 || elementType == DataTypes.Float8E5M2)
         {
             var newType = new VectorType(DataTypes.UInt8, target.Lanes.ToArray());

src/Nncase.Tests/Targets/UnitTestCPUKernels.cs

@@ -474,7 +474,7 @@ public async Task TestReshard(long[] shape, int[] hierarchy, List<int[][]> sbps,
 
         var post = IR.F.Distributed.Boxing(boxed, inputType);
         post.Metadata = new Passes.Distributed.AutoDistributedMetaData() { Skip = true };
-        await RunCases($"Theory{count}", feedDict, new[] { post });
+        await RunCases($"Theory{count}", feedDict, new[] { post }, enableAutoDist: false);
     }
 
     [Theory]
@@ -1898,7 +1898,7 @@ public async Task TestVectorizeScatterND(long[] inShape, long[] indicesShape, lo
         await RunCases($"Theory{count}", feedDict, posts);
     }
 
-    internal async Task RunCases(string dumpDir, Dictionary<IVar, IValue> feedDict, IEnumerable<BaseExpr> posts, Dictionary<IVar, IValue>? feedDictRT = null)
+    internal async Task RunCases(string dumpDir, Dictionary<IVar, IValue> feedDict, IEnumerable<BaseExpr> posts, Dictionary<IVar, IValue>? feedDictRT = null, bool enableAutoDist = true)
     {
         var postArray = posts.ToArray();
         using var pinner = new ExprPinner(postArray);
@@ -1908,11 +1908,11 @@ internal async Task RunCases(string dumpDir, Dictionary<IVar, IValue> feedDict,
             System.Console.WriteLine(CompilerServices.Print(postArray[i]));
 #endif
             var kernelCase = new CpuKernelCase($"Case{i}", new Fusion("kernel", CPUTarget.Kind, postArray[i], feedDict.Keys.ToArray()), feedDict.Keys.ToArray(), feedDict.Values.Select(v => v.AsTensor()).ToArray(), feedDictRT?.Values.Select(v => v.AsTensor()).ToArray() ?? []);
-            await Run(dumpDir, kernelCase);
+            await Run(dumpDir, kernelCase, enableAutoDist: enableAutoDist);
         }
     }
 
-    internal async Task Run(string dumpDir, CpuKernelCase kernelCase)
+    internal async Task Run(string dumpDir, CpuKernelCase kernelCase, bool enableAutoDist = true)
     {
         using var dumpScope = new Diagnostics.DumpScope(Path.Join(dumpDir, kernelCase.Name), CompileOptions.DumpFlags);
 
@@ -1947,7 +1947,7 @@ internal async Task Run(string dumpDir, CpuKernelCase kernelCase)
             }
         }
 #endif
-        await Compile(module);
+        await Compile(module, enableAutoDist: enableAutoDist);
         var (kmodel_path, _) = Testing.BuildKModel("test", module, CompileSession, false);
         Tensor[] actuals;
         if (kernelCase.RTInputs.Any())
@@ -1974,12 +1974,16 @@ internal async Task Run(string dumpDir, CpuKernelCase kernelCase)
         }
     }
 
-    private async Task Compile(IRModule module)
+    private async Task Compile(IRModule module, bool enableAutoDist = true)
     {
         var pmgr = CompileSession.CreatePassManager("pmgr");
         var compiler = (Nncase.Compiler.Compiler)CompileSession.Compiler;
         compiler.TargetIndependentPass(pmgr);
-        compiler.AutoDistributedPass(pmgr);
+        if (enableAutoDist)
+        {
+            compiler.AutoDistributedPass(pmgr);
+        }
+
         compiler.AutoTilingPass(pmgr);
         compiler.TIRPass(pmgr);
         await pmgr.RunAsync(module);