LeetGPU | 01_Vector_Addition

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
}

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
}

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)

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

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