d4r/engine/m_model.py

#!/usr/bin/env python3
"""Preset M (rrlite): plain weights of the Swin layers and a plain numpy model of them.

NGX keeps the weights as FP8 (e4m3) bytes in a layout arranged for its tensor-core fragments, and half of the
activations in a permuted ("pair") channel order. plain_layer() undoes both: every tensor becomes a row-major
matrix W[K][N] (y = x W) over natural channel order, f16 exact (every e4m3 value is an f16 value).
forward() is the layer's mathematics on those matrices. Every activation that feeds a matrix product as its left
operand is rounded to FP8 first, as in NVIDIA's arithmetic (its tensor-core operands are FP8) and in d4r's native
FP8 kernels; sums are accumulated without the intermediate f16 rounding of NVIDIA's k32 steps.

usage: m_model.py CAPTURE_DIR [frame]     compares forward() with the captured outputs of each layer launch
"""
import glob, os, struct, sys
import numpy as np
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'kernels', 'tools'))
import model_enc3 as me
import swin_model as sm

F16 = np.float16
# kernel -> (C, heads, merged channels, low-resolution input channels)
LAYERS = {'enc1': (64, 2, 96, 0), 'enc2': (96, 4, 128, 0), 'enc3_tube': (128, 4, 0, 0), 'dec2': (96, 4, 0, 128), 'dec1': (64, 2, 0, 96)}
ORDER = ['enc1', 'enc2'] + ['enc3_tube'] * 6 + ['dec2', 'dec1']


def gp(n):
    return sm.gp(np.arange(n))


def plain_layer(P, name):
    """dict of plain f16 tensors of one layer launch (P: model_enc3.Params)"""
    C, heads, NPM, CIN = LAYERS[name]
    fp8 = lambda T: me.E4M3[P.wbuf[T]]
    pre = 4 * C * CIN + 8 * C if CIN else 0
    L = sm.layout(C, heads, pre)
    g = gp(C)
    t = {'g1': P.f16vec(L['g1'], C)[g], 'bo': P.f16vec(L['bo'], C)[g], 'g2': P.f16vec(L['g2'], C)[g], 'b2': P.f16vec(L['b2'], C)[g],
         'q': [], 'v': [], 'bias': [], 'w1': [], 'b1': [], 'w2': []}
    wo = np.zeros((32 * heads, C))
    g32 = gp(32)
    for h in range(heads):
        hb = L['head'][h]
        t['q'].append(fp8(sm.woff_table(hb, 1024, 512, C, 32)))
        t['v'].append(fp8(sm.woff_table(hb + 32 * C, 1024, 512, C, 32)))
        t['bias'].append(sm.rel_bias(P, hb + 64 * C))
        raw = fp8(sm.woff_table(hb + 64 * C + 512, 0, 512, 32, C))          # rows: pair-ordered attention outputs
        wo[32 * h + g32] = raw[:, g]                                       # Wo[gp(a)][c] = raw[a][gp(c)]
    t['wo'] = wo
    for B1 in L['b1']:
        t['w1'].append(fp8(sm.woff_table(B1, 512, 16 * C, C, 32)))
        t['b1'].append(P.f16vec(B1 + 32 * C, 32))
        raw = fp8(sm.woff_table(B1 + 32 * C + 64, 0, 512, 32, C))
        w2 = np.zeros((32, C)); w2[g32] = raw[:, g]
        t['w2'].append(w2)
    if NPM:
        PM0 = L['end']
        W = np.concatenate([fp8(sm.woff_table(PM0 + k * 128 * C, 512, 64 * C, 4 * C, 32)) for k in range(NPM // 32)], 1)
        b = np.concatenate([P.f16vec(PM0 + 4 * C * NPM + 64 * k, 32) for k in range(NPM // 32)])
        gm = gp(NPM)
        t['merge_w'], t['merge_b'] = W[:, gm], b[gm]
    if CIN:
        cols = 4 * C // heads
        W = np.concatenate([fp8(sm.woff_table(w * CIN * cols, 512, 16 * CIN, CIN, cols)) for w in range(heads)], 1)   # [CIN][4C]
        b = np.concatenate([P.f16vec(4 * C * CIN + 2 * cols * w, cols) for w in range(heads)])
        idx = (np.arange(4)[:, None] * C + g[None, :]).ravel()            # x0[.., q, c] = E[q C + gp(c)]
        t['expand_w'], t['expand_b'] = W[:, idx], b[idx]
    return {k: ([np.asarray(a, F16) for a in v] if isinstance(v, list) else np.asarray(v, F16)) for k, v in t.items()}


def h(a):
    return np.asarray(a, np.float64).astype(F16).astype(np.float64)


def q8(a):
    return me.q8(a)


def softmax_rows(S):
    t = np.clip(h(S * me.C_SCALE), -float(me.CLAMP), float(me.CLAMP))
    poly = h(t * h(t * (-t) + me.P1) + float(me.P0))
    w = me.exp_trick(poly).astype(np.float64)
    s = h(w.sum(1, keepdims=True))
    return h(w * h(1.0 / s))


def block(t, x0, C, heads):
    """x0: [windows][16][C] f16 values; returns the block output before the FP8 rounding"""
    f = lambda a: np.asarray(a, np.float64)
    r = h(1.0 / np.sqrt(h((x0 * x0).sum(-1, keepdims=True) + 2.0 ** -13)))
    h1 = q8(h(x0 * h(r * f(t['g1']))))
    O = np.zeros(x0.shape[:2] + (32 * heads,))
    for k in range(heads):
        Q = h(h1 @ f(t['q'][k])); V = h(h1 @ f(t['v'][k]))
        for w in range(x0.shape[0]):
            Qq = q8(Q[w])
            O[w, :, 32 * k:32 * k + 32] = q8(h(softmax_rows(h(Qq @ Qq.T + f(t['bias'][k]))) @ V[w]))
    x1 = h(O @ f(t['wo']) + x0 + f(t['bo']))
    r = h(1.0 / np.sqrt(h((x1 * x1).sum(-1, keepdims=True) + 2.0 ** -13)))
    h2 = q8(h(x1 * h(r * f(t['g2']))))
    x2 = x1 + f(t['b2'])
    for w1, b1, w2 in zip(t['w1'], t['b1'], t['w2']):
        x2 = x2 + q8(me.gelu_poly(h(h2 @ f(w1) + f(b1)))) @ f(w2)
    return h(x2)


def tokens(buf, off, W, H, C):
    """planes of 32 e4m3 channels at byte `off` -> [H][W][C] values"""
    a = me.E4M3[buf[off:off + (C // 32) * H * W * 32]].reshape(C // 32, H, W, 32)
    return a.transpose(1, 2, 0, 3).reshape(H, W, C)


def mirror(v, n):
    v = np.abs(v); return np.minimum(v, 2 * n - 2 - v)


def forward(t, name, x, skip, sx, sy):
    """x: [H][W][C] input tokens (decoders: [H/2][W/2][CIN]); returns (out [H][W][C], merged or None), FP8-rounded"""
    C, heads, NPM, CIN = LAYERS[name]
    f = lambda a: np.asarray(a, np.float64)
    if CIN:
        H, W = skip.shape[:2]
    else:
        H, W = x.shape[:2]
    gx, gy = (W + sx + 7) // 8, (H + sy + 7) // 8
    Y, X = np.meshgrid(np.arange(8 * gy) - sy, np.arange(8 * gx) - sx, indexing='ij')
    if CIN:
        e = h(x @ f(t['expand_w']) + f(t['expand_b'])).reshape(x.shape[0], x.shape[1], 4, C)
        # the low-resolution token of block row r is taken at (8 bx - sx) // 2 + r % 4 (floor), then mirrored
        Yb, Xb = (Y + sy) // 8 * 8 - sy, (X + sx) // 8 * 8 - sx
        ly, lx = mirror(Yb // 2 + (Y - Yb) // 2, H // 2), mirror(Xb // 2 + (X - Xb) // 2, W // 2)
        q = ((Y - Yb) & 1) * 2 + ((X - Xb) & 1)
        x0 = h(q8(e[ly, lx, q]) + skip[mirror(Y, H), mirror(X, W)])
    else:
        x0 = x[mirror(Y, H), mirror(X, W)]
    win = x0.reshape(2 * gy, 4, 2 * gx, 4, C).transpose(0, 2, 1, 3, 4).reshape(-1, 16, C)
    y = block(t, win, C, heads).reshape(2 * gy, 2 * gx, 4, 4, C).transpose(0, 2, 1, 3, 4).reshape(8 * gy, 8 * gx, C)
    out = q8(y)
    merged = None
    if NPM:
        A = out.reshape(4 * gy, 2, 4 * gx, 2, C).transpose(0, 2, 1, 3, 4).reshape(4 * gy, 4 * gx, 4 * C)
        m = q8(h(A @ f(t['merge_w']) + f(t['merge_b'])))
        # merged token (my, mx) of a block sits at (8 by - sy) // 2 + my
        my = ((np.arange(4 * gy) // 4) * 8 - sy) // 2 + np.arange(4 * gy) % 4
        mx = ((np.arange(4 * gx) // 4) * 8 - sx) // 2 + np.arange(4 * gx) % 4
        merged = np.zeros((H // 2, W // 2, NPM))
        oky, okx = (my >= 0) & (my < H // 2), (mx >= 0) & (mx < W // 2)
        merged[np.ix_(my[oky], mx[okx])] = m[np.ix_(oky, okx)]
    oy, ox = (Y[:, 0] >= 0) & (Y[:, 0] < H), (X[0] >= 0) & (X[0] < W)
    return out[np.ix_(oy, ox)], merged


def launches(cap, frame=0):
    """the ten layer launches of one frame, in network order, as model_enc3.Params"""
    dirs = sorted(d for d in glob.glob(f'{cap}/replay/replay-*-rrlite_*_4x4') if any(d.endswith(f'rrlite_{k}_4x4') for k in LAYERS))
    if len(dirs) < 10 * (frame + 1): raise ValueError('capture lacks the layer launches of that frame')
    return [me.Params(d) for d in dirs[10 * frame:10 * frame + 10]], dirs[10 * frame:10 * frame + 10]


def psnr(a, b):
    d = (a - b) ** 2
    return 10 * np.log10(np.abs(b).max() ** 2 / max(d.mean(), 1e-30))


def main():
    cap = sys.argv[1]; frame = int(sys.argv[2]) if len(sys.argv) > 2 else 0
    Ps, dirs = launches(cap, frame)
    for i, (P, d, name) in enumerate(zip(Ps, dirs, ORDER)):
        C, heads, NPM, CIN = LAYERS[name]
        seq = int(os.path.basename(d).split('-')[1])
        after = np.fromfile(glob.glob(f'{cap}/dump/launch-{seq:03d}-*-arg024-*')[0], np.uint8)
        t = plain_layer(P, name)
        W, H = P.W, P.H
        if CIN:
            x = tokens(P.arena, P.inp - P.abase, W // 2, H // 2, CIN); skip = tokens(P.arena, P.p32 - P.abase, W, H, C)
        else:
            x = tokens(P.arena, P.inp - P.abase, W, H, C); skip = None
        out, merged = forward(t, name, x, skip, P.sx, P.sy)
        ref = tokens(after, P.out - P.abase, W, H, C)
        line = f'{name:10s} launch {seq:3d} {W}x{H}x{C} shift {P.sx},{P.sy}: output {psnr(out, ref):5.1f} dB, identical {100 * (out == ref).mean():5.1f}%'
        if NPM:
            mref = tokens(after, P.p48 - P.abase, W // 2, H // 2, NPM)
            line += f'; merged {psnr(merged, mref):5.1f} dB, identical {100 * (merged == mref).mean():5.1f}%'
        print(line, flush=True)


if __name__ == '__main__':
    main()