CUDA 算子优化：Transpose

Note

高性能算子系列博客

Transpose 即矩阵转置。它几乎没有计算量，性能主要取决于内存访问效率，因此是理解 CUDA Global Memory 合并访问和 Shared Memory Bank Conflict 的经典算子。

Naive 实现#

假设输入矩阵为 $M \times N$ ：

A \in \mathbb{R}^{M \times N}

Transpose 的结果为：

B = A^T \in \mathbb{R}^{N \times M}

最直接的实现是让一个线程负责一个元素：

1
__global__ void transpose(float *input, float *output, int M, int N) {
2
    int row = blockIdx.y * blockDim.y + threadIdx.y;
3
    int col = blockIdx.x * blockDim.x + threadIdx.x;
4

5
    if (row < M && col < N) {
6
        output[col * M + row] = input[row * N + col];
7
    }
8
}

对于线程(row, col)，它完成：

1
A[row][col] → B[col][row]

看起来很简单，但这个实现存在一个典型问题：Global Memory 的读写访问模式不一致。

读取是合并的#

考虑一个 warp 中的线程，一般 threadIdx.x 是连续变化的，而 threadIdx.y 是保持不变的，这些线程读取：

1
input[row][0]
2
input[row][1]
3
input[row][2]
4
...
5
input[row][31]

也就是连续的内存地址。

因此，读取属于 coalesced access（合并访问）。

Note

关于“一个 warp 中的线程，一般 threadIdx.x 是连续变化的，而 threadIdx.y 是保持不变的”，你可以运行下面的 CUDA 代码进行验证。一般 blockDim.x 设置成 32 的倍数就是为了保持这一良好特性。

1
#include <iostream>
2
#include <cuda_runtime.h>
3

4
__global__ void printThreadInfo() {
5
    // 计算一维线程 ID
6
    int tid = threadIdx.x + threadIdx.y * blockDim.x + threadIdx.z * blockDim.x * blockDim.y;
7
    // 计算所属的 Warp ID
8
    int warpId = tid / 32;
9

10
    // 只打印前两个 warp 的信息，避免输出过多
11
    if (warpId < 2) {
12
        printf("Warp %d | tid=%2d | threadIdx.x=%2d | threadIdx.y=%d\n",
13
               warpId, tid, threadIdx.x, threadIdx.y);
14
    }
15
}
16

17
int main() {
18
    std::cout << "=== Case 1: 1D Block (blockDim.x = 64) ===" << std::endl;
19
    dim3 block1D(64);
20
    printThreadInfo<<<1, block1D>>>();
21
    cudaDeviceSynchronize();
22

23
    std::cout << "\n=== Case 2: 2D Block (blockDim.x = 16, blockDim.y = 4) ===" << std::endl;
24
    dim3 block2D(16, 4);
25
    printThreadInfo<<<1, block2D>>>();
26
    cudaDeviceSynchronize();
27

28
    return 0;
29
}

写入却是不合并的#

写入地址为：

1
output[col * M + row]

当一个 warp 中的 col 持续变化时，实际访问的是：

1
output[0][row]
2
output[1][row]
3
output[2][row]
4
...
5
output[31][row]

也就是矩阵的一列。

如果矩阵较大，相邻线程访问的地址之间会相隔很远：

1
thread 0 → output[0][row]
2
thread 1 → output[1][row]
3
thread 2 → output[2][row]
4
...

因此，这些访问无法很好地合并。

所以 Naive Transpose 的问题可以概括为：

Important

Global Memory 合并读，但非合并写。

Transpose 的第一个优化目标自然也就出现了：

Important

能不能让 Global Memory 的读写都变成合并访问？

调整线程映射#

一种直观的思路是：既然 Naive 实现的问题出在写入，那就调整线程与矩阵元素之间的对应关系，让线程直接按照输出矩阵的布局进行计算。

1
__global__ void transpose(const float* input, float* output, int M, int N) {
2
    int row = blockIdx.y * blockDim.y + threadIdx.y;
3
    int col = blockIdx.x * blockDim.x + threadIdx.x;
4

5
    // output 为 N × M
6
    if (row < N && col < M) {
7
        output[row * M + col] = input[col * N + row];
8
    }
9
}

现在一个 warp 中的线程写入：

1
output[row][0]
2
output[row][1]
3
output[row][2]
4
...
5
output[row][31]

因此写入变成了连续访问。

但是，读取 input 的方式也随之发生了变化：

1
input[0][row]
2
input[1][row]
3
input[2][row]
4
...

读取变成了跨行访问。

也就是说，我们只是把问题从：

1
合并读 + 非合并写

变成了：

1
非合并读 + 合并写

仍然无法同时优化读和写。

这也说明，对于 Transpose 来说，仅仅调整线程映射是不够的。

Shared Memory：在读写之间增加一个“中转站”#

如果希望 Global Memory 同时实现合并读和合并写，就需要将 Global Memory 的访问与矩阵转置操作解耦。

一个经典的方法是使用 Shared Memory：

1
Global Memory
2
      │
3
      │ 合并读取
4
      ↓
5
Shared Memory
6
      │
7
      │ 在片上完成转置
8
      ↓
9
Global Memory
10
      │
11
      │ 合并写入
12
      ↓
13
Output

也就是说：

从 Global Memory 中连续读取一个 tile；
将 tile 放入 Shared Memory；
在 Shared Memory 中完成转置；
将转置后的 tile 连续写回 Global Memory。

这样，Global Memory 的读写都可以保持合并访问。

代码如下：

1
template <const int BLOCK_SIZE>
2
__global__ void transpose(const float* input, float* output, int M, int N)
3
{
4
    __shared__ float s_mem[BLOCK_SIZE][BLOCK_SIZE];
5

6
    int bx = blockIdx.x * BLOCK_SIZE;
7
    int by = blockIdx.y * BLOCK_SIZE;
8

9
    int x1 = bx + threadIdx.x;
10
    int y1 = by + threadIdx.y; // Global Memory → Shared Memory
11

12
    if (x1 < N && y1 < M) {
13
        s_mem[threadIdx.y][threadIdx.x] = input[y1 * N + x1];
14
    }
15
    __syncthreads();
16

17
    int x2 = by + threadIdx.x;
18
    int y2 = bx + threadIdx.y; // Shared Memory → Global Memory
19

20
    if (x2 < M && y2 < N) {
21
        output[y2 * M + x2] = s_mem[threadIdx.x][threadIdx.y];
22
    }
23
}

调用方式：

1
dim3 block(32, 32);
2
dim3 grid(CEIL(N, 32), CEIL(M, 32));
3

4
transpose<32><<<grid, block>>>(input, output, M, N);

到这里，Global Memory 的问题基本解决了：

1
Global Memory
2
    ↓
3
  合并读
4
    ↓
5
Shared Memory
6
    ↓
7
   转置
8
    ↓
9
  合并写
10
    ↓
11
Global Memory

但仔细观察 Shared Memory 的访问，又会发现一个新的问题。

Shared Memory Bank Conflict#

Shared Memory 并不是一块完全平坦的存储空间，而是划分成多个 bank。对于一个 warp，如果不同线程访问不同 bank，这些访问可以并行完成。但如果多个线程同时访问同一个 bank，就会产生 bank conflict，访问需要被拆分，从而降低 Shared Memory 的有效带宽。

Transpose 恰好很容易触发这个问题。

写入 Shared Memory 没有问题#

首先看：

1
s_mem[threadIdx.y][threadIdx.x]

一个 warp 中，threadIdx.x 连续变化，因此访问的是同一行：

1
s_mem[0][0]
2
s_mem[0][1]
3
s_mem[0][2]
4
...
5
s_mem[0][31]

这些元素连续存储，可以映射到不同的 bank。

因此，这一步没有明显的 bank conflict。

转置后读取产生 Bank Conflict#

问题出现在：

1
s_mem[threadIdx.x][threadIdx.y]

此时一个 warp 中的线程访问的是同一列：

1
s_mem[0][0]
2
s_mem[1][0]
3
s_mem[2][0]
4
...
5
s_mem[31][0]

对于：

1
__shared__ float s_mem[32][32];

每行包含 32 个 float。

因此，相邻两行之间相隔：

32\times4=128\text{ bytes}

在 bank 映射上，这意味着相邻行的相同列会落到相同的 bank。

于是一个 warp 的访问可能变成：

1
thread 0  → bank 0
2
thread 1  → bank 0
3
thread 2  → bank 0
4
...
5
thread 31 → bank 0

也就是典型的 32-way bank conflict。 [bank conflict]

所以现在的问题变成了：

1
Global Memory
2
    ↓
3
  合并读
4
    ↓
5
Shared Memory
6
    ↓
7
  Bank Conflict
8
    ↓
9
  合并写

接下来需要解决 Shared Memory 的 bank conflict。

Padding：给每一行增加一个元素#

最经典的解决方法是给 Shared Memory 增加一个 padding：

1
__shared__ float s_mem[32][33];

也就是：

1
原来： 32 × 32
2
    __shared__ float s_mem[BLOCK_SIZE][BLOCK_SIZE];
3

4
现在： 32 × 33
5
    __shared__ float s_mem[BLOCK_SIZE][BLOCK_SIZE + 1];

完整代码：

1
template <const int BLOCK_SIZE>
2
__global__ void transpose(const float* input, float* output, int M, int N)
3
{
4
    __shared__ float s_mem[BLOCK_SIZE][BLOCK_SIZE + 1];
5

6
    int bx = blockIdx.x * BLOCK_SIZE;
7
    int by = blockIdx.y * BLOCK_SIZE;
8

9
    int x1 = bx + threadIdx.x;
10
    int y1 = by + threadIdx.y; // Global Memory → Shared Memory
11

12
    if (x1 < N && y1 < M) {
13
        s_mem[threadIdx.y][threadIdx.x] = input[y1 * N + x1];
14
    }
15
    __syncthreads();
16

17
    int x2 = by + threadIdx.x;
18
    int y2 = bx + threadIdx.y; // Shared Memory → Global Memory
19

20
    if (x2 < M && y2 < N) {
21
        output[y2 * M + x2] = s_mem[threadIdx.x][threadIdx.y];
22
    }
23
}

为什么仅仅增加一个元素就可以解决问题？

原来的行跨度是：32 elements

padding 后变成：33 elements

因此读取：

1
s_mem[0][0]
2
s_mem[1][0]
3
s_mem[2][0]
4
...
5
s_mem[31][0]

时，相邻线程访问的 bank 不再完全重合，而是依次错开，从而避免了严重的 bank conflict。

这种方法虽然简单，但它揭示了一个很重要的思想：

Important

Shared Memory 中的数据布局会影响 bank 映射。改变数据布局，就可以改变线程访问的 bank。

Swizzling：改变数据布局#

除了 padding，还可以直接改变数据在 Shared Memory 中的布局。

一种常见方法是使用 XOR 进行 swizzling：

1
threadIdx.x ^ threadIdx.y

代码如下：

1
template <const int BLOCK_SIZE>
2
__global__ void transpose(float* input, float* output, int M, int N) {
3
    __shared__ float s_mem[BLOCK_SIZE][BLOCK_SIZE];  // 不需要padding
4

5
    int bx = blockIdx.x * BLOCK_SIZE;
6
    int by = blockIdx.y * BLOCK_SIZE;
7

8
    int x1 = bx + threadIdx.x;
9
    int y1 = by + threadIdx.y;
10

11
    if (x1 < N && y1 < M) {
12
        s_mem[threadIdx.y][threadIdx.x ^ threadIdx.y] = input[y1 * N + x1];
13
    }
14
    __syncthreads();
15

16
    int x2 = by + threadIdx.x;
17
    int y2 = bx + threadIdx.y;
18

19
    if (x2 < M && y2 < N) {
20
        output[y2 * M + x2] = s_mem[threadIdx.x][threadIdx.x ^ threadIdx.y];  // swizzling后，此处不存在bank conflict
21
    }
22
}

这里 Shared Memory 仍然是：

1
__shared__ float s_mem[32][32];

并没有额外增加 padding。

区别在于，线程不再简单地使用：

1
s_mem[y][x]

而是使用：

1
s_mem[y][x ^ y]

也就是说，我们并没有改变矩阵转置的逻辑，只是改变了数据在 Shared Memory 中的物理布局。

这样可以重新组织 bank 的映射关系，从而避免转置读取时产生严重的 bank conflict。

因此，Padding 和 Swizzling 的本质其实是一样的：

Important

通过改变 Shared Memory 中的数据布局，避免多个线程访问同一个 bank。

优化思路总结#

Transpose 的优化过程其实非常典型：

从这个过程可以看到，Transpose 本身没有复杂的计算，优化的核心始终围绕着数据如何移动：

Global Memory：尽可能保证 warp 的访问连续，减少非合并访问；
Shared Memory：利用片上高速存储重新组织数据布局；
Bank Conflict：避免一个 warp 中的线程集中访问同一个 bank；
最终目标：让 Global Memory 的读写都保持合并访问，同时让 Shared Memory 的访问尽可能无冲突。

这也是 Transpose 最经典的 CUDA 优化模式：

Important

用 Shared Memory 把两个方向的 Global Memory 访问解耦，再通过数据布局解决 Shared Memory 的 Bank Conflict。

Naive 实现#

读取是合并的#

写入却是不合并的#

调整线程映射#

Shared Memory：在读写之间增加一个“中转站”#

Shared Memory Bank Conflict#

写入 Shared Memory 没有问题#

转置后读取产生 Bank Conflict#

Padding：给每一行增加一个元素#

Swizzling：改变数据布局#

优化思路总结#

支持与分享

目录