Everything I had to do to get p2p setup working on 5060tis, ubuntu

nvidia-p2p-setup.sh

#!/usr/bin/env bash
# nvidia-p2p-setup.sh
# Idempotently install the aikitoria/open-gpu-kernel-modules P2P-patched NVIDIA
# driver (610.57.04-p2p-v2) on this host.
#
# Re-running this script is safe: every step checks whether it has already
# been done and skips if so. Steps that cannot be safely re-run (driver
# install, module swap) are guarded by version/state checks.
#
# Usage:
#   sudo bash nvidia-p2p-setup.sh            # full run
#   sudo bash nvidia-p2p-setup.sh --no-reboot # skip the final reboot
#   sudo bash nvidia-p2p-setup.sh --step grub # run only a named step
#   sudo bash nvidia-p2p-setup.sh --dry-run    # print actions, don't execute
#
# Steps (each idempotent):
#   precheck   - verify env, headers, tools, secure boot, GPU usage
#   download   - fetch the 610.57.04 .run installer
#   purge      - remove apt-managed 595 driver packages
#   install    - install 610.57.04 userspace (no kernel modules)
#   grub       - add amd_iommu=on iommu=pt to GRUB_CMDLINE_LINUX_DEFAULT
#   headless   - add nomodeset + video=...:off (keep nvidia-drm modeset off; simpledrm still loads)
#   build      - clone repo and build kernel modules
#   swap       - unload old modules, install + load new ones
#   reboot     - reboot (unless --no-reboot)

set -euo pipefail

# ---------- config ----------
DRIVER_VERSION="610.57.04"
REPO_URL="https://github.com/aikitoria/open-gpu-kernel-modules.git"
REPO_BRANCH="${DRIVER_VERSION}-p2p-v2"
WORK_DIR="/home/gary/dev/system"
REPO_DIR="${WORK_DIR}/open-gpu-kernel-modules"
# Use the invoking user's home (sudo resets HOME to /root on Ubuntu):
REAL_HOME="$(getent passwd "${SUDO_USER:-$USER}" | cut -d: -f6)"
DOWNLOAD_DIR="${REAL_HOME}/downloads"
RUN_INSTALLER="${DOWNLOAD_DIR}/NVIDIA-Linux-x86_64-${DRIVER_VERSION}.run"
# Standard NVIDIA direct-download URL pattern:
RUN_INSTALLER_URL="https://us.download.nvidia.com/XFree86/Linux-x86_64/${DRIVER_VERSION}/NVIDIA-Linux-x86_64-${DRIVER_VERSION}.run"
DRIVER_DETAILS_URL="https://www.nvidia.com/en-us/drivers/details/274513/"
IOMMU_PARAMS="amd_iommu=on iommu=pt"
HEADLESS_PARAMS="nomodeset video=efifb:off video=vesafb:off video=simplefb:off initcall_blacklist=sysfb_init"
GRUB_FILE="/etc/default/grub"
GRUB_BACKUP="${GRUB_FILE}.bak.$(date +%Y%m%d-%H%M%S)"
APT_PIN_FILE="/etc/apt/preferences.d/no-ubuntu-nvidia.pref"
STEPS_ORDER=(precheck download purge install grub headless build modprobe swap reboot)

# ---------- args ----------
NO_REBOOT=0
DRY_RUN=0
ONLY_STEP=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --no-reboot) NO_REBOOT=1; shift ;;
    --dry-run) DRY_RUN=1; shift ;;
    --step) ONLY_STEP="$2"; shift 2 ;;
    --step=*) ONLY_STEP="${1#--step=}"; shift ;;
    -h|--help)
      sed -n '2,28p' "$0" | sed 's/^# \{0,1\}//'
      exit 0 ;;
    *) echo "unknown arg: $1" >&2; exit 2 ;;
  esac
done

# ---------- helpers ----------
log()  { printf '\033[1;34m[%s]\033[0m %s\n' "$(date +%H:%M:%S)" "$*"; }
ok()   { printf '\033[1;32m[ok]\033[0m %s\n' "$*"; }
warn() { printf '\033[1;33m[warn]\033[0m %s\n' "$*" >&2; }
err()  { printf '\033[1;31m[err]\033[0m %s\n' "$*" >&2; }
skip() { printf '\033[0;90m[skip]\033[0m %s\n' "$*"; }

run() {
  if [[ "$DRY_RUN" -eq 1 ]]; then
    printf '\033[0;90m[dry]\033[0m %s\n' "$*"
  else
    "$@"
  fi
}
# run a shell snippet under dry-run
runsh() {
  if [[ "$DRY_RUN" -eq 1 ]]; then
    printf '\033[0;90m[dry]\033[0m %s\n' "$1"
  else
    bash -c "$1"
  fi
}

step_enabled() {
  if [[ -n "$ONLY_STEP" ]]; then
    [[ "$1" == "$ONLY_STEP" ]]
  else
    return 0
  fi
}

module_loaded() { lsmod | awk '{print $1}' | grep -qx "$1"; }
apt_pkg_installed() { dpkg -s "$1" &>/dev/null; }
# GRUB_CMDLINE_LINUX_DEFAULT accessor (quotes stripped). NB: never use
# awk -F= '{print $2}' here -- it splits on every '=', and params like
# iommu=pt contain '=', so $2 is just the fragment up to the next '=' (e.g.
# pcie_aspm). That bug made the old guard always false and appended
# amd_iommu=on iommu=pt on every run (9x on this host).
grub_cmdline_default() {
  sed -n 's/^GRUB_CMDLINE_LINUX_DEFAULT=//p' "$GRUB_FILE" 2>/dev/null | tr -d '"'
}

# True iff every token in $@ is a whole word in GRUB_CMDLINE_LINUX_DEFAULT
# (whole-word so iommu=pt won't match inside intel_iommu=pt_dbg).
grub_cmdline_has() {
  local cur tok
  cur=$(grub_cmdline_default) || return 1
  [[ -n "$cur" ]] || return 1
  for tok in "$@"; do
    [[ " $cur " == *" $tok "* ]] || return 1
  done
  return 0
}

grub_has_iommu_pt() { grub_cmdline_has amd_iommu=on iommu=pt; }
grub_has_headless() {
  # true iff every HEADLESS_PARAMS token is a whole word in the cmdline
  local tok
  for tok in $HEADLESS_PARAMS; do
    grub_cmdline_has "$tok" || return 1
  done
  return 0
}

# Append $1 before the closing quote of GRUB_CMDLINE_LINUX_DEFAULT. Caller
# guards idempotency. Temp file + cp preserves grub mode and avoids sed escaping.
grub_append_default() {
  local add="$1" tmp
  tmp=$(mktemp) || return 1
  if ! awk -v add="$add" -v q='"' '/^GRUB_CMDLINE_LINUX_DEFAULT=/ {sub(/"$/, " " add q)} 1' "$GRUB_FILE" > "$tmp"; then
    rm -f "$tmp"; return 1
  fi
  run cp "$tmp" "$GRUB_FILE"
  rm -f "$tmp"
}

# Collapse repeated amd_iommu=on iommu=pt runs (left by the pre-fix bug) to one.
# Returns 0 if $GRUB_FILE changed, else 1.
grub_dedup_iommu() {
  local cur new tmp
  cur=$(grub_cmdline_default) || return 1
  [[ -n "$cur" ]] || return 1
  new=$(printf '%s' "$cur" | sed -E ':a; s/(amd_iommu=on iommu=pt)( amd_iommu=on iommu=pt)+/\1/g; ta; s/  +/ /g; s/^ //; s/ $//')
  [[ "$cur" != "$new" ]] || return 1
  tmp=$(mktemp) || return 1
  if ! awk -v new="$new" -v q='"' '/^GRUB_CMDLINE_LINUX_DEFAULT=/ {print "GRUB_CMDLINE_LINUX_DEFAULT=" q new q; next} 1' "$GRUB_FILE" > "$tmp"; then
    rm -f "$tmp"; return 1
  fi
  run cp "$tmp" "$GRUB_FILE"
  rm -f "$tmp"
  return 0
}
run_installer_valid() {
  # Require: exists, at least 50MB (catches truncated downloads), and a shell
  # self-extracting shebang (NVIDIA .run files are makeself archives starting
  # with #!/bin/sh). ELF magic alone passes for any partial download.
  [[ -f "$RUN_INSTALLER" ]] \
    && [[ $(stat -c%s "$RUN_INSTALLER") -gt 52428800 ]] \
    && [[ $(head -c 2 "$RUN_INSTALLER") == '#!' ]]
}

# validate --step name against the known list (after helpers are defined)
if [[ -n "$ONLY_STEP" ]]; then
  valid=0
  for s in "${STEPS_ORDER[@]}"; do
    if [[ "$s" == "$ONLY_STEP" ]]; then valid=1; break; fi
  done
  if [[ $valid -eq 0 ]]; then
    err "unknown step: '$ONLY_STEP'. valid steps: ${STEPS_ORDER[*]}"
    exit 2
  fi
fi

# ---------- step: precheck ----------
step_precheck() {
  log "precheck"
  local fail=0

  # root
  if [[ $EUID -ne 0 ]]; then
    err "must be run as root (use sudo)"; exit 1
  fi

  # kernel headers for running kernel
  if [[ ! -d "/usr/src/linux-headers-$(uname -r)" ]]; then
    warn "linux-headers-$(uname -r) not found; build will fail."
    warn "  install with: sudo apt install linux-headers-$(uname -r)"
    fail=1
  else
    ok "linux-headers-$(uname -r) present"
  fi

  # build tools
  local t missing=()
  for t in make gcc ld objcopy; do
    command -v "$t" &>/dev/null || missing+=("$t")
  done
  if [[ ${#missing[@]} -gt 0 ]]; then
    warn "missing build tools: ${missing[*]}"
    warn "  install with: sudo apt install build-essential"
    fail=1
  else
    ok "build tools present"
  fi

  # secure boot (fail-fast: unsigned modules cannot load under Secure Boot)
  if command -v mokutil &>/dev/null; then
    if mokutil --sb-state 2>/dev/null | grep -qi enabled; then
      err "Secure Boot is ENABLED. Hand-built modules are unsigned and will be"
      err "  refused at modprobe. Disable Secure Boot in firmware, or sign the"
      err "  modules with an enrolled MOK before running this script."
      fail=1
    else
      ok "Secure Boot disabled"
    fi
  else
    warn "mokutil not installed; cannot check Secure Boot state."
  fi

  # CPU vendor -> IOMMU param sanity
  local vendor
  vendor=$(awk '/^vendor_id/ {print $3; exit}' /proc/cpuinfo)
  if [[ "$vendor" == "AuthenticAMD" ]]; then
    ok "AMD CPU -> using amd_iommu=on iommu=pt"
  elif [[ "$vendor" == "GenuineIntel" ]]; then
    warn "Intel CPU detected. This script's IOMMU params are tuned for AMD."
    warn "  Edit IOMMU_PARAMS to 'intel_iommu=on iommu=pt' and re-run."
    fail=1
  fi

  # GPU in use?
  if command -v nvidia-smi &>/dev/null; then
    if nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | grep -q .; then
      warn "GPU has running compute processes. They must be stopped before the"
      warn "  swap step (rmmod will fail while a GPU handle is held)."
      nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
    else
      ok "no GPU compute processes"
    fi
  fi

  # display manager (would block rmmod)
  if systemctl status display-manager &>/dev/null 2>&1; then
    warn "display-manager is active; it may hold the GPU. Stop it before swap:"
    warn "  sudo systemctl stop <display-manager>"
  else
    ok "headless (no display-manager)"
  fi

  if [[ $fail -ne 0 ]]; then
    err "precheck failed; fix the above and re-run."
    exit 1
  fi
  ok "precheck complete"
}

# ---------- step: download ----------
step_download() {
  log "download"
  mkdir -p "$DOWNLOAD_DIR"

  if run_installer_valid; then
    skip "valid .run installer already present: $RUN_INSTALLER"
    return
  fi

  # try the standard NVIDIA direct URL first
  log "fetching installer from $RUN_INSTALLER_URL"
  if run wget -q -c -O "$RUN_INSTALLER" "$RUN_INSTALLER_URL"; then
    if run_installer_valid; then
      run chmod +x "$RUN_INSTALLER"
      ok "downloaded $RUN_INSTALLER"
      return
    fi
  fi
  run rm -f "$RUN_INSTALLER"

  # fallback: ask the user to download manually
  warn "could not auto-download the .run installer."
  warn "  open $DRIVER_DETAILS_URL"
  warn "  download NVIDIA-Linux-x86_64-${DRIVER_VERSION}.run"
  warn "  place it at $RUN_INSTALLER"
  warn "  then re-run this script."
  exit 1
}

# ---------- step: purge ----------
step_purge() {
  log "purge apt-managed 595 driver"

  # only purge if any 595 nvidia package is still installed
  local installed_595
  installed_595=$(dpkg -l 2>/dev/null \
    | awk '/^ii.*nvidia.*595/ {print $2}' || true)
  if [[ -z "$installed_595" ]] \
     && ! dpkg -l 2>/dev/null | grep -q '^ii.*libnvidia.*595'; then
    skip "no 595 nvidia packages installed"
  else
    log "purging 595 nvidia packages:"
    printf '%s\n' "$installed_595" | sed 's/^/  /'
    runsh "apt-get purge -y 'linux-modules-nvidia-595-open-*' 'nvidia-*595*' 'libnvidia-*595*' 'cuda-*595*' && apt-get autoremove --purge -y"
  fi

  # DKMS leftovers
  if command -v dkms &>/dev/null && dkms status 2>/dev/null | grep -qi nvidia; then
    log "removing leftover DKMS nvidia modules"
    dkms status 2>/dev/null | awk -F: '/nvidia/ {print $1}' | while read -r mod; do
      runsh "dkms remove '$mod' --all || true"
    done
  else
    skip "no DKMS nvidia entries"
  fi

  # blacklist apt nvidia packages so they don't come back on upgrade
  if [[ -f "$APT_PIN_FILE" ]]; then
    skip "apt pin already present: $APT_PIN_FILE"
  else
    log "writing apt pin to block Ubuntu-packaged nvidia drivers"
    runsh "cat > '$APT_PIN_FILE' <<'EOF'
Package: nvidia-driver-*
Pin: release o=Ubuntu
Pin-Priority: -1
EOF"
    ok "apt pin written"
  fi
}

# ---------- step: install ----------
step_install() {
  log "install ${DRIVER_VERSION} userspace (no kernel modules)"

  # idempotency: check if the right userspace is already in place.
  # The .run installer drops versioned libs; checking one is enough.
  local marker_lib="/usr/lib/x86_64-linux-gnu/libcuda.so.${DRIVER_VERSION}"
  local already_installed=0
  if [[ -f "$marker_lib" ]] || [[ -f "/usr/lib/x86_64-linux-gnu/libGL.so.${DRIVER_VERSION}" ]]; then
    skip "userspace ${DRIVER_VERSION} libraries already installed"
    already_installed=1
  fi

  if [[ $already_installed -eq 0 ]]; then
    run_installer_valid || { err "installer missing; run download step first"; exit 1; }

    runsh "\"$RUN_INSTALLER\" \
      --no-kernel-modules \
      --accept-license \
      --no-questions \
      --no-dkms \
      --disable-nouveau"
  fi
  # The .run installer's --disable-nouveau may not write a modprobe blacklist
  # in --no-kernel-modules mode. Ensure nouveau is blacklisted so it doesn't
  # grab GPUs on the next boot (causes "only one GPU comes up" after reboot).
  # This must run even if the installer was skipped (idempotency return above).
  local nouveau_conf="/etc/modprobe.d/blacklist-nouveau.conf"
  if [[ -f "$nouveau_conf" ]] && grep -qx 'blacklist nouveau' "$nouveau_conf" 2>/dev/null; then
    skip "nouveau blacklist already present: $nouveau_conf"
  else
    log "writing nouveau blacklist"
    runsh "cat > '$nouveau_conf' <<'EOF'
blacklist nouveau
blacklist snd_hda_intel
options nouveau modeset=0
EOF"
    runsh "update-initramfs -u -k all"
    ok "nouveau blacklisted and initramfs updated"
  fi
  ok "userspace driver ${DRIVER_VERSION} installed"
}

# ---------- step: grub ----------
step_grub() {
  log "grub: IOMMU passthrough"

  if grub_has_iommu_pt; then
    # Already present -- the pre-fix idempotency bug (awk -F= guard always
    # returned false) may have appended amd_iommu=on iommu=pt many times.
    # Collapse the duplicates to a single occurrence.
    if grub_dedup_iommu; then
      runsh "update-grub"
      ok "collapsed duplicate amd_iommu=on iommu=pt entries in $GRUB_FILE"
    else
      skip "iommu=pt already in $GRUB_FILE (no duplicates)"
    fi
    return
  fi

  # back up once
  if [[ ! -f "${GRUB_FILE}.bak" ]]; then
    run cp "$GRUB_FILE" "${GRUB_FILE}.bak"
    ok "backed up $GRUB_FILE -> ${GRUB_FILE}.bak"
  fi

  # append params to GRUB_CMDLINE_LINUX_DEFAULT (idempotent: guarded above)
  runsh "sed -i '/^GRUB_CMDLINE_LINUX_DEFAULT=/ s|\"\\(.*\\)\"|\"\\1 ${IOMMU_PARAMS}\"|' '$GRUB_FILE'"
  runsh "update-grub"
  ok "added '${IOMMU_PARAMS}' and updated grub"
  log "verify: grep GRUB_CMDLINE_LINUX_DEFAULT $GRUB_FILE"
}

# ---------- step: headless ----------
step_headless() {
  log "headless: ensure nomodeset + initcall_blacklist=sysfb_init + video=...:off in GRUB"

  # Adds the HEADLESS_PARAMS to the kernel cmdline. The critical param is
  # initcall_blacklist=sysfb_init: it prevents sysfb from registering the
  # simple-framebuffer platform device, so simpledrm never binds and no GPU's
  # BAR1 holds the firmware framebuffer at boot. This is REQUIRED for
  # RMForceStaticBar1=1 to work -- without the blacklist, simpledrm maps the
  # UEFI-GOP framebuffer into the lowest-bus GPU's BAR1, and RM's static
  # carve-out overflows the occupied BAR1 VA space -> RmInitAdapter fails for
  # that GPU (kbusEnableStaticBar1Mapping_TU102: Failed ... NV_ERR_INVALID_ARGUMENT).
  # Verified: with the blacklist active, all 3 GPUs init successfully with
  # RMForceStaticBar1=1.
  #
  # nomodeset is kept to prevent nvidia-drm from doing modeset (pure compute
  # headless). The video=...:off params are no-ops here (FB_EFI/FB_VESA unset)
  # but kept for portability.
  #
  # Idempotent per-token: only params not already present are appended, so this
  # is safe to re-run after adding a new param to HEADLESS_PARAMS.

  local missing=() tok
  for tok in $HEADLESS_PARAMS; do
    grub_cmdline_has "$tok" || missing+=("$tok")
  done
  if [[ ${#missing[@]} -eq 0 ]]; then
    skip "all headless params already in $GRUB_FILE"
    return
  fi

  if [[ ! -f "${GRUB_FILE}.bak" ]]; then
    run cp "$GRUB_FILE" "${GRUB_FILE}.bak"
    ok "backed up $GRUB_FILE -> ${GRUB_FILE}.bak"
  fi

  grub_append_default "${missing[*]}"
  runsh "update-grub"
  ok "added '${missing[*]}' and updated grub"
  log "verify after reboot: grep -E 'nomodeset|sysfb_init' /proc/cmdline"
  log "note: with initcall_blacklist=sysfb_init the local console goes fully dark"
  log "      (no simpledrm) -- use SSH/journalctl."
}

# ---------- step: build ----------
step_build() {
  log "build kernel modules"

  # clone (idempotent: only fetch/reset if not already on the right branch)
  if [[ -d "$REPO_DIR/.git" ]]; then
    local cur_branch
    cur_branch=$(git -C "$REPO_DIR" rev-parse --abbrev-ref HEAD 2>/dev/null || echo "")
    if [[ "$cur_branch" == "$REPO_BRANCH" ]]; then
      skip "repo already on $REPO_BRANCH at $REPO_DIR"
    else
      log "switching repo to $REPO_BRANCH"
      runsh "cd '$REPO_DIR' && git fetch --quiet origin && git checkout -q '$REPO_BRANCH' && git reset --hard 'origin/$REPO_BRANCH'"
    fi
  else
    run git clone --branch "$REPO_BRANCH" --depth 1 "$REPO_URL" "$REPO_DIR"
  fi

  # build (idempotent: skip if all four .ko for this driver version exist, are
  # valid, and are newer than every source .c/.h in the tree)
  local ko mod missing_kos stale
  local kos=(
    "$REPO_DIR/kernel-open/nvidia.ko"
    "$REPO_DIR/kernel-open/nvidia-uvm.ko"
    "$REPO_DIR/kernel-open/nvidia-modeset.ko"
    "$REPO_DIR/kernel-open/nvidia-drm.ko"
  )
  missing_kos=()
  for ko in "${kos[@]}"; do
    if [[ ! -f "$ko" ]] || ! modinfo "$ko" &>/dev/null \
       || ! modinfo "$ko" 2>/dev/null | grep -q "^version:.*${DRIVER_VERSION}"; then
      missing_kos+=("$ko")
    fi
  done
  # any source newer than the newest .ko means we need to rebuild
  stale=""
  if [[ ${#missing_kos[@]} -eq 0 ]]; then
    local newest_ko_mtime newest_src_mtime
    # avoid SIGPIPE under pipefail: head -1 closes pipe early on big find output
    newest_ko_mtime=$(stat -c%Y "${kos[@]}" 2>/dev/null | sort -rn | head -1)
    newest_src_mtime=$(find "$REPO_DIR" -type f \( -name '*.c' -o -name '*.h' \) \
      -printf '%T@\n' 2>/dev/null | sort -rn | head -1 | cut -d. -f1 || true)
    [[ -n "$newest_src_mtime" && "$newest_src_mtime" -gt "$newest_ko_mtime" ]] \
      && stale="1" || true
  fi
  if [[ ${#missing_kos[@]} -eq 0 && -z "$stale" ]]; then
    skip "all four .ko already built for ${DRIVER_VERSION} and up to date"
    return
  fi
  [[ ${#missing_kos[@]} -gt 0 ]] \
    && log "rebuilding; missing/stale: ${missing_kos[*]}" || true
  [[ -n "$stale" ]] && log "rebuilding; source newer than built modules" || true

  runsh "cd '$REPO_DIR' && make modules -j\$(nproc) NV_VERBOSE=1"
  ok "modules built"

  # post-build: verify all four .ko are valid and match the driver version
  for ko in "${kos[@]}"; do
    if ! modinfo "$ko" &>/dev/null; then
      err "build finished but $ko is not a valid module"; exit 1
    fi
    if ! modinfo "$ko" 2>/dev/null | grep -q "^version:.*${DRIVER_VERSION}"; then
      err "build finished but $ko version != ${DRIVER_VERSION}"; exit 1
    fi
  done
}

# ---------- step: modprobe ----------
# Writes the NVreg_RegistryDwords that force-enable P2P on the 5060 Ti (GB206),
# which the aikitoria patch alone does not grant (CanAccessPeer=1 but enabling
# fails with "mapping of buffer object failed" without these overrides).
# ForceP2P=17      -> reads|writes|atomics all ENABLE (bits 0/4/9 = 0x11 = 17)
# RMForceP2PType=1 -> TYPE_PCIEP2P
# RMPcieP2PType=1  -> TYPE_BAR1 (force BAR1 P2P instead of mailbox)
# PeerMappingOverride=1 -> allow 3rd-party peer mappings (bug 1630288 WAR)
# RMForceStaticBar1=1  -> ENABLED. Forces a static ~BAR1-sized carve-out at
#   RmInitAdapter so peer VAs stay stable across remaps. NOT required for P2P
#   (PeerMappingOverride=1 handles stable mappings); it's a stability aid.
#   REQUIRES initcall_blacklist=sysfb_init (the headless step adds this).
#   Without the blacklist, simpledrm maps the UEFI-GOP firmware framebuffer into
#   the lowest-bus GPU's BAR1 BEFORE RM loads, and the static carve-out's
#   kbusMapFbApertureSingle overflows the BAR1 VA space -> NV_ERR_INVALID_ARGUMENT
#   (0x1f), RmInitAdapter fails for that GPU. The blacklist prevents
#   simple-framebuffer platform device registration, so simpledrm never binds and
#   the BAR1 stays clean. dmesg markers (failure WITHOUT blacklist):
#   kbusEnableStaticBar1Mapping_TU102: Failed to create the static bar1 mapping
#   offset0x20000000 size 0x3e1000000 + total BAR1 VA range Lo=0x0 Hi=0x3ffffffff
#   + Requested BAR1 VA ... Hi=0x400ffffff. Verified: with the blacklist active,
#   all 3 GPUs (05/06/07) init successfully with RMForceStaticBar1=1.
P2P_REGISTRY_DWORDS="ForceP2P=17;RMForceP2PType=1;RMPcieP2PType=1;PeerMappingOverride=1;RMForceStaticBar1=1"
MODPROBE_CONF="/etc/modprobe.d/nvidia-p2p.conf"

step_modprobe() {
  log "modprobe: write P2P registry dwords"

  local want="options nvidia NVreg_RegistryDwords=\"${P2P_REGISTRY_DWORDS}\""
  if [[ -f "$MODPROBE_CONF" ]] && grep -qx "options nvidia NVreg_RegistryDwords=\"${P2P_REGISTRY_DWORDS}\"" "$MODPROBE_CONF" 2>/dev/null; then
    skip "P2P registry dwords already in $MODPROBE_CONF"
    return
  fi

  runsh "cat > '$MODPROBE_CONF' <<'EOF'
options nvidia NVreg_RegistryDwords=\"${P2P_REGISTRY_DWORDS}\"
EOF"
  ok "wrote $MODPROBE_CONF"
  log "contents:"
  runsh "cat '$MODPROBE_CONF'"
}

# ---------- step: swap ----------
step_swap() {
  log "swap kernel modules"

  # idempotency: if the patched module is already loaded for this version, skip
  if module_loaded nvidia; then
    local loaded_ver on_disk_mod
    loaded_ver=$(cat /proc/driver/nvidia/version 2>/dev/null \
                 | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1 || true)
    # locate the installed .ko dynamically (modinfo -n prints the module path)
    on_disk_mod=$(modinfo -n nvidia 2>/dev/null || true)
    if [[ "$loaded_ver" == "$DRIVER_VERSION" ]] \
       && [[ -n "$on_disk_mod" ]] \
       && modinfo "$on_disk_mod" 2>/dev/null \
          | grep -q "^version:.*${DRIVER_VERSION}"; then
      skip "patched nvidia ${DRIVER_VERSION} already loaded"
      return
    fi
  fi

  # safety: refuse if GPU is in use
  if command -v nvidia-smi &>/dev/null \
     && nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | grep -q .; then
    err "GPU has running processes; cannot rmmod. Stop them and re-run."
    nvidia-smi --query-compute-apps=pid,process_name --format=csv
    exit 1
  fi

  # If RMForceStaticBar1 is in the modprobe conf but the running kernel hasn't
  # had simpledrm suppressed at boot yet (no initcall_blacklist=sysfb_init in
  # /proc/cmdline), a live rmmod+modprobe would re-trigger the static BAR1
  # overflow that broke 0000:05:00.0. Install the freshly built modules to disk
  # so the next reboot picks them up, but defer the live swap to that reboot
  # (which loads them with simpledrm suppressed).
  if [[ -f "$MODPROBE_CONF" ]] \
     && grep -q 'RMForceStaticBar1=1' "$MODPROBE_CONF" 2>/dev/null \
     && ! grep -q 'initcall_blacklist=sysfb_init' /proc/cmdline 2>/dev/null; then
    log "installing freshly built modules to disk (deferring live swap)"
    runsh "cd '$REPO_DIR' && make modules_install -j\$(nproc)"
    runsh "depmod -a"
    warn "RMForceStaticBar1 is set but simpledrm isn't suppressed yet; a live swap"
    warn "would overflow 05:00.0's BAR1. Reboot with initcall_blacklist=sysfb_init."
    warn "Skipping live rmmod+modprobe."
    return
  fi

  # unload existing (dependents first); ignore failure if not loaded
  log "unloading current nvidia modules"
  runsh "rmmod nvidia_drm 2>/dev/null || true"
  runsh "rmmod nvidia_modeset 2>/dev/null || true"
  runsh "rmmod nvidia_uvm 2>/dev/null || true"
  runsh "rmmod nvidia 2>/dev/null || true"

  if module_loaded nvidia || module_loaded nvidia_uvm; then
    err "rmmod failed (module still in use). Stop GPU processes and re-run."
    exit 1
  fi

  # install the freshly built modules
  log "installing new modules"
  runsh "cd '$REPO_DIR' && make modules_install -j\$(nproc)"
  runsh "depmod -a"

  # load
  log "loading new modules"
  runsh "modprobe nvidia"
  runsh "modprobe nvidia_uvm"
  # modeset/drm may be optional; capture errors and warn on real failures
  # rather than swallowing all stderr.
  for mod in nvidia_modeset nvidia_drm; do
    if runsh "modprobe $mod 2>/tmp/nvp2p-modprobe.err"; then
      :
    else
      rc=$?
      if grep -qiE 'not found|unknown symbol' /tmp/nvp2p-modprobe.err 2>/dev/null; then
        warn "modprobe $mod failed (see /tmp/nvp2p-modprobe.err); may be optional"
      else
        warn "modprobe $mod exited $rc (see /tmp/nvp2p-modprobe.err)"
      fi
    fi
  done

  if module_loaded nvidia && command -v nvidia-smi &>/dev/null; then
    if nvidia-smi &>/dev/null; then
      ok "nvidia ${DRIVER_VERSION} loaded; nvidia-smi works"
    else
      warn "nvidia loaded but nvidia-smi failed; check dmesg"
    fi
  else
    err "nvidia module did not load. check dmesg for details."
    exit 1
  fi
}

# ---------- step: reboot ----------
step_reboot() {
  log "reboot"
  if [[ "$NO_REBOOT" -eq 1 ]]; then
    skip "reboot skipped (--no-reboot)"
    warn "reboot manually with: sudo reboot"
    return
  fi
  if [[ "$DRY_RUN" -eq 1 ]]; then
    printf '\033[0;90m[dry]\033[0m sudo reboot\n'
    return
  fi
  warn "rebooting in 5 seconds (Ctrl-C to cancel)..."
  sleep 5
  reboot
}

# ---------- main ----------
main() {
  local s
  for s in "${STEPS_ORDER[@]}"; do
    if step_enabled "$s"; then
      "step_$s"
    fi
  done
  ok "all requested steps complete"
}

main