// SPDX-License-Identifier: BSD-3-Clause
/*
 * p2_regs - register-usage / launch-bounds variants.
 *
 * Three kernels cover different register-allocation regimes while staying
 * exact (unsigned wraparound is well defined; no floating point):
 *
 *   regs_lo     __launch_bounds__(256, 4) - strongly capped register use
 *               (<= 64 registers/thread on the 65536-register file)
 *   regs_hi     64 independent accumulators x 128 update steps - high
 *               register pressure (the compiler may spill; correctness is
 *               unaffected by where the values live)
 *   regs_bounds __launch_bounds__(128, 4) - a second occupancy target with a
 *               different block size
 *
 * Every kernel writes one unsigned per element of a 1024-element buffer; the
 * CPU mirrors the exact integer expressions, so the buffer must match
 * element for element. INFO lines report the compiler's register counts and
 * the queried occupancy (informational; ignored by the A/B comparison).
 *
 * Device memory: one 4 KB unsigned buffer (1 page).
 */
#include "common.h"

#define RG_N 1024

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

	if (i < n)
		out[i] = ((unsigned)i ^ 0x5a5au) + 3u;
}

__global__ void regs_hi(unsigned *out, int n, unsigned seed)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;
	unsigned acc[64];
	unsigned s = 0;

	if (i >= n)
		return;
#pragma unroll
	for (int k = 0; k < 64; k++)
		acc[k] = seed + (unsigned)k * ((unsigned)i + 1u);
	for (int t = 0; t < 128; t++) {
#pragma unroll
		for (int k = 0; k < 64; k++)
			acc[k] = acc[k] * 2654435761u + ((unsigned)t ^ (unsigned)k);
	}
#pragma unroll
	for (int k = 0; k < 64; k++)
		s += acc[k];
	out[i] = s;
}

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

	if (i < n)
		out[i] = (unsigned)(i * 5) + 7u;
}

static unsigned cpu_lo(int i)
{
	return ((unsigned)i ^ 0x5a5au) + 3u;
}

static unsigned cpu_hi(int i, unsigned seed)
{
	unsigned acc[64];
	unsigned s = 0;

	for (int k = 0; k < 64; k++)
		acc[k] = seed + (unsigned)k * ((unsigned)i + 1u);
	for (int t = 0; t < 128; t++)
		for (int k = 0; k < 64; k++)
			acc[k] = acc[k] * 2654435761u + ((unsigned)t ^ (unsigned)k);
	for (int k = 0; k < 64; k++)
		s += acc[k];
	return s;
}

static unsigned cpu_bounds(int i)
{
	return (unsigned)(i * 5) + 7u;
}

template <typename T>
static void info_attrs(const char *name, T fn, int block)
{
	cudaFuncAttributes a;
	int occ = -1;

	if (cudaFuncGetAttributes(&a, fn) != cudaSuccess)
		return;
	(void)cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occ, fn, block, 0);
	printf("INFO regs %s numregs=%d maxthreads=%d occ_blocks_per_sm=%d\n",
	       name, a.numRegs, a.maxThreadsPerBlock, occ);
}

int main(int argc, char **argv)
{
	static const unsigned seed = 0x12345678u;
	int bad = 0;

	if (p2_emit_mode(argc, argv)) {
		printf("RESULT regs kernel=lo n=%d bad=0\n", RG_N);
		printf("RESULT regs kernel=hi n=%d bad=0\n", RG_N);
		printf("RESULT regs kernel=bounds n=%d bad=0\n", RG_N);
		printf("RESULT regs kernels=3 bad=0\n");
		return p2_status(0);
	}

	p2_device_info();

	unsigned *d_out = NULL;
	unsigned *got = (unsigned *)malloc(sizeof(unsigned) * RG_N);

	if (!got)
		return 1;
	P2_CHECK(cudaMalloc(&d_out, sizeof(unsigned) * RG_N));

	regs_lo<<<(RG_N + 255) / 256, 256>>>(d_out, RG_N);
	P2_CHECK(cudaGetLastError());
	P2_CHECK(cudaDeviceSynchronize());
	P2_CHECK(cudaMemcpy(got, d_out, sizeof(unsigned) * RG_N,
			    cudaMemcpyDeviceToHost));
	{
		int sbad = 0;

		for (int i = 0; i < RG_N; i++)
			if (got[i] != cpu_lo(i))
				sbad++;
		printf("RESULT regs kernel=lo n=%d bad=%d\n", RG_N, sbad);
		bad += sbad;
	}

	regs_hi<<<(RG_N + 255) / 256, 256>>>(d_out, RG_N, seed);
	P2_CHECK(cudaGetLastError());
	P2_CHECK(cudaDeviceSynchronize());
	P2_CHECK(cudaMemcpy(got, d_out, sizeof(unsigned) * RG_N,
			    cudaMemcpyDeviceToHost));
	{
		int sbad = 0;

		for (int i = 0; i < RG_N; i++)
			if (got[i] != cpu_hi(i, seed))
				sbad++;
		printf("RESULT regs kernel=hi n=%d bad=%d\n", RG_N, sbad);
		bad += sbad;
	}

	regs_bounds<<<(RG_N + 127) / 128, 128>>>(d_out, RG_N);
	P2_CHECK(cudaGetLastError());
	P2_CHECK(cudaDeviceSynchronize());
	P2_CHECK(cudaMemcpy(got, d_out, sizeof(unsigned) * RG_N,
			    cudaMemcpyDeviceToHost));
	{
		int sbad = 0;

		for (int i = 0; i < RG_N; i++)
			if (got[i] != cpu_bounds(i))
				sbad++;
		printf("RESULT regs kernel=bounds n=%d bad=%d\n", RG_N, sbad);
		bad += sbad;
	}

	info_attrs("lo", regs_lo, 256);
	info_attrs("hi", regs_hi, 256);
	info_attrs("bounds", regs_bounds, 128);

	printf("RESULT regs kernels=3 bad=%d\n", bad);

	cudaFree(d_out);
	free(got);
	return p2_status(bad);
}