Lightning-Kokkos installation¶
Standard installation¶
On most Linux systems, Lightning-Kokkos can be installed by following our installation instructions at pennylane.ai/install.
Install Lightning-Kokkos from source¶
Note
The section below contains instructions for installing Lightning-Kokkos from source. For most cases, one can install Lightning-Kokkos via Spack or Docker by the installation instructions at pennylane.ai/install. If those instructions do not work for you, or you have a more complex build environment that requires building from source, then consider reading on.
As Kokkos enables support for many different HPC-targeted hardware platforms, lightning.kokkos
can be built to support any of these platforms when building from source.
Install Kokkos (Optional)¶
Note
Lightning-Kokkos is tested with Kokkos version <= 4.5.00
We suggest first installing Kokkos with the wanted configuration following the instructions found in the Kokkos documentation. For example, the following will build Kokkos for NVIDIA A100 cards:
Download the Kokkos code.
# Replace x, y, and z by the correct version
wget https://github.com/kokkos/kokkos/archive/refs/tags/4.x.yz.tar.gz
tar -xvf 4.x.y.z.tar.gz
cd kokkos-4.x.y.z
Build Kokkos for NVIDIA A100 cards (SM80
architecture), and append the install location to CMAKE_PREFIX_PATH
.
# Replace <install-path> with the path to install Kokkos
# e.g. $HOME/kokkos-install/4.5.0/AMPERE80
export KOKKOS_INSTALL_PATH=<install-path>
mkdir -p ${KOKKOS_INSTALL_PATH}
cmake -S . -B build -G Ninja \
-DCMAKE_BUILD_TYPE=RelWithDebugInfo \
-DCMAKE_INSTALL_PREFIX=${KOKKOS_INSTALL_PATH} \
-DCMAKE_CXX_STANDARD=20 \
-DBUILD_SHARED_LIBS:BOOL=ON \
-DBUILD_TESTING:BOOL=OFF \
-DKokkos_ENABLE_SERIAL:BOOL=ON \
-DKokkos_ENABLE_CUDA:BOOL=ON \
-DKokkos_ARCH_AMPERE80:BOOL=ON \
-DKokkos_ENABLE_COMPLEX_ALIGN:BOOL=OFF \
-DKokkos_ENABLE_EXAMPLES:BOOL=OFF \
-DKokkos_ENABLE_TESTS:BOOL=OFF \
-DKokkos_ENABLE_LIBDL:BOOL=OFF
cmake --build build && cmake --install build
export CMAKE_PREFIX_PATH=:"${KOKKOS_INSTALL_PATH}":$CMAKE_PREFIX_PATH
Note that the C++20 standard is required (enabled via the -DCMAKE_CXX_STANDARD=20
option), hence CUDA 12 is required for the CUDA backend.
Install Lightning-Kokkos¶
If an installation of Kokkos is not found, then our builder will automatically clone and install it during the build process. Lightning-Qubit needs to be ‘installed’ by pip
before Lightning-Kokkos (compilation is not necessary).
The simplest way to install Lightning-Kokkos (OpenMP backend) through pip
.
git clone https://github.com/PennyLaneAI/pennylane-lightning.git
cd pennylane-lightning
pip install -r requirements.txt
pip install git+https://github.com/PennyLaneAI/pennylane.git@master
# Lightning-Qubit needs to be 'installed' by pip before Lightning-Kokkos
# (compilation is not necessary)
PL_BACKEND="lightning_qubit" python scripts/configure_pyproject_toml.py
SKIP_COMPILATION=True pip install -e . --config-settings editable_mode=compat
# Install Lightning-Kokkos with OpenMP backend
PL_BACKEND="lightning_kokkos" python scripts/configure_pyproject_toml.py
CMAKE_ARGS="-DKokkos_ENABLE_OPENMP=ON" python -m pip install -e . --config-settings editable_mode=compat -vv
The supported backend options are
|
|
|
|
|
and the corresponding build options are -DKokkos_ENABLE_XYZ=ON
, where XYZ
needs be replaced by the backend name, for instance OPENMP
.
One can simutaneously activate one serial, one parallel CPU host (e.g. OPENMP
, THREADS
) and one parallel GPU device backend (e.g. HIP
, CUDA
), but not two of any category at the same time.
For HIP
and CUDA
, the appropriate software stacks are required to enable compilation and subsequent use.
Similarly, the CMake option -DKokkos_ARCH_{...}=ON
must also be specified to target a given architecture.
A list of the architectures is found on the Kokkos wiki.
Note that THREADS
backend is not recommended since Kokkos does not guarantee its safety.
Install Lightning-Kokkos with MPI¶
Note
Building Lightning-Kokkos with MPI requires an MPI library and mpi4py
.
If building for GPU, please ensure that MPI is built with GPU support - for example, see the guide to
build OpenMPI with CUDA
and ROCm support.
To install Lightning-Kokkos with MPI support, we recommend first installing Kokkos for your specific architecture such as CPU (SERIAL
, OPENMP
), Nvidia GPU (CUDA
), or AMD GPU (HIP
)
and exporting the install location to CMAKE_PREFIX_PATH
as described above.
Then Lightning-Kokkos with MPI support can be installed in the editable mode by adding the ENABLE_MPI=ON
option to the CMake arguments:
git clone https://github.com/PennyLaneAI/pennylane-lightning.git
cd pennylane-lightning
pip install -r requirements.txt
pip install git+https://github.com/PennyLaneAI/pennylane.git@master
# Lightning-Qubit needs to be 'installed' by pip before Lightning-Kokkos
# (compilation is not necessary)
PL_BACKEND="lightning_qubit" python scripts/configure_pyproject_toml.py
SKIP_COMPILATION=True pip install -e . --config-settings editable_mode=compat
# Install Lightning-Kokkos with MPI support
PL_BACKEND="lightning_kokkos" python scripts/configure_pyproject_toml.py
CMAKE_ARGS="-DENABLE_MPI=ON" python -m pip install -e . --config-settings editable_mode=compat -vv
If required, extra linker flags for MPI (e.g. for GPU Transport Layer) can be added using the MPI_EXTRA_LINKER_FLAGS
environment variable, for example:
# Optional and system dependent
export MPI_EXTRA_LINKER_FLAGS="-lxpmem -L/opt/cray/pe/mpich/8.1.31/gtl/lib -lmpi_gtl_hsa"
For an example of how to install Lightning-Kokkos with MPI on an HPC system, check out the Lightning-Kokkos installation on HPC page
Test Lightning-Kokkos with MPI¶
After installing Lightning-Kokkos with MPI, you can test the Python layer of the MPI enabled plugin as follows (Lightning-Qubit must be installed as well):
pip install -r requirements-tests.txt
PL_DEVICE="lightning_kokkos" mpirun -np 2 python -m pytest mpitests --tb=short
To compile and test the C++ code, you can use the following command:
PL_BACKEND="lightning_kokkos" make test-cpp-mpi
By default this will compile Kokkos with SERIAL
backend. If using a different pre-compiled Kokkos backend, you can export the environment variable CMAKE_PREFIX_PATH
with the Kokkos install location.