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)
|