95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Extract the PTX modules embedded in NVIDIA's nvngx_dlss.dll.
|
|
|
|
usage: extract_dlss_ptx.py NVNGX_DLSS_DLL OUT_DIR
|
|
|
|
Writes OUT_DIR/dlss-NNNN-MM.ptx for PTX entry MM of the NNNN-th CUDA fatbin found in the DLL (in file
|
|
order). The texture-kernel build (kernels/tex) edits some of these modules; the DLL is NVIDIA's and is
|
|
never part of this repository, and neither are the extracted files.
|
|
"""
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
FATBIN_MAGIC = 0xBA55ED50
|
|
|
|
|
|
def lz4_block(src: bytes, size: int) -> bytes:
|
|
dst = bytearray()
|
|
i = 0
|
|
while i < len(src):
|
|
token = src[i]
|
|
i += 1
|
|
literals = token >> 4
|
|
if literals == 15:
|
|
while True:
|
|
b = src[i]
|
|
i += 1
|
|
literals += b
|
|
if b != 255:
|
|
break
|
|
dst += src[i:i + literals]
|
|
i += literals
|
|
if i >= len(src) or len(dst) >= size:
|
|
break
|
|
offset = src[i] | (src[i + 1] << 8)
|
|
i += 2
|
|
match = token & 15
|
|
if match == 15:
|
|
while True:
|
|
b = src[i]
|
|
i += 1
|
|
match += b
|
|
if b != 255:
|
|
break
|
|
match += 4
|
|
for _ in range(match):
|
|
dst.append(dst[-offset])
|
|
return bytes(dst[:size])
|
|
|
|
|
|
def fatbins(data: bytes):
|
|
"""Yields (offset, bytes) of every fatbin container (magic, version 1) in file order."""
|
|
magic = struct.pack('<I', FATBIN_MAGIC)
|
|
start = data.find(magic)
|
|
while start >= 0:
|
|
_, version, header_size, size = struct.unpack_from('<IHHQ', data, start)
|
|
if version == 1 and header_size == 16 and start + header_size + size <= len(data):
|
|
yield start, data[start:start + header_size + size]
|
|
start = data.find(magic, start + header_size + size)
|
|
else:
|
|
start = data.find(magic, start + 4)
|
|
|
|
|
|
def ptx_entries(fatbin: bytes):
|
|
"""Yields the (decompressed) PTX text of each PTX entry of one fatbin."""
|
|
_, _, header_size, size = struct.unpack_from('<IHHQ', fatbin, 0)
|
|
offset = header_size
|
|
while offset < header_size + size:
|
|
kind, _, entry_header, entry_size = struct.unpack_from('<HHIQ', fatbin, offset)
|
|
flags = struct.unpack_from('<Q', fatbin, offset + 40)[0]
|
|
decompressed = struct.unpack_from('<Q', fatbin, offset + 56)[0] if entry_header >= 64 else 0
|
|
payload = fatbin[offset + entry_header:offset + entry_header + entry_size]
|
|
if kind == 1: # PTX
|
|
yield lz4_block(payload, decompressed) if (flags & 0x2000 and decompressed) else payload
|
|
offset += entry_header + entry_size
|
|
|
|
|
|
def main(argv):
|
|
if len(argv) != 3:
|
|
sys.exit(__doc__)
|
|
data = open(argv[1], 'rb').read()
|
|
os.makedirs(argv[2], exist_ok=True)
|
|
total = 0
|
|
for index, (_, fatbin) in enumerate(fatbins(data)):
|
|
for entry, ptx in enumerate(ptx_entries(fatbin)):
|
|
ptx = ptx.rstrip(b'\0')
|
|
with open(os.path.join(argv[2], f'dlss-{index:04d}-{entry:02d}.ptx'), 'wb') as out:
|
|
out.write(ptx)
|
|
total += 1
|
|
print(f'{total} PTX modules written to {argv[2]}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main(sys.argv)
|