// SPDX-License-Identifier: BSD-3-Clause
/*
 * p2_streams - created streams, two kernels per stream, independent buffers,
 * per-stream events.
 *
 *   phase 1  three streams, each: scale (x2) then addv (+1), dependent pair
 *   phase 2  four streams, each: two *independent* fill kernels on separate
 *            buffers - up to eight kernels in flight simultaneously; all
 *            results verified and every stream event required complete
 *
 * Expected value phase 1 per element: ((i % 100) * 2.0f) + 1.0f (exact in
 * fp32). Phase 2 values are small integers (exact in int32).
 * Device memory: phase 1 3 x 4 KB, phase 2 8 x 4 KB (peak 32 KB, 8 pages).
 */
#include "common.h"

#define SN 1024
#define NS 3
#define IN_NS 4
#define IN_N 1024

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

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

__global__ void addv(float *x, float v, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		x[i] = x[i] + v;
}

__global__ void fillk(int *x, int base, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		x[i] = base * 1000 + i;
}

static void fill(float *h, int n)
{
	for (int i = 0; i < n; i++)
		h[i] = (float)(i % 100);
}

static float expected1(int i)
{
	return ((float)(i % 100) * 2.0f) + 1.0f;
}

/* Phase 2: independent kernels submitted across four streams. No kernel
 * depends on any other, so the driver is free to overlap all eight; every
 * buffer is verified element-exactly and every stream event must be complete
 * after the stream is synchronized. */
static int indep_phase(void)
{
	int *got = (int *)malloc(sizeof(int) * IN_N);
	cudaStream_t st[IN_NS];
	cudaEvent_t ev[IN_NS];
	int *d_a[IN_NS], *d_b[IN_NS];
	int vbad = 0, ebad = 0;

	if (!got)
		return 1;

	for (int k = 0; k < IN_NS; k++) {
		P2_CHECK(cudaStreamCreate(&st[k]));
		P2_CHECK(cudaEventCreate(&ev[k]));
		P2_CHECK(cudaMalloc(&d_a[k], sizeof(int) * IN_N));
		P2_CHECK(cudaMalloc(&d_b[k], sizeof(int) * IN_N));
		fillk<<<(IN_N + 255) / 256, 256, 0, st[k]>>>(d_a[k], k + 2,
							     IN_N);
		fillk<<<(IN_N + 255) / 256, 256, 0, st[k]>>>(d_b[k], k + 9,
							     IN_N);
		P2_CHECK(cudaGetLastError());
		P2_CHECK(cudaEventRecord(ev[k], st[k]));
	}

	for (int k = 0; k < IN_NS; k++) {
		P2_CHECK(cudaStreamSynchronize(st[k]));
		P2_CHECK(cudaMemcpy(got, d_a[k], sizeof(int) * IN_N,
				    cudaMemcpyDeviceToHost));
		for (int i = 0; i < IN_N; i++)
			if (got[i] != (k + 2) * 1000 + i)
				vbad++;
		P2_CHECK(cudaMemcpy(got, d_b[k], sizeof(int) * IN_N,
				    cudaMemcpyDeviceToHost));
		for (int i = 0; i < IN_N; i++)
			if (got[i] != (k + 9) * 1000 + i)
				vbad++;
		if (cudaEventQuery(ev[k]) != cudaSuccess)
			ebad++;
		cudaEventDestroy(ev[k]);
		cudaStreamDestroy(st[k]);
		cudaFree(d_a[k]);
		cudaFree(d_b[k]);
	}

	printf("RESULT streams indep nstreams=%d launches=%d n=%d bad=%d\n",
	       IN_NS, IN_NS * 2, IN_N, vbad);
	printf("RESULT streams indep_events=%d bad=%d\n", IN_NS, ebad);
	free(got);
	return vbad + ebad;
}

int main(int argc, char **argv)
{
	float *h = (float *)malloc(sizeof(float) * SN);
	float *got = (float *)malloc(sizeof(float) * SN);
	int bad = 0;

	if (!h || !got)
		return 1;
	fill(h, SN);

	if (p2_emit_mode(argc, argv)) {
		printf("RESULT streams nstreams=%d launches=%d n=%d bad=0\n",
		       NS, NS * 2, SN);
		printf("RESULT streams events=%d bad=0\n", NS);
		printf("RESULT streams indep nstreams=%d launches=%d n=%d bad=0\n",
		       IN_NS, IN_NS * 2, IN_N);
		printf("RESULT streams indep_events=%d bad=0\n", IN_NS);
		free(h); free(got);
		return p2_status(0);
	}

	p2_device_info();

	cudaStream_t st[NS];
	cudaEvent_t ev[NS];
	float *d_x[NS];

	for (int k = 0; k < NS; k++) {
		P2_CHECK(cudaStreamCreate(&st[k]));
		P2_CHECK(cudaEventCreate(&ev[k]));
		P2_CHECK(cudaMalloc(&d_x[k], sizeof(float) * SN));
		P2_CHECK(cudaMemcpyAsync(d_x[k], h, sizeof(float) * SN,
					 cudaMemcpyHostToDevice, st[k]));
		scale<<<(SN + 255) / 256, 256, 0, st[k]>>>(d_x[k], 2.0f, SN);
		addv<<<(SN + 255) / 256, 256, 0, st[k]>>>(d_x[k], 1.0f, SN);
		P2_CHECK(cudaGetLastError());
		P2_CHECK(cudaEventRecord(ev[k], st[k]));
	}

	int ibad = 0;

	for (int k = 0; k < NS; k++) {
		cudaError_t q;
		int sbad = 0;

		P2_CHECK(cudaStreamSynchronize(st[k]));
		P2_CHECK(cudaMemcpy(got, d_x[k], sizeof(float) * SN,
				    cudaMemcpyDeviceToHost));
		for (int i = 0; i < SN; i++)
			if (got[i] != expected1(i))
				sbad++;
		q = cudaEventQuery(ev[k]);
		if (q != cudaSuccess)
			sbad++;
		printf("RESULT streams stream=%d query_rc=%d bad=%d\n", k, (int)q,
		       sbad);
		ibad += sbad;
	}
	printf("RESULT streams nstreams=%d launches=%d n=%d bad=%d\n",
	       NS, NS * 2, SN, ibad);
	printf("RESULT streams events=%d bad=%d\n", NS, ibad == 0 ? 0 : 1);
	bad = ibad;

	for (int k = 0; k < NS; k++) {
		cudaEventDestroy(ev[k]);
		cudaStreamDestroy(st[k]);
		cudaFree(d_x[k]);
	}

	bad += indep_phase();

	free(h); free(got);
	return p2_status(bad);
}