s4fifo_embed_model.py

#!/usr/bin/env python3
"""Convert a .s4m model file into the C source that S4FIFO_model_real.c
compiles in as its built-in default model.

The .s4m format itself (see S4FIFO_model_real.c's header comment, written by
s4fifo_export_model.py) is already just four flat arrays plus a header, so
this is a direct transcription - no re-encoding, no precision loss: every
double is printed with %.17g, which round-trips exactly through strtod.

Only models small enough to be worth vendoring should go through this. The
model currently shipped (1460 trees / 233527 nodes / 5.6MB binary) expands
to ~14MB of C. Anything substantially larger - the 140400-tree ensemble is
>1GB once expanded - must stay a runtime `model-path=` file.

Usage: s4fifo_embed_model.py <model.s4m> <out.h> [--name Embedded]
"""
import argparse
import os
import struct
import sys

# entries per emitted line - keeps lines readable-ish without ballooning the
# file with one entry per line
NODES_PER_LINE = 4
SCALARS_PER_LINE = 8


def fmt(v):
    """%.17g round-trips a double exactly through strtod."""
    return "%.17g" % v


def read_s4m(path):
    with open(path, "rb") as f:
        if f.read(4) != b"S4M1":
            raise SystemExit(f"{path} is not an S4M1 model file")
        n_models, n_classes, n_features, n_trees = struct.unpack("<IIII", f.read(16))
        n_estimators = struct.unpack("<%dI" % n_models, f.read(4 * n_models))
        trees = [struct.unpack("<III", f.read(12)) for _ in range(n_trees)]
        node_off, n_internal, leaf_off = trees[-1]
        n_nodes = node_off + n_internal
        n_leaves = leaf_off + n_internal + 1
        nodes = [struct.unpack("<hhhhd", f.read(16)) for _ in range(n_nodes)]
        leaves = struct.unpack("<%dd" % n_leaves, f.read(8 * n_leaves))
        cost = struct.unpack("<%dd" % (n_classes * n_classes), f.read(8 * n_classes * n_classes))
        if f.read():
            raise SystemExit(f"{path} has trailing bytes - format mismatch?")
    return dict(n_models=n_models, n_classes=n_classes, n_features=n_features,
                n_trees=n_trees, n_estimators=n_estimators, trees=trees,
                nodes=nodes, leaves=leaves, cost=cost)


def emit_rows(out, rows, per_line):
    for i in range(0, len(rows), per_line):
        out.write("    " + " ".join(rows[i:i + per_line]) + "\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("model_s4m")
    ap.add_argument("out_h")
    ap.add_argument("--name", default="Embedded",
                    help="infix for the generated symbols (default: Embedded)")
    args = ap.parse_args()

    m = read_s4m(args.model_s4m)
    n = args.name
    # Deliberately not recording the source .s4m's filename: that file isn't
    # in the repo, so its name says nothing useful to a reader here, and it
    # would pin the generated header to one particular training run.
    upper = n.upper()

    with open(args.out_h, "w") as out:
        # clang-format off must come first: the repo's CI format check runs
        # over this file like any other header, and nothing here is meant to
        # be reformatted.
        out.write(f"""// clang-format off
//
// GENERATED by scripts/s4fifo_embed_model.py - do not hand-edit and do not
// reformat; regenerate instead.
//
// The built-in default model for S4FIFO's learned control plane
// (`auto-tune=1` with no `model-path=`). This is a direct transcription of
// the .s4m binary format documented in S4FIFO_model_real.c's header
// comment: same four arrays, same layout, doubles printed at full %.17g
// precision so they round-trip exactly. Included from S4FIFO_model_real.c
// only - it needs s4fifo_real_node_t / s4fifo_real_tree_t, and these
// arrays must live in exactly one translation unit.
//
// Model: {m['n_models']} booster, {m['n_classes']} classes, {m['n_features']} features,
// {m['n_trees']} trees ({m['n_estimators'][0]} boosting rounds x {m['n_classes']} classes),
// {len(m['nodes'])} internal nodes, {len(m['leaves'])} leaves.
//
#pragma once

#define S4FIFO_{upper}_N_MODELS {m['n_models']}
#define S4FIFO_{upper}_N_CLASSES {m['n_classes']}
#define S4FIFO_{upper}_N_FEATURES {m['n_features']}
#define S4FIFO_{upper}_N_TREES {m['n_trees']}

static const uint32_t kS4FIFO{n}Estimators[S4FIFO_{upper}_N_MODELS] = {{
""")
        emit_rows(out, ["%d," % e for e in m["n_estimators"]], SCALARS_PER_LINE)

        out.write(f"}};\n\nstatic const s4fifo_real_tree_t "
                  f"kS4FIFO{n}Trees[S4FIFO_{upper}_N_TREES] = {{\n")
        emit_rows(out, ["{%d,%d,%d}," % t for t in m["trees"]], NODES_PER_LINE)

        out.write(f"}};\n\nstatic const s4fifo_real_node_t "
                  f"kS4FIFO{n}Nodes[{len(m['nodes'])}] = {{\n")
        emit_rows(out, ["{%d,%d,%d,0,%s}," % (nd[0], nd[1], nd[2], fmt(nd[4]))
                        for nd in m["nodes"]], NODES_PER_LINE)

        out.write(f"}};\n\nstatic const double kS4FIFO{n}Leaves[{len(m['leaves'])}] = {{\n")
        emit_rows(out, ["%s," % fmt(v) for v in m["leaves"]], SCALARS_PER_LINE)

        out.write(f"}};\n\nstatic const double kS4FIFO{n}Cost"
                  f"[S4FIFO_{upper}_N_CLASSES * S4FIFO_{upper}_N_CLASSES] = {{\n")
        emit_rows(out, ["%s," % fmt(v) for v in m["cost"]], SCALARS_PER_LINE)
        out.write("};\n")

    print("wrote %s (%.1f MB): %d trees, %d nodes, %d leaves" %
          (args.out_h, os.path.getsize(args.out_h) / 1e6, m["n_trees"],
           len(m["nodes"]), len(m["leaves"])), file=sys.stderr)


if __name__ == "__main__":
    main()