wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

engine/rr: native preset F (DLSS 4.5 Ray Reconstruction) and the reference's fused roundings

Browse Source 
- 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
This commit is contained in:
woutermans committed 2026-10-13 13:59:17 +02:00
1 parent 065fe2c2a1
commit 8eb7a94533
31 files changed
+3547 -128
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/rr/gen_rr2_enc0.py
+298
View File
@@ -0,0 +1,298 @@
#!/usr/bin/env python3
"""Generate preset F's encoder 0 shader: ptx2rr.py's translation of NVIDIA's per-pixel feature gather and of its guide
tail, around the native block (engine/rr/rr2_m.comp, KIND 4).
usage: gen_rr2_enc0.py ENC0.ptx ENC0.generated.comp OUT.comp
gen_rr2_enc0.py DEC0_1.ptx DEC0_1.generated.comp OUT.comp --dec0_1
NVIDIA's kernel runs one 32-lane warp per 8x8-token window and keeps everything in registers. Its matrix products
(the embedding, the attention block and the guide head) are one contiguous run of the translated main, which is cut
out:
wave w the translation up to the first product: the gather; the first product's A fragments (the window's 64
feature slots per token) go to the shared array XF
4 waves rr2_block(): embedding, attention, MLP, the tokens to the arena, the guide head to the shared array HD
wave w the last product's D fragments are loaded from HD and the translation continues after it: the guide tail
With --group N (2 to 4) a workgroup holds N windows along x: waves 0 ... N - 1 run the translated code of one window
each and the block then runs once per window; the stage's dispatch is 1 / N as wide ("xN" in rr.txt). Measured on
an RX 9070 XT (2026-10-13) that is slower than one window per workgroup (network +0.4 ms at 1280x720 with N = 4).
With --dec0_1 the same for dec0's second kernel (rr2_m.comp KIND 5): nothing before the block (it reads its tokens
from the arena); after it the last product's D fragments (the MLP sum) and the registers holding r are loaded from
MSH / RSH, and the translated tail forms y = r + MLP * s2 and computes the kernel's outputs from it.
The translated parts hold code derived from the user's DLL: the output is model data, do not commit it.
"""
import re
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from phase_dump import mma_statements, phases_of # noqa: E402
def tiles(mm):
"""per mma of a phase: (m-tile, n-tile, k-step), as phase_dump.py orders them"""
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
a_tile, b_tile = {}, {}
for s in range(max(step.values()) + 1):
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
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
return [(a_tile[(step[n], tuple(m[3]))], b_tile[(step[n], tuple(m[4]))], step[n]) for n, m in enumerate(mm)]
def f16vars_of(decls):
names = set()
for l in decls:
if re.match(r'\s*(precise )?float16_t ', l):
names |= set(re.findall(r'\b(rs\d+)\b', l))
return names
REG = r'\b(?:rs|rd|r|f|p|l)\d+(?:_low|_high)?\b'
# what makes a statement of the gather more than index arithmetic on the parameter block
DATA = re.compile(r'fetch0|texture|shfl_|GX\d|GA\d|GW\d|SH\d|sust_|movm_|mma_')
def tail_inputs(pre, tail):
"""The gather's values the tail reads: (recomputed, saved). Recomputed are the statements, in order, of values
that are index arithmetic on the parameter block (assigned once, outside any block); saved are the names of the
others, which the gather pass writes to the window's record."""
assigned, depth = {}, 0
for i, l in enumerate(pre):
m = re.match(r'\s*(\w+) = (.*)', l)
if m:
assigned.setdefault(m.group(1), []).append((i, depth, m.group(2)))
depth += l.count('{') - l.count('}')
used, own = [], set()
for l in tail:
m = re.match(r'\s*(\w+) = (.*)', l)
for v in re.findall(REG, m.group(2) if m else l):
if v not in own and v in assigned and v not in used:
used.append(v)
if m:
own.add(m.group(1))
def closure(v, seen):
"""statement indices v needs, or None when it is not plain index arithmetic"""
if v in seen:
return seen[v]
seen[v] = None
defs = assigned.get(v)
if defs is None:
seen[v] = set() # never assigned in the gather: a declaration's zero
return seen[v]
if len(defs) != 1 or defs[0][1] != 0 or DATA.search(defs[0][2]):
return None
need = {defs[0][0]}
for u in re.findall(REG, defs[0][2]):
c = closure(u, seen)
if c is None:
return None
need |= c
seen[v] = need
return need
lines, saved, seen = set(), [], {}
for v in used:
c = closure(v, seen)
if c is None:
saved.append(v)
else:
lines |= c
return [pre[i] for i in sorted(lines)], saved
def split_layout(xf_word, pre, tail, f16vars):
"""(recomputed statements, saved names, the record's words, the GLSL locating a window's record)"""
recomputed, saved = tail_inputs(pre, tail)
for v in saved:
if not (v in f16vars or re.fullmatch(r'r\d+', v)):
raise SystemExit(f'the tail reads {v} from the gather: not a 16- or 32-bit register')
record = 1024 + 32 * len(saved)
where = (f' const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n'
f' const uint rr2_rec = param.w[{xf_word}u] / 4u + {record}u * (gl_WorkGroupID.y * (param.w[0] / 8u + 1u) + gl_WorkGroupID.x);\n')
return recomputed, saved, record, where
def gather_pass(xf_word, head, decls, pre, tail, first, f16vars):
"""Encoder 0 as two passes (--split gather | block): the translated gather in NVIDIA's own 32-thread groups, which
writes each window's record to the arena, and the native block with the translated tail in 128-thread groups.
One shader holding both runs the gather on one of the block's four waves while the other three wait, which
costs more than the record's trip through memory (RX 9070 XT, 1280x720, 2026-10-13: gather 0.40 -> 0.24 ms).
The record: 1024 words of A fragments ([token slot][64 FP8 feature slots], as the block's XF), then per lane the
gather's values the tail reads. Its address: parameter word XF_WORD (bytes) + 4 * RECORD * (y * (W / 8 + 1) + x)."""
recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars)
var = lambda r: r.lstrip('%')
store = [' if (rr2_valid) { // the first product\'s A fragments: lane 4g + t, register q -> row, 4 k',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
seen = set()
for mm_, (i, j, s) in zip(first, tiles(first)):
if (i, s) in seen:
continue
seen.add((i, s))
for q, r in enumerate(mm_[3]):
store.append(f' rr2_arena.w[rr2_rec + rr2_uninter({16 * i + 8 * (q & 1)}u + g_) * 16u + {(32 * s + 16 * (q >> 1)) // 4}u + t_] = {var(r)};')
for n, v in enumerate(saved):
bits = f'uint(float16BitsToUint16({v}))' if v in f16vars else v
store.append(f' rr2_arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u] = {bits};')
store.append(' }')
print(f'gather pass: record of {record} words, the tail reads {", ".join(saved) or "nothing"} from it')
return (head + '\nlayout(set = 0, binding = 2, std430) buffer Rr2Arena { uint w[]; } rr2_arena;\n'
'uint rr2_uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }\n\n'
'void main()\n{\n' + '\n'.join(decls) + '\n' + where +
' const bool rr2_valid = gl_WorkGroupID.x < rr2_gx;\n'
' if (rr2_valid) { // NVIDIA\'s warp: the translation before the first product\n' + '\n'.join(pre) + '\n }\n' +
'\n'.join(store) + '\n}\n'), record
def main():
ptx, src, out = sys.argv[1:4]
dec = '--dec0_1' in sys.argv
# --group N: windows per workgroup (1 to 4; wave w runs the translated code of window w)
G = int(sys.argv[sys.argv.index('--group') + 1]) if '--group' in sys.argv else 1
kernel = 'rr2_dec0_kernel_1' if dec else 'rr2_enc0_kernel'
lines = Path(ptx).read_text().replace('\r', '').split('\n')
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
phases = phases_of(mma_statements(lines, k, len(lines)), lines)
first, last = phases[0], phases[-1]
var = lambda r: r.lstrip('%')
g = Path(src).read_text().split('\n')
m = next(i for i, l in enumerate(g) if l.startswith('void main()'))
head = '\n'.join(g[:m]).replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;',
'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;')
if dec: # the tail's tanh from rr2_pack.py's table (ptx2rr.py's RR_TANH_TABLE; its offset from OUT's rr2_stage.h)
v = head.index('\n') + 1
table = re.search(r'#define W_TANH (\d+)u', (Path(out).parent / 'rr2_stage.h').read_text()).group(1)
head = head[:v] + f'#define RR_TANH_TABLE {table}u\n' + head[v:]
body = g[m + 2:]
while body and body[-1].strip() in ('', '}'):
last_line = body.pop()
if last_line.strip() == '}':
break
nd = next(i for i, l in enumerate(body) if not re.match(r' (precise |bool |uint64_t |float |uint |int |float16_t |f16vec2 )', l))
calls = [i for i, l in enumerate(body) if 'mma_e4m3_f16' in l]
a = calls[0] - 1 # '{ uint md0_, ...;'
b = next(i for i in range(calls[-1], len(body)) if body[i].strip() == '}')
if not body[a].strip().startswith('{ uint md0_'):
raise SystemExit('unexpected statement before the first product: ' + body[a])
pre, cut, tail = body[nd:a], body[a:b + 1], body[b + 1:]
if sum(l.count('{') - l.count('}') for l in pre):
raise SystemExit('the first product is inside a block of the translated main')
# values the tail takes from the cut region other than the last product's results
assigned = lambda ls: {mm_.group(1) for l in ls for mm_ in [re.match(r'\s*(\w+) = ', l)] if mm_}
restored = {var(r) for mm_ in last for r in mm_[2]}
rmap = {} # dec0_1: the register holding r for (row offset, column offset): y = r + (D half) * s2 in the tail
if dec:
place = {}
for mm_, (i, j, s_) in zip(last, tiles(last)):
for q, r in enumerate(mm_[2]):
place[r] = (16 * i + 8 * q, 8 * j)
text_tail = '\n'.join(lines[last[-1][1]:])
half = {}
for mo in re.finditer(r'mov\.b32\s+\{(%rs\d+),\s*(%rs\d+)\},\s*(%r\d+);', text_tail):
if mo.group(3) in place:
half[mo.group(1)] = (*place[mo.group(3)], 0)
half[mo.group(2)] = (*place[mo.group(3)], 1)
prod = {}
for mo in re.finditer(r'mul\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in half:
prod[mo.group(1)] = half[mo.group(2)]
for mo in re.finditer(r'add\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in prod:
rmap[var(mo.group(3))] = prod[mo.group(2)]
if len(rmap) != 2 * 2 * len(last):
raise SystemExit(f'found {len(rmap)} residual registers, expected {4 * len(last)}')
restored |= set(rmap)
used = set(re.findall(r'\b[a-z]+\d+\b', '\n'.join(tail)))
lost = sorted((assigned(cut) - assigned(pre) - restored) & used)
# (values assigned again in the tail before their use are fine; the kernel has none on the reference build)
keep = []
cut_set = assigned(cut)
for name in lost: # thread indices and the like: their assignment is carried over when it depends on nothing cut
line = next(l for l in reversed(cut) if re.match(rf'\s*{name} = ', l))
if set(re.findall(r'\b[a-z]+\d+\b', line.split('=', 1)[1])) & cut_set:
raise SystemExit(f'the tail needs {name}, which the cut region computes: {line.strip()}')
keep.append(' ' + line)
split = sys.argv[sys.argv.index('--split') + 1] if '--split' in sys.argv else None
xf_word = int(sys.argv[sys.argv.index('--xf-word') + 1]) if split else 0
if split == 'gather':
text, record = gather_pass(xf_word, '\n'.join(g[:m]), body[:nd], pre, tail, first, f16vars_of(body[:nd]))
Path(out).write_text(text)
Path(out).with_suffix('.record').write_text(str(record))
return
if dec:
first = []
store = [' if (rr2_valid) { // the first product\'s A fragments: lane 4g + t, register q -> row, 4 k',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
seen = set()
for mm_, (i, j, s) in zip(first, tiles(first) if first else []):
if (i, s) in seen:
continue
seen.add((i, s))
for q, r in enumerate(mm_[3]):
store.append(f' XF[gl_SubgroupID * 1024u + uninter({16 * i + 8 * (q & 1)}u + g_) * 16u + {(32 * s + 16 * (q >> 1)) // 4}u + t_] = {var(r)};')
store.append(' }')
load = [' { // the guide head as the last product\'s D fragments (rows g, g + 8; columns 2t, 2t + 1)',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
f16vars = set()
for l in body[:nd]:
if re.match(r'\s*(precise )?float16_t ', l):
f16vars |= set(re.findall(r'\b(rs\d+)\b', l))
for mm_, (i, j, s) in zip(last, tiles(last)):
for q, r in enumerate(mm_[2]):
row = f'({16 * i + 8 * q}u + g_)'
load.append(f' {var(r)} = HDV({row}, {8 * j}u + 2u * t_) | (HDV({row}, {8 * j + 1}u + 2u * t_) << 16u);')
for name, (ro, co, hf) in rmap.items():
v = f'RSV({ro}u + g_, {co + hf}u + 2u * t_)'
load.append(f' {name} = ' + (f'uint16BitsToFloat16(uint16_t({v}));' if name in f16vars else f'{v};'))
load.append(' }')
# the native block's source, without its #version line (this shader has its own)
block = (Path(__file__).resolve().parent / 'rr2_m.comp').read_text().split('\n', 1)[1]
if split == 'block': # the second pass: the block on the gather pass's record, then the tail on wave 0
recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars)
restore = [' ' + (f'{v} = uint16BitsToFloat16(uint16_t(arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u]));' if v in f16vars
else f'{v} = arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u];') for n, v in enumerate(saved)]
load[2:2] = restore
Path(out).write_text(
head + '\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP 1\n#define RR2_XF_ARENA 1\n' +
block + '\nuint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n'
'void main()\n{\n' + '\n'.join(body[:nd]) + '\n' + where +
' const bool rr2_valid = gl_SubgroupID < 1u && gl_WorkGroupID.x < rr2_gx;\n'
' rr2_xf = rr2_rec;\n rr2_block();\n barrier();\n if (rr2_valid) {\n' +
' // the gather\'s index arithmetic the tail reads\n' + '\n'.join(recomputed) + '\n' +
'\n'.join(load + keep) + '\n' + '\n'.join(tail) + '\n }\n}\n')
return
wx = lambda t: t.replace('gl_WorkGroupID.x', f'({G}u * gl_WorkGroupID.x + gl_SubgroupID)')
text = (head + f'\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP {G}\n' + block + '\n' +
('uint HDV(uint row, uint col) { return uint(float16BitsToUint16(MSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n'
'uint RSV(uint row, uint col) { return uint(float16BitsToUint16(RSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n\n' if dec else
'uint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n') +
'void main()\n{\n' + '\n'.join(body[:nd]) +
'\n // this wave\'s window, when the level has it\n'
' const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n'
f' const bool rr2_valid = gl_SubgroupID < {G}u && {G}u * gl_WorkGroupID.x + gl_SubgroupID < rr2_gx;\n'
' if (rr2_valid) { // NVIDIA\'s warp: the translation before the first product\n' + wx('\n'.join(pre)) + '\n }\n' +
('' if dec else '\n'.join(store)) +
# (the block runs unconditionally, also for a window past the level's last one: it writes nothing there, and
# inside a branch the driver compiles it much slower)
('\n barrier();\n rr2_block();\n barrier();\n' if G == 1 else
f'\n barrier();\n for (uint w_ = 0u; w_ < {G}u; ++w_) {{\n rr2_win = w_;\n'
' rr2_block();\n barrier();\n }\n') + ' if (rr2_valid) {\n' +
'\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n }\n}\n')
Path(out).write_text(text)
print(f'{out}: {len(pre)} gather lines, {len(cut)} cut, {len(tail)} tail lines, {len(seen)} A fragment sets')
if __name__ == '__main__':
main()