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
858 lines
41 KiB
Python
Executable File
858 lines
41 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Translate one NVIDIA PTX kernel (DLSS texture kernels) into native HIP C++.
|
|
|
|
usage: ptx2hip.py MODULE.ptx KERNEL > KERNEL.hip
|
|
|
|
The output is a plain HIP kernel with the same name and parameter block. Every PTX instruction becomes
|
|
the C expression of ZLUDA's lowering of it (see ptx_rt.h), so the translation computes what ZLUDA's own
|
|
compile of the PTX computes, but it is compiled natively by clang and can then be optimised as C++.
|
|
Registers become locals, basic blocks become labels, branches become gotos.
|
|
"""
|
|
import re
|
|
import sys
|
|
|
|
TYPES = {'f': 'float', 'r': 'uint32_t', 'rd': 'uint64_t', 'rs': 'uint16_t', 'p': 'bool', 'fd': 'double',
|
|
'SP': 'uint64_t', 'SPL': 'uint64_t'}
|
|
SREG = {'%tid.x': '__builtin_amdgcn_workitem_id_x()', '%tid.y': '__builtin_amdgcn_workitem_id_y()',
|
|
'%tid.z': '__builtin_amdgcn_workitem_id_z()', '%ctaid.x': '__builtin_amdgcn_workgroup_id_x()',
|
|
'%ctaid.y': '__builtin_amdgcn_workgroup_id_y()', '%ctaid.z': '__builtin_amdgcn_workgroup_id_z()',
|
|
'%laneid': 'laneid()',
|
|
'%ntid.x': '__builtin_amdgcn_workgroup_size_x()', '%ntid.y': '__builtin_amdgcn_workgroup_size_y()',
|
|
'%nctaid.x': '(__builtin_amdgcn_grid_size_x() / __builtin_amdgcn_workgroup_size_x())',
|
|
'%nctaid.y': '(__builtin_amdgcn_grid_size_y() / __builtin_amdgcn_workgroup_size_y())'}
|
|
REG_RE = re.compile(r'%([a-zA-Z]+)(\d+)$')
|
|
|
|
|
|
class Err(Exception):
|
|
pass
|
|
|
|
|
|
def strip_comments(text):
|
|
return re.sub(r'//[^\n]*', '', text)
|
|
|
|
|
|
def tokenize(body):
|
|
"""-> list of ('open'|'close'|'label'|'stmt', text)"""
|
|
out, i, n = [], 0, len(body)
|
|
while i < n:
|
|
c = body[i]
|
|
if c.isspace():
|
|
i += 1
|
|
continue
|
|
if c == '{':
|
|
out.append(('open', None)); i += 1; continue
|
|
if c == '}':
|
|
out.append(('close', None)); i += 1; continue
|
|
m = re.match(r'(\$?[A-Za-z_][\w$]*):', body[i:])
|
|
if m and not body[i:].startswith('%'):
|
|
out.append(('label', m.group(1))); i += m.end(); continue
|
|
j = body.index(';', i)
|
|
out.append(('stmt', ' '.join(body[i:j].split())))
|
|
i = j + 1
|
|
return out
|
|
|
|
|
|
def split_operands(s):
|
|
ops, depth, cur = [], 0, ''
|
|
for ch in s:
|
|
if ch in '{[':
|
|
depth += 1
|
|
elif ch in '}]':
|
|
depth -= 1
|
|
if ch == ',' and depth == 0:
|
|
ops.append(cur.strip()); cur = ''
|
|
else:
|
|
cur += ch
|
|
if cur.strip():
|
|
ops.append(cur.strip())
|
|
return ops
|
|
|
|
|
|
class Translator:
|
|
def __init__(self, ptx, kernel):
|
|
self.kernel = kernel
|
|
text = strip_comments(ptx)
|
|
m = re.search(r'\.entry\s+' + re.escape(kernel) + r'\s*\((.*?)\)', text, re.S)
|
|
if not m:
|
|
raise Err('kernel not found')
|
|
params = m.group(1)
|
|
pm = re.search(r'\.param\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', params)
|
|
self.param_name, self.param_size, self.param_align = pm.group(2), int(pm.group(3)), int(pm.group(1))
|
|
rest = text[m.end():]
|
|
self.attrs = rest[:rest.index('{')]
|
|
start = rest.index('{')
|
|
depth, k = 0, start
|
|
while True:
|
|
if rest[k] == '{':
|
|
depth += 1
|
|
elif rest[k] == '}':
|
|
depth -= 1
|
|
if depth == 0:
|
|
break
|
|
k += 1
|
|
self.body = rest[start + 1:k]
|
|
self.globals_text = text[:m.start()]
|
|
self.regs = {} # c name -> type
|
|
self.shared = [] # (name, size, align)
|
|
self.locals = []
|
|
self.scope_names = [{}]
|
|
self.scope_id = 0
|
|
self.lines = []
|
|
|
|
# ------------------------------------------------------------ operands
|
|
def reg(self, tok):
|
|
tok = tok.strip()
|
|
for scope in reversed(self.scope_names):
|
|
if tok in scope:
|
|
return scope[tok]
|
|
m = REG_RE.match(tok)
|
|
if not m:
|
|
raise Err('not a register: ' + tok)
|
|
kind, num = m.groups()
|
|
name = f'{kind}{num}'
|
|
if kind not in TYPES:
|
|
raise Err('register kind ' + kind)
|
|
self.regs[name] = TYPES[kind]
|
|
return name
|
|
|
|
def rtype(self, name):
|
|
return self.regs[name]
|
|
|
|
def val(self, tok, as_type):
|
|
"""operand as a C expression of type as_type ('f32','u32','s32','u64','s64','u16','s16','pred','f16x2')"""
|
|
tok = tok.strip()
|
|
if tok in SREG:
|
|
return f'(uint32_t){SREG[tok]}'
|
|
if tok.startswith('%') or any(tok in s for s in self.scope_names):
|
|
name = self.reg(tok)
|
|
t = self.rtype(name)
|
|
if as_type == 'f32':
|
|
return name if t == 'float' else f'asf({name})'
|
|
if as_type in ('u32', 'b32'):
|
|
return f'asu({name})' if t == 'float' else name
|
|
if as_type == 's32':
|
|
return f'(int32_t)asu({name})' if t == 'float' else f'(int32_t){name}'
|
|
if as_type in ('u64', 'b64'):
|
|
return name
|
|
if as_type == 's64':
|
|
return f'(int64_t){name}'
|
|
if as_type in ('u16', 'b16'):
|
|
return name
|
|
if as_type == 's16':
|
|
return f'(int16_t){name}'
|
|
if as_type == 'f16':
|
|
return f'ash({name})'
|
|
if as_type == 'f16x2':
|
|
return f'ash2({name})'
|
|
if as_type == 'pred':
|
|
return name
|
|
raise Err('type ' + as_type)
|
|
if tok in self.shared_names():
|
|
return f'(uint32_t)(uintptr_t)(AS3 uint8_t*){tok}'
|
|
if tok == '__local_depot0':
|
|
# private (address space 5) address of the depot, zero-extended: ld/st.local use it as an AS5 pointer
|
|
return '(uint64_t)(uint32_t)(uintptr_t)(AS5 uint8_t*)__local_depot0'
|
|
# immediates
|
|
if re.match(r'0[fF][0-9A-Fa-f]{8}$', tok):
|
|
bits = int(tok[2:], 16)
|
|
return f'asf(0x{bits:08X}u)' if as_type == 'f32' else f'0x{bits:08X}u'
|
|
if re.match(r'-?\d+$', tok) or re.match(r'0x[0-9a-fA-F]+$', tok):
|
|
v = int(tok, 0)
|
|
if as_type == 'f32':
|
|
raise Err('int immediate as float')
|
|
if as_type in ('u64', 's64', 'b64'):
|
|
return f'(uint64_t){v}ull' if v >= 0 else f'(uint64_t)({v}ll)'
|
|
if as_type in ('u16', 'b16', 's16'):
|
|
return f'(uint16_t){v & 0xffff}'
|
|
if as_type == 's32':
|
|
v = v & 0xffffffff
|
|
v = v - (1 << 32) if v >= (1 << 31) else v
|
|
return f'(int32_t){v}' if v != -(1 << 31) else '(int32_t)0x80000000u'
|
|
return f'{v & 0xffffffff}u'
|
|
raise Err('operand ' + tok)
|
|
|
|
def shared_names(self):
|
|
return [s[0] for s in self.shared]
|
|
|
|
def assign(self, dst, expr, src_type):
|
|
"""dst = expr, expr of C type src_type ('float','u32','u64','u16','bool','f16x2','f16')"""
|
|
name = self.reg(dst)
|
|
t = self.rtype(name)
|
|
if src_type == 'float':
|
|
e = expr if t == 'float' else f'asu({expr})'
|
|
elif src_type == 'u32':
|
|
e = f'asf({expr})' if t == 'float' else f'(uint32_t)({expr})'
|
|
elif src_type == 'u64':
|
|
e = f'(uint64_t)({expr})'
|
|
elif src_type == 'u16':
|
|
e = f'(uint16_t)({expr})'
|
|
elif src_type == 'f16':
|
|
e = f'ashu({expr})'
|
|
elif src_type == 'f16x2':
|
|
e = f'ash2u({expr})'
|
|
elif src_type == 'bool':
|
|
e = f'({expr})'
|
|
else:
|
|
raise Err('src type')
|
|
self.emit(f'{name} = {e};')
|
|
|
|
def emit(self, s):
|
|
self.lines.append(' ' + s)
|
|
|
|
def addr(self, tok):
|
|
"""[reg+off] / [name+off] -> (base_expr, offset_int, space_hint)"""
|
|
m = re.match(r'\[\s*([^\]+]+?)\s*(?:\+\s*(-?\d+))?\s*\]$', tok)
|
|
if not m:
|
|
raise Err('address ' + tok)
|
|
base, off = m.group(1), int(m.group(2) or 0)
|
|
if base == self.param_name:
|
|
return 'KARG', off
|
|
if base in self.shared_names():
|
|
return f'(uint32_t)(uintptr_t)(AS3 uint8_t*){base}', off
|
|
return self.val(base, 'u64' if not base.startswith('%r') or base.startswith('%rd') else 'u32'), off
|
|
|
|
# ------------------------------------------------------------ instructions
|
|
def translate(self):
|
|
for kind, text in tokenize(self.body):
|
|
if kind == 'open':
|
|
self.scope_names.append({})
|
|
continue
|
|
if kind == 'close':
|
|
self.scope_names.pop()
|
|
continue
|
|
if kind == 'label':
|
|
self.lines.append(f'{text.replace("$", "")}:;')
|
|
lab = text.replace('$', '')
|
|
if lab in getattr(self, 'stamps', []):
|
|
self.lines.append(f' D4R_STAMP({self.stamps.index(lab) + 1});')
|
|
if getattr(self, 'bbcount', False):
|
|
self.bblabels.append(lab)
|
|
self.lines.append(f' D4R_BBC({len(self.bblabels) - 1});')
|
|
continue
|
|
self.stmt(text)
|
|
|
|
def stmt(self, s):
|
|
if s.startswith('.reg'):
|
|
m = re.match(r'\.reg\s+\.(\w+)\s+(.*)$', s)
|
|
ty, names = m.groups()
|
|
if '<' in names:
|
|
return # numbered register banks: declared on use
|
|
self.scope_id += 1
|
|
for nm in names.split(','):
|
|
nm = nm.strip()
|
|
cname = f'l{self.scope_id}_{nm.lstrip("%")}'
|
|
ctype = {'f16': 'uint16_t', 'b16': 'uint16_t', 'f32': 'float', 'b32': 'uint32_t', 'pred': 'bool',
|
|
'f16x2': 'uint32_t', 'u32': 'uint32_t', 's32': 'uint32_t', 'b64': 'uint64_t', 'u64': 'uint64_t'}[ty]
|
|
self.scope_names[-1][nm] = cname
|
|
self.regs[cname] = ctype
|
|
return
|
|
if s.startswith('.shared'):
|
|
m = re.match(r'\.shared\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s)
|
|
self.shared.append((m.group(2), int(m.group(3)), int(m.group(1))))
|
|
return
|
|
if s.startswith('.local'):
|
|
m = re.match(r'\.local\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s)
|
|
self.locals.append((m.group(2), int(m.group(3)), int(m.group(1))))
|
|
return
|
|
if s.startswith('.pragma'):
|
|
return
|
|
pred = None
|
|
m = re.match(r'@(!?)(%p\d+)\s+(.*)$', s)
|
|
if m:
|
|
pred = ('!' if m.group(1) else '') + self.reg(m.group(2))
|
|
s = m.group(3)
|
|
parts = s.split(None, 1)
|
|
op = parts[0]
|
|
args = split_operands(parts[1]) if len(parts) > 1 else []
|
|
if pred is not None:
|
|
if op.startswith('bra'):
|
|
self.emit(f'if ({pred}) goto {args[0].replace("$", "")};')
|
|
return
|
|
self.emit(f'if ({pred}) {{')
|
|
self.op(op, args)
|
|
self.emit('}')
|
|
return
|
|
self.op(op, args)
|
|
|
|
def op(self, op, a):
|
|
o = op.split('.')
|
|
base = o[0]
|
|
mods = set(o[1:])
|
|
E, A = self.emit, self.assign
|
|
V = self.val
|
|
|
|
def ftype():
|
|
for t in ('f32', 'f16x2', 'f16', 'f64'):
|
|
if t in mods:
|
|
return t
|
|
return None
|
|
|
|
if base == 'bra':
|
|
E(f'goto {a[0].replace("$", "")};'); return
|
|
if base == 'ret':
|
|
if getattr(self, 'stamps', None):
|
|
E(f'D4R_STAMP({len(self.stamps) + 1});')
|
|
E('return;'); return
|
|
if base == 'bar':
|
|
E('bar_sync();'); return
|
|
if base == 'mov':
|
|
dst, src = a
|
|
if dst.startswith('{'):
|
|
names = [x.strip() for x in dst[1:-1].split(',')]
|
|
srcv = V(src, 'u32')
|
|
A(names[0], f'{srcv} & 0xffffu', 'u16')
|
|
A(names[1], f'{srcv} >> 16', 'u16')
|
|
return
|
|
if src.startswith('{'):
|
|
names = [x.strip() for x in src[1:-1].split(',')]
|
|
if 'b64' in mods:
|
|
A(dst, f'(uint64_t){V(names[0], "u32")} | ((uint64_t){V(names[1], "u32")} << 32)', 'u64'); return
|
|
A(dst, f'pack16({V(names[0], "u16")}, {V(names[1], "u16")})', 'u32'); return
|
|
if src == self.param_name:
|
|
A(dst, 'KARGU', 'u64'); return
|
|
if 'f32' in mods:
|
|
A(dst, V(src, 'f32'), 'float'); return
|
|
if mods & {'u64', 'b64', 's64'}:
|
|
A(dst, V(src, 'u64'), 'u64'); return
|
|
if mods & {'u16', 'b16', 's16'}:
|
|
A(dst, V(src, 'u16'), 'u16'); return
|
|
A(dst, V(src, 'u32'), 'u32'); return
|
|
|
|
ft = ftype()
|
|
# ---------------------------------------------------- float arithmetic
|
|
if ft == 'f32' and base in ('add', 'sub', 'mul'):
|
|
if 'rz' in mods:
|
|
raise Err('add.rz outside the roundf idiom')
|
|
sym = {'add': '+', 'sub': '-', 'mul': '*'}[base]
|
|
A(a[0], f'{V(a[1], "f32")} {sym} {V(a[2], "f32")}', 'float'); return
|
|
if ft == 'f32' and base == 'fma':
|
|
A(a[0], f'__builtin_fmaf({V(a[1], "f32")}, {V(a[2], "f32")}, {V(a[3], "f32")})', 'float'); return
|
|
if ft == 'f32' and base in ('max', 'min'):
|
|
if 'NaN' in mods:
|
|
raise Err('NaN min/max')
|
|
A(a[0], f'{"fmaxn" if base == "max" else "fminn"}({V(a[1], "f32")}, {V(a[2], "f32")})', 'float'); return
|
|
if ft == 'f32' and base == 'neg':
|
|
A(a[0], f'neg_ftz({V(a[1], "f32")})' if 'ftz' in mods else f'-{V(a[1], "f32")}', 'float'); return
|
|
if ft == 'f32' and base == 'abs':
|
|
A(a[0], f'abs_ftz({V(a[1], "f32")})' if 'ftz' in mods else f'__builtin_fabsf({V(a[1], "f32")})', 'float'); return
|
|
if ft == 'f32' and base in ('rcp', 'ex2', 'lg2', 'rsqrt', 'sqrt'):
|
|
if 'approx' not in mods:
|
|
raise Err(op)
|
|
fn = {'rcp': 'rcp_a', 'ex2': 'ex2_a', 'lg2': 'lg2_a', 'rsqrt': 'rsq_a', 'sqrt': '__builtin_amdgcn_sqrtf'}[base]
|
|
A(a[0], f'{fn}({V(a[1], "f32")})', 'float'); return
|
|
if ft == 'f32' and base == 'div':
|
|
if 'approx' not in mods:
|
|
raise Err(op)
|
|
A(a[0], f'div_a({V(a[1], "f32")}, {V(a[2], "f32")})', 'float'); return
|
|
# ---------------------------------------------------- f16 / f16x2
|
|
if ft == 'f16x2' and base in ('add', 'sub', 'mul', 'fma'):
|
|
x = [V(t, 'f16x2') for t in a[1:]]
|
|
e = f'h2fma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})'
|
|
if 'sat' in mods:
|
|
e = f'h2sat({e})'
|
|
A(a[0], e, 'f16x2'); return
|
|
if ft == 'f16x2' and base == 'abs':
|
|
A(a[0], f'h2abs({V(a[1], "f16x2")})', 'f16x2'); return
|
|
if ft == 'f16x2' and base == 'ex2':
|
|
A(a[0], f'h2ex2({V(a[1], "f16x2")})', 'f16x2'); return
|
|
if ft in ('f16', 'f16x2') and base in ('min', 'max'):
|
|
if 'NaN' in mods:
|
|
raise Err('NaN-propagating half min/max is not supported')
|
|
fn = 'hminn' if base == 'min' else 'hmaxn'
|
|
if ft == 'f16':
|
|
A(a[0], f'{fn}({V(a[1], "f16")}, {V(a[2], "f16")})', 'f16'); return
|
|
x, y = V(a[1], 'f16x2'), V(a[2], 'f16x2')
|
|
A(a[0], f'(h2){{{fn}({x}.x, {y}.x), {fn}({x}.y, {y}.y)}}', 'f16x2'); return
|
|
if ft == 'f16' and base in ('add', 'sub', 'mul', 'fma'):
|
|
x = [V(t, 'f16') for t in a[1:]]
|
|
e = f'hfma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'(h1)({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})'
|
|
if 'sat' in mods:
|
|
e = f'hsat({e})'
|
|
A(a[0], e, 'f16'); return
|
|
if op == 'tanh.approx.f16':
|
|
A(a[0], f'__ocml_tanh_f16({V(a[1], "f16")})', 'f16'); return
|
|
if base == 'set' and 'f16x2' in mods and 'u32' in mods:
|
|
cmp = self.cmp_op(mods)
|
|
x, y = V(a[1], 'f16x2'), V(a[2], 'f16x2')
|
|
A(a[0], f'pack16(({cmp.format(f"{x}.x", f"{y}.x")}) ? 0xffffu : 0u, ({cmp.format(f"{x}.y", f"{y}.y")}) ? 0xffffu : 0u)', 'u32'); return
|
|
# ---------------------------------------------------- compares / selects
|
|
if base == 'setp':
|
|
cmp = self.cmp_op(mods)
|
|
if mods & {'f32'}:
|
|
x, y = V(a[1], 'f32'), V(a[2], 'f32')
|
|
elif 'f16' in mods:
|
|
x, y = V(a[1], 'f16'), V(a[2], 'f16')
|
|
elif mods & {'s32'}:
|
|
x, y = V(a[1], 's32'), V(a[2], 's32')
|
|
elif mods & {'u32', 'b32'}:
|
|
x, y = V(a[1], 'u32'), V(a[2], 'u32')
|
|
elif mods & {'s16'}:
|
|
x, y = V(a[1], 's16'), V(a[2], 's16')
|
|
elif mods & {'u16', 'b16'}:
|
|
x, y = V(a[1], 'u16'), V(a[2], 'u16')
|
|
elif mods & {'s64'}:
|
|
x, y = V(a[1], 's64'), V(a[2], 's64')
|
|
elif mods & {'u64'}:
|
|
x, y = V(a[1], 'u64'), V(a[2], 'u64')
|
|
else:
|
|
raise Err(op)
|
|
if len(a) != 3 or '|' in a[0]:
|
|
raise Err('setp form ' + op)
|
|
A(a[0], cmp.format(x, y), 'bool'); return
|
|
if base == 'selp':
|
|
t = 'f32' if 'f32' in mods else 'u16' if mods & {'u16', 'b16', 's16'} else 'u64' if mods & {'u64', 'b64', 's64'} else 'u32'
|
|
ct = {'f32': 'float', 'u16': 'u16', 'u64': 'u64', 'u32': 'u32'}[t]
|
|
A(a[0], f'{V(a[3], "pred")} ? {V(a[1], t)} : {V(a[2], t)}', ct); return
|
|
# ---------------------------------------------------- predicates / integer
|
|
if 'pred' in mods and base in ('and', 'or', 'xor', 'not'):
|
|
if base == 'not':
|
|
A(a[0], f'!{V(a[1], "pred")}', 'bool'); return
|
|
sym = {'and': '&&', 'or': '||', 'xor': '!='}[base]
|
|
A(a[0], f'{V(a[1], "pred")} {sym} {V(a[2], "pred")}', 'bool'); return
|
|
if base in ('and', 'or', 'xor', 'not', 'shl', 'shr', 'add', 'sub', 'mul', 'mad', 'min', 'max', 'div', 'rem', 'neg', 'abs'):
|
|
if ft is not None:
|
|
raise Err('unhandled floating operation: ' + op)
|
|
width = 64 if mods & {'u64', 's64', 'b64'} else 16 if mods & {'u16', 's16', 'b16'} else 32
|
|
signed = bool(mods & {'s32', 's64', 's16'})
|
|
ut = {64: 'u64', 32: 'u32', 16: 'u16'}[width]
|
|
st = {64: 's64', 32: 's32', 16: 's16'}[width]
|
|
ct = {64: 'u64', 32: 'u32', 16: 'u16'}[width]
|
|
if base == 'mul' and 'hi' in mods and width == 32:
|
|
wide = 'int64_t' if signed else 'uint64_t'
|
|
typ = 's32' if signed else 'u32'
|
|
A(a[0], f'(({wide}){V(a[1], typ)} * ({wide}){V(a[2], typ)}) >> 32', 'u32'); return
|
|
if base == 'mul' and 'wide' in mods:
|
|
if mods & {'s32'}:
|
|
A(a[0], f'(int64_t){V(a[1], "s32")} * (int64_t){V(a[2], "s32")}', 'u64'); return
|
|
if mods & {'u32'}:
|
|
A(a[0], f'(uint64_t){V(a[1], "u32")} * (uint64_t){V(a[2], "u32")}', 'u64'); return
|
|
if mods & {'u16'}:
|
|
A(a[0], f'(uint32_t){V(a[1], "u16")} * (uint32_t){V(a[2], "u16")}', 'u32'); return
|
|
raise Err(op)
|
|
if base == 'not':
|
|
A(a[0], f'~{V(a[1], ut)}', ct); return
|
|
if base in ('neg', 'abs'):
|
|
if base == 'abs' and signed:
|
|
A(a[0], f'{V(a[1], st)} < 0 ? (uint{width}_t)(0 - {V(a[1], ut)}) : {V(a[1], ut)}', ct); return
|
|
if base == 'neg':
|
|
A(a[0], f'(uint{width}_t)(0 - {V(a[1], ut)})', ct); return
|
|
raise Err(op)
|
|
x = V(a[1], st if signed else ut)
|
|
if base == 'shl':
|
|
A(a[0], f'{V(a[1], ut)} << ({V(a[2], "u32")} & {width - 1})', ct); return
|
|
if base == 'shr':
|
|
if signed:
|
|
A(a[0], f'{V(a[1], st)} >> ({V(a[2], "u32")} & {width - 1})', ct); return
|
|
A(a[0], f'{V(a[1], ut)} >> ({V(a[2], "u32")} & {width - 1})', ct); return
|
|
y = V(a[2], st if signed else ut)
|
|
if base in ('and', 'or', 'xor'):
|
|
sym = {'and': '&', 'or': '|', 'xor': '^'}[base]
|
|
A(a[0], f'{V(a[1], ut)} {sym} {V(a[2], ut)}', ct); return
|
|
if base in ('add', 'sub'):
|
|
A(a[0], f'{V(a[1], ut)} {"+" if base == "add" else "-"} {V(a[2], ut)}', ct); return
|
|
if base == 'mul':
|
|
if 'lo' not in mods:
|
|
raise Err(op)
|
|
A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)}', ct); return
|
|
if base == 'mad':
|
|
if 'lo' not in mods:
|
|
raise Err(op)
|
|
A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)} + {V(a[3], ut)}', ct); return
|
|
if base in ('min', 'max'):
|
|
fn = 'min' if base == 'min' else 'max'
|
|
A(a[0], f'__builtin_elementwise_{fn}({x}, {y})', ct); return
|
|
if base == 'div':
|
|
A(a[0], f'{x} / {y}', ct); return
|
|
if base == 'rem':
|
|
A(a[0], f'{x} % {y}', ct); return
|
|
if base == 'bfi':
|
|
A(a[0], f'bfi32({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")}, {V(a[4], "u32")})', 'u32'); return
|
|
if op == 'prmt.b32':
|
|
A(a[0], f'prmt32({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")})', 'u32'); return
|
|
# ---------------------------------------------------- conversions
|
|
if base == 'cvt':
|
|
return self.cvt(op, o[1:], a)
|
|
if base == 'cvta':
|
|
A(a[0], V(a[1], 'u64'), 'u64'); return
|
|
# ---------------------------------------------------- memory
|
|
if base == 'ld':
|
|
return self.ld(o[1:], a)
|
|
if base == 'st':
|
|
return self.st(o[1:], a)
|
|
if base == 'tex':
|
|
dst = [x.strip() for x in a[0][1:-1].split(',')]
|
|
m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', a[1])
|
|
h = V(m.group(1), 'u64')
|
|
c = [x.strip() for x in m.group(2).split(',')]
|
|
if 'base' in mods and mods >= {'2d', 'v4', 's32'} and mods & {'f32', 'f16'}:
|
|
E(f'{{ ::f4 t_ = tex_fetch({h}, (int32_t){V(c[0], "u32")}, (int32_t){V(c[1], "u32")});')
|
|
elif mods >= {'level', '2d', 'v4', 'f32'} and op.endswith('f32.f32'):
|
|
E(f'{{ ::f4 t_ = tex_level({h}, {V(c[0], "f32")}, {V(c[1], "f32")}, {V(a[2], "f32")});')
|
|
else:
|
|
raise Err(op)
|
|
for i, d in enumerate(dst):
|
|
A(d, f'(h1)t_[{i}]', 'f16') if 'f16' in mods else A(d, f't_[{i}]', 'float')
|
|
E('}')
|
|
return
|
|
if base == 'tld4':
|
|
if not (mods >= {'r', '2d', 'v4', 'f32'}):
|
|
raise Err(op)
|
|
dst = [x.strip() for x in a[0][1:-1].split(',')]
|
|
m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', a[1])
|
|
c = [x.strip() for x in m.group(2).split(',')]
|
|
E(f'{{ ::f4 t_ = tex_gather_r({V(m.group(1), "u64")}, {V(c[0], "f32")}, {V(c[1], "f32")});')
|
|
for i, d in enumerate(dst):
|
|
A(d, f't_[{i}]', 'float')
|
|
E('}')
|
|
return
|
|
if base == 'sust':
|
|
m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', a[0])
|
|
h = V(m.group(1), 'u64')
|
|
c = [V(x.strip(), 'u32') for x in m.group(2).split(',')]
|
|
v = [V(x.strip(), 'u32') for x in a[1][1:-1].split(',')]
|
|
if op == 'sust.p.2d.v4.b32.zero':
|
|
E(f'd4r_sust_p_v4b32({h}, (int32_t){c[0]}, (int32_t){c[1]}, {v[0]}, {v[1]}, {v[2]}, {v[3]});'); return
|
|
if op == 'sust.b.2d.b32.zero':
|
|
E(f'd4r_sust_b32({h}, (int32_t){c[0]}, (int32_t){c[1]}, {v[0]});'); return
|
|
if op == 'sust.b.2d.v2.b16.zero':
|
|
v = [V(x.strip(), 'u16') for x in a[1][1:-1].split(',')]
|
|
E(f'd4r_sust_v2b16({h}, (int32_t){c[0]}, (int32_t){c[1]}, {v[0]}, {v[1]});'); return
|
|
raise Err(op)
|
|
if base == 'shfl':
|
|
if op == 'shfl.sync.idx.b32':
|
|
d = a[0].split('|')
|
|
pred = f'&{self.reg(d[1])}' if len(d) > 1 else 'nullptr'
|
|
A(d[0], f'shfl_idx({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")}, {pred})', 'u32'); return
|
|
if op != 'shfl.sync.bfly.b32':
|
|
raise Err(op)
|
|
d = a[0].split('|')
|
|
p = f'&{self.reg(d[1])}' if len(d) > 1 else 'nullptr'
|
|
A(d[0], f'shfl_bfly({V(a[1], "u32")}, (int32_t){V(a[2], "u32")}, {V(a[3], "u32")}, {p})', 'u32'); return
|
|
if op == 'vote.sync.ballot.b32':
|
|
A(a[0], f'(uint32_t)(__builtin_amdgcn_ballot_w64({V(a[1], "pred")}) >> (wave_lane() & 32u))', 'u32'); return
|
|
raise Err('unhandled ' + op)
|
|
|
|
def cmp_op(self, mods):
|
|
table = {'lt': '{} < {}', 'gt': '{} > {}', 'le': '{} <= {}', 'ge': '{} >= {}', 'eq': '{} == {}',
|
|
'ne': '{} != {}',
|
|
# unordered: true if either is NaN
|
|
'ltu': '!({} >= {})', 'gtu': '!({} <= {})', 'leu': '!({} > {})', 'geu': '!({} < {})',
|
|
'equ': '!({} != {})', 'neu': '!({} == {})'}
|
|
for k in table:
|
|
if k in mods:
|
|
if k in ('ne', 'equ') and mods & {'f32', 'f16', 'f16x2', 'f64'}:
|
|
# PTX setp.ne.f32 is ordered (false on NaN) and equ is unordered (true on NaN), while C's
|
|
# != is unordered: spell both out. Integer ne is plain.
|
|
return '({0} < {1} || {0} > {1})' if k == 'ne' else '!({0} < {1} || {0} > {1})'
|
|
return table[k]
|
|
raise Err('compare ' + str(mods))
|
|
|
|
def cvt(self, op, o, a):
|
|
A, V = self.assign, self.val
|
|
mods = set(o)
|
|
if op == 'cvt.rn.satfinite.e4m3x2.f16x2':
|
|
A(a[0], f'e4m3x2({V(a[1], "u32")})', 'u16'); return
|
|
if op == 'cvt.rn.f16x2.e4m3x2':
|
|
A(a[0], f'd4r_decode_e4m3x2({V(a[1], "u16")})', 'u32'); return
|
|
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)]
|
|
dt, stp = types[0], types[1]
|
|
if dt == 'f32' and stp == 'f32':
|
|
x = V(a[1], 'f32')
|
|
if 'sat' in mods and 'ftz' in mods and not (mods & {'rzi', 'rmi', 'rni', 'rpi'}):
|
|
A(a[0], f'sat_ftz({x})', 'float'); return
|
|
if mods & {'sat'}:
|
|
raise Err(op)
|
|
if 'rzi' in mods:
|
|
A(a[0], f'__builtin_truncf({x})', 'float'); return
|
|
if 'rmi' in mods:
|
|
A(a[0], f'__builtin_floorf({x})', 'float'); return
|
|
if 'rni' in mods:
|
|
A(a[0], f'__builtin_rintf({x})', 'float'); return
|
|
if 'rpi' in mods:
|
|
A(a[0], f'__builtin_ceilf({x})', 'float'); return
|
|
if 'ftz' in mods:
|
|
A(a[0], f'canon({x})', 'float'); return
|
|
A(a[0], x, 'float'); return
|
|
if dt == 's32' and stp == 'f32':
|
|
x = V(a[1], 'f32')
|
|
if 'rzi' in mods:
|
|
A(a[0], f'(uint32_t)f2i_rz({x})', 'u32'); return
|
|
if 'rmi' in mods:
|
|
A(a[0], f'(uint32_t)f2i_rz(__builtin_floorf({x}))', 'u32'); return
|
|
if 'rni' in mods:
|
|
A(a[0], f'(uint32_t)f2i_rz(__builtin_rintf({x}))', 'u32'); return
|
|
raise Err(op)
|
|
if dt == 'u32' and stp == 'f32' and 'rzi' in mods:
|
|
A(a[0], f'f2u_rz({V(a[1], "f32")})', 'u32'); return
|
|
if dt == 'f32' and stp == 's32':
|
|
A(a[0], f'(float)(int32_t){V(a[1], "u32")}', 'float'); return
|
|
if dt == 'f32' and stp == 'u32':
|
|
A(a[0], f'(float){V(a[1], "u32")}', 'float'); return
|
|
if dt == 'f32' and stp == 'f16':
|
|
A(a[0], f'(float){V(a[1], "f16")}', 'float'); return
|
|
if dt == 'f16' and stp == 'f32':
|
|
if 'rn' not in mods:
|
|
raise Err(op)
|
|
A(a[0], f'(h1){V(a[1], "f32")}', 'f16'); return
|
|
if dt == 'f16x2' and stp == 'f32':
|
|
# d = {cvt(a) in the high half, cvt(b) in the low half}
|
|
A(a[0], f'pack16(ashu((h1){V(a[2], "f32")}), ashu((h1){V(a[1], "f32")}))', 'u32'); return
|
|
if dt == 'f16' and stp == 'f16':
|
|
x = V(a[1], 'f16')
|
|
if 'rni' in mods:
|
|
A(a[0], f'h_roundeven({x})', 'f16'); return
|
|
if 'rmi' in mods:
|
|
A(a[0], f'h_floor({x})', 'f16'); return
|
|
raise Err(op)
|
|
if dt == 's32' and stp == 'f16':
|
|
x = V(a[1], 'f16')
|
|
if 'rni' in mods:
|
|
A(a[0], f'(uint32_t)h2i_sat(h_roundeven({x}))', 'u32'); return
|
|
if 'rmi' in mods:
|
|
A(a[0], f'(uint32_t)h2i_sat(h_floor({x}))', 'u32'); return
|
|
raise Err(op)
|
|
if dt == 'f16' and stp == 's32':
|
|
A(a[0], f'(h1)(int32_t){V(a[1], "u32")}', 'f16'); return
|
|
if dt == 's64' and stp == 's32':
|
|
A(a[0], f'(uint64_t)(int64_t)(int32_t){V(a[1], "u32")}', 'u64'); return
|
|
if dt == 'u64' and stp == 'u32':
|
|
A(a[0], f'(uint64_t){V(a[1], "u32")}', 'u64'); return
|
|
if dt in ('u16', 's16') and stp in ('u32', 's32'):
|
|
A(a[0], f'(uint16_t){V(a[1], "u32")}', 'u16'); return
|
|
if dt in ('u32', 's32') and stp in ('u16',):
|
|
A(a[0], f'(uint32_t){V(a[1], "u16")}', 'u32'); return
|
|
if dt in ('u32', 's32') and stp in ('s16',):
|
|
A(a[0], f'(uint32_t)(int32_t)(int16_t){V(a[1], "u16")}', 'u32'); return
|
|
if dt == 'u32' and stp == 'u64':
|
|
A(a[0], f'(uint32_t){V(a[1], "u64")}', 'u32'); return
|
|
raise Err('cvt ' + op)
|
|
|
|
def mem_type(self, o):
|
|
for t in ('u8', 's8', 'u16', 's16', 'b16', 'u32', 's32', 'b32', 'f32', 'u64', 's64', 'b64', 'f16'):
|
|
if t in o:
|
|
return t
|
|
raise Err('mem type ' + str(o))
|
|
|
|
def ctype_of(self, t):
|
|
return {'u8': 'uint8_t', 's8': 'int8_t', 'u16': 'uint16_t', 's16': 'int16_t', 'b16': 'uint16_t',
|
|
'u32': 'uint32_t', 's32': 'int32_t', 'b32': 'uint32_t', 'f32': 'float', 'u64': 'uint64_t',
|
|
's64': 'int64_t', 'b64': 'uint64_t'}[t]
|
|
|
|
def ptr(self, space, base, off, ct):
|
|
if base == 'KARG':
|
|
return f'*(const AS4 {ct}*)(KARGP + {off})'
|
|
if space == 'param':
|
|
return f'*(const AS4 {ct}*)(uintptr_t)({base} + {off})'
|
|
if space == 'shared':
|
|
return f'*(AS3 {ct}*)(uintptr_t)({base} + {off}u)'
|
|
if space == 'local':
|
|
return f'*(AS5 {ct}*)(uintptr_t)(uint32_t)({base} + {off})'
|
|
if space == 'global' and getattr(self, 'as1', False):
|
|
# --as1: ld/st.global through address-space-1 pointers (global_load / global_store instead of flat)
|
|
return f'*(AS1 {ct}*)(uintptr_t)({base} + {off})'
|
|
if space in ('global', 'local'):
|
|
return f'*({ct}*)(uintptr_t)({base} + {off})'
|
|
raise Err('space ' + space)
|
|
|
|
def ld(self, o, a):
|
|
space = next(s for s in ('param', 'global', 'shared', 'local') if s in o)
|
|
t = self.mem_type(o)
|
|
ct = self.ctype_of(t)
|
|
base, off = self.addr(a[1])
|
|
dsts = [x.strip() for x in a[0][1:-1].split(',')] if a[0].startswith('{') else [a[0]]
|
|
width = {'uint8_t': 1, 'int8_t': 1, 'uint16_t': 2, 'int16_t': 2, 'uint32_t': 4, 'int32_t': 4, 'float': 4,
|
|
'uint64_t': 8, 'int64_t': 8}[ct]
|
|
if len(dsts) in (2, 4) and width == 4 and space in ('global', 'shared') and base != 'KARG':
|
|
# ld.v2/v4 (naturally aligned in PTX) as one vector load instead of per-element loads
|
|
vt = 'u4' if len(dsts) == 4 else 'u2v'
|
|
asq = 'AS3 ' if space == 'shared' else ('AS1 ' if getattr(self, 'as1', False) else '')
|
|
ptr = f'(uintptr_t)({base} + {off}u)' if space == 'shared' else f'(uintptr_t)({base} + {off})'
|
|
self.emit(f'{{ const {vt} v_ = *({asq}const {vt}*){ptr};')
|
|
for i, d in enumerate(dsts):
|
|
self.assign(d, f'v_[{i}]', 'u32')
|
|
self.emit('}')
|
|
return
|
|
for i, d in enumerate(dsts):
|
|
e = self.ptr(space, base, off + i * width, ct)
|
|
if ct == 'float':
|
|
self.assign(d, e, 'float')
|
|
elif width == 8:
|
|
self.assign(d, e, 'u64')
|
|
else:
|
|
name = self.reg(d)
|
|
if self.rtype(name) == 'uint16_t':
|
|
self.assign(d, e, 'u16')
|
|
else:
|
|
self.assign(d, f'(uint32_t){e}' if ct != 'int8_t' and ct != 'int16_t' else f'(uint32_t)(int32_t){e}', 'u32')
|
|
|
|
def st(self, o, a):
|
|
space = next(s for s in ('param', 'global', 'shared', 'local') if s in o)
|
|
t = self.mem_type(o)
|
|
ct = self.ctype_of(t)
|
|
base, off = self.addr(a[0])
|
|
srcs = [x.strip() for x in a[1][1:-1].split(',')] if a[1].startswith('{') else [a[1]]
|
|
width = {'uint8_t': 1, 'uint16_t': 2, 'uint32_t': 4, 'int32_t': 4, 'float': 4, 'uint64_t': 8}[ct]
|
|
for i, s in enumerate(srcs):
|
|
v = self.val(s, 'f32' if ct == 'float' else 'u64' if width == 8 else 'u16' if width == 2 else 'u32')
|
|
self.emit(f'{self.ptr(space, base, off + i * width, ct)} = ({ct})({v});')
|
|
|
|
# ------------------------------------------------------------ roundf idiom
|
|
def rewrite_rz(self):
|
|
"""add.rz.ftz.f32 t, x, h; cvt.rzi.f32.f32 y, t (h = copysign(0.5, x)) -> y = roundf(x)"""
|
|
text = self.body
|
|
pat = re.compile(r'add\.rz\.ftz\.f32\s+(%f\d+),\s*(%f\d+),\s*(%f\d+);\s*(?://[^\n]*\s*)*cvt\.rzi\.f32\.f32\s+(%f\d+),\s*(%f\d+);')
|
|
count = 0
|
|
|
|
def repl(m):
|
|
nonlocal count
|
|
t, x, h, y, t2 = m.groups()
|
|
if t != t2:
|
|
raise Err('rz idiom mismatch')
|
|
uses = len(re.findall(re.escape(t) + r'\b', text))
|
|
if uses != 2:
|
|
raise Err(f'rz temp {t} used {uses} times')
|
|
# check h = copysign(0.5, x): mov.b32 rA, x; and.b32 rB, rA, -2147483648; or.b32 rC, rB, 1056964608; mov.b32 h, rC
|
|
count += 1
|
|
return f'd4r.roundf {y}, {x};'
|
|
self.body = pat.sub(repl, text)
|
|
return count
|
|
|
|
|
|
def main():
|
|
args = [a for a in sys.argv[1:] if not a.startswith('--')]
|
|
flags = {a for a in sys.argv[1:] if a.startswith('--')}
|
|
if len(args) != 2:
|
|
sys.exit(__doc__)
|
|
ptx, kernel = open(args[0]).read(), args[1]
|
|
mpart = None
|
|
if '--m-tail' in flags:
|
|
lines = ptx.replace('\r\n', '\n').split('\n')
|
|
if kernel == 'rrlite_enc0_4x4_mvhi_hdr_folded':
|
|
if lines[2547].strip() != 'bar.warp.sync -1;':
|
|
raise Err('enc0 cut changed')
|
|
lines = lines[:2548] + ['d4r.enc0 %r259, %rd106, %rd78, %r1066, %r1067, %r1230, %r1233;', 'ret;', '}']
|
|
mpart = 'enc0_tail'
|
|
elif kernel == 'rrlite_dec0_4x4_folded':
|
|
if lines[138].strip() != 'mov.u32 %r17, %laneid;' or lines[586].strip() != 'st.shared.v2.u16 [%r570+1168], {%rs63, %rs64};':
|
|
raise Err('dec0 cut changed')
|
|
lines[138:587] = ['mov.u32 %r560, _ZZ22rrlite_dec0_4x4_foldedN6RRLite18Decoder0ParametersEE4smem;',
|
|
'd4r.dec0 %r560, %rd22, %rd24, %r475, %r476, %r3, %r4;']
|
|
mpart = 'dec0_head'
|
|
else:
|
|
raise Err('--m-tail requires folded M enc0 or dec0')
|
|
ptx = '\n'.join(lines)
|
|
tr = Translator(ptx, kernel)
|
|
n_rz = tr.rewrite_rz()
|
|
if '--ifconv' in flags:
|
|
import ifconv, ctile as _ct
|
|
tr.body, n_if = ifconv.apply(tr.body, tokenize, _ct.parse, split_operands)
|
|
sys.stderr.write(f'ifconv: {n_if} diamonds converted\n')
|
|
stamps = [f.split('=', 1)[1].split(',') for f in flags if f.startswith('--stamps=')]
|
|
tr.stamps = stamps[0] if stamps else []
|
|
tr.bbcount = '--bbcount' in flags
|
|
tr.as1 = '--as1' in flags
|
|
tr.bblabels = []
|
|
ct = None
|
|
if '--ctile' in flags:
|
|
import ctile
|
|
ct = ctile.emit(tr, tokenize, Translator)
|
|
if '--sink' in flags:
|
|
import ifconv, ctile as _ct
|
|
tr.body, n_sink = ifconv.sink_stores(tr.body, tokenize, _ct.parse)
|
|
sys.stderr.write(f'sink: {n_sink} stores moved\n')
|
|
orig_op = tr.op
|
|
|
|
def op_with_round(op, a):
|
|
if op in ('d4r.hmin2', 'd4r.hmax2'):
|
|
fn = 'm_min2' if op == 'd4r.hmin2' else 'm_max2'
|
|
left = ', '.join(tr.val(v, 'f16') for v in a[2:4])
|
|
right = ', '.join(tr.val(v, 'f16') for v in a[4:6])
|
|
tr.emit(f'{{ uint32_t packed_ = ash2u({fn}((h2){{{left}}}, (h2){{{right}}}));')
|
|
tr.assign(a[0], 'packed_ & 0xffffu', 'u16')
|
|
tr.assign(a[1], 'packed_ >> 16', 'u16')
|
|
tr.emit('}')
|
|
return
|
|
if op in ('d4r.enc0', 'd4r.dec0'):
|
|
vals = [tr.val(v, t) for v, t in zip(a, ('u32', 'u64', 'u64', 'u32', 'u32', 's32', 's32'))]
|
|
tr.emit(f'd4r_{"enc0_tail" if op == "d4r.enc0" else "dec0_head"}({", ".join(vals)});')
|
|
return
|
|
if op == 'd4r.roundf':
|
|
tr.assign(a[0], f'__builtin_roundf({tr.val(a[1], "f32")})', 'float')
|
|
return
|
|
orig_op(op, a)
|
|
tr.op = op_with_round
|
|
tr.translate()
|
|
out = []
|
|
out.append(f'// Generated by kernels/native/ptx2hip.py from NVIDIA PTX ({kernel}); {n_rz} roundf idioms.')
|
|
out.append('#include "native_rt.h"')
|
|
if kernel.startswith('rrlite_'):
|
|
out.append('#include "m_ptx_rt.h"')
|
|
if 'cvt.rn.f16x2.e4m3x2' in ptx:
|
|
out.append('RT uint32_t d4r_decode_e4m3x2(uint16_t v) { h2 h; for (int i = 0; i < 2; ++i) { '
|
|
'uint32_t b = (v >> (8*i)) & 255u; '
|
|
'h[i] = (b & 127u) == 127u ? ash((uint16_t)(0x7fffu | ((b & 128u) << 8))) : '
|
|
'(h1)(ash((uint16_t)(((b & 127u) << 7) | ((b & 128u) << 8))) * (h1)256.0f); } return ash2u(h); }')
|
|
if mpart:
|
|
out.append(f'#include "../tex/{mpart}.hip"')
|
|
if ct is None:
|
|
for name, size, align in tr.shared:
|
|
out.append(f'__attribute__((shared)) __attribute__((aligned({max(align, 16)}))) static uint8_t {name}[{size}];')
|
|
else:
|
|
# one LDS block: the kernel's own arrays (path B) and the color tile (path A) are never live together
|
|
off = 0
|
|
for name, size, align in tr.shared:
|
|
off = (off + 15) & ~15
|
|
out.append(f'#define {name} (d4r_lds + {off})')
|
|
off += size
|
|
prelude, fill, info = ct
|
|
out.insert(2, f'__attribute__((shared)) __attribute__((aligned(16))) static uint8_t d4r_lds[{max(off, info["lds_bytes"])}];')
|
|
out.insert(2, f'// --ctile: {info["taps"]} color taps through an LDS tile, {len(info["groups"])} conversion shapes')
|
|
out.extend(prelude)
|
|
out.extend(fill)
|
|
out.append(f'struct __attribute__((aligned({tr.param_align}))) {kernel}_params {{ uint8_t b[{tr.param_size}]; }};')
|
|
maxntid = re.search(r'\.maxntid\s+(\d+)', tr.attrs)
|
|
lb = f'__attribute__((amdgpu_flat_work_group_size(1, {maxntid.group(1)})))' if maxntid else ''
|
|
bounds = [f.split('=', 1)[1] for f in flags if f.startswith('--max-block=')]
|
|
if bounds:
|
|
n = int(bounds[0])
|
|
if not 1 <= n <= 1024:
|
|
raise Err('invalid --max-block')
|
|
lb = f'__attribute__((amdgpu_flat_work_group_size(1, {n})))'
|
|
out.append(f'extern "C" __attribute__((global)) {lb} void {kernel}({kernel}_params)')
|
|
out.append('{')
|
|
out.append(' const AS4 uint8_t* KARGP = (const AS4 uint8_t*)__builtin_amdgcn_kernarg_segment_ptr();')
|
|
out.append(' const uint64_t KARGU = (uint64_t)(uintptr_t)KARGP;')
|
|
for name, size, align in tr.locals:
|
|
out.append(f' __attribute__((aligned({align}))) uint8_t {name}[{size}];')
|
|
if ct is not None:
|
|
out.append(' int32_t ct_x0 = 0, ct_y0 = 0, ct_w = 0, ct_h = 0;')
|
|
if tr.bbcount:
|
|
# diagnostic build: d4r_bbc[k] = number of waves that entered basic block k (label list in a comment)
|
|
out.insert(2, '__attribute__((device)) unsigned long long d4r_bbc[512];\n'
|
|
'#define D4R_BBC(k) do { if (wave_lane() == (uint32_t)__builtin_ctzll(__builtin_amdgcn_read_exec())) '
|
|
'__hip_atomic_fetch_add(&d4r_bbc[k], 1ull, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); } while (0)\n'
|
|
'// d4r_bbc labels: ' + ' '.join(tr.bblabels))
|
|
if tr.stamps:
|
|
# diagnostic build: per-region wave cycles (20-bit SHADER_CYCLES deltas) summed into d4r_dbg[k],
|
|
# wave counts into d4r_dbg[32 + k]; region k ends at stamp label k (k = len + 1: kernel return)
|
|
out.insert(2, '__attribute__((device)) unsigned long long d4r_dbg[64];\n'
|
|
'#define D4R_STAMP(k) do { uint32_t n_ = (uint32_t)__builtin_readcyclecounter(); '
|
|
'if (wave_lane() == (uint32_t)__builtin_ctzll(__builtin_amdgcn_read_exec())) { '
|
|
'__hip_atomic_fetch_add(&d4r_dbg[k], (unsigned long long)((n_ - st_last) & 0xfffffu), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); '
|
|
'__hip_atomic_fetch_add(&d4r_dbg[32 + (k)], 1ull, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); } '
|
|
'st_last = n_; } while (0)')
|
|
out.append(' uint32_t st_last = (uint32_t)__builtin_readcyclecounter();')
|
|
by_type = {}
|
|
for name, t in sorted(tr.regs.items()):
|
|
by_type.setdefault(t, []).append(name)
|
|
for t, names in by_type.items():
|
|
for i in range(0, len(names), 16):
|
|
out.append(f' {t} {", ".join(names[i:i + 16])};')
|
|
out.extend(tr.lines)
|
|
out.append('}')
|
|
print('\n'.join(out))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|