// SPDX-License-Identifier: BSD-3-Clause
/*
 * streams2_probe.cu - instrumented streams2 to localize the deterministic
 * failure: per-array mismatch counts, first-bad indices with values, and a
 * GPU-side verify-kernel digest of the device buffers (to separate "device
 * pages never written" from "D2H readback wrong").
 *
 * Modes (argv[1]):
 *   A : exactly streams2's shape (control)
 *   B : like A but the two D2H copy-backs use synchronous cudaMemcpy
 *   C : like A but no events cross the streams; both streams are joined
 *       with cudaStreamSynchronize before any copy-back
 *   D : single stream (everything on s1)
 *
 * Build: nvcc -O2 --cudart=static -o streams2_probe streams2_probe.cu
 */

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

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

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

/* GPU-side digest: counts of matches for c1==a1+b1 and c2==(a2+b2)*2 */
__global__ void digest(const float *a1, const float *b1, const float *c1,
		       const float *a2, const float *b2, const float *c2,
		       int n, unsigned long long *out)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i >= n)
		return;
	if (c1[i] == a1[i] + b1[i])
		atomicAdd(&out[0], 1ull);
	if (c2[i] == (a2[i] + b2[i]) * 2.0f)
		atomicAdd(&out[1], 1ull);
}

static int report(const char *tag, const float *h_c, const float *h_a,
		  const float *h_b, const float *h_a2, const float *h_b2,
		  int n, float scale2)
{
	int bad = 0, i, shown = 0;

	for (i = 0; i < n; i++) {
		float want = h_a[i] + h_b[i];

		if (scale2 != 1.0f)
			want = (h_a2[i] + h_b2[i]) * scale2;
		if (h_c[i] != want) {
			bad++;
			if (shown < 4) {
				printf("FIRSTBAD %s i=%d got=%g want=%g\n",
				       tag, i, h_c[i], want);
				shown++;
			}
		}
	}
	printf("BAD %s n=%d bad=%d\n", tag, n, bad);
	return bad;
}

int main(int argc, char **argv)
{
	const char *mode = argc > 1 ? argv[1] : "A";
	cudaDeviceProp prop;
	const int n = 1 << 15;
	const float s2 = 2.0f;
	size_t bytes = (size_t)n * sizeof(float);
	float *h_a1, *h_b1, *h_c1, *h_a2, *h_b2, *h_c2;
	float *d_a1, *d_b1, *d_c1, *d_a2, *d_b2, *d_c2;
	unsigned long long *d_out, h_out[2];
	cudaStream_t s1, s2s;
	cudaEvent_t ev_add1, ev_scale2, t0, t1;
	float ms = 0.0f;
	int i;
	const int grid = (n + 255) / 256;

	setvbuf(stdout, NULL, _IONBF, 0);

	CHECK(cudaGetDeviceProperties(&prop, 0));
	printf("DEVICE %s mode=%s\n", prop.name, mode);

	h_a1 = (float *)malloc(bytes);
	h_b1 = (float *)malloc(bytes);
	h_c1 = (float *)malloc(bytes);
	h_a2 = (float *)malloc(bytes);
	h_b2 = (float *)malloc(bytes);
	h_c2 = (float *)malloc(bytes);

	for (i = 0; i < n; i++) {
		h_a1[i] = (float)(i % 1000) * 0.5f;
		h_b1[i] = (float)(i % 997) * 0.25f;
		h_a2[i] = (float)(i % 991) * 0.125f;
		h_b2[i] = (float)(i % 983) * 0.0625f;
	}
	memset(h_c1, 0xCD, bytes);
	memset(h_c2, 0xCD, bytes);

	CHECK(cudaMalloc(&d_a1, bytes));
	CHECK(cudaMalloc(&d_b1, bytes));
	CHECK(cudaMalloc(&d_c1, bytes));
	CHECK(cudaMalloc(&d_a2, bytes));
	CHECK(cudaMalloc(&d_b2, bytes));
	CHECK(cudaMalloc(&d_c2, bytes));
	CHECK(cudaMalloc(&d_out, 2 * sizeof(unsigned long long)));
	printf("DPTRS da1=%p db1=%p dc1=%p da2=%p db2=%p dc2=%p\n",
	       (void *)d_a1, (void *)d_b1, (void *)d_c1, (void *)d_a2,
	       (void *)d_b2, (void *)d_c2);

	CHECK(cudaStreamCreate(&s1));
	CHECK(cudaStreamCreate(&s2s));
	CHECK(cudaEventCreate(&ev_add1));
	CHECK(cudaEventCreate(&ev_scale2));
	CHECK(cudaEventCreate(&t0));
	CHECK(cudaEventCreate(&t1));

	CHECK(cudaEventRecord(t0));

	if (!strcmp(mode, "D")) {
		CHECK(cudaMemcpyAsync(d_a1, h_a1, bytes,
				      cudaMemcpyHostToDevice, s1));
		CHECK(cudaMemcpyAsync(d_b1, h_b1, bytes,
				      cudaMemcpyHostToDevice, s1));
		CHECK(cudaMemcpyAsync(d_a2, h_a2, bytes,
				      cudaMemcpyHostToDevice, s1));
		CHECK(cudaMemcpyAsync(d_b2, h_b2, bytes,
				      cudaMemcpyHostToDevice, s1));
		vadd<<<grid, 256, 0, s1>>>(d_a1, d_b1, d_c1, n);
		vadd<<<grid, 256, 0, s1>>>(d_a2, d_b2, d_c2, n);
		CHECK(cudaStreamSynchronize(s1));
		scale<<<grid, 256, 0, s1>>>(d_c2, s2, n);
		CHECK(cudaStreamSynchronize(s1));
		CHECK(cudaMemcpy(h_c1, d_c1, bytes, cudaMemcpyDeviceToHost));
		CHECK(cudaMemcpy(h_c2, d_c2, bytes, cudaMemcpyDeviceToHost));
		goto verify;
	}

	if (!strcmp(mode, "S")) {
		/* staged verification: check device content after every step */
		float *hb1 = (float *)malloc(bytes);
		float *hb2 = (float *)malloc(bytes);
		float *hb3 = (float *)malloc(bytes);

		CHECK(cudaMemcpy(d_a1, h_a1, bytes, cudaMemcpyHostToDevice));
		CHECK(cudaMemcpy(d_b1, h_b1, bytes, cudaMemcpyHostToDevice));
		CHECK(cudaMemcpy(d_a2, h_a2, bytes, cudaMemcpyHostToDevice));
		CHECK(cudaMemcpy(d_b2, h_b2, bytes, cudaMemcpyHostToDevice));
		CHECK(cudaMemcpy(hb1, d_a1, bytes, cudaMemcpyDeviceToHost));
		{ int b = 0; for (i = 0; i < n; i++) if (hb1[i] != h_a1[i]) b++;
		  printf("STAGE-IN da1 bad=%d\n", b); }
		CHECK(cudaMemcpy(hb1, d_b1, bytes, cudaMemcpyDeviceToHost));
		{ int b = 0; for (i = 0; i < n; i++) if (hb1[i] != h_b1[i]) b++;
		  printf("STAGE-IN db1 bad=%d\n", b); }
		CHECK(cudaMemcpy(hb1, d_a2, bytes, cudaMemcpyDeviceToHost));
		{ int b = 0; for (i = 0; i < n; i++) if (hb1[i] != h_a2[i]) b++;
		  printf("STAGE-IN da2 bad=%d\n", b); }
		CHECK(cudaMemcpy(hb1, d_b2, bytes, cudaMemcpyDeviceToHost));
		{ int b = 0; for (i = 0; i < n; i++) if (hb1[i] != h_b2[i]) b++;
		  printf("STAGE-IN db2 bad=%d\n", b); }

		vadd<<<grid, 256, 0, s1>>>(d_a1, d_b1, d_c1, n);
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(hb1, d_c1, bytes, cudaMemcpyDeviceToHost));
		{ int b = 0, first = -1; for (i = 0; i < n; i++)
			if (hb1[i] != h_a1[i] + h_b1[i]) { b++;
				if (first < 0) { first = i;
					printf("S1-FIRSTBAD i=%d got=%g want=%g a2b2=%g\n",
						i, hb1[i], h_a1[i]+h_b1[i], h_a2[i]+h_b2[i]); } }
		  printf("STAGE vadd1->c1 readback bad=%d first=%d\n", b, first); }

		CHECK(cudaDeviceSynchronize());
		vadd<<<grid, 256, 0, s2s>>>(d_a2, d_b2, d_c2, n);
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(hb2, d_c2, bytes, cudaMemcpyDeviceToHost));
		{ int b = 0, first = -1; for (i = 0; i < n; i++)
			if (hb2[i] != h_a2[i] + h_b2[i]) { b++; if (first < 0) first = i; }
		  printf("STAGE vadd2->c2 readback bad=%d first=%d\n", b, first); }

		scale<<<grid, 256, 0, s2s>>>(d_c2, s2, n);
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(hb3, d_c2, bytes, cudaMemcpyDeviceToHost));
		{ int b = 0, first = -1; for (i = 0; i < n; i++)
			if (hb3[i] != (h_a2[i] + h_b2[i]) * 2.0f) { b++; if (first < 0) first = i; }
		  printf("STAGE scale->c2 readback bad=%d first=%d\n", b, first); }
		/* second readback of c1 AFTER everything: does it change? */
		CHECK(cudaMemcpy(hb1, d_c1, bytes, cudaMemcpyDeviceToHost));
		{ int b = 0, first = -1; for (i = 0; i < n; i++)
			if (hb1[i] != h_a1[i] + h_b1[i]) { b++; if (first < 0) first = i; }
		  printf("STAGE c1-late readback bad=%d first=%d\n", b, first); }
		printf("STATUS OK\n");
		return 0;
	}

	if (!strcmp(mode, "C")) {
		CHECK(cudaMemcpyAsync(d_a1, h_a1, bytes,
				      cudaMemcpyHostToDevice, s1));
		CHECK(cudaMemcpyAsync(d_b1, h_b1, bytes,
				      cudaMemcpyHostToDevice, s1));
		CHECK(cudaMemcpyAsync(d_a2, h_a2, bytes,
				      cudaMemcpyHostToDevice, s2s));
		CHECK(cudaMemcpyAsync(d_b2, h_b2, bytes,
				      cudaMemcpyHostToDevice, s2s));
		vadd<<<grid, 256, 0, s1>>>(d_a1, d_b1, d_c1, n);
		vadd<<<grid, 256, 0, s2s>>>(d_a2, d_b2, d_c2, n);
		CHECK(cudaStreamSynchronize(s1));
		CHECK(cudaStreamSynchronize(s2s));
		scale<<<grid, 256, 0, s2s>>>(d_c2, s2, n);
		CHECK(cudaStreamSynchronize(s2s));
		CHECK(cudaMemcpyAsync(h_c1, d_c1, bytes,
				      cudaMemcpyDeviceToHost, s1));
		CHECK(cudaMemcpyAsync(h_c2, d_c2, bytes,
				      cudaMemcpyDeviceToHost, s2s));
		CHECK(cudaStreamSynchronize(s1));
		CHECK(cudaStreamSynchronize(s2s));
		goto verify;
	}

	/* A and B: original shape */
	CHECK(cudaMemcpyAsync(d_a1, h_a1, bytes, cudaMemcpyHostToDevice, s1));
	CHECK(cudaMemcpyAsync(d_b1, h_b1, bytes, cudaMemcpyHostToDevice, s1));
	CHECK(cudaMemcpyAsync(d_a2, h_a2, bytes, cudaMemcpyHostToDevice, s2s));
	CHECK(cudaMemcpyAsync(d_b2, h_b2, bytes, cudaMemcpyHostToDevice, s2s));
	vadd<<<grid, 256, 0, s1>>>(d_a1, d_b1, d_c1, n);
	vadd<<<grid, 256, 0, s2s>>>(d_a2, d_b2, d_c2, n);
	CHECK(cudaGetLastError());

	CHECK(cudaEventRecord(ev_add1, s1));
	CHECK(cudaStreamWaitEvent(s2s, ev_add1, 0));
	scale<<<grid, 256, 0, s2s>>>(d_c2, s2, n);
	CHECK(cudaGetLastError());

	CHECK(cudaEventRecord(ev_scale2, s2s));
	CHECK(cudaStreamWaitEvent(s1, ev_scale2, 0));
	if (!strcmp(mode, "B")) {
		CHECK(cudaStreamSynchronize(s2s));
		CHECK(cudaMemcpy(h_c1, d_c1, bytes, cudaMemcpyDeviceToHost));
		CHECK(cudaMemcpy(h_c2, d_c2, bytes, cudaMemcpyDeviceToHost));
	} else {
		CHECK(cudaMemcpyAsync(h_c1, d_c1, bytes,
				      cudaMemcpyDeviceToHost, s1));
		CHECK(cudaMemcpyAsync(h_c2, d_c2, bytes,
				      cudaMemcpyDeviceToHost, s2s));
	}
	CHECK(cudaEventRecord(t1, s1));

	CHECK(cudaStreamSynchronize(s1));
	CHECK(cudaStreamSynchronize(s2s));
	CHECK(cudaEventElapsedTime(&ms, t0, t1));

verify:
	{
		int bad1, bad2;

		bad1 = report("C1", h_c1, h_a1, h_b1, h_a2, h_b2, n, 1.0f);
		bad2 = report("C2", h_c2, h_a2, h_b2, h_a2, h_b2, n, s2);
		printf("STREAMS n=%d bad=%d\n", n, bad1 + bad2);
		printf("TIMING ms=%.3f\n", ms);

		/* GPU-side digest of the device buffers */
		CHECK(cudaMemset(d_out, 0, 2 * sizeof(unsigned long long)));
		digest<<<grid, 256>>>(d_a1, d_b1, d_c1, d_a2, d_b2, d_c2, n,
				      d_out);
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(h_out, d_out, 2 * sizeof(unsigned long long),
				 cudaMemcpyDeviceToHost));
		printf("GPU-DIGEST c1_ok=%llu c2_ok=%llu\n", h_out[0],
		       h_out[1]);
		printf("STATUS %s\n", (bad1 | bad2) ? "FAIL" : "OK");
	}
	return 0;
}