// SPDX-License-Identifier: BSD-3-Clause
/*
 * p2_shapes - grid/block shape variation with a bound-checked index-hash
 * kernel:
 *   case 1: 1x1x1 grid, 1x1x1 block (minimum launch)
 *   case 2: odd sizes        (n=1013, block 33, 31 blocks)
 *   case 3: non-multiple n   (n=1000, block 64, partial last block)
 *   case 4: 2D grid          (40x20, block 16x16, partial tiles)
 *   case 5: 3D grid          (24x8x8, block 8x4x2)
 *
 * Device memory: one 8 KB int buffer reused across cases (2 pages).
 */
#include "common.h"

#define MAXN 2048

struct shape {
	int nx, ny, nz;   /* logical extent    */
	int bx, by, bz;   /* block dims        */
	int gx, gy, gz;   /* grid dims         */
	int tag;
};

static const struct shape cases[] = {
	{   1,  1,  1,   1,  1,  1,   1,  1,  1, 1 },
	{1013,  1,  1,  33,  1,  1,  31,  1,  1, 2 },
	{1000,  1,  1,  64,  1,  1,  16,  1,  1, 3 },
	{  40, 20,  1,  16, 16,  1,   5,  2,  1, 4 },
	{  24,  8,  8,   8,  4,  2,   3,  2,  4, 5 },
};

#define NCASES ((int)(sizeof(cases) / sizeof(cases[0])))

__global__ void shape_kernel(int *out, int nx, int ny, int nz, int tag)
{
	int ix = blockIdx.x * blockDim.x + threadIdx.x;
	int iy = blockIdx.y * blockDim.y + threadIdx.y;
	int iz = blockIdx.z * blockDim.z + threadIdx.z;

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

		out[idx] = tag * 1000000 + idx;
	}
}

static void shape_expected(int *exp, const struct shape *s)
{
	for (int iz = 0; iz < s->nz; iz++)
		for (int iy = 0; iy < s->ny; iy++)
			for (int ix = 0; ix < s->nx; ix++) {
				int idx = (iz * s->ny + iy) * s->nx + ix;

				exp[idx] = s->tag * 1000000 + idx;
			}
}

int main(int argc, char **argv)
{
	int bad = 0;

	if (p2_emit_mode(argc, argv)) {
		int *exp = (int *)malloc(sizeof(int) * MAXN);

		for (int c = 0; c < NCASES; c++) {
			int n = cases[c].nx * cases[c].ny * cases[c].nz;

			shape_expected(exp, &cases[c]);
			printf("RESULT shapes case=%d n=%d bad=0\n",
			       cases[c].tag, n);
		}
		printf("RESULT shapes cases=%d bad=0\n", NCASES);
		free(exp);
		return p2_status(0);
	}

	p2_device_info();

	int *d_out = NULL, *h_got = (int *)malloc(sizeof(int) * MAXN);
	int *h_exp = (int *)malloc(sizeof(int) * MAXN);

	if (!h_got || !h_exp)
		return 1;
	P2_CHECK(cudaMalloc(&d_out, sizeof(int) * MAXN));

	for (int c = 0; c < NCASES; c++) {
		const struct shape *s = &cases[c];
		int n = s->nx * s->ny * s->nz;
		dim3 grid(s->gx, s->gy, s->gz), block(s->bx, s->by, s->bz);
		int sbad = 0;

		shape_kernel<<<grid, block>>>(d_out, s->nx, s->ny, s->nz,
					      s->tag);
		P2_CHECK(cudaGetLastError());
		P2_CHECK(cudaDeviceSynchronize());
		P2_CHECK(cudaMemcpy(h_got, d_out, sizeof(int) * n,
				    cudaMemcpyDeviceToHost));
		memset(h_exp, 0, sizeof(int) * n);
		shape_expected(h_exp, s);
		for (int i = 0; i < n; i++)
			if (h_got[i] != h_exp[i])
				sbad++;
		printf("RESULT shapes case=%d n=%d bad=%d\n", s->tag, n, sbad);
		bad += sbad;
	}
	printf("RESULT shapes cases=%d bad=%d\n", NCASES, bad);

	cudaFree(d_out);
	free(h_got); free(h_exp);
	return p2_status(bad);
}