// SPDX-License-Identifier: BSD-3-Clause
/*
 * streams2d.cu - streams2 with per-array / per-range diagnostic reporting.
 *
 * Same workload as streams2.cu (two streams, cross-stream event ordering)
 * but instead of one "bad" count it reports, per array and per 4 KB block,
 * how many elements mismatched, and prints the device pointers so a host
 * log can be correlated with what the engine read/wrote.
 *
 * Output lines:
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   PTR <name>=0x<hex>            (d_a1 d_b1 d_c1 d_a2 d_b2 d_c2)
 *   DIFF c1 total=<n> blocks=<first block index with mismatches>
 *   DIFF c2 total=<n> blocks=<...>
 *   STREAMS n=<count> bad=<mismatches>
 *   TIMING ms=<elapsed>
 *   STATUS OK|FAIL
 *
 * Build: nvcc -O2 --cudart=static -o streams2d streams2d.cu
 */

#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 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;
}

static void report_diff(const char *name, const float *got,
			const float *exp, int n)
{
	int i, bad = 0, first = -1, last = -1;
	int blk_hits = 0, blk_first = -1;
	int nb = (n + 1023) / 1024;
	int b;

	for (i = 0; i < n; i++) {
		if (got[i] != exp[i]) {
			bad++;
			if (first < 0)
				first = i;
			last = i;
		}
	}
	for (b = 0; b < nb; b++) {
		int j, hit = 0;

		for (j = b * 1024; j < (b + 1) * 1024 && j < n; j++)
			if (got[j] != exp[j]) {
				hit = 1;
				break;
			}
		if (hit) {
			blk_hits++;
			if (blk_first < 0)
				blk_first = b;
		}
	}
	printf("DIFF %s total=%d first=%d last=%d blocks=%d/%d blk_first=%d\n",
	       name, bad, first, last, blk_hits, nb, blk_first);
	if (first >= 0)
		printf("DIFF %s sample got[%d]=%a exp[%d]=%a got[%d]=%a exp[%d]=%a\n",
		       name, first, got[first], first, exp[first], last,
		       got[last], last, exp[last]);
}

int main(void)
{
	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;
	cudaStream_t s1, s2s;
	cudaEvent_t ev_add1, ev_scale2, t0, t1;
	float ms = 0.0f;
	int bad = 0, i;
	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_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);
	if (!h_a1 || !h_b1 || !h_c1 || !h_a2 || !h_b2 || !h_c2)
		return 1;

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

	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));
	printf("PTR d_a1=%p d_b1=%p d_c1=%p d_a2=%p d_b2=%p d_c2=%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));

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

	report_diff("c1", h_c1, h_a1, n);
	{
		float *exp2 = (float *)malloc(bytes);

		for (i = 0; i < n; i++)
			exp2[i] = (h_a2[i] + h_b2[i]) * s2;
		report_diff("c2", h_c2, exp2, n);
		for (i = 0; i < n; i++) {
			if (h_c1[i] != h_a1[i] + h_b1[i])
				bad++;
			if (h_c2[i] != exp2[i])
				bad++;
		}
		free(exp2);
	}

	printf("STREAMS 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");
	return 0;
}