RoseNNa: A performant library for portable neural network inference with application to CFD

Computational fluid dynamics practitioners have witnessed a dramatic growth in neural-network-based models for traditional closures and numerical methods. The networks are usually trained and tested in high-level languages like Python via PyTorch, TensorFlow, or other packages. Unfortunately, it is challenging to efficiently evaluate such opaque models in current large-scale CFD solvers, which are typically written in derivatives of C or Fortran. As a result, few studies of large machine-learning-assisted fluid dynamics solvers exist. We introduce a Fortran-based library called RoseNNa as a step towards solving this problem. It implements the functionality of the most common neural network architectures used or proposed for CFD. RoseNNa interprets ONNX representations of neural network models, which can be exported from most machine learning libraries and thus promotes usability. RoseNNa is linked to the user’s codebase at compile-time via usual means. The API exposes the neural network inputs and outputs to the user to be minimally invasive to existing code.