// SPDX-License-Identifier: BSD-3-Clause
/*
 * frame_producer.cu - render an animated frame on the guest GPU and publish
 * it into the ivshmem (kvmfr) BAR so the host can see a real GPU-rendered
 * frame from the guest.
 *
 * Path: CUDA kernel -> device buffer -> host buffer -> ivshmem BAR
 *       (/sys/bus/pci/devices/0000:00:03.0/resource2, the 128 MiB kvmfr
 *       device QEMU backs with /dev/shm/kvmfr-<uuid>).
 *
 * Layout in the BAR (offset 0x100000, clear of the kvmfr protocol header):
 *   struct frame_hdr { magic, version, width, height, stride, format, seq }
 *   then width*height*4 bytes of RGBA8.
 *
 * Usage: sudo ./frame_producer [frames] [width] [height] [fps]
 * Prints: FRAME seq=<n> ms=<render+copy ms> (one line per frame, first 5 and
 * every 30th) and FRAME_DONE frames=<n>.
 */

#include <cstdio>
#include <cstdlib>
#include <cstdint>
#include <cstring>
#include <ctime>
#include <fcntl.h>
#include <sys/mman.h>
#include <unistd.h>
#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__);        \
			return 1;                                        \
		}                                                        \
	} while (0)

#define FRAME_OFFSET  0x100000UL
#define HDR_MAGIC     0x4b564d46u	/* "FMVK" little-endian-ish */

struct frame_hdr {
	uint32_t magic;
	uint32_t version;
	uint32_t width;
	uint32_t height;
	uint32_t stride;
	uint32_t format;	/* 0 = RGBA8 */
	uint32_t seq;
	uint32_t reserved;
	uint64_t timestamp_ns;
};

__global__ void render_julia(uchar4 *fb, int w, int h, float t)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;

	if (x >= w || y >= h)
		return;

	float fx = (x - w * 0.5f) / (h * 0.45f);
	float fy = (y - h * 0.5f) / (h * 0.45f);
	float cr = 0.7885f * cosf(t);
	float ci = 0.7885f * sinf(t);
	float zr = fx, zi = fy;
	int it = 0;
	const int maxit = 96;

	while (zr * zr + zi * zi < 4.0f && it < maxit) {
		float nzr = zr * zr - zi * zi + cr;
		float nzi = 2.0f * zr * zi + ci;

		zr = nzr;
		zi = nzi;
		it++;
	}
	{
		float v = (float)it / (float)maxit;
		unsigned char r = (unsigned char)(255.0f * v);
		unsigned char g = (unsigned char)(255.0f * (1.0f - v) * v * 4.0f);
		unsigned char b = (unsigned char)(255.0f * (1.0f - v));
		unsigned char a = 255;

		fb[y * w + x] = make_uchar4(r, g, b, a);
	}
}

static uint64_t now_ns(void)
{
	struct timespec ts;

	clock_gettime(CLOCK_MONOTONIC, &ts);
	return (uint64_t)ts.tv_sec * 1000000000ull + ts.tv_nsec;
}

int main(int argc, char **argv)
{
	int frames = argc > 1 ? atoi(argv[1]) : 60;
	int w = argc > 2 ? atoi(argv[2]) : 800;
	int h = argc > 3 ? atoi(argv[3]) : 600;
	int fps = argc > 4 ? atoi(argv[4]) : 10;
	const char *bar_path =
		"/sys/bus/pci/devices/0000:00:03.0/resource2";
	size_t fb_bytes = (size_t)w * h * 4;
	uchar4 *d_fb = NULL, *h_fb = NULL;
	unsigned char *bar;
	struct frame_hdr *hdr;
	volatile uint32_t *seqp;
	int fd, i;

	setvbuf(stdout, NULL, _IONBF, 0);

	fd = open(bar_path, O_RDWR | O_SYNC);
	if (fd < 0) {
		printf("OPEN_FAIL %s\n", bar_path);
		return 1;
	}
	bar = (unsigned char *)mmap(NULL, 128UL << 20, PROT_READ | PROT_WRITE,
				    MAP_SHARED, fd, 0);
	if (bar == MAP_FAILED) {
		printf("MMAP_FAIL\n");
		return 1;
	}
	hdr = (struct frame_hdr *)(bar + FRAME_OFFSET);
	seqp = &hdr->seq;

	CHECK(cudaMalloc(&d_fb, fb_bytes));
	h_fb = (uchar4 *)malloc(fb_bytes);
	if (!h_fb) {
		printf("HOST_ALLOC_FAIL\n");
		return 1;
	}

	hdr->magic = HDR_MAGIC;
	hdr->version = 1;
	hdr->width = (uint32_t)w;
	hdr->height = (uint32_t)h;
	hdr->stride = (uint32_t)w * 4;
	hdr->format = 0;

	for (i = 0; i < frames; i++) {
		uint64_t t0 = now_ns();
		float phase = (float)i * 0.15f;
		dim3 block(16, 16);
		dim3 grid((w + 15) / 16, (h + 15) / 16);

		render_julia<<<grid, block>>>(d_fb, w, h, phase);
		CHECK(cudaGetLastError());
		CHECK(cudaDeviceSynchronize());
		CHECK(cudaMemcpy(h_fb, d_fb, fb_bytes,
				 cudaMemcpyDeviceToHost));
		memcpy(bar + FRAME_OFFSET + sizeof(struct frame_hdr), h_fb,
		       fb_bytes);
		hdr->timestamp_ns = now_ns();
		__sync_synchronize();
		*seqp = (uint32_t)(i + 1);
		__sync_synchronize();
		if (i < 5 || (i % 30) == 29)
			printf("FRAME seq=%d ms=%.1f\n", i + 1,
			       (double)(now_ns() - t0) / 1e6);
		if (fps > 0)
			usleep((useconds_t)(1000000 / fps));
	}
	printf("FRAME_DONE frames=%d %dx%d\n", frames, w, h);
	CHECK(cudaFree(d_fb));
	free(h_fb);
	munmap(bar, 128UL << 20);
	close(fd);
	return 0;
}