[tmva][sofie] Restructure emitted code to be differentiable with Clad

[guitargeek]

Restructure emitted code to be differentiable with Clad.

[vgvassilev]

Does providing a pullback for std::copy not work?

[guitargeek]

No, I tried a bit but then gave up. This was my approach:

#include <Math/CladDerivator.h>

namespace std {

void copy_pullback(double const *first, double const *last, double *out_first, double *_d_out, double *_d_first,
                   double *_d_last, double *_d_out_first)
{
   // Implementation doesn't matter yet, it doesn't compile anyway
}

} // namespace std

void fooImpl(double const *x, double *y)
{
   std::copy(x, x + 1, y);
}

void foo(double const *x, double *y)
{
   fooImpl(x, y);
}

double g(double *variables)
{
   double out;
   foo(variables, &out);
   return out * variables[1];
}

void clademo()
{

   // Call clad to generate the gradient of g.
   auto g_grad = clad::gradient(g, "variables");

   // Execute the generated gradient function.
   double variables[]{3., 4.};
   double grad_output[]{0., 0.};
   g_grad.execute(variables, grad_output);
   std::cout << "grad_output[0]: " << grad_output[0] << std::endl;
   std::cout << "grad_output[1]: " << grad_output[1] << std::endl;

   // Dump the generated gradient code to standard output.
   g_grad.dump();
}

It segfaults. I think Clad just doesn't play well with the STL algos that take iterators, it's better to avoid this, no?

In any case, supporting this is not crucial for this PR. I was refactoring things to avoid this copy call in the generated code anyway.

Once this PR is functional for our usecase (actually now, but I also want to make the ROOT CI pass again), I'll write up what was not perfect in Clad for this and open issues

[vgvassilev]

Ah, I see. Probably worth opening an issue in clad…

[vgvassilev]

Did that fail to differentiate?

[guitargeek]

No, I didn't even try to differentiate the models in the test. I'm solely focusing on the SBI usecase that we implement with LHCb. The reason why I changed this is because std::vector<bool> is not a good output type parameter. See:

• [TMVA][SOFIE] Use uint8_t instead of bool in return types #18302

[github-actions]

Test Results

    21 files      21 suites   3d 1h 22m 35s ⏱️
 3 810 tests  3 809 ✅ 1 💤 0 ❌
73 036 runs  73 027 ✅ 9 💤 0 ❌

Results for commit 02aa923.

♻️ This comment has been updated with latest results.

[guitargeek]

Proof of concept test for this PR

Take this ONNX file (remove the .txt suffix after downloading):

VRlL_real_500k_evts_model.onnx.txt

Here are the scripts to convert the model to C++ and then to differentiate it with Clad:

// onnx_to_cpp.C

void onnx_to_cpp()
{
   using namespace TMVA::Experimental;
   SOFIE::RModelParser_ONNX parser;
   SOFIE::RModel model = parser.Parse("./VRlL_real_500k_evts_model.onnx");
   model.SetOptimizationLevel(SOFIE::OptimizationLevel::kBasic);
   model.Generate();
   model.PrintRequiredInputTensors();

   model.OutputGenerated("./VRlL_real_500k_evts_model.hxx");
}

// sofie_ad.C

#include "VRlL_real_500k_evts_model.hxx"

#include <Math/CladDerivator.h>

float my_func(TMVA_SOFIE_VRlL_real_500k_evts_model::Session const *session, float const *tensor_x,
              float *tensor_theory_params)
{
   float out = 0.;
   TMVA_SOFIE_VRlL_real_500k_evts_model::doInfer(session, tensor_x, tensor_theory_params, &out);
   return out;
}

void sofie_ad()
{
   std::vector<float> input1{5.0, 2.0, 1.0, -1.0, 1.0};
   std::vector<float> input2{0.0};

   // Generated header file shall contain a Session class which requires
   // initialization to load the corresponding weights.
   TMVA_SOFIE_VRlL_real_500k_evts_model::Session s("VRlL_real_500k_evts_model.dat");

   // Once instantiated the session object's infer method can be used
   // std::vector<float> out = s.infer(input1.data(), input2.data());

   auto func = [&](std::span<float> params) { return s.infer(input1.data(), params.data())[0]; };

   auto numDiff = [&](int i) {
      const float eps = 1e-4;
      std::vector<float> p{input2};
      p[i] = input2[i] - eps;
      float funcValDown = func(p);
      p[i] = input2[i] + eps;
      float funcValUp = func(p);
      return (funcValUp - funcValDown) / (2 * eps);
   };

   for (std::size_t i = 0; i < input2.size(); ++i) {
      std::cout << i << ":" << std::endl;
      std::cout << "  numr : " << numDiff(i) << std::endl;
   }

   float grad_output[]{0., 0., 0., 0., 0.};
   auto g_grad = clad::gradient<clad::opts::disable_tbr>(my_func, "tensor_theory_params");
   g_grad.execute(&s, input1.data(), input2.data(), grad_output);
   std::fill(std::begin(grad_output), std::end(grad_output), 0);
   g_grad.execute(&s, input1.data(), input2.data(), grad_output);

   std::cout << "  clad : " << grad_output[0] << std::endl;

   g_grad.dump();
}

Note that clad::opts::disable_tbr can probably be removed when this Clad issue is fixed:

• Another regression in Clad v1.10 with new crash in code that worked before vgvassilev/clad#1369

Usage with expected output (replace libblas.so location with relevant path for your system):

   ------------------------------------------------------------------
  | Welcome to ROOT 6.35.01                        https://root.cern |
  | (c) 1995-2024, The ROOT Team; conception: R. Brun, F. Rademakers |
  | Built for linuxx8664gcc on Jan 01 1980, 00:00:00                 |
  | From heads/sofie_ad@v6-35-01-2277-g8ddebb98bb                    |
  | With g++ (GCC) 14.2.1 20250322                                   |
  | Try '.help'/'.?', '.demo', '.license', '.credits', '.quit'/'.q'  |
   ------------------------------------------------------------------

root [0] .L /nix/store/6kknwpcf8fl7ihkkxmdb6p764kdn443n-blas-3/lib/libblas.so
root [1] .x onnx_to_cpp.C
Model requires following inputs:
Fully Specified Tensor name: theory_params	type: float	shape: [1]
Fully Specified Tensor name: x	type: float	shape: [5]

root [2] .x sofie_ad.C
0:
  numr : -0.531077
  clad : -0.532437
root [3] .q

[vgvassilev]

Why did we decide to not pursue this?

[guitargeek]

@vgvassilev, sorry that was totally an accident. Maybe I confused it with another PR, or I wanted to close and re-open the PR to run the tests, but apparently I missed the "reopen" button.

[guitargeek]

So it works now, and even includes a test to validate the gradient of a fully-connected multi-layer network! I'll organize the changes a bit better before marking this as not a draft.