config/hardware/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://nixos.wiki/wiki/Nvidia

  #nixpkgs.config.cudaSupport = true; # TODO: TOO SLOW, BREAKS
  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)))
  ];

  hardware.opengl.enable = true;
  hardware.nvidia.package = config.boot.kernelPackages.nvidiaPackages.stable;
  hardware.nvidia.modesetting.enable = true;

  # add this to the host in question:
  #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 = with pkgs; ([
    nvtop

  ] ++ lib.optional config.hardware.nvidia.prime.offload.enable [
    (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
    #(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 "$@"
    #'')
  ]);
}