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

  #https://nixpk.gs/pr-tracker.html?pr=235481
  #nixpkgs.config.cudaSupport = true; # TODO: TOO SLOW, BREAKS
  #nixpkgs.config.nvidiaSupport = true; # TODO: slow? used only by zenith

  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.opengl.driSupport = true;
  hardware.opengl.extraPackages = with pkgs; [
    vaapiVdpau
    libvdpau-va-gl
  ];

  hardware.nvidia.modesetting.enable = lib.mkDefault true; # needed for most wayland compositors
  hardware.nvidia.nvidiaSettings = true;
  #hardware.nvidia.open = true; # open source version of kernel module, only on driver 515.43.04+
  #hardware.nvidia.package = config.boot.kernelPackages.nvidiaPackages.latest; # only do this per-host
  hardware.nvidia.powerManagement.enable = true; # Fix graphical corruption on suspend/resume



  # 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-nvidia

  ] ++ lib.optionals 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 "$@"
    #'')
  ]);
}