[RF] Add RooONNXFunction: ONNX-backed neural function support for RooFit#21928
[RF] Add RooONNXFunction: ONNX-backed neural function support for RooFit#21928guitargeek wants to merge 4 commits intoroot-project:masterfrom
Conversation
guitargeek
force-pushed
the
roofit_sofie_ad
branch
2 times, most recently
from
3e86637 to
bdab5d2
Compare
Test Results 21 files 21 suites 3d 4h 17m 32s ⏱️ Results for commit 65eaa9e. ♻️ This comment has been updated with latest results. |
guitargeek
force-pushed
the
roofit_sofie_ad
branch
3 times, most recently
from
79e32ab to
861c74e
Compare
guitargeek
force-pushed
the
roofit_sofie_ad
branch
3 times, most recently
from
08e980b to
a84d0a7
Compare
guitargeek
force-pushed
the
roofit_sofie_ad
branch
from
a84d0a7 to
5b45216
Compare
vgvassilev
left a comment
vgvassilev
left a comment
There was a problem hiding this comment.
LGTM but I'd rather have somebody else to look at this too.
| }; | ||
|
|
||
| template <class Session_t> | ||
| void doInferWithSessionVoidPtr(void *session, float const *input, float *out) |
There was a problem hiding this comment.
If that void* is to work around some current limitations of clad, let's have a comment.
There was a problem hiding this comment.
Thanks for the comments! Actually, the void* trick here is not for Codegen and Clad. It's for the default likelihood evaluation, where compiled RooFit code needs to use the Session object that is only defined in the code that is generated on-the-fly by SOFIE. So it has to use void*.
| } | ||
| } | ||
|
|
||
| } // namespace %%NAMESPACE%% |
There was a problem hiding this comment.
Why can't we have a fixed namespace name? This way we can also move the custom derivatives into a single file that's outside of generated code.
There was a problem hiding this comment.
The wrapper functions all have to be in the generated code, because the generated SOFIE session class that is only known in the interpreter is declared as <on-the-fly-generated-namespace-for-inference>::Session. By putting the wrappers in the same namespace and using the Session class, the wrappers will call into the SOFIE-generated code thanks to ADL.
guitargeek
force-pushed
the
roofit_sofie_ad
branch
from
5b45216 to
bb33802
Compare
This ensures that we use a consistent GEMM functions by wrapping the GEMM call in the pullback also with the wrapper.
This is helpful if such a test depends on input files that are created by another test.
Add some `constexpr` info in input tensor structure to the generated
code, inside the same namespace as the session struct.
It looks for example like this:
```c++
constexpr std::array<SingleDim, 1> dim_start{SingleDim{1}};
constexpr std::array<SingleDim, 1> dim_limit{SingleDim{1}};
constexpr std::array<SingleDim, 1> dim_delta{SingleDim{1}};
constexpr std::array<TensorDims, 3> inputTensorDims{
makeDims(dim_start),
makeDims(dim_limit),
makeDims(dim_delta)
};
```
This is helpful for validating inputs when using the SOFIE-generated
code without knowing much about about the original ONNX file, and is the
case for the use in RooFit.
Introduce `RooONNXFunction`, a new `RooAbsReal` implementation that enables native inference of ONNX models within RooFit. The class loads ONNX graphs, JIT-compiles them via TMVA SOFIE at runtime, and evaluates them efficiently with support for automatic differentiation through Clad. Key features: * Seamless integration of parametric neural networks into RooFit workflows * Runtime ONNX-to-C++ code generation via SOFIE (no hard dependency at link time, the SOFIE usage is an implementation detail) * Support for analytic gradients (codegen + Clad) * Support for serialization to RooWorkspace by embedding the ONNX payload as a binary blob A unit test is also implemented. This development unlocks new use cases such as neural simulation-based inference (SBI), likelihood surrogate models, and ML-driven parametric models. This commit addresses an item on the ROOT plan of work 2026.
guitargeek
force-pushed
the
roofit_sofie_ad
branch
from
bb33802 to
65eaa9e
Compare
2 participants
Introduce
RooONNXFunction, a newRooAbsRealimplementation that enables native inference of ONNX models within RooFit. The class loads ONNX graphs, JIT-compiles them via TMVA SOFIE at runtime, and evaluates them efficiently with support for automatic differentiation through Clad.Key features:
link time, the SOFIE usage is an implementation detail)
payload as a binary blob
A unit test is also implemented.
This development unlocks new use cases such as neural simulation-based inference (SBI), likelihood surrogate models, and ML-driven parametric models.
This commit addresses an item on the ROOT plan of work 2026.