ayo
7c5d3c9157
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.
109 lines
4.5 KiB
Python
109 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Print the expression tree that defines a PTX register, followed back through arithmetic, pair packs/unpacks,
|
|
conversions and warp shuffles down to loads, constants and registers named as leaves.
|
|
|
|
usage: ptx_expr.py STAGE LINE %REG [--depth N] [--stop %r1,%r2]
|
|
|
|
LINE is a line at or after the register's definition (the search goes backwards from it). Brace-scoped temporaries
|
|
(`{ .reg .b16 hl, hu; ... }`) are resolved within their block. Shuffles print as `shfl.bfly^k(expr)`: the same
|
|
expression on the lane `laneid ^ k`.
|
|
"""
|
|
import argparse
|
|
import re
|
|
from pathlib import Path
|
|
|
|
|
|
def module_of(kernel):
|
|
for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')):
|
|
if f'.entry {kernel}(' in p.read_text(errors='replace'):
|
|
return p
|
|
|
|
|
|
class Tracer:
|
|
def __init__(self, lines, start):
|
|
self.lines = lines
|
|
self.start = start
|
|
self.consts = {}
|
|
self.cache = {}
|
|
|
|
def find_def(self, reg, before, block_lo):
|
|
"""the statement (line index) that last writes REG before line BEFORE"""
|
|
pat = re.compile(r'(^|[\s{,|])' + re.escape(reg) + r'([\s,|}]|$)')
|
|
lo = block_lo if not reg.startswith('%') else self.start
|
|
for i in range(before - 1, lo - 1, -1):
|
|
s = re.sub(r'//.*', '', self.lines[i]).strip().lstrip('{').strip()
|
|
s = re.sub(r'^@!?%p\d+\s+', '', s)
|
|
if not s or s.startswith(('.', '$', '}')):
|
|
continue
|
|
parts = s.split(None, 1)
|
|
if len(parts) < 2:
|
|
continue
|
|
rest = parts[1]
|
|
dsts = ([x.strip() for x in rest[1:rest.index('}')].split(',')] if rest.startswith('{')
|
|
else [x.strip() for x in rest.split(',')[0].split('|')])
|
|
if reg in dsts and not parts[0].startswith(('st.', 'sust', 'bra')):
|
|
return i, parts[0], rest.rstrip(';').strip(), dsts
|
|
return None
|
|
|
|
def block_start(self, line):
|
|
depth = 0
|
|
for i in range(line, self.start, -1):
|
|
depth += self.lines[i].count('}') - self.lines[i].count('{')
|
|
if depth < 0:
|
|
return i
|
|
return self.start
|
|
|
|
def expr(self, reg, before, depth, stop, half=None):
|
|
reg = reg.strip()
|
|
if re.fullmatch(r'-?\d+|0[fFdD][0-9A-Fa-f]+|0x[0-9a-fA-F]+', reg):
|
|
return reg
|
|
if reg in stop or depth <= 0:
|
|
return reg + (f'.{half}' if half is not None else '')
|
|
d = self.find_def(reg, before, self.block_start(before - 1))
|
|
if d is None:
|
|
return reg + '?'
|
|
i, op, rest, dsts = d
|
|
args = [a.strip() for a in re.split(r',(?![^{]*\})', rest)]
|
|
srcs = args[1:] if not rest.startswith('{') else args[1:]
|
|
e = lambda r, h=None: self.expr(r, i, depth - 1, stop, h)
|
|
if op.startswith('mov.b32') and rest.startswith('{'): # unpack: {lo, hi} <- r
|
|
which = 'lo' if dsts.index(reg) == 0 else 'hi'
|
|
return e(srcs[0], which)
|
|
if op.startswith('mov.b32') and srcs and srcs[0].startswith('{'): # pack r <- {lo, hi}
|
|
lo, hi = [x.strip() for x in srcs[0][1:-1].split(',')]
|
|
if half == 'lo':
|
|
return e(lo)
|
|
if half == 'hi':
|
|
return e(hi)
|
|
return f'pack({e(lo)}, {e(hi)})'
|
|
if op.startswith(('mov.', 'cvt.')) and len(srcs) == 1:
|
|
inner = e(srcs[0], half)
|
|
return inner if op.startswith('mov.') else f'{op.split(".")[0]}.{op.split(".")[-2]}({inner})'
|
|
if op.startswith('shfl.sync.bfly'):
|
|
return f'shfl^{srcs[1]}({e(srcs[0], half)})'
|
|
if op.startswith('ld.'):
|
|
return f'load@L{i + 1}[{srcs[0]}]' + (f'.{half}' if half else '')
|
|
name = op.split('.')[0]
|
|
if op.endswith('x2') and half:
|
|
return f'{name}(' + ', '.join(e(s, half) for s in srcs) + ')'
|
|
return f'{name}(' + ', '.join(e(s, half) for s in srcs) + ')'
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('stage')
|
|
ap.add_argument('line', type=int)
|
|
ap.add_argument('reg')
|
|
ap.add_argument('--depth', type=int, default=30)
|
|
ap.add_argument('--stop', default='')
|
|
o = ap.parse_args()
|
|
kernel = f'cuda_dldn_engine_swin_{o.stage}_kernel'
|
|
lines = module_of(kernel).read_text().replace('\r', '').split('\n')
|
|
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
|
|
t = Tracer(lines, k)
|
|
print(t.expr(o.reg, o.line, o.depth, set(filter(None, o.stop.split(',')))))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|