rtl_gen.py

"""Generate all four RTL variants from per_dim_thresholds.json.

    python rtl_gen.py

Kept separate from calibrate.py so the RTL can be regenerated from the committed
thresholds without the feature cache the calibration needs.
"""
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))  # repo root, for `common`
from common import read_artifact  # noqa: E402

HERE = Path(__file__).resolve().parent
N_DIMS = 40
N_POS = 20

HEADER = '''// {title}
// Generated by rtl_gen.py; do not edit by hand.
//
// Inputs are the 40 Stage 0 classifier dims as signed INT8, post-LayerNorm and
// post-max-pool. Output is one bit. Combinational, no multipliers, no memory.
'''


def _ports(name: str, width: int, per_line: int = 10) -> str:
    """Declaration lines for f00..f39 or t00..t39."""
    rows = []
    for start in range(0, N_DIMS, per_line):
        names = ', '.join(f'{name}{i:02d}' for i in range(start, start + per_line))
        rows.append(f'    input  signed [{width - 1}:0] {names},')
    return '\n'.join(rows)


def _sign_extended_sum(lo: int, hi: int, per_line: int = 4) -> str:
    """Sign-extended 16-bit addition of f{lo}..f{hi-1}."""
    terms = [f'{{{{8{{f{i:02d}[7]}}}}, f{i:02d}}}' for i in range(lo, hi)]
    rows = [' + '.join(terms[i:i + per_line]) for i in range(0, len(terms), per_line)]
    return ' +\n        '.join(rows)


def _popcount(lo: int, hi: int, threshold, per_line: int = 3) -> str:
    """Sum of the per-dim comparisons, written inline.

    The comparisons are summed directly rather than collected into a vector and
    indexed. Indexing a vector is the more readable form, but it relies on
    bit-select lowering that the synthesis backend gets wrong at index 0, and a
    decision circuit is not the place to depend on that.
    """
    terms = [f'(f{i:02d} > {threshold(i)})' for i in range(lo, hi)]
    rows = [' + '.join(terms[i:i + per_line]) for i in range(0, len(terms), per_line)]
    return ' +\n        '.join(rows)


def _threshold_bank(values, per_line: int = 5) -> str:
    """localparam bank holding the 40 baked INT8 thresholds."""
    rows = []
    for start in range(0, N_DIMS, per_line):
        rows.append(', '.join(f'T{i:02d} = {values[i]:>4}'
                              for i in range(start, min(start + per_line, N_DIMS))))
    return (',\n' + ' ' * 29).join(rows)


def emit_sum() -> str:
    return HEADER.format(title='Additive 1-parameter person classifier, runtime threshold.') + f'''
module person_classifier_1p (
{_ports('f', 8)}
    input  signed [15:0] threshold,
    output person_present
);
    // Score = sum(f00..f19) - sum(f20..f39); worst case 20 * 127 = 2540 fits in 16 bits.
    wire signed [15:0] pos_sum =
        {_sign_extended_sum(0, 20)};

    wire signed [15:0] neg_sum =
        {_sign_extended_sum(20, 40)};

    wire signed [15:0] score = pos_sum - neg_sum;
    assign person_present = score > threshold;
endmodule
'''


def emit_sum_folded(final_int8: int, final_float: float, quant_scale: int) -> str:
    return HEADER.format(title='Additive 1-parameter person classifier, threshold baked in.') + f'''
module person_classifier_sum_folded (
{_ports('f', 8)}
    output person_present
);
    // Stage 0 threshold {final_float:.4f} at the x{quant_scale} scale used for the per-dim ones.
    localparam signed [15:0] FINAL_T = 16'sd{final_int8};

    wire signed [15:0] pos_sum =
        {_sign_extended_sum(0, 20)};

    wire signed [15:0] neg_sum =
        {_sign_extended_sum(20, 40)};

    wire signed [15:0] score = pos_sum - neg_sum;
    assign person_present = score > FINAL_T;
endmodule
'''


def emit_popcount() -> str:
    title = 'Popcount-reformulated 1-parameter person classifier, runtime thresholds.'
    return HEADER.format(title=title) + f'''
module person_classifier_popcount (
{_ports('f', 8)}
{_ports('t', 8)}
    input  signed [5:0] final_threshold,
    output person_present
);
    wire [5:0] count_pos =
        {_popcount(0, 20, lambda i: f't{i:02d}')};

    wire [5:0] count_neg =
        {_popcount(20, 40, lambda i: f't{i:02d}')};

    wire signed [6:0] diff = {{1'b0, count_pos}} - {{1'b0, count_neg}};
    assign person_present = diff > final_threshold;
endmodule
'''


def emit_popcount_folded(thresholds, final_threshold: int, quant_scale: int) -> str:
    title = 'Popcount-reformulated 1-parameter person classifier, thresholds baked in.'
    return HEADER.format(title=title) + f'''//
// Per-dim thresholds are the calibrated float values scaled by {quant_scale} and rounded.

module person_classifier_popcount_folded (
{_ports('f', 8)}
    output person_present
);
    localparam signed [7:0] {_threshold_bank(thresholds)};
    localparam signed [5:0] FINAL_T = {final_threshold};

    wire [5:0] count_pos =
        {_popcount(0, 20, lambda i: f'T{i:02d}')};

    wire [5:0] count_neg =
        {_popcount(20, 40, lambda i: f'T{i:02d}')};

    wire signed [6:0] diff = {{1'b0, count_pos}} - {{1'b0, count_neg}};
    assign person_present = diff > FINAL_T;
endmodule
'''


def generate(out_dir: Path = None, thresholds_json: Path = None,
             classifier_json: Path = None):
    """Write all four modules; returns the paths written."""
    out_dir = out_dir or HERE / 'rtl'
    cal = read_artifact(thresholds_json or HERE / 'per_dim_thresholds.json')
    classifier = json.loads(
        (classifier_json or HERE / 'classifier.json').read_text())

    scale = cal['quant_scale']
    per_dim = [p['threshold_int8'] for p in cal['per_dim_thresholds']]
    final_pop = cal['popcount']['final_threshold']
    stage_0_thr = float(classifier['threshold'])
    stage_0_int8 = int(round(stage_0_thr * scale))

    out_dir.mkdir(parents=True, exist_ok=True)
    written = {
        'sum.v': emit_sum(),
        'sum_folded.v': emit_sum_folded(stage_0_int8, stage_0_thr, scale),
        'popcount.v': emit_popcount(),
        'popcount_folded.v': emit_popcount_folded(per_dim, final_pop, scale),
    }
    for name, text in written.items():
        (out_dir / name).write_text(text, encoding='utf-8')
    return [out_dir / n for n in written]


if __name__ == '__main__':
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--out', type=Path, default=None)
    args = ap.parse_args()
    for path in generate(args.out):
        print(f'wrote {path}')