Intro-ops/docs/diagrams/grid-stride-loop.md

1.5 KiB
Raw Permalink Blame History

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 尺寸可自由调整,不影响正确性,只影响性能