Clear windows branch; preserve previous contents in win-old
This commit is contained in:
ayo committed
1 parent
f0c881985b
commit
73994d4634
311 files changed
-57018
No files matched your search
@@ -1,440 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Numpy reference of DLSS 4 (preset K) DltssPaddedWinLayer blocks, built stage by stage against the
|
||||
PTX interpreter (kernels/tools/ptxsim.py). f16 arithmetic in the kernel's operation order."""
|
||||
import struct
|
||||
import numpy as np
|
||||
|
||||
F16, F32, F64 = np.float16, np.float32, np.float64
|
||||
|
||||
# Legacy replay reports retain their original approximate norm by default.
|
||||
# Windows strict K validation selects the PTX fragment's rounded reduction.
|
||||
NORM_FRAGMENT_ORDER = False
|
||||
|
||||
|
||||
def norm_rows(x):
|
||||
if not NORM_FRAGMENT_ORDER:
|
||||
return f16(np.sum(x.astype(F32).astype(F64) ** 2, axis=1))
|
||||
tokens, channels = x.shape
|
||||
assert channels in (32, 64, 96, 128, 160)
|
||||
square = f16(x.astype(F64) ** 2)
|
||||
partial = []
|
||||
for lane in range(4):
|
||||
halves = []
|
||||
for parity in range(2):
|
||||
pairs = [add16(square[:, 2 * lane + 8 * parity + 32 * g:2 * lane + 8 * parity + 32 * g + 2],
|
||||
square[:, 2 * lane + 8 * parity + 32 * g + 16:2 * lane + 8 * parity + 32 * g + 18])
|
||||
for g in range(channels // 32)]
|
||||
value = pairs[0]
|
||||
if channels >= 64:
|
||||
value = add16(value, pairs[1])
|
||||
if channels == 96:
|
||||
value = add16(value, pairs[2])
|
||||
if channels >= 128:
|
||||
value = add16(value, add16(pairs[2], pairs[3]))
|
||||
if channels == 160:
|
||||
value = add16(value, pairs[4])
|
||||
halves.append(value)
|
||||
partial.append(add16(halves[0], halves[1]))
|
||||
total = add16(add16(partial[0], partial[1]), add16(partial[2], partial[3]))
|
||||
return add16(total[:, 0], total[:, 1])
|
||||
|
||||
|
||||
class Dump:
|
||||
def __init__(self, d):
|
||||
self.dir = d
|
||||
man = open(d + '/manifest.txt').read().split('\n')
|
||||
self.kernel = man[0].split()[1]
|
||||
la = man[1].split()
|
||||
self.grid, self.block = tuple(map(int, la[1:4])), tuple(map(int, la[4:7]))
|
||||
self.args = open(d + '/args.bin', 'rb').read()
|
||||
self.allocs = {}
|
||||
for l in man:
|
||||
p = l.split()
|
||||
if p and p[0] == 'alloc':
|
||||
self.allocs[int(p[1])] = (int(p[2], 16), np.fromfile(f'{d}/alloc-{p[1]}.bin', np.uint8))
|
||||
|
||||
def u64(self, off):
|
||||
return struct.unpack_from('<Q', self.args, off)[0]
|
||||
|
||||
def i32x2(self, off):
|
||||
return struct.unpack_from('<ii', self.args, off)
|
||||
|
||||
def ptr(self, off):
|
||||
a = self.u64(off)
|
||||
for i, (base, buf) in self.allocs.items():
|
||||
if base <= a < base + buf.size:
|
||||
return buf, a - base
|
||||
raise ValueError(hex(a))
|
||||
|
||||
def f16(self, off_param, byte, n):
|
||||
buf, o = self.ptr(off_param)
|
||||
return buf[o + byte:o + byte + 2 * n].view(F16).copy()
|
||||
|
||||
|
||||
def f16(x):
|
||||
return np.asarray(x, F64).astype(F16)
|
||||
|
||||
|
||||
def mirror(v, n):
|
||||
v = abs(v)
|
||||
return min(v, 2 * n - 2 - v)
|
||||
|
||||
|
||||
def load_window(P, bx, by, C):
|
||||
"""64 tokens (raster order of the 8x8 window) x C channels, f16, mirror padded."""
|
||||
W, H = P.i32x2(0)
|
||||
sx, sy = P.i32x2(56)
|
||||
buf, o = P.ptr(8)
|
||||
x = np.zeros((64, C), F16)
|
||||
for r in range(64):
|
||||
ty, tx = r >> 3, r & 7
|
||||
Y = mirror(8 * by - sy + ty, H)
|
||||
X = mirror(8 * bx - sx + tx, W)
|
||||
a = o + 2 * ((Y * W + X) * C)
|
||||
x[r] = buf[a:a + 2 * C].view(F16)
|
||||
return x
|
||||
|
||||
|
||||
# --- reductions in the kernel's order ---------------------------------------------------------
|
||||
def sumsq_rows(x):
|
||||
"""f16 sum of squares per token, reproducing the kernel's tree for C = 64 (per lane: 16 products
|
||||
over 8 words (t + 4k), pairwise f16x2 adds, butterfly over the 4 lanes of a row, final hi+lo)."""
|
||||
T, C = x.shape
|
||||
assert C == 64
|
||||
sq = f16(x.astype(F64) * x.astype(F64)) # f16 products
|
||||
out = np.zeros(T, F16)
|
||||
for r in range(T):
|
||||
lanes = []
|
||||
for t in range(4):
|
||||
# word k holds channels 2(t+4k), 2(t+4k)+1 as an f16 pair
|
||||
w = [np.array([sq[r, 2 * (t + 4 * k)], sq[r, 2 * (t + 4 * k) + 1]], F16) for k in range(8)]
|
||||
a = [f16(w[2 * i].astype(F64) + w[2 * i + 1].astype(F64)) for i in range(4)] # (0+1),(2+3),(4+5),(6+7)
|
||||
b = [f16(a[0].astype(F64) + a[2].astype(F64)), f16(a[1].astype(F64) + a[3].astype(F64))]
|
||||
lanes.append(f16(b[0].astype(F64) + b[1].astype(F64)))
|
||||
s1 = [f16(lanes[t].astype(F64) + lanes[t ^ 1].astype(F64)) for t in range(4)]
|
||||
s2 = f16(s1[0].astype(F64) + s1[2].astype(F64))
|
||||
out[r] = f16(s2[0].astype(F64) + s2[1].astype(F64))
|
||||
return out
|
||||
|
||||
|
||||
def rsqrt16(s):
|
||||
return f16(1.0 / np.sqrt(s.astype(F32).astype(F64)))
|
||||
|
||||
|
||||
def mma(A, B, C):
|
||||
"""one wmma k16 step: D = f16(A @ B + C) with exact sums (the interpreter's semantics)."""
|
||||
return f16(A.astype(F64) @ B.astype(F64) + C.astype(F64))
|
||||
|
||||
|
||||
def frag_offset(k, n):
|
||||
"""byte offset of logical B[k][n] (k, n < 16) inside one 512-byte pre-swizzled f16 fragment"""
|
||||
return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1)
|
||||
|
||||
|
||||
FRAG = np.array([[frag_offset(k, n) for n in range(16)] for k in range(16)])
|
||||
|
||||
|
||||
def weights(P, byte, K, N, nt_stride=None):
|
||||
"""logical W[K][N] from fragments at weights + byte + 512 * kt + nt_stride * nt"""
|
||||
buf, o = P.ptr(64)
|
||||
nt_stride = nt_stride or 512 * (K // 16)
|
||||
W = np.zeros((K, N), F16)
|
||||
for kt in range(K // 16):
|
||||
for nt in range(N // 16):
|
||||
base = o + byte + 512 * kt + nt_stride * nt
|
||||
vals = buf[(base + FRAG)[..., None] + np.arange(2)].reshape(16, 16, 2).view(F16)[..., 0]
|
||||
W[16 * kt:16 * kt + 16, 16 * nt:16 * nt + 16] = vals
|
||||
return W
|
||||
|
||||
|
||||
def gemm(A, W, C0=None):
|
||||
"""chained wmma k16 steps (f16 rounding after each step), per 16x16 output tile"""
|
||||
M, K = A.shape
|
||||
N = W.shape[1]
|
||||
D = np.zeros((M, N), F16) if C0 is None else C0.astype(F16).copy()
|
||||
for kt in range(K // 16):
|
||||
D = mma(A[:, 16 * kt:16 * kt + 16], W[16 * kt:16 * kt + 16], D)
|
||||
return D
|
||||
|
||||
|
||||
def hv(bits):
|
||||
return np.array(bits, np.uint16).view(F16)
|
||||
|
||||
|
||||
def fma16(a, b, c):
|
||||
return f16(a.astype(F64) * b.astype(F64) + c.astype(F64))
|
||||
|
||||
|
||||
def add16(a, b):
|
||||
return f16(a.astype(F64) + b.astype(F64))
|
||||
|
||||
|
||||
def mul16(a, b):
|
||||
return f16(a.astype(F64) * b.astype(F64))
|
||||
|
||||
|
||||
def softmax_exp(S):
|
||||
"""e = ex2(t * (1 + a t^2) + c), t = clamp(S, +-21.64) (f16x2 ops of the kernel)"""
|
||||
hi, lo = hv(0x4D69), hv(0xCD69)
|
||||
a, one, c = hv(0x91D5), hv(0x3C00), hv(0xC50D)
|
||||
t = np.maximum(np.minimum(S, hi), lo)
|
||||
u = mul16(t, t)
|
||||
v = fma16(u, a, one)
|
||||
w = fma16(t, v, c)
|
||||
return f16(np.exp2(w.astype(F64)))
|
||||
|
||||
|
||||
def gelu16(x):
|
||||
"""x * (0.5 + c (0.41216 - 0.081081 |c|)), c = clamp(x, +-2) (f16x2 ops of the kernel)"""
|
||||
k1, k2 = f16(np.float32(np.uint32(0x3ED306EB).view(F32))), f16(np.float32(np.uint32(0x3DA60DD6).view(F32)))
|
||||
half, two = F16(0.5), F16(2.0)
|
||||
c = np.minimum(np.maximum(x, -two), two)
|
||||
u = f16(k1.astype(F64) - mul16(np.full_like(c, k2), np.abs(c)).astype(F64))
|
||||
return mul16(x, add16(np.full_like(c, half), mul16(c, u)))
|
||||
|
||||
|
||||
def bias_table(P, byte):
|
||||
"""64x64 f16 table stored in C-fragment order, 512 bytes per 16x16 tile (tile 4 mt + nt)"""
|
||||
buf, o = P.ptr(64)
|
||||
B = np.zeros((64, 64), F16)
|
||||
for i in range(64):
|
||||
for j in range(64):
|
||||
mt, nt, r, c = i >> 4, j >> 4, i & 15, j & 15
|
||||
off = 512 * (4 * nt + mt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3)
|
||||
B[i, j] = buf[o + byte + off:o + byte + off + 2].view(F16)[0]
|
||||
return B
|
||||
|
||||
|
||||
def vec(P, byte, n):
|
||||
return P.f16(64, byte, n)
|
||||
|
||||
|
||||
def pwin_enc2(P, bx, by, pm_order=((0, 0), (0, 1), (1, 0), (1, 1)), b2_in_c0=True, dbg=None):
|
||||
"""logical enc2 layer (H 2, C 64 -> merged 96) for one window; returns (y [64 x 64], merged [16 x 96])"""
|
||||
C, H, Dh = 64, 2, 32
|
||||
x0 = load_window(P, bx, by, C)
|
||||
s = norm_rows(x0)
|
||||
r = rsqrt16(s)
|
||||
h1 = mul16(x0, mul16(r[:, None], vec(P, 0, C)[None, :]))
|
||||
Os = []
|
||||
for h in range(H):
|
||||
hb = 128 + 12288 * h
|
||||
q, k, v = [gemm(h1[:, 32:64], weights(P, hb + 2048 * (j + 3), 32, 32, 1024),
|
||||
gemm(h1[:, 0:32], weights(P, hb + 2048 * j, 32, 32, 1024))) for j in range(3)]
|
||||
S = gemm(q, k.T.copy(), bias_table(P, 24704 + 8192 * h))
|
||||
E = softmax_exp(S)
|
||||
rs = f16(np.sum(E.astype(F64), axis=1))
|
||||
rc = f16(1.0 / rs.astype(F32).astype(F64))
|
||||
Pm = mul16(E, rc[:, None])
|
||||
Os.append(gemm(Pm, v))
|
||||
if dbg is not None:
|
||||
dbg[f'q{h}'], dbg[f'k{h}'], dbg[f'v{h}'], dbg[f'S{h}'], dbg[f'P{h}'], dbg[f'O{h}'] = q, k, v, S, Pm, Os[-1]
|
||||
O = np.concatenate(Os, axis=1)
|
||||
bo = vec(P, 41088, C)
|
||||
Y = gemm(O[:, 32:64], weights(P, 45312, 32, 64, 1024), gemm(O[:, 0:32], weights(P, 41216, 32, 64, 1024), np.tile(bo, (64, 1))))
|
||||
x1 = add16(Y, x0)
|
||||
if dbg is not None:
|
||||
dbg['Y'], dbg['x1'] = Y, x1
|
||||
m = mul16(x1, vec(P, 49408, C)[None, :])
|
||||
b2 = vec(P, 49536, C)
|
||||
acc = add16(x1, np.tile(b2, (64, 1))) if b2_in_c0 else x1
|
||||
for c in range(8):
|
||||
Hc = gemm(m, weights(P, 49664 + 4096 * c, 64, 32, 2048), np.tile(vec(P, 82432 + 64 * c, 32), (64, 1)))
|
||||
acc = gemm(gelu16(Hc), weights(P, 82944 + 4096 * c, 32, 64, 1024), acc)
|
||||
if dbg is not None:
|
||||
dbg[f'H{c}'], dbg[f'acc{c}'] = Hc, acc
|
||||
y = acc
|
||||
# patch merge: 16 merged tokens (4x4), K = 4 sub-tokens x 64
|
||||
A = np.zeros((16, 256), F16)
|
||||
for mt in range(16):
|
||||
my, mx = mt >> 2, mt & 3
|
||||
for q, (dy, dx) in enumerate(pm_order):
|
||||
A[mt, 64 * q:64 * q + 64] = y[8 * (2 * my + dy) + 2 * mx + dx]
|
||||
merged = np.concatenate([gemm(A, weights(P, 115712 + 24576 * w, 256, 48, 8192), np.tile(vec(P, 164864 + 96 * w, 48), (16, 1)))
|
||||
for w in range(2)], axis=1)
|
||||
return y, merged
|
||||
|
||||
|
||||
class Layout:
|
||||
"""weight-buffer layout of a plain DltssPaddedWinLayer (encoder with patch merge), bytes"""
|
||||
def __init__(self, H, C, COUT):
|
||||
self.H, self.C, self.COUT = H, C, COUT
|
||||
self.G1 = 0
|
||||
self.QKV = 2 * C
|
||||
self.HEAD = 192 * C
|
||||
self.BIAS = self.QKV + self.HEAD * H
|
||||
self.BO = self.BIAS + 8192 * H
|
||||
self.WO = self.BO + 2 * C
|
||||
self.G2 = self.WO + 64 * C * H
|
||||
self.B2 = self.G2 + 2 * C
|
||||
self.W1 = self.B2 + 2 * C
|
||||
self.B1 = self.W1 + 8 * C * C
|
||||
self.W2 = self.B1 + 8 * C
|
||||
self.PM = self.W2 + 8 * C * C
|
||||
self.NPW = COUT // H # merged channels per original warp
|
||||
self.NPA = -(-self.NPW // 16) * 16 # allocated (padded to whole 16-column tiles)
|
||||
self.PMB = self.PM + 8 * C * self.NPA * H
|
||||
|
||||
|
||||
def swin_core(P, L, x0, dbg=None, posattn=False):
|
||||
"""norm -> attention -> Wo + residual -> MLP (+ residual): the block output y [64 x C]"""
|
||||
C, H = L.C, L.H
|
||||
NCH = C // 32
|
||||
s = norm_rows(x0)
|
||||
r = rsqrt16(s)
|
||||
h1 = mul16(x0, mul16(r[:, None], vec(P, L.G1, C)[None, :]))
|
||||
Os = []
|
||||
for h in range(H):
|
||||
hb = L.QKV + L.HEAD * h
|
||||
proj = []
|
||||
for j in range(3):
|
||||
acc = None
|
||||
for ch in range(NCH):
|
||||
acc = gemm(h1[:, 32 * ch:32 * ch + 32], weights(P, hb + 2048 * (3 * ch + j), 32, 32, 1024), acc)
|
||||
proj.append(acc)
|
||||
q, k, v = proj
|
||||
if posattn:
|
||||
# position-only attention (trait flag 2): the table is the probability matrix itself
|
||||
Os.append(gemm(bias_table(P, L.BIAS + 8192 * h), v))
|
||||
if dbg is not None:
|
||||
dbg[f'v{h}'], dbg[f'O{h}'] = v, Os[-1]
|
||||
continue
|
||||
S = gemm(q, k.T.copy(), bias_table(P, L.BIAS + 8192 * h))
|
||||
E = softmax_exp(S)
|
||||
rs = f16(np.sum(E.astype(F64), axis=1))
|
||||
rc = f16(1.0 / rs.astype(F32).astype(F64))
|
||||
Os.append(gemm(mul16(E, rc[:, None]), v))
|
||||
if dbg is not None:
|
||||
dbg[f'q{h}'], dbg[f'k{h}'], dbg[f'v{h}'], dbg[f'O{h}'], dbg[f'S{h}'], dbg[f'E{h}'], dbg[f'rs{h}'] = q, k, v, Os[-1], S, E, rs
|
||||
Y = np.tile(vec(P, L.BO, C), (64, 1))
|
||||
for h in range(H):
|
||||
Y = gemm(Os[h], weights(P, L.WO + 64 * C * h, 32, C, 1024), Y)
|
||||
x1 = add16(Y, x0)
|
||||
m = mul16(x1, vec(P, L.G2, C)[None, :])
|
||||
acc = add16(x1, np.tile(vec(P, L.B2, C), (64, 1)))
|
||||
if dbg is not None:
|
||||
dbg['Y'], dbg['x1'] = Y, x1
|
||||
for c in range(C // 8):
|
||||
Hc = gemm(m, weights(P, L.W1 + 64 * C * c, C, 32, 32 * C), np.tile(vec(P, L.B1 + 64 * c, 32), (64, 1)))
|
||||
acc = gemm(gelu16(Hc), weights(P, L.W2 + 64 * C * c, 32, C, 1024), acc)
|
||||
if dbg is not None:
|
||||
dbg[f'H{c}'], dbg[f'acc{c}'] = Hc, acc
|
||||
return acc
|
||||
|
||||
|
||||
def pwin_encoder(P, bx, by, H, C, COUT, pm_bias_padded=True, posattn=False):
|
||||
L = Layout(H, C, COUT)
|
||||
x0 = load_window(P, bx, by, C)
|
||||
y = swin_core(P, L, x0, posattn=posattn)
|
||||
A = np.zeros((16, 4 * C), F16)
|
||||
for mt in range(16):
|
||||
my, mx = mt >> 2, mt & 3
|
||||
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
|
||||
A[mt, C * q:C * q + C] = y[8 * (2 * my + dy) + 2 * mx + dx]
|
||||
Wp = weights(P, L.PM, 4 * C, L.NPA * H, 128 * C)
|
||||
bias = vec(P, L.PMB, L.NPA * H) if pm_bias_padded else None
|
||||
cols = []
|
||||
for w in range(H):
|
||||
sl = slice(L.NPA * w, L.NPA * w + L.NPW)
|
||||
b = bias[sl] if pm_bias_padded else vec(P, L.PMB + 2 * L.NPW * w, L.NPW)
|
||||
cols.append(gemm(A, Wp[:, L.NPA * w:L.NPA * (w + 1)], np.tile(np.concatenate([b, np.zeros(L.NPA - L.NPW, F16)]), (16, 1)))[:, :L.NPW])
|
||||
merged = np.concatenate(cols, axis=1)
|
||||
return y, merged
|
||||
|
||||
|
||||
class Shifted:
|
||||
"""view of the weight buffer starting at a byte offset (decoders: core weights follow the expand)"""
|
||||
def __init__(self, P, base):
|
||||
self.P, self.base = P, base
|
||||
|
||||
def ptr(self, off):
|
||||
buf, o = self.P.ptr(off)
|
||||
return buf, o + (self.base if off == 64 else 0)
|
||||
|
||||
def f16(self, off_param, byte, n):
|
||||
return self.P.f16(off_param, byte + (self.base if off_param == 64 else 0), n)
|
||||
|
||||
def i32x2(self, off):
|
||||
return self.P.i32x2(off)
|
||||
|
||||
|
||||
def load_lowres(P, bx, by, CL):
|
||||
W, H = P.i32x2(0)
|
||||
sx, sy = P.i32x2(56)
|
||||
buf, o = P.ptr(8)
|
||||
W2, H2 = W // 2, H // 2
|
||||
lx0, ly0 = int((8 * bx - sx) / 2), int((8 * by - sy) / 2)
|
||||
x = np.zeros((16, CL), F16)
|
||||
for r in range(16):
|
||||
Y = mirror(ly0 + (r >> 2), H2); X = mirror(lx0 + (r & 3), W2)
|
||||
a = o + 2 * ((Y * W2 + X) * CL)
|
||||
x[r] = buf[a:a + 2 * CL].view(F16)
|
||||
return x
|
||||
|
||||
|
||||
def load_skip(P, bx, by, C):
|
||||
W, H = P.i32x2(0)
|
||||
sx, sy = P.i32x2(56)
|
||||
buf, o = P.ptr(16)
|
||||
x = np.zeros((64, C), F16)
|
||||
for r in range(64):
|
||||
Y = mirror(8 * by - sy + (r >> 3), H); X = mirror(8 * bx - sx + (r & 7), W)
|
||||
a = o + 2 * ((Y * W + X) * C)
|
||||
x[r] = buf[a:a + 2 * C].view(F16)
|
||||
return x
|
||||
|
||||
|
||||
def pwin_decoder(P, bx, by, H, C, CL, order=((0, 0), (0, 1), (1, 0), (1, 1))):
|
||||
xl = load_lowres(P, bx, by, CL)
|
||||
EXP = 2 * CL * 4 * C
|
||||
E = gemm(xl, weights(P, 0, CL, 4 * C, 512 * (CL // 16)), np.tile(vec(P, EXP, 4 * C), (16, 1)))
|
||||
skip = load_skip(P, bx, by, C)
|
||||
x0 = np.zeros((64, C), F16)
|
||||
for r in range(16):
|
||||
ly, lx = r >> 2, r & 3
|
||||
for q, (dy, dx) in enumerate(order):
|
||||
t = 8 * (2 * ly + dy) + 2 * lx + dx
|
||||
x0[t] = add16(E[r, q * C:(q + 1) * C], skip[t])
|
||||
L = Layout(H, C, C)
|
||||
return swin_core(Shifted(P, EXP + 8 * C), L, x0)
|
||||
|
||||
|
||||
def pwin_enc0(P, bx, by, H=2, C=32, COUT=64, CIN=16):
|
||||
"""first layer: token embedding x0 = in (CIN) W_e + b_e, then the core (position-only attention) and merge"""
|
||||
xin = load_window(P, bx, by, CIN)
|
||||
x0 = gemm(xin, weights(P, 2 * C, CIN, C, 512 * (CIN // 16)), np.tile(vec(P, 0, C), (64, 1)))
|
||||
x0 = np.maximum(x0, F16(0)) # ReLU after the embedding
|
||||
CB = 2 * C + 2 * CIN * C
|
||||
L = Layout(H, C, COUT)
|
||||
Pw = Shifted(P, CB)
|
||||
y = swin_core(Pw, L, x0, posattn=True)
|
||||
A = np.zeros((16, 4 * C), F16)
|
||||
for mt in range(16):
|
||||
my, mx = mt >> 2, mt & 3
|
||||
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
|
||||
A[mt, C * q:C * q + C] = y[8 * (2 * my + dy) + 2 * mx + dx]
|
||||
Wp = weights(Pw, L.PM, 4 * C, L.NPA * H, 128 * C)
|
||||
bias = vec(Pw, L.PMB, L.NPA * H)
|
||||
cols = [gemm(A, Wp[:, L.NPA * w:L.NPA * (w + 1)], np.tile(bias[L.NPA * w:L.NPA * (w + 1)], (16, 1)))[:, :L.NPW] for w in range(H)]
|
||||
return y, np.concatenate(cols, axis=1)
|
||||
|
||||
|
||||
def pwin_dec0(P, bx, by, H=2, C=32, CL=64, NOUT=40, NOUTA=48):
|
||||
"""last layer: patch expand + skip, core (position-only attention), output head y W_h + b_h (NOUT of NOUTA)"""
|
||||
xl = load_lowres(P, bx, by, CL)
|
||||
EXP = 2 * CL * 4 * C
|
||||
E = gemm(xl, weights(P, 0, CL, 4 * C, 512 * (CL // 16)), np.tile(vec(P, EXP, 4 * C), (16, 1)))
|
||||
skip = load_skip(P, bx, by, C)
|
||||
x0 = np.zeros((64, C), F16)
|
||||
for r in range(16):
|
||||
ly, lx = r >> 2, r & 3
|
||||
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
|
||||
t = 8 * (2 * ly + dy) + 2 * lx + dx
|
||||
x0[t] = add16(E[r, q * C:(q + 1) * C], skip[t])
|
||||
L = Layout(H, C, C)
|
||||
Pw = Shifted(P, EXP + 8 * C)
|
||||
y = swin_core(Pw, L, x0, posattn=True)
|
||||
HB = L.PM # head bias, then the head weights
|
||||
Wh = weights(Pw, HB + 2 * NOUTA, C, NOUTA, 1024)
|
||||
out = gemm(y, Wh, np.tile(vec(Pw, HB, NOUTA), (64, 1)))
|
||||
return y, out[:, :NOUT]
|
||||
Reference in new issue
Block a user