LeetGPU习题01：Matrix Addition

1
import torch
2
import triton
3
import triton.language as tl
4
import time
5
import numpy as np
6

7
# 方案 1：直接加法
8
@triton.jit
9
def matrix_add_kernel(a, b, c, n_elements, BLOCK_SIZE: tl.constexpr):
10
    pid = tl.program_id(axis=0)
11

12
    block_start = pid * BLOCK_SIZE
13
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
14
    mask = offsets < n_elements
15

16
    ga = tl.load(a + offsets, mask=mask)
17
    gb = tl.load(b + offsets, mask=mask)
18
    gc = ga + gb
19
    tl.store(c + offsets, gc, mask=mask)
20

21
def solve_triton_naive(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):
22
    BLOCK_SIZE = 1024
23
    n_elements = N * N
24
    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
25
    matrix_add_kernel[grid](a, b, c, n_elements, BLOCK_SIZE)
26

27

28
# 方案 2：Triton 一维向量化 + Autotune
29
@triton.autotune(
30
    configs=[
31
        triton.Config({'BLOCK_SIZE': 1024, 'VEC_WIDTH': 1}, num_warps=4),
32
        triton.Config({'BLOCK_SIZE': 1024, 'VEC_WIDTH': 2}, num_warps=4),
33
        triton.Config({'BLOCK_SIZE': 2048, 'VEC_WIDTH': 2}, num_warps=8),
34
        triton.Config({'BLOCK_SIZE': 4096, 'VEC_WIDTH': 4}, num_warps=8),
35
        triton.Config({'BLOCK_SIZE': 4096, 'VEC_WIDTH': 8}, num_warps=16),
36
    ],
37
    key=['n_elements'],
38
)
39
@triton.jit
40
def matrix_add_kernel_1d(
41
    a_ptr, b_ptr, c_ptr,
42
    n_elements: tl.constexpr,
43
    BLOCK_SIZE: tl.constexpr,
44
    VEC_WIDTH: tl.constexpr,
45
):
46
    pid = tl.program_id(axis=0)
47
    block_start = pid * BLOCK_SIZE * VEC_WIDTH
48

49
    offsets = block_start + tl.arange(0, BLOCK_SIZE)[:, None] * VEC_WIDTH + tl.arange(0, VEC_WIDTH)[None, :]
50
    offsets = tl.reshape(offsets, (BLOCK_SIZE * VEC_WIDTH,))
51

52
    mask = offsets < n_elements
53

54
    a_vals = tl.load(a_ptr + offsets, mask=mask, other=0.0)
55
    b_vals = tl.load(b_ptr + offsets, mask=mask, other=0.0)
56
    c_vals = a_vals + b_vals
57

58
    tl.store(c_ptr + offsets, c_vals, mask=mask)
59

60
def solve_triton_1d(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
61
    n_elements = N * N
62
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE'] * meta['VEC_WIDTH']),)
63
    matrix_add_kernel_1d[grid](A, B, C, n_elements)
64

65
# 方案 3：Triton 二维块指针 + Autotune
66
@triton.autotune(
67
    configs=[
68
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128}, num_warps=4),
69
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256}, num_warps=4),
70
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128}, num_warps=8),
71
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 256}, num_warps=8),
72
        triton.Config({'BLOCK_M': 512, 'BLOCK_N': 128}, num_warps=8),
73
        triton.Config({'BLOCK_M': 512, 'BLOCK_N': 256}, num_warps=8),
74
        triton.Config({'BLOCK_M': 512, 'BLOCK_N': 512}, num_warps=8),
75
    ],
76
    key=['N'],
77
)
78
@triton.jit
79
def matrix_add_kernel_2d(
80
    a_ptr, b_ptr, c_ptr,
81
    N,
82
    BLOCK_M: tl.constexpr,
83
    BLOCK_N: tl.constexpr,
84
):
85
    pid_m = tl.program_id(axis=0)
86
    pid_n = tl.program_id(axis=1)
87

88
    a_block_ptr = tl.make_block_ptr(
89
        base=a_ptr,
90
        shape=(N, N),
91
        strides=(N, 1),
92
        offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
93
        block_shape=(BLOCK_M, BLOCK_N),
94
        order=(1, 0),
95
    )
96
    b_block_ptr = tl.make_block_ptr(
97
        base=b_ptr,
98
        shape=(N, N),
99
        strides=(N, 1),
100
        offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
101
        block_shape=(BLOCK_M, BLOCK_N),
102
        order=(1, 0),
103
    )
104
    c_block_ptr = tl.make_block_ptr(
105
        base=c_ptr,
106
        shape=(N, N),
107
        strides=(N, 1),
108
        offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
109
        block_shape=(BLOCK_M, BLOCK_N),
110
        order=(1, 0),
111
    )
112

113
    a = tl.load(a_block_ptr, boundary_check=(0, 1))
114
    b = tl.load(b_block_ptr, boundary_check=(0, 1))
115
    c = a + b
116
    tl.store(c_block_ptr, c, boundary_check=(0, 1))
117

118
def solve_triton_2d(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
119
    grid = lambda meta: (
120
        triton.cdiv(N, meta['BLOCK_M']),
121
        triton.cdiv(N, meta['BLOCK_N']),
122
    )
123
    matrix_add_kernel_2d[grid](A, B, C, N)
124

125
# ------------------------------------------------------------
126
# 性能测试工具
127
# ------------------------------------------------------------
128
def benchmark(func, A, B, C, N, warmup=10, repeat=100):
129
    """
130
    运行指定函数多次，返回平均耗时（毫秒）。
131
    """
132
    # 预热
133
    for _ in range(warmup):
134
        func(A, B, C, N)
135
    torch.cuda.synchronize()
136

137
    # 计时
138
    start = time.perf_counter()
139
    for _ in range(repeat):
140
        func(A, B, C, N)
141
    torch.cuda.synchronize()
142
    end = time.perf_counter()
143

144
    avg_time_ms = (end - start) / repeat * 1000
145
    return avg_time_ms
146

147
def verify_results(C_triton_naive, C_triton_1d, C_triton_2d):
148
    """
149
    验证三种方案结果是否一致。
150
    """
151
    if torch.allclose(C_triton_naive, C_triton_1d, atol=1e-5):
152
        print("✅ Naive Triton 与 Triton 1D 结果一致")
153
    else:
154
        print("❌ Naive Triton 与 Triton 1D 结果不一致")
155

156
    if torch.allclose(C_triton_naive, C_triton_2d, atol=1e-5):
157
        print("✅ Naive Triton 与 Triton 2D 结果一致")
158
    else:
159
        print("❌ Naive Triton 与 Triton 2D 结果不一致")
160

161
def main():
162
    # 检查 CUDA 可用性
163
    if not torch.cuda.is_available():
164
        raise RuntimeError("CUDA 不可用，请在有 GPU 的环境下运行")
165
    device = torch.device("cuda")
166
    print(f"运行设备: {torch.cuda.get_device_name(device)}")
167

168
    # 问题规模
169
    N = 4096
170

171
    # 分配 GPU 张量
172
    A = torch.randn(N, N, device=device, dtype=torch.float32)
173
    B = torch.randn(N, N, device=device, dtype=torch.float32)
174
    C_triton_naive = torch.empty_like(A)
175
    C_triton_1d = torch.empty_like(A)
176
    C_triton_2d = torch.empty_like(A)
177

178
    print(f"\n矩阵大小: {N} x {N} ({N*N} 个元素)")
179

180
    # 验证正确性
181
    solve_triton_naive(A, B, C_triton_naive, N)
182
    solve_triton_1d(A, B, C_triton_1d, N)
183
    solve_triton_2d(A, B, C_triton_2d, N)
184
    verify_results(C_triton_naive, C_triton_1d, C_triton_2d)
185

186
    # 性能测试
187
    print("\n开始性能测试 (预热 10 次，计时 100 次取平均)...\n")
188

189
    time_pytorch = benchmark(solve_triton_naive, A, B, C_triton_naive, N)
190
    time_triton_1d = benchmark(solve_triton_1d, A, B, C_triton_1d, N)
191
    time_triton_2d = benchmark(solve_triton_2d, A, B, C_triton_2d, N)
192

193
    # 输出结果
194
    print(f"Triton 直接加法:      {time_pytorch:.4f} ms")
195
    print(f"Triton 1D 向量化:      {time_triton_1d:.4f} ms")
196
    print(f"Triton 2D 块指针:      {time_triton_2d:.4f} ms")
197

198
    # 计算加速比
199
    baseline = time_pytorch
200
    print(f"\n相对于 Naive 的加速比:")
201
    print(f"  Triton 1D: {baseline / time_triton_1d:.2f}x")
202
    print(f"  Triton 2D: {baseline / time_triton_2d:.2f}x")
203

204
    # 计算内存带宽
205
    bytes_per_element = 4  # float32
206
    total_bytes = 3 * N * N * bytes_per_element  # A读 + B读 + C写
207
    bw_pytorch = total_bytes / (time_pytorch / 1000) / 1e9
208
    bw_triton_1d = total_bytes / (time_triton_1d / 1000) / 1e9
209
    bw_triton_2d = total_bytes / (time_triton_2d / 1000) / 1e9
210

211
    print(f"\n估算内存带宽 (GB/s):")
212
    print(f"  Triton Naive:  {bw_pytorch:.2f} GB/s")
213
    print(f"  Triton 1D: {bw_triton_1d:.2f} GB/s")
214
    print(f"  Triton 2D: {bw_triton_2d:.2f} GB/s")
215

216
if __name__ == "__main__":
217
    main()

1
#include <cuda_runtime.h>
2
#include <stdio.h>
3
#include <stdlib.h>
4
#include <math.h>
5

6
__global__ void add_naive(const float* __restrict__ A,
7
                          const float* __restrict__ B,
8
                          float* __restrict__ C, int N) {
9
    int total = N * N;
10
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
11
    if (tid < total) {
12
        C[tid] = A[tid] + B[tid];
13
    }
14
}
15

16
__global__ void add_grid_stride_scalar(const float* __restrict__ A,
17
                                       const float* __restrict__ B,
18
                                       float* __restrict__ C, int N) {
19
    int total = N * N;
20
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
21
    int stride = gridDim.x * blockDim.x;
22
    for (int i = tid; i < total; i += stride) {
23
        C[i] = A[i] + B[i];
24
    }
25
}
26

27
__global__ void add_grid_stride_vec4(const float* __restrict__ A,
28
                                     const float* __restrict__ B,
29
                                     float* __restrict__ C, int N) {
30
    int total = N * N;
31
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
32
    int stride = gridDim.x * blockDim.x;
33

34
    const float4* A4 = reinterpret_cast<const float4*>(A);
35
    const float4* B4 = reinterpret_cast<const float4*>(B);
36
    float4*       C4 = reinterpret_cast<float4*>(C);
37

38
    int num_vec = total / 4;
39
    for (int i = tid; i < num_vec; i += stride) {
40
        float4 a = __ldg(A4 + i);
41
        float4 b = __ldg(B4 + i);
42
        float4 c;
43
        c.x = a.x + b.x;
44
        c.y = a.y + b.y;
45
        c.z = a.z + b.z;
46
        c.w = a.w + b.w;
47
        C4[i] = c;
48
    }
49
    int rem = num_vec * 4;
50
    for (int i = rem + tid; i < total; i += stride) {
51
        C[i] = A[i] + B[i];
52
    }
53
}
54

55
__global__ void add_2d_grid_stride(const float* __restrict__ A,
56
                                   const float* __restrict__ B,
57
                                   float* __restrict__ C, int N) {
58
    int row = blockIdx.y * blockDim.y + threadIdx.y;
59
    int col = blockIdx.x * blockDim.x + threadIdx.x;
60
    int stride_y = gridDim.y * blockDim.y;
61
    int stride_x = gridDim.x * blockDim.x;
62

63
    for (int r = row; r < N; r += stride_y) {
64
        for (int c = col; c < N; c += stride_x) {
65
            int idx = r * N + c;
66
            C[idx] = A[idx] + B[idx];
67
        }
68
    }
69
}
70

71
// ------------------------------------------------------------------
72
// 辅助函数：理论带宽 (GB/s)
73
// ------------------------------------------------------------------
74
float theo_bw_GBs(int dev) {
75
    int clock_kHz, bus_width;
76
    cudaDeviceGetAttribute(&clock_kHz, cudaDevAttrMemoryClockRate, dev);
77
    cudaDeviceGetAttribute(&bus_width, cudaDevAttrGlobalMemoryBusWidth, dev);
78
    return (2.0f * clock_kHz * 1000.0f * (bus_width / 8)) / 1e9f;
79
}
80

81
// ------------------------------------------------------------------
82
// 辅助函数：benchmark（返回单次平均耗时 ms）
83
// ------------------------------------------------------------------
84
typedef void (*kernel_fn)(const float*, const float*, float*, int);
85
float bench(kernel_fn fn, dim3 grid, dim3 block,
86
            const float* dA, const float* dB, float* dC, int N,
87
            int warmup, int iters, cudaStream_t s = 0) {
88
    for (int i = 0; i < warmup; ++i)
89
        fn<<<grid, block, 0, s>>>(dA, dB, dC, N);
90
    cudaDeviceSynchronize();
91

92
    cudaEvent_t start, stop;
93
    cudaEventCreate(&start); cudaEventCreate(&stop);
94
    cudaEventRecord(start, s);
95
    for (int i = 0; i < iters; ++i)
96
        fn<<<grid, block, 0, s>>>(dA, dB, dC, N);
97
    cudaEventRecord(stop, s);
98
    cudaEventSynchronize(stop);
99
    float ms;
100
    cudaEventElapsedTime(&ms, start, stop);
101
    cudaEventDestroy(start); cudaEventDestroy(stop);
102
    return ms / iters;
103
}
104

105
// ------------------------------------------------------------------
106
// 辅助函数：验证
107
// ------------------------------------------------------------------
108
bool verify(const float* ref, const float* out, int N) {
109
    int total = N * N;
110
    for (int i = 0; i < total; ++i) {
111
        if (fabsf(ref[i] - out[i]) > 1e-5f) {
112
            printf("  mismatch at %d: %f vs %f\n", i, ref[i], out[i]);
113
            return false;
114
        }
115
    }
116
    return true;
117
}
118

119
// ------------------------------------------------------------------
120
// 主程序
121
// ------------------------------------------------------------------
122
int main() {
123
    int dev = 0;
124
    cudaSetDevice(dev);
125
    cudaDeviceProp prop;
126
    cudaGetDeviceProperties(&prop, dev);
127
    float peak_bw = theo_bw_GBs(dev);
128

129
    const int N = 4096;
130
    const int total = N * N;
131
    size_t bytes = total * sizeof(float);
132
    int warmup = 5, iters = 10;
133

134
    // 主机数据
135
    float *hA = (float*)malloc(bytes);
136
    float *hB = (float*)malloc(bytes);
137
    float *hRef = (float*)malloc(bytes);
138
    float *hC  = (float*)malloc(bytes);
139
    for (int i = 0; i < total; ++i) {
140
        hA[i] = (float)(rand() % 100) / 10.0f;
141
        hB[i] = (float)(rand() % 100) / 10.0f;
142
        hRef[i] = hA[i] + hB[i];
143
    }
144

145
    // 设备内存
146
    float *dA, *dB, *dC;
147
    cudaMalloc(&dA, bytes);
148
    cudaMalloc(&dB, bytes);
149
    cudaMalloc(&dC, bytes);
150
    cudaMemcpy(dA, hA, bytes, cudaMemcpyHostToDevice);
151
    cudaMemcpy(dB, hB, bytes, cudaMemcpyHostToDevice);
152

153
    printf("Device: %s  |  Peak BW: %.0f GB/s  |  %dx%d (%dM elements)\n\n",
154
           prop.name, peak_bw, N, N, total / (1024*1024));
155

156
    int sm_count = prop.multiProcessorCount;
157

158
    // ---------- 启动参数 ----------
159
    const int block_1d = 256;
160
    dim3 grid_naive((total + block_1d - 1) / block_1d);   // 覆盖全部元素
161
    dim3 grid_gs(sm_count * 8);                           // 一维 Grid‑Stride
162

163
    // 二维 Grid‑Stride： 8x8 个 block，每个 block 16x16 线程
164
    dim3 grid_2d(8, 8);
165
    dim3 block_2d(16, 16);
166

167
    // ---------- 运行 Kernel ----------
168
    const int K = 4;               // 四个 kernel
169
    float ms[K], bw[K];
170
    bool ok[K];
171

172
    // 1) Naive
173
    ms[0] = bench(add_naive, grid_naive, block_1d, dA, dB, dC, N, warmup, iters);
174
    cudaMemcpy(hC, dC, bytes, cudaMemcpyDeviceToHost);
175
    ok[0] = verify(hRef, hC, N);
176
    bw[0] = (3.0 * bytes) / (ms[0] * 1e6);
177

178
    // 2) Grid‑Stride scalar
179
    ms[1] = bench(add_grid_stride_scalar, grid_gs, block_1d, dA, dB, dC, N, warmup, iters);
180
    cudaMemcpy(hC, dC, bytes, cudaMemcpyDeviceToHost);
181
    ok[1] = verify(hRef, hC, N);
182
    bw[1] = (3.0 * bytes) / (ms[1] * 1e6);
183

184
    // 3) Grid‑Stride float4
185
    ms[2] = bench(add_grid_stride_vec4, grid_gs, block_1d, dA, dB, dC, N, warmup, iters);
186
    cudaMemcpy(hC, dC, bytes, cudaMemcpyDeviceToHost);
187
    ok[2] = verify(hRef, hC, N);
188
    bw[2] = (3.0 * bytes) / (ms[2] * 1e6);
189

190
    // 4) 二维 Grid‑Stride
191
    ms[3] = bench(add_2d_grid_stride, grid_2d, block_2d, dA, dB, dC, N, warmup, iters);
192
    cudaMemcpy(hC, dC, bytes, cudaMemcpyDeviceToHost);
193
    ok[3] = verify(hRef, hC, N);
194
    bw[3] = (3.0 * bytes) / (ms[3] * 1e6);
195

196
    // ---------- 输出结果 ----------
197
    printf("%-28s  Time(ms)   BW(GB/s)   %%Peak\n",
198
           "Kernel");
199
    printf("--------------------------------------------"
200
           "------------------\n");
201
    printf("%-28s  %7.4f    %7.1f    %5.1f\n",
202
           "1D Naive", ms[0], bw[0], bw[0]/peak_bw*100);
203
    printf("%-28s  %7.4f    %7.1f    %5.1f\n",
204
           "1D Grid-Stride scalar", ms[1], bw[1], bw[1]/peak_bw*100);
205
    printf("%-28s  %7.4f    %7.1f    %5.1f\n",
206
           "1D Grid-Stride float4", ms[2], bw[2], bw[2]/peak_bw*100);
207
    printf("%-28s  %7.4f    %7.1f    %5.1f\n",
208
           "2D Grid-Stride scalar", ms[3], bw[3], bw[3]/peak_bw*100);
209

210
    printf("\nSpeedup vs 1D Naive:\n");
211
    printf("  1D Grid-Stride scalar:  %.2fx\n", ms[0]/ms[1]);
212
    printf("  1D Grid-Stride float4:  %.2fx\n", ms[0]/ms[2]);
213
    printf("  2D Grid-Stride scalar:  %.2fx\n", ms[0]/ms[3]);
214

215
    bool all_ok = true;
216
    for (int i = 0; i < K; ++i) all_ok = all_ok && ok[i];
217
    printf("\n%s\n", all_ok ? "All kernels verified OK." : "Verification FAILED.");
218

219
    // 清理
220
    cudaFree(dA); cudaFree(dB); cudaFree(dC);
221
    free(hA); free(hB); free(hRef); free(hC);
222
    return 0;
223
}