// SPDX-License-Identifier: BSD-3-Clause
/*
 * victory_matmul.cu - 256x256 tiled matrix multiply with shared memory.
 *
 * C = A * B for three 256x256 float matrices of small integer values. The
 * inputs are integers in [0, 10), so every product and every partial sum is
 * exactly representable in float (< 2^24); whatever order the GPU accumulates
 * in, the result is exact. The host replays the multiply with 64-bit integers
 * and verifies every element exactly, plus a checksum over C.
 *
 * Output lines are stable and greppable:
 *
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   MATMUL n=256 bad=<mismatches> csum=<sum>
 *   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_matmul victory_matmul.cu
 *
 * Then, from the guest:  /tmp/victory_matmul
 */

#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 TILE 16
#define N 256

__global__ void matmul_tiled(const float *A, const float *B, float *C, int n)
{
	__shared__ float As[TILE][TILE];
	__shared__ float Bs[TILE][TILE];
	int row = blockIdx.y * TILE + threadIdx.y;
	int col = blockIdx.x * TILE + threadIdx.x;
	float acc = 0.0f;

	for (int t = 0; t < n / TILE; t++) {
		As[threadIdx.y][threadIdx.x] = A[row * n + t * TILE + threadIdx.x];
		Bs[threadIdx.y][threadIdx.x] =
			B[(t * TILE + threadIdx.y) * n + col];
		__syncthreads();

		for (int k = 0; k < TILE; k++)
			acc += As[threadIdx.y][k] * Bs[k][threadIdx.x];
		__syncthreads();
	}

	C[row * n + col] = acc;
}

int main(void)
{
	cudaDeviceProp prop;
	const int n = N;
	const size_t bytes = (size_t)n * n * sizeof(float);
	float *h_a, *h_b, *h_c;
	float *d_a, *d_b, *d_c;
	double csum = 0.0;
	int bad = 0, i, j, k;
	dim3 grid(n / TILE, n / TILE);
	dim3 block(TILE, TILE);

	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_a = (float *)malloc(bytes);
	h_b = (float *)malloc(bytes);
	h_c = (float *)malloc(bytes);
	if (!h_a || !h_b || !h_c)
		return 1;

	for (i = 0; i < n * n; i++) {
		h_a[i] = (float)(i % 10);
		h_b[i] = (float)((i * 3 + 1) % 10);
	}

	CHECK(cudaMalloc(&d_a, bytes));
	CHECK(cudaMalloc(&d_b, bytes));
	CHECK(cudaMalloc(&d_c, bytes));
	CHECK(cudaMemcpy(d_a, h_a, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(d_b, h_b, bytes, cudaMemcpyHostToDevice));

	matmul_tiled<<<grid, block>>>(d_a, d_b, d_c, n);
	CHECK(cudaGetLastError());
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_c, d_c, bytes, cudaMemcpyDeviceToHost));

	for (i = 0; i < n; i++) {
		for (j = 0; j < n; j++) {
			long long ref = 0;

			for (k = 0; k < n; k++)
				ref += (long long)h_a[i * n + k] *
				       (long long)h_b[k * n + j];
			if (h_c[i * n + j] != (float)ref)
				bad++;
			csum += (double)h_c[i * n + j];
		}
	}

	printf("MATMUL n=%d bad=%d csum=%.6f\n", n, bad, csum);

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

	printf("STATUS OK\n");
	CHECK(cudaFree(d_a));
	CHECK(cudaFree(d_b));
	CHECK(cudaFree(d_c));
	free(h_a);
	free(h_b);
	free(h_c);
	return 0;
}