// SPDX-License-Identifier: BSD-3-Clause
/*
 * perf_workload.cu - the Open-IOV near-native performance workload.
 *
 * One binary, two roles: the bare-metal oracle (run on the host with the
 * distribution driver loaded) and the measured side (run in the guest on the
 * Open-IOV emulated vGPU). Every metric is a stable, machine-readable PERF
 * line; every result is verified exactly on both sides, so a perf run can
 * never silently measure corrupted work.
 *
 * Stable output (greppable):
 *   PERF device=<name> cc=<maj>.<min> sm=<count>
 *   PERF run_cfg n_warm=.. n_malloc=.. n_thru=.. sizes=..
 *   PERF init_ms=<f>
 *   PERF malloc_first_ms=<f>
 *   PERF malloc_ms n=.. mean=.. p50=.. p95=..
 *   PERF copy dir=htod bytes=.. ms=.. gbps=..
 *   PERF copy dir=dtoh bytes=.. ms=.. gbps=..
 *   PERF kernel_fma iters=.. ms=.. gflops=.. bad=..
 *   PERF launch_first_ms=<f>
 *   PERF launch_warm n=.. mean=.. p50=.. p95=..
 *   PERF throughput n=.. wall_ms=.. kps=..
 *   PERF clock64 available=0|1 min=.. median=..
 *   PERF checksum fma=.. pattern=.. copy_bad=.. bad=..
 *   PERF STATUS OK|FAIL
 *
 * Build with a static CUDA runtime so the guest needs only the driver's
 * libcuda.so, exactly like ab.sh/victory_mail:
 *
 *   nvcc -O2 -arch=sm_75 --cudart=static -o perf_workload perf_workload.cu
 *
 * Run (either side):
 *   ./perf_workload [--n-warm=N] [--n-thru=N] [--no-clock64]
 */

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <ctime>
#include <cuda_runtime.h>

#define FMA_THREADS	8192
#define FMA_ITERS	250000
#define FMA_A		0.9999f
#define FMA_B		1.0e-6f

#define PAT_THREADS	32768
#define PAT_SEED	0x9e3779b9u

#define MALLOC_BYTES	(256u * 1024u)
#define N_MALLOC	50
#define N_COPY		3
#define CLK_BLOCKS	48
#define CLK_ITERS	10000

static double now_ms(void)
{
	struct timespec ts;

	clock_gettime(CLOCK_MONOTONIC, &ts);
	return (double)ts.tv_sec * 1e3 + (double)ts.tv_nsec / 1e6;
}

static int cmp_double(const void *a, const void *b)
{
	double x = *(const double *)a;
	double y = *(const double *)b;

	return (x > y) - (x < y);
}

static double pct(double *v, int n, double p)
{
	qsort(v, n, sizeof(double), cmp_double);
	return v[(int)(p / 100.0 * (n - 1) + 0.5)];
}

#define CK(x) do {							\
	cudaError_t e_ = (x);						\
	if (e_ != cudaSuccess) {					\
		printf("PERF ERROR %s at line %d\n",			\
		       cudaGetErrorString(e_), __LINE__);		\
		printf("PERF STATUS FAIL\n");				\
		return 1;						\
	}								\
} while (0)

__global__ void fma_chain(float *out, int iters, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;
	float x = 0.0f;

	if (i >= n)
		return;
	for (int t = 0; t < iters; t++)
		x = __fmaf_rn(x, FMA_A, FMA_B);
	out[i] = x;
}

static float host_chain(int iters)
{
	float x = 0.0f;

	for (int t = 0; t < iters; t++)
		x = fmaf(x, FMA_A, FMA_B);
	return x;
}

__global__ void pattern_kernel(unsigned *out, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i >= n)
		return;
	out[i] = (unsigned)i * 2654435761u ^ PAT_SEED;
}

static unsigned host_pattern(int i)
{
	return (unsigned)i * 2654435761u ^ PAT_SEED;
}

__global__ void clock_probe(unsigned long long *out)
{
	if (threadIdx.x == 0) {
		unsigned long long s = clock64();
		float x = 0.0f;

		for (int t = 0; t < CLK_ITERS; t++)
			x = __fmaf_rn(x, FMA_A, FMA_B);
		out[blockIdx.x * 2 + 0] = s;
		out[blockIdx.x * 2 + 1] = clock64();
	}
}

struct opts {
	int n_warm;
	int n_thru;
	int clock64;
};

static int parse_args(int argc, char **argv, struct opts *o)
{
	o->n_warm = 50;
	o->n_thru = 100;
	o->clock64 = 1;
	for (int i = 1; i < argc; i++) {
		if (!strncmp(argv[i], "--n-warm=", 9))
			o->n_warm = atoi(argv[i] + 9);
		else if (!strncmp(argv[i], "--n-thru=", 9))
			o->n_thru = atoi(argv[i] + 9);
		else if (!strcmp(argv[i], "--no-clock64"))
			o->clock64 = 0;
		else {
			printf("usage: %s [--n-warm=N] [--n-thru=N] [--no-clock64]\n",
			       argv[0]);
			return -1;
		}
	}
	if (o->n_warm < 1)
		o->n_warm = 1;
	if (o->n_thru < 1)
		o->n_thru = 1;
	return 0;
}

int main(int argc, char **argv)
{
	struct opts o;
	double t_start = now_ms();
	double t0;
	float ms;
	double samples[256];
	int bad = 0, pat_bad = 0, copy_bad = 0;
	float fma_ref = host_chain(FMA_ITERS);

	if (parse_args(argc, argv, &o) != 0)
		return 2;

	/* P1: process start -> device ready (the init path). */
	int dev = 0;
	struct cudaDeviceProp prop;

	CK(cudaGetDeviceCount(&dev));
	if (dev < 1) {
		printf("PERF ERROR no CUDA devices\n");
		printf("PERF STATUS FAIL\n");
		return 1;
	}
	CK(cudaSetDevice(0));
	CK(cudaGetDeviceProperties(&prop, 0));
	printf("PERF device=%s cc=%d.%d sm=%d\n", prop.name,
	       prop.major, prop.minor, prop.multiProcessorCount);

	/* Touch the device so init_ms covers the full lazy-init path. */
	CK(cudaFree(0));
	printf("PERF init_ms=%.3f\n", now_ms() - t_start);
	printf("PERF run_cfg n_warm=%d n_malloc=%d n_thru=%d sizes=4096,262144,4194304,16777216\n",
	       o.n_warm, N_MALLOC, o.n_thru);

	/* P2: cudaMalloc first-call and steady-state latency. */
	void *d_m = NULL;

	t0 = now_ms();
	CK(cudaMalloc(&d_m, MALLOC_BYTES));
	printf("PERF malloc_first_ms=%.3f\n", now_ms() - t0);
	CK(cudaFree(d_m));

	for (int i = 0; i < 8; i++) {
		CK(cudaMalloc(&d_m, MALLOC_BYTES));
		CK(cudaFree(d_m));
	}
	for (int i = 0; i < N_MALLOC; i++) {
		t0 = now_ms();
		CK(cudaMalloc(&d_m, MALLOC_BYTES));
		CK(cudaFree(d_m));
		samples[i] = now_ms() - t0;
	}
	{
		double s2[256];
		double a = 0;

		memcpy(s2, samples, sizeof(double) * N_MALLOC);
		for (int i = 0; i < N_MALLOC; i++)
			a += samples[i];
		printf("PERF malloc_ms n=%d mean=%.3f p50=%.3f p95=%.3f\n",
		       N_MALLOC, a / N_MALLOC,
		       pct(s2, N_MALLOC, 50), pct(s2, N_MALLOC, 95));
	}

	/* Verification pattern kernel + first launch. */
	unsigned *d_pat = NULL, *h_pat;

	h_pat = (unsigned *)malloc(sizeof(unsigned) * PAT_THREADS);
	CK(cudaMalloc(&d_pat, sizeof(unsigned) * PAT_THREADS));

	cudaEvent_t e0, e1;

	CK(cudaEventCreate(&e0));
	CK(cudaEventCreate(&e1));
	CK(cudaEventRecord(e0));
	pattern_kernel<<<(PAT_THREADS + 255) / 256, 256>>>(d_pat, PAT_THREADS);
	CK(cudaGetLastError());
	CK(cudaEventRecord(e1));
	CK(cudaDeviceSynchronize());
	CK(cudaEventElapsedTime(&ms, e0, e1));
	printf("PERF launch_first_ms=%.3f\n", (double)ms);
	CK(cudaMemcpy(h_pat, d_pat, sizeof(unsigned) * PAT_THREADS,
		      cudaMemcpyDeviceToHost));
	for (int i = 0; i < PAT_THREADS; i++)
		if (h_pat[i] != host_pattern(i))
			pat_bad++;

	/* P4: warm launch latency (enqueue -> sync completion per launch). */
	if (o.n_warm > 256)
		o.n_warm = 256;
	for (int i = 0; i < o.n_warm; i++) {
		CK(cudaEventRecord(e0));
		pattern_kernel<<<(PAT_THREADS + 255) / 256, 256>>>(d_pat,
								   PAT_THREADS);
		CK(cudaGetLastError());
		CK(cudaEventRecord(e1));
		CK(cudaDeviceSynchronize());
		CK(cudaEventElapsedTime(&ms, e0, e1));
		samples[i] = (double)ms;
	}
	{
		double s2[256];
		double a = 0;

		memcpy(s2, samples, sizeof(double) * o.n_warm);
		for (int i = 0; i < o.n_warm; i++)
			a += samples[i];
		printf("PERF launch_warm n=%d mean=%.3f p50=%.3f p95=%.3f\n",
		       o.n_warm, a / o.n_warm,
		       pct(s2, o.n_warm, 50), pct(s2, o.n_warm, 95));
	}

	/* P5/P9: compute kernel vs the bare-metal oracle (bit-exact). */
	float *d_fma = NULL, *h_fma;

	h_fma = (float *)malloc(sizeof(float) * FMA_THREADS);
	CK(cudaMalloc(&d_fma, sizeof(float) * FMA_THREADS));
	CK(cudaEventRecord(e0));
	fma_chain<<<FMA_THREADS / 256, 256>>>(d_fma, FMA_ITERS, FMA_THREADS);
	CK(cudaGetLastError());
	CK(cudaEventRecord(e1));
	CK(cudaDeviceSynchronize());
	CK(cudaEventElapsedTime(&ms, e0, e1));
	printf("PERF kernel_fma iters=%d ms=%.3f gflops=%.2f\n",
	       FMA_ITERS, (double)ms,
	       2.0 * (double)FMA_THREADS * (double)FMA_ITERS /
	       ((double)ms * 1.0e6));
	CK(cudaMemcpy(h_fma, d_fma, sizeof(float) * FMA_THREADS,
		      cudaMemcpyDeviceToHost));
	for (int i = 0; i < FMA_THREADS; i++)
		if (h_fma[i] != fma_ref)
			bad++;

	/* P10: steady-state throughput (enqueue burst, one sync). */
	CK(cudaEventRecord(e0));
	for (int i = 0; i < o.n_thru; i++) {
		pattern_kernel<<<(PAT_THREADS + 255) / 256, 256>>>(d_pat,
								   PAT_THREADS);
		CK(cudaGetLastError());
	}
	CK(cudaEventRecord(e1));
	CK(cudaDeviceSynchronize());
	CK(cudaEventElapsedTime(&ms, e0, e1));
	printf("PERF throughput n=%d wall_ms=%.3f kps=%.2f\n",
	       o.n_thru, (double)ms,
	       (double)o.n_thru / ((double)ms / 1.0e3));

	/* P7/P8: HtoD and DtoH copy bandwidth, and a small round-trip check. */
	{
		const size_t sizes[] = { 4096u, 262144u, 4194304u, 16777216u };
		unsigned char *h_a, *h_b, *d_a;

		h_a = (unsigned char *)malloc(sizes[3]);
		h_b = (unsigned char *)malloc(sizes[3]);
		if (!h_a || !h_b) {
			printf("PERF ERROR host alloc failed\n");
			printf("PERF STATUS FAIL\n");
			return 1;
		}
		CK(cudaMalloc((void **)&d_a, sizes[3]));
		for (int s = 0; s < 4; s++) {
			size_t n = sizes[s];
			double best_h = 1e30, best_d = 1e30;

			for (int i = 0; i < (int)n; i++)
				h_a[i] = (unsigned char)(i * 131 + s);
			for (int rep = 0; rep <= N_COPY; rep++) {
				CK(cudaEventRecord(e0));
				CK(cudaMemcpy(d_a, h_a, n, cudaMemcpyHostToDevice));
				CK(cudaEventRecord(e1));
				CK(cudaDeviceSynchronize());
				CK(cudaEventElapsedTime(&ms, e0, e1));
				if (rep > 0 && (double)ms < best_h)
					best_h = (double)ms;
			}
			for (int rep = 0; rep <= N_COPY; rep++) {
				CK(cudaEventRecord(e0));
				CK(cudaMemcpy(h_b, d_a, n, cudaMemcpyDeviceToHost));
				CK(cudaEventRecord(e1));
				CK(cudaDeviceSynchronize());
				CK(cudaEventElapsedTime(&ms, e0, e1));
				if (rep > 0 && (double)ms < best_d)
					best_d = (double)ms;
			}
			printf("PERF copy dir=htod bytes=%zu ms=%.4f gbps=%.3f\n",
			       n, best_h, (double)n / (best_h * 1.0e6));
			printf("PERF copy dir=dtoh bytes=%zu ms=%.4f gbps=%.3f\n",
			       n, best_d, (double)n / (best_d * 1.0e6));
			if (s == 0 && memcmp(h_a, h_b, n) != 0)
				copy_bad++;
		}
		free(h_a);
		free(h_b);
		cudaFree(d_a);
	}

	/* Optional clock64 cross-check (SM cycle counter, not wall time). */
	if (o.clock64) {
		unsigned long long *d_clk, *h_clk;

		h_clk = (unsigned long long *)malloc(
			sizeof(unsigned long long) * CLK_BLOCKS * 2);
		CK(cudaMalloc((void **)&d_clk,
			      sizeof(unsigned long long) * CLK_BLOCKS * 2));
		CK(cudaMemset(d_clk, 0,
			      sizeof(unsigned long long) * CLK_BLOCKS * 2));
		clock_probe<<<CLK_BLOCKS, 32>>>(d_clk);
		CK(cudaGetLastError());
		CK(cudaDeviceSynchronize());
		CK(cudaMemcpy(h_clk, d_clk,
			      sizeof(unsigned long long) * CLK_BLOCKS * 2,
			      cudaMemcpyDeviceToHost));
		{
			double cyc[CLK_BLOCKS];
			int n = 0;

			for (int i = 0; i < CLK_BLOCKS; i++) {
				if (h_clk[i * 2 + 1] > h_clk[i * 2] &&
				    h_clk[i * 2] != 0)
					cyc[n++] = (double)(h_clk[i * 2 + 1] -
							    h_clk[i * 2]);
			}
			if (n > 0)
				printf("PERF clock64 available=1 min=%.0f median=%.0f\n",
				       pct(cyc, n, 0), pct(cyc, n, 50));
			else
				printf("PERF clock64 available=0 min=0 median=0\n");
		}
		free(h_clk);
		cudaFree(d_clk);
	} else {
		printf("PERF clock64 available=0 min=0 median=0\n");
	}

	printf("PERF checksum fma=%.9g pattern=%u copy_bad=%d bad=%d\n",
	       (double)fma_ref, host_pattern(PAT_THREADS - 1), copy_bad,
	       bad + pat_bad);
	printf("PERF STATUS %s\n",
	       (bad + pat_bad + copy_bad) == 0 ? "OK" : "FAIL");

	cudaEventDestroy(e0);
	cudaEventDestroy(e1);
	cudaFree(d_pat);
	cudaFree(d_fma);
	free(h_pat);
	free(h_fma);
	return (bad + pat_bad + copy_bad) == 0 ? 0 : 1;
}