// SPDX-License-Identifier: BSD-3-Clause
/*
 * multikernel.cu - dependent three-kernel sequence over 1M floats.
 *
 * Three different kernels run one after another with a data dependency, each
 * result verified on the host before the next stage starts:
 *
 *   stage 1  add    c[i]   = a[i] + b[i]
 *   stage 2  scale  c[i]   = c[i] * 3
 *   stage 3  saxpy  out[i] = 1.5 * c[i] + d[i]
 *
 * All inputs are exact binary fractions, so the host model is bit-exact and
 * any mismatch is a real device-side divergence.
 *
 * Output lines are stable and greppable:
 *
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   STAGE idx=<n> name=<add|scale|saxpy> bad=<mismatches>
 *   MULTI stages=3 bad=<mismatches>
 *   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 multikernel multikernel.cu
 *
 * Then, from the guest:  /tmp/multikernel
 */

#include <cstdio>
#include <cstdlib>
#include <cstdint>
#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 k_add(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];
}

__global__ void k_scale(float *c, float s, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		c[i] = c[i] * s;
}

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

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

static int verify_stage(const float *got, const float *a, const float *b,
			const float *d, float scale, float alpha, int n,
			int stage)
{
	float *exp = (float *)malloc((size_t)n * sizeof(float));
	int i, bad = 0;
	const char *name = stage == 1 ? "add" : stage == 2 ? "scale" : "saxpy";

	if (!exp)
		return -1;

	/* Host reference model, matching the kernel chain exactly. */
	for (i = 0; i < n; i++) {
		float v = a[i] + b[i];

		if (stage >= 2)
			v = v * scale;
		if (stage >= 3)
			v = alpha * v + d[i];
		exp[i] = v;
	}
	for (i = 0; i < n; i++) {
		if (got[i] != exp[i])
			bad++;
	}
	printf("STAGE idx=%d name=%s bad=%d\n", stage, name, bad);
	free(exp);
	return bad;
}

int main(void)
{
	cudaDeviceProp prop;
	const int n = 1 << 20; /* 1M floats */
	const float scale = 3.0f, alpha = 1.5f;
	size_t bytes = (size_t)n * sizeof(float);
	float *h_a, *h_b, *h_c, *h_d, *h_out;
	float *d_a, *d_b, *d_c, *d_d, *d_out;
	int bad = 0, i, s;
	const int grid = (n + 255) / 256;

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

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

	CHECK(cudaMalloc(&d_a, bytes));
	CHECK(cudaMalloc(&d_b, bytes));
	CHECK(cudaMalloc(&d_c, bytes));
	CHECK(cudaMalloc(&d_d, bytes));
	CHECK(cudaMalloc(&d_out, bytes));
	CHECK(cudaMemcpy(d_a, h_a, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(d_b, h_b, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(d_d, h_d, bytes, cudaMemcpyHostToDevice));

	/* stage 1: add */
	k_add<<<grid, 256>>>(d_a, d_b, d_c, n);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_c, d_c, bytes, cudaMemcpyDeviceToHost));
	s = verify_stage(h_c, h_a, h_b, h_d, scale, alpha, n, 1);
	if (s < 0)
		return 1;
	bad += s;

	/* stage 2: scale (depends on stage 1's c) */
	k_scale<<<grid, 256>>>(d_c, scale, n);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_c, d_c, bytes, cudaMemcpyDeviceToHost));
	s = verify_stage(h_c, h_a, h_b, h_d, scale, alpha, n, 2);
	if (s < 0)
		return 1;
	bad += s;

	/* stage 3: saxpy out = alpha * c + d (depends on stage 2's c) */
	k_saxpy<<<grid, 256>>>(d_c, d_d, d_out, alpha, n);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_out, d_out, bytes, cudaMemcpyDeviceToHost));
	s = verify_stage(h_out, h_a, h_b, h_d, scale, alpha, n, 3);
	if (s < 0)
		return 1;
	bad += s;

	printf("MULTI stages=3 bad=%d\n", bad);

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

	printf("STATUS OK\n");
	CHECK(cudaFree(d_a));
	CHECK(cudaFree(d_b));
	CHECK(cudaFree(d_c));
	CHECK(cudaFree(d_d));
	CHECK(cudaFree(d_out));
	free(h_a);
	free(h_b);
	free(h_c);
	free(h_d);
	free(h_out);
	return 0;
}