// SPDX-License-Identifier: BSD-3-Clause
/*
 * victory_saxpy.cu - SAXPY (y = a*x + y) over 1M floats, two alphas.
 *
 * Two back-to-back launches over a 4 MB buffer:
 *
 *   y = 3.0  * x + y
 *   y = 0.25 * x + y
 *
 * The host replays the same two steps from the same inputs and verifies the
 * device result element by element with a small tolerance (the device may
 * contract multiply-add to FMA; the host may not).
 *
 * Output lines are stable and greppable:
 *
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   SAXPY n=<count> launches=2 bad=<mismatches> csum=<sum>
 *   STATUS OK|FAIL
 *
 * Build with a static CUDA runtime (--cudart=static) so the guest only needs
 * the driver's libcuda.so, not a CUDA toolkit installation:
 *
 *   nvcc -O2 --cudart=static -o victory_saxpy victory_saxpy.cu
 *
 * Then, from the guest:  /tmp/victory_saxpy
 */

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

#define CHECK(x)                                                         \
	do {                                                             \
		cudaError_t e_ = (x);                                    \
		if (e_ != cudaSuccess) {                                 \
			printf("CUDA_ERROR %s at line %d\n",             \
			       cudaGetErrorString(e_), __LINE__);        \
			printf("STATUS FAIL\n");                         \
			return 1;                                        \
		}                                                        \
	} while (0)

__global__ void saxpy(const float *x, float *y, float alpha, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		y[i] = alpha * x[i] + y[i];
}

int main(void)
{
	cudaDeviceProp prop;
	const int n = 1 << 20; /* 1M floats = 4 MB */
	const float alpha1 = 3.0f, alpha2 = 0.25f;
	size_t bytes = (size_t)n * sizeof(float);
	float *h_x, *h_y0, *h_ref, *h_y;
	float *d_x, *d_y;
	double csum = 0.0;
	int bad = 0, i;

	setvbuf(stdout, NULL, _IONBF, 0);

	CHECK(cudaGetDeviceProperties(&prop, 0));
	printf("DEVICE %s cc=%d.%d sm=%d mem=%llu\n", prop.name,
	       prop.major, prop.minor, prop.multiProcessorCount,
	       (unsigned long long)prop.totalGlobalMem);

	h_x = (float *)malloc(bytes);
	h_y0 = (float *)malloc(bytes);
	h_ref = (float *)malloc(bytes);
	h_y = (float *)malloc(bytes);
	if (!h_x || !h_y0 || !h_ref || !h_y)
		return 1;

	for (i = 0; i < n; i++) {
		h_x[i] = (float)(i % 1000) * 0.5f;
		h_y0[i] = (float)(i % 997) * 0.25f;
		h_ref[i] = h_y0[i];
	}

	/* Host reference: the same two SAXPY steps, in the same order. */
	for (i = 0; i < n; i++)
		h_ref[i] = alpha1 * h_x[i] + h_ref[i];
	for (i = 0; i < n; i++)
		h_ref[i] = alpha2 * h_x[i] + h_ref[i];

	CHECK(cudaMalloc(&d_x, bytes));
	CHECK(cudaMalloc(&d_y, bytes));
	CHECK(cudaMemcpy(d_x, h_x, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(d_y, h_y0, bytes, cudaMemcpyHostToDevice));

	saxpy<<<(n + 255) / 256, 256>>>(d_x, d_y, alpha1, n);
	saxpy<<<(n + 255) / 256, 256>>>(d_x, d_y, alpha2, n);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_y, d_y, bytes, cudaMemcpyDeviceToHost));

	for (i = 0; i < n; i++) {
		float tol = 1e-3f + 1e-5f * fabsf(h_ref[i]);

		if (fabsf(h_y[i] - h_ref[i]) > tol)
			bad++;
		csum += (double)h_y[i];
	}

	printf("SAXPY n=%d launches=2 bad=%d csum=%.6f\n", n, bad, csum);

	if (bad != 0) {
		printf("STATUS FAIL\n");
		return 1;
	}

	printf("STATUS OK\n");
	CHECK(cudaFree(d_x));
	CHECK(cudaFree(d_y));
	free(h_x);
	free(h_y0);
	free(h_ref);
	free(h_y);
	return 0;
}