warp.sim and Multi-GPU

[knauth]

I'm not seeing anything in the Warp docs about multi-GPU simulations. Is there a useful primitive, system, or technique to leverage multiple GPUs for accelerating sims?

[shi-eric]

I think it's highly dependent on how you envision leveraging multiple GPUs in your application. You won't be able to set a few flags and have a previously single-GPU Warp application automatically decide how to best make use of two GPUs. Instead, you currently need to be explicit about which arrays are allocated on which devices (e.g. cuda:0, cuda:1, cuda:2, etc.), which devices are used for kernel launches, when streams from different devices need to be synchronized with each other, and when copies are needed between Warp arrays allocated on different devices.

Some examples:

• https://github.com/NVIDIA/warp/blob/main/warp/tests/test_streams.py#L361-L465
• https://github.com/NVIDIA/warp/blob/main/warp/tests/test_multigpu.py
• https://github.com/NVIDIA/warp/blob/main/warp/examples/fem/example_diffusion_mgpu.py

Personally my interests are in domain-decomposed simulations, so I have used Warp in combination with mpi4py to leverage multiple GPUs (e.g. launch one process per GPU and inter-GPU communications are done using MPI) and mpi4py's support for __cuda_array_interface__ so that Warp arrays can be passed as arguments to MPI API calls. More recently I've been experimenting with Warp + NCCL for multi-GPU programming.

[knauth]

Interesting! This is a lot to look into, thanks. We may end up leveraging a GPU per collision group with some barycenter logic to make sure they don't overlap.