Getting Started#

Install nvmath-python#

nvmath-python, like most modern Python packages, provides pre-built binaries (wheels and later conda packages) to the end users. The full source code is hosted in the NVIDIA/nvmath-python repository.

In terms of CUDA Toolkit (CTK) choices, nvmath-python is designed and implemented to allow building and running against 1. pip-wheel, 2. conda, or 3. system installation of CTK. Having a full CTK installation at either build- or run- time is not necessary; just a small fraction as explained below is enough.

Host & device APIs (see Overview) have different run-time dependencies and requirements. Even among host APIs the needed underlying libraries are different (for example, fft() on GPUs only needs cuFFT and not cuBLAS). Libraries are loaded when only needed. Therefore, nvmath-python is designed to have most of its dependencies optional, but provides convenient installation commands for users to quickly spin up a working Python environment.

The cheatsheet below captures nvmath-python’s required/optional, build-/run- time dependencies. Using the installation commands from the sections below should support most of your needs.

Install from PyPI#

The pre-built wheels can be pip-installed from the public PyPI. There are several optional dependencies expressible in the standard “extras” bracket notation. The following assumes that CTK components are also installed via pip (so no extra step from users is needed; the dependencies are pulled via extras).

Command

Description

pip install nvmath-python[cu11]

Install nvmath-python along with all CUDA 11 optional dependencies (wheels for cuBLAS/cuFFT/… and CuPy) to support nvmath host APIs.

pip install nvmath-python[cu12]

Install nvmath-python along with all CUDA 12 optional dependencies (wheels for cuBLAS/cuFFT/… and CuPy) to support nvmath host APIs.

pip install nvmath-python[cu12,dx]

Install nvmath-python along with all CUDA 12 optional dependencies (wheels for cuBLAS/cuFFT/…, CuPy, Numba, pynvjitlink, …) to support nvmath host & device APIs (which only supports CUDA 12) [8].

pip install nvmath-python[cpu]

Install nvmath-python along with all CPU optional dependencies (wheels for NVPL or MKL) to support optimized CPU FFT APIs. [1]

Note:

  1. NVPL is for ARM architecture only. MKL or another FFTW3 [9] compatible library may be substituted for x86 architecture.

  2. The environment variable NVMATH_FFT_CPU_LIBRARY may be used to provide the path to an alternate shared object which implements the FFTW3 (non-guru) API. LD_LIBRARY_PATH should be set properly to include this library if it is not already in the PATH.

The options below are for adventurous users who want to manage most of the dependencies themselves. The following assumes that system CTK is installed.

Command

Description

pip install nvmath-python[sysctk11]

Install nvmath-python along with CuPy for CUDA 11 to support nvmath host APIs.

NoteLD_LIBRARY_PATH should be set properly to include CUDA libraries.

pip install nvmath-python[sysctk12]

Install nvmath-python along with CuPy for CUDA 12 to support nvmath host APIs.

NoteLD_LIBRARY_PATH should be set properly to include CUDA libraries.

pip install nvmath-python[sysctk12-dx]

Install nvmath-python along with CuPy for CUDA 12 to support nvmath host & device APIs.

Note:

  1. LD_LIBRARY_PATH should be set properly to include CUDA libraries.

  2. For using nvmath.device APIs, CUDA_HOME (or CUDA_PATH) should be set to point to the system CTK.

For system admins or expert users, pip install nvmath-python would be a bare minimal installation (very lightweight). This allows fully explicit control of all dependencies.

Install from conda#

Conda packages can be installed from the conda-forge channel.

Command

Description

conda install -c conda-forge nvmath-python cuda-version=11

Install nvmath-python along with all CUDA 11 optional dependencies (wheels for cuBLAS/cuFFT/… and CuPy) to support nvmath host APIs.

conda install -c conda-forge nvmath-python cuda-version=12

Install nvmath-python along with all CUDA 12 optional dependencies (wheels for cuBLAS/cuFFT/… and CuPy) to support nvmath host APIs.

conda install -c conda-forge -c rapidsai nvmath-python-dx "pynvjitlink>=0.2" cuda-version=12

Install nvmath-python along with all CUDA 12 optional dependencies (wheels for cuBLAS/cuFFT/…, CuPy, Numba, pynvjitlink, …) to support nvmath host & device APIs (which only supports CUDA 12).

Note:

  1. nvmath-python-dx is a meta-package for ease of installing nvmath-python and other dependencies.

  2. pynvjitlink currently only lives on the rapidsai channel, not the conda-forge channel.

conda install -c conda-forge nvmath-python-cpu

Install nvmath-python along with all CPU optional dependencies (NVPL or other) to support optimized CPU FFT APIs. [1]

Note:

  1. nvmath-python-cpu is a meta-package for ease of installing nvmath-python and other dependencies.

  2. NVPL is for ARM architecture only. MKL or another FFTW3 [9] compatible library may be substituted for x86 architecture.

  3. The environment variable NVMATH_FFT_CPU_LIBRARY may be used to provide the path to an alternate shared object which implements the FFTW3 (non-guru) API. LD_LIBRARY_PATH should be set properly to include this library if it is not already in the PATH.

Notes:

  • For expert users, conda install -c conda-forge nvmath-python=*=core* would be a bare minimal installation (very lightweight). This allows fully explicit control of all dependencies.

  • If you installed conda from miniforge, most likely the conda-forge channel is already set as the default, then the -c conda-forge part in the above instruction can be omitted.

Build from source#

Once you clone the repository and go into the root directory, you can build the project from source. There are several ways to build it since we need some CUDA headers at build time.

Command

Description

pip install -v .

Set up a build isolation (as per PEP 517), install CUDA wheels and other build-time dependencies to the build environment, build the project, and install it to the current user environment together with the run-time dependencies.

Note: in this case we get CUDA headers by installing pip wheels to the isolated build environment.

CUDA_PATH=/path/to/your/cuda/installation pip install --no-build-isolation -v .

Skip creating a build isolation (it’d use CUDA headers from $CUDA_PATH/include instead), build the project, and install it to the current user environment together with the run-time dependencies. One can use:

  • conda: After installing CUDA 12 conda packages, set the environment variable CUDA_PATH

    • linux-64: CUDA_PATH=$CONDA_PREFIX/targets/x86_64-linux/

    • linux-aarch64: CUDA_PATH=$CONDA_PREFIX/targets/sbsa-linux/

    • win-64: CUDA_PATH=$CONDA_PREFIX\Library

  • system CTK: Just set CUDA_PATH to the system CTK location.

Notes:

  • If you add the “extras” notation after the dot ., e.g., .[cu11], .[cu12,dx], …, it has the same meaning as explained in the previous section.

  • If you don’t want the run-time dependencies to be automatically handled, add --no-deps after the pip install command above; in this case, however, it’s your responsibility to make sure that all the run-time requirements are met.

  • By replacing install by wheel, a wheel can be built targeting the current OS and CPython version.

  • If you want inplace/editable install, add the -e flag to the command above (before the dot .). This is suitable for local development with a system-installed CTK. However, our wheels rely on non-editable builds so that the RPATH hack can kick in. DO NOT pass the -e flag when building wheels!

  • All optional run-time dependencies as listed below need to be manually installed.

Cheatsheet#

Below we provide a summary of requirements to support all nvmath-python functionalities. A dependency is required unless stated otherwise.

When Building

When Running - host APIs

When Running - device APIs

When Running - host API callbacks

CPU architecture & OS

linux-64, linux-aarch64, win-64

linux-64, linux-aarch64, win-64

linux-64, linux-aarch64 [1]

linux-64, linux-aarch64

GPU hardware

All hardware supported by the underlying CUDA Toolkit [5]

Optional: needed if the execution space is GPU.

Compute Capability 7.0+ (Volta and above)

Compute Capability 7.0+ (Volta and above)

CUDA driver [2]

450.80.02+ (Linux) / 450.39+ (Windows) with CUDA >=11.2

525.60.13+ (Linux) / 527.41+ (Windows) with CUDA >=12.0

Optional: needed if the execution space is GPU or for loading any CUDA library.

525.60.13+ (Linux) with CUDA 12.x

525.60.13+ (Linux) with CUDA 12.x

Python

3.10-3.12

3.10-3.12

3.10-3.12

3.10-3.12

pip

22.3.1+

setuptools

>=61.0.0

wheel

>=0.34.0

Cython

>=0.29.22,<3

CUDA

CUDA >=11.2
(only need headers from NVCC & CUDART [6])
CUDA >=11.2

Optional: depending on the math operations in use
CUDA >=12.0,!=12.4.*,!=12.5.0 [7]
(NVRTC, NVVM, CCCL [8], CUDART)

CUDA 12.x

NumPy

>=1.24

>=1.24

>=1.24

>=10.0.0 [4]

>=10.0.0 [4]

>=1.10 (optional)

>=1.10 (optional)

MathDx (cuBLASDx, cuFFTDx, …)

24.04

Numba

0.60

0.60

pynvjitlink

>=0.2

Math Kernel Library (MKL)

2024.4 (optional)

NVIDIA Performance Libraries (NVPL)

24.7 (optional)

Test Configuration#

nvmath-python is tested in the following environments:

CUDA

11.x (latest), 12.x (latest)

Driver

R450, R520, R525, R560

GPU model

A100, H100, RTX 4090, CG1 (Grace-Hopper)

Python

3.10, 3.11, 3.12

CPU architecture

x86_64, aarch64

Operating system

Ubuntu 22.04, Ubuntu 20.04, RHEL 9, Windows 11

Run nvmath-python#

As mentioned earlier, nvmath-python can be run with all methods of CUDA installation, including wheels, conda packages, and system CTK. As a result, there is detection logic to discover shared libraries (for host APIs) and headers (for device APIs to do JIT compilation).

Shared libraries#

  • pip wheels: Will be auto-discovered if installed

  • conda packages: Will be auto-discovered if installed, after wheel

  • system CTK: On Linux, the users needs to ensure the shared libraries are discoverable by the dynamic linker, say by setting LD_LIBRARY_PATH or updating system search paths to include the DSO locations.

Headers#

This includes libraries such as CCCL and MathDx.

  • pip wheels: Will be auto-discovered if installed

  • conda packages: Will be auto-discovered if installed, after wheel

  • system CTK: Need to set CUDA_HOME (or CUDA_PATH) and MATHDX_HOME (for MathDx headers)

Host APIs#

This terminlogy is explained in the Host APIs.

Examples#

See the examples directory in the repo. Currently we have:

  • examples/fft

  • examples/linalg

Tests#

The requirements/pip/tests.txt file lists dependencies required for pip-controlled environments to run tests. These requirements are installed via the main requirements/pip-dev-<name>.txt files.

Running functionality tests#
pytest tests/example_tests tests/nvmath_tests/fft tests/nvmath_tests/linalg
Running performance tests#

This will currently run two tests for fft and one test for linalg:

pytest -v -s -k 'perf' tests/nvmath_tests/fft/
pytest -v -s -k 'perf' tests/nvmath_tests/linalg/

Device APIs#

This terminlogy is explained in the Device APIs.

Examples#

See the examples/device directory in the repo.

Tests#

Running functionality tests#
pytest tests/nvmath_tests/device examples/device
Running performance tests#
pytest -v -s -k 'perf' tests/nvmath_tests/device/

Troubleshooting#

For pip-users, there are known limitations (many of which are nicely captured in the pypackaging community project) in Python packaging tools. For a complex library such as nvmath-python that interacts with many native libraries, there are user-visible caveats.

  1. Be sure that there are no packages with both -cu11 (for CUDA 11) and -cu12 (for CUDA 12) suffices coexisting in your Python environment. For example, this is a corrupted environment:

    $ pip list
    Package            Version
    ------------------ ---------
    nvidia-cublas-cu11 11.11.3.6
    nvidia-cublas-cu12 12.5.2.13
    pip                24.0
    setuptools         70.0.0
    wheel              0.43.0
    

    Some times such conflicts could come from a dependency of the libraries that you use, so pay extra attention to what’s installed.

  2. pip does not attempt to check if the installed packages can actually be run against the installed GPU driver (CUDA GPU driver cannot be installed by pip), so make sure your GPU driver is new enough to support the installed -cuXX packages [2]. The driver version can be checked by executing nvidia-smi and inspecting the Driver Version field on the output table.

  3. CuPy installed from pip currently (as of v13.3.0) only supports conda and system CTK, and not pip-installed CUDA wheels. nvmath-python can help CuPy use the CUDA libraries installed to site-packages (where wheels are installed to) if nvmath is imported. As of beta 2 (v0.2.0) the libraries are “soft-loaded” (no error is raised if a library is not installed) when import nvmath happens. This behavior may change in a future release.

  4. Numba installed from pip currently (as of v0.60.0) only supports conda and system CTK, and not pip-installed CUDA wheels. nvmath-python can also help Numba use the CUDA compilers installed to site-packages if nvmath is imported. Same as above, this behavior may change in a future release.

In general, mixing-and-matching CTK packages from pip, conda, and the system is possible but can be very fragile, so please understand what you’re doing. The nvmath-python internals are designed to work with everything installed either via pip, conda, or local system (system CTK, including tarball extractions, are the fallback solution in the detection logic), but mix-n-match makes the detection logic impossible to get right.

To help you perform an integrity check, the rule of thumb is that every single package should only come from one place (either pip, or conda, or local system). For example, if both nvidia-cufft-cu11 (which is from pip) and libcufft (from conda) appear in the output of conda list, something is almost certainly wrong. Below is the package name mapping between pip and conda, with XX={11,12} denoting CUDA’s major version:

pip

conda (cuda-version>=12)

conda (cuda-version<12)

nvidia-cuda-nvcc-cuXX

cuda-nvcc

n/a

nvidia-cuda-nvrtc-cuXX

cuda-nvrtc

cudatoolkit

nvidia-cuda-runtime-cuXX

cuda-cudart-dev

cudatoolkit

nvidia-cuda-cccl-cuXX

cuda-cccl

n/a

pynvjitlink-cuXX

pynvjitlink

n/a

nvidia-cublas-cuXX

libcublas

cudatoolkit

nvidia-cusolver-cuXX

libcusolver

cudatoolkit

nvidia-cusparse-cuXX

libcusparse

cudatoolkit

nvidia-cufft-cuXX

libcufft

cudatoolkit

nvidia-curand-cuXX

libcurand

cudatoolkit

Note that system packages (by design) do not show up in the output of conda list or pip list. Linux users should check the installation list from your distro package manager (apt, yum, dnf, …). See also the Linux Package Manager Installation Guide for additional information.

For more information with regard to the new CUDA 12+ package layout on conda-forge, see the CUDA recipe README.

Footnotes