5090 GPUDirect RDMA support

auto_patch_libcuda_dmabuf.py

#!/usr/bin/env python3
import argparse
import hashlib
import os
import shutil
import struct
import sys

DEFAULT_SRC = "/usr/lib/x86_64-linux-gnu/libcuda.so.590.48.01"

KNOWN_FLAG_OFFSET = 0x5F14
KNOWN_GENERATION_OFFSET = 0xC60
REQUIRED_EXISTING_BIT = 0x40
BIT_TO_ADD = 0x20


class Candidate:
    def __init__(self, insn_offset, patch_offset, flag_offset, generation_offset,
                 generation_threshold, old_imm):
        self.insn_offset = insn_offset
        self.patch_offset = patch_offset
        self.flag_offset = flag_offset
        self.generation_offset = generation_offset
        self.generation_threshold = generation_threshold
        self.old_imm = old_imm


def sha256_short(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()[:16]


def elf_vaddr_for_offset(data, file_offset):
    if len(data) < 64 or data[:4] != b"\x7fELF":
        return None
    if data[4] != 2 or data[5] != 1:
        return None

    try:
        header = struct.unpack_from("<HHIQQQIHHHHHH", data, 16)
        e_phoff = header[4]
        e_phentsize = header[8]
        e_phnum = header[9]
    except struct.error:
        return None

    for i in range(e_phnum):
        off = e_phoff + i * e_phentsize
        try:
            ph = struct.unpack_from("<IIQQQQQQ", data, off)
        except struct.error:
            return None
        p_type, _p_flags, p_offset, p_vaddr, _p_paddr, p_filesz, _p_memsz, _p_align = ph
        if p_type != 1:
            continue
        if p_offset <= file_offset < p_offset + p_filesz:
            return p_vaddr + (file_offset - p_offset)
    return None


def read_u32le(data, offset):
    return struct.unpack_from("<I", data, offset)[0]


def find_candidates(data):
    candidates = []

    # Look for this x86_64 instruction shape:
    #
    #   0f b6 87 xx xx xx xx    movzbl flag(%rdi),%eax
    #   89 c2                   mov    %eax,%edx
    #   83 ca ii                or     $imm8,%edx
    #   83 bf yy yy yy yy 08    cmpl   $0x8,generation(%rdi)
    #   88 97 xx xx xx xx       mov    %dl,flag(%rdi)
    #
    # In 590.48.01 the immediate is 0x40. Changing it to 0x60 preserves that
    # bit and additionally sets the GPUDirect RDMA/DMA-BUF gate bit 0x20.
    limit = len(data) - 25
    for i in range(max(0, limit)):
        if data[i:i + 3] != b"\x0f\xb6\x87":
            continue
        if data[i + 7:i + 9] != b"\x89\xc2":
            continue
        if data[i + 9:i + 11] != b"\x83\xca":
            continue
        if data[i + 12:i + 14] != b"\x83\xbf":
            continue
        if data[i + 19:i + 21] != b"\x88\x97":
            continue

        flag_offset = read_u32le(data, i + 3)
        old_imm = data[i + 11]
        generation_offset = read_u32le(data, i + 14)
        generation_threshold = data[i + 18]
        stored_flag_offset = read_u32le(data, i + 21)

        if stored_flag_offset != flag_offset:
            continue
        if generation_threshold != 8:
            continue
        if (old_imm & REQUIRED_EXISTING_BIT) == 0:
            continue

        candidates.append(Candidate(
            insn_offset=i,
            patch_offset=i + 11,
            flag_offset=flag_offset,
            generation_offset=generation_offset,
            generation_threshold=generation_threshold,
            old_imm=old_imm,
        ))

    return candidates


def print_candidate(prefix, data, cand):
    insn_vaddr = elf_vaddr_for_offset(data, cand.insn_offset)
    patch_vaddr = elf_vaddr_for_offset(data, cand.patch_offset)
    new_imm = cand.old_imm | BIT_TO_ADD
    layout = "known-layout" if (
        cand.flag_offset == KNOWN_FLAG_OFFSET and
        cand.generation_offset == KNOWN_GENERATION_OFFSET
    ) else "layout-drift"

    print(f"{prefix}:")
    print(f"  instruction file offset: 0x{cand.insn_offset:x}")
    if insn_vaddr is not None:
        print(f"  instruction ELF vaddr:   0x{insn_vaddr:x}")
    print(f"  patch byte file offset:  0x{cand.patch_offset:x}")
    if patch_vaddr is not None:
        print(f"  patch byte ELF vaddr:    0x{patch_vaddr:x}")
    print(f"  flag offset:             0x{cand.flag_offset:x}")
    print(f"  generation offset:       0x{cand.generation_offset:x}")
    print(f"  immediate:               0x{cand.old_imm:02x} -> 0x{new_imm:02x}")
    print(f"  classification:          {layout}")
    print("  expected instruction:    or $imm8,%edx")


def choose_candidate(data, allow_layout_drift):
    candidates = find_candidates(data)
    if not candidates:
        raise RuntimeError("no matching libcuda device-init pattern found")

    for idx, cand in enumerate(candidates):
        print_candidate(f"candidate[{idx}]", data, cand)

    if len(candidates) != 1:
        raise RuntimeError(
            f"found {len(candidates)} candidates; refusing to patch automatically"
        )

    cand = candidates[0]
    if ((cand.flag_offset != KNOWN_FLAG_OFFSET or
         cand.generation_offset != KNOWN_GENERATION_OFFSET) and
            not allow_layout_drift):
        raise RuntimeError(
            "candidate was found, but private struct offsets differ from the "
            "590.48.01 layout; re-run with --allow-layout-drift only after "
            "manual disassembly review"
        )
    return cand


def prepare_target(src, out_dir, in_place, force):
    if in_place:
        return src
    if out_dir is None:
        raise RuntimeError("use --out-dir for a patched copy, or --in-place")

    os.makedirs(out_dir, exist_ok=True)
    dst = os.path.join(out_dir, os.path.basename(src))
    if os.path.exists(dst) and not force:
        raise RuntimeError(f"output exists, use --force: {dst}")
    shutil.copy2(src, dst)

    for link_name in ("libcuda.so.1", "libcuda.so"):
        link = os.path.join(out_dir, link_name)
        if os.path.lexists(link):
            os.unlink(link)
        os.symlink(os.path.basename(dst), link)
    return dst


def patch_file(path, cand):
    new_imm = cand.old_imm | BIT_TO_ADD
    with open(path, "r+b") as f:
        f.seek(cand.patch_offset)
        old = f.read(1)
        if len(old) != 1:
            raise RuntimeError("short read while patching")
        if old[0] != cand.old_imm:
            raise RuntimeError(
                f"byte changed before patch: saw 0x{old[0]:02x}, "
                f"expected 0x{cand.old_imm:02x}"
            )
        if old[0] == new_imm:
            print(f"already patched: {path}")
            return
        f.seek(cand.patch_offset)
        f.write(bytes([new_imm]))
    print(
        f"patched {path}: offset 0x{cand.patch_offset:x} "
        f"0x{cand.old_imm:02x}->0x{new_imm:02x}"
    )


def main():
    parser = argparse.ArgumentParser(
        description="Locate and optionally patch the private libcuda device-init "
                    "instruction that gates GPUDirect RDMA/DMA-BUF support."
    )
    parser.add_argument("--src", default=DEFAULT_SRC,
                        help=f"source libcuda, default: {DEFAULT_SRC}")
    parser.add_argument("--out-dir",
                        help="write a patched copy into this directory")
    parser.add_argument("--in-place", action="store_true",
                        help="patch --src directly; not recommended")
    parser.add_argument("--force", action="store_true",
                        help="overwrite output when using --out-dir")
    parser.add_argument("--dry-run", action="store_true",
                        help="only locate the patch byte")
    parser.add_argument("--allow-layout-drift", action="store_true",
                        help="allow patching when the private struct offsets "
                             "differ from the verified 590.48.01 layout")
    args = parser.parse_args()

    src = os.path.realpath(args.src)
    if not os.path.isfile(src):
        print(f"source not found: {src}", file=sys.stderr)
        return 1
    if args.in_place and args.out_dir:
        print("use either --in-place or --out-dir, not both", file=sys.stderr)
        return 1

    with open(src, "rb") as f:
        data = f.read()

    print(f"source: {src}")
    print(f"sha256: {sha256_short(src)}...")
    try:
        cand = choose_candidate(data, args.allow_layout_drift)
    except RuntimeError as e:
        print(f"error: {e}", file=sys.stderr)
        return 2

    if args.dry_run:
        print("dry-run: no file modified")
        return 0

    try:
        target = prepare_target(src, args.out_dir, args.in_place, args.force)
        patch_file(target, cand)
    except RuntimeError as e:
        print(f"error: {e}", file=sys.stderr)
        return 3

    if args.out_dir:
        print("use with:")
        print(f"  export LD_LIBRARY_PATH={os.path.realpath(args.out_dir)}:${{LD_LIBRARY_PATH:-}}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

666