gistfile1.txt

{
  config,
  lib,
  pkgs,
  ...
}: let
  cfg = config.kernel.autofdo;

  profileExists =
    cfg.profilePath
    != null
    && builtins.pathExists cfg.profilePath;

  helper = pkgs.writeShellApplication {
    name = "kernel-afdo-collect";
    runtimeInputs = [
      pkgs.linuxPackages_latest.perf
      pkgs.llvmPackages_latest.llvm
      pkgs.coreutils
      pkgs.gawk
      pkgs.gnugrep
    ];
    text = ''
      set -euo pipefail

      duration=60
      output=""
      vmlinux="${config.boot.kernelPackages.kernel.dev}/vmlinux"
      while [ $# -gt 0 ]; do
        case "$1" in
          --duration)
            duration="$2"; shift 2 ;;
          --output)
            output="$2"; shift 2 ;;
          --vmlinux)
            vmlinux="$2"; shift 2 ;;
          -h|--help)
            cat <<EOF
      Usage: kernel-afdo-collect [--duration SECS] [--output PATH] [--vmlinux PATH]

      Records branch-sample perf data against the running kernel,
      converts it via llvm-profgen, and writes an .afdo profile
      suitable for the kernel.autofdo NixOS profile.

      Defaults:
        --duration 60
        --output   /tmp/kernel-\$(hostname)-\$(date +%Y%m%d-%H%M%S).afdo
        --vmlinux  ${config.boot.kernelPackages.kernel.dev}/vmlinux
                   (the dev output of the kernel package in the current system;
                    must match the booted kernel — reboot into the new system
                    before collecting, or pass --vmlinux explicitly)

      After collection, copy the file into your flake repo at
      hosts/<host>/afdo/kernel.afdo and set kernel.autofdo.enable = true.
      EOF
            exit 0 ;;
          *)
            echo "unknown argument: $1" >&2; exit 2 ;;
        esac
      done

      if [ "$(id -u)" -ne 0 ]; then
        echo "error: kernel-afdo-collect must run as root (system-wide perf + sysctl)" >&2
        echo "       try: sudo kernel-afdo-collect ..." >&2
        exit 1
      fi

      if [ -z "$output" ]; then
        output="/tmp/kernel-$(hostname)-$(date +%Y%m%d-%H%M%S).afdo"
      fi

      vendor="$(awk -F': ' '/^vendor_id/ {print $2; exit}' /proc/cpuinfo)"
      case "$vendor" in
        AuthenticAMD)
          if ! grep -qE 'amd_lbr_v2|brs' /proc/cpuinfo; then
            echo "error: AMD CPU lacks amd_lbr_v2 or BRS — branch sampling unsupported" >&2
            exit 1
          fi
          event_args=(-e 'ex_ret_brn_tkn:k')
          ;;
        GenuineIntel)
          event_args=(-e 'BR_INST_RETIRED.NEAR_TAKEN:k')
          ;;
        *)
          echo "error: unsupported CPU vendor: $vendor" >&2
          exit 1
          ;;
      esac

      if [ ! -e "$vmlinux" ]; then
        echo "error: vmlinux not found at $vmlinux" >&2
        echo "       pass --vmlinux PATH or rebuild+reboot so the kernel.dev" >&2
        echo "       output of the current system matches the booted kernel" >&2
        exit 1
      fi

      booted_release="$(uname -r)"
      vmlinux_release="$(strings "$vmlinux" | grep -m1 -oE 'Linux version [0-9][^ ]*' | awk '{print $3}' || true)"
      if [ -n "$vmlinux_release" ] && [ "$vmlinux_release" != "$booted_release" ]; then
        echo "warning: booted kernel is $booted_release but vmlinux reports $vmlinux_release" >&2
        echo "         profile attribution will be wrong; reboot or pass --vmlinux" >&2
      fi

      old_kptr="$(sysctl -n kernel.kptr_restrict)"
      old_paranoid="$(sysctl -n kernel.perf_event_paranoid)"

      tmpdir="$(mktemp -d)"
      trap 'rm -rf "$tmpdir"; sysctl -q -w kernel.kptr_restrict="$old_kptr"; sysctl -q -w kernel.perf_event_paranoid="$old_paranoid"' EXIT

      sysctl -q -w kernel.kptr_restrict=0
      sysctl -q -w kernel.perf_event_paranoid=0

      echo "recording $duration s of branch samples..."
      perf record -a -N -b -c 500009 \
        "''${event_args[@]}" \
        -o "$tmpdir/perf.data" \
        -- sleep "$duration"

      echo "converting samples to AutoFDO profile..."
      llvm-profgen \
        --kernel \
        --binary="$vmlinux" \
        --perfdata="$tmpdir/perf.data" \
        --output="$output"

      if [ ! -s "$output" ]; then
        echo "error: llvm-profgen produced an empty profile (no LBR samples?)" >&2
        exit 1
      fi

      cat <<EOF

      wrote $output

      To use it:
        cp "$output" hosts/$(hostname)/afdo/kernel.afdo
      then set kernel.autofdo.enable = true and rebuild.
      EOF
    '';
  };
in {
  options.kernel.autofdo = {
    enable = lib.mkEnableOption "AutoFDO sample-use for the kernel build";

    profilePath = lib.mkOption {
      type = lib.types.nullOr lib.types.path;
      default = null;
      description = ''
        Path to an .afdo profile generated by kernel-afdo-collect.
        Required when enable = true. Build fails loudly if missing.
      '';
    };
  };

  config = {
    environment.systemPackages = [helper];

    assertions = [
      {
        assertion = !cfg.enable || profileExists;
        message = ''
          kernel.autofdo.enable is true but kernel.autofdo.profilePath
          (${toString cfg.profilePath}) does not exist. Run
          kernel-afdo-collect on the booted machine, copy the resulting
          .afdo into the repo, then re-enable.
        '';
      }
      {
        assertion = !cfg.enable || config.kernel.tuned.lto == "thin";
        message = ''
          kernel.autofdo.enable requires kernel.tuned.lto = "thin".
          AutoFDO is incompatible with Full LTO.
        '';
      }
    ];

    kernel.tuned.extraMakeFlags = lib.mkIf (cfg.enable && profileExists) [
      "CLANG_AUTOFDO_PROFILE=${cfg.profilePath}"
    ];

    boot.kernelPatches = lib.mkIf (cfg.enable && profileExists) (lib.singleton {
      name = "kernel-autofdo";
      patch = null;
      structuredExtraConfig = with lib.kernel; {
        AUTOFDO_CLANG = lib.mkForce yes;
      };
    });
  };
}