2、张量（Tensor）的内存布局：连续内存与非连续内存、Stride与内存对齐、NCHW与NHWC格式对显存的影响

大家好，我是你们的老朋友。今天我们来聊聊张量在显存里的“住法”。

很多人写CUDA代码，觉得只要把数据拷进去、算完、拷出来就完事了。嗯，我以前也这么想。直到有一次，我调一个卷积算子，性能死活上不去，profiling一看，显存带宽利用率只有30%。后来发现，问题出在内存布局上——数据在显存里是“散装”的，GPU读起来自然费劲。

说白了，张量的内存布局，决定了你的计算是“高速公路”还是“乡间小路”。

2.1 连续内存 vs 非连续内存

先问个问题：什么是连续内存？

连续内存，就是数据在显存里一个挨一个地排着，中间没有“空位”。比如一个形状为 [3, 4] 的矩阵，如果它是行优先存储的，那么第0行的4个元素紧挨着，第1行的4个元素紧挨着，第2行也是。GPU在读取时，可以一次性把一整块数据搬进寄存器或共享内存。

非连续内存呢？就是数据之间有“缝隙”。最常见的场景是切片操作。比如你从一个 [10, 10] 的矩阵里取第2行到第5行，每行只取第3列到第7列。在PyTorch里，这通常返回一个“视图”，底层数据还是原来的大矩阵，但访问时得跳着读。

核心观点：连续内存是GPU高效计算的前提。非连续内存会导致非合并访问，严重降低带宽利用率。

我在项目中遇到过这样一个坑：用PyTorch写一个自定义算子，输入是一个经过多次切片和转置后的张量。我直接把它传给CUDA kernel，结果kernel里用 ldg 指令读数据，每次只读4字节，带宽利用率惨不忍睹。后来我加了一行 .contiguous()，性能直接翻倍。

// 非连续内存访问示例（性能差）
__global__ void bad_access(float* data, int stride, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float val = data[idx * stride];  // 每次访问跨步很大
        // ... 计算 ...
    }
}

// 连续内存访问示例（性能好）
__global__ void good_access(float* data, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float val = data[idx];  // 连续访问
        // ... 计算 ...
    }
}

小技巧：在PyTorch里，可以用 tensor.is_contiguous() 检查张量是否连续。如果不连续，调用 tensor.contiguous() 会重新分配一块连续内存并拷贝数据。虽然多了一次拷贝，但往往比非连续访问快得多。

2.2 Stride 与内存对齐

Stride，中文叫“步长”，是理解内存布局的关键。每个维度都有自己的stride，表示从该维度的一个元素到下一个元素，需要跳过多少个元素。

举个例子：一个形状为 [3, 4] 的矩阵，行优先存储时，第0维的stride是4（因为一行有4个元素），第1维的stride是1（因为同一行内相邻元素紧挨着）。

但如果是转置后的矩阵，stride就变了。比如形状为 [4, 3] 的转置矩阵，第0维的stride可能是1，第1维的stride可能是4。这时候数据就不是连续的了。

布局	形状	Stride (dim0, dim1)	是否连续
行优先	[3, 4]	(4, 1)	是
列优先	[3, 4]	(1, 3)	否（对行优先架构而言）
切片后	[2, 3]	(10, 1)	否

内存对齐呢？简单说，就是数据的起始地址要满足一定的对齐要求。比如，float类型的数据最好对齐到4字节边界，double对齐到8字节边界。为什么？因为GPU的缓存行（cache line）通常是128字节，如果数据不对齐，一次读取可能跨越两个缓存行，导致额外的访存开销。

我曾经踩过的坑：有一次写一个自定义Allocator，没注意对齐。结果在A100上跑，某些kernel偶尔会出现“unaligned memory access”的警告，性能波动很大。后来我强制所有分配都按256字节对齐，问题就消失了。记住：对齐是免费的优化。

2.3 NCHW 与 NHWC 格式

这两个格式是深度学习里绕不开的话题。N是batch size，C是通道数，H是高度，W是宽度。

NCHW：通道优先。所有通道的同一位置的数据挨在一起。比如一张3通道的图片，先存所有R像素，再存所有G像素，再存所有B像素。
NHWC：像素优先。每个像素的所有通道数据挨在一起。比如先存像素(0,0)的RGB，再存像素(0,1)的RGB，以此类推。

为什么会有这两种格式？

说白了，是CPU和GPU的“口味”不同。CPU上，NCHW更友好，因为矩阵乘法（比如im2col+GEMM）在NCHW下更容易实现连续访问。而GPU上，尤其是Tensor Core出现后，NHWC逐渐成为主流。

我个人习惯是：如果算子主要用卷积，用NCHW；如果算子主要用逐元素操作或全连接，用NHWC。但这不是绝对的，得看具体硬件和库的优化情况。

举个例子，在NVIDIA的cuDNN库中，卷积操作对NHWC格式有专门的优化路径。我记得有一次，把一个ResNet-50的训练从NCHW改成NHWC，吞吐量提升了15%左右。当然，这取决于模型和batch size。

关键点：NHWC格式在GPU上更容易实现合并访问。因为每个线程处理一个像素时，它需要的所有通道数据都在连续的内存区域里。而NCHW格式下，同一个线程可能需要跨很大的stride去访问不同通道的数据。

// NCHW 访问模式（每个线程处理一个像素的所有通道）
// 假设线程 (x, y) 处理像素 (x, y)
// 在 NCHW 中，访问通道 c 的数据需要：data[c * H * W + y * W + x]
// 不同通道的地址相差 H*W，不连续

// NHWC 访问模式（每个线程处理一个像素的所有通道）
// 在 NHWC 中，访问通道 c 的数据需要：data[y * W * C + x * C + c]
// 同一像素的不同通道地址连续，容易合并

你想想看，如果每个线程要读3个float（RGB），在NHWC下，这3个float是紧挨着的，一次 float4 读取就能搞定。但在NCHW下，你得读三个不同的位置，每个位置相隔H*W个元素，这得多慢？

2.4 知识体系图

下面我用一张SVG图来总结本章的核心逻辑。这张图展示了张量内存布局的各个维度以及它们之间的关联。

这张图把本章的三个核心概念串起来了。你看，从“张量内存布局”出发，分出三个分支：连续与非连续、Stride与对齐、NCHW与NHWC。每个分支下面又有具体的子知识点。最下面那句话是我个人的总结：连续 + 对齐 + 合适格式 = 高带宽利用。

实战建议：如果你在写一个新的CUDA kernel，先检查输入张量是否连续。如果不连续，要么在host端调用 .contiguous()，要么在kernel里手动处理stride。我个人倾向于前者，因为代码更干净，而且cuDNN和cuBLAS这些库对连续张量有更好的优化。

好了，关于张量内存布局，今天就聊这么多。记住，显存里的数据怎么“住”，直接决定了你的GPU能跑多快。下次写代码前，先想想你的数据是“排着队”还是“散着步”。

📚 课程目录

2、张量（Tensor）的内存布局：连续内存与非连续内存、Stride与内存对齐、NCHW与NHWC格式对显存的影响

2.1 连续内存 vs 非连续内存

2.2 Stride 与内存对齐

2.3 NCHW 与 NHWC 格式

2.4 知识体系图