LeetGPU习题02：Matrix Copy

1
import torch
2
import triton
3
import triton.language as tl
4
@triton.autotune(
5
    configs=[
6
        triton.Config({'BLOCK_SIZE': 256}, num_warps=4),
7
        triton.Config({'BLOCK_SIZE': 512}, num_warps=4),
8
        triton.Config({'BLOCK_SIZE': 1024}, num_warps=4),
9
        triton.Config({'BLOCK_SIZE': 2048}, num_warps=8),
10
    ],
11
    key=['n_elements'],
12
)
13
@triton.jit
14
def copy_kernel(src_ptr, dst_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
15
    pid = tl.program_id(axis=0)
16
    block_start = pid * BLOCK_SIZE
17
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
18
    mask = offsets < n_elements
19
    x = tl.load(src_ptr + offsets, mask=mask)
20
    tl.store(dst_ptr + offsets, x, mask=mask)
21

22
def solve_triton(A, B, N):
23
    n_elements = A.numel()
24
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
25
    copy_kernel[grid](A, B, n_elements)

1
#include <cuda_runtime.h>
2
#include <stdio.h>
3

4
__global__ void copy_naive(const float* src, float* dst, int N) {
5
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
6

7
    if (tid < N) {
8
        dst[tid] = src[tid];
9
    }
10
}
11

12
__global__ void copy_stride(const float* src, float* dst, int N) {
13
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
14
    int stride = gridDim.x * blockDim.x;
15
    for (int i = tid; i < N; i += stride) {
16
        dst[i] = src[i];
17
    }
18
}
19

20
__global__ void copy_vec4(const float* src, float* dst, int N) {
21
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
22
    int stride = gridDim.x * blockDim.x;
23
    int vecN = N / 4;
24

25
    const float4* src4 = (const float4*)src;
26
    float4* dst4 = (float4*) dst;
27

28
    for (int i = tid; i < vecN; i += stride) {
29
        dst4[i] = src4[i];
30
    }
31

32
    int tail_idx = vecN * 4 + tid;
33
    if (tail_idx < N) {
34
        dst[tail_idx] = src[tail_idx];
35
    }
36
}
37

38
// 使用模板 T 接受任意可调用对象
39
template <typename T>
40
void run_bench(const char* name, T func, size_t bytes) {
41
    cudaEvent_t start, stop;
42
    cudaEventCreate(&start); cudaEventCreate(&stop);
43

44
    func(); // Warmup
45
    cudaEventRecord(start);
46
    for (int i = 0; i < 100; i++) func();
47
    cudaEventRecord(stop);
48
    cudaEventSynchronize(stop);
49

50
    float ms;
51
    cudaEventElapsedTime(&ms, start, stop);
52
    float gib_s = (bytes * 2.0f * 100) / (ms * 1e6f * 1.073741824f);
53
    printf("%-12s : %.2f GiB/s | Avg: %.3f ms\n", name, gib_s, ms / 100.0f);
54
}
55

56
int main() {
57
    const int N = 4097 * 4097, block = 256;
58
    const size_t bytes = N * sizeof(float);
59
    float *d_src, *d_dst;
60
    cudaMalloc(&d_src, bytes); cudaMalloc(&d_dst, bytes);
61

62
    int sms;
63
    cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0);
64
    int g_stride = sms * 32;
65

66
    run_bench("Naive", [=](){ copy_naive<<<(N+block-1)/block, block>>>(d_src, d_dst, N); }, bytes);
67
    run_bench("Stride", [=](){ copy_stride<<<g_stride, block>>>(d_src, d_dst, N); }, bytes);
68
    run_bench("Vec4", [=](){
69
        copy_vec4<<<g_stride, block>>>(d_src,d_dst, N);
70
    }, bytes);
71

72
    cudaFree(d_src); cudaFree(d_dst);
73
    return 0;
74
}