// SPDX-License-Identifier: BSD-3-Clause
/*
 * p2_smem - shared-memory kernels:
 *   reduce_sum: 4096 int32 -> single sum, 16 blocks x 256 threads, shared
 *               tree reduction plus a global atomicAdd.
 *   gemm_tiled: 32x32 int32 GEMM with 16x16 shared-memory tiles,
 *               grid(2,2) x block(16,16).
 *
 * All arithmetic is exact in int32 (inputs are small non-negative ints), so
 * device results must equal the CPU reference bit for bit.
 * Device memory: in 16 KB, out 4 B, then A/B/C 4 KB each (peak 5 pages).
 */
#include "common.h"

#define R_N 4096
#define TS 16
#define M 32
#define K 32
#define NN 32

__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 gemm_tiled(const int *A, const int *B, int *C)
{
	__shared__ int As[TS][TS], Bs[TS][TS];
	int row = blockIdx.y * TS + threadIdx.y;
	int col = blockIdx.x * TS + threadIdx.x;
	int sum = 0;

	for (int t = 0; t < K / TS; t++) {
		As[threadIdx.y][threadIdx.x] = A[row * K + t * TS + threadIdx.x];
		Bs[threadIdx.y][threadIdx.x] = B[(t * TS + threadIdx.y) * NN + col];
		__syncthreads();
		for (int j = 0; j < TS; j++)
			sum += As[threadIdx.y][j] * Bs[j][threadIdx.x];
		__syncthreads();
	}
	C[row * NN + col] = sum;
}

int main(int argc, char **argv)
{
	int *h_in = (int *)malloc(sizeof(int) * R_N);
	int *h_A = (int *)malloc(sizeof(int) * M * K);
	int *h_B = (int *)malloc(sizeof(int) * K * NN);
	int *h_C = (int *)malloc(sizeof(int) * M * NN);
	int red_sum = 0, gemm_sum = 0, bad = 0;

	if (!h_in || !h_A || !h_B || !h_C)
		return 1;
	for (int i = 0; i < R_N; i++)
		h_in[i] = i;
	for (int i = 0; i < M * K; i++)
		h_A[i] = i % 5;
	for (int i = 0; i < K * NN; i++)
		h_B[i] = i % 3;
	for (int i = 0; i < R_N; i++)
		red_sum += h_in[i];
	for (int r = 0; r < M; r++)
		for (int c = 0; c < NN; c++) {
			int s = 0;

			for (int t = 0; t < K; t++)
				s += h_A[r * K + t] * h_B[t * NN + c];
			h_C[r * NN + c] = s;
			gemm_sum += s;
		}

	if (p2_emit_mode(argc, argv)) {
		printf("RESULT smem reduce n=%d sum=%d bad=0\n", R_N, red_sum);
		printf("RESULT smem gemm m=%d n=%d k=%d checksum=%d bad=0\n",
		       M, NN, K, gemm_sum);
		free(h_in); free(h_A); free(h_B); free(h_C);
		return p2_status(0);
	}

	p2_device_info();

	int *d_in = NULL, *d_out = NULL, *d_A = NULL, *d_B = NULL, *d_C = NULL;
	int got_sum = 0, got_gemm = 0;

	P2_CHECK(cudaMalloc(&d_in, sizeof(int) * R_N));
	P2_CHECK(cudaMalloc(&d_out, sizeof(int)));
	P2_CHECK(cudaMemcpy(d_in, h_in, sizeof(int) * R_N,
			    cudaMemcpyHostToDevice));
	P2_CHECK(cudaMemset(d_out, 0, sizeof(int)));
	reduce_sum<<<(R_N + 255) / 256, 256>>>(d_in, d_out, R_N);
	P2_CHECK(cudaGetLastError());
	P2_CHECK(cudaDeviceSynchronize());
	P2_CHECK(cudaMemcpy(&got_sum, d_out, sizeof(int),
			    cudaMemcpyDeviceToHost));
	printf("RESULT smem reduce n=%d sum=%d bad=%d\n", R_N, got_sum,
	       got_sum == red_sum ? 0 : 1);
	bad += (got_sum == red_sum) ? 0 : 1;
	cudaFree(d_in); cudaFree(d_out);

	P2_CHECK(cudaMalloc(&d_A, sizeof(int) * M * K));
	P2_CHECK(cudaMalloc(&d_B, sizeof(int) * K * NN));
	P2_CHECK(cudaMalloc(&d_C, sizeof(int) * M * NN));
	P2_CHECK(cudaMemcpy(d_A, h_A, sizeof(int) * M * K,
			    cudaMemcpyHostToDevice));
	P2_CHECK(cudaMemcpy(d_B, h_B, sizeof(int) * K * NN,
			    cudaMemcpyHostToDevice));
	gemm_tiled<<<dim3(2, 2), dim3(TS, TS)>>>(d_A, d_B, d_C);
	P2_CHECK(cudaGetLastError());
	P2_CHECK(cudaDeviceSynchronize());

	int *got_C = (int *)malloc(sizeof(int) * M * NN);

	if (!got_C)
		return 1;
	P2_CHECK(cudaMemcpy(got_C, d_C, sizeof(int) * M * NN,
			    cudaMemcpyDeviceToHost));
	int cbad = 0;

	for (int i = 0; i < M * NN; i++) {
		got_gemm += got_C[i];
		if (got_C[i] != h_C[i])
			cbad++;
	}
	printf("RESULT smem gemm m=%d n=%d k=%d checksum=%d bad=%d\n",
	       M, NN, K, got_gemm, cbad);
	bad += cbad;

	cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
	free(got_C);
	free(h_in); free(h_A); free(h_B); free(h_C);
	return p2_status(bad);
}