wai-quantum 0.6.0

"""The device side of wai-quantum's tensor-engine member: one kernel on a
tensor processor's matrix engine, driven over a pipe.

The kernel is one bf16 matrix multiply per tile of TILE pairs, accumulated in
fp32: the gate's limb weights (16 x 112, the stationary matrix) times the
pairs' limbs (16 x TILE, the moving matrix). Each output is one limb
coefficient of one of the gate's sixteen products per pair; the member
assembles the products and does everything else. Every input is an integer
bf16 holds exactly and every output an integer below 2^24, so fp32 holds each
partial sum exactly in any order (see quantum_fabric::tensor_engine).

    python -I tensor_engine_helper.py [--target simulator|<device platform>]

The target is the kernel toolchain's simulator (the default), or a device
platform the toolchain names (trn1, trn2, ...): then the kernel is compiled
for that device and runs on it.

Protocol (little-endian). On start, one line: "wq-tensor-engine 1 <target>
<package> <version> ...". Then requests until end of input: u32 0x45545157,
u32 n (a multiple of TILE), the weights as 16*112 u16 (bf16 bits, row-major),
the pairs' limbs as 16*n u16 (row-major); each answered by u32 0 and the
coefficients as 112*n f32 (row-major), or by u32 1, u32 length and a message.
"""

import struct
import sys

TILE = 512
K = 16
M = 112
MAGIC = 0x45545157


def main():
    import os

    target = "simulator"
    if len(sys.argv) == 3 and sys.argv[1] == "--target":
        target = sys.argv[2]
    elif len(sys.argv) != 1:
        print(__doc__, file=sys.stderr)
        return 2
    if not target.replace("-", "").isalnum():
        print(f"a target is letters, digits and hyphens: {target!r}", file=sys.stderr)
        return 2
    if target != "simulator":
        os.environ.setdefault("NEURON_PLATFORM_TARGET_OVERRIDE", target)

    import numpy as np
    import ml_dtypes
    import nki
    import nki.isa as nisa
    import nki.language as nl

    @nki.jit
    def coefficients(w, a):
        tiles = a.shape[1] // TILE
        out = nl.ndarray((M, a.shape[1]), dtype=nl.float32, buffer=nl.shared_hbm)
        ws = nl.load(w)
        for t in nl.affine_range(tiles):
            xs = nl.load(a[:, t * TILE:(t + 1) * TILE])
            acc = nl.ndarray((M, TILE), dtype=nl.float32, buffer=nl.psum)
            nisa.nc_matmul(acc, ws, xs)
            nl.store(out[:, t * TILE:(t + 1) * TILE], nl.copy(acc))
        return out

    inp, out = sys.stdin.buffer, sys.stdout.buffer
    version = getattr(nki, "__version__", "unknown")
    run = nki.simulate(coefficients) if target == "simulator" else coefficients
    out.write(f"wq-tensor-engine 1 {target} nki {version} numpy {np.__version__} ml_dtypes {ml_dtypes.__version__}\n".encode())
    out.flush()

    def read(n):
        b = inp.read(n)
        if len(b) != n:
            raise EOFError
        return b

    while True:
        try:
            magic, n = struct.unpack("<II", read(8))
        except EOFError:
            return
        try:
            if magic != MAGIC or n == 0 or n % TILE:
                raise ValueError(f"a request of {n} pairs (magic {magic:#x})")
            w = np.frombuffer(read(K * M * 2), "<u2").view(ml_dtypes.bfloat16).reshape(K, M)
            a = np.frombuffer(read(K * n * 2), "<u2").view(ml_dtypes.bfloat16).reshape(K, n)
            c = np.asarray(run(w, a), dtype=np.float32)
            out.write(struct.pack("<I", 0))
            out.write(np.ascontiguousarray(c, dtype="<f4").tobytes())
        except EOFError:
            return
        except Exception as e:  # report, keep serving
            msg = f"{type(e).__name__}: {e}".encode()
            out.write(struct.pack("<II", 1, len(msg)) + msg)
        out.flush()


if __name__ == "__main__":
    sys.exit(main())