CodeJeNN: A simple C++ neural network generator for physics applications

Technical Summary: CodeJeNN

Problem Statement
Machine learning has demonstrated significant potential for accelerating numerical methods in physics, particularly in computational fluid dynamics (CFD) for tasks such as turbulence modeling and surrogate constitutive models. However, a critical bottleneck exists in integrating Python-based machine learning libraries (e.g., TensorFlow, PyTorch, Keras) into high-performance, compiled C++ or Fortran solvers. This disparity introduces computational overhead and complicates the deployment of neural networks in established physics frameworks. Furthermore, relying on framework-provided C++ APIs often limits optimization capabilities, as it forces tensor operations into external libraries that may conflict with existing parallelization and vectorization strategies. While previous solutions like RoseNNa exist for Fortran and ONNX models, there remains a need for a tool that generates self-contained, dependency-free C++ code specifically for Keras models to facilitate seamless integration without external runtime dependencies.

Methodology
The authors present CodeJeNN, a C++ neural network generator that bridges the gap between trained Keras models and compiled physics solvers. The workflow operates as follows:

1. Model Training: Users train a neural network (MLP or CNN) using Keras and save the architecture and parameters to a .keras or .h5 file.
2. Code Generation: CodeJeNN loads the model, extracts parameters into a Python dictionary, and reconstructs the model logic.
3. C++ Synthesis: The tool generates a self-contained C++ header file. Key design principles include:
   • Inlining: Aggressive inlining of activation and propagation functions to minimize function call overhead.
   • Dependency Reduction: The generated code uses only standard C++ libraries (e.g., <iostream>, <array>, <cmath>), eliminating the need for external linear algebra libraries like LAPACK or Eigen.
   • Compile-Time Optimization: The use of constexpr, static, and inline keywords ensures that array sizes and operations are resolved at compile time, optimizing cache utilization.
   • Extensibility: The source structure is designed to allow Large Language Models (LLMs) to generate logic for unsupported architectures with minimal token consumption, facilitating non-invasive framework extension.

The tool supports a wide range of Keras layers (Dense, Convolutional, Pooling, Normalization, etc.) and activation functions (ReLU, GELU, Swish, etc.), including user-defined custom activations provided they are serialized and defined in a Python script.

Key Contributions

• Self-Contained Code Generation: CodeJeNN produces a single C++ header file containing all necessary weights, biases, and logic, allowing direct inclusion into existing solvers without external runtime dependencies.
• Performance Optimization: By avoiding external API calls and leveraging aggressive inlining, the generated code is designed to execute efficiently on local hardware.
• Debugging Capabilities: The tool includes a debugging API that compares the extracted model against Keras' model.summary() and, when enabled, prints intermediate layer outputs to help pinpoint logic or algebraic errors during inference.
• Extensibility via LLMs: The codebase is structured to minimize the token cost and effort required to add support for new architectures using LLMs.

Results
The paper validates CodeJeNN through two primary benchmarks:

1. Inference Speed Benchmarks: Tests were conducted on an Apple M1 Pro comparing CodeJeNN against Keras running on TensorFlow and PyTorch backends (both eager and JIT-compiled modes).
   • For a Multi-Layer Perceptron (MLP), CodeJeNN achieved a 25x speedup over standard TensorFlow and a 38x speedup over standard PyTorch.
   • Even against JIT-compiled Keras models, CodeJeNN remained approximately 2.5x faster (TensorFlow) and 9x faster (PyTorch).
   • For Convolutional Neural Networks (CNNs), speedups were substantial (18x–26x over standard Keras), though performance degraded for larger models as they transitioned from compute-bound to memory-bound due to cache limitations.
2. CFD Application (Viscosity Modeling): A neural network was embedded into an OpenFOAM-based solver (DetonationFOAM) to replace Wilke's mixing rule for calculating viscosity in a hydrogen-air mixing layer.
   • The network, trained on 10,000 states, was embedded directly into the solver's C++ headers.
   • In simulations, the neural network surrogate introduced a maximum error of 0.25% in the velocity field compared to the native Wilke model, preserving the overall flow evolution and instability development.

Significance and Claims
The authors claim that CodeJeNN significantly lowers the barrier for physics researchers to integrate neural networks into high-performance solvers. By generating raw, inlined source code, it eliminates the need for complex infrastructure or reliance on external tensor libraries, which is crucial for codes that utilize all available hardware resources. The tool addresses the inefficiency of using LLMs to repeatedly generate infrastructure code by providing a dedicated, repeatable pipeline. The paper positions CodeJeNN as a practical solution for scientific computing where models must remain compact, efficient, and free of external dependencies to satisfy real-time and hardware constraints. Future work includes integrating Graph Neural Network (GNN) support via Spektral and exploring low-level library exports (e.g., Eigen) for further acceleration.