High-performance in-situ ML Inference with dalotia: A Lightweight Tensor Loader API for Science Codes

With the increasing capabilities of machine-learning (ML) systems, the domains of simulation and ML are growing closer together. In particular, the evolving field of ML-based surrogate modeling within turbulence simulations often relies on explicit interfacing between scientific codes and ML frameworks. In the face of increasing hardware and software complexity, it is important to critically assess this approach and to explore alternatives that allow for a tighter integration of simulation and ML. We develop a tensor loading library, dalotia, that allows researchers to embed ML inference directly into their HPC code (“in-situ”) using C, C++, or Fortran. Using dalotia, we evaluate the in-situ approach against inference with ML frameworks (PyTorch, libTorch, and oneDNN). Our experiments using state-of-the-art ML surrogate architectures present complex tradeoffs between convenience, performance, memory usage, energy, and dependency complexity. The data shows that small-batch neural networks can save a factor of 5x in memory overheads when using the in-situ approach. We further observe in-situ speedups of up to 8x compared to the libTorch implementation on the node scale (48 cores). Thus it can be beneficial, in many regards, to rely on native in-situ implementations, without the need for bloated ML frameworks and their dependencies.