Reduce amount of code JITted by RDataFrame #17282
Conversation
|
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. |
Test Results 22 files 22 suites 3d 12h 35m 30s ⏱️ Results for commit c91b372. ♻️ This comment has been updated with latest results. |
vepadulano
force-pushed
the
rdf_jit_recall_forpr
branch
from
8379eb8 to
8136e27
Compare
|
I've taken the liberty of rebasing this PR on master. Locally, I have tried to exploit these changes as a baseline to remove the controlled leaks of #15520 , but I've hit a blocker because of |
vepadulano
force-pushed
the
rdf_jit_recall_forpr
branch
3 times, most recently
from
70d919c to
e5f2ff9
Compare
vepadulano
force-pushed
the
rdf_jit_recall_forpr
branch
from
c677e78 to
3c3b1ee
Compare
|
@gpetruc I have squashed your proposed changes in one commit, the second commit introduces improvements that also lead to fixing #15520. Our CI is running well with the changes, let me know if you have time to check this PR again with your analysis, I'll be running some tests with valgrind locally as well. |
|
Hi @vepadulano, Thanks for reviving and updating this PR, I've run a quick check with one of my analysis benchmarks and it looks ok. I'm failing to do more tests because for reasons I don't yet understand the el9 docker containers I'm building with this root branch or the master one fail to access eos ( |
|
Okay, after fixing my docker builds now I managed to run also more tests with this PR, and confirm that results look ok. |
These changes introduce a refactoring of the way JIT calls for RDataFrame nodes are handled. Previously, every time a certain node creation corresponded to a JIT call, RDataFrame would declare a line of code to JIT that node. This would possibly lead to calling the same JIT code multiple times, if the same function was to be JITted by many computation graphs. This commit introduces the concept of a "deferred JIT call", which refactors the approach in three steps: 1. Every time a new node needs to be JITted, RDataFrame will declare a registrator function, in the body of which the true function to JIT will be called. The registrator is identified by a sequentially increasing number, together with its whole function body. 2. In RLoopManager::Jit, a cache is accessed to retrieve the function pointers of the registrator functions. 3. Run the registrators, once per computation graph. This practically achieves a JIT reduction by disambiguating the same registrator function that can be reused by multiple computation graphs. 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).
vepadulano
force-pushed
the
rdf_jit_recall_forpr
branch
from
3c3b1ee to
1a82815
Compare
|
Tested the changes of this PR with our rootbench RDataFrame benchmarks at https://github.com/root-project/rootbench/blob/master/root/tree/dataframe/RDataFrameBenchmarks.cxx . On my machine, the overall execution time for all the benchmarks lowers: master: patch: In particular, certain cases see a very large improvement. This example with JITting of a large graph (O(10K) JIT nodes with reusage of the same 2 functions) sees a factor 2 reduction, going from 14 seconds to 7 seconds. master patch: I tried offline and the difference would be even more dramatic if the computation graph functions would be reused not only within the same computation graph, but also across computation graphs, which is the main purpose of this PR. I attach the CSV outputs from rootbench in case anyone is interested I believe this PR is good for a final review, adding the other RDF devs @martamaja10 @hageboeck |
vepadulano
changed the title
|
@hageboeck I have addressed your comments in a separate commit, will squash later on 👍 |
|
Following offline discussion, here is an example showing the difference in terms of amount of code seen by the interpreter before/after this PR. Take the following code cat jit_large_graphs.cpp
#include <ROOT/RDataFrame.hxx>
#include <ROOT/RLogger.hxx>
#include <iostream>
#include <vector>
#include <string>
auto verbosity = ROOT::RLogScopedVerbosity(ROOT::Detail::RDF::RDFLogChannel(), ROOT::ELogLevel::kDebug + 10);
void jit_large_graph()
{
ROOT::RDataFrame root(1);
std::vector<ROOT::RDF::RResultPtr<double>> dfs;
const auto n = 2;
for (auto i = 0u; i < n; i++) {
const auto column = "x" + std::to_string(i);
auto define = root.Define(column, "42.f");
auto filter = define.Filter(column + " > 0.f");
auto sum = filter.Sum(column);
dfs.emplace_back(sum);
}
// Trigger event loop
dfs[0].GetValue();
}
int main()
{
// Run twice to see the benefit of reusing the same JITted node registrators
jit_large_graph();
jit_large_graph();
}Executing the code will print, among other things, the code that is JITted. When running this program with ROOT master, we can see (just extracting interesting parts here) These come from the strings passed to the These lines are repeated for every new RDataFrame. Instead, when running the same program after this PR, we have These are the same as before. Then we have That is, there is one JIT node registrator line per function passed by the user (or per type of action, e.g. in the case of I have repeated the test with extracted_patch.txt And just for completeness, here is the full list of commands to get to those results |
This commit leverages the newly created DeferredJitCall infrastructure in RLoopManager to avoid the need for propagating heap-allocated objects through the JITted nodes flow, thus fixing root-project#15520. * When creating a new JITted node, the (shared) ownership of the previous node is now passed to the JITted node itself, which will then release it to the concrete node after JITting. This in particular happens in JitFilterHelper, CallBuildAction. The DeferredJitCall objects have no knowledge about the previous node anymore. * The column register is owned by DeferredJitCall via std::unique_ptr. * DeferredJitCall now has shared ownership of the JITted node itself and eventually extra arguments passed to the call. Thanks to these changes, it is not necessary to call RLoopManager::Jit in the RDataFrame destructor anymore, thus that call is removed. A few improvements were added to the new logic that runs JIT calls lazily, namely: * The function id is now stored as a std::size_t value. * The function body is hashed before being inserted in the map.
When accumulating code to JIT for Define, Filter and Vary nodes, the
current logic embeds the name of e.g. the new column to define as part
of the function body that will go into the accumulated code to JIT. This
in turn means that every time we call a Define, there will be a
corresponding code to JIT, as well as a corresponding function body that
will go into the newly created cache for the JIT node registrators.
This can be optimised by avoiding to store the names as part of the
function body, thus allowing a much more aggressive caching. In a simple
example such as:
```
ROOT::RDataFrame root(1);
std::vector<ROOT::RDF::RResultPtr<double>> dfs;
const auto n = 10000;
for (auto i = 0u; i < n; i++) {
const auto column = "x" + std::to_string(i);
auto define = root.Define(column, "42.f");
auto filter = define.Filter(column + " > 0.f");
auto sum = filter.Sum(column);
dfs.emplace_back(sum);
}
dfs[n - 1].GetValue();
```
This commit changes the previous situation that would have had 10K JIT
function bodies, i.e. one per Define with a different column name, down
to one function body to JIT for the lambda returning 42.f.
vepadulano
force-pushed
the
rdf_jit_recall_forpr
branch
from
d9356ec to
c91b372
Compare
4 participants
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
BookDefineJitthat craft the code as a string and pass it toRLoopManager::toJitExec. Calls are accumulated, and executed in one go byRLoopManager::Jit()With this PR
Methods like
BookDefineJitnow call a new functionRLoopManager::RegisterJitHelperCalland the whole operation is refactored in three stepsgInterpreterSimilarly 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 insideRLoopManager::Jit()The deferred function calls (3) are done both in
RLoopManager::Jit()and also inRLoopManager::Runin 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?
BookDefineJitand deleted in the JITed calls likeJitDefineHelpercould be instead owned by the deferred function call object via RAIIvoid *for the weak pointer, the extra column names for variations could be in a smart pointer, ...