config/hardware/gpu/cuda-common.nix

{ config, pkgs, lib, ... }:
# assumes common-gpu-nvidia from nixos-hardware is also added
# https://github.com/NixOS/nixos-hardware/blob/master/common/gpu/nvidia/default.nix
# https://github.com/NixOS/nixos-hardware/blob/master/common/gpu/nvidia/prime.nix
lib.mkMerge [

(lib.mkIf (lib.versionOlder lib.version "24.11") {
  hardware.opengl.enable = true;
  hardware.opengl.driSupport = true;
  hardware.opengl.driSupport32Bit = true;
})
(lib.mkIf (lib.versionAtLeast lib.version "24.11") {
  hardware.graphics.enable = true;
  hardware.graphics.enable32Bit = true;

  # does not work on pascal
  # hardware.nvidia.open = lib.mkDefault true;

  # Experimental and only works on modern Nvidia GPUs (Turing or newer).
  /* powerManagement.finegrained = true; */
})
{

  nixpkgs.config.cudaSupport= true;

  nixpkgs.config.allowUnfreePredicate =
    pkg:
    let
      pkgname = (lib.toLower (lib.getName pkg));
    in
    lib.any (x: x) [
      (lib.hasInfix "nvidia" pkgname)
      (lib.hasInfix "cuda"   pkgname)
      (lib.hasInfix "cudnn"  pkgname)
      (lib.hasInfix "cublas" pkgname)
    ];

  # avoid expensive rebuilds
  nixpkgs.config.packageOverrides = pkgs: {
    firefox-unwrapped = pkgs.firefox-unwrapped.override (old: {
      onnxruntime = old.onnxruntime.override {
        cudaSupport = false;
        ncclSupport = false;
      };
    });
  };

}
/* lib.mkIf (lib.versionAtLeast lib.version "24.11") { */
(lib.optionalAttrs (lib.versionAtLeast lib.version "24.11") {
  programs.nix-required-mounts.enable = true;
  programs.nix-required-mounts.presets.nvidia-gpu.enable = true;
  # adding "cuda", "opengl" and "nvidia-gpu" ^ overrides the implicit defaults
  nix.settings.system-features = lib.mkIf (pkgs.stdenv.hostPlatform.system == "x86_64-linux") ["kvm" "nixos-test"];
})
{

  # https://wiki.nixos.org/wiki/Nvidia

  hardware.opengl.extraPackages = [
    #pkgs.vaapiVdpau # added by nixos-hardware
    #pkgs.libvdpau-va-gl
    pkgs.nvidia-vaapi-driver
  ];

  hardware.nvidia.modesetting.enable = lib.mkDefault true; # needed for most wayland compositors
  hardware.nvidia.nvidiaSettings = lib.mkDefault true;

  # https://discourse.nixos.org/t/nvidia-docker-container-runtime-doesnt-detect-my-gpu/51336/2?u=pbsds
  /* virtualisation.docker.enableNvidia = lib.mkDefault true; # deprecated */
  /* virtualisation.podman.enableNvidia = lib.mkDefault true; # deprecated */
  # this works, but you have to use `--device nvidia.com/gpu=all` rather than --gpus
  hardware.nvidia-container-toolkit.enable = lib.mkDefault true;

  services.ollama.acceleration = lib.mkDefault "cuda";

  # only do these per-host

  #hardware.nvidia.open = lib.mkDefault true; # open source version of kernel module, only on driver 515.43.04+
  #hardware.nvidia.powerManagement.enable = lib.mkDefault true; # Fix graphical corruption on suspend/resume
  #hardware.nvidia.powerManagement.finegrained = lib.mkDefault false; # Turns off GPU when not in use, Turing or newer

  # https://www.nvidia.com/en-us/drivers/unix/legacy-gpu/
  # https://github.com/NixOS/nixpkgs/blob/nixos-unstable/pkgs/os-specific/linux/nvidia-x11/default.nix
  #hardware.nvidia.package = config.boot.kernelPackages.nvidiaPackages.stable;
  #hardware.nvidia.package = config.boot.kernelPackages.nvidiaPackages.production;
  #hardware.nvidia.package = config.boot.kernelPackages.nvidiaPackages.latest;
  #hardware.nvidia.package = config.boot.kernelPackages.nvidiaPackages.beta;

  environment.systemPackages = [
    (pkgs.nvtopPackages.nvidia or pkgs.nvtop-nvidia)
  ];

  # nvidia bad
  nixpkgs.config.packageOverrides = pkgs: {
    /** /
    teams = pkgs.teams.overrideAttrs (old: rec {
      postInstall = (old.postInstall or "") + ''
        substituteInPlace $out/share/applications/teams.desktop \
          --replace 'Exec=teams' 'Exec=teams --use-gl=desktop'
      '';
      #  --replace 'Exec=teams' 'Exec=teams --disable-gpu'
    });
    /**/
    discord = pkgs.discord.overrideAttrs (old: rec {
      desktopItem = old.desktopItem.override (old: {
        #exec = "env NIXOS_OZONE_WL=1 ${old.exec}";
        exec = "${old.exec} --use-gl=desktop";
      });
      postInstall = ''
        ln -sf "${desktopItem}/share/applications" $out/share/
      '' + old.postInstall;
    });
    /** /
    slack = pkgs.slack.overrideAttrs (old: {
      postInstall = old.postInstall or "" + ''
        substituteInPlace $out/share/applications/slack.desktop \
          --replace "$out/bin/slack " "$out/bin/slack --use-gl=desktop "
      '';
      #  --replace "$out/bin/slack " "$out/bin/slack --disable-gpu --in-process-gpu "
    });
    element-desktop = pkgs.element-desktop.overrideAttrs (old: {
      desktopItem = old.desktopItem.override (old: {
        #exec = "env NIXOS_OZONE_WL=1 element-desktop %u";
        #exec = "element-desktop --disable-gpu --in-process-gpu %u";
        exec = "element-desktop --use-gl=desktop %u";
      });
    });
    /**/
  };
}

]