[df] [for discussion] JIT graph creation functions only once

[gpetruc]

This PR

This is some code, that could also be improved further, to reduce the per-sample JIT cost in RDataFrame, and potentially a way to avoid the "controlled memory leaks" #15520 (though this code currently does not yet address them in full).

I don't propose to merge this code as is, but I'm opening the PR as a reference for future discussion

Before this PR

Currently JITed nodes of the computation graph, RDF JITs the function code just once, but JITs code lines to create the computational graph once per graph, with code like

ROOT::Internal::RDF::JitDefineHelper<ROOT::Internal::RDF::DefineTypes::RDefineTag>(
        R_rdf::func45, 
        new const char*[3]{"Jet_isSelected", "Jet_area", "Jet_pt"}, 3, 
        "SelJet_area",
        reinterpret_cast<ROOT::Detail::RDF::RLoopManager*>(0x562b52fe7cb0),
        reinterpret_cast<std::weak_ptr<ROOT::Detail::RDF::RJittedDefine>*>(0x562b5530d120),
        reinterpret_cast<ROOT::Internal::RDF::RColumnRegister*>(0x562b552e1900),
        reinterpret_cast<std::shared_ptr<ROOT::Detail::RDF::RNodeBase>*>(0x562b552e0d20));

This is done by methods like BookDefineJit that craft the code as a string and pass it to RLoopManager::toJitExec. Calls are accumulated, and executed in one go by RLoopManager::Jit()

With this PR

Methods like BookDefineJit now call a new function RLoopManager::RegisterJitHelperCall and the whole operation is refactored in three steps

1. a definition via JIT of a registrator function, done once per function, kind of action (e.g. Define) and column name

namespace R_rdf {
  void jitNodeRegistrator_70(
        ROOT::Detail::RDF::RLoopManager *lm, 
        std::shared_ptr<ROOT::Detail::RDF::RNodeBase> *prevNodeOnHeap,
        ROOT::Internal::RDF::RColumnRegister* colRegister, 
        const std::vector<std::string> & colNames, 
        void *wkJittedDefine, void *) {
                 ROOT::Internal::RDF::JitDefineHelper<ROOT::Internal::RDF::DefineTypes::RDefineTag>(R_rdf::func45,
                 colNames, 
                 "SelJet_area", 
                 lm, reinterpret_cast<std::weak_ptr<ROOT::Detail::RDF::RJittedDefine>*>(wkJittedDefine),
                 colRegister, 
                 prevNodeOnHeap);
        }
}

2. lookup of the function address in memory, via calls to gInterpreter
3. from C++ call the registrator function once per computation graph

Similarly to the current setup, the declarations from point (1) are accumulated in a single string and passed to the interpreter in one go by RLoopManager::Jit(), and the address lookup (2) is also done inside RLoopManager::Jit()

The deferred function calls (3) are done both in RLoopManager::Jit() and also in RLoopManager::Run in order to also allow doing them in multiple threads when RunGraphs is used (RunGraphs calls Jit only on one of the loop managers, and then calls Run on all multithreaded)

What next?

• I tested the code on several RDF analyses, but it could be tested more extensively, and also not just on x86_64.
• The column name could also be gathered in the registration, to reduce the number of registrator functions to JIT
• The objects created on the heap by BookDefineJit and deleted in the JITed calls like JitDefineHelper could be instead owned by the deferred function call object via RAII
• Code could be cleaned further, e.g. the registration key could become an integer and the map a concurrent container to avoid the read lock, a c++ std::variant could also be used instead of a void * for the weak pointer, the extra column names for variations could be in a smart pointer, ...

[vepadulano]

Thank you @gpetruc for kickstarting this absolutely non-trivial change to a core part of RDataFrame! It is definitely worth discussing, I will take a look at this.

[github-actions]

Test Results

    18 files      18 suites   4d 3h 24m 3s ⏱️
 2 679 tests  2 674 ✅ 0 💤  5 ❌
46 502 runs  46 481 ✅ 0 💤 21 ❌

For more details on these failures, see this check.

Results for commit 8379eb8.