nix-shell environment for CUDA and DeepJavaLibrary / published by https://github.com/dacr/code-examples-manager #d84d6a0d-e5d2-4a42-aed4-d785231ef7fe/56c95a483944ae80465393759ed18997eb04750b

nix-shell-cuda.nix

## summary : nix-shell environment for CUDA and DeepJavaLibrary
## keywords : nix, djl, nix-shell
## publish : gist
## authors : David Crosson
## license : Apache NON-AI License Version 2.0 (https://raw.githubusercontent.com/non-ai-licenses/non-ai-licenses/main/NON-AI-APACHE2)
## id : d84d6a0d-e5d2-4a42-aed4-d785231ef7fe
## created-on : 2024-01-27T15:49:43+01:00
## managed-by : https://github.com/dacr/code-examples-manager
## run-with : nix-shell $file

# ---------------------------------------
# Run with `nix-shell nix-shell-cuda.nix`
#
# check the status : nvidia-smi  (how much gpu memory is used by process, ...)
# ---------------------------------------
{ pkgs ? import <nixpkgs> {
    config.allowUnfree = true;
    config.cudaSupport = true;
    config.cudaVersion = "12";
  }
}:
pkgs.mkShell {
   name = "cuda-env-shell";
   buildInputs = with pkgs; [
     git gitRepo gnupg autoconf curl
     procps gnumake util-linux m4 gperf unzip
     linuxPackages.nvidia_x11
     cudatoolkit
     cudaPackages.cuda_cudart
     cudaPackages.cudnn
     #cudaPackages.cuda_documentation
     #cudaPackages.cuda-samples
     #cudaPackages.cuda_demo_suite
     libGLU libGL
     xorg.libXi xorg.libXmu freeglut
     xorg.libXext xorg.libX11 xorg.libXv xorg.libXrandr zlib 
     ncurses5 stdenv.cc binutils
   ];
   shellHook = ''
      export CUDA_PATH=${pkgs.cudatoolkit}
      export LD_LIBRARY_PATH=${pkgs.linuxPackages.nvidia_x11}/lib:${pkgs.ncurses5}/lib:${pkgs.cudaPackages.cuda_cudart}/lib
      export EXTRA_LDFLAGS="-L/lib -L${pkgs.linuxPackages.nvidia_x11}/lib"
      export EXTRA_CCFLAGS="-I/usr/include"
   '';
}