fix(kernel): 修正 SliceInfo 优化计算

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

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

@@ -39,41 +39,36 @@ namespace refactor::kernel {
             }
             dims_.resize(rank = shape.size());
         }
-        dims.reserve(rank);
+        // 合并末尾的连续维度
+        for (auto i : range0_(rank).rev()) {
+            if (auto d = shape[i]; dims_[i].length == d) {
+                blockSize *= d;
+                shape.pop_back();
+                dims_.pop_back();
+            } else {
+                dims.resize(rank = shape.size());
+                if (auto &dim = dims_[i]; dim.step == 1) {
+                    if (auto times = std::gcd(std::gcd(dim.start, dim.length), shape[i]); times > 1) {
+                        blockSize *= times;
+                        dim.start /= times;
+                        dim.length /= times;
+                        shape[i] /= times;
+                    }
+                }
+                break;
+            }
+        }
         dim_t strideI = 1;
         for (auto i : range0_(rank).rev()) {
             auto const &dim = dims_[i];
-            dims.push_back({
+            dims[i] = {
                 .strideO = blockCount,
                 .skip = static_cast<dim_t>(strideI * dim.start),
                 .strideI = static_cast<sdim_t>(strideI * dim.step),
-            });
+            };
             blockCount *= dim.length;
             strideI *= shape[i];
         }
-        std::reverse(dims.begin(), dims.end());
-
-        while (!dims.empty()) {
-            auto const &dim = dims.back();
-            if (dim.strideI == static_cast<sdim_t>(dim.strideO) && !dim.skip) {
-                dims.pop_back();
-            } else {
-                long times = std::gcd(std::gcd(dim.strideI, dim.strideO), dim.skip);
-                blockCount /= times;
-                blockSize *= times;
-                if (!dims.empty()) {
-                    for (auto &dim : dims) {
-                        dim.strideO /= times;
-                        dim.skip /= times;
-                        dim.strideI /= times;
-                    }
-                    if (dims.back().strideO != 1) {
-                        dims.push_back({1, 0, 1});
-                    }
-                }
-                break;
-            }
-        }
     }
 
     SliceInfo SliceInfo::reform(dim_t maxblockSize) const noexcept {
@@ -97,5 +92,4 @@ namespace refactor::kernel {
         dims.back() = {1, 0, 1};
     }
 
-
 }// namespace refactor::kernel

src/04kernel/test/kernels/slice/test_cpu.cpp

@@ -59,7 +59,8 @@ TEST(kernel, SliceCpu) {
     };
     auto input = Tensor::share(DataType::F32, Shape{7, 6, 5, 1, 2, 3}),
          output = Tensor::share(DataType::F32, Shape{3, 2, 3, 1, 2, 3});
-    auto kernel = SliceCpu::build(SliceInfo(dims, *input));
+    auto info = SliceInfo(dims, *input);
+    auto kernel = SliceCpu::build(info);
     ASSERT_TRUE(kernel);
     auto res = runtime::Resources();
     auto routine = kernel->lower(res).routine;
@@ -94,7 +95,7 @@ TEST(kernel, SliceCpu) {
         }
     }
     // test reform
-    auto kernelReformed = SliceCpu::build(SliceInfo(dims, *input).reform(16));
+    auto kernelReformed = SliceCpu::build(info.reform(16));
     ASSERT_TRUE(kernelReformed);
     auto routineReformed = kernelReformed->lower(res).routine;
     std::vector<float> resultReformed(result.size());