// SPDX-License-Identifier: BSD-3-Clause
/*
 * wedge_ab.cu - discriminate the wedge trigger: is it the second launch of
 * the SAME program, or the second launch of the context regardless of
 * program?
 *
 * Sequence (each launch verified against mutated inputs):
 *   1. vadd  (program A)
 *   2. vsub  (program B, different code)
 *   3. vadd  (program A again)
 *   4. vsub  (program B again)
 *
 * If launch 2 (B) executes -> the trigger is same-program reuse.
 * If launch 2 (B) wedges -> the trigger is the second launch in the context.
 *
 * Build: nvcc -O2 --cudart=static -o wedge_ab wedge_ab.cu
 */

#include <cstdio>
#include <cstdlib>
#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 vsub(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(void)
{
	const int n = 1 << 15;
	size_t bytes = (size_t)n * sizeof(float);
	float *ha, *hb, *hc;
	float *da, *db, *dc;
	int i, bad, step;

	setvbuf(stdout, NULL, _IONBF, 0);
	ha = (float *)malloc(bytes);
	hb = (float *)malloc(bytes);
	hc = (float *)malloc(bytes);
	if (!ha || !hb || !hc)
		return 1;
	for (i = 0; i < n; i++) {
		ha[i] = (float)(i % 977) * 0.25f;
		hb[i] = (float)(i % 251) * 0.5f;
	}
	CHECK(cudaMalloc(&da, bytes));
	CHECK(cudaMalloc(&db, bytes));
	CHECK(cudaMalloc(&dc, bytes));

	for (step = 0; step < 4; step++) {
		const char *which = (step & 1) ? "vsub" : "vadd";
		int is_sub = step & 1;

		if (step > 0) {
			ha[step] += 1.0f;
			hb[step * 7 + 3] += 2.0f;
		}
		CHECK(cudaMemcpy(da, ha, bytes, cudaMemcpyHostToDevice));
		CHECK(cudaMemcpy(db, hb, bytes, cudaMemcpyHostToDevice));
		if (is_sub)
			vsub<<<(n + 255) / 256, 256>>>(da, db, dc, n);
		else
			vadd<<<(n + 255) / 256, 256>>>(da, db, dc, n);
		CHECK(cudaGetLastError());
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(hc, dc, bytes, cudaMemcpyDeviceToHost));
		bad = 0;
		for (i = 0; i < n; i++) {
			float expect = is_sub ? ha[i] - hb[i] : ha[i] + hb[i];

			if (hc[i] != expect)
				bad++;
		}
		printf("STEP %d %s bad=%d/%d\n", step, which, bad, n);
	}
	printf("STATUS OK\n");
	return 0;
}