feat: add d4r inference engines and Windows parity tooling
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.
This commit is contained in:
ayo committed
1 parent
25493b3198
commit
7c5d3c9157
430 files changed
+76191
No files matched your search
@@ -0,0 +1,206 @@
|
||||
#!/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()
|
||||
Reference in new issue
Block a user