#!/usr/bin/env python3
"""Pack a denoiser stage's weights (engine/rr/enc1_ref.py --extract) for the native stage shaders.

usage: native_pack.py WEIGHTS.npz OUT.bin OUT.h

Every FP8 matrix W (K x N, rows in NVIDIA's k order) becomes 16x16 tiles, tile (kt, nt) at
256 * (kt * N/16 + nt) bytes, row-major inside the tile (stride 16): a B operand of the cooperative matrix loads.
Every f16 matrix is stored as is (row-major); an f16 bias vector (a 1-D entry whose name starts with 'b') as 16
equal rows, so that it loads as an accumulator tile. OUT.h lists the byte offset of every entry as a #define,
for the shader's include.
"""
import sys

import numpy as np


def tiles(w):
    k, n = w.shape
    assert k % 16 == 0 and n % 16 == 0, w.shape
    return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1)


def main():
    src, out_bin, out_h = sys.argv[1:4]
    w = dict(np.load(src))
    blob = bytearray()
    defines = []
    for name in sorted(w):
        a = w[name]
        while len(blob) % 256:
            blob.append(0)
        off = len(blob)
        if a.dtype == np.uint8:
            blob += tiles(a).tobytes()
        elif a.ndim == 1 and name.startswith('b'):
            blob += np.tile(np.asarray(a, np.float16), (16, 1)).tobytes()
        else:
            blob += np.asarray(a, np.float16).tobytes()
        defines.append(f'#define W_{name.upper()} {off}u   // {a.dtype} {a.shape}')
    open(out_bin, 'wb').write(bytes(blob))
    open(out_h, 'w').write('// generated by engine/rr/native_pack.py\n' + '\n'.join(defines) + '\n')
    print(f'{len(blob)} bytes, {len(w)} entries')


if __name__ == '__main__':
    main()