// SPDX-License-Identifier: BSD-3-Clause
/*
 * ctaid_probe.cu - CTA-coordinate + engine-read reality probe (r229 revision).
 *
 * Revisions since the first r229 run:
 *  - RAWA: the A buffer is copied back D2H through a plain cudaMemcpy BEFORE
 *    any kernel and printed (first8/last8), so the ground truth "does the
 *    buffer itself hold i" is on record with no engine involvement.
 *  - RAWA2: the same readback is repeated AFTER both phases, so a launch
 *    that corrupts the input buffer would show.
 *  - Every band prints its raw first 16 values (was 4).
 *  - Phase 2 is run TWICE (2 and 2b) into two different output buffers and
 *    both are printed: a first-attempt-only failure is the engine's
 *    per-launch lottery, not addressing.
 *
 *  Phase 1  coord map : out[(by*16+ty)*256 + (bx*16+tx)] =
 *                          by*10000 + bx*100 + ty*10 + tx
 *  Phase 2  A-row read: out2[(by*16+ty)*256 + (bx*16+tx)] =
 *                          A[(by*16+ty)*256 + (bx*16+tx)], A[i]=i
 *
 * Output lines are stable and greppable:
 *   DEVICE <name> ...
 *   PTRS d_A=0x... d_out=0x...
 *   RAWA first8: ... last8: ...
 *   PROBE1 band=<by> row0: <16 values>
 *   PROBE1 bad=<count>
 *   PROBE2 band=<by> row0: <16 values>
 *   PROBE2 bad=<count>
 *   PROBE2B band=<by> row0: <16 values>
 *   PROBE2B bad=<count>
 *   RAWA2 first8: ... last8: ...
 *   CTAID-PROBE DONE
 *
 * Build: nvcc -O2 --cudart=static -arch=sm_75 -o ctaid_probe ctaid_probe.cu
 * Run  : timeout -k 5 20 /tmp/ctaid_probe
 */

#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("CTAID-PROBE DONE (FAIL)\n");             \
			return 1;                                        \
		}                                                        \
	} while (0)

#define SEQ 64
#define DH 256
#define KK 256
#define TILE 16

__global__ void coord_map(int *out)
{
	int row = blockIdx.y * TILE + threadIdx.y;
	int col = blockIdx.x * TILE + threadIdx.x;

	out[row * DH + col] = blockIdx.y * 10000 + blockIdx.x * 100 +
			      threadIdx.y * 10 + threadIdx.x;
}

__global__ void arow_read(const int *A, int *out)
{
	int row = blockIdx.y * TILE + threadIdx.y;
	int col = blockIdx.x * TILE + threadIdx.x;

	out[row * DH + col] = A[row * KK + col];
}

static void print_band(const char *tag, const int *h, int by)
{
	int base = by * TILE * DH, tx;

	printf("%s band=%d row0:", tag, by);
	for (tx = 0; tx < 16; tx++)
		printf(" %d", h[base + tx]);
	printf("\n");
}

static int check_map(const int *h)
{
	int bad = 0, row, col;

	for (row = 0; row < SEQ; row++)
		for (col = 0; col < DH; col++) {
			int exp = (row / TILE) * 10000 + (col / TILE) * 100 +
				  (row % TILE) * 10 + (col % TILE);

			if (h[row * DH + col] != exp)
				bad++;
		}
	return bad;
}

static int check_arow(const int *h)
{
	int bad = 0, row, col;

	for (row = 0; row < SEQ; row++)
		for (col = 0; col < DH; col++)
			if (h[row * DH + col] != row * KK + col)
				bad++;
	return bad;
}

int main(void)
{
	cudaDeviceProp prop;
	const int n = SEQ * DH;
	const size_t abytes = (size_t)SEQ * KK * sizeof(int);
	const size_t obytes = (size_t)n * sizeof(int);
	int *d_out, *d_out2, *d_A;
	int *h_out, *h_A;
	dim3 grid(DH / TILE, SEQ / TILE); /* (16,4,1) */
	dim3 block(TILE, TILE);           /* (16,16,1) */
	int by;

	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_out = (int *)malloc(obytes);
	h_A = (int *)malloc(abytes);
	if (!h_out || !h_A)
		return 1;
	for (by = 0; by < SEQ * KK; by++)
		h_A[by] = by;

	CHECK(cudaMalloc(&d_out, obytes));
	CHECK(cudaMalloc(&d_out2, obytes));
	CHECK(cudaMalloc(&d_A, abytes));
	printf("PTRS d_A=0x%llx d_out=0x%llx d_out2=0x%llx\n",
	       (unsigned long long)(uintptr_t)d_A,
	       (unsigned long long)(uintptr_t)d_out,
	       (unsigned long long)(uintptr_t)d_out2);
	printf("ATAG h_A=0x%llx h_out=0x%llx A=0x%llx out=0x%llx out2=0x%llx\n",
	       (unsigned long long)(uintptr_t)h_A,
	       (unsigned long long)(uintptr_t)h_out,
	       (unsigned long long)(uintptr_t)d_A,
	       (unsigned long long)(uintptr_t)d_out,
	       (unsigned long long)(uintptr_t)d_out2);

	/* Ground truth: upload A, read it straight back with no kernel. */
	printf("ATAG phase=RAWA src=A dst=h_out buf=A\n");
	CHECK(cudaMemcpy(d_A, h_A, abytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(h_out, d_A, abytes, cudaMemcpyDeviceToHost));
	printf("RAWA first8: %d %d %d %d %d %d %d %d last8: %d %d %d %d %d %d %d %d\n",
	       h_out[0], h_out[1], h_out[2], h_out[3], h_out[4], h_out[5],
	       h_out[6], h_out[7], h_out[16376], h_out[16377], h_out[16378],
	       h_out[16379], h_out[16380], h_out[16381], h_out[16382],
	       h_out[16383]);

	/* ---- Phase 1: CTA coordinate map ---- */
	printf("ATAG phase=PROBE1 src=out dst=h_out buf=out\n");
	CHECK(cudaMemset(d_out, 0, obytes));
	coord_map<<<grid, block>>>(d_out);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_out, d_out, obytes, cudaMemcpyDeviceToHost));
	for (by = 0; by < SEQ / TILE; by++)
		print_band("PROBE1", h_out, by);
	printf("PROBE1 bad=%d\n", check_map(h_out));

	/* ---- Phase 2: A read, first attempt ---- */
	printf("ATAG phase=PROBE2 src=A dst=h_out buf=out\n");
	CHECK(cudaMemset(d_out, 0, obytes));
	arow_read<<<grid, block>>>(d_A, d_out);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_out, d_out, obytes, cudaMemcpyDeviceToHost));
	for (by = 0; by < SEQ / TILE; by++)
		print_band("PROBE2", h_out, by);
	printf("PROBE2 bad=%d\n", check_arow(h_out));

	/* ---- Phase 2b: A read, second attempt, fresh output buffer ---- */
	printf("ATAG phase=PROBE2B src=A dst=h_out buf=out2\n");
	CHECK(cudaMemset(d_out2, 0, obytes));
	arow_read<<<grid, block>>>(d_A, d_out2);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_out, d_out2, obytes, cudaMemcpyDeviceToHost));
	for (by = 0; by < SEQ / TILE; by++)
		print_band("PROBE2B", h_out, by);
	printf("PROBE2B bad=%d\n", check_arow(h_out));

	/* Post-launch A integrity (did any launch corrupt the input?). */
	printf("ATAG phase=RAWA2 src=A dst=h_out buf=A\n");
	CHECK(cudaMemcpy(h_out, d_A, abytes, cudaMemcpyDeviceToHost));
	printf("RAWA2 first8: %d %d %d %d %d %d %d %d last8: %d %d %d %d %d %d %d %d\n",
	       h_out[0], h_out[1], h_out[2], h_out[3], h_out[4], h_out[5],
	       h_out[6], h_out[7], h_out[16376], h_out[16377], h_out[16378],
	       h_out[16379], h_out[16380], h_out[16381], h_out[16382],
	       h_out[16383]);

	printf("CTAID-PROBE DONE\n");
	CHECK(cudaFree(d_out));
	CHECK(cudaFree(d_out2));
	CHECK(cudaFree(d_A));
	free(h_out);
	free(h_A);
	return 0;
}