Exploring CPU and GPU strategies for solving some coupled nonlinear Schrödinger systems in Julia

Abstract

We developed the SchrodingerToolkit.jl package in the Julia programming language to solve systems of coupled nonlinear Schrödinger equations in both 2D and 3D. The package builds upon the theories, methods, and strategies presented in [1-3]. It supports parallel computations with GPU acceleration and is designed for extensibility, making it a valuable research tool for exploring parallel techniques for these problems. The package features high-order spatial schemes using finite differences and high-order temporal schemes employing composition methods. It also includes tools for plotting, constructing and modeling abstract grids, tracking runtime statistics for execution and solver analysis, and logging key physical quantities such as mass per component, system energy, and total mass perturbation. GPU parallel implementation is achieved through CUDA.jl, allowing the generation of CUDA kernels directly from Julia code. This approach results in significant speedups of 5x to 10x, depending on the method used, particularly for 3D problems where matrix factorizations are often impractical due to time and memory constraints, even for smaller problems. The main bottlenecks in the implementation are related to solving linear systems, especially due to the use of implicit time-stepping methods, which require global communication or some form of internal state.

Desarrollamos el paquete SchrodingerToolkit.jl en el lenguaje de programación Julia para resolver sistemas de ecuaciones de Schrödinger no lineales acopladas en 2D y 3D. El paquete se basa en las teorías, métodos y estrategias presentados en \cite{aguilera2021structure,Aguilera_Castillo_Gómez_2022,Castillo2023}. Soporta cálculos paralelos con aceleración por GPU y está diseñado para ser extensible, convirtiéndose en una valiosa herramienta de investigación para explorar técnicas paralelas en estos problemas. El paquete ofrece esquemas espaciales de alto orden utilizando diferencias finitas y esquemas temporales de alto orden empleando métodos de composición. También incluye herramientas para la visualización, construcción y modelado de cuadrículas abstractas, seguimiento de estadísticas de ejecución y registro de cantidades físicas clave como la masa por componente, la energía del sistema y la perturbación de la masa total. La implementación paralela en GPU se logra mediante CUDA.jl, lo que permite la generación de kernels CUDA directamente desde código en Julia. Este enfoque resulta en aceleraciones significativas de 5x a 10x, dependiendo del método utilizado, especialmente en problemas 3D donde las factorizaciones de matrices a menudo son impracticables debido a las limitaciones de tiempo y memoria, incluso para problemas más pequeños. Las principales limitaciones en la implementación están relacionados con la resolución de sistemas lineales, especialmente debido al uso de métodos de integración temporal implícitos, que requieren comunicación global o alguna forma de estado interno.