#include <cuda_runtime.h>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <vector>

inline void cuda_check_impl(cudaError_t status, const char* statement,
                            const char* file, int line) {
    if (status != cudaSuccess) {
        fprintf(stderr, "CUDA error in %s:%d (%s): %s: %s\n", file, line,
                statement, cudaGetErrorName(status), cudaGetErrorString(status));
        exit(EXIT_FAILURE);
    }
}
#define CUDA_CHECK(statement) cuda_check_impl((statement), #statement, __FILE__, __LINE__)

static constexpr int EL_PER_THREAD = 4;

__global__ void saxpy_v2(int n, float alpha,
                         const float* x, float* y) {
    // Base index for this thread; stride by blockDim.x so consecutive threads
    // always access consecutive addresses (coalescing preserved).
    int base   = blockIdx.x * blockDim.x * EL_PER_THREAD + threadIdx.x;
    int stride = blockDim.x;

    #pragma unroll
    for (int k = 0; k < EL_PER_THREAD; ++k) {
        int i = base + k * stride;
        if (i < n)
            y[i] = alpha * x[i] + y[i];
    }
}

void launch_saxpy(int n, float alpha, float* x, float* y) {
    constexpr int BLOCK = 256;
    // Each block covers BLOCK*EL_PER_THREAD elements, so shrink the grid accordingly.
    const int grid = (n + BLOCK * EL_PER_THREAD - 1) / (BLOCK * EL_PER_THREAD);
    saxpy_v2<<<grid, BLOCK>>>(n, alpha, x, y);
}

int main() {
    const int N     = 1 << 25;   // 32 M elements
    const float A   = 2.0f;
    const size_t sz = N * sizeof(float);

    std::vector<float> hX(N), hY(N);
    for (int i = 0; i < N; i++) {
        hX[i] = 0.001f * i;
        hY[i] = 1.0f;
    }

    float *dX, *dY;
    CUDA_CHECK(cudaMalloc(&dX, sz));
    CUDA_CHECK(cudaMalloc(&dY, sz));
    CUDA_CHECK(cudaMemcpy(dX, hX.data(), sz, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dY, hY.data(), sz, cudaMemcpyHostToDevice));

    launch_saxpy(N, A, dX, dY);
    CUDA_CHECK(cudaGetLastError());      // errors from the launch itself
    CUDA_CHECK(cudaDeviceSynchronize()); // errors from kernel execution

    CUDA_CHECK(cudaMemcpy(hY.data(), dY, sz, cudaMemcpyDeviceToHost));
    for (int i : {0, 1, 12345, N - 1})
        printf("y[%8d] = %g\n", i, (double)hY[i]);

    CUDA_CHECK(cudaFree(dX));
    CUDA_CHECK(cudaFree(dY));
    return 0;
}