// GPU vs CPU parity tests for the DT3 dual-plane ternary format
//
// The CPU path (dequantize_row_dt3) is the validated reference. This test
// checks the GPU backend against it in two steps:
//
//   1. dequantization: GET_ROWS on the GPU must reproduce the CPU reference
//      bit by bit — same fp16 scales, exact products by {-1, 0, +1}, one
//      float rounding per element on both sides.
//   2. matrix multiplication: MUL_MAT with a small number of destination
//      columns takes the MMVQ path (vec_dot_dt3_q8_1). The activations are
//      chosen so that their q8_1 quantization is exact (integer values with
//      amax 127 in every 32-element chunk), which makes a double precision
//      reference computed from the dequantized weights valid to float
//      rounding of the accumulation. One case is also checked against a
//      manual sum over trits stored by the test, with non-trivial qh trits.
//
// Errors are judged on the relative Frobenius norm, ||gpu - ref|| / ||ref||;
// the elementwise maximum is reported as information only, since it explodes
// on cancellation whenever a true output value is near zero.
//
// MUL_MAT with more destination columns than the MMVQ limit falls back to
// dequantization + cuBLAS GEMM, which on fast-fp16 hardware rounds the
// dequantized weights to fp16. DT3 weights (d1*t1 + d2*t2, the sum of two
// fp16-scaled terms) are generally NOT fp16-representable, so that path is
// judged against a reference computed from fp16-rounded weights (taking the
// better of the two references, so the test also passes when
// GGML_CUDA_CUBLAS_COMPUTE_TYPE=f32 disables the rounding). Q4_1 (same
// regime: d*q + m not fp16-exact) and Q4_0 (weights fp16-exact) go through
// the identical comparison as controls, reported but gated loosely.
//
// The directed blocks exercise the three packing regions, the 79/80 and
// 119/120 region boundaries, and negative scales. The random blocks use raw
// random bytes: every byte value 0..255 must decode identically on both
// sides, including values >= 243 that never come out of the packer.
//
// Without a GPU backend the test is skipped and succeeds.

#include "ggml.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml-cpu.h"

#undef NDEBUG
#include <assert.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <vector>

constexpr int    QK_DT3         = 128;
constexpr size_t DT3_QS_BYTES   = 24; // per plane
constexpr size_t DT3_QH_BYTES   = 2;  // per plane
constexpr size_t DT3_BLOCK_SIZE = 2*DT3_QS_BYTES + 2*DT3_QH_BYTES + 2*sizeof(uint16_t);

// byte offsets inside a block (spec: qs[2][24] | qh[2][2] | d[2])
constexpr size_t OFF_QS = 0;
constexpr size_t OFF_QH = 2*DT3_QS_BYTES;
constexpr size_t OFF_D  = 2*DT3_QS_BYTES + 2*DT3_QH_BYTES;

// independent packer, written from the format specification (same as in
// test-dt3.cpp): element i of a plane goes to
//   region A: qs[m],    m in [0,16), digit n: elements m + n*16     (0..79)
//   region B: qs[16+m], m in [0,8),  digit n: elements 80 + m + n*8 (80..119)
//   region C: qh[j],    j in [0,2),  digit n: elements 120 + j + n*2 (120..127)
static void ref_pack_plane(const int8_t * t, uint8_t * qs, uint8_t * qh) {
    for (int m = 0; m < 16; ++m) {
        uint32_t q = 0;
        for (int n = 0; n < 5; ++n) {
            q = q*3 + (uint32_t)(t[m + n*16] + 1);
        }
        qs[m] = (uint8_t)((q*256 + 242)/243);
    }
    for (int m = 0; m < 8; ++m) {
        uint32_t q = 0;
        for (int n = 0; n < 5; ++n) {
            q = q*3 + (uint32_t)(t[80 + m + n*8] + 1);
        }
        qs[16 + m] = (uint8_t)((q*256 + 242)/243);
    }
    for (int j = 0; j < 2; ++j) {
        uint32_t q = 0;
        for (int n = 0; n < 4; ++n) {
            q = q*3 + (uint32_t)(t[120 + j + n*2] + 1);
        }
        q *= 3; // shift the first value to the most significant trit
        qh[j] = (uint8_t)((q*256 + 242)/243);
    }
}

static void ref_pack_block(const int8_t * t1, float d1, const int8_t * t2, float d2, uint8_t * block) {
    ref_pack_plane(t1, block + OFF_QS,                block + OFF_QH);
    ref_pack_plane(t2, block + OFF_QS + DT3_QS_BYTES, block + OFF_QH + DT3_QH_BYTES);
    const uint16_t h1 = ggml_fp32_to_fp16(d1);
    const uint16_t h2 = ggml_fp32_to_fp16(d2);
    memcpy(block + OFF_D,     &h1, sizeof(h1));
    memcpy(block + OFF_D + 2, &h2, sizeof(h2));
}

// deterministic PRNG so failures are reproducible
static uint32_t rng_state = 0x2b992ddf;
static uint32_t rng_next(void) {
    rng_state ^= rng_state << 13;
    rng_state ^= rng_state >> 17;
    rng_state ^= rng_state << 5;
    return rng_state;
}
static int8_t rng_trit(void) {
    return (int8_t)(rng_next() % 3) - 1;
}

constexpr int NROWS   = 16;
constexpr int NCOLS   = 896; // 7 blocks per row; deliberately not a multiple of 256
constexpr int NBLOCKS = NROWS*NCOLS/QK_DT3;
constexpr int ROW0_NB = NCOLS/QK_DT3;

// trits and scales of row 0, kept for the manual MUL_MAT reference
static int8_t row0_t1[ROW0_NB][QK_DT3];
static int8_t row0_t2[ROW0_NB][QK_DT3];
static float  row0_d1[ROW0_NB];
static float  row0_d2[ROW0_NB];

static void build_dt3_data(std::vector<uint8_t> & data) {
    data.resize((size_t)NBLOCKS*DT3_BLOCK_SIZE);

    // row 0: known trits with non-trivial qh region and mixed-sign scales
    for (int j = 0; j < ROW0_NB; ++j) {
        for (int i = 0; i < QK_DT3; ++i) {
            row0_t1[j][i] = rng_trit();
            row0_t2[j][i] = rng_trit();
        }
        // make sure the qh-packed elements are not all zero
        row0_t1[j][127] = -1;
        row0_t2[j][120] = +1;

        row0_d1[j] = j % 2 == 0 ?  1.5f  : -0.75f;  // exact in fp16
        row0_d2[j] = j % 2 == 0 ? -0.625f:  0.375f; // exact in fp16

        ref_pack_block(row0_t1[j], row0_d1[j], row0_t2[j], row0_d2[j], data.data() + (size_t)j*DT3_BLOCK_SIZE);
    }

    // directed single-trit blocks at the region boundaries, negative d2
    const int special_pos[] = {0, 15, 16, 79, 80, 87, 88, 119, 120, 121, 126, 127};
    const int n_special     = (int)(sizeof(special_pos)/sizeof(special_pos[0]));
    for (int c = 0; c < n_special; ++c) {
        int8_t t1[QK_DT3] = {0};
        int8_t t2[QK_DT3] = {0};
        t1[special_pos[c]] = +1;
        t2[special_pos[c]] = -1;
        ref_pack_block(t1, 1.0f, t2, -0.25f, data.data() + (size_t)(ROW0_NB + c)*DT3_BLOCK_SIZE);
    }

    // the rest: raw random bytes (any byte value is decodable) and random
    // small scales, some negative
    for (int b = ROW0_NB + n_special; b < NBLOCKS; ++b) {
        uint8_t * block = data.data() + (size_t)b*DT3_BLOCK_SIZE;
        for (size_t k = 0; k < OFF_D; ++k) {
            block[k] = (uint8_t)(rng_next() & 0xFF);
        }
        const uint16_t h1 = ggml_fp32_to_fp16(((int)(rng_next() % 2001) - 1000)/500.0f);
        const uint16_t h2 = ggml_fp32_to_fp16(((int)(rng_next() % 2001) - 1000)/500.0f);
        memcpy(block + OFF_D,     &h1, sizeof(h1));
        memcpy(block + OFF_D + 2, &h2, sizeof(h2));
    }
}

// run a single-output graph on the backend and read the result back
static void compute_graph(ggml_backend_t backend, ggml_context * ctx, ggml_tensor * out, float * result) {
    ggml_cgraph * gf = ggml_new_graph(ctx);
    ggml_build_forward_expand(gf, out);

    ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
    const bool ok = ggml_gallocr_alloc_graph(galloc, gf);
    GGML_ASSERT(ok);

    const ggml_status status = ggml_backend_graph_compute(backend, gf);
    GGML_ASSERT(status == GGML_STATUS_SUCCESS);

    ggml_backend_tensor_get(out, result, 0, ggml_nbytes(out));
    ggml_gallocr_free(galloc);
}

// GET_ROWS over all rows on the GPU vs the CPU reference dequantization
static int test_dequant(ggml_backend_t backend, const std::vector<uint8_t> & data, const std::vector<float> & ref) {
    ggml_init_params params = {
        /*.mem_size   =*/ ggml_tensor_overhead()*8 + ggml_graph_overhead(),
        /*.mem_buffer =*/ nullptr,
        /*.no_alloc   =*/ true,
    };
    ggml_context * ctx = ggml_init(params);

    ggml_tensor * a    = ggml_new_tensor_2d(ctx, GGML_TYPE_DT3, NCOLS, NROWS);
    ggml_tensor * rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, NROWS);
    ggml_tensor * out  = ggml_get_rows(ctx, a, rows);

    if (!ggml_backend_supports_op(backend, out)) {
        printf("FAILED: backend does not support GET_ROWS on DT3\n");
        ggml_free(ctx);
        return 1;
    }

    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
    GGML_ASSERT(buf != nullptr);

    std::vector<int32_t> row_idx(NROWS);
    for (int r = 0; r < NROWS; ++r) {
        row_idx[r] = r;
    }
    ggml_backend_tensor_set(a,    data.data(),    0, data.size());
    ggml_backend_tensor_set(rows, row_idx.data(), 0, NROWS*sizeof(int32_t));

    std::vector<float> gpu((size_t)NROWS*NCOLS);
    compute_graph(backend, ctx, out, gpu.data());

    int    num_failed = 0;
    double max_diff   = 0.0;
    for (size_t i = 0; i < gpu.size(); ++i) {
        const double diff = fabs((double)gpu[i] - (double)ref[i]);
        max_diff = diff > max_diff ? diff : max_diff;
        if (gpu[i] != ref[i]) {
            if (num_failed < 8) {
                printf("FAILED: dequant mismatch at block %zu elem %zu: gpu %.9g, cpu %.9g\n",
                       i/QK_DT3, i%QK_DT3, gpu[i], ref[i]);
            }
            num_failed++;
        }
    }
    printf("%s: dequant GPU vs CPU on %d blocks: %d mismatches, max |diff| = %g\n",
           num_failed == 0 ? "OK" : "FAILED", NBLOCKS, num_failed, max_diff);

    ggml_backend_buffer_free(buf);
    ggml_free(ctx);
    return num_failed == 0 ? 0 : 1;
}

struct mat_err {
    double norm_rel; // ||gpu - ref|| / ||ref||
    double max_rel;  // max elementwise |gpu - ref| / max(|ref|, 1) — information only
};

static mat_err compare_mat(const std::vector<float> & gpu, const std::vector<double> & ref) {
    double num  = 0.0;
    double den  = 0.0;
    double mrel = 0.0;
    for (size_t i = 0; i < gpu.size(); ++i) {
        const double diff = (double)gpu[i] - ref[i];
        num += diff*diff;
        den += ref[i]*ref[i];
        const double rel = fabs(diff) / (fabs(ref[i]) > 1.0 ? fabs(ref[i]) : 1.0);
        mrel = rel > mrel ? rel : mrel;
    }
    return { sqrt(num/den), mrel };
}

// MUL_MAT on the GPU vs double precision references from the dequantized
// weights (exact, and rounded to fp16 as the GEMM fallback does).
// strict = tight gates (DT3); controls are gated loosely at 1e-2.
static int test_mul_mat(ggml_backend_t backend, ggml_type type, const std::vector<uint8_t> & data,
                        const std::vector<float> & ref_w, const std::vector<float> & y, bool strict) {
    int num_failed = 0;

    const int ncols_dst[] = {1, 2, 5, 8, 16};

    // the same weights as the fp16 GEMM fallback sees them
    std::vector<float> ref_w16(ref_w.size());
    for (size_t i = 0; i < ref_w.size(); ++i) {
        ref_w16[i] = ggml_fp16_to_fp32(ggml_fp32_to_fp16(ref_w[i]));
    }

    std::vector<std::vector<float>> results;

    for (int c = 0; c < (int)(sizeof(ncols_dst)/sizeof(ncols_dst[0])); ++c) {
        const int n = ncols_dst[c];

        ggml_init_params params = {
            /*.mem_size   =*/ ggml_tensor_overhead()*8 + ggml_graph_overhead(),
            /*.mem_buffer =*/ nullptr,
            /*.no_alloc   =*/ true,
        };
        ggml_context * ctx = ggml_init(params);

        ggml_tensor * a   = ggml_new_tensor_2d(ctx, type, NCOLS, NROWS);
        ggml_tensor * b   = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, NCOLS, n);
        ggml_tensor * out = ggml_mul_mat(ctx, a, b);

        if (!ggml_backend_supports_op(backend, out)) {
            printf("FAILED: backend does not support MUL_MAT on %s\n", ggml_type_name(type));
            ggml_free(ctx);
            return 1;
        }

        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
        GGML_ASSERT(buf != nullptr);

        ggml_backend_tensor_set(a, data.data(), 0, data.size());
        ggml_backend_tensor_set(b, y.data(),    0, (size_t)NCOLS*n*sizeof(float));

        std::vector<float> gpu((size_t)NROWS*n);
        compute_graph(backend, ctx, out, gpu.data());
        results.push_back(gpu);

        // references in double from the exact and the fp16-rounded weights
        std::vector<double> ref((size_t)NROWS*n);
        std::vector<double> ref16((size_t)NROWS*n);
        for (int j = 0; j < n; ++j) {
            for (int r = 0; r < NROWS; ++r) {
                double sum   = 0.0;
                double sum16 = 0.0;
                for (int k = 0; k < NCOLS; ++k) {
                    sum   += (double)ref_w  [(size_t)r*NCOLS + k] * (double)y[(size_t)j*NCOLS + k];
                    sum16 += (double)ref_w16[(size_t)r*NCOLS + k] * (double)y[(size_t)j*NCOLS + k];
                }
                ref  [(size_t)j*NROWS + r] = sum;
                ref16[(size_t)j*NROWS + r] = sum16;
            }
        }

        const mat_err err   = compare_mat(gpu, ref);
        const mat_err err16 = compare_mat(gpu, ref16);

        // n <= 8 is the MMVQ path with exact integer dot products, judged
        // against the exact reference; larger n is the dequantize + GEMM
        // fallback, judged against the better of the exact and fp16-rounded
        // references (which one applies depends on the hardware and on
        // GGML_CUDA_CUBLAS_COMPUTE_TYPE)
        const bool   is_mmvq = n <= 8;
        const double err_gate = is_mmvq ? err.norm_rel : (err.norm_rel < err16.norm_rel ? err.norm_rel : err16.norm_rel);
        const double tol      = strict ? 1e-5 : 1e-2;
        printf("%s: %s mul_mat GPU, ncols_dst = %2d (%s): norm rel err vs exact ref = %g, vs fp16 ref = %g (max elem rel: %g)\n",
               err_gate <= tol ? "OK" : "FAILED", ggml_type_name(type), n, is_mmvq ? "MMVQ" : "GEMM",
               err.norm_rel, err16.norm_rel, err.max_rel);
        if (err_gate > tol) {
            num_failed++;
        }

        ggml_backend_buffer_free(buf);
        ggml_free(ctx);
    }

    // MMVQ vs the dequantization-based path: the first 8 columns of the GEMM
    // run must match the ncols_dst = 8 MMVQ run to fp16 weight rounding
    {
        const std::vector<float> & mmvq = results[3]; // n = 8
        const std::vector<float> & gemm = results[4]; // n = 16
        double num = 0.0;
        double den = 0.0;
        for (int j = 0; j < 8; ++j) {
            for (int r = 0; r < NROWS; ++r) {
                const double diff = (double)mmvq[(size_t)j*NROWS + r] - (double)gemm[(size_t)j*NROWS + r];
                num += diff*diff;
                den += (double)mmvq[(size_t)j*NROWS + r]*(double)mmvq[(size_t)j*NROWS + r];
            }
        }
        const double norm_rel = sqrt(num/den);
        const double tol      = strict ? 2e-3 : 1e-2;
        printf("%s: %s MMVQ vs GEMM path on shared columns: norm rel err = %g\n",
               norm_rel <= tol ? "OK" : "FAILED", ggml_type_name(type), norm_rel);
        if (norm_rel > tol) {
            num_failed++;
        }
    }

    // manual sum over the trits stored by the test for row 0, column 0 —
    // computed from the trits themselves, not from any dequantization, with
    // non-trivial qh trits in every block of the row
    if (type == GGML_TYPE_DT3) {
        double sum = 0.0;
        for (int j = 0; j < ROW0_NB; ++j) {
            for (int i = 0; i < QK_DT3; ++i) {
                sum += (double)y[(size_t)j*QK_DT3 + i] *
                       ((double)row0_d1[j]*row0_t1[j][i] + (double)row0_d2[j]*row0_t2[j][i]);
            }
        }
        const double got = results[0][0]; // ncols_dst = 1, row 0
        const double rel = fabs(got - sum) / (fabs(sum) > 1.0 ? fabs(sum) : 1.0);
        printf("%s: MMVQ vs manual trit sum (row 0, col 0): gpu %.9g, manual %.9g, rel err = %g\n",
               rel <= 1e-5 ? "OK" : "FAILED", got, sum, rel);
        if (rel > 1e-5) {
            num_failed++;
        }
    }

    return num_failed;
}

// quantize random floats to a control type and return raw data + dequantized
// reference weights
static void build_control_data(ggml_type type, std::vector<uint8_t> & data, std::vector<float> & ref_w) {
    std::vector<float> src((size_t)NROWS*NCOLS);
    for (size_t i = 0; i < src.size(); ++i) {
        src[i] = ((int)(rng_next() % 2001) - 1000)/1000.0f;
    }

    data.resize(ggml_row_size(type, NCOLS)*NROWS);
    const size_t written = ggml_quantize_chunk(type, src.data(), data.data(), 0, NROWS, NCOLS, nullptr);
    GGML_ASSERT(written == data.size());

    ref_w.resize(src.size());
    ggml_get_type_traits(type)->to_float(data.data(), ref_w.data(), (int64_t)NROWS*NCOLS);
}

int main(void) {
    ggml_backend_t backend = nullptr;
    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
        if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU) {
            backend = ggml_backend_dev_init(dev, nullptr);
            printf("using GPU backend: %s\n", ggml_backend_dev_name(dev));
            break;
        }
    }
    if (backend == nullptr) {
        printf("no GPU backend available, skipping\n");
        return 0;
    }

    std::vector<uint8_t> data;
    build_dt3_data(data);

    // CPU reference dequantization — the validated path
    std::vector<float> ref((size_t)NROWS*NCOLS);
    const ggml_type_traits * qfns = ggml_get_type_traits(GGML_TYPE_DT3);
    qfns->to_float(data.data(), ref.data(), (int64_t)NROWS*NCOLS);

    // activations: integers with amax 127 in every 32-element chunk of every
    // column, so their q8_1 quantization is exact
    std::vector<float> y((size_t)NCOLS*16);
    for (size_t i = 0; i < y.size(); ++i) {
        y[i] = i % 32 == 0 ? 127.0f : (float)((int)(rng_next() % 255) - 127);
    }

    int num_failed = 0;
    num_failed += test_dequant(backend, data, ref);
    num_failed += test_mul_mat(backend, GGML_TYPE_DT3, data, ref, y, /*strict =*/ true);

    // controls through the identical comparison: Q4_1 shares DT3's regime
    // (dequantized weights not fp16-exact), Q4_0's weights are fp16-exact
    // and show the pure GEMM error floor
    for (ggml_type control : {GGML_TYPE_Q4_1, GGML_TYPE_Q4_0}) {
        std::vector<uint8_t> cdata;
        std::vector<float>   cref;
        build_control_data(control, cdata, cref);
        num_failed += test_mul_mat(backend, control, cdata, cref, y, /*strict =*/ false);
    }

    ggml_backend_free(backend);

    if (num_failed > 0) {
        printf("%d tests FAILED\n", num_failed);
        return 1;
    }
    printf("all tests OK\n");
    return 0;
}