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.
211 lines
8.7 KiB
Python
211 lines
8.7 KiB
Python
#!/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)
|