File size: 5,542 Bytes
01860b0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """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
|