// SPDX-License-Identifier: BSD-3-Clause
/*
 * scale_probe.cu - localize the deterministic "in-place scale writes only the
 * first 4 KB page" failure.
 *
 * Modes (argv[1]):
 *   V : vadd(a,b->c); sync; D2H+verify; vadd again; D2H+verify
 *   I : vadd(a,b->c); sync; D2H+verify; scale(c,*2 in place); sync;
 *       D2H+verify (with page map); scale again; sync; D2H+verify; D2H again
 *   O : vadd(a,b->c); sync; D2H+verify; scale_out(c->d,*2); sync;
 *       D2H+verify (with page map)
 *
 * Every stage also runs a GPU-side verify kernel that counts matching
 * elements reading the device buffers (engine's own view) and writes the
 * counts to a small buffer, D2H'd afterwards: separates "device pages never
 * written" (GPU count low) from "D2H readback stale" (GPU count high, host
 * count low).
 *
 * Build: nvcc -O2 --cudart=static -o scale_probe scale_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;
}

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

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

/* verify kernel: out[0]=count(c==a+b), out[1]=count(c==2*(a+b)) */
__global__ void vdigest(const float *a, const float *b, const float *c, int n,
			unsigned long long *out)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i >= n)
		return;
	if (c[i] == a[i] + b[i])
		atomicAdd(&out[0], 1ull);
	if (c[i] == (a[i] + b[i]) * 2.0f)
		atomicAdd(&out[1], 1ull);
}

int main(int argc, char **argv)
{
	const char *mode = argc > 1 ? argv[1] : "V";
	cudaDeviceProp prop;
	const int n = 1 << 15;
	const int grid = (n + 255) / 256;
	size_t bytes = (size_t)n * sizeof(float);
	float *ha, *hb, *hc, *hd;
	float *da, *db, *dc, *dd;
	unsigned long long *d_out, h_out[2];
	int i;
	float k = argc > 2 ? atof(argv[2]) : 0.0f;

	setvbuf(stdout, NULL, _IONBF, 0);

	CHECK(cudaGetDeviceProperties(&prop, 0));
	{
		FILE *f = fopen("/proc/uptime", "r");
		double up = -1.0;

		if (f) { if (fscanf(f, "%lf", &up) != 1) up = -1.0; fclose(f); }
		printf("UPTIME %.3f\n", up);
	}
	printf("DEVICE %s mode=%s seed=%g\n", prop.name, mode, (double)k);

	ha = (float *)malloc(bytes);
	hb = (float *)malloc(bytes);
	hc = (float *)malloc(bytes);
	hd = (float *)malloc(bytes);
	for (i = 0; i < n; i++) {
		ha[i] = (float)(i % 1000) * 0.5f + k;
		hb[i] = (float)(i % 997) * 0.25f + 2.0f * k;
	}
	memset(hc, 0xCD, bytes);
	memset(hd, 0xCD, bytes);

	CHECK(cudaMalloc(&da, bytes));
	CHECK(cudaMalloc(&db, bytes));
	CHECK(cudaMalloc(&dc, bytes));
	CHECK(cudaMalloc(&dd, bytes));
	CHECK(cudaMalloc(&d_out, 2 * sizeof(unsigned long long)));
	printf("DPTRS da=%p db=%p dc=%p dd=%p\n", (void *)da, (void *)db,
	       (void *)dc, (void *)dd);

	CHECK(cudaMemcpy(da, ha, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(db, hb, bytes, cudaMemcpyHostToDevice));

	vadd<<<grid, 256>>>(da, db, dc, n);
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(hc, dc, bytes, cudaMemcpyDeviceToHost));
	{
		int b = 0;

		for (i = 0; i < n; i++)
			if (hc[i] != ha[i] + hb[i])
				b++;
		printf("S1 vadd readback bad=%d/%d\n", b, n);
		CHECK(cudaMemset(d_out, 0, 2 * sizeof(unsigned long long)));
		vdigest<<<grid, 256>>>(da, db, dc, n, d_out);
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(h_out, d_out,
				 2 * sizeof(unsigned long long),
				 cudaMemcpyDeviceToHost));
		printf("S1 GPU-DIGEST c_eq_ab=%llu c_eq_2ab=%llu\n", h_out[0],
		       h_out[1]);
	}

	if (!strcmp(mode, "V")) {
		vadd<<<grid, 256>>>(da, db, dc, n);
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(hc, dc, bytes, cudaMemcpyDeviceToHost));
		{
			int b = 0;

			for (i = 0; i < n; i++)
				if (hc[i] != ha[i] + hb[i])
					b++;
			printf("S2 vadd-again readback bad=%d/%d\n", b, n);
		}
		printf("STATUS DONE\n");
		return 0;
	}

	if (!strcmp(mode, "I")) {
		int stage;

		for (stage = 0; stage < 2; stage++) {
			float mult = stage == 0 ? 2.0f : 4.0f;

			scale<<<grid, 256>>>(dc, 2.0f, n);
			CHECK(cudaDeviceSynchronize());
			CHECK(cudaMemcpy(hc, dc, bytes, cudaMemcpyDeviceToHost));
			{
				int b = 0, page;

				for (i = 0; i < n; i++)
					if (hc[i] != (ha[i] + hb[i]) * mult)
						b++;
				printf("S%d inplace-scale readback bad=%d/%d (mult=%g)\n",
				       2 + stage, b, n, mult);
				if (stage == 0)
					for (page = 0; page < 32; page++) {
						int pb = 0, k;

						for (k = 0; k < 256; k++)
							if (hc[page * 256 + k] !=
							    (ha[page * 256 + k] +
							     hb[page * 256 + k]) * mult)
								pb++;
						if (pb)
							printf("S%d PAGE %d bad=%d/256 first=%g want=%g\n",
							       2 + stage, page,
							       pb, hc[page * 256],
							       (ha[page * 256] +
								hb[page * 256]) * mult);
					}
			}
			if (stage == 0) {
				CHECK(cudaMemset(d_out, 0,
						 2 * sizeof(unsigned long long)));
				vdigest<<<grid, 256>>>(da, db, dc, n, d_out);
				CHECK(cudaDeviceSynchronize());
				CHECK(cudaMemcpy(h_out, d_out,
						 2 * sizeof(unsigned long long),
						 cudaMemcpyDeviceToHost));
				printf("S2 GPU-DIGEST c_eq_ab=%llu c_eq_2ab=%llu\n",
				       h_out[0], h_out[1]);
			}
		}
		printf("STATUS DONE\n");
		return 0;
	}

	if (!strcmp(mode, "O")) {
		scale_out<<<grid, 256>>>(dc, dd, 2.0f, n);
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(hd, dd, bytes, cudaMemcpyDeviceToHost));
		{
			int b = 0, page;

			for (i = 0; i < n; i++)
				if (hd[i] != (ha[i] + hb[i]) * 2.0f)
					b++;
			printf("S2 outofplace-scale readback bad=%d/%d\n", b, n);
			for (page = 0; page < 32; page++) {
				int pb = 0, k;

				for (k = 0; k < 256; k++)
					if (hd[page * 256 + k] !=
					    (ha[page * 256 + k] +
					     hb[page * 256 + k]) * 2.0f)
						pb++;
				if (pb)
					printf("S2 PAGE %d bad=%d/256\n", page,
					       pb);
			}
		}
		CHECK(cudaMemset(d_out, 0, 2 * sizeof(unsigned long long)));
		vdigest<<<grid, 256>>>(da, db, dd, n, d_out);
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(h_out, d_out, 2 * sizeof(unsigned long long),
				 cudaMemcpyDeviceToHost));
		printf("S2 GPU-DIGEST d_eq_ab=%llu d_eq_2ab=%llu\n", h_out[0],
		       h_out[1]);
		printf("STATUS DONE\n");
		return 0;
	}
	return 1;
}