// SPDX-License-Identifier: BSD-3-Clause
/*
 * victory_memset2d.cu - cudaMemset + pitched cudaMemcpy2D round trip.
 *
 * A 4096-column x 1024-row byte image is allocated with cudaMallocPitch
 * (row pitch >= width), then exercised three ways:
 *
 *   1. cudaMemset the whole allocation (including pitch padding) to 0xA5 and
 *      read it back row by row with cudaMemcpy2D; every byte must match.
 *   2. H2D cudaMemcpy2D of a host pattern, D2H cudaMemcpy2D of the same
 *      image; every byte must match the host pattern.
 *   3. cudaMemset2D of the upper-left quadrant to 0x3C; read back and verify
 *      that exactly that rectangle changed.
 *
 * Output lines are stable and greppable:
 *
 *   DEVICE <name> cc=<major>.<minor> sm=<count> mem=<bytes>
 *   MEMSET2D w=<width> h=<height> pitch=<pitch> bad=<mismatches>
 *   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_memset2d victory_memset2d.cu
 *
 * Then, from the guest:  /tmp/victory_memset2d
 */

#include <cstdio>
#include <cstdlib>
#include <cstdint>
#include <cstring>
#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 IMG_W 4096
#define IMG_H 1024
#define FILL_BYTE 0xA5
#define QUAD_BYTE 0x3C

int main(void)
{
	cudaDeviceProp prop;
	const size_t w = IMG_W, h = IMG_H;
	unsigned char *h_img, *h_back;
	unsigned char *d_img;
	size_t pitch = 0;
	int bad = 0, i, row, col;

	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_img = (unsigned char *)malloc(w * h);
	h_back = (unsigned char *)malloc(w * h);
	if (!h_img || !h_back)
		return 1;

	CHECK(cudaMallocPitch(&d_img, &pitch, w, h));

	/* 1. Full cudaMemset over the pitched allocation, read back with 2D. */
	CHECK(cudaMemset(d_img, FILL_BYTE, pitch * h));
	CHECK(cudaMemcpy2D(h_back, w, d_img, pitch, w, h,
			   cudaMemcpyDeviceToHost));
	for (i = 0; i < (int)(w * h); i++) {
		if (h_back[i] != (unsigned char)FILL_BYTE)
			bad++;
	}

	/* 2. Pitched H2D -> D2H round trip of a host pattern. */
	for (row = 0; row < (int)h; row++)
		for (col = 0; col < (int)w; col++)
			h_img[row * w + col] = (unsigned char)((row + col) & 0xFF);
	CHECK(cudaMemcpy2D(d_img, pitch, h_img, w, w, h,
			   cudaMemcpyHostToDevice));
	memset(h_back, 0, w * h);
	CHECK(cudaMemcpy2D(h_back, w, d_img, pitch, w, h,
			   cudaMemcpyDeviceToHost));
	for (i = 0; i < (int)(w * h); i++) {
		if (h_back[i] != h_img[i])
			bad++;
	}

	/* 3. cudaMemset2D of the upper-left quadrant only. */
	CHECK(cudaMemset2D(d_img, pitch, QUAD_BYTE, w / 2, h / 2));
	memset(h_back, 0, w * h);
	CHECK(cudaMemcpy2D(h_back, w, d_img, pitch, w, h,
			   cudaMemcpyDeviceToHost));
	for (row = 0; row < (int)h; row++) {
		for (col = 0; col < (int)w; col++) {
			unsigned char expect = h_img[row * w + col];

			if (row < (int)(h / 2) && col < (int)(w / 2))
				expect = QUAD_BYTE;
			if (h_back[row * w + col] != expect)
				bad++;
		}
	}

	printf("MEMSET2D w=%zu h=%zu pitch=%zu bad=%d\n", w, h, pitch, bad);

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

	printf("STATUS OK\n");
	CHECK(cudaFree(d_img));
	free(h_img);
	free(h_back);
	return 0;
}