[RF] Add RooONNXFunction: ONNX-backed neural function support for RooFit

README/ReleaseNotes/v640/index.md

@@ -239,6 +239,8 @@ This is new and efficient bracketing root-finding algorithm. It combines bisecti
 
 ## RooFit
 
+### General changes
+
 - A new RooAbsPdf has been added: `RooStudentT`, which describes the location-scale student's t-distribution.
 - The `RooNumber::setRangeEpsRel()` and `RooNumber::setRangeEpsAbs()` have been
   introduced 2 years ago in 48637270a9113aa to customize range check behavior
@@ -249,6 +251,22 @@ This is new and efficient bracketing root-finding algorithm. It combines bisecti
 - The constructors of **RooDataSet** and **RooDataHist** that combine datasets via `Index()` and `Import()` now validate that the import names correspond to existing states of the index category. If an imported data slice refers to a category label that is not defined in the index category, the constructor now throws an error.
   Previously, such labels were silently added as new category states, which could lead to inconsistent datasets when the state names were not synchronized with the model definition. This change prevents the creation of invalid combined datasets and surfaces configuration problems earlier.
 
+### ONNX model integration via RooONNXFunction
+
+A new class `RooONNXFunction` has been introduced to enable the use of machine learning models in ONNX format directly within RooFit workflows.
+
+`RooONNXFunction` wraps an ONNX model as a `RooAbsReal`, allowing it to be used as a building block in likelihoods, fits, and statistical analyses without additional boilerplate code. The class supports models with one or more statically-shaped input tensors and a single scalar output.
+The class was designed to share workspaces with neural functions for combined fits in RooFit-based frameworks written in C++.
+Therefore, the `RooONNXFunction` doesn't depend on any Python packages and fully supports ROOT IO,
+
+**Key features:**
+
+  * **Compiled inference via TMVA SOFIE:** The ONNX graph is translated into optimized C++ code at runtime using SOFIE, avoiding external runtime dependencies.
+
+  * **Automatic differentiation with Clad:** Gradients of the model output with respect to RooFit parameters are generated automatically for efficient gradient-based minimization with RooFits `"codegen"` backend.
+
+  * **Portable serialization:** The ONNX model is stored as part of the `RooONNXFunction` object and serialized with ROOT I/O. Upon reading a workspace, the inference code is regenerated automatically.
+
 ### Deprecation of the the constant term optimization for legacy test statistic classes
 
 The **RooFit::Optimize()** option (constant term optimization) has been deprecated and will be removed in ROOT 6.42.

cmake/modules/RootMacros.cmake

@@ -1932,6 +1932,7 @@ endfunction()
 #                        [TIMEOUT seconds]
 #                        [COPY_TO_BUILDDIR file1 file2] Copy listed files when ctest invokes the test.
 #                        [LIBRARIES lib1 lib2...] -- Libraries to link against
+#                        [FIXTURES_SETUP ...] [FIXTURES_CLEANUP ...] [FIXTURES_REQUIRED ...]
 #                        [LABELS label1 label2...] -- Labels to annotate the test
 #                        [INCLUDE_DIRS label1 label2...] -- Extra target include directories
 #                        [REPEATS number] -- Repeats testsuite `number` times, stopping at the first failure.
@@ -1943,7 +1944,7 @@ function(ROOT_ADD_GTEST test_suite)
   cmake_parse_arguments(ARG
     "WILLFAIL"
     "TIMEOUT;REPEATS;FAILREGEX"
-    "COPY_TO_BUILDDIR;LIBRARIES;LABELS;INCLUDE_DIRS;ENVIRONMENT" ${ARGN})
+    "COPY_TO_BUILDDIR;LIBRARIES;LABELS;FIXTURES_SETUP;FIXTURES_CLEANUP;FIXTURES_REQUIRED;INCLUDE_DIRS;ENVIRONMENT" ${ARGN})
 
   ROOT_GET_SOURCES(source_files . ${ARG_UNPARSED_ARGUMENTS})
   # Note we cannot use ROOT_EXECUTABLE without user-specified set of LIBRARIES to link with.
@@ -1979,6 +1980,18 @@ function(ROOT_ADD_GTEST test_suite)
     set(extra_command --gtest_repeat=${ARG_REPEATS} --gtest_break_on_failure)
   endif()
 
+  if (ARG_FIXTURES_SETUP)
+    set(fixtures_setup ${ARG_FIXTURES_SETUP})
+  endif()
+
+  if (ARG_FIXTURES_CLEANUP)
+    set(fixtures_cleanup ${ARG_FIXTURES_CLEANUP})
+  endif()
+
+  if (ARG_FIXTURES_REQUIRED)
+    set(fixtures_required ${ARG_FIXTURES_REQUIRED})
+  endif()
+
   ROOT_PATH_TO_STRING(name_with_path ${test_suite} PATH_SEPARATOR_REPLACEMENT "-")
   string(REPLACE "-test-" "-" clean_name_with_path ${name_with_path})
   ROOT_ADD_TEST(
@@ -1989,6 +2002,9 @@ function(ROOT_ADD_GTEST test_suite)
     ${willfail}
     TIMEOUT "${ARG_TIMEOUT}"
     LABELS "${ARG_LABELS}"
+    FIXTURES_SETUP ${fixtures_setup}
+    FIXTURES_CLEANUP ${fixtures_cleanup}
+    FIXTURES_REQUIRED ${fixtures_required}
     FAILREGEX "${ARG_FAILREGEX}"
     ENVIRONMENT "${ARG_ENVIRONMENT}"
   )

math/mathcore/inc/Math/CladDerivator.h

@@ -1114,45 +1114,34 @@ inline void Gemm_Call_pullback(float *output, bool transa, bool transb, int m, i
                                bool *, int *, int *, int *, float *_d_alpha, float *_d_A, float *_d_B, float *_d_beta,
                                float *_d_C)
 {
+   using ::TMVA::Experimental::SOFIE::Gemm_Call;
+
    // TODO:
    //    - fix and test the implementation for alpha != 1.0
    if (alpha != 1.0f) {
       return;
    }
 
-   char ct = 't';
-   char cn = 'n';
-
    // beta needs to be one because we want to add to _d_A and _d_B instead of
    // overwriting it.
    float one = 1.;
 
-   // Leading dimensions for the original storage (must match how sgemm_ is called in the primal)
-   const int lda_opA = transa ? k : m; // lda used with transa flag as in primal
-   const int ldb_opB = transb ? n : k; // ldb used with transb flag as in primal
-
-   // Flags for op(A), op(B)
-   const char TA = transa ? ct : cn;
-   const char TB = transb ? ct : cn;
-
-   // Flags for op(A)^T and op(B)^T
-   const char TAT = transa ? cn : ct; // (A^T)^T = A, A^T if A
-   const char TBT = transb ? cn : ct; // (B^T)^T = B, B^T if B
-
+   // ---- dA ----
    if (!transa) {
-      // dA += alpha * dY * op(B)^T    (m x n) * (n x k) -> (m x k)
-      ::sgemm_(&cn, &TBT, &m, &k, &n, &alpha, _d_output, &m, B, &ldb_opB, &one, _d_A, &m);
+      // dA += dY * op(B)^T
+      Gemm_Call(_d_A, false, !transb, m, k, n, one, _d_output, B, one, _d_A);
    } else {
-      // dA (shape k x m) += alpha * op(B) * dY^T   (k x n) * (n x m) -> (k x m)
-      ::sgemm_(&TB, &ct, &k, &m, &n, &alpha, B, &ldb_opB, _d_output, &m, &one, _d_A, &k);
+      // dA += op(B) * dY^T
+      Gemm_Call(_d_A, transb, true, k, m, n, one, B, _d_output, one, _d_A);
    }
 
+   // ---- dB ----
    if (!transb) {
-      // dB += alpha * op(A)^T * dY   (k x m) * (m x n) -> (k x n)
-      ::sgemm_(&TAT, &cn, &k, &n, &m, &alpha, A, &lda_opA, _d_output, &m, &one, _d_B, &k);
+      // dB += op(A)^T * dY
+      Gemm_Call(_d_B, !transa, false, k, n, m, one, A, _d_output, one, _d_B);
    } else {
-      // dB (shape n x k) += alpha * dY^T * op(A)   (n x m) * (m x k) -> (n x k)
-      ::sgemm_(&ct, &TA, &n, &k, &m, &alpha, _d_output, &m, A, &lda_opA, &one, _d_B, &n);
+      // dB += dY^T * op(A)
+      Gemm_Call(_d_B, true, transa, n, k, m, one, _d_output, A, one, _d_B);
    }
 
    int sizeC = n * m;

roofit/codegen/inc/RooFit/CodegenImpl.h

@@ -48,6 +48,7 @@ class RooLandau;
 class RooLognormal;
 class RooMultiPdf;
 class RooMultiVarGaussian;
+class RooONNXFunction;
 class RooParamHistFunc;
 class RooPoisson;
 class RooPolyVar;
@@ -111,12 +112,13 @@ void codegenImpl(RooHistFunc &arg, CodegenContext &ctx);
 void codegenImpl(RooHistPdf &arg, CodegenContext &ctx);
 void codegenImpl(RooLandau &arg, CodegenContext &ctx);
 void codegenImpl(RooLognormal &arg, CodegenContext &ctx);
+void codegenImpl(RooMultiPdf &arg, CodegenContext &ctx);
 void codegenImpl(RooMultiVarGaussian &arg, CodegenContext &ctx);
+void codegenImpl(RooONNXFunction &arg, CodegenContext &ctx);
 void codegenImpl(RooParamHistFunc &arg, CodegenContext &ctx);
 void codegenImpl(RooPoisson &arg, CodegenContext &ctx);
 void codegenImpl(RooPolyVar &arg, CodegenContext &ctx);
 void codegenImpl(RooPolynomial &arg, CodegenContext &ctx);
-void codegenImpl(RooMultiPdf &arg, CodegenContext &ctx);
 void codegenImpl(RooProduct &arg, CodegenContext &ctx);
 void codegenImpl(RooRatio &arg, CodegenContext &ctx);
 void codegenImpl(RooRealIntegral &arg, CodegenContext &ctx);

roofit/codegen/src/CodegenImpl.cxx

@@ -42,6 +42,7 @@
 #include <RooLognormal.h>
 #include <RooMultiPdf.h>
 #include <RooMultiVarGaussian.h>
+#include <RooONNXFunction.h>
 #include <RooParamHistFunc.h>
 #include <RooPoisson.h>
 #include <RooPolyVar.h>
@@ -946,6 +947,21 @@ std::string codegenIntegralImpl(RooMultiVarGaussian &arg, int code, const char *
    return doubleToString(arg.analyticalIntegral(code, rangeName));
 }
 
+void codegenImpl(RooONNXFunction &arg, CodegenContext &ctx)
+{
+   std::stringstream ss;
+   ss << arg.outerWrapperName() << "(";
+   for (std::size_t i = 0; i < arg.nInputTensors(); ++i) {
+      ss << ctx.buildArg(arg.inputTensorList(i)) << std::endl;
+      if (i != arg.nInputTensors() - 1) {
+         ss << ", ";
+      }
+   }
+   ss << ")";
+
+   ctx.addResult(&arg, ss.str());
+}
+
 std::string codegenIntegralImpl(RooPoisson &arg, int code, const char *rangeName, CodegenContext &ctx)
 {
    assert(code == 1 || code == 2);

roofit/roofit/CMakeLists.txt

@@ -21,52 +21,54 @@ ROOT_STANDARD_LIBRARY_PACKAGE(RooFit
     RooBCPEffDecay.h
     RooBCPGenDecay.h
     RooBDecay.h
+    RooBMixDecay.h
     RooBernstein.h
     RooBifurGauss.h
     RooBlindTools.h
-    RooBMixDecay.h
     RooBreitWigner.h
     RooBukinPdf.h
     RooCBShape.h
-    RooCrystalBall.h
     RooCFunction1Binding.h
     RooCFunction2Binding.h
     RooCFunction3Binding.h
     RooCFunction4Binding.h
     RooChebychev.h
     RooChi2MCSModule.h
     RooChiSquarePdf.h
+    RooCrystalBall.h
     RooDecay.h
     RooDstD0BG.h
     RooExponential.h
-    RooLegacyExpPoly.h
-    RooPowerSum.h
     RooFunctor1DBinding.h
     RooFunctorBinding.h
+    RooGExpModel.h
     RooGamma.h
     RooGaussExpTails.h
-    RooGaussian.h
     RooGaussModel.h
-    RooGExpModel.h
+    RooGaussian.h
     RooHistConstraint.h
     RooIntegralMorph.h
     RooJeffreysPrior.h
+    RooJohnson.h
     RooKeysPdf.h
     RooLagrangianMorphFunc.h
     RooLandau.h
+    RooLegacyExpPoly.h
     RooLognormal.h
     RooMathCoreReg.h
+    RooMomentMorph.h
     RooMomentMorphFunc.h
     RooMomentMorphFuncND.h
-    RooMomentMorph.h
     RooMultiBinomial.h
     RooNDKeysPdf.h
     RooNonCPEigenDecay.h
     RooNovosibirsk.h
-    RooParametricStepFunction.h
+    RooONNXFunction.h
     RooParamHistFunc.h
+    RooParametricStepFunction.h
     RooPoisson.h
     RooPolynomial.h
+    RooPowerSum.h
     RooPyBind.h
     RooSpline.h
     RooStepFunction.h
@@ -80,46 +82,45 @@ ROOT_STANDARD_LIBRARY_PACKAGE(RooFit
     RooUnblindUniform.h
     RooUniform.h
     RooVoigtian.h
-    RooJohnson.h
   SOURCES
     src/Roo2DKeysPdf.cxx
     src/RooArgusBG.cxx
     src/RooBCPEffDecay.cxx
     src/RooBCPGenDecay.cxx
     src/RooBDecay.cxx
+    src/RooBMixDecay.cxx
     src/RooBernstein.cxx
     src/RooBifurGauss.cxx
     src/RooBlindTools.cxx
-    src/RooBMixDecay.cxx
     src/RooBreitWigner.cxx
     src/RooBukinPdf.cxx
     src/RooCBShape.cxx
-    src/RooCrystalBall.cxx
     src/RooCFunction1Binding.cxx
     src/RooCFunction2Binding.cxx
     src/RooCFunction3Binding.cxx
     src/RooCFunction4Binding.cxx
     src/RooChebychev.cxx
     src/RooChi2MCSModule.cxx
     src/RooChiSquarePdf.cxx
+    src/RooCrystalBall.cxx
     src/RooDecay.cxx
     src/RooDstD0BG.cxx
     src/RooExponential.cxx
-    src/RooLegacyExpPoly.cxx
-    src/RooPowerSum.cxx
     src/RooFunctor1DBinding.cxx
     src/RooFunctorBinding.cxx
+    src/RooGExpModel.cxx
     src/RooGamma.cxx
     src/RooGaussExpTails.cxx
-    src/RooGaussian.cxx
     src/RooGaussModel.cxx
-    src/RooGExpModel.cxx
+    src/RooGaussian.cxx
     src/RooHistConstraint.cxx
     src/RooIntegralMorph.cxx
     src/RooJeffreysPrior.cxx
+    src/RooJohnson.cxx
     src/RooKeysPdf.cxx
     src/RooLagrangianMorphFunc.cxx
     src/RooLandau.cxx
+    src/RooLegacyExpPoly.cxx
     src/RooLognormal.cxx
     src/RooMathCoreReg.cxx
     src/RooMomentMorph.cxx
@@ -129,10 +130,12 @@ ROOT_STANDARD_LIBRARY_PACKAGE(RooFit
     src/RooNDKeysPdf.cxx
     src/RooNonCPEigenDecay.cxx
     src/RooNovosibirsk.cxx
-    src/RooParametricStepFunction.cxx
+    src/RooONNXFunction.cxx
     src/RooParamHistFunc.cxx
+    src/RooParametricStepFunction.cxx
     src/RooPoisson.cxx
     src/RooPolynomial.cxx
+    src/RooPowerSum.cxx
     src/RooSpline.cxx
     src/RooStepFunction.cxx
     src/RooStudentT.cxx
@@ -145,7 +148,6 @@ ROOT_STANDARD_LIBRARY_PACKAGE(RooFit
     src/RooUnblindUniform.cxx
     src/RooUniform.cxx
     src/RooVoigtian.cxx
-    src/RooJohnson.cxx
   DICTIONARY_OPTIONS
     "-writeEmptyRootPCM"
   LINKDEF

roofit/roofit/inc/LinkDef1.h

@@ -77,6 +77,8 @@
                 _nonInterfering.back().emplace_back(arg->GetName()); \
              }                                                       \
           } }";
+#pragma link C++ class RooONNXFunction+ ;
+
 #pragma link C++ class RooFunctorBinding+ ;
 #pragma link C++ class RooFunctor1DBinding+ ;
 #pragma link C++ class RooFunctorPdfBinding+ ;