-
NeuralPDR: Neural Differential Equations as surrogate models for Photodissociation Regions
Authors:
Gijs Vermariën,
Thomas G. Bisbas,
Serena Viti,
Yue Zhao,
Xuefei Tang,
Rahul Ravichandran
Abstract:
Computational astrochemical models are essential for helping us interpret and understand the observations of different astrophysical environments. In the age of high-resolution telescopes such as JWST and ALMA, the substructure of many objects can be resolved, raising the need for astrochemical modeling at these smaller scales, meaning that the simulations of these objects need to include both the…
▽ More
Computational astrochemical models are essential for helping us interpret and understand the observations of different astrophysical environments. In the age of high-resolution telescopes such as JWST and ALMA, the substructure of many objects can be resolved, raising the need for astrochemical modeling at these smaller scales, meaning that the simulations of these objects need to include both the physics and chemistry to accurately model the observations. The computational cost of the simulations coupling both the three-dimensional hydrodynamics and chemistry is enormous, creating an opportunity for surrogate models that can effectively substitute the chemical solver. In this work we present surrogate models that can replace the original chemical code, namely Latent Augmented Neural Ordinary Differential Equations. We train these surrogate architectures on three datasets of increasing physical complexity, with the last dataset derived directly from a three-dimensional simulation of a molecular cloud using a Photodissociation Region (PDR) code, 3D-PDR. We show that these surrogate models can provide speedup and reproduce the original observable column density maps of the dataset. This enables the rapid inference of the chemistry (on the GPU), allowing for the faster statistical inference of observations or increasing the resolution in hydrodynamical simulations of astrophysical environments.
△ Less
Submitted 17 June, 2025;
originally announced June 2025.
-
3D-PDR Orion dataset and NeuralPDR: Neural Differential Equations for Photodissociation Regions
Authors:
Gijs Vermariën,
Serena Viti,
Rahul Ravichandran,
Thomas G. Bisbas
Abstract:
We present a novel dataset of simulations of the photodissociation region (PDR) in the Orion Bar and provide benchmarks of emulators for the dataset. Numerical models of PDRs are computationally expensive since the modeling of these changing regions requires resolving the thermal balance and chemical composition along a line-of-sight into an interstellar cloud. This often makes it a bottleneck for…
▽ More
We present a novel dataset of simulations of the photodissociation region (PDR) in the Orion Bar and provide benchmarks of emulators for the dataset. Numerical models of PDRs are computationally expensive since the modeling of these changing regions requires resolving the thermal balance and chemical composition along a line-of-sight into an interstellar cloud. This often makes it a bottleneck for 3D simulations of these regions. In this work, we provide a dataset of 8192 models with different initial conditions simulated with 3D-PDR. We then benchmark different architectures, focusing on Augmented Neural Ordinary Differential Equation (ANODE) based models (Code be found at https://github.com/uclchem/neuralpdr). Obtaining fast and robust emulators that can be included as preconditioners of classical codes or full emulators into 3D simulations of PDRs.
△ Less
Submitted 1 December, 2024;
originally announced December 2024.
-
A shared compilation stack for distributed-memory parallelism in stencil DSLs
Authors:
George Bisbas,
Anton Lydike,
Emilien Bauer,
Nick Brown,
Mathieu Fehr,
Lawrence Mitchell,
Gabriel Rodriguez-Canal,
Maurice Jamieson,
Paul H. J. Kelly,
Michel Steuwer,
Tobias Grosser
Abstract:
Domain Specific Languages (DSLs) increase programmer productivity and provide high performance. Their targeted abstractions allow scientists to express problems at a high level, providing rich details that optimizing compilers can exploit to target current- and next-generation supercomputers. The convenience and performance of DSLs come with significant development and maintenance costs. The siloe…
▽ More
Domain Specific Languages (DSLs) increase programmer productivity and provide high performance. Their targeted abstractions allow scientists to express problems at a high level, providing rich details that optimizing compilers can exploit to target current- and next-generation supercomputers. The convenience and performance of DSLs come with significant development and maintenance costs. The siloed design of DSL compilers and the resulting inability to benefit from shared infrastructure cause uncertainties around longevity and the adoption of DSLs at scale. By tailoring the broadly-adopted MLIR compiler framework to HPC, we bring the same synergies that the machine learning community already exploits across their DSLs (e.g. Tensorflow, PyTorch) to the finite-difference stencil HPC community. We introduce new HPC-specific abstractions for message passing targeting distributed stencil computations. We demonstrate the sharing of common components across three distinct HPC stencil-DSL compilers: Devito, PSyclone, and the Open Earth Compiler, showing that our framework generates high-performance executables based upon a shared compiler ecosystem.
△ Less
Submitted 7 March, 2025; v1 submitted 2 April, 2024;
originally announced April 2024.
-
Automated MPI-X code generation for scalable finite-difference solvers
Authors:
George Bisbas,
Rhodri Nelson,
Mathias Louboutin,
Fabio Luporini,
Paul H. J. Kelly,
Gerard Gorman
Abstract:
Partial differential equations (PDEs) are crucial in modeling diverse phenomena across scientific disciplines, including seismic and medical imaging, computational fluid dynamics, image processing, and neural networks. Solving these PDEs at scale is an intricate and time-intensive process that demands careful tuning. This paper introduces automated code-generation techniques specifically tailored…
▽ More
Partial differential equations (PDEs) are crucial in modeling diverse phenomena across scientific disciplines, including seismic and medical imaging, computational fluid dynamics, image processing, and neural networks. Solving these PDEs at scale is an intricate and time-intensive process that demands careful tuning. This paper introduces automated code-generation techniques specifically tailored for distributed memory parallelism (DMP) to execute explicit finite-difference (FD) stencils at scale, a fundamental challenge in numerous scientific applications. These techniques are implemented and integrated into the Devito DSL and compiler framework, a well-established solution for automating the generation of FD solvers based on a high-level symbolic math input. Users benefit from modeling simulations for real-world applications at a high-level symbolic abstraction and effortlessly harnessing HPC-ready distributed-memory parallelism without altering their source code. This results in drastic reductions both in execution time and developer effort. A comprehensive performance evaluation of Devito's DMP via MPI demonstrates highly competitive strong and weak scaling on CPU and GPU clusters, proving its effectiveness and capability to meet the demands of large-scale scientific simulations.
△ Less
Submitted 6 January, 2025; v1 submitted 20 December, 2023;
originally announced December 2023.
-
Temporal blocking of finite-difference stencil operators with sparse "off-the-grid" sources
Authors:
George Bisbas,
Fabio Luporini,
Mathias Louboutin,
Rhodri Nelson,
Gerard Gorman,
Paul H. J. Kelly
Abstract:
Stencil kernels dominate a range of scientific applications, including seismic and medical imaging, image processing, and neural networks. Temporal blocking is a performance optimization that aims to reduce the required memory bandwidth of stencil computations by re-using data from the cache for multiple time steps. It has already been shown to be beneficial for this class of algorithms. However,…
▽ More
Stencil kernels dominate a range of scientific applications, including seismic and medical imaging, image processing, and neural networks. Temporal blocking is a performance optimization that aims to reduce the required memory bandwidth of stencil computations by re-using data from the cache for multiple time steps. It has already been shown to be beneficial for this class of algorithms. However, applying temporal blocking to practical applications' stencils remains challenging. These computations often consist of sparsely located operators not aligned with the computational grid ("off-the-grid"). Our work is motivated by modeling problems in which source injections result in wavefields that must then be measured at receivers by interpolation from the grided wavefield. The resulting data dependencies make the adoption of temporal blocking much more challenging. We propose a methodology to inspect these data dependencies and reorder the computation, leading to performance gains in stencil codes where temporal blocking has not been applicable. We implement this novel scheme in the Devito domain-specific compiler toolchain. Devito implements a domain-specific language embedded in Python to generate optimized partial differential equation solvers using the finite-difference method from high-level symbolic problem definitions. We evaluate our scheme using isotropic acoustic, anisotropic acoustic, and isotropic elastic wave propagators of industrial significance. After auto-tuning, performance evaluation shows that this enables substantial performance improvement through temporal blocking over highly-optimized vectorized spatially-blocked code of up to 1.6x.
△ Less
Submitted 25 February, 2021; v1 submitted 20 October, 2020;
originally announced October 2020.
-
Scaling through abstractions -- high-performance vectorial wave simulations for seismic inversion with Devito
Authors:
Mathias Louboutin,
Fabio Luporini,
Philipp Witte,
Rhodri Nelson,
George Bisbas,
Jan Thorbecke,
Felix J. Herrmann,
Gerard Gorman
Abstract:
[Devito] is an open-source Python project based on domain-specific language and compiler technology. Driven by the requirements of rapid HPC applications development in exploration seismology, the language and compiler have evolved significantly since inception. Sophisticated boundary conditions, tensor contractions, sparse operations and features such as staggered grids and sub-domains are all su…
▽ More
[Devito] is an open-source Python project based on domain-specific language and compiler technology. Driven by the requirements of rapid HPC applications development in exploration seismology, the language and compiler have evolved significantly since inception. Sophisticated boundary conditions, tensor contractions, sparse operations and features such as staggered grids and sub-domains are all supported; operators of essentially arbitrary complexity can be generated. To accommodate this flexibility whilst ensuring performance, data dependency analysis is utilized to schedule loops and detect computational-properties such as parallelism. In this article, the generation and simulation of MPI-parallel propagators (along with their adjoints) for the pseudo-acoustic wave-equation in tilted transverse isotropic media and the elastic wave-equation are presented. Simulations are carried out on industry scale synthetic models in a HPC Cloud system and reach a performance of 28TFLOP/s, hence demonstrating Devito's suitability for production-grade seismic inversion problems.
△ Less
Submitted 22 April, 2020;
originally announced April 2020.
