"""LM head speed on the NPU, one token: INT8 weight as a constant (baked in) vs as a runtime input (shared).
usage: python prof_head.py V H"""
import sys, time
import numpy as np
import openvino as ov
import openvino.opset13 as ops
from onw import ir


def main():
    V, H = int(sys.argv[1]), int(sys.argv[2])
    core = ov.Core()
    rng = np.random.default_rng(0)
    wq = rng.integers(-127, 128, size=(V, H), dtype=np.int8)
    ws = (rng.uniform(0.5, 1.5, V) / 1000).astype(np.float16)
    x = rng.standard_normal((1, 1, H)).astype(np.float16)
    ctx = core.get_default_context("NPU")
    for kind in ("constant", "input"):
        h = ops.parameter([1, 1, H], ov.Type.f32, name="h")
        if kind == "constant":
            w = ops.convert(ops.multiply(ops.convert(ops.constant(wq), ov.Type.f16), ops.constant(ws[:, None])), ov.Type.f32)
            m = ir.finish({"y": ops.matmul(h, w, False, True)}, [h], "c")
        else:
            lg, p = ir.shared_head(h, V, H)
            m = ir.finish({"y": lg}, [h] + p, "i")
        t0 = time.time()
        req = core.compile_model(m, "NPU").create_infer_request()
        tc = time.time() - t0
        if kind == "input":
            tq = ctx.create_host_tensor(ov.Type.i8, ov.Shape([V, H])); np.copyto(tq.data, wq)
            tsc = ctx.create_host_tensor(ov.Type.f16, ov.Shape([V])); np.copyto(tsc.data, ws)
            req.set_tensor("shared.head_q", tq); req.set_tensor("shared.head_s", tsc)
        for _ in range(3):
            req.infer({"h": x})
        n = 20
        t0 = time.time()
        for _ in range(n):
            req.infer({"h": x})
        print(f"{kind:9s} {(time.time()-t0)/n*1000:6.1f} ms/token  (compile {tc:.0f}s)", flush=True)


if __name__ == "__main__":
    main()