ayo
22026ede12
Preset K, byte-identical to the previous output in all captured buffers and the harness image, measured on an RX 7700 XT: - LDS row padding so WMMA operand loads hit distinct banks - WGP mode for enc1/enc2/dec1/dec2, ds_bpermute half-wave exchange for enc0/dec0/dec3, more waves per window in enc4/dec4/dec5 - cheaper operand plumbing in common/wmma_layout.h (packed f16 conversions, no dynamic vector indexing), raw buffer weight loads, accumulator-init vectors prepared once - enc1 as a single-buffer recompute layer, dec1 Q/K/V chains split - dec5 as four dependent phase kernels
235 lines
9.0 KiB
Python
Executable File
235 lines
9.0 KiB
Python
Executable File
"""If-conversion of pure branch diamonds in PTX (used by ptx2hip.py --ifconv).
|
|
|
|
NVIDIA's DLSS kernels branch around small pure arithmetic alternatives (e.g. the inverse tonemap
|
|
`v < 0.998 ? 1 / (1 - v) : (exp2(...) - 1) / (9 v)`), laid out as
|
|
|
|
@%p bra $A; (or @!%p)
|
|
bra.uni $B;
|
|
$A: <pure ops> bra.uni $M;
|
|
$B: <pure ops> (falls through)
|
|
$M:
|
|
|
|
Branches like these split the code into basic blocks that no compiler schedules memory operations across, so
|
|
independent fetch chains on either side are serialized. This pass rewrites such a diamond into straight-line code:
|
|
both arms compute into fresh registers and every register an arm assigns is then selected by the predicate. Each lane
|
|
keeps exactly the values of the arm it would have taken (GPU arithmetic has no traps), so results are bit-identical.
|
|
Only arms made of register arithmetic are converted (no memory, texture, barrier, shuffle, branch or return; no
|
|
integer division, which is undefined behaviour in C for some inputs).
|
|
"""
|
|
import re
|
|
|
|
REG = re.compile(r'%([a-zA-Z]+)(\d+)\b')
|
|
PURE_BASES = {'add', 'sub', 'mul', 'fma', 'mad', 'div', 'rcp', 'sqrt', 'rsqrt', 'ex2', 'lg2', 'min', 'max', 'neg',
|
|
'abs', 'cvt', 'mov', 'selp', 'setp', 'and', 'or', 'xor', 'not', 'shl', 'shr', 'set', 'bfi', 'copysign',
|
|
'sin', 'cos', 'cvta'}
|
|
SEL = {'f': 'selp.f32', 'r': 'selp.b32', 'rd': 'selp.b64', 'rs': 'selp.b16'}
|
|
|
|
|
|
def _pure(it):
|
|
if it['kind'] != 'stmt' or it['pred'] is not None:
|
|
return False
|
|
op = it['op']
|
|
base = op.split('.')[0]
|
|
if base not in PURE_BASES:
|
|
return False
|
|
if base == 'div' and not op.endswith('.f32') and 'approx' not in op:
|
|
return False # integer division (and IEEE f32 division, which the translator rejects anyway)
|
|
if base == 'cvta':
|
|
return False
|
|
return True
|
|
|
|
|
|
def _target(it):
|
|
if it['kind'] == 'stmt' and it['op'].startswith('bra'):
|
|
return it['rest'].strip().rstrip(';').strip()
|
|
return None
|
|
|
|
|
|
def apply(body, tokenize, parse, split_operands):
|
|
items = parse(body, tokenize)
|
|
# label -> number of branches targeting it
|
|
refs = {}
|
|
for it in items:
|
|
t = _target(it)
|
|
if t:
|
|
refs[t] = refs.get(t, 0) + 1
|
|
# fresh register numbers per kind
|
|
top = {}
|
|
for m in REG.finditer(body):
|
|
k, n = m.group(1), int(m.group(2))
|
|
top[k] = max(top.get(k, 0), n)
|
|
fresh = {k: v + 100000 for k, v in top.items()}
|
|
|
|
def new_reg(kind):
|
|
fresh[kind] = fresh.get(kind, 100000) + 1
|
|
return f'%{kind}{fresh[kind]}'
|
|
|
|
use_sites = {}
|
|
for x in items:
|
|
if x['kind'] == 'stmt':
|
|
for r in x['src']:
|
|
use_sites.setdefault(r, []).append(id(x))
|
|
|
|
def uses_outside(r, arm_ids):
|
|
return any(u not in arm_ids for u in use_sites.get(r, []))
|
|
|
|
defined_before = set()
|
|
out, i, n_conv = [], 0, 0
|
|
while i < len(items):
|
|
it = items[i]
|
|
conv = None
|
|
if it['kind'] == 'stmt' and it['op'].startswith('bra') and it['pred'] is not None:
|
|
conv = _match(items, i, refs)
|
|
if conv is None:
|
|
if it['kind'] == 'stmt':
|
|
defined_before.update(it['dst'])
|
|
out.append(_text(it))
|
|
i += 1
|
|
continue
|
|
neg, pred, arm_a, arm_b, merge_idx = conv
|
|
lines, last = [], [{}, {}]
|
|
for side, arm in enumerate((arm_a, arm_b)):
|
|
ren = {}
|
|
for st in arm:
|
|
op, rest = st['text'].split(None, 1)
|
|
ops = split_operands(rest)
|
|
if len(st['dst']) != 1 or ops[0] != st['dst'][0]:
|
|
raise ValueError('unexpected destination form: ' + st['text'])
|
|
srcs = [REG.sub(lambda mm: ren.get(mm.group(0), mm.group(0)), o) for o in ops[1:]]
|
|
d = st['dst'][0]
|
|
kind = REG.match(d).group(1)
|
|
if kind not in SEL:
|
|
raise ValueError('predicate assigned in a diamond arm: ' + st['text'])
|
|
nr = new_reg(kind)
|
|
ren[d] = nr
|
|
lines.append(f"{op} {', '.join([nr] + srcs)}")
|
|
last[side] = ren
|
|
# only registers read outside the two arms need a select (arm temporaries are dropped)
|
|
arm_ids = {id(x) for x in arm_a} | {id(x) for x in arm_b}
|
|
assigned = sorted(r for r in set(last[0]) | set(last[1]) if uses_outside(r, arm_ids))
|
|
cond = pred
|
|
for r in assigned:
|
|
kind = REG.match(r).group(1)
|
|
a = last[0].get(r, r if r in defined_before else last[1].get(r))
|
|
b = last[1].get(r, r if r in defined_before else last[0].get(r))
|
|
# selp d, a, b, p: d = p ? a : b ; arm A is taken when (p xor neg)
|
|
if neg:
|
|
a, b = b, a
|
|
lines.append(f'{SEL[kind]} {r}, {a}, {b}, {cond}')
|
|
defined_before.update(assigned)
|
|
out.extend(l + ';' for l in lines)
|
|
n_conv += 1
|
|
i = merge_idx # the merge label itself is emitted by the main loop
|
|
return '\n'.join(out), n_conv
|
|
|
|
|
|
def _text(it):
|
|
if it['kind'] == 'open':
|
|
return '{'
|
|
if it['kind'] == 'close':
|
|
return '}'
|
|
if it['kind'] == 'label':
|
|
return it['text'] + ':'
|
|
return it['text'] + ';'
|
|
|
|
|
|
def _match(items, i, refs):
|
|
"""@p bra A; bra.uni B; A: ... bra.uni M; B: ... M: -> (neg, pred, armA, armB, index of M)"""
|
|
br = items[i]
|
|
m = re.match(r'@(!?)(%p\d+)\s', br['text'])
|
|
if not m:
|
|
return None
|
|
neg, pred = m.group(1) == '!', m.group(2)
|
|
a_lab = _target(br)
|
|
if i + 2 >= len(items):
|
|
return None
|
|
j = items[i + 1]
|
|
if not (j['kind'] == 'stmt' and j['op'].startswith('bra') and j['pred'] is None):
|
|
return None
|
|
b_lab = _target(j)
|
|
if refs.get(a_lab) != 1 or refs.get(b_lab) != 1:
|
|
return None
|
|
|
|
def arm(start, label):
|
|
if not (items[start]['kind'] == 'label' and items[start]['text'] == label):
|
|
return None
|
|
body, k = [], start + 1
|
|
while k < len(items):
|
|
x = items[k]
|
|
if x['kind'] == 'label':
|
|
return body, k, None
|
|
if x['kind'] in ('open', 'close'):
|
|
return None
|
|
if x['kind'] == 'stmt' and x['op'].startswith('bra'):
|
|
if x['pred'] is not None:
|
|
return None
|
|
return body, k + 1, _target(x)
|
|
if not _pure(x):
|
|
return None
|
|
body.append(x)
|
|
k += 1
|
|
return None
|
|
|
|
first = arm(i + 2, a_lab)
|
|
order = 'AB'
|
|
if first is None:
|
|
first = arm(i + 2, b_lab)
|
|
order = 'BA'
|
|
if first is None:
|
|
return None
|
|
body1, k, jump1 = first
|
|
if jump1 is None:
|
|
return None
|
|
second = arm(k, b_lab if order == 'AB' else a_lab)
|
|
if second is None:
|
|
return None
|
|
body2, k2, jump2 = second
|
|
if jump2 is not None and jump2 != jump1:
|
|
return None # the second arm may also jump to the merge label explicitly
|
|
if not (k2 < len(items) and items[k2]['kind'] == 'label' and items[k2]['text'] == jump1):
|
|
return None
|
|
arm_a, arm_b = (body1, body2) if order == 'AB' else (body2, body1)
|
|
return neg, pred, arm_a, arm_b, k2
|
|
|
|
|
|
def sink_stores(body, tokenize, parse):
|
|
"""Drop labels nothing branches to, then move each st.shared / st.local down to just before the next statement
|
|
that could observe or reorder it: a shared/local load or any other store, a barrier, shuffle, atomic, label,
|
|
branch or return, or a statement that writes one of the store's registers. Stores keep their relative order, so
|
|
every memory location sees the same sequence of writes; texture loads (read-only data) and arithmetic move ahead
|
|
of them, which lets independent fetch chains issue together. Returns (body, number of stores moved)."""
|
|
items = parse(body, tokenize)
|
|
refs = set()
|
|
for it in items:
|
|
t = _target(it)
|
|
if t:
|
|
refs.add(t)
|
|
items = [it for it in items if not (it['kind'] == 'label' and it['text'] not in refs)]
|
|
|
|
def is_fence(it):
|
|
if it['kind'] != 'stmt':
|
|
return True # labels and scope braces
|
|
op = it['op']
|
|
base = op.split('.')[0]
|
|
if base in ('bar', 'shfl', 'atom', 'red', 'bra', 'ret', 'membar', 'fence', 'st', 'sust', 'd4r'):
|
|
return True
|
|
if base == 'ld' and ('.shared' in op or '.local' in op):
|
|
return True
|
|
return False
|
|
|
|
moved = 0
|
|
i = len(items) - 1
|
|
while i >= 0:
|
|
it = items[i]
|
|
if it['kind'] == 'stmt' and it['pred'] is None and it['op'].startswith('st.') and ('.shared' in it['op'] or '.local' in it['op']):
|
|
regs = set(it['src'])
|
|
j = i + 1
|
|
while j < len(items) and not is_fence(items[j]) and not (regs & set(items[j]['dst'])):
|
|
j += 1
|
|
if j > i + 1:
|
|
# after the pop the fence that stopped the scan is at j - 1: the store goes right before it
|
|
items.insert(j - 1, items.pop(i))
|
|
moved += 1
|
|
i -= 1
|
|
return '\n'.join(_text(it) for it in items), moved
|