LeetGPU | 01_Vector_Addition

关注 LeetGPU repository，了解更多内容。

Problem Description#

Write a GPU program that performs element-wise addition of two vectors containing 32-bit floating point numbers. The program should take two input vectors of equal length and produce a single output vector containing their sum.

Implementation Requirements#

External libraries are not permitted.
The solve function signature must remain unchanged.
The final result must be stored in vector C.

Example 1#

1
Input:  A = [1.0, 2.0, 3.0, 4.0]
2
        B = [5.0, 6.0, 7.0, 8.0]
3
Output: C = [6.0, 8.0, 10.0, 12.0]

Example 2#

1
Input:  A = [1.5, 1.5, 1.5]
2
        B = [2.3, 2.3, 2.3]
3
Output: C = [3.8, 3.8, 3.8]

Constraints#

Input vectors A and B have identical lengths.
1 <= N <= 100,000,000.
Performance is measured with N = 25,000,000.

CUDA#

Approach#

CUDA 采用“一元素一线程”的方式，为每个向量元素分配一个 GPU 线程。

线程的全局索引由 blockIdx.x、blockDim.x 和 threadIdx.x 共同计算得出。
每个索引有效的线程分别从 A 和 B 读取一个元素，完成加法后将结果写入 C。
每个线程块包含 256 个线程，线程块数量通过向上取整除法确定。
idx < N 的边界检查可避免最后一个线程块发生越界访问。
内核使用 CUDA C++ 编写，并通过 nvcc 编译。

Solution#

1
#include <cuda_runtime.h>
2

3
__global__ void vector_add(const float* A, const float* B, float* C, int N) {
4
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
5
    if (idx < N) {
6
        C[idx] = A[idx] + B[idx];
7
    }
8
}
9

10
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
11
extern "C" void solve(const float* A, const float* B, float* C, int N) {
12
    int threadsPerBlock = 256;
13
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
14

15
    vector_add<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, N);
16
}

Test Code#

1
#include <cuda_runtime.h>
2

3
#include <array>
4
#include <cmath>
5
#include <iostream>
6

7
extern "C" void solve(const float* a, const float* b, float* c, int n);
8

9
int main() {
10
    constexpr int n = 4;
11
    const std::array<float, n> input_a = {1.0f, 2.0f, 3.0f, 4.0f};
12
    const std::array<float, n> input_b = {5.0f, 6.0f, 7.0f, 8.0f};
13
    const std::array<float, n> expected = {6.0f, 8.0f, 10.0f, 12.0f};
14
    std::array<float, n> actual = {};
15
    constexpr size_t bytes = n * sizeof(float);
16

17
    float* device_a = nullptr;
18
    float* device_b = nullptr;
19
    float* device_c = nullptr;
20

21
    auto check_cuda = [](cudaError_t status, const char* operation) {
22
        if (status == cudaSuccess) {
23
            return true;
24
        }
25
        std::cerr << "Test failed. CUDA error in " << operation << ": "
26
                  << cudaGetErrorString(status) << '\n';
27
        return false;
28
    };
29

30
    bool success = check_cuda(cudaMalloc(&device_a, bytes), "cudaMalloc(device_a)") &&
31
                   check_cuda(cudaMalloc(&device_b, bytes), "cudaMalloc(device_b)") &&
32
                   check_cuda(cudaMalloc(&device_c, bytes), "cudaMalloc(device_c)") &&
33
                   check_cuda(cudaMemcpy(device_a, input_a.data(), bytes, cudaMemcpyHostToDevice), "copy input_a") &&
34
                   check_cuda(cudaMemcpy(device_b, input_b.data(), bytes, cudaMemcpyHostToDevice), "copy input_b");
35

36
    if (success) {
37
        solve(device_a, device_b, device_c, n);
38
        success = check_cuda(cudaGetLastError(), "solve") &&
39
                  check_cuda(cudaDeviceSynchronize(), "synchronize") &&
40
                  check_cuda(cudaMemcpy(actual.data(), device_c, bytes, cudaMemcpyDeviceToHost), "copy result");
41
    }
42

43
    cudaFree(device_a);
44
    cudaFree(device_b);
45
    cudaFree(device_c);
46

47
    if (!success) {
48
        return 1;
49
    }
50

51
    for (int index = 0; index < n; ++index) {
52
        if (std::fabs(actual[index] - expected[index]) > 1e-6f) {
53
            std::cerr << "Test failed at index " << index
54
                      << ". Expected: " << expected[index]
55
                      << ", Actual: " << actual[index] << '\n';
56
            return 1;
57
        }
58
    }
59

60
    std::cout << "Test passed.\n";
61
    return 0;
62
}

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	N=25000000	20	0.826 ms	0.748 ms	1.325 ms	363.144 GB/s
NVIDIA GeForce RTX 4090	PASS	N=25000000	20	0.324 ms	0.322 ms	0.327 ms	926.118 GB/s
NVIDIA GeForce RTX 5090	PASS	N=25000000	20	0.193 ms	0.190 ms	0.195 ms	1557.852 GB/s

Triton#

Approach#

Triton 让每个程序实例负责处理一段连续的向量元素。

tl.program_id(0) 用于标识当前程序实例。
每个程序实例最多处理由 tl.arange 生成的 1024 个元素。
内核分别加载两个数据块，逐元素相加后存储结果。
掩码可避免处理最后一个数据块时发生越界加载和存储。
内核使用 Python 编写，并由 Triton 即时编译（JIT）为 GPU 代码。

Solution#

1
import torch
2
import triton
3
import triton.language as tl
4

5

6
@triton.jit
7
def vector_add_kernel(a, b, c, n_elements, BLOCK_SIZE: tl.constexpr):
8
    pid = tl.program_id(0)
9

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

14
    a_block = tl.load(a+offsets, mask=mask)
15
    b_block = tl.load(b+offsets, mask=mask)
16

17
    c_block = a_block+b_block
18

19
    tl.store(c+offsets, c_block, mask=mask)
20

21

22
# a, b, c are tensors on the GPU
23
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):
24
    BLOCK_SIZE = 1024
25
    grid = (triton.cdiv(N, BLOCK_SIZE),)
26
    vector_add_kernel[grid](a, b, c, N, BLOCK_SIZE)

Test Code#

1
import importlib.util
2
import sys
3
from pathlib import Path
4

5
import torch
6

7

8
def load_implementation():
9
    source_path = Path(__file__).resolve().parents[2] / "src" / "triton" / "01_Vector_Addition.py"
10
    spec = importlib.util.spec_from_file_location("vector_addition_triton_minimal", source_path)
11
    if spec is None or spec.loader is None:
12
        raise RuntimeError(f"Unable to load implementation: {source_path}")
13
    module = importlib.util.module_from_spec(spec)
14
    spec.loader.exec_module(module)
15
    return module
16

17

18
def main():
19
    input_a = torch.tensor([1.0, 2.0, 3.0, 4.0], device="cuda")
20
    input_b = torch.tensor([5.0, 6.0, 7.0, 8.0], device="cuda")
21
    expected = torch.tensor([6.0, 8.0, 10.0, 12.0], device="cuda")
22
    actual = torch.empty_like(input_a)
23

24
    load_implementation().solve(input_a, input_b, actual, input_a.numel())
25
    torch.cuda.synchronize()
26
    torch.testing.assert_close(actual, expected)
27
    print("Test passed.")
28

29

30
if __name__ == "__main__":
31
    try:
32
        main()
33
    except Exception as error:
34
        print(f"Test failed: {error}", file=sys.stderr)
35
        raise

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	N=25000000	20	0.809 ms	0.764 ms	1.085 ms	371.021 GB/s
NVIDIA GeForce RTX 4090	PASS	N=25000000	20	0.341 ms	0.338 ms	0.350 ms	880.608 GB/s
NVIDIA GeForce RTX 5090	PASS	N=25000000	20	0.223 ms	0.216 ms	0.251 ms	1344.240 GB/s

PyTorch#

Approach#

PyTorch 通过内置的 torch.add 算子完成向量加法。

torch.add(A, B, out=C) 将对应位置的元素相加，并把结果直接写入 C。
PyTorch 会自动选择合适的 GPU 内核并发起执行。
线程调度、索引计算和边界处理均由 PyTorch 在内部完成。
为满足接口要求，函数签名中保留了 N，但该实现不会直接使用它。

Solution#

1
import torch
2

3

4
# A, B, C are tensors on the GPU
5
def solve(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
6
    torch.add(A,B,out = C);

Test Code#

1
import importlib.util
2
import sys
3
from pathlib import Path
4

5
import torch
6

7

8
def load_implementation():
9
    source_path = Path(__file__).resolve().parents[2] / "src" / "pytorch" / "01_Vector_Addition.py"
10
    spec = importlib.util.spec_from_file_location("vector_addition_pytorch_minimal", source_path)
11
    if spec is None or spec.loader is None:
12
        raise RuntimeError(f"Unable to load implementation: {source_path}")
13
    module = importlib.util.module_from_spec(spec)
14
    spec.loader.exec_module(module)
15
    return module
16

17

18
def main():
19
    input_a = torch.tensor([1.0, 2.0, 3.0, 4.0], device="cuda")
20
    input_b = torch.tensor([5.0, 6.0, 7.0, 8.0], device="cuda")
21
    expected = torch.tensor([6.0, 8.0, 10.0, 12.0], device="cuda")
22
    actual = torch.empty_like(input_a)
23

24
    load_implementation().solve(input_a, input_b, actual, input_a.numel())
25
    torch.cuda.synchronize()
26
    torch.testing.assert_close(actual, expected)
27
    print("Test passed.")
28

29

30
if __name__ == "__main__":
31
    try:
32
        main()
33
    except Exception as error:
34
        print(f"Test failed: {error}", file=sys.stderr)
35
        raise

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	N=25000000	20	0.817 ms	0.731 ms	1.077 ms	367.031 GB/s
NVIDIA GeForce RTX 4090	PASS	N=25000000	20	0.330 ms	0.329 ms	0.334 ms	909.717 GB/s
NVIDIA GeForce RTX 5090	PASS	N=25000000	20	0.207 ms	0.204 ms	0.216 ms	1450.363 GB/s

References#

LeetGPU Challenge

关注 LeetGPU repository，了解更多内容。

音乐

音乐

Problem Description#

Implementation Requirements#

Example 1#

Example 2#

Constraints#

CUDA#

Approach#

Solution#

Test Code#

Test Result#

Triton#

Approach#

Solution#

Test Code#

Test Result#

PyTorch#

Approach#

Solution#

Test Code#

Test Result#

References#

文章分享

评论区

音乐

文章目录