// SPDX-License-Identifier: BSD-3-Clause
/*
 * victory_reduce.cu - block reduction to shared memory + global atomics.
 *
 * Every block reduces its own slice of the values 1..N with a shared-memory
 * tree reduction; each block's partial sum is then folded into one global
 * counter with 64-bit atomicAdd. The host verifies the exact closed form
 * N*(N+1)/2, so any lost or double-counted block fails the client.
 *
 * Output lines are stable and greppable:
 *
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   REDUCE n=<count> sum=<gpu_total> expected=<closed_form> bad=<0|1>
 *   STATUS OK|FAIL
 *
 * Build with a static CUDA runtime (--cudart=static) so the guest only needs
 * the driver's libcuda.so, not a CUDA toolkit installation:
 *
 *   nvcc -O2 --cudart=static -o victory_reduce victory_reduce.cu
 *
 * Then, from the guest:  /tmp/victory_reduce
 */

#include <cstdio>
#include <cstdlib>
#include <cstdint>
#include <cuda_runtime.h>

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

#define NTHREADS 256

__global__ void reduce_sum(const unsigned int *in, unsigned long long *out,
			   int n)
{
	__shared__ unsigned long long s[NTHREADS];
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	s[threadIdx.x] = (i < n) ? (unsigned long long)in[i] : 0ULL;
	__syncthreads();

	for (int off = NTHREADS / 2; off > 0; off >>= 1) {
		if (threadIdx.x < off)
			s[threadIdx.x] += s[threadIdx.x + off];
		__syncthreads();
	}

	if (threadIdx.x == 0)
		atomicAdd(out, s[0]);
}

int main(void)
{
	cudaDeviceProp prop;
	const int n = 1 << 20; /* 1M values */
	size_t bytes = (size_t)n * sizeof(unsigned int);
	unsigned int *h_in;
	unsigned long long h_sum = 0;
	unsigned int *d_in;
	unsigned long long *d_sum;
	unsigned long long expected, bad;
	int i;

	setvbuf(stdout, NULL, _IONBF, 0);

	CHECK(cudaGetDeviceProperties(&prop, 0));
	printf("DEVICE %s cc=%d.%d sm=%d mem=%llu\n", prop.name,
	       prop.major, prop.minor, prop.multiProcessorCount,
	       (unsigned long long)prop.totalGlobalMem);

	h_in = (unsigned int *)malloc(bytes);
	if (!h_in)
		return 1;
	for (i = 0; i < n; i++)
		h_in[i] = (unsigned int)(i + 1);

	CHECK(cudaMalloc(&d_in, bytes));
	CHECK(cudaMalloc(&d_sum, sizeof(unsigned long long)));
	CHECK(cudaMemcpy(d_in, h_in, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemset(d_sum, 0, sizeof(unsigned long long)));

	reduce_sum<<<(n + NTHREADS - 1) / NTHREADS, NTHREADS>>>(d_in, d_sum, n);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(&h_sum, d_sum, sizeof(h_sum), cudaMemcpyDeviceToHost));

	expected = (unsigned long long)n * (unsigned long long)(n + 1) / 2ULL;
	bad = (h_sum == expected) ? 0 : 1;

	printf("REDUCE n=%d sum=%llu expected=%llu bad=%llu\n", n, h_sum,
	       expected, bad);

	if (bad != 0) {
		printf("STATUS FAIL\n");
		return 1;
	}

	printf("STATUS OK\n");
	CHECK(cudaFree(d_in));
	CHECK(cudaFree(d_sum));
	free(h_in);
	return 0;
}