feat(kernel): 实现 slice info 整合

Signed-off-by: YdrMaster <[email protected]>

Author: YdrMaster

src/00common/include/common.h

@@ -17,6 +17,8 @@ namespace refactor {
     using ddim_t = int16_t;
     /// @brief 用于表示形状的数值。
     using dim_t = uint32_t;
+    /// @brief 用于表示带符号的形状的数值。
+    using sdim_t = int32_t;
     /// @brief 用于表示对象的数量。
     using count_t = uint32_t;
 

src/04kernel/include/kernel/attributes/slice_info.h

@@ -4,12 +4,27 @@
 #include "../tensor.h"
 
 namespace refactor::kernel {
+    namespace slice {
+        struct Dim {
+            int64_t start, step, length;
+        };
+    }// namespace slice
+
+    using Dimensions = std::vector<slice::Dim>;
 
     /// @brief 优化用于计算的 Slice 描述。
     struct SliceInfo {
         struct Dim {
-            int64_t start, step, length;
+            dim_t countStride, sizeStart;
+            sdim_t sizeStride;
+
+            bool operator==(Dim const &) const noexcept;
+            bool operator!=(Dim const &) const noexcept;
         };
+        std::vector<Dim> dims;
+        dim_t blockSize;
+
+        SliceInfo(Dimensions const &, Tensor const &) noexcept;
     };
 
 }// namespace refactor::kernel

src/04kernel/include/kernel/collectors/slice.h

@@ -6,7 +6,6 @@
 #include "../target.h"
 
 namespace refactor::kernel {
-    using Dimensions = std::vector<SliceInfo::Dim>;
 
     struct SliceCollector final : public InfoCollector {
         Target target;

src/04kernel/src/attributes/slice_info.cc

@@ -2,4 +2,46 @@
 
 namespace refactor::kernel {
 
+    bool SliceInfo::Dim::operator==(Dim const &rhs) const noexcept {
+        return countStride == rhs.countStride &&
+               sizeStart == rhs.sizeStart &&
+               sizeStride == rhs.sizeStride;
+    }
+    bool SliceInfo::Dim::operator!=(Dim const &rhs) const noexcept {
+        return !operator==(rhs);
+    }
+
+    SliceInfo::SliceInfo(Dimensions const &dims_, Tensor const &input) noexcept
+        : blockSize(input.dataType.size()), dims(1) {
+        ASSERT(dims_.size() == input.rank(), "Unreachable");
+
+        auto continuous = true;
+        auto stride = blockSize;
+        dims[0] = {1, 0, static_cast<sdim_t>(stride)};
+        for (auto i : range0_(input.rank()).rev()) {
+            auto l = input.shape[i];
+            auto const &d = dims_[i];
+            if (continuous && d.step == 1) {
+                auto &it = dims.back();
+                it.countStride *= d.length;
+                it.sizeStart = d.start * stride;
+                it.sizeStride *= l;
+            } else {
+                dims.push_back(Dim{
+                    static_cast<dim_t>(dims.back().countStride * d.length),
+                    static_cast<dim_t>(d.start * stride),
+                    static_cast<sdim_t>(d.step * stride),
+                });
+            }
+            continuous = d.length == l;
+            stride *= l;
+        }
+        auto blockCount = dims[0].countStride;
+        blockSize *= blockCount;
+        for (auto &d : dims) {
+            d.countStride /= blockCount;
+        }
+        std::reverse(dims.begin(), dims.end());
+    }
+
 }// namespace refactor::kernel

src/04kernel/test/attributes/test_slice_info.cpp

@@ -0,0 +1,28 @@
+#include "kernel/attributes/slice_info.h"
+#include <gtest/gtest.h>
+
+using namespace refactor;
+using namespace kernel;
+
+TEST(kernel, SliceInfo) {
+    auto input = Tensor::share(DataType::F32, Shape{7, 6, 5, 1, 2, 3});
+    Dimensions dims{
+        {5, -2, 3},// 7 -> {5, 3, 1} -> {108, 900, -360}
+        {2, 3, 2}, // 6 -> {2, 5}    -> { 36,  60,   90}
+        {1, 1, 3}, // 5 -> {1, 2, 3} -> { 18,   6,   30}
+        {0, 1, 1}, // 1 -> {0}
+        {0, 1, 2}, // 2 -> {0, 1}
+        {0, 1, 3}, // 3 -> {0, 1, 2}
+    };
+    SliceInfo info(dims, *input);
+    EXPECT_EQ(info.blockSize, 72);
+    EXPECT_EQ(info.dims,
+              // clang-format off
+              (decltype(info.dims){
+                  {108 / 18, 900 * 4, -360 * 4},
+                  { 36 / 18,  60 * 4,   90 * 4},
+                  { 18 / 18,   6 * 4,   30 * 4},
+              })
+              // clang-format on
+    );
+}