woutermans
8eb7a94533
- engine/compile_rr2.py, runtime_rr2.inc, engine/rr/rr2_m.comp, exposure_f.comp, gen_rr2_enc0.py, rr2_pack.py, rr2_ref.py: the rr2_* model (nvngx_dlssd.dll 310.9's third network) as d4r's own shaders - exposure, enc0 gather/block split, enc1-5 and dec4-0 through one matrix shader, translated output passes; weights packed from the DLL, argument blocks from the capture - ptx2rr: match the driver's fused roundings - opaque-zero OR around cvt.rn.f16.f32 against v_fma_mix, half-away-from-zero formatted stores, UNORM 10:10:10:2 as RGBA32F of q/1023, OCML tanhF with the AMDGPU exp lowering, binary16 fma and fetch roundings; --probe for register bisection against ZLUDA - dec0_m/swin_m/enc0_*: RR_MIX build option (D4R_RR_MIX=1), one v_fma_mix rounding per 32-deep step, ~5% faster, not bit-exact; default build stays exact - runtime_rr.inc: RModel/REngine dispatch to the D4RR0002 model; replay_rr: RGBA32F and A2B10G10R1010 texel sizes, roughnessPacked bit, larger debug staging - shim: engine\rr-f model slot, 310.9 version probe so preset F is only claimed when the CUDA path has it; d4r.ini / d4r_config.py Model = F; run_rr_engine_compare.sh D4R_RR_DLL - kernels/rr/swin_native.h + rr_swin_native_test.cpp: HIP port of encoder 1 for the FP8-product and rounding cost measurements - engine/rr/README: preset F's block, options, capture-by-capture results, timings and the open camera-matrix difference
207 lines
10 KiB
Python
207 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Dump every mma of a denoiser layer as matrices, from a ZLUDA replay of a captured launch: the oracle for porting
|
|
a layer to native shaders.
|
|
|
|
usage: phase_dump.py CAPTURE STAGE FRAME BX,BY OUT.npz [--phases A-B]
|
|
|
|
For one thread block, debug stores added to the PTX record every mma's A, B, C and D registers (each mma's last
|
|
execution) into padding behind the plane arena; the replay runs through ZLUDA (engine/rr/zluda_replay.py). The
|
|
fragments are then put back together with the m16n8k32 layouts (A row-major 16x32 e4m3, B column-major 32x8 e4m3,
|
|
C/D 16x8 f16) into one matrix per phase: phase P's mma are grouped into m-tiles by A fragment and into n-tiles by
|
|
B fragment, k-steps follow the C <- D chains. OUT.npz holds, per phase p and wave w: p{p}_w{w}_A (M x K, e4m3 codes),
|
|
_B (K x N codes), _C (M x N f16 of the first k-step), _D (M x N f16 of the last).
|
|
"""
|
|
import argparse
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
ROOT = HERE.parents[1]
|
|
sys.path.insert(0, str(HERE))
|
|
from rr_capture import Capture # noqa: E402
|
|
|
|
REG = r'%r\d+'
|
|
|
|
|
|
def kernel_module(kernel):
|
|
corpus = '~/.cache/d4r-rr-corpus-f' if kernel.startswith('rr2_') else '~/.cache/d4r-rr-corpus'
|
|
for p in sorted(Path(corpus).expanduser().glob('*.ptx')):
|
|
if f'.entry {kernel}(' in p.read_text(errors='replace'):
|
|
return p
|
|
raise SystemExit(f'no corpus module for {kernel}')
|
|
|
|
|
|
def mma_statements(lines, start, end):
|
|
"""(first line, last line, D, A, B, C) of every mma in [start, end)"""
|
|
out, i = [], start
|
|
while i < end:
|
|
if 'mma.sync.aligned.m16n8k32' in lines[i]:
|
|
j, text = i, ''
|
|
while True:
|
|
text += lines[j]
|
|
if lines[j].rstrip().endswith(';'):
|
|
break
|
|
j += 1
|
|
groups = [re.findall(REG, g) for g in re.findall(r'\{([^}]*)\}', text)]
|
|
out.append((i + 1, j + 1, *groups[:4]))
|
|
i = j + 1
|
|
else:
|
|
i += 1
|
|
return out
|
|
|
|
|
|
def phases_of(mmas, lines):
|
|
"""maximal runs of mma separated only by blank or comment lines"""
|
|
phases, cur = [], []
|
|
for m in mmas:
|
|
if cur:
|
|
gap = [l.strip() for l in lines[cur[-1][1]:m[0] - 1]]
|
|
if any(g and not g.startswith('//') for g in gap):
|
|
phases.append(cur)
|
|
cur = []
|
|
cur.append(m)
|
|
if cur:
|
|
phases.append(cur)
|
|
return phases
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('capture')
|
|
ap.add_argument('stage')
|
|
ap.add_argument('frame', type=int)
|
|
ap.add_argument('block')
|
|
ap.add_argument('out')
|
|
ap.add_argument('--phases', help='first-last (1-based), default all')
|
|
o = ap.parse_args()
|
|
# the capture decides the kernel: a preset D capture runs cuda_dldn_engine_swin_{enc0,dec0}_diamond_wallaby_kernel
|
|
cap = Capture(os.path.expanduser(o.capture))
|
|
# preset F: STAGE is the kernel's full name (rr2_enc1_kernel, rr2_dec1_kernel_0)
|
|
rr2 = o.stage.startswith('rr2_')
|
|
launch = [l for l in cap.launches if (l.kernel == o.stage if rr2 else f'swin_{o.stage}' in l.kernel)][o.frame - 1]
|
|
kernel = launch.kernel
|
|
module = kernel_module(kernel)
|
|
lines = module.read_text().replace('\r', '').split('\n')
|
|
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
|
|
e = next((i for i in range(k + 1, len(lines)) if '.entry ' in lines[i]), len(lines))
|
|
phases = phases_of(mma_statements(lines, k, e), lines)
|
|
first, last = (int(x) for x in o.phases.split('-')) if o.phases else (1, len(phases))
|
|
T = launch.block[0] * launch.block[1] * launch.block[2]
|
|
arena_bytes = len(launch.pre(launch.arena_alloc))
|
|
# any argument pointing into the arena locates its base (enc0 has none at 8)
|
|
aptr = next(p for p in launch.pointers if p.alloc == launch.arena_alloc)
|
|
delta8, off8 = aptr.delta, aptr.offset
|
|
bx, by = (int(v) for v in o.block.split(','))
|
|
# one probe slot per register per mma: after the mma, all of A, B, C, D (C read after the mma is still C
|
|
# unless C and D share registers, which the k-step chains never do for the first step)
|
|
slots, inserts = [], {}
|
|
pname = f'{kernel}_param_0'
|
|
for p in range(first, last + 1):
|
|
for (l0, l1, D, A, B, C) in phases[p - 1]:
|
|
code = []
|
|
for role, regs in (('D', D), ('A', A), ('B', B), ('C', C)):
|
|
for r in regs:
|
|
s = len(slots)
|
|
slots.append((p, l0, role, r))
|
|
code.append(f'mov.b32 %dv, {r}; mad.wide.u32 %da, %dt, 4, %db; '
|
|
f'add.s64 %da, %da, {4 * T * s}; @%dq2 st.global.u32 [%da], %dv;')
|
|
inserts[l1] = ('{ .reg .b64 %da, %db; .reg .b32 %dv, %dt, %dx, %dy; .reg .pred %dq, %dq2;\n'
|
|
f'mov.u32 %dx, %ctaid.x; mov.u32 %dy, %ctaid.y; setp.eq.u32 %dq, %dx, {bx}; '
|
|
f'setp.eq.and.u32 %dq2, %dy, {by}, %dq;\n'
|
|
f'mov.u32 %dt, %tid.z; mov.u32 %dx, %tid.y; mov.u32 %dy, %tid.x; '
|
|
f'mad.lo.u32 %dt, %dt, {launch.block[1]}, %dx; mad.lo.u32 %dt, %dt, {launch.block[0]}, %dy;\n'
|
|
f'ld.param.u64 %db, [{pname}+{off8}]; cvta.to.global.u64 %db, %db; '
|
|
f'add.s64 %db, %db, {arena_bytes - delta8};\n' + '\n'.join(code) + '\n}')
|
|
out = list(lines)
|
|
for l1 in sorted(inserts, reverse=True):
|
|
out[l1:l1] = [inserts[l1]]
|
|
work = ROOT / 'build/rr' / f'phase-{o.stage}'
|
|
work.mkdir(parents=True, exist_ok=True)
|
|
(work / 'probe.ptx').write_text('\n'.join(out))
|
|
pad = 4 * T * len(slots)
|
|
env = dict(os.environ, D4R_ZLUDA_WMMA='1', D4R_ZLUDA_WMMA_FP8='1', D4R_ZLUDA_WMMA_FP8_NATIVE='1',
|
|
D4R_ZLUDA_IMPLICIT_MAX_BLOCK='256')
|
|
cmd = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0']
|
|
cmd += [f'--setenv={k_}={v_}' for k_, v_ in env.items() if k_.startswith('D4R_ZLUDA')]
|
|
r = subprocess.run(cmd + [sys.executable, str(HERE / 'zluda_replay.py'), os.path.expanduser(o.capture),
|
|
kernel if rr2 else f'swin_{o.stage}' + ('_diamond' if 'diamond' in kernel else ''), str(o.frame), str(work),
|
|
'--ptx', str(work / 'probe.ptx'),
|
|
'--pad', str(pad)], capture_output=True, text=True)
|
|
if 'failed' in r.stdout + r.stderr:
|
|
sys.exit(r.stdout[-2000:] + r.stderr[-3000:])
|
|
raw = np.frombuffer((work / 'arena.out.bin').read_bytes()[arena_bytes:], np.uint32).reshape(len(slots), T)
|
|
val = {(p, l0, role, r): raw[s] for s, (p, l0, role, r) in enumerate(slots)}
|
|
result = {}
|
|
waves = T // 32
|
|
for p in range(first, last + 1):
|
|
mm = phases[p - 1]
|
|
# k-step chains: C of a later mma = D of an earlier one
|
|
dset = {tuple(m[2]): n for n, m in enumerate(mm)}
|
|
prev = {n: dset.get(tuple(m[5])) for n, m in enumerate(mm)}
|
|
step = {}
|
|
for n in range(len(mm)):
|
|
s_, q = 0, n
|
|
while prev[q] is not None:
|
|
q = prev[q]
|
|
s_ += 1
|
|
step[n] = s_
|
|
ksteps = max(step.values()) + 1
|
|
a_keys = list(dict.fromkeys(tuple(m[3]) for m in mm))
|
|
b_keys = list(dict.fromkeys(tuple(m[4]) for m in mm))
|
|
# m-tile of an A fragment: its index among the A fragments of the same k-step
|
|
a_tile = {}
|
|
for s_ in range(ksteps):
|
|
for t_, key in enumerate(dict.fromkeys(tuple(m[3]) for n, m in enumerate(mm) if step[n] == s_)):
|
|
a_tile[(s_, key)] = t_
|
|
b_tile = {}
|
|
for s_ in range(ksteps):
|
|
for t_, key in enumerate(dict.fromkeys(tuple(m[4]) for n, m in enumerate(mm) if step[n] == s_)):
|
|
b_tile[(s_, key)] = t_
|
|
mt = max(t for (s_, _), t in a_tile.items()) + 1
|
|
nt = max(t for (s_, _), t in b_tile.items()) + 1
|
|
for w in range(waves):
|
|
lanes = np.arange(32) + 32 * w
|
|
A = np.zeros((16 * mt, 32 * ksteps), np.uint8)
|
|
B = np.zeros((32 * ksteps, 8 * nt), np.uint8)
|
|
C = np.zeros((16 * mt, 8 * nt), np.uint16)
|
|
D = np.zeros((16 * mt, 8 * nt), np.uint16)
|
|
for n, (l0, l1, Dr, Ar, Br, Cr) in enumerate(mm):
|
|
s_ = step[n]
|
|
i, j = a_tile[(s_, tuple(Ar))], b_tile[(s_, tuple(Br))]
|
|
for lane in range(32):
|
|
g, t = lane >> 2, lane & 3
|
|
av = [int(val[(p, l0, 'A', r)][lanes[lane]]) for r in Ar]
|
|
bv = [int(val[(p, l0, 'B', r)][lanes[lane]]) for r in Br]
|
|
for q in range(4): # A regs: (g, k 4t..), (g+8, 4t..), (g, 16+4t..), (g+8, 16+4t..)
|
|
row = 16 * i + g + 8 * (q & 1)
|
|
col = 32 * s_ + 4 * t + 16 * (q >> 1)
|
|
for byte in range(4):
|
|
A[row, col + byte] = (av[q] >> (8 * byte)) & 0xff
|
|
for q in range(2): # B regs: column g, k 4t.. and 16+4t..
|
|
for byte in range(4):
|
|
B[32 * s_ + 4 * t + 16 * q + byte, 8 * j + g] = (bv[q] >> (8 * byte)) & 0xff
|
|
for q in range(2): # C/D regs: row g (+8), columns 2t, 2t+1
|
|
row = 16 * i + g + 8 * q
|
|
if s_ == 0:
|
|
cv = int(val[(p, l0, 'C', Cr[q])][lanes[lane]])
|
|
C[row, 8 * j + 2 * t] = cv & 0xffff
|
|
C[row, 8 * j + 2 * t + 1] = cv >> 16
|
|
if s_ == ksteps - 1:
|
|
dv = int(val[(p, l0, 'D', Dr[q])][lanes[lane]])
|
|
D[row, 8 * j + 2 * t] = dv & 0xffff
|
|
D[row, 8 * j + 2 * t + 1] = dv >> 16
|
|
for name, mat in (('A', A), ('B', B), ('C', C), ('D', D)):
|
|
result[f'p{p}_w{w}_{name}'] = mat
|
|
print(f'phase {p}: {len(mm)} mma, {mt} m-tiles x {nt} n-tiles x {ksteps} k-steps -> '
|
|
f'{16 * mt}x{8 * nt}x{32 * ksteps}, lines {mm[0][0]}-{mm[-1][1]}')
|
|
np.savez_compressed(o.out, **result)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|