// SPDX-License-Identifier: BSD-3-Clause
/*
 * ip_probe.cu - how many CTAs actually executed for an in-place kernel?
 *
 * scale_mark marks progress[blockIdx.x]=0xCAFEBABE after the c store; D2H of
 * progress counts executed CTAs exactly. c (the in-place buffer) is read back
 * and verified. seed (argv[1]) makes the run content-unique.
 *
 * Build: nvcc -O2 --cudart=static -o ip_probe ip_probe.cu
 */

#include <cstdio>
#include <cstdlib>
#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)

__global__ void vadd(const float *a, const float *b, float *c, int n)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		c[i] = a[i] + b[i];
}

__global__ void scale_mark(float *c, float s, int n, unsigned *prog)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i < n)
		c[i] = c[i] * s;
	if (threadIdx.x == 0)
		prog[blockIdx.x] = 0xCAFEBABEu;
}

int main(int argc, char **argv)
{
	cudaDeviceProp prop;
	const int n = 1 << 15;
	const int grid = (n + 255) / 256;
	size_t bytes = (size_t)n * sizeof(float);
	float k = argc > 1 ? atof(argv[1]) : 0.0f;
	float *ha, *hb, *hc;
	float *da, *db, *dc;
	unsigned *d_prog, h_prog[128];
	int i;

	setvbuf(stdout, NULL, _IONBF, 0);
	CHECK(cudaGetDeviceProperties(&prop, 0));
	printf("DEVICE %s seed=%g grid=%d\n", prop.name, (double)k, grid);

	ha = (float *)malloc(bytes);
	hb = (float *)malloc(bytes);
	hc = (float *)malloc(bytes);
	for (i = 0; i < n; i++) {
		ha[i] = (float)(i % 1000) * 0.5f + k;
		hb[i] = (float)(i % 997) * 0.25f + 2.0f * k;
	}
	memset(hc, 0xCD, bytes);

	CHECK(cudaMalloc(&da, bytes));
	CHECK(cudaMalloc(&db, bytes));
	CHECK(cudaMalloc(&dc, bytes));
	CHECK(cudaMalloc(&d_prog, 128 * sizeof(unsigned)));
	printf("DPTRS da=%p db=%p dc=%p prog=%p\n", (void *)da, (void *)db,
	       (void *)dc, (void *)d_prog);

	CHECK(cudaMemcpy(da, ha, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemcpy(db, hb, bytes, cudaMemcpyHostToDevice));
	CHECK(cudaMemset(d_prog, 0xEE, 128 * sizeof(unsigned)));

	vadd<<<grid, 256>>>(da, db, dc, n);
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(hc, dc, bytes, cudaMemcpyDeviceToHost));
	{
		int b = 0;

		for (i = 0; i < n; i++)
			if (hc[i] != ha[i] + hb[i])
				b++;
		printf("VADD readback bad=%d/%d\n", b, n);
	}

	/* in-place scale with CTA markers */
	scale_mark<<<grid, 256>>>(dc, 2.0f, n, d_prog);
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_prog, d_prog, 128 * sizeof(unsigned),
			 cudaMemcpyDeviceToHost));
	{
		int marked = 0, first_missing = -1;

		for (i = 0; i < grid; i++) {
			if (h_prog[i] == 0xCAFEBABEu)
				marked++;
			else if (first_missing < 0)
				first_missing = i;
		}
		printf("CTA-MARK marked=%d/%d first_missing=%d\n", marked, grid,
		       first_missing);
	}
	CHECK(cudaMemcpy(hc, dc, bytes, cudaMemcpyDeviceToHost));
	{
		int b = 0, first = -1;

		for (i = 0; i < n; i++)
			if (hc[i] != (ha[i] + hb[i]) * 2.0f) {
				b++;
				if (first < 0)
					first = i;
			}
		printf("SCALE readback bad=%d/%d first=%d\n", b, n, first);
		if (first >= 0)
			printf("SCALE firstbad got=%g want=%g\n", hc[first],
			       (ha[first] + hb[first]) * 2.0f);
	}
	printf("STATUS DONE\n");
	return 0;
}