We demonstrate a high-performance vendor-agnostic method for massively parallel solving of ensembles of ordinary differential equations (ODEs) and stochastic differential equations (SDEs) on GPUs. The method is integrated with a widely used differential equation solver library in a high-level language (Julia's DifferentialEquations.jl) and enables GPU acceleration without requiring code changes by the user. Our approach achieves state-of-the-art performance compared to hand-optimized CUDA-C++ kernels, while performing $20-100\times$ faster than the vectorized-map (vmap) approach implemented in JAX and PyTorch. Performance evaluation on NVIDIA, AMD, Intel, and Apple GPUs demonstrates performance portability and vendor-agnosticism. We show composability with MPI to enable distributed multi-GPU workflows. The implemented solvers are fully featured, supporting event handling, automatic differentiation, and incorporating of datasets via the GPU's texture memory, allowing scientists to take advantage of GPU acceleration on all major current architectures without changing their model code and without loss of performance.
翻译:我们展示了一种高性能、供应商无关的方法,用于在GPU上对常微分方程(ODEs)和随机微分方程(SDEs)集合进行大规模并行求解。该方法集成于高级语言中广泛使用的微分方程求解器库(Julia的DifferentialEquations.jl),且用户无需修改代码即可启用GPU加速。与手写优化的CUDA-C++内核相比,我们的方法实现了最先进的性能,同时比JAX和PyTorch中实现的向量化映射(vmap)方法快20-100倍。在NVIDIA、AMD、Intel和Apple GPU上的性能评估证明了其可移植性和供应商无关性。我们展示了该方法与MPI的组合能力,以支持分布式多GPU工作流。实现的求解器功能完备,支持事件处理、自动微分以及通过GPU纹理内存整合数据集,使科学家能够在无需修改模型代码且不损失性能的前提下,利用当前所有主流架构上的GPU加速。