shells/cuda.nix


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

{
  pkgs,
  lib,
}: let
  nvidiaPackage = pkgs.linuxPackages_latest.nvidiaPackages.stable;
  libs = with pkgs; [
    # CUDA
    cudaPackages.cuda_cudart
    cudaPackages.cuda_nvcc
    cudaPackages.cuda_cccl
    cudaPackages.cudnn
    cudatoolkit
    nvidiaPackage

    # OpenGL
    libGLU
    libGL
    freeglut

    # X11
    xorg.libXi
    xorg.libXmu
    xorg.libXext
    xorg.libX11
    xorg.libXv
    xorg.libXrandr

    # Dev Tools
    ffmpeg
    zlib
    gcc
    binutils
  ];
in {
  packages = libs;
  LD_LIBRARY_PATH = lib.makeLibraryPath libs;

  shellHook = ''
    # Nvidia Driver setup
    export NVIDIA_VISIBLE_DEVICES=all
    export NVIDIA_DRIVER_CAPABILITIES=compute,utility
    export CUDA_VISIBLE_DEVICES=0
    export CUDA_PATH=${pkgs.cudatoolkit}
  '';
}