// SPDX-License-Identifier: BSD-3-Clause
/*
 * ip_probe3.cu - what does the SM actually READ for the scale kernel?
 *
 * scale_report writes, per CTA: prog[2b]   = bits of the float param s,
 *                                  prog[2b+1] = bits of c[b*256] (the load).
 * The D2H of prog then shows exactly what the SM saw, independent of any
 * store-visibility path for c.
 *
 * Build: nvcc -O2 --cudart=static -o ip_probe3 ip_probe3.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_report(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[2 * blockIdx.x] = __float_as_uint(s);
		prog[2 * blockIdx.x + 1] = __float_as_uint(c[blockIdx.x * 256]);
	}
}

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[256];
	int i;

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

	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, 256 * 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, 256 * 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);
	}

	scale_report<<<grid, 256>>>(dc, 2.0f, n, d_prog);
	CHECK(cudaDeviceSynchronize());
	CHECK(cudaMemcpy(h_prog, d_prog, 256 * sizeof(unsigned),
			 cudaMemcpyDeviceToHost));
	{
		int s_ok = 0, s_1 = 0, c_ok = 0, n_seen = 0;
		unsigned first_s = 0, first_c = 0;

		for (i = 0; i < grid; i++) {
			unsigned sb = h_prog[2 * i];
			unsigned cb = h_prog[2 * i + 1];

			if (i == 0) {
				first_s = sb;
				first_c = cb;
			}
			if (sb == 0x40000000u)
				s_ok++;
			if (sb == 0x3F800000u)
				s_1++;
			{
				float want = ha[i * 256] + hb[i * 256];

				if (cb == *(unsigned *)&want)
					c_ok++;
			}
			n_seen++;
		}
		printf("REPORT ctas=%d s_eq_2.0=%d s_eq_1.0=%d c_load_ok=%d first_s=0x%08x first_c=0x%08x\n",
		       n_seen, s_ok, s_1, c_ok, first_s, first_c);
	}
	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);
	}
	printf("STATUS DONE\n");
	return 0;
}