untitled

#include <cuda_runtime.h>

__global__ void vector_add(const float4* __restrict__ A4, const float4* __restrict__ B4, float4* __restrict__ C4, int N4) {
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    if (idx < N4) {
        float4 a = A4[idx];
        float4 b = B4[idx];
        C4[idx] = make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
    }
}

__global__ void tail_add(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int start_offset, int N) {
    int idx = start_offset + blockDim.x * blockIdx.x + threadIdx.x;
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int N) {
    int N4 = N >> 2;
    int tail = N & 3;

    if (N4 > 0) {
        const float4* A4 = reinterpret_cast<const float4*>(A);
        const float4* B4 = reinterpret_cast<const float4*>(B);
        float4*       C4 = reinterpret_cast<float4*>(C);
        int threadsPerBlock = 256;
        int blocksPerGrid = (N4 + threadsPerBlock - 1) / threadsPerBlock;
        vector_add<<<blocksPerGrid, threadsPerBlock>>>(A4, B4, C4, N4);
    }

    if (tail) {
        int start = N & ~3;
        tail_add<<<1, 32>>>(A, B, C, start, N);
    }

    // vector_add<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, N);
    cudaDeviceSynchronize();
}