Automated Translation and Accelerated Solving of Differential Equations on Multiple GPU Platforms

Abstract: 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 vectorizing 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 incorporation 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. We distribute the software as an open-source library https://github.com/SciML/DiffEqGPU.jl

• The paper introduces custom GPU kernel generation for ODE/SDE solving, eliminating kernel launch overhead and achieving significant performance gains.
• It integrates seamlessly with Julia’s DifferentialEquations.jl, offering vendor-agnostic support across NVIDIA, AMD, Intel, and Apple GPUs.
• Benchmarks reveal the approach outperforms traditional vectorized methods in speed, scalability, and flexibility for large-scale computational applications.

Automated Translation and Accelerated Solving of Differential Equations on Multiple GPU Platforms

This paper presents a method for efficiently solving ensembles of Ordinary Differential Equations (ODEs) and Stochastic Differential Equations (SDEs) on Graphics Processing Units (GPUs) without vendor specificity. By integrating with Julia's DifferentialEquations.jl, the approach allows users to leverage GPU acceleration seamlessly, delivering performance comparable to hand-optimized CUDA-C++ kernels and achieving significant speed enhancements—20 to 100 times faster—over traditional vectorized approaches in JAX and PyTorch.

Methodology

The solution revolves around two core strategies for GPU parallelism:

1. EnsembleGPUArray: Implements GPU vectorization and works directly with existing solvers, promoting easy integration within Julia's SciML ecosystem. Despite an intuitive design, this method suffers from notable overheads due to frequent kernel launches, hence impacting performance.
2. EnsembleGPUKernel: Generates custom GPU kernels for entire ODE integrations. This method eliminates the kernel launch overhead and supports adaptive time-stepping and automatic differentiation, providing state-of-the-art performance.

Performance and Benchmarks

• Scalability and Efficiency: The paper showcases the method's ability to solve billions of ODEs across various GPU clusters, highlighting its potential for large-scale applications such as parameter sweeps and uncertainty quantification.
• Vendor-Agnosticism: Implementations were tested across major brands like NVIDIA, AMD, Intel, and Apple GPUs, maintaining consistent performance. The benchmarks reveal that NVIDIA GPUs outperform due to mature ecosystem support.

Numerical Methods and Applications

The paper discusses several ODE and SDE solving methods, notably the Rosenbrock and Runge-Kutta methods, demonstrating both non-stiff and stiff problem-solving capabilities.

• Advanced Features: The approach supports features like event handling, automatic differentiation, and dataset incorporation via texture memory, enabling comprehensive study of complex dynamical systems and stochastic processes.

Practical Implications

This method democratizes high-performance GPU computing for differential equations, allowing scientists and engineers without CUDA expertise to employ GPUs efficiently.

• Comparison with Existing Tools: Demonstrates superior performance compared to MPGOS, JAX, and PyTorch, suggesting that custom kernel generation is essential for achieving maximal efficiency in solving differential equations.

Future Directions

The results highlight limitations in array abstraction tools like PyTorch and JAX for GPU parallelism, urging further research into more tailored GPU kernel generation. Future research may focus on expanding the approach to solve other classes of differential equations and enhancing support for varied complex models.

Conclusion

The integration of GPU-based differential equation solvers in Julia through DiffEqGPU.jl represents a substantial advancement in numerical computing, combining flexibility, cross-platform compatibility, and leading performance. This approach is poised to facilitate breakthroughs in computational science, offering significant implications for future developments in parallel computing and scientific modeling.