深度学习张量reshape操作的内存布局与NPU优化

1. 理解张量reshape操作的内存布局机制

在深度学习模型优化和部署过程中，reshape操作是最基础也是最容易引发性能问题的操作之一。特别是在RKNN这类NPU加速器上，理解reshape操作何时会在CPU执行、何时能在NPU执行，对模型性能优化至关重要。

1.1 连续内存与非连续内存的本质区别

连续内存指的是张量元素在物理内存中按照逻辑顺序线性排列。例如一个形状为(2,3)的张量，其6个元素在内存中的排列顺序就是[0,0]、[0,1]、[0,2]、[1,0]、[1,1]、[1,2]。这种布局下，只需知道起始地址和每个维度的步长(stride)，就能快速定位任意元素。

而非连续内存则打破了这种线性关系。常见产生非连续内存的操作包括：

• transpose/permute：交换维度顺序
• 跨步切片：如x[:, ::2, :]
• 高级索引：如x[[1,3,5]]
• 特定维度的expand/expand_as操作

这些操作不会实际移动数据，而是通过修改元数据（shape和stride）来创建新的视图(view)。此时物理内存中的元素顺序与逻辑顺序不再一致。

1.2 reshape操作的两种执行路径

当对张量执行reshape时，框架会根据内存连续性决定执行路径：

1. 仅修改元数据（NPU友好）：

   • 条件：原始张量是连续的，且新形状可以兼容现有内存布局
   • 操作：仅调整shape描述符，不移动数据
   • 示例：(2,3,4)→(24,)、(2,3,4)→(6,4)

2. 实际内存重排（需CPU介入）：

   • 条件：原始张量不连续，或新形状与现有布局不兼容
   • 操作：先创建连续副本，再应用新形状
   • 示例：transpose后的reshape、切片后的reshape

提示：在PyTorch中，contiguous()方法会显式执行内存重排，返回连续副本。而reshape()会尽可能避免拷贝，只在必要时才触发。

2. RKNN模型中的reshape处理机制

2.1 NPU对reshape操作的支持限制

Rockchip NPU对张量操作的支持有其硬件限制：

• 支持简单的shape调整（总元素数不变）
• 不支持复杂的内存重排操作
• 对非连续输入