The solution of partial differential equations plays a crucial role in research across many scientific fields. Physics-informed DeepONet (PI-DeepONet) is an advanced deep learning model for solving partial differential equations. This paper focuses on PI-DeepONet and performs multilevel optimization to improve its inference and training speeds. We have customized the single-precision matrix multiplication (SGEMM, C = alpha * \(A \cdot B\) + beta * bias) kernels, within the commonly used model size range of PI-DeepONet, our customized SGEMM kernels achieves 1.1–1.5 times the acceleration compared to the best single-precision matrix multiplication in cuBLAS. Additionally, we designed fusion kernels (integrating multiple computations into one GPU kernel to improve computational efficiency) specifically for the unique structure of the modified MLP used in PI-DeepONet, resulting in 1.7–2.3 times acceleration in the model’s end-to-end inference (includes data transfer from memory to global memory) speed, with GPU computation time being reduced by 3.0–5.7 times. Finally, we employed data parallelism (divide data into multiple subsets and process them in parallel on multiple GPUs) for distributed training of the model and used communication masking strategies, maintaining scalability at 83–85%.
