// SPDX-License-Identifier: BSD-3-Clause
/*
 * victory_streams_events.cu - 3 streams, overlapped copies, event ordering.
 *
 * Three non-default streams each copy their own inputs, run a vectorAdd, and
 * are then chained with events:
 *
 *   stream 0: vadd c0 = a0 + b0
 *   stream 1: waits event e0 from stream 0, then c1 += c0
 *   stream 2: waits event e1 from stream 1, then c2 += c1
 *
 * Every stream also does its own async D2H copy; the whole chain is timed
 * with cudaEventRecord/cudaEventElapsedTime. The host verifies all three
 * outputs; inputs are multiples of 0.25 so the arithmetic is exact in float.
 *
 * Output lines are stable and greppable:
 *
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   STREAMS3 n=<count> bad=<mismatches>
 *   TIMING ms=<elapsed>
 *   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_streams_events victory_streams_events.cu
 *
 * Then, from the guest:  /tmp/victory_streams_events
 */

#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)

#define NSTREAMS 3

__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];
}

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

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

int main(void)
{
	cudaDeviceProp prop;
	const int n = 1 << 16; /* 256 KB per stream */
	size_t bytes = (size_t)n * sizeof(float);
	float *h_a[NSTREAMS], *h_b[NSTREAMS], *h_c[NSTREAMS];
	float *d_a[NSTREAMS], *d_b[NSTREAMS], *d_c[NSTREAMS];
	cudaStream_t st[NSTREAMS];
	cudaEvent_t ev[NSTREAMS], t0, t1;
	float ms = 0.0f;
	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);

	for (s = 0; s < NSTREAMS; s++) {
		h_a[s] = (float *)malloc(bytes);
		h_b[s] = (float *)malloc(bytes);
		h_c[s] = (float *)malloc(bytes);
		if (!h_a[s] || !h_b[s] || !h_c[s])
			return 1;
		for (i = 0; i < n; i++) {
			h_a[s][i] = (float)((i + s) % 1000) * 0.5f;
			h_b[s][i] = (float)((i + 2 * s) % 997) * 0.25f;
		}
		CHECK(cudaMalloc(&d_a[s], bytes));
		CHECK(cudaMalloc(&d_b[s], bytes));
		CHECK(cudaMalloc(&d_c[s], bytes));
		CHECK(cudaStreamCreate(&st[s]));
		CHECK(cudaEventCreate(&ev[s]));
	}
	CHECK(cudaEventCreate(&t0));
	CHECK(cudaEventCreate(&t1));

	CHECK(cudaEventRecord(t0));

	for (s = 0; s < NSTREAMS; s++) {
		CHECK(cudaMemcpyAsync(d_a[s], h_a[s], bytes,
				      cudaMemcpyHostToDevice, st[s]));
		CHECK(cudaMemcpyAsync(d_b[s], h_b[s], bytes,
				      cudaMemcpyHostToDevice, st[s]));
		vadd<<<grid, 256, 0, st[s]>>>(d_a[s], d_b[s], d_c[s], n);
		CHECK(cudaGetLastError());
		CHECK(cudaEventRecord(ev[s], st[s]));
	}

	/* Cross-stream chain: 1 waits on 0, 2 waits on 1. */
	CHECK(cudaStreamWaitEvent(st[1], ev[0], 0));
	vadd_into<<<grid, 256, 0, st[1]>>>(d_c[1], d_c[0], n);
	CHECK(cudaGetLastError());
	CHECK(cudaEventRecord(ev[1], st[1]));

	CHECK(cudaStreamWaitEvent(st[2], ev[1], 0));
	vadd_into<<<grid, 256, 0, st[2]>>>(d_c[2], d_c[1], n);
	CHECK(cudaGetLastError());

	/* Copies back overlap across the three streams. */
	for (s = 0; s < NSTREAMS; s++)
		CHECK(cudaMemcpyAsync(h_c[s], d_c[s], bytes,
				      cudaMemcpyDeviceToHost, st[s]));
	CHECK(cudaEventRecord(t1, st[2]));

	for (s = 0; s < NSTREAMS; s++)
		CHECK(cudaStreamSynchronize(st[s]));
	CHECK(cudaEventElapsedTime(&ms, t0, t1));

	/* c0 = a0+b0; c1 = (a1+b1)+c0; c2 = (a2+b2)+c1 (all exact floats). */
	for (i = 0; i < n; i++) {
		float c0 = h_a[0][i] + h_b[0][i];
		float c1 = h_a[1][i] + h_b[1][i] + c0;
		float c2 = h_a[2][i] + h_b[2][i] + c1;

		if (h_c[0][i] != c0)
			bad++;
		if (h_c[1][i] != c1)
			bad++;
		if (h_c[2][i] != c2)
			bad++;
	}

	printf("STREAMS3 n=%d bad=%d\n", n, bad);
	printf("TIMING ms=%.3f\n", ms);

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

	printf("STATUS OK\n");
	for (s = 0; s < NSTREAMS; s++) {
		CHECK(cudaEventDestroy(ev[s]));
		CHECK(cudaStreamDestroy(st[s]));
		CHECK(cudaFree(d_a[s]));
		CHECK(cudaFree(d_b[s]));
		CHECK(cudaFree(d_c[s]));
		free(h_a[s]);
		free(h_b[s]);
		free(h_c[s]);
	}
	CHECK(cudaEventDestroy(t0));
	CHECK(cudaEventDestroy(t1));
	return 0;
}