// SPDX-License-Identifier: BSD-3-Clause
/*
 * wedge_recover.cu - reproduce the per-launch no-execute (wedge) and test
 * whether an idle gap heals a wedged context.
 *
 * v2 (session 10, r416 finding): the v1 loop reused one deterministic
 * kernel+buffers, so a no-write launch was MASKED - the stale buffer held
 * exactly the previous iteration's (identical) values and the verify passed.
 * r416 ran 60/60 "OK" while 58 consecutive launches no-wrote underneath.
 *
 * v2 mutates the inputs every iteration (H2D of fresh ha/hb), so a stale
 * readback cannot match the new expected values: a no-write launch must
 * report BAD. On BAD the same iteration is retried after a sleep, and the
 * readback is classified: STALE-PREV (byte-identical to the previous
 * iteration's output = launch never executed) vs PARTIAL/OTHER (engine wrote
 * something). Latency and per-iteration outcome are printed.
 *
 * Build: nvcc -O2 --cudart=static -o wedge_recover wedge_recover.cu
 */

#include <cstdio>
#include <cstdlib>
#include <cuda_runtime.h>
#include <unistd.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];
}

int main(int argc, char **argv)
{
	const int n = 1 << 15;
	size_t bytes = (size_t)n * sizeof(float);
	float *ha, *hb, *hc, *hc_prev;
	float *da, *db, *dc;
	int iters = argc > 1 ? atoi(argv[1]) : 40;
	int sleep_s = argc > 2 ? atoi(argv[2]) : 20;
	int fails = 0, heals = 0, i;

	setvbuf(stdout, NULL, _IONBF, 0);
	ha = (float *)malloc(bytes);
	hb = (float *)malloc(bytes);
	hc = (float *)malloc(bytes);
	hc_prev = (float *)malloc(bytes);
	if (!ha || !hb || !hc || !hc_prev)
		return 1;
	/* iteration 0 baseline; mutated each iteration below (v2) */
	for (i = 0; i < n; i++) {
		ha[i] = (float)(i % 977) * 0.25f;
		hb[i] = (float)(i % 251) * 0.5f;
	}
	memset(hc_prev, 0, bytes);
	CHECK(cudaMalloc(&da, bytes));
	CHECK(cudaMalloc(&db, bytes));
	CHECK(cudaMalloc(&dc, bytes));
	CHECK(cudaMemcpy(da, ha, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(db, hb, bytes, cudaMemcpyHostToDevice));

	for (i = 0; i < iters; i++) {
		int attempt = 0;
		long long ms[2] = { 0, 0 };

		/*
		 * v2: fresh inputs every iteration so a stale (never
		 * executed) launch cannot match. Mutate after iteration 0.
		 */
		if (i > 0) {
			ha[i % n] += 1.0f;
			hb[(i * 7 + 3) % n] += 2.0f;
			CHECK(cudaMemcpy(da, ha, bytes,
					 cudaMemcpyHostToDevice));
			CHECK(cudaMemcpy(db, hb, bytes,
					 cudaMemcpyHostToDevice));
		}

		for (;;) {
			int bad = 0, stale_prev = 1, j;
			cudaEvent_t e0, e1;
			float ems = 0;

			cudaEventCreate(&e0);
			cudaEventCreate(&e1);
			cudaEventRecord(e0);
			vadd<<<(n + 255) / 256, 256>>>(da, db, dc, n);
			CHECK(cudaGetLastError());
			cudaEventRecord(e1);
			CHECK(cudaDeviceSynchronize());
			cudaEventElapsedTime(&ems, e0, e1);
			cudaEventDestroy(e0);
			cudaEventDestroy(e1);
			CHECK(cudaMemcpy(hc, dc, bytes, cudaMemcpyDeviceToHost));
			for (j = 0; j < n; j++) {
				float expect = ha[j] + hb[j];

				if (hc[j] != expect) {
					bad++;
					if (hc[j] != hc_prev[j])
						stale_prev = 0;
				}
			}
			ms[attempt & 1] = (long long)(ems * 1000.0f);
			if (bad == 0) {
				if (attempt > 0)
					heals++;
				printf("ITER %d OK attempt=%d launch_us=%lld\n",
				       i, attempt, ms[attempt & 1]);
				memcpy(hc_prev, hc, bytes);
				break;
			}
			fails++;
			printf("ITER %d BAD bad=%d/%d attempt=%d launch_us=%lld class=%s\n",
			       i, bad, n, attempt, ms[attempt & 1],
			       stale_prev ? "STALE-PREV(NOWRITE)"
					  : "OTHER/PARTIAL");
			if (attempt >= 3) {
				printf("ITER %d GIVEUP\n", i);
				break;
			}
			attempt++;
			printf("ITER %d SLEEP %ds before retry\n", i, sleep_s);
			sleep(sleep_s);
		}
	}
	printf("SUMMARY iters=%d fails=%d heals=%d sleep_s=%d\n",
	       iters, fails, heals, sleep_s);
	printf("STATUS %s\n", fails == 0 ? "OK" : "FAIL");
	return fails == 0 ? 0 : 1;
}