// SPDX-License-Identifier: BSD-3-Clause
/*
 * transformer_mt.cu - attention forward on 3 independent streams + a
 * 5-launch sequential loop, in one binary.
 *
 * Same single-head attention op as transformer.cu (X[64x256],
 * Wq/Wk/Wv/Wo[256x256], head dim 256, fp32, softmax(Q*K^T/16), Y=(A*V)*Wo),
 * exercised two ways:
 *
 *  1. Multi-stream class: three independent streams each run the full op on
 *     their own input (X_s[i] = xval(i + s)). Every stream records an event
 *     after its chain; stream 2 waits on streams 0 and 1 via cudaStreamWaitEvent
 *     and the whole span is timed with cudaEventElapsedTime.
 *  2. Repeated-launch class: the op is launched 5 times in a loop on one
 *     stream with per-iteration inputs and per-iteration outputs, one
 *     synchronization at the end.
 *
 * Both classes are verified against an internal double-precision CPU
 * reference (tolerance 1e-3 relative, 1e-6 absolute floor).
 *
 * Output lines are stable and greppable:
 *
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   TRANSFORMER_MT streams=3 bad=<mismatches> ysum=<sum>
 *   TRANSFORMER_LOOP launches=5 bad=<mismatches> ysum=<sum>
 *   TIMING ms=<elapsed>
 *   STATUS OK|FAIL
 *
 * Build with a static CUDA runtime (--cudart=static) so the guest only needs
 * the driver's libcuda.so, not a CUDA toolkit installation:
 *
 *   nvcc -O2 --cudart=static -arch=sm_75 -o transformer_mt transformer_mt.cu
 *
 * Then, from the guest:  /tmp/transformer_mt
 */

#include <cstdio>
#include <cstdlib>
#include <cstdint>
#include <cmath>
#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 SEQ 64
#define DH 256
#define TILE 16
#define TOL 1e-3f
#define NST 3
#define NLOOP 5

__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;
}

__global__ void gemm_nt(const float *A, const float *B, float *C,
			int M, int N, int K)
{
	/* C[i][j] = sum_k A[i][k] * B[j][k]  (B is [N][K], row-major). */
	__shared__ float As[TILE][TILE];
	__shared__ float Bs[TILE][TILE];
	int row = blockIdx.y * TILE + threadIdx.y;
	int tcol = 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[(blockIdx.x * TILE + threadIdx.y) * K +
			  t * TILE + threadIdx.x];
		__syncthreads();
		for (int k = 0; k < TILE; k++)
			acc += As[threadIdx.y][k] * Bs[threadIdx.x][k];
		__syncthreads();
	}
	C[row * N + tcol] = acc;
}

__global__ void softmax_rows(const float *S, float *P, int rows, int cols,
			     float scale)
{
	extern __shared__ float sm[];
	int row = blockIdx.x;
	int t = threadIdx.x;
	float v = S[row * cols + t] * scale;
	int stride;

	sm[t] = v;
	__syncthreads();
	for (stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
		if (t < stride)
			sm[t] = fmaxf(sm[t], sm[t + stride]);
		__syncthreads();
	}
	v = expf(v - sm[0]);
	sm[t] = v;
	__syncthreads();
	for (stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
		if (t < stride)
			sm[t] += sm[t + stride];
		__syncthreads();
	}
	P[row * cols + t] = v / sm[0];
}

/* ------------------------------------------------------------------ host */

struct attn_bufs {
	float *q, *k, *v, *s, *a, *av, *y;
};

static void launch_chain(cudaStream_t st, const float *x, struct attn_bufs b,
			 const float *wq, const float *wk, const float *wv,
			 const float *wo)
{
	dim3 g1(DH / TILE, SEQ / TILE);
	dim3 g2(SEQ / TILE, SEQ / TILE);
	dim3 tb(TILE, TILE);
	static int chain_n;

	/* r226 GRID-WIT: the exact geometry every chain launch uses, so the
	 * host-side QMD geometry capture (GRID-WIT in vgsp) can be compared
	 * against what the client actually launched. Chains 0..2 are the
	 * multi-stream class, 3..7 the 5-launch loop. */
	printf("GEO chain=%d gemm_nn grid=(%u,%u,1) block=(%u,%u,1) gemm_nt grid=(%u,%u,1) softmax grid=(64,1,1) block=(64,1,1)\n",
	       chain_n++, g1.x, g1.y, tb.x, tb.y, g2.x, g2.y);

	gemm_nn<<<g1, tb, 0, st>>>(x, wq, b.q, SEQ, DH, DH);
	gemm_nn<<<g1, tb, 0, st>>>(x, wk, b.k, SEQ, DH, DH);
	gemm_nn<<<g1, tb, 0, st>>>(x, wv, b.v, SEQ, DH, DH);
	gemm_nt<<<g2, tb, 0, st>>>(b.q, b.k, b.s, SEQ, SEQ, DH);
	softmax_rows<<<SEQ, SEQ, (size_t)SEQ * sizeof(float), st>>>(b.s, b.a,
								    SEQ, SEQ,
								    1.0f / 16.0f);
	gemm_nn<<<g1, tb, 0, st>>>(b.a, b.v, b.av, SEQ, DH, SEQ);
	gemm_nn<<<g1, tb, 0, st>>>(b.av, wo, b.y, SEQ, DH, DH);
}

static float xval(int i)
{
	return (float)(((i % 7) - 3)) * 0.25f;
}

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

static void ref_gemm(const double *A, const double *B, double *C,
		     int M, int N, int K)
{
	for (int i = 0; i < M; i++)
		for (int j = 0; j < N; j++) {
			double acc = 0.0;

			for (int k = 0; k < K; k++)
				acc += A[i * K + k] * B[k * N + j];
			C[i * N + j] = acc;
		}
}

static void ref_gemm_nt(const double *A, const double *B, double *C,
			int M, int N, int K)
{
	for (int i = 0; i < M; i++)
		for (int j = 0; j < N; j++) {
			double acc = 0.0;

			for (int k = 0; k < K; k++)
				acc += A[i * K + k] * B[j * K + k];
			C[i * N + j] = acc;
		}
}

static void ref_softmax(const double *S, double *P, int rows, int cols,
			double scale)
{
	for (int i = 0; i < rows; i++) {
		const double *s = S + i * cols;
		double *p = P + i * cols;
		double m = s[0] * scale, sum = 0.0;

		for (int j = 1; j < cols; j++)
			if (s[j] * scale > m)
				m = s[j] * scale;
		for (int j = 0; j < cols; j++) {
			p[j] = exp(s[j] * scale - m);
			sum += p[j];
		}
		for (int j = 0; j < cols; j++)
			p[j] /= sum;
	}
}

static int count_bad(const float *gpu, const double *ref, int n)
{
	int bad = 0;

	for (int i = 0; i < n; i++) {
		double g = (double)gpu[i];
		double r = ref[i];

		if (fabs(g - r) > (double)TOL * (fabs(r) + 1e-6))
			bad++;
	}
	return bad;
}

static double dsum(const float *h, int n)
{
	double s = 0.0;

	for (int i = 0; i < n; i++)
		s += (double)h[i];
	return s;
}

/* Verify one attention result against the CPU reference for input offset o. */
static int verify_one(const float *h_y, int o, double *ysum_out)
{
	const int seqdh = SEQ * DH, seq2 = SEQ * SEQ;
	double *x, *q, *k, *v, *s, *a, *av, *y, *w;
	int i, bad;

	x = (double *)malloc((size_t)seqdh * sizeof(double));
	q = (double *)malloc((size_t)seqdh * sizeof(double));
	k = (double *)malloc((size_t)seqdh * sizeof(double));
	v = (double *)malloc((size_t)seqdh * sizeof(double));
	s = (double *)malloc((size_t)seq2 * sizeof(double));
	a = (double *)malloc((size_t)seq2 * sizeof(double));
	av = (double *)malloc((size_t)seqdh * sizeof(double));
	y = (double *)malloc((size_t)seqdh * sizeof(double));
	w = (double *)malloc((size_t)DH * DH * sizeof(double));
	if (!x || !q || !k || !v || !s || !a || !av || !y || !w)
		exit(1);

	for (i = 0; i < seqdh; i++)
		x[i] = (double)xval(i + o);
	for (i = 0; i < DH * DH; i++)
		w[i] = (double)wval(i, 1);
	ref_gemm(x, w, q, SEQ, DH, DH);
	for (i = 0; i < DH * DH; i++)
		w[i] = (double)wval(i, 2);
	ref_gemm(x, w, k, SEQ, DH, DH);
	for (i = 0; i < DH * DH; i++)
		w[i] = (double)wval(i, 3);
	ref_gemm(x, w, v, SEQ, DH, DH);
	ref_gemm_nt(q, k, s, SEQ, SEQ, DH);
	ref_softmax(s, a, SEQ, SEQ, 1.0 / 16.0);
	ref_gemm(a, v, av, SEQ, DH, SEQ);
	for (i = 0; i < DH * DH; i++)
		w[i] = (double)wval(i, 4);
	ref_gemm(av, w, y, SEQ, DH, DH);

	bad = count_bad(h_y, y, seqdh);
	*ysum_out = dsum(h_y, seqdh);
	free(x);
	free(q);
	free(k);
	free(v);
	free(s);
	free(a);
	free(av);
	free(y);
	free(w);
	return bad;
}

int main(void)
{
	cudaDeviceProp prop;
	const int seqdh = SEQ * DH, dh2 = DH * DH;
	const size_t b_seqdh = (size_t)seqdh * sizeof(float);
	const size_t b_dh2 = (size_t)dh2 * sizeof(float);
	const size_t b_seq2 = (size_t)SEQ * SEQ * sizeof(float);
	cudaStream_t st[NST], st_loop;
	cudaEvent_t ev[NST], t0, t1;
	struct attn_bufs ab[NST], lb;
	float *d_x[NST], *h_y[NST];
	float *d_xl[NLOOP], *d_yl[NLOOP], *h_yl[NLOOP];
	float *d_wq, *d_wk, *d_wv, *d_wo, *wl;
	float ms = 0.0f;
	int bad_mt = 0, bad_loop = 0, s, i;

	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);

	CHECK(cudaMalloc(&d_wq, b_dh2));
	CHECK(cudaMalloc(&d_wk, b_dh2));
	CHECK(cudaMalloc(&d_wv, b_dh2));
	CHECK(cudaMalloc(&d_wo, b_dh2));
	wl = (float *)malloc(b_dh2);
	if (!wl)
		return 1;
	for (i = 0; i < dh2; i++)
		wl[i] = wval(i, 1);
	CHECK(cudaMemcpy(d_wq, wl, b_dh2, cudaMemcpyHostToDevice));
	for (i = 0; i < dh2; i++)
		wl[i] = wval(i, 2);
	CHECK(cudaMemcpy(d_wk, wl, b_dh2, cudaMemcpyHostToDevice));
	for (i = 0; i < dh2; i++)
		wl[i] = wval(i, 3);
	CHECK(cudaMemcpy(d_wv, wl, b_dh2, cudaMemcpyHostToDevice));
	for (i = 0; i < dh2; i++)
		wl[i] = wval(i, 4);
	CHECK(cudaMemcpy(d_wo, wl, b_dh2, cudaMemcpyHostToDevice));

	/* ------------- class 1: 3 independent streams ------------- */
	for (s = 0; s < NST; s++) {
		float *hx = (float *)malloc(b_seqdh);

		if (!hx)
			return 1;
		for (i = 0; i < seqdh; i++)
			hx[i] = xval(i + s);
		h_y[s] = (float *)malloc(b_seqdh);
		if (!h_y[s])
			return 1;

		CHECK(cudaMalloc(&d_x[s], b_seqdh));
		CHECK(cudaMalloc(&ab[s].q, b_seqdh));
		CHECK(cudaMalloc(&ab[s].k, b_seqdh));
		CHECK(cudaMalloc(&ab[s].v, b_seqdh));
		CHECK(cudaMalloc(&ab[s].s, b_seq2));
		CHECK(cudaMalloc(&ab[s].a, b_seq2));
		CHECK(cudaMalloc(&ab[s].av, b_seqdh));
		CHECK(cudaMalloc(&ab[s].y, b_seqdh));
		CHECK(cudaStreamCreate(&st[s]));
		CHECK(cudaEventCreate(&ev[s]));
		CHECK(cudaMemcpyAsync(d_x[s], hx, b_seqdh,
				      cudaMemcpyHostToDevice, st[s]));
		free(hx);
	}
	CHECK(cudaEventCreate(&t0));
	CHECK(cudaEventCreate(&t1));
	CHECK(cudaEventRecord(t0));

	for (s = 0; s < NST; s++) {
		launch_chain(st[s], d_x[s], ab[s], d_wq, d_wk, d_wv, d_wo);
		CHECK(cudaGetLastError());
		CHECK(cudaEventRecord(ev[s], st[s]));
		CHECK(cudaMemcpyAsync(h_y[s], ab[s].y, b_seqdh,
				      cudaMemcpyDeviceToHost, st[s]));
	}
	/* event-ordered join: stream 2 waits on streams 0 and 1 */
	CHECK(cudaStreamWaitEvent(st[2], ev[0], 0));
	CHECK(cudaStreamWaitEvent(st[2], ev[1], 0));
	CHECK(cudaEventRecord(t1, st[2]));
	for (s = 0; s < NST; s++)
		CHECK(cudaStreamSynchronize(st[s]));
	CHECK(cudaEventElapsedTime(&ms, t0, t1));

	for (s = 0; s < NST; s++) {
		double ys;

		bad_mt += verify_one(h_y[s], s, &ys);
	}
	printf("TRANSFORMER_MT streams=%d bad=%d ysum=%.6f\n", NST, bad_mt,
	       dsum(h_y[0], seqdh));

	/* ----------------- class 2: 5-launch loop ----------------- */
	CHECK(cudaStreamCreate(&st_loop));
	CHECK(cudaMalloc(&lb.q, b_seqdh));
	CHECK(cudaMalloc(&lb.k, b_seqdh));
	CHECK(cudaMalloc(&lb.v, b_seqdh));
	CHECK(cudaMalloc(&lb.s, b_seq2));
	CHECK(cudaMalloc(&lb.a, b_seq2));
	CHECK(cudaMalloc(&lb.av, b_seqdh));
	CHECK(cudaMalloc(&lb.y, b_seqdh)); /* unused; op writes per-iter y */
	for (i = 0; i < NLOOP; i++) {
		float *hx = (float *)malloc(b_seqdh);

		if (!hx)
			return 1;
		for (s = 0; s < seqdh; s++)
			hx[s] = xval(s + 11 * i);
		CHECK(cudaMalloc(&d_xl[i], b_seqdh));
		CHECK(cudaMalloc(&d_yl[i], b_seqdh));
		CHECK(cudaMemcpy(d_xl[i], hx, b_seqdh,
				 cudaMemcpyHostToDevice));
		h_yl[i] = (float *)malloc(b_seqdh);
		if (!h_yl[i])
			return 1;
		free(hx);
	}

	for (i = 0; i < NLOOP; i++) {
		struct attn_bufs b = lb;

		b.y = d_yl[i];
		launch_chain(st_loop, d_xl[i], b, d_wq, d_wk, d_wv, d_wo);
		CHECK(cudaGetLastError());
	}
	CHECK(cudaStreamSynchronize(st_loop));
	for (i = 0; i < NLOOP; i++) {
		double ys;

		CHECK(cudaMemcpy(h_yl[i], d_yl[i], b_seqdh,
				 cudaMemcpyDeviceToHost));
		bad_loop += verify_one(h_yl[i], 11 * i, &ys);
	}
	printf("TRANSFORMER_LOOP launches=%d bad=%d ysum=%.6f\n", NLOOP,
	       bad_loop, dsum(h_yl[NLOOP - 1], seqdh));
	printf("TIMING ms=%.3f\n", ms);

	if (bad_mt != 0 || bad_loop != 0) {
		printf("STATUS FAIL\n");
		return 1;
	}

	printf("STATUS OK\n");
	for (s = 0; s < NST; s++) {
		CHECK(cudaEventDestroy(ev[s]));
		CHECK(cudaStreamDestroy(st[s]));
		CHECK(cudaFree(d_x[s]));
		CHECK(cudaFree(ab[s].q));
		CHECK(cudaFree(ab[s].k));
		CHECK(cudaFree(ab[s].v));
		CHECK(cudaFree(ab[s].s));
		CHECK(cudaFree(ab[s].a));
		CHECK(cudaFree(ab[s].av));
		CHECK(cudaFree(ab[s].y));
		free(h_y[s]);
	}
	for (i = 0; i < NLOOP; i++) {
		CHECK(cudaFree(d_xl[i]));
		CHECK(cudaFree(d_yl[i]));
		free(h_yl[i]);
	}
	CHECK(cudaEventDestroy(t0));
	CHECK(cudaEventDestroy(t1));
	CHECK(cudaStreamDestroy(st_loop));
	CHECK(cudaFree(lb.q));
	CHECK(cudaFree(lb.k));
	CHECK(cudaFree(lb.v));
	CHECK(cudaFree(lb.s));
	CHECK(cudaFree(lb.a));
	CHECK(cudaFree(lb.av));
	CHECK(cudaFree(lb.y));
	CHECK(cudaFree(d_wq));
	CHECK(cudaFree(d_wk));
	CHECK(cudaFree(d_wv));
	CHECK(cudaFree(d_wo));
	free(wl);
	return 0;
}