// SPDX-License-Identifier: BSD-3-Clause
/*
 * rt_probe.cu - instrumented memcpy_roundtrip variant isolating the
 * memset-phase D2H readback failure.
 *
 * Phases (argv[1]):
 *   "full"   : H2D -> kernel -> D2H -> memset -> D2H (the original order)
 *   "sync"   : same but cudaDeviceSynchronize() before the second D2H
 *   "memset" : memset -> D2H only (memset_probe shape, 3-buffer alloc)
 *   "rev"    : memset -> D2H first, then H2D -> kernel -> D2H
 *
 * Forensics on the memset readback: first bad offset, byte histogram
 * (0x5A / 0x00 / other), per-1MB chunk histogram, and the first 32 bytes.
 *
 * Build: nvcc -O2 --cudart=static -o rt_probe rt_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)

#define MemsetByte 0x5A

__global__ void bump(const float *in, float *out, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		out[i] = in[i] + 0.5f;
}

static void memset_forensics(const unsigned char *h, size_t bytes,
			     const char *tag)
{
	size_t i, first_bad = (size_t)-1;
	int bad = 0, n5a = 0, n00 = 0, nother = 0;
	size_t chunk;

	for (i = 0; i < bytes; i++) {
		if (h[i] == 0x5A)
			n5a++;
		else if (h[i] == 0x00)
			n00++;
		else
			nother++;
		if (h[i] != 0x5A) {
			if (first_bad == (size_t)-1)
				first_bad = i;
			bad++;
		}
	}
	printf("MEMSET-READ %s bad=%d first_bad=%zu n5a=%d n00=%d nother=%d\n",
	       tag, bad, first_bad, n5a, n00, nother);
	printf("MEMSET-HEAD %s", tag);
	for (i = 0; i < 32; i++)
		printf(" %02x", h[i]);
	printf("\n");
	if (first_bad != (size_t)-1) {
		printf("MEMSET-ATBAD %s off=%zu:", tag, first_bad);
		for (i = first_bad; i < first_bad + 16 && i < bytes; i++)
			printf(" %02x", h[i]);
		printf("\n");
	}
	for (chunk = 0; chunk < bytes / (1u << 20); chunk++) {
		size_t lo = chunk << 20, hi = lo + (1u << 20);
		int c5a = 0, c00 = 0, cother = 0;

		for (i = lo; i < hi; i++) {
			if (h[i] == 0x5A)
				c5a++;
			else if (h[i] == 0x00)
				c00++;
			else
				cother++;
		}
		printf("MEMSET-MB %s mb=%zu 5a=%d 00=%d other=%d\n",
		       tag, chunk, c5a, c00, cother);
	}
}

int main(int argc, char **argv)
{
	const char *mode = argc > 1 ? argv[1] : "full";
	cudaDeviceProp prop;
	const size_t bytes = 8u << 20;
	const int n = (int)(bytes / sizeof(float));
	unsigned char *h_m;
	float *h_in, *h_out;
	float *d_in, *d_out;
	unsigned char *d_m;
	double csum = 0.0;
	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 mode=%s\n", prop.name,
	       prop.major, prop.minor, prop.multiProcessorCount,
	       (unsigned long long)prop.totalGlobalMem, mode);

	h_in = (float *)malloc(bytes);
	h_out = (float *)malloc(bytes);
	h_m = (unsigned char *)malloc(bytes);
	if (!h_in || !h_out || !h_m)
		return 1;

	for (i = 0; i < n; i++)
		h_in[i] = (float)(i % 1024) * 0.25f;

	CHECK(cudaMalloc(&d_in, bytes));
	CHECK(cudaMalloc(&d_out, bytes));
	CHECK(cudaMalloc(&d_m, bytes));
	printf("DPTRS d_in=%p d_out=%p d_m=%p\n", (void *)d_in,
	       (void *)d_out, (void *)d_m);

	if (!strcmp(mode, "memset") || !strcmp(mode, "rev")) {
		/* memset phase first */
		CHECK(cudaMemset(d_m, MemsetByte, bytes));
		if (!strcmp(mode, "sync"))
			CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(h_m, d_m, bytes, cudaMemcpyDeviceToHost));
		memset_forensics(h_m, bytes, "early");
	}

	if (!strcmp(mode, "full") || !strcmp(mode, "sync") ||
	    !strcmp(mode, "rev")) {
		CHECK(cudaMemcpy(d_in, h_in, bytes, cudaMemcpyHostToDevice));
		bump<<<grid, 256>>>(d_in, d_out, n);
		CHECK(cudaGetLastError());
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(h_out, d_out, bytes, cudaMemcpyDeviceToHost));

		for (i = 0; i < n; i++) {
			if (h_out[i] != h_in[i] + 0.5f)
				bad++;
			csum += (double)h_out[i];
		}
		printf("BUMP bad=%d csum=%.6f\n", bad, csum);
	}

	if (!strcmp(mode, "full") || !strcmp(mode, "sync")) {
		CHECK(cudaMemset(d_m, MemsetByte, bytes));
		if (!strcmp(mode, "sync"))
			CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(h_m, d_m, bytes, cudaMemcpyDeviceToHost));
		memset_forensics(h_m, bytes, "late");
	}

	printf("STATUS OK\n");
	return 0;
}