refactor(kernel): transpose 支持合并连续访存

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Jan 29, 2024 · bbed31e · bbed31e
1 parent f4d063a
commit bbed31e
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/src/04kernel/src/attributes/transpose_info.cc b/src/04kernel/src/attributes/transpose_info.cc
@@ -35,7 +35,12 @@ namespace refactor::kernel {
                 }
             }
         }
-        if (rank == 0) { return; }
+        if (rank == 0) {
+            dims = {{1, 1}};
+            blockSize *= blockCount;
+            blockCount = 1;
+            return;
+        }
         // 合并连续的维度
         {
             std::vector<ddim_t> mapDim(rank, 0);
@@ -68,6 +73,14 @@ namespace refactor::kernel {
             }
             perm.resize(rank);
         }
+        // 合并末尾连续访存
+        if (perm.back() == rank - 1) {
+            blockSize *= shape.back();
+            blockCount /= shape.back();
+            shape.pop_back();
+            perm.pop_back();
+            --rank;
+        }
         // 计算 stride
         struct StrideI {
             dim_t strideI;