// SPDX-License-Identifier: BSD-3-Clause
/*
 * p2_seq5 - five sequential kernel launches in one process on the same
 * buffers, verifying the device state after every launch.
 *
 *   step 1: c = a + b
 *   step 2: c = c + a
 *   step 3: c = c + b
 *   step 4: c = c * 2
 *   step 5: c = c + 1
 *
 * Inputs are small integers so every operation is exact in fp32; the CPU
 * reference reproduces the same chain and must match element for element.
 * Device memory: 3 x 4 KB pointers (3 pages total).
 */
#include "common.h"

#define N 1024

__global__ void kadd(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 kmul2(float *c, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		c[i] = c[i] * 2.0f;
}

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

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

static void make_inputs(float *a, float *b)
{
	for (int i = 0; i < N; i++) {
		a[i] = (float)(i % 7);
		b[i] = (float)(i % 5);
	}
}

static void step_cpu(float *c, const float *a, const float *b, int step)
{
	for (int i = 0; i < N; i++) {
		switch (step) {
		case 1: c[i] = a[i] + b[i]; break;
		case 2: c[i] = c[i] + a[i]; break;
		case 3: c[i] = c[i] + b[i]; break;
		case 4: c[i] = c[i] * 2.0f; break;
		case 5: c[i] = c[i] + 1.0f; break;
		}
	}
}

int main(int argc, char **argv)
{
	float *h_a, *h_b, *h_c, *h_exp;
	int bad = 0;

	if (p2_emit_mode(argc, argv)) {
		h_a = (float *)malloc(sizeof(float) * N);
		h_b = (float *)malloc(sizeof(float) * N);
		h_exp = (float *)malloc(sizeof(float) * N);
		make_inputs(h_a, h_b);
		for (int step = 1; step <= 5; step++) {
			step_cpu(h_exp, h_a, h_b, step);
			printf("RESULT seq step=%d n=%d bad=0\n", step, N);
		}
		printf("RESULT seq launches=5 n=%d bad=0\n", N);
		free(h_a); free(h_b); free(h_exp);
		return p2_status(0);
	}

	p2_device_info();
	h_a = (float *)malloc(sizeof(float) * N);
	h_b = (float *)malloc(sizeof(float) * N);
	h_c = (float *)malloc(sizeof(float) * N);
	h_exp = (float *)malloc(sizeof(float) * N);
	if (!h_a || !h_b || !h_c || !h_exp)
		return 1;
	make_inputs(h_a, h_b);

	float *d_a = NULL, *d_b = NULL, *d_c = NULL;

	P2_CHECK(cudaMalloc(&d_a, sizeof(float) * N));
	P2_CHECK(cudaMalloc(&d_b, sizeof(float) * N));
	P2_CHECK(cudaMalloc(&d_c, sizeof(float) * N));
	P2_CHECK(cudaMemcpy(d_a, h_a, sizeof(float) * N, cudaMemcpyHostToDevice));
	P2_CHECK(cudaMemcpy(d_b, h_b, sizeof(float) * N, cudaMemcpyHostToDevice));

	for (int step = 1; step <= 5; step++) {
		int sbad = 0;

		switch (step) {
		case 1: kadd<<<(N + 255) / 256, 256>>>(d_a, d_b, d_c, N); break;
		case 2: kadd<<<(N + 255) / 256, 256>>>(d_c, d_a, d_c, N); break;
		case 3: kadd<<<(N + 255) / 256, 256>>>(d_c, d_b, d_c, N); break;
		case 4: kmul2<<<(N + 255) / 256, 256>>>(d_c, N); break;
		case 5: kaddv<<<(N + 255) / 256, 256>>>(d_c, 1.0f, N); break;
		}
		P2_CHECK(cudaGetLastError());
		P2_CHECK(cudaDeviceSynchronize());
		P2_CHECK(cudaMemcpy(h_c, d_c, sizeof(float) * N,
				    cudaMemcpyDeviceToHost));
		step_cpu(h_exp, h_a, h_b, step);
		for (int i = 0; i < N; i++)
			if (h_c[i] != h_exp[i])
				sbad++;
		printf("RESULT seq step=%d n=%d bad=%d\n", step, N, sbad);
		bad += sbad;
	}
	printf("RESULT seq launches=5 n=%d bad=%d\n", N, bad);

	cudaFree(d_a); cudaFree(d_b); cudaFree(d_c);
	free(h_a); free(h_b); free(h_c); free(h_exp);
	return p2_status(bad);
}