103 lines
4.1 KiB
C++
103 lines
4.1 KiB
C++
// test_ss_flow_forward — validate the C++ SS-flow DiT forward pass against the
|
|
// PyTorch float32 reference produced by ref_ss_flow.py.
|
|
//
|
|
// usage: test_ss_flow_forward <ss_flow_dit_f32.gguf> <ss_flow_ref.bin> [rel_tol]
|
|
//
|
|
// Reads x/t/cond and the reference output from ss_flow_ref.bin, runs the C++
|
|
// forward with the f32 weights, and reports max abs / relative error. Exits
|
|
// nonzero if the relative L2 error exceeds the tolerance (default 2e-3).
|
|
|
|
#include "trellis2.h"
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
static bool read_exact(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 <f32.gguf> <ss_flow_ref.bin> [rel_tol]\n", argv[0]);
|
|
return 2;
|
|
}
|
|
const std::string gguf_path = argv[1];
|
|
const std::string 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]) : 2e-3;
|
|
|
|
// ── read the reference bundle ────────────────────────────────────────────
|
|
std::ifstream f(ref_path, std::ios::binary);
|
|
char magic[8];
|
|
if (!f || !read_exact(f, magic, 8) || std::memcmp(magic, "SSFREF01", 8) != 0) {
|
|
std::fprintf(stderr, "error: bad/missing ref file %s\n", ref_path.c_str());
|
|
return 1;
|
|
}
|
|
int32_t dims[5];
|
|
float t = 0.0f;
|
|
read_exact(f, dims, sizeof(dims));
|
|
read_exact(f, &t, sizeof(t));
|
|
const int R = dims[0], Cin = dims[1], Cout = dims[2], Lkv = dims[3], Cctx = dims[4];
|
|
const size_t N = (size_t) R * R * R;
|
|
|
|
std::vector<float> x( (size_t) Cin * N);
|
|
std::vector<float> cond((size_t) Lkv * Cctx);
|
|
std::vector<float> ref( (size_t) Cout * N);
|
|
read_exact(f, x.data(), x.size() * sizeof(float));
|
|
read_exact(f, cond.data(), cond.size() * sizeof(float));
|
|
if (!read_exact(f, ref.data(), ref.size() * sizeof(float))) {
|
|
std::fprintf(stderr, "error: ref file truncated\n");
|
|
return 1;
|
|
}
|
|
std::printf("ref: R=%d Cin=%d Cout=%d Lkv=%d Cctx=%d t=%.3f\n", R, Cin, Cout, Lkv, Cctx, t);
|
|
|
|
// ── load weights + run the C++ forward ───────────────────────────────────
|
|
std::string err;
|
|
trellis2_ss_flow_model * m = trellis2_ss_flow_load(gguf_path, /*load_tensors*/ true, &err);
|
|
if (!m) { std::fprintf(stderr, "load error: %s\n", err.c_str()); return 1; }
|
|
std::printf("backend: %s\n", trellis2_ss_flow_backend_name(m));
|
|
|
|
std::vector<float> out((size_t) Cout * N, 0.0f);
|
|
if (!trellis2_ss_flow_forward(m, x.data(), t, cond.data(), Lkv, Cctx, out.data(), &err)) {
|
|
std::fprintf(stderr, "forward error: %s\n", err.c_str());
|
|
trellis2_ss_flow_free(m);
|
|
return 1;
|
|
}
|
|
trellis2_ss_flow_free(m);
|
|
|
|
// ── compare ──────────────────────────────────────────────────────────────
|
|
double max_abs = 0.0, sse = 0.0, ref_sq = 0.0;
|
|
int max_i = 0;
|
|
for (size_t i = 0; i < out.size(); ++i) {
|
|
const double d = (double) out[i] - (double) ref[i];
|
|
if (std::fabs(d) > max_abs) { max_abs = std::fabs(d); max_i = (int) i; }
|
|
sse += d * d;
|
|
ref_sq += (double) ref[i] * (double) ref[i];
|
|
}
|
|
const double rel_l2 = std::sqrt(sse) / (std::sqrt(ref_sq) + 1e-30);
|
|
|
|
std::printf("out: min/max checked over %zu elems\n", out.size());
|
|
std::printf("max abs err : %.3e (at %d: cpp=%.6f ref=%.6f)\n",
|
|
max_abs, max_i, out[max_i], ref[max_i]);
|
|
std::printf("rel L2 err : %.3e (tol %.1e)\n", rel_l2, rel_tol);
|
|
|
|
if (rel_l2 > rel_tol) {
|
|
std::printf("RESULT: FAIL\n");
|
|
return 1;
|
|
}
|
|
std::printf("RESULT: PASS\n");
|
|
return 0;
|
|
}
|