fix(kernel): 为 TransposeInfo 应对更多 coner case 并增加一个单元测试

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

Author: YdrMaster

src/04kernel/src/attributes/transpose_info.cc

@@ -73,6 +73,12 @@ namespace refactor::kernel {
             }
             perm.resize(rank);
         }
+        if (rank <= 1) {
+            dims = {{1, 1}};
+            blockSize *= blockCount;
+            blockCount = 1;
+            return;
+        }
         // 合并末尾连续访存
         if (perm.back() == rank - 1) {
             blockSize *= shape.back();
@@ -81,7 +87,6 @@ namespace refactor::kernel {
             perm.pop_back();
             --rank;
         }
-
         // 计算 stride
         struct StrideI {
             dim_t strideI;

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

@@ -0,0 +1,39 @@
+#include "kernel/attributes/transpose_info.h"
+#include <gtest/gtest.h>
+
+using namespace refactor;
+using namespace kernel;
+
+TEST(kernel, TransposeInfo) {
+    {
+        TransposeInfo info(
+            DataType::F32,
+            {1, 2, 3, 2, 1},
+            {1, 2, 3, 0, 4});
+        EXPECT_EQ(info.blockSize, 48);
+        EXPECT_EQ(info.blockCount, 1);
+        EXPECT_EQ(info.dims.size(), 1);
+    }
+    {
+        TransposeInfo info(
+            DataType::F32,
+            {1, 1, 2, 1, 1},
+            {1, 2, 3, 0, 4});
+        EXPECT_EQ(info.blockSize, 8);
+        EXPECT_EQ(info.blockCount, 1);
+        EXPECT_EQ(info.dims.size(), 1);
+    }
+    {
+        TransposeInfo info(
+            DataType::F32,
+            {1, 2, 3, 4, 5},
+            {2, 3, 1, 0, 4});
+        EXPECT_EQ(info.blockSize, 20);
+        EXPECT_EQ(info.blockCount, 24);
+        EXPECT_EQ(info.dims.size(), 2);
+        EXPECT_EQ(info.dims[1].strideI, 12);
+        EXPECT_EQ(info.dims[1].strideO, 1);
+        EXPECT_EQ(info.dims[0].strideI, 1);
+        EXPECT_EQ(info.dims[0].strideO, 2);
+    }
+}