forked from ccf-ai-infra/Intro-ops
1.5 KiB
1.5 KiB
Grid-Stride Loop 图示
概念
Grid-stride loop 是 GPU kernel 处理任意大小 tensor 的核心技术。每个线程不只处理一个元素,而是在循环中负责多个元素(跨步为 grid 总线程数),从而覆盖任意 N。
Mermaid 图示
graph TB
subgraph "Global Memory (N=10 elements)"
M0["[0]"] --- M1["[1]"] --- M2["[2]"] --- M3["[3]"] --- M4["[4]"]
M5["[5]"] --- M6["[6]"] --- M7["[7]"] --- M8["[8]"] --- M9["[9]"]
end
subgraph "Grid = 2 Blocks × 2 Threads = 4 Threads"
subgraph "Block 0"
T0["Thread 0<br/>idx=0"]
T1["Thread 1<br/>idx=1"]
end
subgraph "Block 1"
T2["Thread 2<br/>idx=2"]
T3["Thread 3<br/>idx=3"]
end
end
T0 -->|"i=0"| M0
T0 -->|"i=4"| M4
T0 -->|"i=8"| M8
T1 -->|"i=1"| M1
T1 -->|"i=5"| M5
T1 -->|"i=9"| M9
T2 -->|"i=2"| M2
T2 -->|"i=6"| M6
T3 -->|"i=3"| M3
T3 -->|"i=7"| M7
代码对应
// stride = gridDim.x * blockDim.x = 2 * 2 = 4
for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; // idx = 0,1,2,3
i < N; // N = 10
i += gridDim.x * blockDim.x) { // i += 4
dst[i] = src[i];
}
要点
- 每个线程负责多个元素,间隔为 stride(本例中 stride=4)
- 无论 N 多大,只要
i < N条件在,就不会越界 - grid/block 尺寸可自由调整,不影响正确性,只影响性能