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
73994d4634
commit
25493b3198
1 file changed
+1000
+1000
@@ -0,0 +1,1000 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Translate one NVIDIA PTX texture kernel into a Vulkan GLSL compute shader.
|
||||
|
||||
usage: ptx2glsl.py MODULE.ptx KERNEL [--tex OFF:linear|nearest ...] > KERNEL.comp
|
||||
|
||||
Every PTX instruction becomes the GLSL expression of the same operation that kernels/native/ptx2hip.py
|
||||
emits for HIP, so the shader computes what the translated kernel computes. Registers become locals and
|
||||
basic blocks become guarded sections of a program-counter loop (GLSL has no goto). The 8-byte handles in
|
||||
the parameter block become bindings: textures by their parameter offset (T<off>), surfaces (S<off>);
|
||||
pointers stay 64-bit addresses and are read through buffer references.
|
||||
"""
|
||||
import re
|
||||
import sys
|
||||
|
||||
REG_RE = re.compile(r'%([a-zA-Z]+)(\d+)$')
|
||||
GT = {'f': 'float', 'r': 'uint', 'rd': 'uint64_t', 'rs': 'uint', 'p': 'bool'}
|
||||
SREG = {'%tid.x': 'gl_LocalInvocationID.x', '%tid.y': 'gl_LocalInvocationID.y', '%tid.z': 'gl_LocalInvocationID.z',
|
||||
'%ctaid.x': 'gl_WorkGroupID.x', '%ctaid.y': 'gl_WorkGroupID.y', '%ctaid.z': 'gl_WorkGroupID.z',
|
||||
'%ntid.x': 'gl_WorkGroupSize.x', '%ntid.y': 'gl_WorkGroupSize.y',
|
||||
'%nctaid.x': 'gl_NumWorkGroups.x', '%nctaid.y': 'gl_NumWorkGroups.y', '%laneid': 'gl_SubgroupInvocationID'}
|
||||
|
||||
|
||||
class Err(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def tokenize(body):
|
||||
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:
|
||||
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 T:
|
||||
def __init__(self, ptx, kernel):
|
||||
text = re.sub(r'//[^\n]*', '', ptx)
|
||||
m = re.search(r'\.entry\s+' + re.escape(kernel) + r'\s*\((.*?)\)', text, re.S)
|
||||
if not m:
|
||||
raise Err('kernel not found')
|
||||
pm = re.search(r'\.param\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', m.group(1))
|
||||
self.kernel, self.pname, self.psize = kernel, pm.group(2), int(pm.group(3))
|
||||
rest = text[m.end():]
|
||||
mt = re.search(r'\.maxntid\s+(\d+)\s*,\s*(\d+)\s*,\s*(\d+)', rest[:rest.index('{')])
|
||||
self.maxntid = [int(x) for x in mt.groups()] if mt else [256, 1, 1]
|
||||
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.toks = tokenize(rest[start + 1:k])
|
||||
self.regs = {}
|
||||
self.shared = {} # name -> byte base
|
||||
self.shared_bytes = 0
|
||||
self.scopes = [{}]
|
||||
self.scope_id = 0
|
||||
self.blocks = [[]] # lists of lines
|
||||
self.labels = {} # label -> block index
|
||||
self.textures, self.surfaces = set(), {}
|
||||
self.dyn_tex = False
|
||||
self.fast = False # --fast: no 'precise', the compiler may contract and reassociate
|
||||
self.collect = False # first pass: gather how every 16/32-bit register is used
|
||||
self.ctx, self.links, self.width, self.half = {}, [], {}, set()
|
||||
self.exp_lut = False
|
||||
self.prescan()
|
||||
|
||||
# ---------------------------------------------------------------- static facts about registers
|
||||
def prescan(self):
|
||||
defs = {}
|
||||
for kind, s in self.toks:
|
||||
if kind != 'stmt' or s.startswith('.'):
|
||||
continue
|
||||
s = re.sub(r'^@!?%p\d+\s+', '', s)
|
||||
parts = s.split(None, 1)
|
||||
if len(parts) < 2:
|
||||
continue
|
||||
a = split_operands(parts[1])
|
||||
for d in re.findall(r'%rd\d+', a[0]) if not a[0].startswith('[') else []:
|
||||
defs.setdefault(d, []).append((parts[0], a))
|
||||
self.pbase = {} # reg -> constant offset from the parameter block
|
||||
self.handle = {} # reg -> parameter offset its 8-byte value was loaded from
|
||||
changed = True
|
||||
while changed:
|
||||
changed = False
|
||||
for r, ds in defs.items():
|
||||
if len(ds) != 1:
|
||||
continue
|
||||
op, a = ds[0]
|
||||
if r not in self.pbase:
|
||||
v = None
|
||||
if op.startswith('mov') and a[1] == self.pname:
|
||||
v = 0
|
||||
elif op in ('add.s64', 'add.u64') and a[1] in self.pbase and re.match(r'-?\d+$', a[2]):
|
||||
v = self.pbase[a[1]] + int(a[2])
|
||||
if v is not None:
|
||||
self.pbase[r] = v; changed = True
|
||||
if r not in self.handle:
|
||||
v = None
|
||||
if op == 'ld.param.u64':
|
||||
v = self.param_off(a[1])
|
||||
elif op.startswith('mov') and a[1] in self.handle:
|
||||
v = self.handle[a[1]]
|
||||
if v is not None:
|
||||
self.handle[r] = v; changed = True
|
||||
|
||||
def param_off(self, tok):
|
||||
m = re.match(r'\[\s*([^\]+]+?)\s*(?:\+\s*(-?\d+))?\s*\]$', tok)
|
||||
base, off = m.group(1), int(m.group(2) or 0)
|
||||
if base == self.pname:
|
||||
return off
|
||||
if base in self.pbase:
|
||||
return self.pbase[base] + off
|
||||
return None
|
||||
|
||||
# ---------------------------------------------------------------- operands
|
||||
def reg(self, tok):
|
||||
tok = tok.strip()
|
||||
for sc in reversed(self.scopes):
|
||||
if tok in sc:
|
||||
return sc[tok]
|
||||
m = REG_RE.match(tok)
|
||||
if not m or m.group(1) not in GT:
|
||||
raise Err('not a register: ' + tok)
|
||||
name = m.group(1) + m.group(2)
|
||||
if m.group(1) in ('r', 'rs'):
|
||||
self.width[name] = 32 if m.group(1) == 'r' else 16
|
||||
self.regs[name] = self.htype(name) or GT[m.group(1)]
|
||||
return name
|
||||
|
||||
def htype(self, name):
|
||||
return ('f16vec2' if self.width[name] == 32 else 'float16_t') if name in self.half else None
|
||||
|
||||
def note(self, name, cls):
|
||||
if self.collect:
|
||||
self.ctx.setdefault(name, set()).add(cls)
|
||||
|
||||
def hk(self, name):
|
||||
return 'f16x2' if self.width.get(name) == 32 else 'f16'
|
||||
|
||||
def is_reg(self, tok):
|
||||
return tok.startswith('%') and tok not in SREG or any(tok in sc for sc in self.scopes)
|
||||
|
||||
def val(self, tok, t):
|
||||
tok = tok.strip()
|
||||
if tok in SREG:
|
||||
return f'uint({SREG[tok]})' if t != 's32' else f'int({SREG[tok]})'
|
||||
if tok in ('%SP', '%SPL'):
|
||||
return 'uint64_t(0ul)' # the local depot is its own array; addresses in it start at 0
|
||||
if self.is_reg(tok):
|
||||
n = self.reg(tok); g = self.regs[n]
|
||||
self.note(n, 'h' if t in ('f16', 'f16x2') else 'u')
|
||||
if g in ('float16_t', 'f16vec2'):
|
||||
if t in ('f16', 'f16x2'):
|
||||
return n
|
||||
bits = f'HU({n})' if g == 'float16_t' else f'H2U({n})'
|
||||
if t == 's16':
|
||||
return f's16({bits})'
|
||||
if t == 's32':
|
||||
return f'int({bits})'
|
||||
if t in ('u64', 'b64'):
|
||||
return f'uint64_t({bits})'
|
||||
return bits
|
||||
if t == 'f32':
|
||||
return n if g == 'float' else f'uintBitsToFloat({n})'
|
||||
if t in ('u32', 'b32'):
|
||||
return f'floatBitsToUint({n})' if g == 'float' else (f'uint({n})' if g == 'uint64_t' else n)
|
||||
if t == 's32':
|
||||
return f'floatBitsToInt({n})' if g == 'float' else f'int({n})'
|
||||
if t in ('u64', 'b64'):
|
||||
return n if g == 'uint64_t' else f'uint64_t({n})'
|
||||
if t == 's64':
|
||||
return f'int64_t({n})'
|
||||
if t in ('u16', 'b16'):
|
||||
return n
|
||||
if t == 's16':
|
||||
return f's16({n})'
|
||||
if t == 'f16':
|
||||
return f'H({n})'
|
||||
if t == 'f16x2':
|
||||
return f'H2({n})'
|
||||
if t == 'pred':
|
||||
return n
|
||||
raise Err('type ' + t)
|
||||
if tok in self.shared:
|
||||
return f'{self.shared[tok]}u'
|
||||
if re.match(r'0[fF][0-9A-Fa-f]{8}$', tok):
|
||||
bits = int(tok[2:], 16)
|
||||
return f'uintBitsToFloat(0x{bits:08X}u)' if t == '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 t == 'f32':
|
||||
raise Err('int immediate as float')
|
||||
if t in ('u64', 'b64'):
|
||||
return f'uint64_t({v & 0xffffffffffffffff}ul)'
|
||||
if t == 's64':
|
||||
return f'int64_t({v}l)'
|
||||
if t in ('u16', 'b16'):
|
||||
return f'{v & 0xffff}u'
|
||||
if t == 's16':
|
||||
v &= 0xffff
|
||||
return f'({v - 0x10000 if v >= 0x8000 else v})'
|
||||
if t == 'f16':
|
||||
return f'H({v & 0xffff}u)'
|
||||
if t == 'f16x2':
|
||||
return f'H2({v & 0xffffffff}u)'
|
||||
if t == 's32':
|
||||
v &= 0xffffffff
|
||||
return f'int(0x{v:08X}u)' if v >= 1 << 31 else f'{v}'
|
||||
return f'{v & 0xffffffff}u'
|
||||
raise Err('operand ' + tok)
|
||||
|
||||
def A(self, dst, expr, st):
|
||||
n = self.reg(dst); g = self.regs[n]
|
||||
self.note(n, 'h' if st in ('f16', 'f16x2') else 'u')
|
||||
if g == 'bool' and st in ('u32', 'u16', 'u64'):
|
||||
self.E(f'{n} = ({expr}) != 0u;'); return # mov.pred with an immediate
|
||||
if g in ('float16_t', 'f16vec2'):
|
||||
if st in ('f16', 'f16x2'):
|
||||
self.E(f'{n} = {expr};')
|
||||
else:
|
||||
self.E(f'{n} = {"H" if g == "float16_t" else "H2"}(uint({expr}));')
|
||||
return
|
||||
if st == 'float':
|
||||
e = expr if g == 'float' else f'floatBitsToUint({expr})'
|
||||
elif st == 'u32':
|
||||
e = f'uintBitsToFloat({expr})' if g == 'float' else (f'uint64_t({expr})' if g == 'uint64_t' else f'uint({expr})')
|
||||
elif st == 'u64':
|
||||
e = f'uint64_t({expr})'
|
||||
elif st == 'u16':
|
||||
e = f'(uint({expr}) & 0xffffu)'
|
||||
elif st == 'f16':
|
||||
e = f'HU({expr})'
|
||||
elif st == 'f16x2':
|
||||
e = f'H2U({expr})'
|
||||
elif st == 'bool':
|
||||
e = f'({expr})'
|
||||
else:
|
||||
raise Err('src type')
|
||||
self.E(f'{n} = {e};')
|
||||
|
||||
def E(self, s):
|
||||
self.blocks[-1].append(s)
|
||||
|
||||
# ---------------------------------------------------------------- driver
|
||||
def translate(self):
|
||||
for kind, text in self.toks:
|
||||
if kind == 'open':
|
||||
self.scopes.append({}); continue
|
||||
if kind == 'close':
|
||||
self.scopes.pop(); continue
|
||||
if kind == 'label':
|
||||
self.blocks.append([])
|
||||
self.labels[text] = len(self.blocks) - 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
|
||||
self.scope_id += 1
|
||||
for nm in names.split(','):
|
||||
nm = nm.strip()
|
||||
c = f'l{self.scope_id}_' + re.sub(r'\W', '_', nm)
|
||||
self.scopes[-1][nm] = c
|
||||
if ty in ('f16', 'b16', 'f16x2', 'b32', 'u32', 's32'):
|
||||
self.width[c] = 16 if ty in ('f16', 'b16') else 32
|
||||
self.regs[c] = self.htype(c) if c in self.half else {'f16': 'uint', 'b16': 'uint', 'f32': 'float', 'b32': 'uint', 'pred': 'bool', 'f16x2': 'uint',
|
||||
'u32': 'uint', 's32': 'uint', 'b64': 'uint64_t', 'u64': 'uint64_t'}[ty]
|
||||
return
|
||||
if s.startswith('.local'):
|
||||
m = re.match(r'\.local\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s)
|
||||
self.local_bytes = max(getattr(self, 'local_bytes', 0), int(m.group(3)))
|
||||
return
|
||||
if re.match(r'mov\.u64 %SPL?, __local_depot', s):
|
||||
return
|
||||
if s.startswith('.shared'):
|
||||
m = re.match(r'\.shared\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s)
|
||||
self.shared[m.group(2)] = self.shared_bytes
|
||||
self.shared_bytes += (int(m.group(3)) + 3) & ~3
|
||||
return
|
||||
if s.startswith('.'):
|
||||
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]
|
||||
a = split_operands(parts[1]) if len(parts) > 1 else []
|
||||
if op.startswith('bra'):
|
||||
tgt = ('L', a[0])
|
||||
if pred is not None:
|
||||
self.E(('BR', pred, tgt)) # conditional: falls through into a new block
|
||||
else:
|
||||
self.E(('BR', None, tgt))
|
||||
self.blocks.append([])
|
||||
return
|
||||
if op == 'ret':
|
||||
self.E(('BR', pred, ('END', None)))
|
||||
self.blocks.append([])
|
||||
return
|
||||
if pred is not None:
|
||||
self.E(f'if ({pred}) {{')
|
||||
self.op(op, a)
|
||||
self.E('}')
|
||||
return
|
||||
self.op(op, a)
|
||||
|
||||
def cmp(self, mods, isf):
|
||||
table = {'lt': '{0} < {1}', 'gt': '{0} > {1}', 'le': '{0} <= {1}', 'ge': '{0} >= {1}', 'eq': '{0} == {1}', 'ne': '{0} != {1}',
|
||||
'ltu': '!({0} >= {1})', 'gtu': '!({0} <= {1})', 'leu': '!({0} > {1})', 'geu': '!({0} < {1})',
|
||||
'equ': '!({0} != {1})', 'neu': '!({0} == {1})'}
|
||||
for k in table:
|
||||
if k in mods:
|
||||
if isf and k == 'ne':
|
||||
return '({0} < {1} || {0} > {1})'
|
||||
if isf and k == 'equ':
|
||||
return '!({0} < {1} || {0} > {1})'
|
||||
return table[k]
|
||||
raise Err('compare')
|
||||
|
||||
def texref(self, tok):
|
||||
m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', tok)
|
||||
return m.group(1), [x.strip() for x in m.group(2).split(',')]
|
||||
|
||||
def op(self, op, a):
|
||||
o = op.split('.'); base = o[0]; mods = set(o[1:])
|
||||
E, A, V = self.E, self.A, self.val
|
||||
ft = next((t for t in ('f32', 'f16x2', 'f16') if t in mods), None)
|
||||
if base == 'bar':
|
||||
E('memoryBarrierShared(); barrier();'); return
|
||||
if base == 'mov':
|
||||
dst, src = a
|
||||
if dst.startswith('{'):
|
||||
n = [x.strip() for x in dst[1:-1].split(',')]
|
||||
sr = self.reg(src) if self.is_reg(src) else None
|
||||
if self.collect and sr:
|
||||
for x in [sr] + [self.reg(y) for y in n]:
|
||||
self.ctx.setdefault(x, set()).add('h')
|
||||
return
|
||||
if sr and (sr in self.half or any(self.reg(y) in self.half for y in n)):
|
||||
A(n[0], f'{V(src, "f16x2")}.x', 'f16'); A(n[1], f'{V(src, "f16x2")}.y', 'f16'); return
|
||||
A(n[0], f'{V(src, "u32")} & 0xffffu', 'u16'); A(n[1], f'{V(src, "u32")} >> 16', 'u16'); return
|
||||
if src.startswith('{'):
|
||||
n = [x.strip() for x in src[1:-1].split(',')]
|
||||
dr = self.reg(dst)
|
||||
if self.collect:
|
||||
for x in [dr] + [self.reg(y) for y in n if self.is_reg(y)]:
|
||||
self.ctx.setdefault(x, set()).add('h')
|
||||
return
|
||||
if dr in self.half:
|
||||
A(dst, f'f16vec2({V(n[0], "f16")}, {V(n[1], "f16")})', 'f16x2'); return
|
||||
A(dst, f'({V(n[0], "u16")} & 0xffffu) | ({V(n[1], "u16")} << 16)', 'u32'); return
|
||||
if self.is_reg(dst) and self.reg(dst) in self.width and not (mods & {'f32', 'u64', 'b64', 's64'}):
|
||||
dr = self.reg(dst)
|
||||
srcreg = self.reg(src) if self.is_reg(src) and src not in SREG else None
|
||||
isimm = bool(re.match(r'-?\d+$', src) or re.match(r'0x[0-9a-fA-F]+$', src))
|
||||
if self.collect and srcreg and srcreg in self.width:
|
||||
self.links.append((dr, srcreg)); return
|
||||
if self.collect and isimm:
|
||||
return
|
||||
if dr in self.half and (srcreg or isimm):
|
||||
A(dst, V(src, self.hk(dr)), self.hk(dr)); return
|
||||
if src == self.pname:
|
||||
A(dst, '0ul', '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
|
||||
# ------------------------------------------------ f32
|
||||
if ft == 'f32' and base in ('add', 'sub', 'mul'):
|
||||
if 'rz' in mods:
|
||||
A(a[0], f'round_half_away({V(a[1], "f32")})', 'float'); return # the roundf idiom: + copysign(0.5), then trunc
|
||||
A(a[0], f'{V(a[1], "f32")} {"+-*"[["add", "sub", "mul"].index(base)]} {V(a[2], "f32")}', 'float'); return
|
||||
if ft == 'f32' and base == 'fma':
|
||||
A(a[0], f'fma({V(a[1], "f32")}, {V(a[2], "f32")}, {V(a[3], "f32")})', 'float'); return
|
||||
if ft == 'f32' and base in ('max', 'min'):
|
||||
A(a[0], f'{base}({V(a[1], "f32")}, {V(a[2], "f32")})', 'float'); return
|
||||
if ft == 'f32' and base == 'neg':
|
||||
A(a[0], f'-{V(a[1], "f32")}', 'float'); return
|
||||
if ft == 'f32' and base == 'abs':
|
||||
A(a[0], f'abs({V(a[1], "f32")})', 'float'); return
|
||||
if ft == 'f32' and base in ('rcp', 'ex2', 'lg2', 'rsqrt', 'sqrt'):
|
||||
fn = {'rcp': 'rcp_a', 'ex2': 'exp2', 'lg2': 'log2', 'rsqrt': 'inversesqrt', 'sqrt': 'sqrt'}[base]
|
||||
A(a[0], f'{fn}({V(a[1], "f32")})', 'float'); return
|
||||
if ft == 'f32' and base == 'div':
|
||||
A(a[0], f'{V(a[1], "f32")} * rcp_a({V(a[2], "f32")})', 'float'); return
|
||||
# ------------------------------------------------ f16
|
||||
if ft == 'f16x2' and base in ('add', 'sub', 'mul', 'fma'):
|
||||
x = [V(t, 'f16x2') for t in a[1:]]
|
||||
e = f'fma({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'min(max({e}, f16vec2(0.0hf)), f16vec2(1.0hf))'
|
||||
A(a[0], e, 'f16x2'); return
|
||||
if ft == 'f16x2' and base == 'abs':
|
||||
A(a[0], f'abs({V(a[1], "f16x2")})', 'f16x2'); return
|
||||
if ft == 'f16x2' and base == 'ex2':
|
||||
if self.exp_lut:
|
||||
A(a[0], f'ex2_lut({V(a[1], "u32")})', 'u32'); return
|
||||
A(a[0], f'f16vec2(exp2(vec2({V(a[1], "f16x2")})))', 'f16x2'); return
|
||||
if ft in ('f16', 'f16x2') and base in ('min', 'max'):
|
||||
A(a[0], f'{base}({V(a[1], ft)}, {V(a[2], ft)})', ft); return
|
||||
if ft == 'f16' and base in ('add', 'sub', 'mul', 'fma'):
|
||||
x = [V(t, 'f16') for t in a[1:]]
|
||||
e = f'fma({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'min(max({e}, 0.0hf), 1.0hf)'
|
||||
A(a[0], e, 'f16'); return
|
||||
if base == 'set' and 'f16x2' in mods and 'u32' in mods:
|
||||
c = self.cmp(mods, True); x, y = V(a[1], 'f16x2'), V(a[2], 'f16x2')
|
||||
A(a[0], f'(({c.format(x + ".x", y + ".x")}) ? 0xffffu : 0u) | (({c.format(x + ".y", y + ".y")}) ? 0xffff0000u : 0u)', 'u32'); return
|
||||
# ------------------------------------------------ compares / selects
|
||||
if base == 'setp':
|
||||
isf = bool(mods & {'f32', 'f16'})
|
||||
t = next(t for t in ('f32', 'f16', 's32', 'u32', 'b32', 's16', 'u16', 'b16', 's64', 'u64') if t in mods)
|
||||
if '|' in a[0]:
|
||||
raise Err('setp form')
|
||||
A(a[0], self.cmp(mods, isf).format(V(a[1], t), V(a[2], t)), 'bool'); return
|
||||
if base == 'selp' and not (mods & {'f32', 'u64', 'b64', 's64'}) and self.reg(a[0]) in self.width:
|
||||
dr = self.reg(a[0])
|
||||
srcs = [self.reg(x) for x in a[1:3] if self.is_reg(x) and x not in SREG]
|
||||
if self.collect:
|
||||
self.links += [(dr, x) for x in srcs if x in self.width]
|
||||
V(a[3], 'pred')
|
||||
if any(x not in self.width for x in srcs):
|
||||
self.note(dr, 'u')
|
||||
return
|
||||
if dr in self.half:
|
||||
k = self.hk(dr)
|
||||
A(a[0], f'{V(a[3], "pred")} ? {V(a[1], k)} : {V(a[2], k)}', k); 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'
|
||||
A(a[0], f'{V(a[3], "pred")} ? {V(a[1], t)} : {V(a[2], t)}', {'f32': 'float', 'u16': 'u16', 'u64': 'u64', 'u32': 'u32'}[t]); return
|
||||
if 'pred' in mods and base in ('and', 'or', 'xor', 'not'):
|
||||
if base == 'not':
|
||||
A(a[0], f'!{V(a[1], "pred")}', 'bool'); return
|
||||
A(a[0], f'{V(a[1], "pred")} {{"and": "&&", "or": "||", "xor": "!="}}'.replace('{"and": "&&", "or": "||", "xor": "!="}', {'and': '&&', 'or': '||', 'xor': '!='}[base]) + f' {V(a[2], "pred")}', 'bool'); return
|
||||
# ------------------------------------------------ integer
|
||||
if base in ('and', 'or', 'xor', 'not', 'shl', 'shr', 'add', 'sub', 'mul', 'mad', 'min', 'max', 'div', 'rem', 'neg', 'abs') and ft is None:
|
||||
w = 64 if mods & {'u64', 's64', 'b64'} else 16 if mods & {'u16', 's16', 'b16'} else 32
|
||||
signed = bool(mods & {'s32', 's64', 's16'})
|
||||
ut, st = {64: ('u64', 's64'), 32: ('u32', 's32'), 16: ('u16', 's16')}[w]
|
||||
if base == 'mul' and 'wide' in mods:
|
||||
if 's32' in mods:
|
||||
A(a[0], f'uint64_t(int64_t({V(a[1], "s32")}) * int64_t({V(a[2], "s32")}))', 'u64'); return
|
||||
if 'u32' in mods:
|
||||
A(a[0], f'uint64_t({V(a[1], "u32")}) * uint64_t({V(a[2], "u32")})', 'u64'); return
|
||||
if 'u16' in mods:
|
||||
A(a[0], f'{V(a[1], "u16")} * {V(a[2], "u16")}', 'u32'); return
|
||||
raise Err(op)
|
||||
if base == 'not':
|
||||
A(a[0], f'~{V(a[1], ut)}', ut); return
|
||||
if base == 'neg':
|
||||
A(a[0], f'0u - {V(a[1], ut)}', ut); return
|
||||
if base == 'shl':
|
||||
A(a[0], f'{V(a[1], ut)} << ({V(a[2], "u32")} & {w - 1}u)', ut); return
|
||||
if base == 'shr':
|
||||
x = V(a[1], st if signed else ut)
|
||||
A(a[0], f'{x} >> ({V(a[2], "u32")} & {w - 1}u)', ut); return
|
||||
if base in ('and', 'or', 'xor'):
|
||||
A(a[0], f'{V(a[1], ut)} {{"and": "&", "or": "|", "xor": "^"}}'.replace('{"and": "&", "or": "|", "xor": "^"}', {'and': '&', 'or': '|', 'xor': '^'}[base]) + f' {V(a[2], ut)}', ut); return
|
||||
if base in ('add', 'sub'):
|
||||
A(a[0], f'{V(a[1], ut)} {"+" if base == "add" else "-"} {V(a[2], ut)}', ut); return
|
||||
if base == 'mul' and 'hi' in mods:
|
||||
if signed:
|
||||
A(a[0], f'uint(int((int64_t({V(a[1], "s32")}) * int64_t({V(a[2], "s32")})) >> 32))', 'u32'); return
|
||||
A(a[0], f'uint((uint64_t({V(a[1], "u32")}) * uint64_t({V(a[2], "u32")})) >> 32)', 'u32'); return
|
||||
if base == 'mul':
|
||||
A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)}', ut); return
|
||||
if base == 'mad':
|
||||
A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)} + {V(a[3], ut)}', ut); return
|
||||
if base == 'abs':
|
||||
A(a[0], f'uint(abs({V(a[1], st)}))', ut); return
|
||||
if base == 'neg' and len(a) == 2:
|
||||
A(a[0], f'uint(-{V(a[1], st)})', ut); return
|
||||
x, y = V(a[1], st if signed else ut), V(a[2], st if signed else ut)
|
||||
if base in ('min', 'max'):
|
||||
A(a[0], f'{base}({x}, {y})', ut); return
|
||||
if base == 'div':
|
||||
A(a[0], f'{x} / {y}', ut); return
|
||||
if base == 'rem':
|
||||
A(a[0], f'{x} % {y}', ut); 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 base == 'cvt':
|
||||
return self.cvt(op, o[1:], a)
|
||||
if base == 'cvta':
|
||||
A(a[0], V(a[1], 'u64'), 'u64'); return
|
||||
if base == 'ld':
|
||||
return self.ld(o[1:], a)
|
||||
if base == 'st':
|
||||
return self.st(o[1:], a)
|
||||
if base in ('tex', 'tld4'):
|
||||
dst = [x.strip() for x in a[0][1:-1].split(',')]
|
||||
h, c = self.texref(a[1])
|
||||
off = self.handle.get(h)
|
||||
if off is None:
|
||||
raise Err('texture handle of ' + h + ' is not a single parameter load')
|
||||
self.textures.add(off)
|
||||
if base == 'tld4':
|
||||
E(f'{{ vec4 t_ = textureGather(T{off}, vec2({V(c[0], "f32")}, {V(c[1], "f32")}), 0);')
|
||||
elif 'base' in mods:
|
||||
E(f'{{ vec4 t_ = texelFetch(T{off}, ivec2({V(c[0], "s32")}, {V(c[1], "s32")}), 0);')
|
||||
elif 'level' in mods:
|
||||
E(f'{{ vec4 t_ = textureLod(T{off}, vec2({V(c[0], "f32")}, {V(c[1], "f32")}), {V(a[2], "f32")});')
|
||||
else:
|
||||
raise Err(op)
|
||||
for i, d in enumerate(dst):
|
||||
if 'f16' in o:
|
||||
A(d, f'float16_t(t_[{i}])', 'f16') # the texel as a half float
|
||||
else:
|
||||
A(d, f't_[{i}]', 'float')
|
||||
E('}')
|
||||
return
|
||||
if base == 'sust':
|
||||
h, c = self.texref(a[0])
|
||||
off = self.handle.get(h)
|
||||
if off is None:
|
||||
raise Err('surface handle of ' + h)
|
||||
v = [x.strip() for x in a[1][1:-1].split(',')]
|
||||
if op == 'sust.p.2d.v4.b32.zero':
|
||||
self.surfaces.setdefault(off, 'p'); assert self.surfaces[off] == 'p'
|
||||
E(f'sust_p(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, vec4({", ".join(V(x, "f32") for x in v)}));'); return
|
||||
if op == 'sust.b.2d.b32.zero':
|
||||
self.surfaces.setdefault(off, 'b'); assert self.surfaces[off] == 'b'
|
||||
E(f'sust_b(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, {V(v[0], "u32")});'); return
|
||||
if op == 'sust.b.2d.v2.b16.zero': # reading only: two 16-bit values at a byte-addressed position
|
||||
self.surfaces.setdefault(off, 'b')
|
||||
E(f'sust_b2x16(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, {V(v[0], "u16")}, {V(v[1], "u16")});'); return
|
||||
raise Err(op)
|
||||
if op == 'shfl.sync.bfly.b32':
|
||||
d = a[0].split('|')
|
||||
E('{ bool ib_;')
|
||||
A(d[0], f'shfl_bfly({V(a[1], "u32")}, {V(a[2], "s32")}, {V(a[3], "u32")}, ib_)', 'u32')
|
||||
if len(d) > 1:
|
||||
A(d[1], 'ib_', 'bool')
|
||||
E('}')
|
||||
return
|
||||
if op == 'shfl.sync.idx.b32':
|
||||
d = a[0].split('|')
|
||||
A(d[0], f'subgroupShuffle({V(a[1], "u32")}, {V(a[2], "u32")} & 31u)', 'u32')
|
||||
if len(d) > 1:
|
||||
A(d[1], 'true', 'bool')
|
||||
return
|
||||
if op == 'tanh.approx.f16':
|
||||
A(a[0], f'tanh({V(a[1], "f16")})', 'f16'); return
|
||||
if op == 'prmt.b32':
|
||||
A(a[0], f'prmt({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")})', 'u32'); return
|
||||
if op.startswith('mma.sync'):
|
||||
# not executable here: the kernel is translated for reading only (tensor-core fragments are not emulated)
|
||||
self.E(f'// {op} ' + ' '.join(a))
|
||||
for d in a[0][1:-1].split(','):
|
||||
A(d.strip(), '0u', 'u32')
|
||||
return
|
||||
if op == 'vote.sync.ballot.b32':
|
||||
# the warp is the subgroup (32 lanes in thread order), as for the shuffles
|
||||
A(a[0], f'subgroupBallot({V(a[1], "pred")}).x', 'u32'); return
|
||||
raise Err('unhandled ' + op)
|
||||
|
||||
def cvt(self, op, o, a):
|
||||
A, V = self.A, self.val
|
||||
mods = set(o)
|
||||
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)]
|
||||
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$|^e4m3x2$', t)]
|
||||
dt, st = types[0], types[1]
|
||||
# FP8 pairs (reading only: the helpers are not defined)
|
||||
if dt == 'f16x2' and st == 'e4m3x2':
|
||||
A(a[0], f'e4m3x2_to_f16x2({V(a[1], "u16")})', 'f16x2'); return
|
||||
if dt == 'e4m3x2' and st == 'f16x2':
|
||||
A(a[0], f'f16x2_to_e4m3x2({V(a[1], "f16x2")})', 'u16'); return
|
||||
if dt == 'f32' and st == 'f32':
|
||||
x = V(a[1], 'f32')
|
||||
if 'sat' in mods:
|
||||
A(a[0], f'min(max({x}, 0.0), 1.0)', 'float'); return
|
||||
fn = 'trunc' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven' if 'rni' in mods else 'ceil' if 'rpi' in mods else ''
|
||||
A(a[0], f'{fn}({x})', 'float'); return
|
||||
if dt == 's32' and st == 'f32':
|
||||
x = V(a[1], 'f32')
|
||||
fn = '' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven'
|
||||
A(a[0], f'uint(int({fn}({x})))', 'u32'); return
|
||||
if dt == 'u32' and st == 'f32':
|
||||
x = V(a[1], 'f32')
|
||||
fn = 'trunc' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven'
|
||||
A(a[0], f'uint(clamp({fn}({x}), 0.0, 4294967040.0))', 'u32'); return # CUDA saturates
|
||||
if dt == 'f32' and st == 's32':
|
||||
A(a[0], f'float({V(a[1], "s32")})', 'float'); return
|
||||
if dt == 'f32' and st == 'u32':
|
||||
A(a[0], f'float({V(a[1], "u32")})', 'float'); return
|
||||
if dt == 'f32' and st == 'f16':
|
||||
A(a[0], f'float({V(a[1], "f16")})', 'float'); return
|
||||
if dt == 'f16' and st == 'f32':
|
||||
A(a[0], f'float16_t({V(a[1], "f32")})', 'f16'); return
|
||||
if dt == 'f16x2' and st == 'f32':
|
||||
A(a[0], f'f16vec2(float16_t({V(a[2], "f32")}), float16_t({V(a[1], "f32")}))', 'f16x2'); return
|
||||
if dt == 'f16' and st == 'f16':
|
||||
fn = 'roundEven' if 'rni' in mods else 'floor'
|
||||
A(a[0], f'{fn}({V(a[1], "f16")})', 'f16'); return
|
||||
if dt == 's32' and st == 'f16':
|
||||
fn = 'roundEven' if 'rni' in mods else 'floor'
|
||||
A(a[0], f'uint(int(float({fn}({V(a[1], "f16")}))))', 'u32'); return
|
||||
if dt == 'f16' and st == 's32':
|
||||
A(a[0], f'float16_t({V(a[1], "s32")})', 'f16'); return
|
||||
if dt == 's64' and st == 's32':
|
||||
A(a[0], f'uint64_t(int64_t({V(a[1], "s32")}))', 'u64'); return
|
||||
if dt == 'u64' and st == 'u32':
|
||||
A(a[0], f'uint64_t({V(a[1], "u32")})', 'u64'); return
|
||||
if dt in ('u16', 's16') and st in ('u32', 's32'):
|
||||
A(a[0], V(a[1], 'u32'), 'u16'); return
|
||||
if dt in ('u32', 's32') and st in ('u16',):
|
||||
A(a[0], V(a[1], 'u16'), 'u32'); return
|
||||
if dt == 's32' and st == 's16':
|
||||
A(a[0], f'uint({V(a[1], "s16")})', 'u32'); return
|
||||
raise Err('cvt ' + op)
|
||||
|
||||
def addr(self, tok):
|
||||
m = re.match(r'\[\s*([^\]+]+?)\s*(?:\+\s*(-?\d+))?\s*\]$', tok)
|
||||
return m.group(1), int(m.group(2) or 0)
|
||||
|
||||
def ld(self, o, a):
|
||||
A, V, E = self.A, self.val, self.E
|
||||
mods = set(o)
|
||||
ty = next(t for t in o if re.match(r'^(f|s|u|b)\d+$', t))
|
||||
vec = next((int(t[1]) for t in o if re.match(r'^v\d$', t)), 1)
|
||||
size = int(ty[1:]) // 8
|
||||
dsts = [x.strip() for x in a[0][1:-1].split(',')] if a[0].startswith('{') else [a[0]]
|
||||
base, off = self.addr(a[1])
|
||||
for i, d in enumerate(dsts):
|
||||
if 'param' in mods:
|
||||
po = self.param_off(a[1])
|
||||
if po is None:
|
||||
raise Err('param address ' + a[1])
|
||||
e = f'PARAM{size * 8}({po + i * size}u)'
|
||||
elif 'shared' in mods and size == 2 and vec > 1:
|
||||
# a vector access is aligned to its whole size (PTX ISA): its 16-bit halves have fixed places in words
|
||||
e = f'(sh[({V(base, "u32")} + {off + i * size}u) >> 2] ' + ('& 0xffffu)' if i % 2 == 0 else '>> 16u)')
|
||||
elif 'shared' in mods:
|
||||
e = f'SH{size * 8}({V(base, "u32")} + {off + i * size}u)'
|
||||
elif 'local' in mods:
|
||||
e = f'loc[(uint({V(base, "u64")}) + {off + i * size}u) >> 2]'
|
||||
else:
|
||||
e = f'G{size * 8}({V(base, "u64")} + uint64_t({off + i * size}l))'
|
||||
if size == 8:
|
||||
A(d, e, 'u64')
|
||||
elif size == 2:
|
||||
A(d, e, 'u16')
|
||||
elif size == 1:
|
||||
A(d, e, 'u16')
|
||||
else:
|
||||
A(d, e, 'u32')
|
||||
|
||||
def st(self, o, a):
|
||||
V, E = self.val, self.E
|
||||
mods = set(o)
|
||||
ty = next(t for t in o if re.match(r'^(f|s|u|b)\d+$', t))
|
||||
size = int(ty[1:]) // 8
|
||||
base, off = self.addr(a[0])
|
||||
if 'shared' in mods:
|
||||
# vector stores go word by word; 16-bit pairs share a word (the kernels keep them 4-byte aligned)
|
||||
vals = [x.strip() for x in a[1][1:-1].split(',')] if a[1].startswith('{') else [a[1]]
|
||||
if size == 4:
|
||||
for i, v in enumerate(vals):
|
||||
E(f'sh[({V(base, "u32")} + {off + 4 * i}u) >> 2] = {V(v, "u32")};')
|
||||
return
|
||||
if size == 2 and len(vals) % 2 == 0:
|
||||
for i in range(0, len(vals), 2):
|
||||
E(f'sh[({V(base, "u32")} + {off + 2 * i}u) >> 2] = (uint({V(vals[i], "u16")}) & 0xFFFFu) | (uint({V(vals[i + 1], "u16")}) << 16);')
|
||||
return
|
||||
if 'local' in mods and size == 4:
|
||||
vals = [x.strip() for x in a[1][1:-1].split(',')] if a[1].startswith('{') else [a[1]]
|
||||
for i, v in enumerate(vals):
|
||||
E(f'loc[(uint({V(base, "u64")}) + {off + 4 * i}u) >> 2] = {V(v, "u32")};')
|
||||
return
|
||||
if 'global' in mods and size == 2 and not a[1].startswith('{'):
|
||||
E(f'W16({V(base, "u64")} + uint64_t({off}l)).v[0] = uint16_t({V(a[1], "u16")});'); return
|
||||
if 'global' in mods and size == 2 and a[1].startswith('{'):
|
||||
vals = [x.strip() for x in a[1][1:-1].split(',')]
|
||||
E(f'{{ W16 w_ = W16({V(base, "u64")} + uint64_t({off}l)); ' + ' '.join(f'w_.v[{i}] = uint16_t({V(v, "u16")});' for i, v in enumerate(vals)) + ' }'); return
|
||||
if 'global' in mods and size == 4:
|
||||
E(f'{{ W16 w_ = W16({V(base, "u64")} + uint64_t({off}l)); uint v_ = {V(a[1], "u32")}; w_.v[0] = uint16_t(v_ & 0xffffu); w_.v[1] = uint16_t(v_ >> 16); }}'); return
|
||||
raise Err('store ' + '.'.join(o))
|
||||
|
||||
# ---------------------------------------------------------------- output
|
||||
def preamble(self, w, tex_filter):
|
||||
w('#version 460')
|
||||
# scripts/check_redistributable.py refuses files carrying this line: they are NVIDIA's code.
|
||||
w(f'// Generated by vulkan/ptx2glsl.py from NVIDIA PTX ({self.kernel}); do not commit or redistribute.')
|
||||
for e in ('GL_EXT_shader_explicit_arithmetic_types', 'GL_EXT_buffer_reference', 'GL_EXT_buffer_reference2', 'GL_EXT_scalar_block_layout',
|
||||
'GL_KHR_shader_subgroup_basic', 'GL_KHR_shader_subgroup_shuffle', 'GL_KHR_shader_subgroup_ballot', 'GL_EXT_samplerless_texture_functions'):
|
||||
w(f'#extension {e} : require')
|
||||
bx = 16 if self.maxntid[0] == 256 else self.maxntid[0]
|
||||
by = self.maxntid[0] // bx
|
||||
w(f'layout(local_size_x = {bx}, local_size_y = {by}) in;')
|
||||
w(f'layout(set = 0, binding = 0, std430) readonly buffer ParamBlock {{ uint w[{(self.psize + 3) // 4}]; }} param;')
|
||||
b = 1
|
||||
for off in sorted(self.textures):
|
||||
w(f'layout(set = 0, binding = {b}) uniform sampler2D T{off}; // {tex_filter.get(off, "?")}')
|
||||
b += 1
|
||||
for off in sorted(self.surfaces):
|
||||
fmt = 'r16ui' if self.surfaces[off] == 'b' else 'rgba16f'
|
||||
typ = 'uimage2D' if self.surfaces[off] == 'b' else 'image2D'
|
||||
w(f'layout(set = 0, binding = {b}, {fmt}) uniform writeonly {typ} S{off};')
|
||||
b += 1
|
||||
if self.exp_lut:
|
||||
w(f'layout(set = 0, binding = {b}, scalar) readonly buffer ExpLut {{ uint16_t e[65536]; }} explut;')
|
||||
w('uint ex2_lut(uint u) { return uint(explut.e[u & 0xffffu]) | (uint(explut.e[u >> 16]) << 16); }')
|
||||
w(f'shared uint sh[{max(self.shared_bytes // 4, 1)}];')
|
||||
def runtime(self):
|
||||
return RUNTIME
|
||||
|
||||
def emit(self, tex_filter):
|
||||
nb = len(self.blocks)
|
||||
out = []
|
||||
w = out.append
|
||||
self.preamble(w, tex_filter)
|
||||
w(self.runtime())
|
||||
w('void main()\n{')
|
||||
if getattr(self, 'local_bytes', 0):
|
||||
w(f' uint loc[{(self.local_bytes + 3) // 4}];')
|
||||
for g in ('float', 'float16_t', 'f16vec2', 'uint', 'uint64_t', 'bool'):
|
||||
names = sorted(n for n, t in self.regs.items() if t == g)
|
||||
for i in range(0, len(names), 24):
|
||||
init = {'float': '0.0', 'float16_t': '0.0hf', 'f16vec2': 'f16vec2(0.0hf)', 'uint': '0u', 'uint64_t': '0ul', 'bool': 'false'}[g]
|
||||
w(f' {"precise " if (g in ("float", "float16_t", "f16vec2") or getattr(self, "precise_all", False) and g != "bool") and not self.fast else ""}{g} ' + ', '.join(f'{n} = {init}' for n in names[i:i + 24]) + ';')
|
||||
# control flow graph over the blocks; END is block nb
|
||||
succ = []
|
||||
for i, blk in enumerate(self.blocks):
|
||||
br = blk[-1] if blk and isinstance(blk[-1], tuple) else None
|
||||
if br is None:
|
||||
succ.append([min(i + 1, nb)])
|
||||
else:
|
||||
tgt = nb if br[2][0] == 'END' else self.labels[br[2][1]]
|
||||
succ.append([tgt] if br[1] is None else [tgt, i + 1])
|
||||
# ---- structured emission: if/else arms that rejoin at the branch's immediate post-dominator --------------
|
||||
END = nb
|
||||
# reachable nodes, back edges (to a node on the DFS stack)
|
||||
color, back = {}, set()
|
||||
st = [(0, iter(succ[0]))]; color[0] = 1
|
||||
while st:
|
||||
n, it = st[-1]
|
||||
for m in it:
|
||||
if m == END:
|
||||
continue
|
||||
if color.get(m) == 1:
|
||||
back.add((n, m))
|
||||
elif m not in color:
|
||||
color[m] = 1; st.append((m, iter(succ[m]))); break
|
||||
else:
|
||||
color[n] = 2; st.pop()
|
||||
nodes = sorted(color)
|
||||
# each loop header h gets a synthetic "continue" node: back edges go there, and it flows to the loop's exit
|
||||
headers = sorted({h for _, h in back})
|
||||
cont = {h: END + 1 + i for i, h in enumerate(headers)}
|
||||
g = {n: [(cont[m] if (n, m) in back else m) for m in succ[n]] for n in nodes}
|
||||
g[END] = []
|
||||
for h in headers:
|
||||
g[cont[h]] = [END] # provisional, replaced below
|
||||
def postorder():
|
||||
rpo, seen = [], set()
|
||||
def dfs(n):
|
||||
stack = [(n, iter(g[n]))]; seen.add(n)
|
||||
while stack:
|
||||
x, it = stack[-1]
|
||||
for m in it:
|
||||
if m not in seen:
|
||||
seen.add(m); stack.append((m, iter(g[m]))); break
|
||||
else:
|
||||
rpo.append(x); stack.pop()
|
||||
dfs(0)
|
||||
return rpo
|
||||
# A loop may leave through its header as well as through its latches. Its body is the set of nodes that
|
||||
# reach a latch without passing the header; the loop region ends at the nearest node that post-dominates
|
||||
# every exit of the body (for a loop with one latch exit this is that exit, as before).
|
||||
preds = {}
|
||||
for n in nodes:
|
||||
for m in succ[n]:
|
||||
preds.setdefault(m, []).append(n)
|
||||
pdom = {END: {END}}
|
||||
for n in postorder(): # successors first
|
||||
if n != END:
|
||||
pdom[n] = {n} | set.intersection(*(pdom[m] for m in g[n]))
|
||||
for h in headers:
|
||||
body, work = {h}, [n for n, m in back if m == h]
|
||||
while work:
|
||||
n = work.pop()
|
||||
if n not in body:
|
||||
body.add(n); work += preds.get(n, [])
|
||||
exits = {m for n in body for m in succ[n] if m not in body}
|
||||
if not exits:
|
||||
raise Err('loop without an exit')
|
||||
common = set.intersection(*(pdom[e] for e in exits))
|
||||
g[cont[h]] = [max(common, key=lambda c: len(pdom[c]))]
|
||||
rpo = postorder()
|
||||
topo = rpo[::-1] # entry first
|
||||
index = {n: i for i, n in enumerate(topo)}
|
||||
ipdom = {END: END}
|
||||
def inter(x, y):
|
||||
while x != y:
|
||||
while index[x] < index[y]:
|
||||
x = ipdom[x]
|
||||
while index[y] < index[x]:
|
||||
y = ipdom[y]
|
||||
return x
|
||||
for n in reversed(topo): # END first
|
||||
if n == END:
|
||||
continue
|
||||
d = None
|
||||
for m in g[n]:
|
||||
d = m if d is None else inter(d, m)
|
||||
ipdom[n] = d
|
||||
emitted = [0]
|
||||
def body(n, ind):
|
||||
if n > END:
|
||||
w(ind + f'again_{n - END - 1} = true;'); return
|
||||
for line in self.blocks[n]:
|
||||
if not isinstance(line, tuple):
|
||||
w(ind + line)
|
||||
emitted[0] += 1
|
||||
def region(n, stop, ind):
|
||||
while n != stop:
|
||||
if n in cont and (n, 'open') not in region.active:
|
||||
# loop: body once per iteration until no back edge is taken
|
||||
e = g[cont[n]][0]
|
||||
w(ind + f'bool again_{cont[n] - END - 1};'); w(ind + 'do\n' + ind + '{')
|
||||
w(ind + f' again_{cont[n] - END - 1} = false;')
|
||||
region.active.add((n, 'open'))
|
||||
region(n, e, ind + ' ')
|
||||
region.active.discard((n, 'open'))
|
||||
w(ind + f'}} while (again_{cont[n] - END - 1});')
|
||||
n = e
|
||||
continue
|
||||
body(n, ind)
|
||||
sc = g[n]
|
||||
if len(sc) == 1:
|
||||
n = sc[0]; continue
|
||||
br = self.blocks[n][-1]
|
||||
m = ipdom[n]
|
||||
t_empty, f_empty = sc[0] == m, sc[1] == m
|
||||
if t_empty and f_empty:
|
||||
pass
|
||||
elif f_empty:
|
||||
w(ind + f'if ({br[1]})\n' + ind + '{'); region(sc[0], m, ind + ' '); w(ind + '}')
|
||||
elif t_empty:
|
||||
w(ind + f'if (!({br[1]}))\n' + ind + '{'); region(sc[1], m, ind + ' '); w(ind + '}')
|
||||
else:
|
||||
w(ind + f'if ({br[1]})\n' + ind + '{'); region(sc[0], m, ind + ' ')
|
||||
w(ind + '}\n' + ind + 'else\n' + ind + '{'); region(sc[1], m, ind + ' '); w(ind + '}')
|
||||
n = m
|
||||
region.active = set()
|
||||
import sys as _s
|
||||
_s.setrecursionlimit(20000)
|
||||
region(0, END, ' ')
|
||||
w('}')
|
||||
sys.stderr.write(f'reachable blocks={len(nodes)} emitted block copies={emitted[0]} loops={len(headers)}\n')
|
||||
return '\n'.join(out) + '\n'
|
||||
|
||||
|
||||
RUNTIME = r'''
|
||||
layout(buffer_reference, scalar, buffer_reference_align = 2) readonly buffer R16 { uint16_t v[4]; };
|
||||
layout(buffer_reference, scalar, buffer_reference_align = 2) writeonly buffer W16 { uint16_t v[4]; };
|
||||
float16_t H(uint u) { return uint16BitsToFloat16(uint16_t(u)); }
|
||||
uint HU(float16_t h) { return uint(float16BitsToUint16(h)); }
|
||||
f16vec2 H2(uint u) { return uint16BitsToFloat16(u16vec2(uint16_t(u & 0xffffu), uint16_t(u >> 16))); }
|
||||
uint HU2_(u16vec2 u) { return uint(u.x) | (uint(u.y) << 16); }
|
||||
uint H2U(f16vec2 h) { return HU2_(float16BitsToUint16(h)); }
|
||||
int s16(uint u) { return (int(u) << 16) >> 16; }
|
||||
float rcp_a(float x) { return 1.0 / x; }
|
||||
// x + copysign(0.5, x) rounded toward zero and then truncated: round half away from zero, exactly
|
||||
float round_half_away(float x) { float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; }
|
||||
uint PARAM32(uint o) { return param.w[o >> 2]; }
|
||||
uint64_t PARAM64(uint o) { return uint64_t(param.w[o >> 2]) | (uint64_t(param.w[(o >> 2) + 1u]) << 32); }
|
||||
uint PARAM16(uint o) { return (param.w[o >> 2] >> ((o & 2u) * 8u)) & 0xffffu; }
|
||||
uint PARAM8(uint o) { return (param.w[o >> 2] >> ((o & 3u) * 8u)) & 0xffu; }
|
||||
uint SH32(uint a) { return sh[a >> 2]; }
|
||||
uint SH16(uint a) { return (sh[a >> 2] >> ((a & 2u) * 8u)) & 0xffffu; }
|
||||
uint G16(uint64_t a) { return uint(R16(a).v[0]); }
|
||||
uint G32(uint64_t a) { R16 r = R16(a); return uint(r.v[0]) | (uint(r.v[1]) << 16); }
|
||||
uint bfi32(uint ins, uint base, uint pos, uint len)
|
||||
{
|
||||
pos &= 0xffu; len &= 0xffu;
|
||||
if (pos >= 32u) return base;
|
||||
uint mask = len >= 32u ? (0xffffffffu << pos) : (((1u << len) - 1u) << pos);
|
||||
return (~mask & base) | (mask & (ins << pos));
|
||||
}
|
||||
// prmt.b32 (default mode): byte i of the result is byte (selector >> 4i) & 7 of the eight bytes b:a
|
||||
uint prmt(uint a, uint b, uint sel)
|
||||
{
|
||||
uint r = 0u;
|
||||
for (uint i = 0u; i < 4u; ++i) { const uint k = (sel >> (4u * i)) & 7u; r |= (((k < 4u ? a : b) >> (8u * (k & 3u))) & 0xffu) << (8u * i); }
|
||||
return r;
|
||||
}
|
||||
// shfl.sync.bfly within a 32-thread warp (ZLUDA's SHFL_SYNC_IMPL)
|
||||
uint shfl_bfly(uint v, int delta, uint opts, out bool in_bounds)
|
||||
{
|
||||
int section_mask = int((opts >> 8) & 31u), warp_end = int(opts & 31u);
|
||||
int self = int(gl_SubgroupInvocationID & 31u);
|
||||
int sub_end = (section_mask & self) | (~section_mask & warp_end);
|
||||
int idx = self ^ (delta & 31);
|
||||
bool oob = idx > sub_end;
|
||||
if (oob) idx = self;
|
||||
in_bounds = !oob;
|
||||
return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));
|
||||
}
|
||||
void sust_p(writeonly image2D s, int x, int y, vec4 v)
|
||||
{
|
||||
ivec2 n = imageSize(s);
|
||||
if (uint(x) < uint(n.x) && uint(y) < uint(n.y)) imageStore(s, ivec2(x, y), v);
|
||||
}
|
||||
// byte-addressed 32-bit store into a 16-bit single-channel surface: two texels
|
||||
void sust_b(writeonly uimage2D s, int xb, int y, uint v)
|
||||
{
|
||||
ivec2 n = imageSize(s);
|
||||
int x = xb >> 1;
|
||||
if (uint(y) < uint(n.y))
|
||||
{
|
||||
if (uint(x) < uint(n.x)) imageStore(s, ivec2(x, y), uvec4(v & 0xffffu));
|
||||
if (uint(x + 1) < uint(n.x)) imageStore(s, ivec2(x + 1, y), uvec4(v >> 16));
|
||||
}
|
||||
}
|
||||
'''
|
||||
|
||||
def build(make, lut, fast, no_half):
|
||||
t = make()
|
||||
t.exp_lut = lut
|
||||
if not no_half:
|
||||
t.collect = True
|
||||
t.translate()
|
||||
parent = {}
|
||||
def find(x):
|
||||
while parent.setdefault(x, x) != x:
|
||||
parent[x] = parent[parent[x]]; x = parent[x]
|
||||
return x
|
||||
for x, y in t.links:
|
||||
if t.width.get(x) == t.width.get(y):
|
||||
parent[find(x)] = find(y)
|
||||
else:
|
||||
t.ctx.setdefault(x, set()).add('u'); t.ctx.setdefault(y, set()).add('u')
|
||||
groups = {}
|
||||
for x in set(t.ctx) | set(parent):
|
||||
groups.setdefault(find(x), set()).update(t.ctx.get(x, set()))
|
||||
half = {x for x in set(t.ctx) | set(parent) if x in t.width and groups[find(x)] == {'h'}}
|
||||
sys.stderr.write(f'half-typed registers: {len(half)} of {len(t.width)}\n')
|
||||
lut2 = t.exp_lut
|
||||
t = make(); t.exp_lut = lut2; t.half = half
|
||||
t.fast = fast
|
||||
t.translate()
|
||||
return t
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
filt = {}
|
||||
args = sys.argv[1:]
|
||||
while '--tex' in args:
|
||||
i = args.index('--tex'); k, v = args[i + 1].split(':'); filt[int(k)] = v; del args[i:i + 2]
|
||||
lut = '--exp-lut' in args
|
||||
if lut:
|
||||
args.remove('--exp-lut')
|
||||
fast = '--fast' in args
|
||||
if fast:
|
||||
args.remove('--fast')
|
||||
src = open(args[0]).read()
|
||||
t = build(lambda: T(src, args[1]), lut, fast, '--no-half-types' in args)
|
||||
sys.stdout.write(t.emit(filt))
|
||||
sys.stderr.write(f'blocks={len(t.blocks)} textures={sorted(t.textures)} surfaces={t.surfaces} shared={t.shared_bytes}\n')
|
||||
Reference in new issue
Block a user