If you’ve ever had trouble running a CUDA application, then you might want to try out this shell hook.
CUDA on Nix is a…
CUDA can be an interesting beast to set up properly. This is not at all helped by CUDA Toolkit that is intended to be setup per project - making a one-time system install not so easy as different software relies on different versions. Thankfully, Nix Flakes can come to the rescue once you have figured out a decent setup. Here’s what I’ve arrived at.
Do you not have an understanding of Nix Flakes?. If not, have a read of my previous post here.
Shell Hooks
Let’s start with a simple flake with a simple shell hook:
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
{
inputs = {
nixpkgs.url = "github:NixOS/nixpkgs/nixos-23.11";
flake-utils.url = "github:numtide/flake-utils";
};
outputs = { self, nixpkgs, flake-utils, ... }:
flake-utils.lib.eachDefaultSystem (system:
let
pkgs = import nixpkgs {
system = system;
config = { allowUnfree = true; };
};
lib = pkgs.lib;
in
{
devShell = pkgs.mkShell {
shellHook = ''
echo YOUR SHELL HOOK HERE
'';
};
});
}
If you stick this in a directory and type nix develop then you will end up in a shell with “YOUR SHELL HOOK HERE” before your usual prompt. A shell hook here allows you to specify a script to run every time you run nix develop. This can be useful for all sorts of things; in this case, setting up our CUDA environment.
The Nix Wiki has a recommended solution as follows:
1
2
3
4
5
6
7
8
devShell = pkgs.mkShell {
shellHook = ''
export CUDA_PATH=${pkgs.cudatoolkit}
# export LD_LIBRARY_PATH=${pkgs.linuxPackages.nvidia_x11}/lib
export EXTRA_LDFLAGS="-L/lib -L${pkgs.linuxPackages.nvidia_x11}/lib"
export EXTRA_CCFLAGS="-I/usr/include"
'';
};
This would do the job, if you uncomment out the LD_LIBRARY_PATH. Of course this may have a major flaw - if your flake has a majorly different version of graphics drivers in your system, this may have some issues. I’m not sure to be honest - if you do then you may have to setup the LD_LIBRARY_PATH on a system level and create an impure devshell (by running nix develop --impure). But that’s not a topic for now.
My CUDA Shell Hook
My shell hook is expanded on this. That’s because at work I’m not allowed to use NixOS, but I have managed to install Nix on Linux and on MacOS (not that CUDA will work on MacOS). Hence I wanted a way to detect whether I’m on NixOS or not - and the great people of the NixOS discourse told me that you would have to do it at a script level with an if statement. I don’t like this, but nevertheless it’s a solution.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
devShell = pkgs.mkShell {
shellHook = ''
export CUDA_PATH=${pkgs.cudatoolkit}
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${lib.makeLibraryPath lib-packages} # Use this if you have other libraries to link
# Check for presence of /etc/nixos - if so, then add pkgs.linuxPackages.nvidia_x11 to LD_LIBRARY_PATH
if [ -d /etc/nixos ]; then
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${pkgs.linuxPackages.nvidia_x11}/lib
export EXTRA_LDFLAGS="-L/lib -L${pkgs.linuxPackages.nvidia_x11}/lib"
else
echo "NixOS not detected - use nixGL to pass in nvidia drivers to scripts for cuda."
fi
export EXTRA_CCFLAGS="-I/usr/include"
'';
};
Here we can see that I check for the presence of /etc/nixos - which is only true if you are on NixOS directly. In that case, we setup the graphics drivers for the linker. Otherwise, I print a nice message on what to do otherwise.
NixGL is a nice community project that solves the issue of getting all your drivers to the linker for running an application. It’s best to install this into your user profile rather than into a project, but it works all the same. You just have to run your program within nixGL and everything magically works!
Some Final Thoughts
CUDA can be a pain to setup right, and hopefully this will help you set up CUDA to your satisfaction. This could be better in a few ways, but this is where I leave it for now. The important thing is it works for me, and hopefully for you as well.
If you found this useful, then share this around to those that may find it useful! Please leave a comment down below if you have any questions or something else to say!