config/hardware/gpu/cuda.nix

{ config, pkgs, lib, ... }:
{
  # assumes common-gpu-nvidia from nixos-hardware is also added
  # TODO: should we move it from flake.nix to here?
  # 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

  hardware.opengl.enable = true;
  hardware.opengl.driSupport = true;
  hardware.opengl.driSupport32Bit = true;

  nixpkgs.config.openglSupport = true; # why is this not set by hardware.opengl.enable ?
  nixpkgs.config.cudaSupport = true;
  nixpkgs.config.vdpauSupport  = true;
  nixpkgs.config.vaapiSupport  = true;
  nixpkgs.config.nvidiaSupport = true;

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

  programs.nix-required-mounts.enable = true;
  programs.nix-required-mounts.presets.nvidia-gpu.enable = true;

  # https://nixos.wiki/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://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;

  #hardware.nvidia.prime = {
  #  offload.enable = true;
  #  # Bus IDs. You can find them using lspci, grepping for "3D" or "VGA"
  #  intelBusId  = "PCI:0:2:0";
  #  nvidiaBusId = "PCI:1:0:0";
  #};

  environment.systemPackages = [
    (pkgs.nvtopPackages.nvidia or pkgs.nvtop-nvidia)
  ] ++ lib.optionals config.hardware.nvidia.prime.offload.enable [
    (pkgs.writeShellScriptBin "prime-run" ''
      export __NV_PRIME_RENDER_OFFLOAD=1
      export __NV_PRIME_RENDER_OFFLOAD_PROVIDER=NVIDIA-G0
      export __GLX_VENDOR_LIBRARY_NAME=nvidia
      export __VK_LAYER_NV_optimus=NVIDIA_only
      exec "$@"
    '')
    # already added by nixos-hardware.common-gpu-nvidia
    #(pkgs.writeShellScriptBin "nvidia-offload" ''
    #  export __NV_PRIME_RENDER_OFFLOAD=1
    #  export __NV_PRIME_RENDER_OFFLOAD_PROVIDER=NVIDIA-G0
    #  export __GLX_VENDOR_LIBRARY_NAME=nvidia
    #  export __VK_LAYER_NV_optimus=NVIDIA_only
    #  exec "$@"
    #'')
  ];

  # 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";
      });
    });
    /**/
  };
}