// SPDX-License-Identifier: BSD-3-Clause
/*
 * tgemm_probe.cu - minimal isolation for the transformer K/V divergence.
 *
 * Reproduces exactly the transformer's Q/K/V GEMM shape (X[64x256] x
 * W[256x256], tile 16) for three weight matrices uploaded back-to-back
 * (wq, wk, wv), then verifies each against a CPU reference and prints
 * per-GEMM bad counts plus a bad-position histogram (first bad index per
 * 16-row band and per 16-col band) so a copy-side hole is distinguishable
 * from a kernel-side error.
 *
 * Output:
 *   DEVICE ...
 *   GEMM name=<q|k|v> bad=<n> first=(row,col) got=.. ref=..
 *   BADPOS name=.. rows=<r0,r1,..> cols=<c0,c1,..>
 *   STATUS OK|FAIL
 */

#include <cstdio>
#include <cstdlib>
#include <cstdint>
#include <cuda_runtime.h>

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

__global__ void gemm_nn(const float *A, const float *B, float *C,
			int M, int N, int K)
{
	__shared__ float As[TILE][TILE];
	__shared__ float Bs[TILE][TILE];
	int row = blockIdx.y * TILE + threadIdx.y;
	int col = blockIdx.x * TILE + threadIdx.x;
	float acc = 0.0f;

	for (int t = 0; t < K / TILE; t++) {
		As[threadIdx.y][threadIdx.x] = A[row * K + t * TILE + threadIdx.x];
		Bs[threadIdx.y][threadIdx.x] =
			B[(t * TILE + threadIdx.y) * N + col];
		__syncthreads();
		for (int k = 0; k < TILE; k++)
			acc += As[threadIdx.y][k] * Bs[k][threadIdx.x];
		__syncthreads();
	}
	C[row * N + col] = acc;
}

static float wval(int i, int seed)
{
	return (float)(((i * 13 + seed * 101 + 7) % 5) - 2) * 0.125f;
}

static float xval(int i, int j)
{
	return (float)(((i * 7 + j * 3 + 1) % 9) - 4) * 0.25f;
}

int main(void)
{
	cudaDeviceProp prop;
	float *h_x, *h_w, *h_c, *h_ref;
	float *d_x, *d_w[3], *d_c[3];
	const int seqdh = SEQ * DH, dh2 = DH * DH;
	int g, i, j, k;

	setvbuf(stdout, NULL, _IONBF, 0);
	if (cudaGetDeviceProperties(&prop, 0) != cudaSuccess) {
		printf("STATUS FAIL\n");
		return 1;
	}
	printf("DEVICE %s cc=%d.%d\n", prop.name, prop.major, prop.minor);

	h_x = (float *)malloc(seqdh * 4);
	h_w = (float *)malloc(dh2 * 4);
	h_c = (float *)malloc(seqdh * 4);
	h_ref = (float *)malloc(seqdh * 4);

	if (cudaMalloc(&d_x, seqdh * 4) != cudaSuccess)
		return 1;
	for (g = 0; g < 3; g++) {
		if (cudaMalloc(&d_w[g], dh2 * 4) != cudaSuccess)
			return 1;
		if (cudaMalloc(&d_c[g], seqdh * 4) != cudaSuccess)
			return 1;
	}

	for (i = 0; i < seqdh; i++)
		h_x[i] = xval(i / DH, i % DH);

	/* Upload X once (the transformer's shared input). */
	if (cudaMemcpy(d_x, h_x, seqdh * 4, cudaMemcpyHostToDevice) !=
	    cudaSuccess)
		return 1;
	/* Verify the X upload itself (first row + full buffer). */
	{
		float *h_xb = (float *)malloc(seqdh * 4);
		int xbad = 0, xr0 = 0;

		if (!h_xb)
			return 1;
		if (cudaMemcpy(h_xb, d_x, seqdh * 4,
			       cudaMemcpyDeviceToHost) != cudaSuccess)
			return 1;
		for (i = 0; i < seqdh; i++)
			if (h_xb[i] != h_x[i]) {
				xbad++;
				if (i < DH)
					xr0++;
			}
		printf("XVERIFY bad=%d row0_bad=%d\n", xbad, xr0);
		free(h_xb);
	}

	for (g = 0; g < 3; g++) {
		int bad = 0, first_r = -1, first_c = -1;
		float got = 0, ref = 0;
		int rows[64] = {0}, cols[64] = {0};
		int nr = 0, nc = 0;
		int pass;

		for (i = 0; i < dh2; i++)
			h_w[i] = wval(i, g + 1);
		if (cudaMemcpy(d_w[g], h_w, dh2 * 4,
			       cudaMemcpyHostToDevice) != cudaSuccess)
			return 1;

		dim3 blk(TILE, TILE);
		dim3 grd(DH / TILE, SEQ / TILE);

		for (pass = 0; pass < 3; pass++) {
			float *hc = h_c, *hxb = h_ref;
			int xbad = 0;

			gemm_nn<<<grd, blk>>>(d_x, d_w[g], d_c[g], SEQ, DH, DH);
			if (cudaDeviceSynchronize() != cudaSuccess) {
				printf("SYNC FAIL g=%d\n", g);
				return 1;
			}
			if (cudaMemcpy(hc, d_c[g], seqdh * 4,
				       cudaMemcpyDeviceToHost) != cudaSuccess)
				return 1;
			/* re-verify X after the launch */
			if (cudaMemcpy(hxb, d_x, seqdh * 4,
				       cudaMemcpyDeviceToHost) != cudaSuccess)
				return 1;
			for (i = 0; i < seqdh; i++)
				if (hxb[i] != h_x[i]) {
					printf("XDIFF pass=%d i=%d got=%.6f want=%.6f\n",
					       pass, i, hxb[i], h_x[i]);
					xbad++;
				}
			bad = 0;
			nr = 0;
			nc = 0;
			for (i = 0; i < SEQ; i++)
				for (j = 0; j < DH; j++) {
					double acc = 0.0;
					float r, a, diff;

					for (k = 0; k < DH; k++)
						acc += (double)h_x[i * DH + k] *
						       (double)h_w[k * DH + j];
					r = (float)acc;
					a = hc[i * DH + j];
					diff = a - r;
					if (diff < 0)
						diff = -diff;
					if (diff > 1e-4f) {
						int rr = i, cc = j;

						if (bad == 0) {
							first_r = rr;
							first_c = cc;
							got = a;
							ref = r;
						}
						bad++;
						if (nr < 64 && (bad == 1 ||
								rr != rows[nr - 1]))
							rows[nr++] = rr;
						if (nc < 64 && (bad == 1 ||
								cc != cols[nc - 1]))
							cols[nc++] = cc;
					}
				}
			printf("PASS name=%c pass=%d bad=%d xbad=%d first=(%d,%d) got=%.6f ref=%.6f\n",
			       "qkv"[g], pass, bad, xbad, first_r, first_c,
			       got, ref);
			if (bad) {
				printf("ROW0 name=%c got=", "qkv"[g]);
				for (i = 0; i < 8; i++)
					printf("%.6f ", hc[i]);
				printf(" ref=");
				for (i = 0; i < 8; i++) {
					double acc = 0.0;

					for (k = 0; k < DH; k++)
						acc += (double)h_x[k] *
						       (double)h_w[k * DH + i];
					printf("%.6f ", (float)acc);
				}
				printf("\n");
			}
		}
		(void)got;
		(void)ref;
		(void)first_r;
		(void)first_c;
		(void)nr;
		(void)nc;
	}

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