[RF] RooFit AD refactoring

[guitargeek]

1. Avoid referencing RooFuncWrapper inside code generation context
2. Take out RooFit code generation context outside of Detail namespace
3. Use dispatching via interpreter for code generation, so that the code generation methods can be defined in any place (also user frameworks)

All of this cumulates in an improved developer experience for working on the RooFit code generation, which will be valuable next week at the CMS hackathon.

Here is an example of how to override the code generation function for a class in RooFit:

// To demonstrate that even when the class is a derived class, the dispatching
// will work.
class MyGaussian : public RooGaussian {
public:
    using RooGaussian::RooGaussian;

    ClassDefOverride(MyGaussian, 0);
};

namespace RooFit {

using MyPrio = Prio<3>;

// RooFit internal functions have prio 9, so we have to use a lower number to
// override it. The prio parameter is optional and only relevant for
// overriding: if you implement codegen for a new class you don't need it.

void codegenImpl(RooGaussian &arg,
                 CodegenContext &ctx,
                 MyPrio p)
{

   std::cout << "User override" << std::endl;

   // Fall through to next implementation
   return codegenImpl(arg, ctx, p.next());
}

std::string codegenIntegralImpl(RooGaussian &arg,
                                int code,
                                const char *rangeName,
                                CodegenContext &ctx,
                                MyPrio p)
{
   std::cout << "User override (integral)" << std::endl;

   // Fall through to next implementation
   return codegenIntegralImpl(arg, code, rangeName, ctx, p.next());
}

} // namespace RooFit

void repro()
{
   RooRealVar x("x", "x", -10, 10);
   RooRealVar mean("mean", "mean of gaussian", 1, -10, 10);
   RooRealVar sigma("sigma", "width of gaussian", 1, 0.1, 10);
   MyGaussian gauss("gauss", "gaussian PDF", x, mean, sigma);

   std::unique_ptr<RooAbsReal> integ{gauss.createIntegral(x)};

   // Generate function for unnormalized Gaussian and the integral
   std::string funcName = RooFit::CodegenContext{}.buildFunction(gauss);
   std::string funcNameInteg = RooFit::CodegenContext{}.buildFunction(*integ);

   // Check out generated code
   gInterpreter->ProcessLine(funcName.c_str());
   gInterpreter->ProcessLine(funcNameInteg.c_str());
}

Output:

User override
User override (integral)
.... some code dump ....

[github-actions]

Test Results

    18 files      18 suites   4d 2h 32m 0s ⏱️
 2 678 tests  2 678 ✅ 0 💤 0 ❌
46 342 runs  46 342 ✅ 0 💤 0 ❌

Results for commit 63c0cb8.

♻️ This comment has been updated with latest results.

[egpbos]

Just thinking about this a bit after the Stats meeting yesterday, I was wondering: does the dispatching actually work well with class hierarchies? Or would a user need to override for every subclass separately as well? Possibly this is not a problem in practice for RooFit, if the user would mainly be interested in modifying the codegen implementation functions for "final" classes. But I was just considering it in the broader context of emulating multiple dispatch :)

Note btw that there are also existing libraries that do multiple dispatch emulation, like https://github.com/jll63/yomm2. That one also requires some boilerplate (e.g. registration), but your solution does too (the priority flags), so may be interesting to weigh those against each other. Note that I have no strong opinion either way; I'm just interested to see how this turns out, because multiple dispatch is awesome 😄

[guitargeek]

Hi @egpbos, thank you very much for your comment!

I would like to not introduce further dependencies to ROOT, so let's try to implement this with the interpreter that we already have :)

does the dispatching actually work well with class hierarchies?

No, it doesn't. I have thought about this quite a bit, and realized that this is sort of hard to implement at the same time, at least how it is done now.

For example, if you have MyGaussian deriving from RooGaussian, which overload would you expect to hit?

Priority<4>, RooGaussian&
Priority<5>, MyGaussian& // higher number means lower priority

I would expect it to resolve to the second one. However, right now in my implementation, "Prio" takes priority over argument type matching, so this is unintuitive. This is hard to unify without additional boilerblate....

One solution would be to just leave how things are in this PR, which would mean that you have to write some boilerplate when deriving classes where you want to keep the code generation the same:

void codegenImpl(Prio<0> p, MyGaussian &arg, CodgenContext &ctx) {
   codegenImpl(p, static_cast<RooGaussian&>(arg), ctx);
}

The other solution is to drop the priority mechanism, but I think it is quite useful...

[guitargeek]

...that one also requires some boilerplate (e.g. registration), but your solution does too (the priority flags)

Just to clarify: the priority flags are not boilerplate but a feature. They can also be dropped, but then it would not be possible anymore to "override" behavior for existing classes in usercode.

[guitargeek]

Or the ambiguity I mentioned can be resolved by requiring that priority has to decrease (e.g. higher priority) with depth in the class hierarchy. This can't be enforced, but it can be documented.

[guitargeek]

I have now implemented the comments I have received by @egpbos and also @vgvassilev in private:

1. If you create a new derived class, it will correctly dispatch to the function that takes the base class (see demo in the PR description)
2. The Prio argument is now at the final position and optional

[lmoneta]

Thank you Jonas for this big refactoring and improvement in the code generation for AD.

LGTM!

[vgvassilev]

Can we add a few tests checking the expected scenarios. This would help us see the bigger picture.

[egpbos]

Ah that's right, the priority mechanism is a separate thing. Your idea about documenting "sensible" priority ordering makes sense too. Btw I think it's actually not even a thing for Julia, because there is no inheritance there.

Edit: oh weird, I only now see all your later comments. The above was a response only to the comments between 1 and 2 PM 😄

[egpbos]

I really like this concept. Maybe there are some details to iron out. Perhaps consider putting this in an Experimental namespace, or is it already?

[egpbos]

Do I understand correctly that the Caller first tries to call the prio 1 (highest) implementation, and then will loop through all 10 using the Impl template until it finds one that is actually implemented for the Arg_t already so doesn't need to use the template?

Few questions:

1. Do you actually need all 10 priority levels? How about replacing that with just an overload boolean? It will save some looping through empty functions.
2. I'm not sure, but there may be some pathological situations. For instance, what if in an interactive session you call some function, the tree of templates gets instantiated and then afterwards you try to define a new overload? Won't that overload then already have been instantiated from the Impl template? Another possible situation would be something similar to the above but in separate compilation units.

[guitargeek]

Great comments!

1. I was hoping the optimizer gets rid of these loop levels, which are "statically" clear to the interpreter. So in principle they should be optimized out. But it's an interesting question: I'll benchmark how expensive the levels are
2. You're right, this is a clear limitation, but design: the function pointer that was resolved is cached in an unordered_map, so speed up the lookup for the same class later. I'll check what would happen if the caching is disabled, that's an interesting question!