// SPDX-License-Identifier: BSD-3-Clause
/*
 * p2_mixed - mail + vadd + int reduction + 2D grid + 2 streams in one
 * process, freeing each phase before the next so the peak live device
 * footprint stays at the vadd phase (3 x 32 KB = 24 pages).
 *
 *   phase 1: device-authored mail via a parameter-buffer kernel
 *   phase 2: vectorAdd n=8192 (element-exact fp32)
 *   phase 3: 4096 -> 1 int reduction (shared memory + atomicAdd)
 *   phase 4: 2D grid hash kernel (40x20, block 16x16, partial tiles)
 *   phase 5: 2 streams, addv + scale per stream (exact fp32)
 */
#include "common.h"

#define MIX_MAIL "the little gpu man says p2 mixed"
#define VN 8192
#define RD_N 4096
#define G2_NX 40
#define G2_NY 20
#define ST_N 1024
#define ST_NS 2

__global__ void mail_kernel(char *out, const char *msg, int len)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < len)
		out[i] = msg[i];
}

__global__ void vadd(const float *a, const float *b, float *c, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		c[i] = a[i] + b[i];
}

__global__ void reduce_sum(const int *in, int *out, int n)
{
	__shared__ int sdata[256];
	int tid = threadIdx.x;
	int i = blockIdx.x * blockDim.x + tid;

	sdata[tid] = (i < n) ? in[i] : 0;
	__syncthreads();
	for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
		if (tid < s)
			sdata[tid] += sdata[tid + s];
		__syncthreads();
	}
	if (tid == 0)
		atomicAdd(out, sdata[0]);
}

__global__ void grid2d(int *out, int nx, int ny)
{
	int ix = blockIdx.x * blockDim.x + threadIdx.x;
	int iy = blockIdx.y * blockDim.y + threadIdx.y;

	if (ix < nx && iy < ny) {
		int idx = iy * nx + ix;

		out[idx] = 1000000 + idx;
	}
}

__global__ void addv(float *x, float v, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		x[i] = x[i] + v;
}

__global__ void scale(float *x, float s, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		x[i] = x[i] * s;
}

static int phase_mail(void)
{
	char h[128];
	char *d_out = NULL, *d_msg = NULL;
	int len = (int)strlen(MIX_MAIL);
	int bad = 0;

	if (cudaMalloc(&d_out, len) != cudaSuccess)
		return 1;
	if (cudaMalloc(&d_msg, len) != cudaSuccess) {
		cudaFree(d_out);
		return 1;
	}
	cudaMemcpy(d_msg, MIX_MAIL, len, cudaMemcpyHostToDevice);
	mail_kernel<<<1, 64>>>(d_out, d_msg, len);
	cudaGetLastError();
	cudaDeviceSynchronize();
	cudaMemcpy(h, d_out, len, cudaMemcpyDeviceToHost);
	bad = (memcmp(h, MIX_MAIL, len) != 0) ? 1 : 0;
	if (bad == 0)
		printf("MAIL %s\n", h);
	cudaFree(d_out);
	cudaFree(d_msg);
	printf("RESULT mixed mail bad=%d\n", bad);
	return bad;
}

static int phase_vadd(void)
{
	float *h_a = (float *)malloc(sizeof(float) * VN);
	float *h_b = (float *)malloc(sizeof(float) * VN);
	float *h_c = (float *)malloc(sizeof(float) * VN);
	float *d_a = NULL, *d_b = NULL, *d_c = NULL;
	int bad = 0;

	if (!h_a || !h_b || !h_c)
		return 1;
	for (int i = 0; i < VN; i++) {
		h_a[i] = (float)(i % 1000) * 0.5f;
		h_b[i] = (float)(i % 997) * 0.25f;
	}
	if (cudaMalloc(&d_a, sizeof(float) * VN) != cudaSuccess)
		return 1;
	if (cudaMalloc(&d_b, sizeof(float) * VN) != cudaSuccess)
		return 1;
	if (cudaMalloc(&d_c, sizeof(float) * VN) != cudaSuccess)
		return 1;
	cudaMemcpy(d_a, h_a, sizeof(float) * VN, cudaMemcpyHostToDevice);
	cudaMemcpy(d_b, h_b, sizeof(float) * VN, cudaMemcpyHostToDevice);
	vadd<<<(VN + 255) / 256, 256>>>(d_a, d_b, d_c, VN);
	cudaGetLastError();
	cudaDeviceSynchronize();
	cudaMemcpy(h_c, d_c, sizeof(float) * VN, cudaMemcpyDeviceToHost);
	for (int i = 0; i < VN; i++)
		if (h_c[i] != h_a[i] + h_b[i])
			bad++;
	cudaFree(d_a);
	cudaFree(d_b);
	cudaFree(d_c);
	free(h_a);
	free(h_b);
	free(h_c);
	printf("RESULT mixed vadd n=%d bad=%d\n", VN, bad);
	return bad;
}

static int phase_reduce(void)
{
	int *h_in = (int *)malloc(sizeof(int) * RD_N);
	int *d_in = NULL, *d_out = NULL;
	int got = 0, exp = 0, bad = 0;

	if (!h_in)
		return 1;
	for (int i = 0; i < RD_N; i++) {
		h_in[i] = i;
		exp += i;
	}
	if (cudaMalloc(&d_in, sizeof(int) * RD_N) != cudaSuccess)
		return 1;
	if (cudaMalloc(&d_out, sizeof(int)) != cudaSuccess)
		return 1;
	cudaMemcpy(d_in, h_in, sizeof(int) * RD_N, cudaMemcpyHostToDevice);
	cudaMemset(d_out, 0, sizeof(int));
	reduce_sum<<<(RD_N + 255) / 256, 256>>>(d_in, d_out, RD_N);
	cudaGetLastError();
	cudaDeviceSynchronize();
	cudaMemcpy(&got, d_out, sizeof(int), cudaMemcpyDeviceToHost);
	bad = (got != exp) ? 1 : 0;
	cudaFree(d_in);
	cudaFree(d_out);
	free(h_in);
	printf("RESULT mixed red n=%d sum=%d bad=%d\n", RD_N, got, bad);
	return bad;
}

static int phase_grid2d(void)
{
	int n = G2_NX * G2_NY;
	int *d_out = NULL, *got = (int *)malloc(sizeof(int) * n);
	int bad = 0;

	if (!got)
		return 1;
	if (cudaMalloc(&d_out, sizeof(int) * n) != cudaSuccess) {
		free(got);
		return 1;
	}
	grid2d<<<dim3(5, 2), dim3(16, 16)>>>(d_out, G2_NX, G2_NY);
	cudaGetLastError();
	cudaDeviceSynchronize();
	cudaMemcpy(got, d_out, sizeof(int) * n, cudaMemcpyDeviceToHost);
	for (int i = 0; i < n; i++)
		if (got[i] != 1000000 + i)
			bad++;
	cudaFree(d_out);
	free(got);
	printf("RESULT mixed grid2d nx=%d ny=%d bad=%d\n", G2_NX, G2_NY,
	       bad);
	return bad;
}

static int phase_streams(void)
{
	float *h = (float *)malloc(sizeof(float) * ST_N);
	float *got = (float *)malloc(sizeof(float) * ST_N);
	cudaStream_t st[ST_NS];
	float *d_x[ST_NS];
	int bad = 0;

	if (!h || !got)
		return 1;
	for (int i = 0; i < ST_N; i++)
		h[i] = (float)(i % 100);
	for (int k = 0; k < ST_NS; k++) {
		if (cudaStreamCreate(&st[k]) != cudaSuccess)
			return 1;
		if (cudaMalloc(&d_x[k], sizeof(float) * ST_N) != cudaSuccess)
			return 1;
		cudaMemcpyAsync(d_x[k], h, sizeof(float) * ST_N,
				cudaMemcpyHostToDevice, st[k]);
		addv<<<(ST_N + 255) / 256, 256, 0, st[k]>>>(d_x[k], 7.0f, ST_N);
		scale<<<(ST_N + 255) / 256, 256, 0, st[k]>>>(d_x[k], 3.0f, ST_N);
		cudaGetLastError();
	}
	for (int k = 0; k < ST_NS; k++) {
		cudaStreamSynchronize(st[k]);
		cudaMemcpy(got, d_x[k], sizeof(float) * ST_N,
			   cudaMemcpyDeviceToHost);
		for (int i = 0; i < ST_N; i++)
			if (got[i] != ((float)(i % 100) + 7.0f) * 3.0f)
				bad++;
		cudaStreamDestroy(st[k]);
		cudaFree(d_x[k]);
	}
	free(h);
	free(got);
	printf("RESULT mixed stream nstreams=%d n=%d bad=%d\n", ST_NS, ST_N,
	       bad);
	return bad;
}

int main(int argc, char **argv)
{
	if (p2_emit_mode(argc, argv)) {
		printf("MAIL %s\n", MIX_MAIL);
		printf("RESULT mixed mail bad=0\n");
		printf("RESULT mixed vadd n=%d bad=0\n", VN);
		printf("RESULT mixed red n=%d sum=%d bad=0\n", RD_N,
		       RD_N * (RD_N - 1) / 2);
		printf("RESULT mixed grid2d nx=%d ny=%d bad=0\n", G2_NX, G2_NY);
		printf("RESULT mixed stream nstreams=%d n=%d bad=0\n", ST_NS,
		       ST_N);
		printf("RESULT mixed total bad=0\n");
		return p2_status(0);
	}

	p2_device_info();

	int bad = 0;

	bad += phase_mail();
	bad += phase_vadd();
	bad += phase_reduce();
	bad += phase_grid2d();
	bad += phase_streams();

	printf("RESULT mixed total bad=%d\n", bad);
	return p2_status(bad);
}