ayo
7c5d3c9157
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
142 lines
4.9 KiB
Python
142 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Numeric primitives and replay-dump access for swin_model.py (the rrlite Swin blocks of DLSS 4.5, preset M).
|
|
|
|
Everything here is NVIDIA's arithmetic as the native kernels reproduce it (kernels/m/swin_common.h,
|
|
swin_block.h): f16 rounding after every operation, e4m3 (OCP, satfinite, RNE) re-quantisation, one f16 rounding
|
|
per k32 MMA step, and the softmax's bit-level exponent trick.
|
|
"""
|
|
import struct
|
|
|
|
import numpy as np
|
|
|
|
F16, F32, F64 = np.float16, np.float32, np.float64
|
|
|
|
|
|
def f16(x):
|
|
"""round to f16; the result stays float64 (on the f16 grid), so products and sums of such values are exact
|
|
until the next f16() (numpy would otherwise do f16 arithmetic, e.g. in matmul)"""
|
|
return np.asarray(x, F64).astype(F16).astype(F64)
|
|
|
|
|
|
def _e4m3_table():
|
|
t = np.zeros(256, F64)
|
|
for c in range(256):
|
|
s, e, m = c >> 7, (c >> 3) & 15, c & 7
|
|
if e == 15 and m == 7:
|
|
v = np.nan
|
|
elif e == 0:
|
|
v = m * 2.0 ** -9
|
|
else:
|
|
v = (1 + m / 8) * 2.0 ** (e - 7)
|
|
t[c] = -v if s else v
|
|
return t
|
|
|
|
|
|
E4M3 = _e4m3_table()
|
|
_POS = E4M3[:127] # codes 0..126: the finite non-negative values, ascending
|
|
|
|
|
|
def q8(x):
|
|
"""e4m3 RNE satfinite round trip of f16 values (NaN stays NaN), as f16"""
|
|
h = f16(x).astype(F64)
|
|
a = np.abs(h)
|
|
i = np.clip(np.searchsorted(_POS, a), 1, 126)
|
|
lo, hi = _POS[i - 1], _POS[i]
|
|
# nearest; ties to the even code
|
|
pick_hi = (hi - a < a - lo) | ((hi - a == a - lo) & ((i & 1) == 0))
|
|
v = np.where(pick_hi, hi, lo)
|
|
v = np.where(a >= _POS[126], _POS[126], v)
|
|
v = np.where(a <= 0, 0.0, v)
|
|
v = np.where(np.isnan(h), np.nan, np.copysign(v, h))
|
|
return f16(v)
|
|
|
|
|
|
# softmax constants (swin_block.h stage 2)
|
|
C_SCALE = f16(0.0972222164273262 / 5.656854152679443)
|
|
CLAMP = F16(0.55615234375)
|
|
P1 = f16(0.92730712890625)
|
|
P0 = F16(1.375)
|
|
|
|
|
|
def exp_trick(poly):
|
|
"""per row, columns (2j, 2j+1): the 32-bit word even | odd << 16 of the f16 bits, << 5, + 0x7FF88000; the even
|
|
column takes its low half, the odd column its high half (as f16 bits)"""
|
|
b = poly.astype(F16).view(np.uint16).astype(np.uint64)
|
|
out = np.empty_like(b)
|
|
even, odd = b[:, 0::2], b[:, 1::2]
|
|
word = ((even | (odd << np.uint64(16))) << np.uint64(5)) + np.uint64(0x7FF88000)
|
|
word &= np.uint64(0xFFFFFFFF)
|
|
out[:, 0::2] = word & np.uint64(0xFFFF)
|
|
out[:, 1::2] = word >> np.uint64(16)
|
|
return out.astype(np.uint16).view(F16)
|
|
|
|
|
|
def row_sum16(w):
|
|
"""f16 butterfly over the 16 columns: partner c ^ 8, then ^ 2, ^ 4, ^ 1 (every column ends with the same sum)"""
|
|
s = w.astype(F16)
|
|
c = np.arange(16)
|
|
for m in (8, 2, 4, 1):
|
|
s = f16(s.astype(F64) + s[:, c ^ m].astype(F64))
|
|
return s
|
|
|
|
|
|
def gelu_poly(z):
|
|
z = f16(z)
|
|
cz = np.clip(z, F16(-2.0), F16(2.0))
|
|
inner = f16(f16(0.41216981).astype(F64) - f16(f16(0.08108133).astype(F64) * np.abs(cz).astype(F64)).astype(F64))
|
|
return f16(z.astype(F64) * f16(0.5 + f16(cz.astype(F64) * inner.astype(F64)).astype(F64)).astype(F64))
|
|
|
|
|
|
def mma_chain(A, W, C0=None):
|
|
"""chained k32 MMA steps with f16 accumulation: D = f16(A[:, 32k:32k+32] W[32k:32k+32] + D), exact sums"""
|
|
M, K = A.shape
|
|
D = np.zeros((M, W.shape[1]), F64) if C0 is None else f16(C0)
|
|
A = np.asarray(A, F64)
|
|
W = np.asarray(W, F64)
|
|
for k in range(0, K, 32):
|
|
D = f16(A[:, k:k + 32] @ W[k:k + 32] + D.astype(F64))
|
|
return D
|
|
|
|
|
|
class Params:
|
|
"""one rrlite Swin-block launch from a replay dump (CommonParams: w 0, sx sy 8, tw th 16, in 24, p32 32,
|
|
out 40, p48 48)"""
|
|
|
|
def __init__(self, d):
|
|
man = open(d + '/manifest.txt').read().split('\n')
|
|
self.kernel = man[0].split()[1]
|
|
la = man[1].split()
|
|
self.grid = tuple(map(int, la[1:4]))
|
|
self.args = open(d + '/args.bin', 'rb').read()
|
|
self.allocs = {}
|
|
for line in man:
|
|
p = line.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))
|
|
self.w_ptr, = struct.unpack_from('<Q', self.args, 0)
|
|
self.sx, self.sy, self.W, self.H = struct.unpack_from('<iiii', self.args, 8)
|
|
self.inp, self.p32, self.out, self.p48 = struct.unpack_from('<QQQQ', self.args, 24)
|
|
wbase, wbuf = self.alloc_of(self.w_ptr)
|
|
self.wbuf = wbuf[self.w_ptr - wbase:]
|
|
self.abase, self.arena = self.alloc_of(self.inp)
|
|
|
|
def alloc_of(self, a):
|
|
for base, buf in self.allocs.values():
|
|
if base <= a < base + buf.size:
|
|
return base, buf
|
|
raise ValueError(hex(a))
|
|
|
|
def f16vec(self, off, n):
|
|
return self.wbuf[off:off + 2 * n].view(F16).copy()
|
|
|
|
def token_xy(self, bx, by):
|
|
T = np.arange(64)
|
|
tx = (T & 3) + 4 * ((T >> 4) & 1)
|
|
ty = ((T >> 2) & 3) + 4 * (T >> 5)
|
|
return 8 * bx - self.sx + tx, 8 * by - self.sy + ty
|
|
|
|
@staticmethod
|
|
def mirror(v, n):
|
|
v = np.abs(v)
|
|
return np.minimum(v, 2 * n - 2 - v)
|