Installation#

Building Connector#

The connector interfaces XLA and PJRT with MD applications such as LAMMPS, which might use a different building system and MPI.

In general, compiling the connector requires clang with C++14 support as well as the NVCC cuda compiler.

The connector can be built using the following command:

python build.py

Additionally, the PjRt plugin for CUDA enabled GPUs can be built using

python build.py --build_gpu_pjrt_plugin --enable_cuda --cuda_version 12.6.0

Alternatively, a prebuilt PjRt plugin can be fetched from JAX. Therefore, a JAX version compatible to the installed CUDA version and compatible to the XLA library must be installed. Then, the plugin can be fetched via

python build.py --load_gpu_pjrt_plugin

Building LAMMPS Plugin#

In the connector directory create and cd into a build directory and compile the plugin with the following commands:

mkdir build && cd build
cmake -D LAMMPS_HEADER_DIR=<path/to/lammps/src> ../lammps_plugin
cmake --build .

Note: When the plugin is changed, it must be recompiled via

cmake --build . --clean-first

Building LAMMPS with Plugin Support#

To build lammps with plugin support, run:

cmake -D PKG_PLUGIN=yes ../cmake
cmake --build . -j <number_of_cores>

“Installing” LAMMPS and the plugin#

To “install” LAMMPS and the plugin, we can create a script to set the correct environment variables. The script should look like this:

activate:

#! /bin/bash

export PATH=<path/to/lammps/build>:$PATH
export LAMMPS_PLUGIN_PATH=<path/to/chemtrain-deploy/build>
export JCN_PJRT_PATH=<path/to/chemtrain-deploy/lib>

Calling the script with source ./activate will set all necessary variables to discover the LAMMPS executable, the plugin, and the PJRT library.

Docker Container#

Note: Using chemtrain-deploy within a docker container requires the NVIDIA Container Toolkit.

Before compiling the connector, we have to determine the compute capabilities of the GPUs. Therefore, we can run the following command

nvidia-smi --query-gpu=compute_cap --format=csv,noheader

We then set the compute capabilities as environment variable and build the docker container.

CUDA_COMPUTE_CAPABILITIES="8.6,8.0"

docker build -t chemtrain-deploy \
    --build-arg CUDA_COMPUTE_CAPABILITIES=${CUDA_COMPUTE_CAPABILITIES} \
    -f Dockerfile .

Afterward, simulations can be run inside the container:

docker run --gpus all -it --rm -v /home/ga27pej/myjaxmd/examples/spice:/workspace chemtrain-deploy

Installation

Contents