// SPDX-License-Identifier: BSD-3-Clause
/*
 * vector_add.cu - CUDA A/B workload for the Open-IOV harness.
 *
 * One deterministic, machine-readable workload that can run on the host (bare
 * metal) and in a guest (emulated vGPU). The output lines are stable so the
 * two runs can be diffed:
 *
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   RESULT n=<count> bad=<mismatches> checksum=<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.
 */

#include <cstdio>
#include <cstdlib>
#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 vadd(const float *a, const float *b, float *c, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		c[i] = a[i] + b[i];
}

int main(void)
{
	cudaDeviceProp prop;
	const int n = 1 << 20;
	size_t bytes = (size_t)n * sizeof(float);
	float *h_a, *h_b, *h_c;
	float *d_a, *d_b, *d_c;
	double sum = 0.0;
	int bad = 0, i;

	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_a = (float *)malloc(bytes);
	h_b = (float *)malloc(bytes);
	h_c = (float *)malloc(bytes);
	if (!h_a || !h_b || !h_c)
		return 1;

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

	memset(h_c, 0xCC, bytes);	/* marker: reveal unwritten spans */
	CHECK(cudaMalloc(&d_a, bytes));
	CHECK(cudaMalloc(&d_b, bytes));
	CHECK(cudaMalloc(&d_c, bytes));
	CHECK(cudaMemcpy(d_a, h_a, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(d_b, h_b, bytes, cudaMemcpyHostToDevice));

	vadd<<<(n + 255) / 256, 256>>>(d_a, d_b, d_c, n);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_c, d_c, bytes, cudaMemcpyDeviceToHost));

	{
		int first = -1, last = -1, cb = 0, k;

		for (i = 0; i < n; i++) {
			if (h_c[i] != h_a[i] + h_b[i]) {
				bad++;
				if (first < 0)
					first = i;
				last = i;
			}
			sum += h_c[i];
		}
		printf("VADD first=%d last=%d\n", first, last);
		for (k = 0; k < 8 && k * 131072 < n; k++) {
			int lo = k * 131072, hi = lo + 131072, j, b2 = 0;

			if (hi > n)
				hi = n;
			for (j = lo; j < hi; j++)
				if (h_c[j] != h_a[j] + h_b[j])
					b2++;
			cb += b2;
			printf("VADD chunk%d [%d..%d) bad=%d%s\n", k, lo, hi, b2,
			       b2 ? "  <== BAD" : "");
		}
		{
			int mrun = 0, mfirst = -1, mlast = -1, j;

			for (j = 0; j < n; j++) {
				unsigned char b0 = ((unsigned char *)h_c)[j * 4];

				if (b0 == 0xCC) {
					mrun++;
					if (mfirst < 0)
						mfirst = j;
					mlast = j;
				}
			}
			printf("VADD marker-untouched=%d first=%d last=%d\n",
			       mrun, mfirst, mlast);
		}
		if (first >= 0) {
			printf("VADD at first: a=%.4f b=%.4f c=%.4f exp=%.4f\n",
			       h_a[first], h_b[first], h_c[first],
			       h_a[first] + h_b[first]);
			printf("VADD near %d: a=%.4f b=%.4f c=%.4f exp=%.4f\n",
			       (first / 1024) * 1024, h_a[(first / 1024) * 1024],
			       h_b[(first / 1024) * 1024], h_c[(first / 1024) * 1024],
			       h_a[(first / 1024) * 1024] + h_b[(first / 1024) * 1024]);
		}
	}

	printf("RESULT n=%d bad=%d checksum=%.6f\n", n, bad, sum);
	printf("STATUS %s\n", bad == 0 ? "OK" : "FAIL");

	CHECK(cudaFree(d_a));
	CHECK(cudaFree(d_b));
	CHECK(cudaFree(d_c));
	free(h_a);
	free(h_b);
	free(h_c);
	return bad == 0 ? 0 : 1;
}