| """Kernel build and ctypes bindings for the Mach Boom runtime.""" |
| import ctypes, os, subprocess, sys |
| import numpy as np |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
|
|
|
|
| def tiles_from_qidxs(Q, m, n, K): |
| Q = np.asarray(Q).view(np.uint16).reshape(m, -1) |
| nb, K = n // 16, int(K) |
| assert Q.shape[1] == K * nb, (Q.shape, K, nb) |
| t = Q.reshape(m // 16, 16, nb, K).transpose(0, 2, 1, 3) |
| return t.astype(">u2").tobytes() |
|
|
|
|
| def signbits(s): |
| return np.packbits(np.asarray(s).reshape(-1) > 0, bitorder="little").tobytes() |
|
|
|
|
| def tlut10(tl): |
| tl = np.asarray(tl, np.float32).reshape(512, 2) |
| t = np.concatenate([tl, tl]).copy() |
| t[512:, 0] *= -1 |
| return np.ascontiguousarray(t) |
|
|
|
|
| def build(out=None, isa=None): |
| src = [os.path.join(HERE, f) for f in ("tmk.c", "tmk_act.c", "tmk_model.c", "tmk_amx.c")] |
| if out is None: |
| out = os.path.join(HERE, f"tmk_{isa or 'native'}.so") |
| if not os.access(HERE, os.W_OK): |
| import hashlib |
| h = hashlib.sha256(b"".join(open(f, "rb").read() for f in src)).hexdigest()[:12] |
| d = os.path.join(os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")), "mach_boom") |
| os.makedirs(d, exist_ok=True) |
| out = os.path.join(d, f"tmk_{isa or 'native'}_{h}.so") |
| if os.path.exists(out) and os.path.getmtime(out) >= max(os.path.getmtime(f) for f in src): |
| return out |
| cc = os.environ.get("CC", "cc") |
| flags = {"native": "-march=native", "avx512": "-march=skylake-avx512", "avx2": "-march=haswell", "scalar": "", |
| None: "-march=native"}[isa] |
| tmp = out + f".{os.getpid()}.tmp" |
| objs = [tmp + ".act.o", tmp + ".model.o", tmp + ".amx.o"] |
|
|
| def run(flags, libs, amx): |
| for f, o in zip(src[1:3], objs): |
| subprocess.check_call([cc, "-O3", "-ffast-math", *flags.split(), "-fPIC", "-c", f, "-o", o]) |
| subprocess.check_call([cc, "-O3", *flags.split(), *([] if amx else ["-DTMK_NO_AMX"]), "-fPIC", "-c", src[3], |
| "-o", objs[2]], stderr=subprocess.DEVNULL if amx else None) |
| subprocess.check_call([cc, "-O3", *flags.split(), "-shared", "-fPIC", src[0], *objs, "-o", tmp, *libs.split()]) |
| mv, nm = "-lmvec -lm -pthread", "-lm -pthread" |
| tries = [(flags, mv, 1), (flags, mv, 0), (flags, nm, 1), (flags, nm, 0), ("", nm, 0)] |
| try: |
| for i, (fl, libs, amx) in enumerate(tries): |
| try: |
| run(fl, libs, amx) |
| break |
| except subprocess.CalledProcessError: |
| if i == len(tries) - 1: |
| raise |
| os.replace(tmp, out) |
| finally: |
| for o in objs: |
| if os.path.exists(o): |
| os.remove(o) |
| return out |
|
|
|
|
| P_, I_, F_ = ctypes.c_void_p, ctypes.c_int, ctypes.c_float |
| MAXL = 8 |
|
|
|
|
| class Chunk(ctypes.Structure): |
| _fields_ = [("codes", P_), ("su", P_), ("sv", P_), ("ws", F_), ("m", I_), ("n", I_)] |
|
|
|
|
| class Lin(ctypes.Structure): |
| _fields_ = [("ch", ctypes.POINTER(Chunk)), ("nch", I_), ("step", I_), ("tlut", P_), ("bias", P_), ("dense", P_), |
| ("m", I_), ("n", I_)] |
|
|
|
|
| class LMS(ctypes.Structure): |
| _fields_ = [("D", I_), ("NL", I_), ("NH", I_), ("HD", I_), ("FF", I_), ("Pmax", I_), ("eps", F_), ("inv_freq", P_), |
| ("ln", (P_ * 4) * MAXL), ("lin", (Lin * 6) * MAXL), ("lnf_w", P_), ("lnf_b", P_), ("in_proj", Lin), |
| ("in_add", P_)] |
|
|
|
|
| class FlowS(ctypes.Structure): |
| _fields_ = [("D", I_), ("NL", I_), ("LD", I_), ("eps", F_), ("inp", Lin), ("cond", Lin), ("ada", Lin * MAXL), |
| ("fc1", Lin * MAXL), ("fc2", Lin * MAXL), ("fada", Lin), ("fproj", Lin), ("norm", (P_ * 2) * MAXL), |
| ("time", P_)] |
|
|
|
|
| class K: |
| def __init__(self, so=None): |
| lib = ctypes.CDLL(so or build()) |
| P, I, F, Lg = ctypes.c_void_p, ctypes.c_int, ctypes.c_float, ctypes.c_long |
| lib.tmk_isa.restype = ctypes.c_char_p |
| lib.tmk_qgemv.argtypes = [P, I, I, I, P, P, P] |
| lib.tmk_qgemv_rows.argtypes = [P, I, I, P, P, P, I, I] |
| lib.tmk_tlut10.argtypes = [P, P] |
| lib.tmk_threads.argtypes = [I]; lib.tmk_threads.restype = I |
| lib.tmk_qdense.argtypes = [P, I, I, I, P, F, P] |
| lib.tmk_fwht.argtypes = [P, I, I] |
| lib.tmk_signs.argtypes = [P, I, I, P] |
| lib.tmk_rot.argtypes = [P, I, I, P, P, F] |
| lib.tmk_linear_vec.argtypes = [P, I, I, I, P, P, P, F, P, P, P] |
| lib.tmk_lm_work.argtypes = [ctypes.POINTER(LMS)]; lib.tmk_lm_work.restype = ctypes.c_size_t |
| lib.tmk_lm_step.argtypes = [ctypes.POINTER(LMS), P, P, I, P, P, P, P] |
| lib.tmk_flow_step.argtypes = [ctypes.POINTER(FlowS), P, P, P, P] |
| lib.tmk_softmax_band.argtypes = [P, Lg, I, I, I, I] |
| lib.tmk_bf16_pack_k.argtypes = [P, Lg, I, I, P, I, I] |
| lib.tmk_bf16_pack_j.argtypes = [P, Lg, I, I, P, I, I] |
| lib.tmk_bf16_x.argtypes = [P, Lg, I, I, I, I, P, I] |
| lib.tmk_colrot.argtypes = [P, I, I, P, F] |
| lib.tmk_dtf_work.argtypes = [I, I, I, I]; lib.tmk_dtf_work.restype = Lg |
| lib.tmk_dtf_layer.argtypes = [P, I, I, I, I, F, P, P, P, P, P, P, P, P, P, I, I, P, Lg, P] |
| lib.tmk_wpack.argtypes = [P, I, I, I, P, F, P, P, P, P, P, I, I] |
| lib.tmk_amx_gemm.argtypes = [P, I, I, P, I, P, Lg, I] |
| for f in ("tmk_gelu", "tmk_silu", "tmk_elu"): |
| getattr(lib, f).argtypes = [P, Lg] |
| self.lib, self.isa = lib, lib.tmk_isa().decode() |
|
|
|
|
| def ptr(a): |
| return a.ctypes.data if isinstance(a, np.ndarray) else ctypes.cast(a, ctypes.c_void_p).value |
|
|
