# gpu-unlock (github.com/Zagitta/gpu-unlock): full-rate FP16/BF16 on GA102.
# Loads from the initrd before nvidia binds; nothing is flashed, a power-off undoes it.
# Disable by removing gpu_unlock.targets from boot.kernelParams.
{ config, lib, pkgs, ... }:

let
  nv = "610.43.03";
  kernel = config.boot.kernelPackages.kernel;

  src = pkgs.fetchFromGitHub {
    owner = "Zagitta";
    repo = "gpu-unlock";
    rev = "af28cbbfd3d922a659227ecfb7369adfd924cf03";  # 1.0.0, 2026-10-24
    hash = "sha256-6FaDuPojB707ds7gkQY/g/+m4Tnz1xHZ95UjNfFxBJc=";
  };

  nvSource = pkgs.fetchurl {
    url = "https://download.nvidia.com/XFree86/NVIDIA-kernel-module-source/NVIDIA-kernel-module-source-${nv}.tar.xz";
    hash = "sha256-fhGJI8eiPtw2EU1jJzpG4+BOmvmGlaQiA+esLf6fwdw=";
  };
  nvUtils = pkgs.fetchurl {
    url = "https://archive.cachyos.org/archive/cachyos-v3/nvidia-utils-${nv}-1-x86_64_v3.pkg.tar.zst";
    hash = "sha256-4X58ku7/oXiw/h9A6xekcJX4/Dzkp3vxFc9b1aTfhuk=";
  };

  # The generators check every input and output against pinned SHA-256s.
  firmware = pkgs.stdenvNoCC.mkDerivation {
    pname = "gpu-unlock-firmware";
    version = "1.0.0";
    inherit src;
    nativeBuildInputs = [ pkgs.python3 pkgs.zstd pkgs.xz ];
    dontConfigure = true;
    buildPhase = ''
      mkdir nvsrc assets
      tar -xf ${nvSource} -C nvsrc
      python3 firmware/generate_v67.py ${nvUtils} assets
      python3 firmware/extract_boot_assets.py nvsrc/NVIDIA-kernel-module-source-${nv} assets
    '';
    installPhase = ''
      install -Dm644 -t $out/lib/firmware/gpu-unlock/${nv} \
        assets/booter_load_ga102.bin assets/gsp_ga10x.bin assets/gsp_boot_ga102.bin \
        assets/gsp_boot_desc_ga102.bin assets/v67-ga102-${nv}.bin
    '';
  };

  module = pkgs.stdenv.mkDerivation {
    pname = "gpu-unlock";
    version = "1.0.0-${kernel.modDirVersion}";
    inherit src;
    nativeBuildInputs = kernel.moduleBuildDependencies;
    hardeningDisable = [ "pic" "format" ];
    buildPhase = ''
      make -C ${kernel.dev}/lib/modules/${kernel.modDirVersion}/build \
        M=$PWD/kernel/gpu_unlock ARCH=${pkgs.stdenv.hostPlatform.linuxArch} modules
    '';
    installPhase = ''
      install -Dm644 kernel/gpu_unlock/gpu_unlock.ko \
        $out/lib/modules/${kernel.modDirVersion}/extra/gpu_unlock.ko
    '';
  };
in
{
  boot.extraModulePackages = [ module ];
  hardware.firmware = [ firmware ];
  # Must run before nvidia binds the cards, so it lives in the initrd with its firmware.
  boot.initrd.kernelModules = lib.mkBefore [ "gpu_unlock" ];
  boot.initrd.extraFirmwarePaths = map (f: "gpu-unlock/${nv}/${f}") [
    "booter_load_ga102.bin" "gsp_ga10x.bin" "gsp_boot_ga102.bin"
    "gsp_boot_desc_ga102.bin" "v67-ga102-${nv}.bin"
  ];
  boot.kernelParams = [ "gpu_unlock.targets=all" ];
}