Intro-ops/docs/diagrams/thread-grid-layout.md

线程网格布局 — vector_add

概念

vector_add 是最经典的逐元素并行模式。输入 a 和 b 是两个等长一维 tensor，输出 c[i] = a[i] + b[i]。每个线程处理一组独立元素，线程之间无需通信。

Mermaid 图示

graph TB
    subgraph "输入 A (N=8)"
        A0["a[0]"] --- A1["a[1]"] --- A2["a[2]"] --- A3["a[3]"]
        A4["a[4]"] --- A5["a[5]"] --- A6["a[6]"] --- A7["a[7]"]
    end

    subgraph "输入 B (N=8)"
        B0["b[0]"] --- B1["b[1]"] --- B2["b[2]"] --- B3["b[3]"]
        B4["b[4]"] --- B5["b[5]"] --- B6["b[6]"] --- B7["b[7]"]
    end

    subgraph "Grid (4 Threads, stride=4)"
        T0["Thread 0<br/>i=0,4"]
        T1["Thread 1<br/>i=1,5"]
        T2["Thread 2<br/>i=2,6"]
        T3["Thread 3<br/>i=3,7"]
    end

    subgraph "输出 C (N=8)"
        C0["c[0]"] --- C1["c[1]"] --- C2["c[2]"] --- C3["c[3]"]
        C4["c[4]"] --- C5["c[5]"] --- C6["c[6]"] --- C7["c[7]"]
    end

    A0 --> T0 --> C0
    B0 --> T0
    A4 --> T0 --> C4
    B4 --> T0
    A1 --> T1 --> C1
    B1 --> T1
    A5 --> T1 --> C5
    B5 --> T1
    A2 --> T2 --> C2
    B2 --> T2
    A6 --> T2 --> C6
    B6 --> T2
    A3 --> T3 --> C3
    B3 --> T3
    A7 --> T3 --> C7
    B7 --> T3

代码对应

template <typename T>
__global__ void vector_add_contiguous_kernel(
    T *out, const T *a, const T *b, int64_t n) {

    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = gridDim.x * blockDim.x;

    for (int64_t i = idx; i < n; i += stride) {
        out[i] = a[i] + b[i];
    }
}

要点

• 每个 (a[i], b[i]) 对独立计算，线程间零通信
• 相邻线程访问相邻内存地址 → 合并访问（coalesced access）
• TileLang 版本用 T.Parallel 表达同样的并行语义