Add RDNA4 (gfx12) support and a 0.1.2 optimization round
RDNA4 (experimental; untested on RDNA4 hardware) - kernels/common/wmma_layout.h picks the WMMA operand/accumulator layout per target; K (pwin_*) and M (swin_*) index through it, so the same sources build for gfx11 and gfx12. D4R_WMMA_LAYOUT=12 on gfx11 is a test shim that runs the gfx12 indexing on gfx11 WMMA. - M can use native e4m3 FP8 WMMA on gfx12 (D4R_NATIVE_FP8=1). Texture tails (enc0_tail, dec0_head) get gfx12 f16 and FP8 paths that match ZLUDA's gfx12 lowering. - ZLUDA patch 0007: gfx12 WMMA helpers, optional native FP8 MMA (D4R_ZLUDA_WMMA_FP8_NATIVE), a layout override for testing, and an architecture argument for d4r_emit. Texture kernels are now compiled offline per target through d4r_emit (D4R_ZLUDA_EMIT) instead of on the local GPU. - New d4r.ini option [Kernels] NativeFp8 (default on), wired through the shim and d4r_config.py. On gfx12 the bridge serves kernels/<arch>-fp8 when it is on. - The bridge picks the KFD GPU with the most SIMDs (so an iGPU is not chosen over a discrete GPU); D4R_GPU_ARCH overrides it. - kernels/build.sh and package_release.sh build gfx1100-gfx1103 and gfx1200/gfx1201 (plus -fp8 variants), and D4R_BUNDLE_TEX takes per-target folders. Performance (RDNA3; harness byte-identical unless noted) - ZLUDA patch 0006: m16n8k8 f16 MMAs on WMMA (E; not bit-exact, 52-55 dB per layer, D4R_ZLUDA_WMMA_K8=0 disables it), per-key weight-image slots for native prep kernels (d4r_prep_key_at, d4r_prep_key_slots), and wave64 builds of MMA-free kernels (D4R_ZLUDA_WAVE64). - M: packed FP8 encoding, conflict-free patch-merge reads, scalar weight loads in enc1, and weight prep once per weight set for the six enc3 tube blocks. - K: CU mode for the six non-wide pwin layers. - Texture kernels: every DLSS flag variant is built natively (24 K output, 8 M post, 4 M downsample, 4 M enc0), not only Townfall's; post and K output run as wave64. - M direct output: rrlite_downsample_kernel writes the shim's output buffer; the shim's default direct-output presets are now 11,13. Tools - dump_runner takes the prep grid from d4r_prep_blocks and refuses texture/surface dumps. - New pwin_check.py, swin_check.py, model_enc3.py and tail_check.hip. Docs: README GPU table and RDNA4 notes, building (patches 0002-0007, d4r_emit), native-kernels (RDNA4 section), performance (new numbers), packaging README. Version 0.1.2.
This commit is contained in:
ayo committed
1 parent
f310d52b69
commit
4334a0c3c4
42 files changed
+3068
-529
No files matched your search
@@ -44,7 +44,18 @@ int main(int argc, char** argv)
|
||||
hipFunction_t prep = nullptr;
|
||||
if (hipModuleGetFunction(&prep, mod, (std::string(argv[2]) + "_prep").c_str()) != hipSuccess)
|
||||
prep = nullptr;
|
||||
const unsigned prepGrid = getenv("PREP_GRID") ? atoi(getenv("PREP_GRID")) : 88;
|
||||
// the prep grid: PREP_GRID, else the module's d4r_prep_blocks (what the ZLUDA hook launches), else 88
|
||||
unsigned prepGrid = 88;
|
||||
{
|
||||
hipDeviceptr_t pb;
|
||||
size_t pbs = 0;
|
||||
uint32_t v = 0;
|
||||
if (hipModuleGetGlobal(&pb, &pbs, mod, "d4r_prep_blocks") == hipSuccess && pbs == 4 &&
|
||||
hipMemcpyDtoH(&v, pb, 4) == hipSuccess && v != 0)
|
||||
prepGrid = v;
|
||||
}
|
||||
if (getenv("PREP_GRID"))
|
||||
prepGrid = atoi(getenv("PREP_GRID"));
|
||||
|
||||
std::vector<char> args = read_file(dump + "/args.bin");
|
||||
std::ifstream man(dump + "/manifest.txt");
|
||||
@@ -58,6 +69,13 @@ int main(int argc, char** argv)
|
||||
std::istringstream s(line);
|
||||
std::string kind;
|
||||
s >> kind;
|
||||
if (kind == "texture" || kind == "surface")
|
||||
{
|
||||
// texture/surface object handles in the arguments would point at nothing here: the kernel faults
|
||||
fprintf(stderr, "%s: the dump uses texture or surface objects, which dump_runner cannot recreate\n",
|
||||
dump.c_str());
|
||||
return 3;
|
||||
}
|
||||
if (kind == "launch")
|
||||
s >> grid[0] >> grid[1] >> grid[2] >> block[0] >> block[1] >> block[2] >> shared;
|
||||
else if (kind == "alloc")
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
#!/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)
|
||||
@@ -0,0 +1,127 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compares native preset-K layer outputs with the numpy model (pwin_model.py) on sampled windows of a replay dump.
|
||||
|
||||
usage: pwin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--windows N]
|
||||
DUMP_DIR replay dump of one dltss_pwin_* launch (inputs)
|
||||
OUT_DIR dump_runner output of a native kernel on that dump (alloc-N.bin after the launch)
|
||||
For every sampled window (the four corners, border windows and random interior ones) the model output and the
|
||||
kernel's stored values are compared per f16 element: exact matches, and the error in f16 ulps at the model
|
||||
value's magnitude (values below 2^-10 count with the ulp of 2^-10, so near-zero sign flips do not dominate).
|
||||
The kernels round as NVIDIA's f16 wmma does except in the PWIN_F32ACC layers, and gfx12 WMMA groups its f32
|
||||
sums differently from gfx11, so neither build is expected to match the model exactly.
|
||||
"""
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import pwin_model as pm # noqa: E402
|
||||
|
||||
F16 = np.float16
|
||||
|
||||
# layer -> (kind, H, C, COUT/CL, extra)
|
||||
LAYERS = {
|
||||
'enc0': ('enc0', 2, 32, 64, None), 'enc1': ('enc', 2, 64, 64, True), 'enc2': ('enc', 2, 64, 96, False),
|
||||
'enc3': ('enc', 4, 96, 128, False), 'enc4': ('enc', 4, 128, 160, False), 'dec5': ('plain', 8, 160, 16, None),
|
||||
'dec4': ('dec', 4, 128, 160, None), 'dec3': ('dec', 4, 96, 128, None), 'dec2': ('dec', 2, 64, 96, None),
|
||||
'dec1': ('dec', 2, 64, 64, None), 'dec0': ('dec0', 2, 32, 64, None),
|
||||
}
|
||||
|
||||
|
||||
class Outputs:
|
||||
"""the dump's pointer map over the allocations a native run saved"""
|
||||
def __init__(self, dump, out_dir):
|
||||
self.d = dump
|
||||
self.allocs = {i: (base, np.fromfile(f'{out_dir}/alloc-{i}.bin', np.uint8)) for i, (base, _) in dump.allocs.items()}
|
||||
|
||||
def rows(self, off_param, index, width):
|
||||
a = self.d.u64(off_param)
|
||||
for base, buf in self.allocs.values():
|
||||
if base <= a < base + buf.size:
|
||||
o = a - base + 2 * index * width
|
||||
return buf[o:o + 2 * width].view(F16)
|
||||
raise ValueError(hex(a))
|
||||
|
||||
|
||||
def ulps(model, got):
|
||||
m = model.astype(np.float64)
|
||||
g = got.astype(np.float64)
|
||||
scale = np.maximum(np.abs(m), 2.0 ** -10)
|
||||
ulp = 2.0 ** (np.floor(np.log2(scale)) - 10)
|
||||
return np.abs(m - g) / ulp
|
||||
|
||||
|
||||
def windows(grid, n, seed=1):
|
||||
gx, gy = grid[0], grid[1]
|
||||
w = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
|
||||
r = random.Random(seed)
|
||||
while len(w) < min(n, gx * gy):
|
||||
w.add((r.randrange(gx), r.randrange(gy)))
|
||||
return sorted(w)
|
||||
|
||||
|
||||
def model_window(P, layer, bx, by):
|
||||
kind, H, C, X, extra = LAYERS[layer]
|
||||
if kind == 'enc0':
|
||||
return pm.pwin_enc0(P, bx, by)
|
||||
if kind == 'enc':
|
||||
return pm.pwin_encoder(P, bx, by, H, C, X, posattn=extra)
|
||||
if kind == 'plain':
|
||||
return pm.swin_core(P, pm.Layout(H, C, C), pm.load_window(P, bx, by, C)), None
|
||||
if kind == 'dec':
|
||||
return pm.pwin_decoder(P, bx, by, H, C, X), None
|
||||
return pm.pwin_dec0(P, bx, by)
|
||||
|
||||
|
||||
def stored(P, O, layer, bx, by, y, extra_out):
|
||||
"""(model values, kernel values) of everything the kernel stores for window (bx, by)"""
|
||||
kind, H, C, X, _ = LAYERS[layer]
|
||||
W, Hh = P.i32x2(0)
|
||||
sx, sy = P.i32x2(56)
|
||||
ms, gs = [], []
|
||||
full_off = 32 if kind in ('enc0', 'enc') else 24
|
||||
for t in range(64):
|
||||
Y0, X0 = 8 * by - sy + (t >> 3), 8 * bx - sx + (t & 7)
|
||||
if not (0 <= Y0 < Hh and 0 <= X0 < W):
|
||||
continue
|
||||
if kind == 'dec0':
|
||||
ms.append(extra_out[t]); gs.append(O.rows(24, Y0 * W + X0, 40))
|
||||
else:
|
||||
ms.append(y[t]); gs.append(O.rows(full_off, Y0 * W + X0, C))
|
||||
if kind in ('enc0', 'enc'):
|
||||
cout = extra_out.shape[1]
|
||||
for r in range(16):
|
||||
Ym, Xm = int((8 * by - sy) / 2) + (r >> 2), int((8 * bx - sx) / 2) + (r & 3)
|
||||
if 0 <= Ym < Hh // 2 and 0 <= Xm < W // 2:
|
||||
ms.append(extra_out[r]); gs.append(O.rows(24, Ym * (W // 2) + Xm, cout))
|
||||
return np.concatenate(ms), np.concatenate(gs)
|
||||
|
||||
|
||||
def main():
|
||||
args = [a for a in sys.argv[1:] if not a.startswith('--')]
|
||||
n = 12
|
||||
if '--windows' in sys.argv:
|
||||
n = int(sys.argv[sys.argv.index('--windows') + 1])
|
||||
args.remove(str(n))
|
||||
dump_dir, outs = args[0], args[1:]
|
||||
P = pm.Dump(dump_dir)
|
||||
layer = P.kernel.replace('dltss_pwin_', '').replace('_layer', '')
|
||||
wins = windows(P.grid, n)
|
||||
models = {w: model_window(P, layer, *w) for w in wins}
|
||||
for out in outs:
|
||||
O = Outputs(P, out)
|
||||
m, g = zip(*(stored(P, O, layer, bx, by, *models[(bx, by)]) for bx, by in wins))
|
||||
m, g = np.concatenate(m), np.concatenate(g)
|
||||
u = ulps(m, g)
|
||||
exact = (m.view(np.uint16) == g.view(np.uint16)).mean()
|
||||
md, gd = m.astype(np.float64), g.astype(np.float64)
|
||||
psnr = 10 * np.log10(np.abs(md).max() ** 2 / max(((md - gd) ** 2).mean(), 1e-30))
|
||||
print(f'{layer:5s} {os.path.basename(os.path.normpath(out)):10s} {len(wins)} windows {m.size} values: exact {100 * exact:5.1f}% '
|
||||
f'<=1ulp {100 * (u <= 1).mean():5.1f}% <=4ulp {100 * (u <= 4).mean():6.2f}% max {u.max():7.1f} ulp '
|
||||
f'mean {u.mean():.3f} ulp psnr {psnr:.1f} dB finite {np.isfinite(g.astype(np.float32)).mean() * 100:.0f}%')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,105 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compares native preset-M Swin-block outputs with the numpy model (swin_model.py) on sampled blocks of a dump.
|
||||
|
||||
usage: swin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--blocks N]
|
||||
DUMP_DIR replay dump of one rrlite_{enc1,enc2,enc3_tube,dec1,dec2}_4x4 launch
|
||||
OUT_DIR dump_runner output of a native kernel on that dump
|
||||
The stored e4m3 outputs (block output, and the patch merge of encoders that have one) are compared code by code:
|
||||
the share of identical codes, of codes one e4m3 step apart, and the PSNR of the decoded values. The model is
|
||||
NVIDIA's arithmetic (e4m3 activations, f16 accumulation per k32 step); the production kernels keep activations
|
||||
f16 and accumulate in f32 (SWIN_NO_Q8, SWIN_F32ACC), so they are not expected to match it exactly.
|
||||
"""
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import model_enc3 as me # noqa: E402
|
||||
import swin_model as sm # noqa: E402
|
||||
|
||||
# kernel -> (C, heads, NPM, CIN)
|
||||
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)}
|
||||
|
||||
|
||||
def blocks(grid, n, seed=1):
|
||||
gx, gy = grid[0], grid[1]
|
||||
b = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
|
||||
r = random.Random(seed)
|
||||
while len(b) < min(n, gx * gy):
|
||||
b.add((r.randrange(gx), r.randrange(gy)))
|
||||
return sorted(b)
|
||||
|
||||
|
||||
def model_block(P, layer, bx, by):
|
||||
C, heads, NPM, CIN = LAYERS[layer]
|
||||
pre = 4 * C * CIN + 8 * C if CIN else 0
|
||||
x0 = sm.decoder_input(P, bx, by, C, CIN, heads, P.p32) if CIN else sm.load_tokens(P, bx, by, C // 32)
|
||||
x2, out, L = sm.swin_block(P, x0, C, heads, pre)
|
||||
merged = None
|
||||
if NPM:
|
||||
A = sm.patch_merge_input(out, C)
|
||||
PM0 = L['end']
|
||||
cols = []
|
||||
for g in range(NPM // 32):
|
||||
Wg = me.E4M3[P.wbuf[sm.woff_table(PM0 + g * 128 * C, 512, 64 * C, 4 * C, 32)]]
|
||||
bias = P.f16vec(PM0 + 4 * C * NPM + 64 * g, 32)
|
||||
cols.append(me.mma_chain(A, Wg, C0=np.broadcast_to(bias, (16, 32)).copy()))
|
||||
merged = me.q8(sm.pair_to_nat(np.concatenate(cols, axis=1)))
|
||||
return out, merged
|
||||
|
||||
|
||||
def codes_at(buf, base, ptr, W, H, X, Y, C):
|
||||
"""e4m3 codes of tokens (X, Y) (in bounds) from planes of 32 channels at ptr"""
|
||||
c = np.arange(C)
|
||||
off = ptr - base + (((c[None, :] >> 5) * H + Y[:, None]) * W + X[:, None]) * 32 + (c[None, :] & 31)
|
||||
return buf[off]
|
||||
|
||||
|
||||
def main():
|
||||
args = [a for a in sys.argv[1:] if not a.startswith('--')]
|
||||
n = 8
|
||||
if '--blocks' in sys.argv:
|
||||
n = int(sys.argv[sys.argv.index('--blocks') + 1])
|
||||
args.remove(str(n))
|
||||
dump, outs = args[0], args[1:]
|
||||
P = me.Params(dump)
|
||||
layer = P.kernel.replace('rrlite_', '').replace('_4x4', '')
|
||||
C, heads, NPM, CIN = LAYERS[layer]
|
||||
bl = blocks(P.grid, n)
|
||||
models = {b: model_block(P, layer, *b) for b in bl}
|
||||
for out_dir in outs:
|
||||
got, want = [], []
|
||||
for (bx, by) in bl:
|
||||
base, _ = P.alloc_of(P.out)
|
||||
aid = next(i for i, (b0, buf) in P.allocs.items() if b0 == base)
|
||||
buf = np.fromfile(f'{out_dir}/alloc-{aid}.bin', np.uint8)
|
||||
X, Y = P.token_xy(bx, by)
|
||||
ok = (X >= 0) & (X < P.W) & (Y >= 0) & (Y < P.H)
|
||||
out, merged = models[(bx, by)]
|
||||
got.append(codes_at(buf, base, P.out, P.W, P.H, X[ok], Y[ok], C).ravel())
|
||||
want.append(out[ok].ravel())
|
||||
if NPM:
|
||||
W2, H2 = P.W // 2, P.H // 2
|
||||
r = np.arange(16)
|
||||
MX, MY = (8 * bx - P.sx) // 2 + (r & 3), (8 * by - P.sy) // 2 + (r >> 2)
|
||||
okm = (MX >= 0) & (MX < W2) & (MY >= 0) & (MY < H2)
|
||||
mbase, _ = P.alloc_of(P.p48)
|
||||
mid = next(i for i, (b0, _) in P.allocs.items() if b0 == mbase)
|
||||
mbuf = buf if mid == aid else np.fromfile(f'{out_dir}/alloc-{mid}.bin', np.uint8)
|
||||
got.append(codes_at(mbuf, mbase, P.p48, W2, H2, MX[okm], MY[okm], NPM).ravel())
|
||||
want.append(merged[okm].ravel())
|
||||
g = me.E4M3[np.concatenate(got)]
|
||||
w = np.concatenate(want).astype(np.float64)
|
||||
rank = lambda v: np.searchsorted(me._POS, np.abs(v)) * np.sign(v) # noqa: E731
|
||||
steps = np.abs(rank(g) - rank(w))
|
||||
psnr = 10 * np.log10(np.nanmax(np.abs(w)) ** 2 / max(np.nanmean((g - w) ** 2), 1e-30))
|
||||
print(f'{layer:9s} {os.path.basename(os.path.normpath(out_dir)):10s} {len(bl)} blocks {g.size} codes: '
|
||||
f'identical {100 * (g == w).mean():5.1f}% <=1 step {100 * (steps <= 1).mean():6.2f}% '
|
||||
f'max {int(steps.max())} steps psnr {psnr:.1f} dB')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,141 @@
|
||||
// Standalone check of the texture-kernel GEMM tails (kernels/tex/enc0_tail.hip, dec0_head.hip) outside the
|
||||
// ZLUDA-compiled kernels they are linked into: random e4m3 weights and activations, several tile positions
|
||||
// (including partial tiles at the image border), all outputs written to one file. Builds of different WMMA
|
||||
// layouts (gfx11, the gfx12-layout shim, gfx12, gfx12 native FP8) are then compared byte for byte (same
|
||||
// arithmetic) or by value (gfx12 rounding). The tails touch no images, so this also runs in rocjitsu.
|
||||
//
|
||||
// build: clang++ -x hip --offload-arch=ARCH -O3 [-DD4R_WMMA_LAYOUT=12] [-DD4R_TEX_FP8] -I kernels/tex
|
||||
// kernels/tools/tail_check.hip -o tail_check usage: tail_check OUT.bin
|
||||
#include <hip/hip_runtime.h>
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
|
||||
// the surface-store helpers of tex_common.h reference ZLUDA's helper library, which the tails do not use
|
||||
extern "C" __device__ void __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(uint64_t, int __attribute__((ext_vector_type(2))),
|
||||
short __attribute__((ext_vector_type(2))))
|
||||
{
|
||||
}
|
||||
|
||||
// both tails define the same hook globals
|
||||
#define D4R_KERNEL_NAME tail_enc0
|
||||
#define d4r_prep_blocks d4r_prep_blocks_enc0
|
||||
#define d4r_prep_key_at d4r_prep_key_at_enc0
|
||||
#include "enc0_tail.hip"
|
||||
#undef d4r_prep_blocks
|
||||
#undef d4r_prep_key_at
|
||||
#define d4r_prep_blocks d4r_prep_blocks_dec0
|
||||
#define d4r_prep_key_at d4r_prep_key_at_dec0
|
||||
#include "dec0_head.hip"
|
||||
|
||||
extern "C" __global__ void __launch_bounds__(32) run_enc0(const uint32_t* lds_in, const uint8_t* w, uint32_t* out,
|
||||
uint32_t dimx, uint32_t dimy, int ox, int oy)
|
||||
{
|
||||
__shared__ uint32_t s[512 + 4];
|
||||
for (int i = threadIdx.x; i < 512; i += 32)
|
||||
s[i] = lds_in[i];
|
||||
__syncthreads();
|
||||
d4r_enc0_tail((uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)s, (uint64_t)w, (uint64_t)out, dimx,
|
||||
dimy, ox, oy);
|
||||
}
|
||||
|
||||
extern "C" __global__ void __launch_bounds__(32) run_dec0(const uint8_t* w, const uint32_t* in, uint32_t* staged,
|
||||
uint32_t dimx, uint32_t dimy, int ox, int oy)
|
||||
{
|
||||
__shared__ uint32_t s[512 + 4];
|
||||
for (int i = threadIdx.x; i < 512; i += 32)
|
||||
s[i] = 0xdeadbeefu;
|
||||
__syncthreads();
|
||||
d4r_dec0_head((uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)s, (uint64_t)w, (uint64_t)in, dimx,
|
||||
dimy, ox, oy);
|
||||
__syncthreads();
|
||||
for (int i = threadIdx.x; i < 512; i += 32)
|
||||
staged[i] = s[i];
|
||||
}
|
||||
|
||||
#define CHECK(x) \
|
||||
do \
|
||||
{ \
|
||||
hipError_t e_ = (x); \
|
||||
if (e_ != hipSuccess) \
|
||||
{ \
|
||||
fprintf(stderr, "%s: %s\n", #x, hipGetErrorString(e_)); \
|
||||
return 1; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static uint8_t e4m3(uint32_t r)
|
||||
{
|
||||
// Keep the matrix values in a useful range so saturation cannot hide a wrong WMMA result.
|
||||
const uint8_t magnitude = (r & 31) == 0 ? 0 : (uint8_t)(0x18 + (r & 31));
|
||||
return (uint8_t)(magnitude | ((r >> 8) & 0x80));
|
||||
}
|
||||
static uint16_t f16_bits(float v)
|
||||
{
|
||||
return __builtin_bit_cast(uint16_t, (_Float16)v);
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
if (argc < 2)
|
||||
{
|
||||
fprintf(stderr, "usage: %s OUT.bin\n", argv[0]);
|
||||
return 2;
|
||||
}
|
||||
srand(12345);
|
||||
auto rnd = []() { return (uint32_t)rand() * 2654435761u + (uint32_t)rand(); };
|
||||
// weights: e4m3 fragments, f16 biases (enc0 at 12288, dec0 at 8192)
|
||||
std::vector<uint8_t> w(16384 + 512);
|
||||
for (auto& b : w)
|
||||
b = e4m3(rnd());
|
||||
for (int i = 0; i < 64; ++i)
|
||||
*(uint16_t*)&w[12288 + 2 * i] = f16_bits((float)(int)(rnd() % 2001 - 1000) / 256.0f);
|
||||
for (int i = 0; i < 128; ++i)
|
||||
*(uint16_t*)&w[8192 + 2 * i] = f16_bits((float)(int)(rnd() % 2001 - 1000) / 256.0f);
|
||||
const uint32_t dimx = 20, dimy = 14, W2 = dimx / 2, H2 = dimy / 2, rowstride = W2 * 8, plane = H2 * rowstride;
|
||||
std::vector<uint32_t> lds(512), in(2 * plane);
|
||||
for (auto& v : lds)
|
||||
v = (uint32_t)e4m3(rnd()) | (uint32_t)e4m3(rnd()) << 8 | (uint32_t)e4m3(rnd()) << 16 | (uint32_t)e4m3(rnd()) << 24;
|
||||
for (auto& v : in)
|
||||
v = (uint32_t)e4m3(rnd()) | (uint32_t)e4m3(rnd()) << 8 | (uint32_t)e4m3(rnd()) << 16 | (uint32_t)e4m3(rnd()) << 24;
|
||||
|
||||
uint8_t *dw, *dummy;
|
||||
uint32_t *dlds, *din, *dout, *dstaged;
|
||||
CHECK(hipMalloc(&dw, w.size()));
|
||||
CHECK(hipMalloc(&dlds, lds.size() * 4));
|
||||
CHECK(hipMalloc(&din, in.size() * 4));
|
||||
CHECK(hipMalloc(&dout, 2 * plane * 4));
|
||||
CHECK(hipMalloc(&dstaged, 512 * 4));
|
||||
CHECK(hipMalloc(&dummy, 512));
|
||||
CHECK(hipMemcpy(dw, w.data(), w.size(), hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(dlds, lds.data(), lds.size() * 4, hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(din, in.data(), in.size() * 4, hipMemcpyHostToDevice));
|
||||
// the prep kernels read the weight pointer at parameter offset 0
|
||||
Enc0Params pe{};
|
||||
pe.w = (uint64_t)dw;
|
||||
Dec0Params pd{};
|
||||
pd.w = (uint64_t)dw;
|
||||
hipLaunchKernelGGL(tail_enc0_prep, dim3(4), dim3(128), 0, 0, pe);
|
||||
hipLaunchKernelGGL(rrlite_dec0_4x4_folded_prep, dim3(4), dim3(128), 0, 0, pd);
|
||||
CHECK(hipDeviceSynchronize());
|
||||
|
||||
FILE* f = fopen(argv[1], "wb");
|
||||
const int offsets[][2] = {{0, 0}, {8, 4}, {12, 6}, {16, 8}, {-4, -2}, {18, 12}};
|
||||
std::vector<uint32_t> host(2 * plane);
|
||||
for (const auto& o : offsets)
|
||||
{
|
||||
CHECK(hipMemset(dout, 0x5a, 2 * plane * 4));
|
||||
hipLaunchKernelGGL(run_enc0, dim3(1), dim3(32), 0, 0, dlds, dw, dout, dimx, dimy, o[0], o[1]);
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(host.data(), dout, 2 * plane * 4, hipMemcpyDeviceToHost));
|
||||
fwrite(host.data(), 4, 2 * plane, f);
|
||||
hipLaunchKernelGGL(run_dec0, dim3(1), dim3(32), 0, 0, dw, din, dstaged, dimx, dimy, o[0], o[1]);
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(host.data(), dstaged, 512 * 4, hipMemcpyDeviceToHost));
|
||||
fwrite(host.data(), 4, 512, f);
|
||||
}
|
||||
fclose(f);
|
||||
printf("wrote %s\n", argv[1]);
|
||||
return 0;
|
||||
}
|
||||
Reference in new issue
Block a user