【AI】CUDA编程中的维度

grid、block、thread的维度主要是为了让 CUDA 的逻辑线程坐标和矩阵的二维下标自然对应起来，但这不是硬性要求。

更准确地说：

CUDA 的一维、二维、三维只是给 thread 和 block 提供一种方便的"坐标编号方式"。

GPU 硬件最终仍然会把 block 内所有线程线性排列，再按每 32 个线程组成 warp 执行。

CUDA 官方也明确说明：grid 和 thread block 都可以是一维、二维或三维，这些维度主要用于简化"线程与数据/任务之间的映射"。多维 block 本身不会凭空提高性能。(NVIDIA Docs)

---

一、首先澄清：线程本身并没有真的变成"二维线程"

你写：

dim3 block(16, 16);

含义不是创建了某种特殊的二维硬件线程，而是创建：

16 × 16 = 256 个普通线程

只是给每个线程两个逻辑坐标：

threadIdx.x
threadIdx.y

例如：

(0,0) (1,0) (2,0) ... (15,0)
(0,1) (1,1) (2,1) ... (15,1)
...
(0,15)                 (15,15)

它们在硬件内部仍会被线性化：

linear_thread_id =
    threadIdx.x
  + blockDim.x * threadIdx.y
  + blockDim.x * blockDim.y * threadIdx.z;

也可以写成：

linear_thread_id =
    threadIdx.x
  + blockDim.x *
      (threadIdx.y + blockDim.y * threadIdx.z);

其中：

x 变化最快
然后是 y
最后是 z

硬件再按连续的线性线程编号，每 32 个线程组成一个 warp。(NVIDIA Docs)

所以：

一维、二维、三维

主要是程序员看到的逻辑坐标，而不是 GPU 物理上真的存在一维阵列、二维阵列、三维阵列。

---

二、为什么朴素 GEMM 通常使用二维 grid 和二维 block

设：

A[M × K]
B[K × N]
C[M × N]

每个输出元素：

C\[m,n\]=\\sum_{k=0}\^{K-1} A\[m,k\]B\[k,n

]

输出矩阵 C 有两个下标：

行 m
列 n

朴素 GEMM 通常规定：

一个 CUDA thread 计算一个 C[m,n]。

因此最自然的映射是：

threadIdx.x → C 的列 n
threadIdx.y → C 的行 m

代码：

__global__ void gemm(
    const float* A,
    const float* B,
    float* C,
    int M,
    int N,
    int K)
{
    int col =
        blockIdx.x * blockDim.x + threadIdx.x;

    int row =
        blockIdx.y * blockDim.y + threadIdx.y;

    if (row < M && col < N) {
        float sum = 0.0f;

        for (int k = 0; k < K; ++k) {
            sum += A[row * K + k]
                 * B[k * N + col];
        }

        C[row * N + col] = sum;
    }
}

启动：

dim3 block(16, 16);

dim3 grid(
    (N + block.x - 1) / block.x,
    (M + block.y - 1) / block.y
);

gemm<<<grid, block>>>(A, B, C, M, N, K);

这里：

block.x = 16：一个 block 覆盖 16 列
block.y = 16：一个 block 覆盖 16 行

所以一个 block 负责：

C 中一个 16 × 16 的输出区域

即：

block(0,0) → C[0:16,   0:16]
block(1,0) → C[0:16,  16:32]
block(2,0) → C[0:16,  32:48]

block(0,1) → C[16:32,  0:16]
block(1,1) → C[16:32, 16:32]
...

而 block 内：

thread(0,0) → tile 内第 0 行第 0 列
thread(1,0) → tile 内第 0 行第 1 列
thread(0,1) → tile 内第 1 行第 0 列

因此：

grid 的二维坐标
    → 当前处理 C 的哪个 tile

block 内线程的二维坐标
    → 当前处理 tile 内的哪个位置

---

三、所以二维 grid 和二维 block 分别代表什么

二者虽然都是二维，但意义不一样。

1. 二维 grid：定位整个输出矩阵中的 tile

blockIdx.x
blockIdx.y

表示：

这个 block 负责 C 的哪一块区域

例如：

blockIdx.x = tile column
blockIdx.y = tile row

---

2. 二维 block：定位 tile 内部的线程任务

threadIdx.x
threadIdx.y

表示：

这个线程负责 tile 内的哪一行、哪一列

因此层次是：

整个 C 矩阵
    ↓ 由 grid 切分
多个 C tile
    ↓ 每个 tile 由一个 block 处理
tile 中的多个元素
    ↓ 由 block 内线程处理

示意：

C matrix
┌──────────┬──────────┬──────────┐
│ block0,0 │ block1,0 │ block2,0 │
├──────────┼──────────┼──────────┤
│ block0,1 │ block1,1 │ block2,1 │
├──────────┼──────────┼──────────┤
│ block0,2 │ block1,2 │ block2,2 │
└──────────┴──────────┴──────────┘

每个 block 内部：

16 × 16 threads

---

四、GEMM 有 M、N、K 三个维度，为什么通常不是三维？

这是最关键的一点。

GEMM 确实有三个循环：

for (int m = 0; m < M; ++m)
    for (int n = 0; n < N; ++n)
        for (int k = 0; k < K; ++k)
            C[m][n] += A[m][k] * B[k][n];

但是三个维度的性质不同。

M、N 是输出空间维度

不同的：

C[m0,n0]
C[m1,n1]

彼此是独立的。

因此它们特别适合映射到不同：

blocks
threads

也就是做空间并行。

---

K 是归约维度

同一个输出：

C\[m,n

]

需要所有 K 项相加：

A\[m,0\]B\[0,n

Am,1B1,n

+\cdots

]

不同 K 值并不是产生不同输出，而是产生同一个输出的不同 partial sum。

因此 K 通常写成线程内部的循环：

for (int k = 0; k < K; ++k) {
    sum += A[row * K + k]
         * B[k * N + col];
}

或者在 tiled GEMM 中：

for (int kTile = 0; kTile < K; kTile += TILE_K) {
    load A tile;
    load B tile;
    accumulate partial sum;
}

所以常见 GEMM 映射是：

M → grid/block 的 y 方向
N → grid/block 的 x 方向
K → 时间循环和累加

从 AI 加速器角度说，就是：

M、N：空间展开
K：时间累加 / reduction

这和你之前学习的 MAC 阵列映射完全一致。

---

五、K 能不能映射到 grid.z？

可以，但会产生新的问题。

假设：

blockIdx.z

代表不同的 K 分段：

z = 0 → K[0:256]
z = 1 → K[256:512]
z = 2 → K[512:768]

那么不同 z block 会分别计算：

partial_C_0[m,n]
partial_C_1[m,n]
partial_C_2[m,n]

最后还需要：

C[m,n] =
    partial_C_0[m,n]
  + partial_C_1[m,n]
  + partial_C_2[m,n]

这称为类似 split-K 的思路。

好处是：

当 M、N 很小而 K 很大时
可以增加并行 block 数量

代价是：

需要额外 reduction
需要 atomic add 或额外 kernel
产生更多 partial sum 写回

所以普通 GEMM 不会直接把 K 当作第三个完全独立的输出维度。

---

六、GEMM 必须用二维 block 吗？

完全不必须。

可以把整个输出矩阵展平成一维：

C[M × N]
    ↓
C[0 ... M*N-1]

然后使用一维 block：

__global__ void gemm1D(
    const float* A,
    const float* B,
    float* C,
    int M,
    int N,
    int K)
{
    int id =
        blockIdx.x * blockDim.x + threadIdx.x;

    if (id < M * N) {
        int row = id / N;
        int col = id % N;

        float sum = 0.0f;

        for (int k = 0; k < K; ++k) {
            sum += A[row * K + k]
                 * B[k * N + col];
        }

        C[id] = sum;
    }
}

启动：

int threads = 256;
int blocks = (M * N + threads - 1) / threads;

gemm1D<<<blocks, threads>>>(A, B, C, M, N, K);

功能完全正确。

它与二维映射本质上一样：

一维：
id → row、col

二维：
threadIdx.x/y 直接得到 row、col

二维的好处主要是：

• 代码更直观；
• block 天然对应矩阵 tile；
• 行列边界判断更清楚；
• shared memory tile 更容易表达；
• 避免显式 /N 和 %N；
• 更容易把 x 方向对应连续内存。

所以你说得没错：

二维主要是为了匹配矩阵二维结构，方便程序员做任务映射，而不是 GPU 必须二维执行。

---

七、高性能 GEMM 甚至经常使用"二维 grid + 一维 block"

高性能 GEMM 中，通常不是：

一个 thread 只算一个 C 元素

而可能是：

一个 block 算 128 × 128 C tile
一个 warp 算 64 × 32 warp tile
一个 thread 算 8 × 8 micro tile

这时程序员有时会写：

dim3 grid(
    ceil_div(N, 128),
    ceil_div(M, 128)
);

dim3 block(256);

即：

grid 是二维
block 是一维

然后在线程内部计算：

int tid = threadIdx.x;

int warpId = tid / 32;
int laneId = tid % 32;

再把：

warpId、laneId

映射到二维矩阵 tile。

所以：

GEMM 的数学任务是二维的，不代表 block 声明必须是二维的。

一维或二维只是两种编号方式：

二维 block：
编程接口直接给你 x/y

一维 block：
你自己从 tid 推导 warp row、warp col、lane row、lane col

高性能代码有时更喜欢一维 block，因为其 warp 分工更显式，也更容易做复杂的 warp-level mapping。

---

八、什么时候使用一维？

典型情况是任务天然只有一个独立下标：

1. 向量运算

C[i] = A[i] + B[i];

映射：

int i =
    blockIdx.x * blockDim.x + threadIdx.x;

---

2. 一维数组处理

例如：

• 激活函数；
• elementwise 算子；
• 数组初始化；
• 向量缩放；
• histogram 的输入遍历；
• reduction 输入遍历；
• scan；
• 稀疏数据项处理。

---

3. 数据虽然多维，但你决定展平

例如四维 tensor：

[N, C, H, W]

总元素数：

N × C × H × W

可以直接展平：

int i = ...;

int w = i % W;
int h = (i / W) % H;
int c = (i / (W * H)) % C;
int n = i / (W * H * C);

因此：

数据是四维，不代表 CUDA launch 也必须是四维。

CUDA 最多提供三维坐标，但任意高维 tensor 都可以线性化。

---

九、什么时候使用二维？

当任务有两个天然独立的空间坐标时。

常见例子：

1. 图像

pixel[y][x]

int x =
    blockIdx.x * blockDim.x + threadIdx.x;

int y =
    blockIdx.y * blockDim.y + threadIdx.y;

适用于：

• 图像滤波；
• resize；
• blur；
• 边缘检测；
• 二维卷积；
• 图像颜色变换。

---

2. 矩阵

matrix[row][col]

适用于：

• 矩阵加法；
• GEMM；
• transpose；
• softmax 的行处理；
• attention score 矩阵。

---

3. 二维网格计算

例如：

• 2D stencil；
• PDE；
• heat diffusion；
• 网格仿真。

---

十、什么时候使用三维？

当任务有三个相对独立的逻辑坐标时。

1. 三维体数据

例如：

volume[z][y][x]

适用于：

• CT / MRI；
• voxel；
• 三维 stencil；
• 3D convolution；
• 流体仿真。

代码：

int x =
    blockIdx.x * blockDim.x + threadIdx.x;

int y =
    blockIdx.y * blockDim.y + threadIdx.y;

int z =
    blockIdx.z * blockDim.z + threadIdx.z;

启动：

dim3 block(8, 8, 4);

dim3 grid(
    ceil_div(X, block.x),
    ceil_div(Y, block.y),
    ceil_div(Z, block.z)
);

---

2. Batch 矩阵乘法

假设：

A[B, M, K]
B[B, K, N]
C[B, M, N]

可以映射为：

grid.x → N tile
grid.y → M tile
grid.z → batch

例如：

dim3 grid(
    ceil_div(N, TILE_N),
    ceil_div(M, TILE_M),
    batchSize
);

kernel：

int batch = blockIdx.z;

这是很自然的三维 grid：

x = 输出列 tile
y = 输出行 tile
z = 第几个矩阵

---

3. Attention 的 batch/head

Attention score：

Score[B, H, Lq, Lk]

CUDA grid 只有三个坐标，因此可能把：

batch 和 head 合并

例如：

grid.x → Lk tile
grid.y → Lq tile
grid.z → batch × head

kernel 内：

int bh = blockIdx.z;

int batch = bh / numHeads;
int head  = bh % numHeads;

这再次说明：

CUDA 坐标维数不需要与 tensor rank 完全一致，可以合并或展平维度。

---

十一、grid 和 block 不需要使用相同维数

可以有很多组合：

// 一维 grid + 一维 block
kernel<<<blocks, 256>>>();

// 二维 grid + 二维 block
kernel<<<dim3(gx, gy), dim3(16, 16)>>>();

// 二维 grid + 一维 block
kernel<<<dim3(gx, gy), 256>>>();

// 三维 grid + 二维 block
kernel<<<dim3(gx, gy, batch), dim3(32, 8)>>>();

// 三维 grid + 三维 block
kernel<<<dim3(gx, gy, gz), dim3(8, 8, 4)>>();

维数的选择是独立的：

grid 维度
    → block 在总任务空间中的组织

block 维度
    → 线程在局部任务中的组织

例如 batched GEMM 很可能是：

3D grid
2D 或 1D block

而不是 3D block。

---

十二、二维只是坐标方便，但坐标如何映射会影响性能

官方文档说，多维 block/grid 本身只是方便，并不天然改变性能；但由于线程会按照 x → y → z 的顺序线性化，维度和内存下标的对应方式会影响 warp 内存访问。(NVIDIA Docs)

对于 C/C++ row-major 矩阵：

C[row][col]

连续内存方向是 col：

C[row][0]
C[row][1]
C[row][2]
...

因此通常让：

threadIdx.x → col
threadIdx.y → row

这样 warp 中连续线程更可能访问连续内存：

lane 0 → C[row][0]
lane 1 → C[row][1]
lane 2 → C[row][2]
...

有利于 global-memory coalescing。

如果反过来：

threadIdx.x → row

连续线程可能访问：

C[0][col]
C[1][col]
C[2][col]

地址间隔是整整一行：

stride = N

这通常不利于内存合并。

所以更准确地说：

一维、二维、三维标签本身不带性能；真正影响性能的是你如何把这些坐标映射到数据地址、warp 和协作任务。

---

十三、16×16 和 32×8 虽然都是 256 threads，也不完全一样

例如：

dim3 blockA(16, 16);  // 256 threads
dim3 blockB(32, 8);   // 256 threads

总线程数相同，但 warp 的形状不同。

因为 x 变化最快。

对于：

block(16,16)

warp 0 包含：

y=0, x=0..15
y=1, x=0..15

即一个 warp 跨两行。

对于：

block(32,8)

warp 0 包含：

y=0, x=0..31

即一个 warp 正好覆盖一整行。

所以即使总线程数相同：

block 的逻辑形状

也会改变 warp 中线程与二维数据的对应关系，从而可能影响：

• global memory 合并；
• shared-memory bank access；
• 边界利用率；
• 每个 warp 的任务分布。

但这不是因为 GPU 有一个真实的 32×8 物理线程阵列，而是因为线性化顺序不同。

---

十四、选择一维、二维、三维的实际原则

不要单纯看输入 tensor 是几维，而应该看：

第一步：找到独立输出任务的下标

例如 GEMM：

独立输出下标 = (m,n)
reduction 下标 = k

因此天然是二维输出并行。

---

第二步：让连续内存方向对应 threadIdx.x

对于 row-major：

最后一维连续

所以通常：

threadIdx.x → 最后一维

例如：

矩阵：x → column
图像：x → width
三维体：x → width

---

第三步：把需要通信和复用的线程放进同一个 block

例如 tiled GEMM：

一个 block 内线程
共同加载同一个 A tile 和 B tile

因为同一个 block 内线程可以：

• 使用 shared memory；
• 使用 __syncthreads()；
• 协作处理同一 tile。(NVIDIA Docs)

---

第四步：不要轻易把 reduction 维度当作普通 grid 维度

例如：

sum
dot product
GEMM 的 K
softmax reduction

这些维度的结果需要合并。

映射到多个线程可以，但必须设计：

warp reduction
block reduction
atomic
第二个 kernel

而不是仅仅增加一个 z 维度就结束。

---

十五、最核心的理解

对于朴素 GEMM：

A[M,K] × B[K,N] → C[M,N]

CUDA 映射通常是：

grid.x       → C 的 N 方向 tile
grid.y       → C 的 M 方向 tile

threadIdx.x  → tile 内的 N 方向
threadIdx.y  → tile 内的 M 方向

K            → 每个 thread/block 内部循环累加

也就是：

M、N 是并行输出空间
K 是 reduction 空间

所以二维并不是因为"矩阵有二维，所以 GPU 必须二维"，而是因为：

GEMM 的独立输出任务构成一个 M×N 的二维空间，用二维 grid/block 表达最自然；K 则不是独立输出维度，而是生成每个输出所需要的累加维度。

最后压缩成一句：

CUDA 的 1D/2D/3D 是任务坐标系，不是硬件形状。
选择几维，取决于你想怎样给并行任务编号、怎样映射数据和复用，而不只是看输入数据有几维。