d4r: official NVIDIA DLSS on AMD GPUs under Linux and Proton
This commit is contained in:
ayo committed
commit
118173ecd7
128 files changed
+31333
No files matched your search
@@ -0,0 +1,210 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Instrument a PTX kernel to report where non-finite values first appear.
|
||||
|
||||
Adds a trailing `.param .u64 d4r_dbg` to the entry and, after every
|
||||
floating-point result written to a %-register (MMA outputs, f16/f16x2
|
||||
arithmetic, conversions, rsqrt/rcp, ...), ORs a flag into
|
||||
d4r_dbg[source_line] when the value is Inf or NaN (all-ones exponent):
|
||||
bit 0 for an f16/f16x2 result, bit 1 for an f32 result.
|
||||
|
||||
With --max, each slot instead records the largest magnitude seen on that
|
||||
line (f16 results as f16 bits, f32 results as f32 bits with bit 31 set to
|
||||
tell them apart), which shows where values grow out of range.
|
||||
|
||||
Replay the instrumented kernel with replay_dlss_layer, passing a zero-filled
|
||||
buffer of 4 * (line_count + 1) bytes as the extra argument, then list the
|
||||
flagged lines with --report.
|
||||
|
||||
usage: instrument_ptx_nonfinite.py [--max] IN.ptx OUT.ptx KERNEL
|
||||
instrument_ptx_nonfinite.py --report [--max] IN.ptx DEBUG_BUFFER
|
||||
"""
|
||||
import re
|
||||
import struct
|
||||
import sys
|
||||
|
||||
# Instructions whose first operand (or brace list) is a floating-point result.
|
||||
F16X2_OPS = re.compile(r"^(add|sub|mul|fma\.rn|min|max|neg|abs)\.f16x2$|^cvt\.rn\.f16x2\.e[45]m[23]x2$")
|
||||
F16_OPS = re.compile(r"^(add|sub|mul|fma\.rn|min|max|neg|abs)\.f16$|^cvt\.rn\.f16\.f(32|64)$")
|
||||
F32_OPS = re.compile(r"^(add|sub|mul|fma\.rn|div\.approx|rsqrt\.approx|rcp\.approx|sqrt\.approx|lg2\.approx|ex2\.approx|cvt\.f32\.f16)(\.ftz)?(\.f32)?$")
|
||||
MMA_F16 = re.compile(r"^mma\.sync\.aligned\.m16n8k(16|32)\.row\.col\.f16\.")
|
||||
|
||||
INSTRUCTION = re.compile(r"^(\s*\{?\s*)(@!?%p\d+\s+)?([a-z][\w.:]*)\s+(.*)$")
|
||||
|
||||
|
||||
MAX_MODE = False
|
||||
|
||||
|
||||
def max_checks(offset, register, kind):
|
||||
"""Record the largest magnitude written by one instruction."""
|
||||
if kind == "f16x2":
|
||||
return [
|
||||
f"and.b32 %d4rdbg_r0, {register}, 32767;",
|
||||
f"shr.u32 %d4rdbg_r1, {register}, 16;",
|
||||
"and.b32 %d4rdbg_r1, %d4rdbg_r1, 32767;",
|
||||
"max.u32 %d4rdbg_r0, %d4rdbg_r0, %d4rdbg_r1;",
|
||||
f"atom.global.max.u32 %d4rdbg_r1, [%d4rdbg_rd0+{offset}], %d4rdbg_r0;",
|
||||
]
|
||||
if kind == "f16":
|
||||
return [
|
||||
f"cvt.u32.u16 %d4rdbg_r0, {register};",
|
||||
"and.b32 %d4rdbg_r0, %d4rdbg_r0, 32767;",
|
||||
f"atom.global.max.u32 %d4rdbg_r1, [%d4rdbg_rd0+{offset}], %d4rdbg_r0;",
|
||||
]
|
||||
return [
|
||||
f"mov.b32 %d4rdbg_r0, {register};",
|
||||
"and.b32 %d4rdbg_r0, %d4rdbg_r0, 2147483647;",
|
||||
"shr.u32 %d4rdbg_r0, %d4rdbg_r0, 1;",
|
||||
"or.b32 %d4rdbg_r0, %d4rdbg_r0, 2147483648;",
|
||||
f"atom.global.max.u32 %d4rdbg_r1, [%d4rdbg_rd0+{offset}], %d4rdbg_r0;",
|
||||
]
|
||||
|
||||
|
||||
def checks_for(line_number, opcode, operands):
|
||||
"""Return PTX lines checking the destination of one instruction."""
|
||||
offset = 4 * line_number
|
||||
if MMA_F16.match(opcode):
|
||||
destination = re.match(r"\{([^}]*)\}", operands)
|
||||
registers = [r.strip() for r in destination.group(1).split(",")] if destination else []
|
||||
kind = "f16x2"
|
||||
else:
|
||||
first = operands.split(",")[0].strip().rstrip(";")
|
||||
registers = [first]
|
||||
if F16X2_OPS.match(opcode):
|
||||
kind = "f16x2"
|
||||
elif F16_OPS.match(opcode):
|
||||
kind = "f16"
|
||||
elif F32_OPS.match(opcode) and opcode.endswith(("f32", "f16")):
|
||||
kind = "f32"
|
||||
else:
|
||||
return []
|
||||
lines = []
|
||||
for register in registers:
|
||||
if not register.startswith("%"):
|
||||
continue # registers declared in nested inline-asm scopes
|
||||
if MAX_MODE:
|
||||
lines += max_checks(offset, register, kind)
|
||||
continue
|
||||
if kind == "f16x2":
|
||||
lines += [
|
||||
f"and.b32 %d4rdbg_r0, {register}, 2080406528;", # 0x7c007c00
|
||||
"and.b32 %d4rdbg_r1, %d4rdbg_r0, 31744;", # 0x7c00
|
||||
"setp.eq.u32 %d4rdbg_p0, %d4rdbg_r1, 31744;",
|
||||
"and.b32 %d4rdbg_r1, %d4rdbg_r0, 2080374784;", # 0x7c000000
|
||||
"setp.eq.or.u32 %d4rdbg_p0, %d4rdbg_r1, 2080374784, %d4rdbg_p0;",
|
||||
f"@%d4rdbg_p0 atom.global.or.b32 %d4rdbg_r1, [%d4rdbg_rd0+{offset}], 1;",
|
||||
]
|
||||
elif kind == "f16":
|
||||
lines += [
|
||||
f"cvt.u32.u16 %d4rdbg_r0, {register};",
|
||||
"and.b32 %d4rdbg_r0, %d4rdbg_r0, 31744;",
|
||||
"setp.eq.u32 %d4rdbg_p0, %d4rdbg_r0, 31744;",
|
||||
f"@%d4rdbg_p0 atom.global.or.b32 %d4rdbg_r1, [%d4rdbg_rd0+{offset}], 1;",
|
||||
]
|
||||
else:
|
||||
lines += [
|
||||
f"mov.b32 %d4rdbg_r0, {register};",
|
||||
"and.b32 %d4rdbg_r0, %d4rdbg_r0, 2139095040;", # 0x7f800000
|
||||
"setp.eq.u32 %d4rdbg_p0, %d4rdbg_r0, 2139095040;",
|
||||
f"@%d4rdbg_p0 atom.global.or.b32 %d4rdbg_r1, [%d4rdbg_rd0+{offset}], 2;",
|
||||
]
|
||||
return lines
|
||||
|
||||
|
||||
def instrument(source, kernel):
|
||||
lines = source.split("\n")
|
||||
output = []
|
||||
in_kernel = False
|
||||
entry_params_open = False
|
||||
body_started = False
|
||||
pending_statement = None # multi-line mma statements
|
||||
for number, line in enumerate(lines, start=1):
|
||||
if re.match(rf"^\.visible \.entry {re.escape(kernel)}\(", line):
|
||||
in_kernel = True
|
||||
entry_params_open = True
|
||||
output.append(line)
|
||||
continue
|
||||
if in_kernel and entry_params_open and line.strip() == ")":
|
||||
output[-1] = output[-1] + ","
|
||||
output.append(".param .u64 d4r_dbg")
|
||||
output.append(line)
|
||||
entry_params_open = False
|
||||
continue
|
||||
if in_kernel and not body_started and line.strip() == "{":
|
||||
output.append(line)
|
||||
output.append(".reg .pred %d4rdbg_p0;")
|
||||
output.append(".reg .b32 %d4rdbg_r<2>;")
|
||||
output.append(".reg .b64 %d4rdbg_rd0;")
|
||||
output.append("ld.param.u64 %d4rdbg_rd0, [d4r_dbg];")
|
||||
output.append("cvta.to.global.u64 %d4rdbg_rd0, %d4rdbg_rd0;")
|
||||
body_started = True
|
||||
continue
|
||||
output.append(line)
|
||||
if not body_started:
|
||||
continue
|
||||
if line.strip() == "}" and pending_statement is None and number == len(lines):
|
||||
in_kernel = False
|
||||
statement = line if pending_statement is None else pending_statement[1] + " " + line
|
||||
start = number if pending_statement is None else pending_statement[0]
|
||||
if ";" not in line:
|
||||
match = INSTRUCTION.match(statement)
|
||||
if match and MMA_F16.match(match.group(3)):
|
||||
pending_statement = (start, statement)
|
||||
continue
|
||||
pending_statement = None
|
||||
match = INSTRUCTION.match(statement)
|
||||
if not match or match.group(2):
|
||||
continue # skip predicated results: the check would run unconditionally
|
||||
checks = checks_for(start, match.group(3), match.group(4))
|
||||
if checks:
|
||||
output.extend(checks)
|
||||
return "\n".join(output), len(lines)
|
||||
|
||||
|
||||
def half_value(bits):
|
||||
return struct.unpack("<e", struct.pack("<H", bits & 0xFFFF))[0]
|
||||
|
||||
|
||||
def report_max(lines, values, limit):
|
||||
rows = [(index, value) for index, value in enumerate(values) if value]
|
||||
print(f"{len(rows)} source lines recorded magnitudes")
|
||||
for index, value in rows[:limit]:
|
||||
if value & 0x80000000:
|
||||
magnitude = struct.unpack("<f", struct.pack("<I", (value & 0x7FFFFFFF) << 1))[0]
|
||||
kind = "f32"
|
||||
else:
|
||||
magnitude = half_value(value)
|
||||
kind = "f16"
|
||||
text = lines[index - 1].strip() if 0 < index <= len(lines) else "?"
|
||||
print(f" line {index:6d} [{kind}] max|x|={magnitude:<12.6g} {text[:90]}")
|
||||
|
||||
|
||||
def report(ptx_path, debug_path, max_mode=False, limit=40):
|
||||
lines = open(ptx_path).read().split("\n")
|
||||
data = open(debug_path, "rb").read()
|
||||
flags = struct.unpack(f"<{len(data) // 4}I", data)
|
||||
if max_mode:
|
||||
report_max(lines, flags, limit)
|
||||
return
|
||||
hits = [(index, flag) for index, flag in enumerate(flags) if flag]
|
||||
print(f"{len(hits)} source lines produced non-finite values")
|
||||
for index, flag in hits[:40]:
|
||||
kind = "+".join(k for bit, k in ((1, "f16"), (2, "f32")) if flag & bit)
|
||||
text = lines[index - 1].strip() if 0 < index <= len(lines) else "?"
|
||||
print(f" line {index:6d} [{kind}] {text[:110]}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
arguments = sys.argv[1:]
|
||||
reporting = "--report" in arguments
|
||||
MAX_MODE = "--max" in arguments
|
||||
arguments = [a for a in arguments if a not in ("--report", "--max")]
|
||||
if reporting and len(arguments) == 2:
|
||||
report(arguments[0], arguments[1], MAX_MODE, limit=400)
|
||||
elif len(arguments) == 3:
|
||||
sys.argv = [sys.argv[0]] + arguments
|
||||
text, count = instrument(open(sys.argv[1]).read(), sys.argv[3])
|
||||
open(sys.argv[2], "w").write(text)
|
||||
print(f"instrumented {sys.argv[3]}: {count} source lines; debug buffer needs {4 * (count + 1)} bytes")
|
||||
else:
|
||||
print(__doc__)
|
||||
sys.exit(2)
|
||||
Reference in new issue
Block a user