Files
2026-08-16 18:33:03 +07:00

101 lines
3.7 KiB
C++

// test_ss_dec — validate the C++ SS decoder against the PyTorch reference
// produced by ref_ss_dec.py.
//
// usage: test_ss_dec <ss_dec_f32.gguf> <ss_dec_ref.bin> [rel_tol]
//
// Reads the input latent z_s and the reference occupancy logits, runs the C++
// decoder on the identical latent, and reports max abs / relative L2 error. Also
// checks sign agreement at 0 — the occupancy scaffold is logit>0, so the sign
// map is what ultimately matters. Default tolerance 3e-2.
#include "trellis2.h"
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <string>
#include <vector>
static bool rd(std::ifstream & f, void * p, size_t n) {
return (bool) f.read(reinterpret_cast<char *>(p), (std::streamsize) n);
}
int main(int argc, char ** argv) {
if (argc < 3) {
std::fprintf(stderr, "usage: %s <ss_dec_f32.gguf> <ss_dec_ref.bin> [rel_tol]\n", argv[0]);
return 2;
}
const std::string gguf_path = argv[1], ref_path = argv[2];
{
std::ifstream _a(gguf_path), _b(ref_path);
if (!_a.good() || !_b.good()) {
std::fprintf(stderr, "missing input file(s), skipping\n");
return 77;
}
}
const double rel_tol = (argc > 3) ? std::atof(argv[3]) : 3e-2;
std::ifstream f(ref_path, std::ios::binary);
char magic[8];
if (!f || !rd(f, magic, 8) || std::memcmp(magic, "SSDEC001", 8) != 0) {
std::fprintf(stderr, "error: bad/missing ref file %s\n", ref_path.c_str());
return 1;
}
int32_t hdr[4];
rd(f, hdr, sizeof(hdr));
const int Cin = hdr[0], Rin = hdr[1], Oc = hdr[2], Rout = hdr[3];
const size_t n_in = (size_t) Cin * Rin * Rin * Rin;
const size_t n_out = (size_t) Oc * Rout * Rout * Rout;
std::vector<float> latent(n_in), ref(n_out);
rd(f, latent.data(), latent.size() * sizeof(float));
if (!rd(f, ref.data(), ref.size() * sizeof(float))) {
std::fprintf(stderr, "error: ref truncated\n");
return 1;
}
std::printf("ref: latent[%d,%d^3] -> logits[%d,%d^3]\n", Cin, Rin, Oc, Rout);
std::string err;
trellis2_ss_dec_model * m = trellis2_ss_dec_load(gguf_path, true, &err);
if (!m) { std::fprintf(stderr, "load error: %s\n", err.c_str()); return 1; }
std::printf("backend: %s\n", trellis2_ss_dec_backend_name(m));
const trellis2_ss_dec_hparams hp = trellis2_ss_dec_hparams_of(m);
if (hp.latent_channels != Cin || hp.res_in() != Rin || hp.out_channels != Oc || hp.res_out() != Rout) {
std::fprintf(stderr, "error: model/ref shape mismatch\n");
trellis2_ss_dec_free(m);
return 1;
}
std::vector<float> out(n_out, 0.0f);
if (!trellis2_ss_dec_decode(m, latent.data(), out.data(), &err)) {
std::fprintf(stderr, "decode error: %s\n", err.c_str());
trellis2_ss_dec_free(m);
return 1;
}
trellis2_ss_dec_free(m);
double max_abs = 0.0, sse = 0.0, ref_sq = 0.0;
size_t sign_agree = 0;
for (size_t i = 0; i < n_out; ++i) {
const double d = (double) out[i] - (double) ref[i];
max_abs = std::fmax(max_abs, std::fabs(d));
sse += d * d;
ref_sq += (double) ref[i] * (double) ref[i];
if ((out[i] > 0.0f) == (ref[i] > 0.0f)) ++sign_agree;
}
const double rel_l2 = std::sqrt(sse) / (std::sqrt(ref_sq) + 1e-30);
const double sign_pct = 100.0 * (double) sign_agree / (double) n_out;
std::printf("max abs err : %.3e\n", max_abs);
std::printf("rel L2 err : %.3e (tol %.1e)\n", rel_l2, rel_tol);
std::printf("sign agree : %.3f%% (occupancy is logit>0)\n", sign_pct);
if (rel_l2 > rel_tol) { std::printf("RESULT: FAIL\n"); return 1; }
std::printf("RESULT: PASS\n");
return 0;
}