// Dense mul_mat check for NVFP4: tiled path (M=64 columns) vs vec_dot path (M=1) vs float reference.
#include "ggml.h"
#include "ggml-cpu.h"
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

void ggml_cpu_init(void);
void dequantize_row_nvfp4(const void * x, float * y, int64_t k);

static uint32_t st = 777;
static uint32_t rnd(void) { st = st * 1664525u + 1013904223u; return st >> 8; }
static float frnd(void) { return ((int) (rnd() % 20001) - 10000) / 10000.0f; }

int main(int argc, char ** argv) {
    const int K = argc > 1 ? atoi(argv[1]) : 2048;
    const int N = argc > 2 ? atoi(argv[2]) : 512;
    const int M = 64;

    ggml_cpu_init();
    struct ggml_init_params ip = { .mem_size = 512u * 1024 * 1024, .mem_buffer = NULL, .no_alloc = false };
    struct ggml_context * ctx = ggml_init(ip);

    struct ggml_tensor * w  = ggml_new_tensor_2d(ctx, GGML_TYPE_NVFP4, K, N);
    struct ggml_tensor * xM = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, M);
    struct ggml_tensor * x1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, 1);

    // random valid nvfp4 blocks (UE4M3 codes below the NaN code)
    uint8_t * wd = (uint8_t *) w->data;
    const size_t blocks = (size_t) N * (K / 64);
    for (size_t b = 0; b < blocks; b++) {
        uint8_t * blk = wd + b * 36;
        for (int s = 0; s < 4; s++) blk[s] = (uint8_t) (rnd() % 0x70);
        for (int j = 0; j < 32; j++) blk[4 + j] = (uint8_t) rnd();
    }
    float * xd = (float *) xM->data;
    for (int m = 0; m < M; m++) for (int k = 0; k < K; k++) xd[(size_t) m * K + k] = frnd();
    for (int k = 0; k < K; k++) ((float *) x1->data)[k] = xd[k];  // column 0

    struct ggml_tensor * yM = ggml_mul_mat(ctx, w, xM);
    struct ggml_tensor * y1 = ggml_mul_mat(ctx, w, x1);
    struct ggml_cgraph * gM = ggml_new_graph(ctx);
    ggml_build_forward_expand(gM, yM);
    ggml_graph_compute_with_ctx(ctx, gM, 16);
    struct ggml_cgraph * g1 = ggml_new_graph(ctx);
    ggml_build_forward_expand(g1, y1);
    ggml_graph_compute_with_ctx(ctx, g1, 16);

    // float reference from the dequantized weights
    float * wf = malloc(sizeof(float) * (size_t) K * N);
    for (int n = 0; n < N; n++) dequantize_row_nvfp4(wd + (size_t) n * (K / 64) * 36, wf + (size_t) n * K, K);

    double e_tiled = 0, n_tiled = 0, e_ref = 0, n_ref = 0, e_ref1 = 0;
    const float * yMd = (const float *) yM->data;
    const float * y1d = (const float *) y1->data;
    for (int n = 0; n < N; n++) {
        double ref = 0;
        for (int k = 0; k < K; k++) ref += (double) wf[(size_t) n * K + k] * xd[k];
        // column 0 of the tiled result vs the vec_dot result
        const double a = yMd[n], b = y1d[n];
        e_tiled += (a - b) * (a - b); n_tiled += b * b;
        e_ref += (b - ref) * (b - ref); n_ref += ref * ref;
        for (int m = 1; m < M; m++) {
            double r = 0;
            for (int k = 0; k < K; k++) r += (double) wf[(size_t) n * K + k] * xd[(size_t) m * K + k];
            const double t = yMd[(size_t) m * N + n];
            e_ref1 += (t - r) * (t - r);
            n_ref  += r * r;
        }
    }
    printf("K=%d N=%d M=%d\n", K, N, M);
    printf("tiled(M=64) vs vec_dot(M=1), col0: rel L2 = %.3e\n", sqrt(e_tiled / n_tiled));
    printf("vec_dot(M=1) vs f32 ref (col0):    rel L2 = %.3e\n", sqrt(e_ref / n_ref));
    printf("tiled(M=64) vs f32 ref (cols 1..): rel L2 = %.3e\n", sqrt(e_ref1 / n_ref));
    free(wf);
    ggml_free(ctx);
    return 0;
}