From d8f850574b74a355410a7cb2e66082a4e33639a0 Mon Sep 17 00:00:00 2001 From: Giovanni Date: Sat, 7 Dec 2024 13:43:13 +0100 Subject: [PATCH 1/3] [df] Register JIT calls only once across many computation graphs 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). --- .../dataframe/inc/ROOT/RDF/InterfaceUtils.hxx | 36 ++-- .../dataframe/inc/ROOT/RDF/RInterfaceBase.hxx | 9 +- tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx | 26 +++ tree/dataframe/src/RDFInterfaceUtils.cxx | 158 ++++++++---------- tree/dataframe/src/RLoopManager.cxx | 109 +++++++++++- 5 files changed, 211 insertions(+), 127 deletions(-) diff --git a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx b/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx index acfe71b5d9f12..4fdc846a701fb 100644 --- a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx +++ b/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx @@ -385,10 +385,9 @@ BookVariationJit(const std::vector &colNames, std::string_view vari RDataSource *ds, const RColumnRegister &colRegister, std::shared_ptr *upcastNodeOnHeap, bool isSingleColumn); -std::string JitBuildAction(const ColumnNames_t &bl, std::shared_ptr *prevNode, - const std::type_info &art, const std::type_info &at, void *rOnHeap, TTree *tree, +std::string JitBuildAction(const ColumnNames_t &bl, const std::type_info &art, const std::type_info &at, TTree *tree, const unsigned int nSlots, const RColumnRegister &colRegister, RDataSource *ds, - std::weak_ptr *jittedActionOnHeap, const bool vector2RVec = true); + const bool vector2RVec = true); // Allocate a weak_ptr on the heap, return a pointer to it. The user is responsible for deleting this weak_ptr. // This function is meant to be used by RInterface's methods that book code for jitting. @@ -473,7 +472,7 @@ void AddDSColumns(const std::vector &requiredCols, ROOT::Detail::RD // this function is meant to be called by the jitted code generated by BookFilterJit template -void JitFilterHelper(F &&f, const char **colsPtr, std::size_t colsSize, std::string_view name, +void JitFilterHelper(F &&f, const ColumnNames_t &cols, std::string_view name, std::weak_ptr *wkJittedFilter, std::shared_ptr *prevNodeOnHeap, RColumnRegister *colRegister) noexcept { @@ -486,9 +485,6 @@ void JitFilterHelper(F &&f, const char **colsPtr, std::size_t colsSize, std::str return; } - const ColumnNames_t cols(colsPtr, colsPtr + colsSize); - delete[] colsPtr; - const auto jittedFilter = wkJittedFilter->lock(); // mock Filter logic -- validity checks and Define-ition of RDataSource columns @@ -538,7 +534,7 @@ auto MakeDefineNode(DefineTypes::RDefinePerSampleTag, std::string_view name, std // This function is meant to be called by jitted code right before starting the event loop. // If colsPtr is null, build a RDefinePerSample (it has no input columns), otherwise a RDefine. template -void JitDefineHelper(F &&f, const char **colsPtr, std::size_t colsSize, std::string_view name, RLoopManager *lm, +void JitDefineHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RLoopManager *lm, std::weak_ptr *wkJittedDefine, RColumnRegister *colRegister, std::shared_ptr *prevNodeOnHeap) noexcept { @@ -547,7 +543,6 @@ void JitDefineHelper(F &&f, const char **colsPtr, std::size_t colsSize, std::str delete wkJittedDefine; delete colRegister; delete prevNodeOnHeap; - delete[] colsPtr; }; if (wkJittedDefine->expired()) { @@ -557,15 +552,13 @@ void JitDefineHelper(F &&f, const char **colsPtr, std::size_t colsSize, std::str return; } - const ColumnNames_t cols(colsPtr, colsPtr + colsSize); - auto jittedDefine = wkJittedDefine->lock(); using Callable_t = std::decay_t; using ColTypes_t = typename TTraits::CallableTraits::arg_types; auto ds = lm->GetDataSource(); - if (ds != nullptr && colsPtr) + if (ds != nullptr) AddDSColumns(cols, *lm, *ds, ColTypes_t(), *colRegister); // will never actually be used (trumped by jittedDefine->GetTypeName()), but we set it to something meaningful @@ -580,18 +573,14 @@ void JitDefineHelper(F &&f, const char **colsPtr, std::size_t colsSize, std::str } template -void JitVariationHelper(F &&f, const char **colsPtr, std::size_t colsSize, const char **variedCols, - std::size_t variedColsSize, const char **variationTags, std::size_t variationTagsSize, - std::string_view variationName, RLoopManager *lm, - std::weak_ptr *wkJittedVariation, RColumnRegister *colRegister, - std::shared_ptr *prevNodeOnHeap) noexcept +void JitVariationHelper(F &&f, const ColumnNames_t &inputColNames, const ColumnNames_t &variedColNames, + const char **variationTags, std::size_t variationTagsSize, std::string_view variationName, + RLoopManager *lm, std::weak_ptr *wkJittedVariation, + RColumnRegister *colRegister, std::shared_ptr *prevNodeOnHeap) noexcept { // a helper to delete objects allocated before jitting, so that the jitter can share data with lazily jitted code auto doDeletes = [&] { - delete[] colsPtr; - delete[] variedCols; delete[] variationTags; - delete wkJittedVariation; delete colRegister; delete prevNodeOnHeap; @@ -604,8 +593,6 @@ void JitVariationHelper(F &&f, const char **colsPtr, std::size_t colsSize, const return; } - const ColumnNames_t inputColNames(colsPtr, colsPtr + colsSize); - std::vector variedColNames(variedCols, variedCols + variedColsSize); std::vector tags(variationTags, variationTags + variationTagsSize); auto jittedVariation = wkJittedVariation->lock(); @@ -628,13 +615,12 @@ void JitVariationHelper(F &&f, const char **colsPtr, std::size_t colsSize, const /// Convenience function invoked by jitted code to build action nodes at runtime template -void CallBuildAction(std::shared_ptr *prevNodeOnHeap, const char **colsPtr, std::size_t colsSize, +void CallBuildAction(std::shared_ptr *prevNodeOnHeap, const ColumnNames_t &cols, const unsigned int nSlots, std::shared_ptr *helperArgOnHeap, std::weak_ptr *wkJittedActionOnHeap, RColumnRegister *colRegister) noexcept { // a helper to delete objects allocated before jitting, so that the jitter can share data with lazily jitted code auto doDeletes = [&] { - delete[] colsPtr; delete helperArgOnHeap; delete wkJittedActionOnHeap; // colRegister must be deleted before prevNodeOnHeap because their dtor needs the RLoopManager to be alive @@ -650,8 +636,6 @@ void CallBuildAction(std::shared_ptr *prevNodeOnHeap, const char * return; } - const ColumnNames_t cols(colsPtr, colsPtr + colsSize); - auto jittedActionOnHeap = wkJittedActionOnHeap->lock(); // if we are here it means we are jitting, if we are jitting the loop manager must be alive diff --git a/tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx b/tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx index 7c9a0985cd591..03546dbaf25c2 100644 --- a/tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx @@ -201,10 +201,11 @@ protected: fColRegister, proxiedPtr->GetVariations()); auto jittedActionOnHeap = RDFInternal::MakeWeakOnHeap(jittedAction); - auto toJit = RDFInternal::JitBuildAction(validColumnNames, upcastNodeOnHeap, typeid(HelperArgType), - typeid(ActionTag), helperArgOnHeap, nullptr, nSlots, fColRegister, - GetDataSource(), jittedActionOnHeap, vector2RVec); - fLoopManager->ToJitExec(toJit); + auto definesCopy = new RDFInternal::RColumnRegister(fColRegister); // deleted in jitted call + auto funcBody = RDFInternal::JitBuildAction(validColumnNames, typeid(HelperArgType), typeid(ActionTag), nullptr, + nSlots, fColRegister, GetDataSource(), vector2RVec); + fLoopManager->RegisterJitHelperCall(funcBody, upcastNodeOnHeap, definesCopy, validColumnNames, jittedActionOnHeap, + helperArgOnHeap); return MakeResultPtr(r, *fLoopManager, std::move(jittedAction)); } diff --git a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx index 1c5414b1a6c9f..07351f8f1065f 100644 --- a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx @@ -50,6 +50,7 @@ class RActionBase; class RVariationBase; class RDefinesWithReaders; class RVariationsWithReaders; +class RColumnRegister; namespace GraphDrawing { class GraphCreatorHelper; @@ -201,6 +202,27 @@ class RLoopManager : public RNodeBase { std::set>> fUniqueVariationsWithReaders; + // deferred function calls to Jitted functions + struct DeferredJitCall { + std::string functionId; + std::shared_ptr *prevNodeOnHeap; + ROOT::Internal::RDF::RColumnRegister *colRegister; + std::vector colNames; + void *wkJittedNode, *argument; + DeferredJitCall(const std::string &id, std::shared_ptr *prevNode, + ROOT::Internal::RDF::RColumnRegister *cols, const std::vector &colnames, + void *wkNodePtr, void *arg) + : functionId(id), + prevNodeOnHeap(prevNode), + colRegister(cols), + colNames(colnames), + wkJittedNode(wkNodePtr), + argument(arg) + { + } + }; + std::vector fJitHelperCalls; + public: RLoopManager(const ColumnNames_t &defaultColumns = {}); RLoopManager(TTree *tree, const ColumnNames_t &defaultBranches); @@ -217,6 +239,7 @@ public: ~RLoopManager() override; void Jit(); + void RunDeferredCalls(); RLoopManager *GetLoopManagerUnchecked() final { return this; } void Run(bool jit = true); const ColumnNames_t &GetDefaultColumnNames() const; @@ -240,6 +263,9 @@ public: void IncrChildrenCount() final { ++fNChildren; } void StopProcessing() final { ++fNStopsReceived; } void ToJitExec(const std::string &) const; + void RegisterJitHelperCall(const std::string &funcBody, std::shared_ptr *prevNodeOnHeap, + ROOT::Internal::RDF::RColumnRegister *colRegister, + const std::vector &colNames, void *wkJittedPtr, void *argument = nullptr); void RegisterCallback(ULong64_t everyNEvents, std::function &&f); unsigned int GetNRuns() const { return fNRuns; } bool HasDataSourceColumnReaders(std::string_view col, const std::type_info &ti) const; diff --git a/tree/dataframe/src/RDFInterfaceUtils.cxx b/tree/dataframe/src/RDFInterfaceUtils.cxx index 13762da519c46..701d0ed2cd67a 100644 --- a/tree/dataframe/src/RDFInterfaceUtils.cxx +++ b/tree/dataframe/src/RDFInterfaceUtils.cxx @@ -17,6 +17,7 @@ #include #include #include +#include "ROOT/RLogger.hxx" #include #include #include @@ -631,35 +632,29 @@ BookFilterJit(std::shared_ptr *prevNodeOnHeap, std::string // definesOnHeap is deleted by the jitted call to JitFilterHelper ROOT::Internal::RDF::RColumnRegister *definesOnHeap = new ROOT::Internal::RDF::RColumnRegister(colRegister); - const auto definesOnHeapAddr = PrettyPrintAddr(definesOnHeap); - const auto prevNodeAddr = PrettyPrintAddr(prevNodeOnHeap); const auto jittedFilter = std::make_shared( (*prevNodeOnHeap)->GetLoopManagerUnchecked(), name, Union(colRegister.GetVariationDeps(parsedExpr.fUsedCols), (*prevNodeOnHeap)->GetVariations())); // Produce code snippet that creates the filter and registers it with the corresponding RJittedFilter - // Windows requires std::hex << std::showbase << (size_t)pointer to produce notation "0x1234" - std::stringstream filterInvocation; - filterInvocation << "ROOT::Internal::RDF::JitFilterHelper(" << funcName << ", new const char*[" - << parsedExpr.fUsedCols.size() << "]{"; - for (const auto &col : parsedExpr.fUsedCols) - filterInvocation << "\"" << col << "\", "; - if (!parsedExpr.fUsedCols.empty()) - filterInvocation.seekp(-2, filterInvocation.cur); // remove the last ", // lifetime of pointees: // - jittedFilter: heap-allocated weak_ptr to the actual jittedFilter that will be deleted by JitFilterHelper // - prevNodeOnHeap: heap-allocated shared_ptr to the actual previous node that will be deleted by JitFilterHelper // - definesOnHeap: heap-allocated, will be deleted by JitFilterHelper - filterInvocation << "}, " << parsedExpr.fUsedCols.size() << ", \"" << name << "\", " - << "reinterpret_cast*>(" - << PrettyPrintAddr(MakeWeakOnHeap(jittedFilter)) << "), " - << "reinterpret_cast*>(" << prevNodeAddr << ")," - << "reinterpret_cast(" << definesOnHeapAddr << ")" - << ");\n"; - + std::stringstream filterInvocation; + filterInvocation << "(ROOT::Detail::RDF::RLoopManager *lm, " + << "std::shared_ptr *prevNodeOnHeap," + << "ROOT::Internal::RDF::RColumnRegister* colRegister, " + << "const std::vector & colNames, " + << "void *wkJittedFilter, void *) {\n"; + filterInvocation << " ROOT::Internal::RDF::JitFilterHelper(" << funcName << ", " + << " colNames, \"" << name << "\", " + << "reinterpret_cast*>(wkJittedFilter)," + << "prevNodeOnHeap, colRegister);\n}\n"; auto lm = jittedFilter->GetLoopManagerUnchecked(); - lm->ToJitExec(filterInvocation.str()); + lm->RegisterJitHelperCall(filterInvocation.str(), prevNodeOnHeap, definesOnHeap, parsedExpr.fUsedCols, + MakeWeakOnHeap(jittedFilter)); return jittedFilter; } @@ -678,30 +673,28 @@ std::shared_ptr BookDefineJit(std::string_view name, std::string_ const auto type = RetTypeOfFunc(funcName); auto definesCopy = new RColumnRegister(colRegister); - auto definesAddr = PrettyPrintAddr(definesCopy); auto jittedDefine = std::make_shared(name, type, lm, colRegister, parsedExpr.fUsedCols); - std::stringstream defineInvocation; - defineInvocation << "ROOT::Internal::RDF::JitDefineHelper(" << funcName - << ", new const char*[" << parsedExpr.fUsedCols.size() << "]{"; - for (const auto &col : parsedExpr.fUsedCols) { - defineInvocation << "\"" << col << "\", "; - } - if (!parsedExpr.fUsedCols.empty()) - defineInvocation.seekp(-2, defineInvocation.cur); // remove the last ", // lifetime of pointees: // - lm is the loop manager, and if that goes out of scope jitting does not happen at all (i.e. will always be valid) // - jittedDefine: heap-allocated weak_ptr that will be deleted by JitDefineHelper after usage // - definesAddr: heap-allocated, will be deleted by JitDefineHelper after usage - defineInvocation << "}, " << parsedExpr.fUsedCols.size() << ", \"" << name - << "\", reinterpret_cast(" << PrettyPrintAddr(&lm) - << "), reinterpret_cast*>(" - << PrettyPrintAddr(MakeWeakOnHeap(jittedDefine)) - << "), reinterpret_cast(" << definesAddr - << "), reinterpret_cast*>(" - << PrettyPrintAddr(upcastNodeOnHeap) << "));\n"; - - lm.ToJitExec(defineInvocation.str()); + std::stringstream defineInvocation; + defineInvocation << "(ROOT::Detail::RDF::RLoopManager *lm, " + << "std::shared_ptr *prevNodeOnHeap," + << "ROOT::Internal::RDF::RColumnRegister* colRegister, " + << "const std::vector & colNames, " + << "void *wkJittedDefine, void *) {\n"; + defineInvocation << "ROOT::Internal::RDF::JitDefineHelper(" << funcName + << ", colNames, \"" << name << "\", " + << "lm, " + << "reinterpret_cast*>(wkJittedDefine)," + << "colRegister, " + << "prevNodeOnHeap);\n}\n"; + + lm.RegisterJitHelperCall(defineInvocation.str(), upcastNodeOnHeap, definesCopy, parsedExpr.fUsedCols, + MakeWeakOnHeap(jittedDefine)); + return jittedDefine; } @@ -718,21 +711,21 @@ std::shared_ptr BookDefinePerSampleJit(std::string_view name, std auto definesAddr = PrettyPrintAddr(definesCopy); auto jittedDefine = std::make_shared(name, retType, lm, colRegister, ColumnNames_t{}); - std::stringstream defineInvocation; - defineInvocation << "ROOT::Internal::RDF::JitDefineHelper(" - << funcName << ", nullptr, 0, "; // lifetime of pointees: // - lm is the loop manager, and if that goes out of scope jitting does not happen at all (i.e. will always be valid) // - jittedDefine: heap-allocated weak_ptr that will be deleted by JitDefineHelper after usage // - definesAddr: heap-allocated, will be deleted by JitDefineHelper after usage - defineInvocation << "\"" << name << "\", reinterpret_cast(" << PrettyPrintAddr(&lm) - << "), reinterpret_cast*>(" - << PrettyPrintAddr(MakeWeakOnHeap(jittedDefine)) - << "), reinterpret_cast(" << definesAddr - << "), reinterpret_cast*>(" - << PrettyPrintAddr(upcastNodeOnHeap) << "));\n"; - - lm.ToJitExec(defineInvocation.str()); + std::stringstream defineInvocation; + defineInvocation << "(ROOT::Detail::RDF::RLoopManager *lm, " + << "std::shared_ptr *prevNodeOnHeap," + << "ROOT::Internal::RDF::RColumnRegister* colRegister, " + << "const std::vector & colNames, " + << "void *wkJittedDefine, void *) {\n"; + defineInvocation << "ROOT::Internal::RDF::JitDefineHelper(" + << funcName << ", colNames, \"" << name << "\", lm, " + << "reinterpret_cast*>(wkJittedDefine), " + << "colRegister, prevNodeOnHeap);\n}\n"; + lm.RegisterJitHelperCall(defineInvocation.str(), upcastNodeOnHeap, definesCopy, {}, MakeWeakOnHeap(jittedDefine)); return jittedDefine; } @@ -764,50 +757,43 @@ BookVariationJit(const std::vector &colNames, std::string_view vari const auto colRegisterAddr = PrettyPrintAddr(colRegisterCopy); auto jittedVariation = std::make_shared(colNames, variationName, variationTags, type, colRegister, lm, parsedExpr.fUsedCols); + auto variedColsOnHeap = new ColumnNames_t(colNames); // build invocation to JitVariationHelper - // arrays of strings are passed as const char** plus size. + // variation tag (array of strings) passed as const char** plus size. // lifetime of pointees: // - lm is the loop manager, and if that goes out of scope jitting does not happen at all (i.e. will always be valid) // - jittedVariation: heap-allocated weak_ptr that will be deleted by JitDefineHelper after usage // - definesAddr: heap-allocated, will be deleted by JitDefineHelper after usage + // - variedColsOnHeap: deleted by registration function std::stringstream varyInvocation; + varyInvocation << "(ROOT::Detail::RDF::RLoopManager *lm, " + << "std::shared_ptr *prevNodeOnHeap," + << "ROOT::Internal::RDF::RColumnRegister* colRegister, " + << "const std::vector & inputColNames, " + << "void *wkJittedVariation, void *variedColsOnHeap) {\n"; + varyInvocation << "auto * variedColNames = reinterpret_cast*>(variedColsOnHeap);\n"; varyInvocation << "ROOT::Internal::RDF::JitVariationHelper<" << (isSingleColumn ? "true" : "false") << ">(" - << funcName << ", new const char*[" << parsedExpr.fUsedCols.size() << "]{"; - for (const auto &col : parsedExpr.fUsedCols) { - varyInvocation << "\"" << col << "\", "; - } - if (!parsedExpr.fUsedCols.empty()) - varyInvocation.seekp(-2, varyInvocation.cur); // remove the last ", " - varyInvocation << "}, " << parsedExpr.fUsedCols.size(); - varyInvocation << ", new const char*[" << colNames.size() << "]{"; - for (const auto &col : colNames) { - varyInvocation << "\"" << col << "\", "; - } - varyInvocation.seekp(-2, varyInvocation.cur); // remove the last ", " - varyInvocation << "}, " << colNames.size() << ", new const char*[" << variationTags.size() << "]{"; + << funcName << ", inputColNames, *variedColNames, "; + varyInvocation << "new const char*[" << variationTags.size() << "]{"; for (const auto &tag : variationTags) { varyInvocation << "\"" << tag << "\", "; } varyInvocation.seekp(-2, varyInvocation.cur); // remove the last ", " - varyInvocation << "}, " << variationTags.size() << ", \"" << variationName - << "\", reinterpret_cast(" << PrettyPrintAddr(&lm) - << "), reinterpret_cast*>(" - << PrettyPrintAddr(MakeWeakOnHeap(jittedVariation)) - << "), reinterpret_cast(" << colRegisterAddr - << "), reinterpret_cast*>(" - << PrettyPrintAddr(upcastNodeOnHeap) << "));\n"; - - lm.ToJitExec(varyInvocation.str()); + varyInvocation << "}, " << variationTags.size() << ", \"" << variationName << "\", lm, " + << "reinterpret_cast*>(wkJittedVariation)," + << "colRegister, prevNodeOnHeap);\n" + << "delete variedColNames;\n}\n"; + lm.RegisterJitHelperCall(varyInvocation.str(), upcastNodeOnHeap, colRegisterCopy, parsedExpr.fUsedCols, + MakeWeakOnHeap(jittedVariation), variedColsOnHeap); return jittedVariation; } // Jit and call something equivalent to "this->BuildAndBook(params...)" // (see comments in the body for actual jitted code) -std::string JitBuildAction(const ColumnNames_t &cols, std::shared_ptr *prevNode, - const std::type_info &helperArgType, const std::type_info &at, void *helperArgOnHeap, +std::string JitBuildAction(const ColumnNames_t &cols, const std::type_info &helperArgType, const std::type_info &at, TTree *tree, const unsigned int nSlots, const RColumnRegister &colRegister, RDataSource *ds, - std::weak_ptr *jittedActionOnHeap, const bool vector2RVec) + const bool vector2RVec) { // retrieve type of action as a string auto actionTypeClass = TClass::GetClass(at); @@ -824,30 +810,22 @@ std::string JitBuildAction(const ColumnNames_t &cols, std::shared_ptr *prevNodeOnHeap," + << "ROOT::Internal::RDF::RColumnRegister* colRegister, " + << "const std::vector & colNames, " + << "void *wkJittedAction, void *actionArg) {\n"; createAction_str << "ROOT::Internal::RDF::CallBuildAction<" << actionTypeName; const auto columnTypeNames = GetValidatedArgTypes(cols, colRegister, tree, ds, actionTypeNameBase, vector2RVec); for (auto &colType : columnTypeNames) createAction_str << ", " << colType; - // on Windows, to prefix the hexadecimal value of a pointer with '0x', - // one need to write: std::hex << std::showbase << (size_t)pointer - createAction_str << ">(reinterpret_cast*>(" - << PrettyPrintAddr(prevNode) << "), new const char*[" << cols.size() << "]{"; - for (auto i = 0u; i < cols.size(); ++i) { - if (i != 0u) - createAction_str << ", "; - createAction_str << '"' << cols[i] << '"'; - } - createAction_str << "}, " << cols.size() << ", " << nSlots << ", reinterpret_cast*>(" << PrettyPrintAddr(helperArgOnHeap) - << "), reinterpret_cast*>(" - << PrettyPrintAddr(jittedActionOnHeap) - << "), reinterpret_cast(" << definesAddr << "));"; + createAction_str << ">(prevNodeOnHeap, colNames," << nSlots << ", " + << " reinterpret_cast*>(actionArg)," + << " reinterpret_cast*>(wkJittedAction)," + << "colRegister);\n}\n"; return createAction_str.str(); } diff --git a/tree/dataframe/src/RLoopManager.cxx b/tree/dataframe/src/RLoopManager.cxx index f8247eb9ac8a2..6ea764f87cbf5 100644 --- a/tree/dataframe/src/RLoopManager.cxx +++ b/tree/dataframe/src/RLoopManager.cxx @@ -30,6 +30,7 @@ #include "TEntryList.h" #include "TFile.h" #include "TFriendElement.h" +#include "TInterpreter.h" #include "TROOT.h" // IsImplicitMTEnabled, gCoreMutex, R__*_LOCKGUARD #include "TTreeReader.h" #include "TTree.h" // For MaxTreeSizeRAII. Revert when #6640 will be solved. @@ -90,6 +91,24 @@ std::string &GetCodeToJit() return code; } +std::string &GetCodeToDeclare() +{ + static std::string code; + return code; +} +using JitHelperFunc = void(RLoopManager *, std::shared_ptr *, RColumnRegister *, const ColumnNames_t &, + void *, void *); +std::unordered_map &GetJitHelperFuncMap() +{ + static std::unordered_map map; + return map; +} +std::unordered_map &GetJitHelperNameMap() +{ + static std::unordered_map map; + return map; +} + void ThrowIfNSlotsChanged(unsigned int nSlots) { const auto currentSlots = RDFInternal::GetNSlots(); @@ -766,24 +785,73 @@ void RLoopManager::Jit() { { R__READ_LOCKGUARD(ROOT::gCoreMutex); - if (GetCodeToJit().empty()) { + if (GetCodeToJit().empty() && GetCodeToDeclare().empty()) { + RunDeferredCalls(); R__LOG_INFO(RDFLogChannel()) << "Nothing to jit and execute."; return; } } - const std::string code = []() { + std::string codeToDeclare, code; + { R__WRITE_LOCKGUARD(ROOT::gCoreMutex); - return std::move(GetCodeToJit()); - }(); + codeToDeclare.swap(GetCodeToDeclare()); + code.swap(GetCodeToJit()); + }; TStopwatch s; s.Start(); - RDFInternal::InterpreterCalc(code, "RLoopManager::Run"); + if (!codeToDeclare.empty()) { + ROOT::Internal::RDF::InterpreterDeclare(codeToDeclare); + auto &funcMap = GetJitHelperFuncMap(); + auto &nameMap = GetJitHelperNameMap(); + auto clinfo = gInterpreter->ClassInfo_Factory("R_rdf"); + assert(gInterpreter->ClassInfo_IsValid(clinfo)); + for (auto &codeAndName : nameMap) { + JitHelperFunc *&addr = funcMap[codeAndName.second]; + if (!addr) { + // fast fetch of the address via gInterpreter + // (faster than gInterpreter->Evaluate(function name, ret), ret->GetAsPointer()) + auto declid = gInterpreter->GetFunction(clinfo, codeAndName.second.c_str()); + assert(declid); + auto minfo = gInterpreter->MethodInfo_Factory(declid); + assert(gInterpreter->MethodInfo_IsValid(minfo)); + auto mname = gInterpreter->MethodInfo_GetMangledName(minfo); + addr = reinterpret_cast(gInterpreter->FindSym(mname)); + gInterpreter->MethodInfo_Delete(minfo); + } + } + gInterpreter->ClassInfo_Delete(clinfo); + } + if (!code.empty()) { + RDFInternal::InterpreterCalc(code, "RLoopManager::Run"); + } s.Stop(); R__LOG_INFO(RDFLogChannel()) << "Just-in-time compilation phase completed" << (s.RealTime() > 1e-3 ? " in " + std::to_string(s.RealTime()) + " seconds." : " in less than 1ms."); + + RunDeferredCalls(); +} + +void RLoopManager::RunDeferredCalls() +{ + if (!fJitHelperCalls.empty()) { + R__READ_LOCKGUARD(ROOT::gCoreMutex); // methods are thread-safe but funcMap isn't (yet) + TStopwatch s2; + s2.Start(); + auto &funcMap = GetJitHelperFuncMap(); + for (auto &call : fJitHelperCalls) { + assert(funcMap.find(call.functionId) != funcMap.end()); + funcMap[call.functionId](this, call.prevNodeOnHeap, call.colRegister, call.colNames, call.wkJittedNode, + call.argument); + } + s2.Stop(); + R__LOG_INFO(RDFLogChannel()) << "Deferred calls (" << fJitHelperCalls.size() << ") completed" + << (s2.RealTime() > 1e-3 ? " in " + std::to_string(s2.RealTime()) + " seconds." + : " in less than 1ms."); + fJitHelperCalls.clear(); + } } /// Trigger counting of number of children nodes for each node of the functional graph. @@ -808,13 +876,13 @@ void RLoopManager::Run(bool jit) // Change value of TTree::GetMaxTreeSize only for this scope. Revert when #6640 will be solved. MaxTreeSizeRAII ctxtmts; - R__LOG_INFO(RDFLogChannel()) << "Starting event loop number " << fNRuns << '.'; - ThrowIfNSlotsChanged(GetNSlots()); if (jit) Jit(); + RunDeferredCalls(); + InitNodes(); // Exceptions can occur during the event loop. In order to ensure proper cleanup of nodes @@ -934,6 +1002,33 @@ void RLoopManager::ToJitExec(const std::string &code) const GetCodeToJit().append(code); } +void RLoopManager::RegisterJitHelperCall(const std::string &funcCode, std::shared_ptr *prevNodeOnHeap, + ROOT::Internal::RDF::RColumnRegister *colRegister, + const std::vector &colNames, void *wkJittedNode, void *argument) +{ + auto &nameMap = GetJitHelperNameMap(); + { + R__READ_LOCKGUARD(ROOT::gCoreMutex); + auto match = nameMap.find(funcCode); + if (match != nameMap.end()) { + R__LOG_DEBUG(0, RDFLogChannel()) << "JitHelper " << match->second << " already defined"; + fJitHelperCalls.emplace_back(match->second, prevNodeOnHeap, colRegister, colNames, wkJittedNode, argument); + return; + } + } + + { + R__WRITE_LOCKGUARD(ROOT::gCoreMutex); + std::string registerId = "jitNodeRegistrator_" + std::to_string(nameMap.size()); + nameMap[funcCode] = registerId; + R__LOG_DEBUG(0, RDFLogChannel()) << "JitHelper new " << registerId << " defined for funcCode " << funcCode; + // step 1: register function (now) + std::string toDeclare = "namespace R_rdf {\n void " + registerId + funcCode + "\n}\n"; + GetCodeToDeclare().append(toDeclare); + fJitHelperCalls.emplace_back(registerId, prevNodeOnHeap, colRegister, colNames, wkJittedNode, argument); + } +} + void RLoopManager::RegisterCallback(ULong64_t everyNEvents, std::function &&f) { if (everyNEvents == 0ull) From 022b7432aa3b153d1689b845baf73fcc1e70f55b Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Sun, 30 Nov 2025 11:49:27 +0100 Subject: [PATCH 2/3] [df] Rework memory ownership model of JITted nodes 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 https://github.com/root-project/root/issues/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. --- .../dataframe/inc/ROOT/RDF/InterfaceUtils.hxx | 147 +++++---------- tree/dataframe/inc/ROOT/RDF/RInterface.hxx | 23 +-- .../dataframe/inc/ROOT/RDF/RInterfaceBase.hxx | 15 +- tree/dataframe/inc/ROOT/RDF/RJittedAction.hxx | 6 +- tree/dataframe/inc/ROOT/RDF/RJittedFilter.hxx | 5 +- tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx | 42 +++-- tree/dataframe/inc/ROOT/RDataFrame.hxx | 8 - tree/dataframe/src/RDFInterfaceUtils.cxx | 174 +++++++++--------- tree/dataframe/src/RDataFrame.cxx | 11 -- tree/dataframe/src/RJittedAction.cxx | 10 +- tree/dataframe/src/RJittedFilter.cxx | 11 +- tree/dataframe/src/RLoopManager.cxx | 162 +++++++++++----- 12 files changed, 300 insertions(+), 314 deletions(-) diff --git a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx b/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx index 4fdc846a701fb..83c992aad5bb8 100644 --- a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx +++ b/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx @@ -367,23 +367,20 @@ void CheckForNoVariations(const std::string &where, std::string_view definedColV std::string PrettyPrintAddr(const void *const addr); -std::shared_ptr BookFilterJit(std::shared_ptr *prevNodeOnHeap, std::string_view name, +std::shared_ptr BookFilterJit(std::shared_ptr prevNode, std::string_view name, std::string_view expression, const RColumnRegister &colRegister, TTree *tree, RDataSource *ds); std::shared_ptr BookDefineJit(std::string_view name, std::string_view expression, RLoopManager &lm, - RDataSource *ds, const RColumnRegister &colRegister, - std::shared_ptr *prevNodeOnHeap); + RDataSource *ds, const RColumnRegister &colRegister); std::shared_ptr BookDefinePerSampleJit(std::string_view name, std::string_view expression, - RLoopManager &lm, const RColumnRegister &colRegister, - std::shared_ptr *upcastNodeOnHeap); + RLoopManager &lm, const RColumnRegister &colRegister); std::shared_ptr BookVariationJit(const std::vector &colNames, std::string_view variationName, const std::vector &variationTags, std::string_view expression, RLoopManager &lm, - RDataSource *ds, const RColumnRegister &colRegister, std::shared_ptr *upcastNodeOnHeap, - bool isSingleColumn); + RDataSource *ds, const RColumnRegister &colRegister, bool isSingleColumn); std::string JitBuildAction(const ColumnNames_t &bl, const std::type_info &art, const std::type_info &at, TTree *tree, const unsigned int nSlots, const RColumnRegister &colRegister, RDataSource *ds, @@ -471,42 +468,32 @@ void AddDSColumns(const std::vector &requiredCols, ROOT::Detail::RD ROOT::Internal::RDF::RColumnRegister &colRegister); // this function is meant to be called by the jitted code generated by BookFilterJit -template -void JitFilterHelper(F &&f, const ColumnNames_t &cols, std::string_view name, - std::weak_ptr *wkJittedFilter, std::shared_ptr *prevNodeOnHeap, - RColumnRegister *colRegister) noexcept +template +void JitFilterHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RColumnRegister &colRegister, + ROOT::Detail::RDF::RLoopManager &lm, ROOT::Detail::RDF::RJittedFilter *jittedFilter) noexcept { - if (wkJittedFilter->expired()) { + if (!jittedFilter) { // The branch of the computation graph that needed this jitted code went out of scope between the type // jitting was booked and the time jitting actually happened. Nothing to do other than cleaning up. - delete wkJittedFilter; - delete colRegister; - delete prevNodeOnHeap; return; } - const auto jittedFilter = wkJittedFilter->lock(); - // mock Filter logic -- validity checks and Define-ition of RDataSource columns using Callable_t = std::decay_t; - using F_t = RFilter; + auto prevNode = jittedFilter->MoveOutPrevNode(); + using PrevNode_t = typename decltype(prevNode)::element_type; + using F_t = RFilter; using ColTypes_t = typename TTraits::CallableTraits::arg_types; constexpr auto nColumns = ColTypes_t::list_size; CheckFilter(f); - auto &lm = *jittedFilter->GetLoopManagerUnchecked(); // RLoopManager must exist at this time auto ds = lm.GetDataSource(); - if (ds != nullptr) - AddDSColumns(cols, lm, *ds, ColTypes_t(), *colRegister); + if (ds != nullptr && !cols.empty()) + AddDSColumns(cols, lm, *ds, ColTypes_t(), colRegister); jittedFilter->SetFilter( - std::unique_ptr(new F_t(std::forward(f), cols, *prevNodeOnHeap, *colRegister, name))); - // colRegister points to the columns structure in the heap, created before the jitted call so that the jitter can - // share data after it has lazily compiled the code. Here the data has been used and the memory can be freed. - delete colRegister; - delete prevNodeOnHeap; - delete wkJittedFilter; + std::unique_ptr(new F_t(std::forward(f), cols, prevNode, colRegister, name))); } namespace DefineTypes { @@ -534,124 +521,80 @@ auto MakeDefineNode(DefineTypes::RDefinePerSampleTag, std::string_view name, std // This function is meant to be called by jitted code right before starting the event loop. // If colsPtr is null, build a RDefinePerSample (it has no input columns), otherwise a RDefine. template -void JitDefineHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RLoopManager *lm, - std::weak_ptr *wkJittedDefine, RColumnRegister *colRegister, - std::shared_ptr *prevNodeOnHeap) noexcept +void JitDefineHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RColumnRegister &colRegister, + ROOT::Detail::RDF::RLoopManager &lm, ROOT::Detail::RDF::RJittedDefine *jittedDefine) noexcept { - // a helper to delete objects allocated before jitting, so that the jitter can share data with lazily jitted code - auto doDeletes = [&] { - delete wkJittedDefine; - delete colRegister; - delete prevNodeOnHeap; - }; - - if (wkJittedDefine->expired()) { + + if (!jittedDefine) { // The branch of the computation graph that needed this jitted code went out of scope between the type // jitting was booked and the time jitting actually happened. Nothing to do other than cleaning up. - doDeletes(); return; } - auto jittedDefine = wkJittedDefine->lock(); - using Callable_t = std::decay_t; using ColTypes_t = typename TTraits::CallableTraits::arg_types; - auto ds = lm->GetDataSource(); - if (ds != nullptr) - AddDSColumns(cols, *lm, *ds, ColTypes_t(), *colRegister); + auto ds = lm.GetDataSource(); + if (ds != nullptr && !cols.empty()) + AddDSColumns(cols, lm, *ds, ColTypes_t(), colRegister); // will never actually be used (trumped by jittedDefine->GetTypeName()), but we set it to something meaningful // to help devs debugging const auto dummyType = "jittedCol_t"; // use unique_ptr instead of make_unique to reduce jit/compile-times std::unique_ptr newCol{ - MakeDefineNode(RDefineTypeTag{}, name, dummyType, std::forward(f), cols, *colRegister, *lm)}; + MakeDefineNode(RDefineTypeTag{}, name, dummyType, std::forward(f), cols, colRegister, lm)}; jittedDefine->SetDefine(std::move(newCol)); - - doDeletes(); } template -void JitVariationHelper(F &&f, const ColumnNames_t &inputColNames, const ColumnNames_t &variedColNames, - const char **variationTags, std::size_t variationTagsSize, std::string_view variationName, - RLoopManager *lm, std::weak_ptr *wkJittedVariation, - RColumnRegister *colRegister, std::shared_ptr *prevNodeOnHeap) noexcept +void JitVariationHelper(F &&f, const ColumnNames_t &inputColNames, std::string_view variationName, + RColumnRegister &colRegister, ROOT::Detail::RDF::RLoopManager &lm, + RJittedVariation *jittedVariation, const ColumnNames_t &variedColNames, + const ColumnNames_t &variationTags) noexcept { - // a helper to delete objects allocated before jitting, so that the jitter can share data with lazily jitted code - auto doDeletes = [&] { - delete[] variationTags; - delete wkJittedVariation; - delete colRegister; - delete prevNodeOnHeap; - }; - - if (wkJittedVariation->expired()) { + + if (!jittedVariation) { // The branch of the computation graph that needed this jitted variation went out of scope between the type // jitting was booked and the time jitting actually happened. Nothing to do other than cleaning up. - doDeletes(); return; } - std::vector tags(variationTags, variationTags + variationTagsSize); - - auto jittedVariation = wkJittedVariation->lock(); - using Callable_t = std::decay_t; using ColTypes_t = typename TTraits::CallableTraits::arg_types; - auto ds = lm->GetDataSource(); - if (ds != nullptr) - AddDSColumns(inputColNames, *lm, *ds, ColTypes_t(), *colRegister); + auto ds = lm.GetDataSource(); + if (ds != nullptr && !inputColNames.empty()) + AddDSColumns(inputColNames, lm, *ds, ColTypes_t(), colRegister); // use unique_ptr instead of make_unique to reduce jit/compile-times - std::unique_ptr newVariation{new RVariation, IsSingleColumn>( - std::move(variedColNames), variationName, std::forward(f), std::move(tags), jittedVariation->GetTypeName(), - *colRegister, *lm, inputColNames)}; + std::unique_ptr newVariation{ + new RVariation, IsSingleColumn>(variedColNames, variationName, std::forward(f), variationTags, + jittedVariation->GetTypeName(), colRegister, lm, inputColNames)}; jittedVariation->SetVariation(std::move(newVariation)); - - doDeletes(); } /// Convenience function invoked by jitted code to build action nodes at runtime -template -void CallBuildAction(std::shared_ptr *prevNodeOnHeap, const ColumnNames_t &cols, - const unsigned int nSlots, std::shared_ptr *helperArgOnHeap, - std::weak_ptr *wkJittedActionOnHeap, RColumnRegister *colRegister) noexcept +template +void CallBuildAction(const ColumnNames_t &cols, RColumnRegister &colRegister, ROOT::Detail::RDF::RLoopManager &lm, + RJittedAction *jittedAction, unsigned int nSlots, + std::shared_ptr *helperArg) noexcept { - // a helper to delete objects allocated before jitting, so that the jitter can share data with lazily jitted code - auto doDeletes = [&] { - delete helperArgOnHeap; - delete wkJittedActionOnHeap; - // colRegister must be deleted before prevNodeOnHeap because their dtor needs the RLoopManager to be alive - // and prevNodeOnHeap is what keeps it alive if the rest of the computation graph is already out of scope - delete colRegister; - delete prevNodeOnHeap; - }; - - if (wkJittedActionOnHeap->expired()) { + if (!jittedAction) { // The branch of the computation graph that needed this jitted variation went out of scope between the type // jitting was booked and the time jitting actually happened. Nothing to do other than cleaning up. - doDeletes(); return; } - auto jittedActionOnHeap = wkJittedActionOnHeap->lock(); - - // if we are here it means we are jitting, if we are jitting the loop manager must be alive - auto &prevNodePtr = *prevNodeOnHeap; - auto &loopManager = *prevNodePtr->GetLoopManagerUnchecked(); using ColTypes_t = TypeList; constexpr auto nColumns = ColTypes_t::list_size; - auto ds = loopManager.GetDataSource(); - if (ds != nullptr) - AddDSColumns(cols, loopManager, *ds, ColTypes_t(), *colRegister); - - auto actionPtr = BuildAction(cols, std::move(*helperArgOnHeap), nSlots, std::move(prevNodePtr), - ActionTag{}, *colRegister); - jittedActionOnHeap->SetAction(std::move(actionPtr)); + auto ds = lm.GetDataSource(); + if (ds != nullptr && !cols.empty()) + AddDSColumns(cols, lm, *ds, ColTypes_t(), colRegister); - doDeletes(); + auto actionPtr = + BuildAction(cols, *helperArg, nSlots, jittedAction->MoveOutPrevNode(), ActionTag{}, colRegister); + jittedAction->SetAction(std::move(actionPtr)); } /// The contained `type` alias is `double` if `T == RInferredType`, `U` if `T == std::container`, `T` otherwise. diff --git a/tree/dataframe/inc/ROOT/RDF/RInterface.hxx b/tree/dataframe/inc/ROOT/RDF/RInterface.hxx index e13447db384f7..417351f6b325d 100644 --- a/tree/dataframe/inc/ROOT/RDF/RInterface.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RInterface.hxx @@ -290,12 +290,8 @@ public: /// ~~~ RInterface Filter(std::string_view expression, std::string_view name = "") { - // deleted by the jitted call to JitFilterHelper - auto upcastNodeOnHeap = RDFInternal::MakeSharedOnHeap(RDFInternal::UpcastNode(fProxiedPtr)); - using BaseNodeType_t = typename std::remove_pointer_t::element_type; - RInterface upcastInterface(*upcastNodeOnHeap, *fLoopManager, fColRegister); - const auto jittedFilter = - RDFInternal::BookFilterJit(upcastNodeOnHeap, name, expression, fColRegister, nullptr, GetDataSource()); + const auto jittedFilter = RDFInternal::BookFilterJit(RDFInternal::UpcastNode(fProxiedPtr), name, expression, + fColRegister, nullptr, GetDataSource()); return RInterface(std::move(jittedFilter), *fLoopManager, fColRegister); } @@ -538,9 +534,7 @@ public: RDFInternal::CheckForRedefinition(where, name, fColRegister, GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); - auto upcastNodeOnHeap = RDFInternal::MakeSharedOnHeap(RDFInternal::UpcastNode(fProxiedPtr)); - auto jittedDefine = - RDFInternal::BookDefineJit(name, expression, *fLoopManager, GetDataSource(), fColRegister, upcastNodeOnHeap); + auto jittedDefine = RDFInternal::BookDefineJit(name, expression, *fLoopManager, GetDataSource(), fColRegister); RDFInternal::RColumnRegister newCols(fColRegister); newCols.AddDefine(std::move(jittedDefine)); @@ -628,9 +622,7 @@ public: GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); RDFInternal::CheckForNoVariations(where, name, fColRegister); - auto upcastNodeOnHeap = RDFInternal::MakeSharedOnHeap(RDFInternal::UpcastNode(fProxiedPtr)); - auto jittedDefine = - RDFInternal::BookDefineJit(name, expression, *fLoopManager, GetDataSource(), fColRegister, upcastNodeOnHeap); + auto jittedDefine = RDFInternal::BookDefineJit(name, expression, *fLoopManager, GetDataSource(), fColRegister); RDFInternal::RColumnRegister newCols(fColRegister); newCols.AddDefine(std::move(jittedDefine)); @@ -805,9 +797,7 @@ public: RDFInternal::CheckForRedefinition("DefinePerSample", name, fColRegister, GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); - auto upcastNodeOnHeap = RDFInternal::MakeSharedOnHeap(RDFInternal::UpcastNode(fProxiedPtr)); - auto jittedDefine = - RDFInternal::BookDefinePerSampleJit(name, expression, *fLoopManager, fColRegister, upcastNodeOnHeap); + auto jittedDefine = RDFInternal::BookDefinePerSampleJit(name, expression, *fLoopManager, fColRegister); RDFInternal::RColumnRegister newCols(fColRegister); newCols.AddDefine(std::move(jittedDefine)); @@ -3415,10 +3405,9 @@ private: throw std::logic_error("A column name was passed to the same Vary invocation multiple times."); } - auto upcastNodeOnHeap = RDFInternal::MakeSharedOnHeap(RDFInternal::UpcastNode(fProxiedPtr)); auto jittedVariation = RDFInternal::BookVariationJit(colNames, variationName, variationTags, expression, *fLoopManager, - GetDataSource(), fColRegister, upcastNodeOnHeap, isSingleColumn); + GetDataSource(), fColRegister, isSingleColumn); RDFInternal::RColumnRegister newColRegister(fColRegister); newColRegister.AddVariation(std::move(jittedVariation)); diff --git a/tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx b/tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx index 03546dbaf25c2..5b7fd7334999d 100644 --- a/tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx @@ -193,19 +193,14 @@ protected: const auto validColumnNames = GetValidatedColumnNames(realNColumns, columns); const unsigned int nSlots = fLoopManager->GetNSlots(); - auto *helperArgOnHeap = RDFInternal::MakeSharedOnHeap(helperArg); + const auto jittedAction = std::make_shared( + *fLoopManager, validColumnNames, fColRegister, proxiedPtr->GetVariations(), proxiedPtr); - auto upcastNodeOnHeap = RDFInternal::MakeSharedOnHeap(RDFInternal::UpcastNode(proxiedPtr)); - - const auto jittedAction = std::make_shared(*fLoopManager, validColumnNames, - fColRegister, proxiedPtr->GetVariations()); - auto jittedActionOnHeap = RDFInternal::MakeWeakOnHeap(jittedAction); - - auto definesCopy = new RDFInternal::RColumnRegister(fColRegister); // deleted in jitted call auto funcBody = RDFInternal::JitBuildAction(validColumnNames, typeid(HelperArgType), typeid(ActionTag), nullptr, nSlots, fColRegister, GetDataSource(), vector2RVec); - fLoopManager->RegisterJitHelperCall(funcBody, upcastNodeOnHeap, definesCopy, validColumnNames, jittedActionOnHeap, - helperArgOnHeap); + fLoopManager->RegisterJitHelperCall(funcBody, + std::make_unique(fColRegister), + validColumnNames, jittedAction, helperArg); return MakeResultPtr(r, *fLoopManager, std::move(jittedAction)); } diff --git a/tree/dataframe/inc/ROOT/RDF/RJittedAction.hxx b/tree/dataframe/inc/ROOT/RDF/RJittedAction.hxx index 621079bde40db..9d31bf25d2591 100644 --- a/tree/dataframe/inc/ROOT/RDF/RJittedAction.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RJittedAction.hxx @@ -23,6 +23,7 @@ namespace ROOT { namespace Detail { namespace RDF { class RMergeableValueBase; +class RNodeBase; } // namespace RDF } // namespace Detail } // namespace ROOT @@ -39,10 +40,12 @@ class GraphNode; class RJittedAction : public RActionBase { private: std::unique_ptr fConcreteAction; + std::shared_ptr fPrevNode; public: RJittedAction(RLoopManager &lm, const ROOT::RDF::ColumnNames_t &columns, const RColumnRegister &colRegister, - const std::vector &prevVariations); + const std::vector &prevVariations, + std::shared_ptr prevNode = nullptr); ~RJittedAction(); void SetAction(std::unique_ptr a) { fConcreteAction = std::move(a); } @@ -67,6 +70,7 @@ public: std::unique_ptr MakeVariedAction(std::vector &&results) final; std::unique_ptr CloneAction(void *newResult) final; + std::shared_ptr MoveOutPrevNode(); }; } // ns RDF diff --git a/tree/dataframe/inc/ROOT/RDF/RJittedFilter.hxx b/tree/dataframe/inc/ROOT/RDF/RJittedFilter.hxx index d3b3e1802c9dc..b01325ed2e68e 100644 --- a/tree/dataframe/inc/ROOT/RDF/RJittedFilter.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RJittedFilter.hxx @@ -38,9 +38,11 @@ namespace RDFGraphDrawing = ROOT::Internal::RDF::GraphDrawing; /// at a later time, from jitted code. class RJittedFilter final : public RFilterBase { std::unique_ptr fConcreteFilter = nullptr; + std::shared_ptr fPrevNode; public: - RJittedFilter(RLoopManager *lm, std::string_view name, const std::vector &variations); + RJittedFilter(RLoopManager *lm, std::string_view name, const std::vector &variations, + std::shared_ptr prevNode = nullptr); // Rule of five @@ -68,6 +70,7 @@ public: std::shared_ptr GetGraph(std::unordered_map> &visitedMap) final; std::shared_ptr GetVariedFilter(const std::string &variationName) final; + std::shared_ptr MoveOutPrevNode(); }; } // ns RDF diff --git a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx index 07351f8f1065f..6b6af72f84331 100644 --- a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx @@ -204,24 +204,25 @@ class RLoopManager : public RNodeBase { // deferred function calls to Jitted functions struct DeferredJitCall { - std::string functionId; - std::shared_ptr *prevNodeOnHeap; - ROOT::Internal::RDF::RColumnRegister *colRegister; - std::vector colNames; - void *wkJittedNode, *argument; - DeferredJitCall(const std::string &id, std::shared_ptr *prevNode, - ROOT::Internal::RDF::RColumnRegister *cols, const std::vector &colnames, - void *wkNodePtr, void *arg) - : functionId(id), - prevNodeOnHeap(prevNode), - colRegister(cols), - colNames(colnames), - wkJittedNode(wkNodePtr), - argument(arg) - { - } + std::size_t fFunctionId{}; + std::unique_ptr fColRegister; + std::vector fColNames; + std::shared_ptr fJittedNode; + // Extra arguments to be passed to the jitted function, in one value. Each function will need to unpack it + // accordingly. + std::shared_ptr fExtraArgs; + DeferredJitCall(std::size_t id, std::unique_ptr cols, + const std::vector &colNamesArg, std::shared_ptr jittedNode, + std::shared_ptr arg); + + DeferredJitCall(const DeferredJitCall &) = delete; + DeferredJitCall &operator=(const DeferredJitCall &) = delete; + DeferredJitCall(DeferredJitCall &&) noexcept; + DeferredJitCall &operator=(DeferredJitCall &&) noexcept; + ~DeferredJitCall(); }; - std::vector fJitHelperCalls; + std::vector fJitHelperCalls{}; + std::hash fStringHasher{}; public: RLoopManager(const ColumnNames_t &defaultColumns = {}); @@ -263,9 +264,10 @@ public: void IncrChildrenCount() final { ++fNChildren; } void StopProcessing() final { ++fNStopsReceived; } void ToJitExec(const std::string &) const; - void RegisterJitHelperCall(const std::string &funcBody, std::shared_ptr *prevNodeOnHeap, - ROOT::Internal::RDF::RColumnRegister *colRegister, - const std::vector &colNames, void *wkJittedPtr, void *argument = nullptr); + void RegisterJitHelperCall(const std::string &funcBody, + std::unique_ptr colRegister, + const std::vector &colnames, std::shared_ptr jittedNode, + std::shared_ptr argument = nullptr); void RegisterCallback(ULong64_t everyNEvents, std::function &&f); unsigned int GetNRuns() const { return fNRuns; } bool HasDataSourceColumnReaders(std::string_view col, const std::type_info &ti) const; diff --git a/tree/dataframe/inc/ROOT/RDataFrame.hxx b/tree/dataframe/inc/ROOT/RDataFrame.hxx index 3ea7c16e224ba..b69e15f91a9e0 100644 --- a/tree/dataframe/inc/ROOT/RDataFrame.hxx +++ b/tree/dataframe/inc/ROOT/RDataFrame.hxx @@ -55,14 +55,6 @@ public: RDataFrame(ULong64_t numEntries); RDataFrame(std::unique_ptr, const ColumnNames_t &defaultColumns = {}); RDataFrame(ROOT::RDF::Experimental::RDatasetSpec spec); - - // Rule of five - - RDataFrame(const RDataFrame &) = default; - RDataFrame &operator=(const RDataFrame &) = default; - RDataFrame(RDataFrame &&) = default; - RDataFrame &operator=(RDataFrame &&) = default; - ~RDataFrame(); }; namespace RDF { diff --git a/tree/dataframe/src/RDFInterfaceUtils.cxx b/tree/dataframe/src/RDFInterfaceUtils.cxx index 701d0ed2cd67a..9db800d6bed8b 100644 --- a/tree/dataframe/src/RDFInterfaceUtils.cxx +++ b/tree/dataframe/src/RDFInterfaceUtils.cxx @@ -617,7 +617,7 @@ std::string PrettyPrintAddr(const void *const addr) /// Book the jitting of a Filter call std::shared_ptr -BookFilterJit(std::shared_ptr *prevNodeOnHeap, std::string_view name, std::string_view expression, +BookFilterJit(std::shared_ptr prevNode, std::string_view name, std::string_view expression, const RColumnRegister &colRegister, TTree *tree, RDataSource *ds) { const auto &dsColumns = ds ? ds->GetColumnNames() : ColumnNames_t{}; @@ -630,39 +630,34 @@ BookFilterJit(std::shared_ptr *prevNodeOnHeap, std::string if (type != "bool") std::runtime_error("Filter: the following expression does not evaluate to bool:\n" + std::string(expression)); - // definesOnHeap is deleted by the jitted call to JitFilterHelper - ROOT::Internal::RDF::RColumnRegister *definesOnHeap = new ROOT::Internal::RDF::RColumnRegister(colRegister); - + auto *lm = prevNode->GetLoopManagerUnchecked(); const auto jittedFilter = std::make_shared( - (*prevNodeOnHeap)->GetLoopManagerUnchecked(), name, - Union(colRegister.GetVariationDeps(parsedExpr.fUsedCols), (*prevNodeOnHeap)->GetVariations())); + lm, name, Union(colRegister.GetVariationDeps(parsedExpr.fUsedCols), prevNode->GetVariations()), prevNode); // Produce code snippet that creates the filter and registers it with the corresponding RJittedFilter - // lifetime of pointees: - // - jittedFilter: heap-allocated weak_ptr to the actual jittedFilter that will be deleted by JitFilterHelper - // - prevNodeOnHeap: heap-allocated shared_ptr to the actual previous node that will be deleted by JitFilterHelper - // - definesOnHeap: heap-allocated, will be deleted by JitFilterHelper std::stringstream filterInvocation; - filterInvocation << "(ROOT::Detail::RDF::RLoopManager *lm, " - << "std::shared_ptr *prevNodeOnHeap," - << "ROOT::Internal::RDF::RColumnRegister* colRegister, " - << "const std::vector & colNames, " - << "void *wkJittedFilter, void *) {\n"; - filterInvocation << " ROOT::Internal::RDF::JitFilterHelper(" << funcName << ", " - << " colNames, \"" << name << "\", " - << "reinterpret_cast*>(wkJittedFilter)," - << "prevNodeOnHeap, colRegister);\n}\n"; - auto lm = jittedFilter->GetLoopManagerUnchecked(); - lm->RegisterJitHelperCall(filterInvocation.str(), prevNodeOnHeap, definesOnHeap, parsedExpr.fUsedCols, - MakeWeakOnHeap(jittedFilter)); + filterInvocation << "(const std::vector &colNames, " + << "ROOT::Internal::RDF::RColumnRegister &colRegister, " + << "ROOT::Detail::RDF::RLoopManager &lm, " + << "void *jittedFilter, " + << "std::shared_ptr *) {\n"; + filterInvocation << " ROOT::Internal::RDF::JitFilterHelper(" << funcName << ", " + << "colNames, " + << "\"" << name << "\", " + << "colRegister, " + << "lm, " + << "reinterpret_cast(jittedFilter)" + << ");\n}\n"; + lm->RegisterJitHelperCall(filterInvocation.str(), + std::make_unique(colRegister), parsedExpr.fUsedCols, + jittedFilter); return jittedFilter; } /// Book the jitting of a Define call std::shared_ptr BookDefineJit(std::string_view name, std::string_view expression, RLoopManager &lm, - RDataSource *ds, const RColumnRegister &colRegister, - std::shared_ptr *upcastNodeOnHeap) + RDataSource *ds, const RColumnRegister &colRegister) { const auto &dsColumns = ds ? ds->GetColumnNames() : ColumnNames_t{}; @@ -672,7 +667,6 @@ std::shared_ptr BookDefineJit(std::string_view name, std::string_ const auto funcName = DeclareFunction(parsedExpr.fExpr, parsedExpr.fVarNames, exprVarTypes); const auto type = RetTypeOfFunc(funcName); - auto definesCopy = new RColumnRegister(colRegister); auto jittedDefine = std::make_shared(name, type, lm, colRegister, parsedExpr.fUsedCols); // lifetime of pointees: @@ -680,35 +674,33 @@ std::shared_ptr BookDefineJit(std::string_view name, std::string_ // - jittedDefine: heap-allocated weak_ptr that will be deleted by JitDefineHelper after usage // - definesAddr: heap-allocated, will be deleted by JitDefineHelper after usage std::stringstream defineInvocation; - defineInvocation << "(ROOT::Detail::RDF::RLoopManager *lm, " - << "std::shared_ptr *prevNodeOnHeap," - << "ROOT::Internal::RDF::RColumnRegister* colRegister, " - << "const std::vector & colNames, " - << "void *wkJittedDefine, void *) {\n"; - defineInvocation << "ROOT::Internal::RDF::JitDefineHelper(" << funcName - << ", colNames, \"" << name << "\", " - << "lm, " - << "reinterpret_cast*>(wkJittedDefine)," + defineInvocation << "(const std::vector &colNames, " + << "ROOT::Internal::RDF::RColumnRegister &colRegister, " + << "ROOT::Detail::RDF::RLoopManager &lm, " + << "void *jittedDefine, " + << "std::shared_ptr *) {\n"; + defineInvocation << " ROOT::Internal::RDF::JitDefineHelper(" + << funcName << ", " + << "colNames, " + << "\"" << name << "\", " << "colRegister, " - << "prevNodeOnHeap);\n}\n"; - - lm.RegisterJitHelperCall(defineInvocation.str(), upcastNodeOnHeap, definesCopy, parsedExpr.fUsedCols, - MakeWeakOnHeap(jittedDefine)); + << "lm, " + << "reinterpret_cast(jittedDefine)" + << ");\n}\n"; + lm.RegisterJitHelperCall(defineInvocation.str(), std::make_unique(colRegister), + parsedExpr.fUsedCols, jittedDefine); return jittedDefine; } /// Book the jitting of a DefinePerSample call std::shared_ptr BookDefinePerSampleJit(std::string_view name, std::string_view expression, - RLoopManager &lm, const RColumnRegister &colRegister, - std::shared_ptr *upcastNodeOnHeap) + RLoopManager &lm, const RColumnRegister &colRegister) { const auto funcName = DeclareFunction(std::string(expression), {"rdfslot_", "rdfsampleinfo_"}, {"unsigned int", "const ROOT::RDF::RSampleInfo"}); const auto retType = RetTypeOfFunc(funcName); - auto definesCopy = new RColumnRegister(colRegister); - auto definesAddr = PrettyPrintAddr(definesCopy); auto jittedDefine = std::make_shared(name, retType, lm, colRegister, ColumnNames_t{}); // lifetime of pointees: @@ -716,16 +708,21 @@ std::shared_ptr BookDefinePerSampleJit(std::string_view name, std // - jittedDefine: heap-allocated weak_ptr that will be deleted by JitDefineHelper after usage // - definesAddr: heap-allocated, will be deleted by JitDefineHelper after usage std::stringstream defineInvocation; - defineInvocation << "(ROOT::Detail::RDF::RLoopManager *lm, " - << "std::shared_ptr *prevNodeOnHeap," - << "ROOT::Internal::RDF::RColumnRegister* colRegister, " - << "const std::vector & colNames, " - << "void *wkJittedDefine, void *) {\n"; - defineInvocation << "ROOT::Internal::RDF::JitDefineHelper(" - << funcName << ", colNames, \"" << name << "\", lm, " - << "reinterpret_cast*>(wkJittedDefine), " - << "colRegister, prevNodeOnHeap);\n}\n"; - lm.RegisterJitHelperCall(defineInvocation.str(), upcastNodeOnHeap, definesCopy, {}, MakeWeakOnHeap(jittedDefine)); + defineInvocation << "(const std::vector &colNames, " + << "ROOT::Internal::RDF::RColumnRegister &colRegister, " + << "ROOT::Detail::RDF::RLoopManager &lm, " + << "void *jittedDefine, " + << "std::shared_ptr *) {\n"; + defineInvocation << " ROOT::Internal::RDF::JitDefineHelper(" + << funcName << ", " + << "colNames, " + << "\"" << name << "\", " + << "colRegister, " + << "lm, " + << "reinterpret_cast(jittedDefine)" + << ");\n}\n"; + lm.RegisterJitHelperCall(defineInvocation.str(), std::make_unique(colRegister), + {}, jittedDefine); return jittedDefine; } @@ -733,8 +730,7 @@ std::shared_ptr BookDefinePerSampleJit(std::string_view name, std std::shared_ptr BookVariationJit(const std::vector &colNames, std::string_view variationName, const std::vector &variationTags, std::string_view expression, RLoopManager &lm, - RDataSource *ds, const RColumnRegister &colRegister, std::shared_ptr *upcastNodeOnHeap, - bool isSingleColumn) + RDataSource *ds, const RColumnRegister &colRegister, bool isSingleColumn) { const auto &dsColumns = ds ? ds->GetColumnNames() : ColumnNames_t{}; @@ -745,19 +741,13 @@ BookVariationJit(const std::vector &colNames, std::string_view vari const auto type = RetTypeOfFunc(funcName); if (type.rfind("ROOT::VecOps::RVec", 0) != 0) { - // Avoid leak - delete upcastNodeOnHeap; - upcastNodeOnHeap = nullptr; throw std::runtime_error( "Jitted Vary expressions must return an RVec object. The following expression returns a " + type + " instead:\n" + parsedExpr.fExpr); } - auto colRegisterCopy = new RColumnRegister(colRegister); - const auto colRegisterAddr = PrettyPrintAddr(colRegisterCopy); auto jittedVariation = std::make_shared(colNames, variationName, variationTags, type, colRegister, lm, parsedExpr.fUsedCols); - auto variedColsOnHeap = new ColumnNames_t(colNames); // build invocation to JitVariationHelper // variation tag (array of strings) passed as const char** plus size. @@ -767,25 +757,28 @@ BookVariationJit(const std::vector &colNames, std::string_view vari // - definesAddr: heap-allocated, will be deleted by JitDefineHelper after usage // - variedColsOnHeap: deleted by registration function std::stringstream varyInvocation; - varyInvocation << "(ROOT::Detail::RDF::RLoopManager *lm, " - << "std::shared_ptr *prevNodeOnHeap," - << "ROOT::Internal::RDF::RColumnRegister* colRegister, " - << "const std::vector & inputColNames, " - << "void *wkJittedVariation, void *variedColsOnHeap) {\n"; - varyInvocation << "auto * variedColNames = reinterpret_cast*>(variedColsOnHeap);\n"; - varyInvocation << "ROOT::Internal::RDF::JitVariationHelper<" << (isSingleColumn ? "true" : "false") << ">(" - << funcName << ", inputColNames, *variedColNames, "; - varyInvocation << "new const char*[" << variationTags.size() << "]{"; - for (const auto &tag : variationTags) { - varyInvocation << "\"" << tag << "\", "; - } - varyInvocation.seekp(-2, varyInvocation.cur); // remove the last ", " - varyInvocation << "}, " << variationTags.size() << ", \"" << variationName << "\", lm, " - << "reinterpret_cast*>(wkJittedVariation)," - << "colRegister, prevNodeOnHeap);\n" - << "delete variedColNames;\n}\n"; - lm.RegisterJitHelperCall(varyInvocation.str(), upcastNodeOnHeap, colRegisterCopy, parsedExpr.fUsedCols, - MakeWeakOnHeap(jittedVariation), variedColsOnHeap); + varyInvocation << "(const std::vector &inputColNames, " + << "ROOT::Internal::RDF::RColumnRegister &colRegister, " + << "ROOT::Detail::RDF::RLoopManager &lm, " + << "void *jittedVariation, " + << "std::shared_ptr *helperArg) {\n"; + varyInvocation + << " auto *variedColNamesAndTags = reinterpret_cast, " + "std::vector>> *>(helperArg);" + << " ROOT::Internal::RDF::JitVariationHelper<" << (isSingleColumn ? "true" : "false") << ">(" << funcName + << ", " + << "inputColNames, " + << "\"" << variationName << "\", " + << "colRegister, " + << "lm, " + << "reinterpret_cast(jittedVariation), " + << "(*variedColNamesAndTags)->first, " + << "(*variedColNamesAndTags)->second" + << ");\n}\n"; + lm.RegisterJitHelperCall( + varyInvocation.str(), std::make_unique(colRegister), parsedExpr.fUsedCols, + jittedVariation, + std::make_shared, std::vector>>(colNames, variationTags)); return jittedVariation; } @@ -813,19 +806,22 @@ std::string JitBuildAction(const ColumnNames_t &cols, const std::type_info &help // Build a call to CallBuildAction with the appropriate argument. When run through the interpreter, this code will // just-in-time create an RAction object and it will assign it to its corresponding RJittedAction. std::stringstream createAction_str; - createAction_str << "(ROOT::Detail::RDF::RLoopManager *, " - << "std::shared_ptr *prevNodeOnHeap," - << "ROOT::Internal::RDF::RColumnRegister* colRegister, " - << "const std::vector & colNames, " - << "void *wkJittedAction, void *actionArg) {\n"; - createAction_str << "ROOT::Internal::RDF::CallBuildAction<" << actionTypeName; + createAction_str << "(const std::vector &colNames, " + << "ROOT::Internal::RDF::RColumnRegister &colRegister, " + << "ROOT::Detail::RDF::RLoopManager &lm, " + << "void *jittedAction, " + << "std::shared_ptr *helperArg) {\n"; + createAction_str << " ROOT::Internal::RDF::CallBuildAction<" << actionTypeName; const auto columnTypeNames = GetValidatedArgTypes(cols, colRegister, tree, ds, actionTypeNameBase, vector2RVec); for (auto &colType : columnTypeNames) createAction_str << ", " << colType; - createAction_str << ">(prevNodeOnHeap, colNames," << nSlots << ", " - << " reinterpret_cast*>(actionArg)," - << " reinterpret_cast*>(wkJittedAction)," - << "colRegister);\n}\n"; + createAction_str << ">(" + << "colNames, " + << "colRegister, " + << "lm, " + << "reinterpret_cast(jittedAction), " << nSlots << ", " + << "reinterpret_cast *>(helperArg)" + << ");\n}\n"; return createAction_str.str(); } diff --git a/tree/dataframe/src/RDataFrame.cxx b/tree/dataframe/src/RDataFrame.cxx index ba19560d1b143..e96dc6b8e3943 100644 --- a/tree/dataframe/src/RDataFrame.cxx +++ b/tree/dataframe/src/RDataFrame.cxx @@ -2226,17 +2226,6 @@ RDataFrame::RDataFrame(ROOT::RDF::Experimental::RDatasetSpec spec) { } -RDataFrame::~RDataFrame() -{ - // If any node of the computation graph associated with this RDataFrame - // declared code to jit, we need to make sure the compilation actually - // happens. For example, a jitted Define could have been booked but - // if the computation graph is not actually run then the code of the - // Define node is not jitted. This in turn would cause memory leaks. - // See https://github.com/root-project/root/issues/15399 - fLoopManager->Jit(); -} - namespace RDF { namespace Experimental { diff --git a/tree/dataframe/src/RJittedAction.cxx b/tree/dataframe/src/RJittedAction.cxx index cbd6d21aa9b41..65b6f392b6360 100644 --- a/tree/dataframe/src/RJittedAction.cxx +++ b/tree/dataframe/src/RJittedAction.cxx @@ -60,8 +60,9 @@ class GraphNode; RJittedAction::RJittedAction(RLoopManager &lm, const ROOT::RDF::ColumnNames_t &columns, const ROOT::Internal::RDF::RColumnRegister &colRegister, - const std::vector &prevVariations) - : RActionBase(&lm, columns, colRegister, prevVariations) + const std::vector &prevVariations, + std::shared_ptr prevNode) + : RActionBase(&lm, columns, colRegister, prevVariations), fPrevNode(prevNode) { } @@ -159,3 +160,8 @@ std::unique_ptr RJittedAction::CloneAction(voi assert(fConcreteAction != nullptr); return fConcreteAction->CloneAction(newResult); } + +std::shared_ptr RJittedAction::MoveOutPrevNode() +{ + return std::move(fPrevNode); +} diff --git a/tree/dataframe/src/RJittedFilter.cxx b/tree/dataframe/src/RJittedFilter.cxx index 6b755b446b424..9e94e74e1c414 100644 --- a/tree/dataframe/src/RJittedFilter.cxx +++ b/tree/dataframe/src/RJittedFilter.cxx @@ -17,8 +17,10 @@ using namespace ROOT::Detail::RDF; -RJittedFilter::RJittedFilter(RLoopManager *lm, std::string_view name, const std::vector &variations) - : RFilterBase(lm, name, lm->GetNSlots(), RDFInternal::RColumnRegister(lm), /*columnNames*/ {}, variations) +RJittedFilter::RJittedFilter(RLoopManager *lm, std::string_view name, const std::vector &variations, + std::shared_ptr prevNode) + : RFilterBase(lm, name, lm->GetNSlots(), RDFInternal::RColumnRegister(lm), /*columnNames*/ {}, variations), + fPrevNode(prevNode) { // Jitted nodes of the computation graph (e.g. RJittedAction, RJittedDefine) usually don't need to register // themselves with the RLoopManager: the _concrete_ nodes will be registered with the RLoopManager right before @@ -143,3 +145,8 @@ std::shared_ptr RJittedFilter::GetVariedFilter(const std::string &var assert(fConcreteFilter != nullptr); return fConcreteFilter->GetVariedFilter(variationName); } + +std::shared_ptr RJittedFilter::MoveOutPrevNode() +{ + return std::move(fPrevNode); +} diff --git a/tree/dataframe/src/RLoopManager.cxx b/tree/dataframe/src/RLoopManager.cxx index 6ea764f87cbf5..3f533303b1409 100644 --- a/tree/dataframe/src/RLoopManager.cxx +++ b/tree/dataframe/src/RLoopManager.cxx @@ -96,19 +96,64 @@ std::string &GetCodeToDeclare() static std::string code; return code; } -using JitHelperFunc = void(RLoopManager *, std::shared_ptr *, RColumnRegister *, const ColumnNames_t &, - void *, void *); -std::unordered_map &GetJitHelperFuncMap() + +// Signature of all helper functions that are created by JIT helpers, see +// Book*Jit and JitBuildAction in RDFInterfaceUtils.cxx +using JitHelperFunc_t = void (*)(const std::vector &, ROOT::Internal::RDF::RColumnRegister &, + ROOT::Detail::RDF::RLoopManager &, void *, std::shared_ptr *); +std::unordered_map &GetJitHelperFuncMap() { - static std::unordered_map map; + static std::unordered_map map; return map; } -std::unordered_map &GetJitHelperNameMap() +std::unordered_map &GetJitFuncBodyToFuncIdMap() { - static std::unordered_map map; + static std::unordered_map map; return map; } +void DeclareAndRetrieveDeferredJitCalls(const std::string &codeToDeclare) +{ + // This function uses the interpreter and writes to the caches. + R__WRITE_LOCKGUARD(ROOT::gCoreMutex); + + // Step 1: Declare the DeferredJitCall functions to the interpreter + // We use ProcessLine to ensure meta functionality (e.g. autoloading) is + // processed when needed. + // If instead we used Declare, builds with runtime_cxxmodules=OFF would fail + // in jitted actions with custom helpers with errors like: + // error: 'MyHelperType' is an incomplete type + // return std::make_unique(Helper_t(std::move(*h)), bl, std::move(prevNode), colRegister); + // ^ + gInterpreter->ProcessLine(codeToDeclare.c_str()); + + // Step 2: Retrieve the declared functions as function pointers, cache them + // for later use in RunDeferredCalls + auto &funcIdToFuncPointersMap = GetJitHelperFuncMap(); + auto &funcBodyToFuncIdMap = GetJitFuncBodyToFuncIdMap(); + auto clinfo = gInterpreter->ClassInfo_Factory("R_rdf"); + assert(gInterpreter->ClassInfo_IsValid(clinfo)); + + for (auto &codeAndId : funcBodyToFuncIdMap) { + if (auto it = funcIdToFuncPointersMap.find(codeAndId.second); it == funcIdToFuncPointersMap.end()) { + // fast fetch of the address via gInterpreter + // (faster than gInterpreter->Evaluate(function name, ret), ret->GetAsPointer()) + // Retrieve the JIT helper function we registered via RegisterJitHelperCall + auto declid = + gInterpreter->GetFunction(clinfo, ("jitNodeRegistrator_" + std::to_string(codeAndId.second)).c_str()); + assert(declid); + auto minfo = gInterpreter->MethodInfo_Factory(declid); + assert(gInterpreter->MethodInfo_IsValid(minfo)); + auto mname = gInterpreter->MethodInfo_GetMangledName(minfo); + [[maybe_unused]] auto res = funcIdToFuncPointersMap.insert( + {codeAndId.second, reinterpret_cast(gInterpreter->FindSym(mname))}); + assert(res.second); + gInterpreter->MethodInfo_Delete(minfo); + } + } + gInterpreter->ClassInfo_Delete(clinfo); +} + void ThrowIfNSlotsChanged(unsigned int nSlots) { const auto currentSlots = RDFInternal::GetNSlots(); @@ -802,26 +847,7 @@ void RLoopManager::Jit() TStopwatch s; s.Start(); if (!codeToDeclare.empty()) { - ROOT::Internal::RDF::InterpreterDeclare(codeToDeclare); - auto &funcMap = GetJitHelperFuncMap(); - auto &nameMap = GetJitHelperNameMap(); - auto clinfo = gInterpreter->ClassInfo_Factory("R_rdf"); - assert(gInterpreter->ClassInfo_IsValid(clinfo)); - for (auto &codeAndName : nameMap) { - JitHelperFunc *&addr = funcMap[codeAndName.second]; - if (!addr) { - // fast fetch of the address via gInterpreter - // (faster than gInterpreter->Evaluate(function name, ret), ret->GetAsPointer()) - auto declid = gInterpreter->GetFunction(clinfo, codeAndName.second.c_str()); - assert(declid); - auto minfo = gInterpreter->MethodInfo_Factory(declid); - assert(gInterpreter->MethodInfo_IsValid(minfo)); - auto mname = gInterpreter->MethodInfo_GetMangledName(minfo); - addr = reinterpret_cast(gInterpreter->FindSym(mname)); - gInterpreter->MethodInfo_Delete(minfo); - } - } - gInterpreter->ClassInfo_Delete(clinfo); + DeclareAndRetrieveDeferredJitCalls(codeToDeclare); } if (!code.empty()) { RDFInternal::InterpreterCalc(code, "RLoopManager::Run"); @@ -837,19 +863,22 @@ void RLoopManager::Jit() void RLoopManager::RunDeferredCalls() { if (!fJitHelperCalls.empty()) { - R__READ_LOCKGUARD(ROOT::gCoreMutex); // methods are thread-safe but funcMap isn't (yet) - TStopwatch s2; - s2.Start(); - auto &funcMap = GetJitHelperFuncMap(); + // funcMap is not thread-safe + R__READ_LOCKGUARD(ROOT::gCoreMutex); + TStopwatch s; + s.Start(); + const auto &funcMap = GetJitHelperFuncMap(); for (auto &call : fJitHelperCalls) { - assert(funcMap.find(call.functionId) != funcMap.end()); - funcMap[call.functionId](this, call.prevNodeOnHeap, call.colRegister, call.colNames, call.wkJittedNode, - call.argument); + funcMap.at(call.fFunctionId)(call.fColNames, *call.fColRegister, *this, call.fJittedNode.get(), + &call.fExtraArgs); } - s2.Stop(); - R__LOG_INFO(RDFLogChannel()) << "Deferred calls (" << fJitHelperCalls.size() << ") completed" - << (s2.RealTime() > 1e-3 ? " in " + std::to_string(s2.RealTime()) + " seconds." - : " in less than 1ms."); + s.Stop(); + const auto realTime = s.RealTime(); + R__LOG_INFO(RDFLogChannel()) << fJitHelperCalls.size() << " deferred calls completed" + << (realTime > 1e-3 ? " in " + std::to_string(realTime) + " seconds." + : " in less than 1ms."); + // Promoting to write lock to clear the vector + R__WRITE_LOCKGUARD(ROOT::gCoreMutex); fJitHelperCalls.clear(); } } @@ -876,11 +905,15 @@ void RLoopManager::Run(bool jit) // Change value of TTree::GetMaxTreeSize only for this scope. Revert when #6640 will be solved. MaxTreeSizeRAII ctxtmts; + R__LOG_INFO(RDFLogChannel()) << "Starting event loop number " << fNRuns << '.'; + ThrowIfNSlotsChanged(GetNSlots()); if (jit) Jit(); + // Called here since in a RunGraphs run, multiple RLoopManager runs could be + // triggered from different threads. RunDeferredCalls(); InitNodes(); @@ -1002,30 +1035,37 @@ void RLoopManager::ToJitExec(const std::string &code) const GetCodeToJit().append(code); } -void RLoopManager::RegisterJitHelperCall(const std::string &funcCode, std::shared_ptr *prevNodeOnHeap, - ROOT::Internal::RDF::RColumnRegister *colRegister, - const std::vector &colNames, void *wkJittedNode, void *argument) +void RLoopManager::RegisterJitHelperCall(const std::string &funcBody, + std::unique_ptr colRegister, + const std::vector &colNames, std::shared_ptr jittedNode, + std::shared_ptr argument) { - auto &nameMap = GetJitHelperNameMap(); + auto &funcBodyToFuncIdMap = GetJitFuncBodyToFuncIdMap(); { R__READ_LOCKGUARD(ROOT::gCoreMutex); - auto match = nameMap.find(funcCode); - if (match != nameMap.end()) { - R__LOG_DEBUG(0, RDFLogChannel()) << "JitHelper " << match->second << " already defined"; - fJitHelperCalls.emplace_back(match->second, prevNodeOnHeap, colRegister, colNames, wkJittedNode, argument); + auto match = funcBodyToFuncIdMap.find(fStringHasher(funcBody)); + if (match != funcBodyToFuncIdMap.end()) { + R__WRITE_LOCKGUARD(ROOT::gCoreMutex); // modifying fJitHelperCalls + std::string funcName = "jitNodeRegistrator_" + std::to_string(match->second); + R__LOG_DEBUG(0, RDFLogChannel()) << "JIT helper " << funcName << " was already registered."; + fJitHelperCalls.emplace_back(match->second, std::move(colRegister), colNames, jittedNode, argument); return; } } { + // Register lazily a JIT helper R__WRITE_LOCKGUARD(ROOT::gCoreMutex); - std::string registerId = "jitNodeRegistrator_" + std::to_string(nameMap.size()); - nameMap[funcCode] = registerId; - R__LOG_DEBUG(0, RDFLogChannel()) << "JitHelper new " << registerId << " defined for funcCode " << funcCode; - // step 1: register function (now) - std::string toDeclare = "namespace R_rdf {\n void " + registerId + funcCode + "\n}\n"; + auto registratorId = funcBodyToFuncIdMap.size(); + std::string funcName = "jitNodeRegistrator_" + std::to_string(registratorId); + [[maybe_unused]] auto res = funcBodyToFuncIdMap.insert({fStringHasher(funcBody), registratorId}); + assert(res.second); + + std::string toDeclare = "namespace R_rdf {\n void " + funcName + funcBody + "\n}\n"; + R__LOG_DEBUG(0, RDFLogChannel()) << "Registering deferred JIT helper:\n" << toDeclare; + GetCodeToDeclare().append(toDeclare); - fJitHelperCalls.emplace_back(registerId, prevNodeOnHeap, colRegister, colNames, wkJittedNode, argument); + fJitHelperCalls.emplace_back(registratorId, std::move(colRegister), colNames, jittedNode, argument); } } @@ -1332,3 +1372,23 @@ void ROOT::Detail::RDF::RLoopManager::TTreeThreadTask(TTreeReader &treeReader, R (void)entryCount; #endif } + +ROOT::Detail::RDF::RLoopManager::DeferredJitCall::DeferredJitCall( + ROOT::Detail::RDF::RLoopManager::DeferredJitCall &&) noexcept = default; + +ROOT::Detail::RDF::RLoopManager::DeferredJitCall &ROOT::Detail::RDF::RLoopManager::DeferredJitCall::operator=( + ROOT::Detail::RDF::RLoopManager::DeferredJitCall &&) noexcept = default; + +ROOT::Detail::RDF::RLoopManager::DeferredJitCall::~DeferredJitCall() = default; + +ROOT::Detail::RDF::RLoopManager::DeferredJitCall::DeferredJitCall( + std::size_t id, std::unique_ptr colRegisterPtr, + const std::vector &colNamesArg, std::shared_ptr jittedNode, std::shared_ptr argPtr) + : fFunctionId(id), + fColRegister(std::move(colRegisterPtr)), + fColNames(colNamesArg), + fJittedNode(jittedNode), + fExtraArgs(argPtr) +{ + assert(fJittedNode != nullptr); +} From c91b37214da5d19e98bb33ae85e4ef36db987679 Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Thu, 8 Jan 2026 17:31:01 +0100 Subject: [PATCH 3/3] [df] Avoid unnecessary duplication of JIT lines 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> 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. --- .../dataframe/inc/ROOT/RDF/InterfaceUtils.hxx | 21 +++++++++---------- .../dataframe/inc/ROOT/RDF/RVariationBase.hxx | 4 +++- tree/dataframe/src/RDFInterfaceUtils.cxx | 4 ---- tree/dataframe/src/RVariationBase.cxx | 19 +++++++++++++---- 4 files changed, 28 insertions(+), 20 deletions(-) diff --git a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx b/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx index 83c992aad5bb8..854cfb59ff5d2 100644 --- a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx +++ b/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx @@ -469,7 +469,7 @@ void AddDSColumns(const std::vector &requiredCols, ROOT::Detail::RD // this function is meant to be called by the jitted code generated by BookFilterJit template -void JitFilterHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RColumnRegister &colRegister, +void JitFilterHelper(F &&f, const ColumnNames_t &cols, RColumnRegister &colRegister, ROOT::Detail::RDF::RLoopManager &lm, ROOT::Detail::RDF::RJittedFilter *jittedFilter) noexcept { if (!jittedFilter) { @@ -493,7 +493,7 @@ void JitFilterHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RC AddDSColumns(cols, lm, *ds, ColTypes_t(), colRegister); jittedFilter->SetFilter( - std::unique_ptr(new F_t(std::forward(f), cols, prevNode, colRegister, name))); + std::unique_ptr(new F_t(std::forward(f), cols, prevNode, colRegister, jittedFilter->GetName()))); } namespace DefineTypes { @@ -521,7 +521,7 @@ auto MakeDefineNode(DefineTypes::RDefinePerSampleTag, std::string_view name, std // This function is meant to be called by jitted code right before starting the event loop. // If colsPtr is null, build a RDefinePerSample (it has no input columns), otherwise a RDefine. template -void JitDefineHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RColumnRegister &colRegister, +void JitDefineHelper(F &&f, const ColumnNames_t &cols, RColumnRegister &colRegister, ROOT::Detail::RDF::RLoopManager &lm, ROOT::Detail::RDF::RJittedDefine *jittedDefine) noexcept { @@ -543,15 +543,14 @@ void JitDefineHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RC const auto dummyType = "jittedCol_t"; // use unique_ptr instead of make_unique to reduce jit/compile-times std::unique_ptr newCol{ - MakeDefineNode(RDefineTypeTag{}, name, dummyType, std::forward(f), cols, colRegister, lm)}; + MakeDefineNode(RDefineTypeTag{}, jittedDefine->GetName(), dummyType, std::forward(f), cols, colRegister, lm)}; jittedDefine->SetDefine(std::move(newCol)); } template -void JitVariationHelper(F &&f, const ColumnNames_t &inputColNames, std::string_view variationName, - RColumnRegister &colRegister, ROOT::Detail::RDF::RLoopManager &lm, - RJittedVariation *jittedVariation, const ColumnNames_t &variedColNames, - const ColumnNames_t &variationTags) noexcept +void JitVariationHelper(F &&f, const ColumnNames_t &inputColNames, RColumnRegister &colRegister, + ROOT::Detail::RDF::RLoopManager &lm, RJittedVariation *jittedVariation, + const ColumnNames_t &variedColNames, const ColumnNames_t &variationTags) noexcept { if (!jittedVariation) { @@ -568,9 +567,9 @@ void JitVariationHelper(F &&f, const ColumnNames_t &inputColNames, std::string_v AddDSColumns(inputColNames, lm, *ds, ColTypes_t(), colRegister); // use unique_ptr instead of make_unique to reduce jit/compile-times - std::unique_ptr newVariation{ - new RVariation, IsSingleColumn>(variedColNames, variationName, std::forward(f), variationTags, - jittedVariation->GetTypeName(), colRegister, lm, inputColNames)}; + std::unique_ptr newVariation{new RVariation, IsSingleColumn>( + variedColNames, jittedVariation->GetVariationName(), std::forward(f), variationTags, + jittedVariation->GetTypeName(), colRegister, lm, inputColNames)}; jittedVariation->SetVariation(std::move(newVariation)); } diff --git a/tree/dataframe/inc/ROOT/RDF/RVariationBase.hxx b/tree/dataframe/inc/ROOT/RDF/RVariationBase.hxx index da695f9e730da..34b54a9b0a245 100644 --- a/tree/dataframe/inc/ROOT/RDF/RVariationBase.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RVariationBase.hxx @@ -39,7 +39,8 @@ namespace RDF { class RVariationBase { protected: std::vector fColNames; ///< The names of the varied columns. - std::vector fVariationNames; ///< The names of the systematic variation. + std::vector fVariationNames; ///< The tags of the systematic variation. + std::string fVariationName; ///< The name of the systematic variation. std::string fType; ///< The type of the custom column as a text string. std::vector fLastCheckedEntry; RColumnRegister fColumnRegister; @@ -66,6 +67,7 @@ public: virtual const std::type_info &GetTypeId() const = 0; const std::vector &GetColumnNames() const; const std::vector &GetVariationNames() const; + const std::string &GetVariationName() const; std::string GetTypeName() const; /// Update the value at the address returned by GetValuePtr with the content corresponding to the given entry virtual void Update(unsigned int slot, Long64_t entry) = 0; diff --git a/tree/dataframe/src/RDFInterfaceUtils.cxx b/tree/dataframe/src/RDFInterfaceUtils.cxx index 9db800d6bed8b..f5b2db30cf7e1 100644 --- a/tree/dataframe/src/RDFInterfaceUtils.cxx +++ b/tree/dataframe/src/RDFInterfaceUtils.cxx @@ -643,7 +643,6 @@ BookFilterJit(std::shared_ptr prevNode, std::string_view n << "std::shared_ptr *) {\n"; filterInvocation << " ROOT::Internal::RDF::JitFilterHelper(" << funcName << ", " << "colNames, " - << "\"" << name << "\", " << "colRegister, " << "lm, " << "reinterpret_cast(jittedFilter)" @@ -682,7 +681,6 @@ std::shared_ptr BookDefineJit(std::string_view name, std::string_ defineInvocation << " ROOT::Internal::RDF::JitDefineHelper(" << funcName << ", " << "colNames, " - << "\"" << name << "\", " << "colRegister, " << "lm, " << "reinterpret_cast(jittedDefine)" @@ -716,7 +714,6 @@ std::shared_ptr BookDefinePerSampleJit(std::string_view name, std defineInvocation << " ROOT::Internal::RDF::JitDefineHelper(" << funcName << ", " << "colNames, " - << "\"" << name << "\", " << "colRegister, " << "lm, " << "reinterpret_cast(jittedDefine)" @@ -768,7 +765,6 @@ BookVariationJit(const std::vector &colNames, std::string_view vari << " ROOT::Internal::RDF::JitVariationHelper<" << (isSingleColumn ? "true" : "false") << ">(" << funcName << ", " << "inputColNames, " - << "\"" << variationName << "\", " << "colRegister, " << "lm, " << "reinterpret_cast(jittedVariation), " diff --git a/tree/dataframe/src/RVariationBase.cxx b/tree/dataframe/src/RVariationBase.cxx index fa0d9a613b51e..9c04f94e4f160 100644 --- a/tree/dataframe/src/RVariationBase.cxx +++ b/tree/dataframe/src/RVariationBase.cxx @@ -27,13 +27,19 @@ namespace RDF { RVariationBase::RVariationBase(const std::vector &colNames, std::string_view variationName, const std::vector &variationTags, std::string_view type, const RColumnRegister &colRegister, RLoopManager &lm, const ColumnNames_t &inputColNames) - : fColNames(colNames), fVariationNames(variationTags), fType(type), - fLastCheckedEntry(lm.GetNSlots() * CacheLineStep(), -1), fColumnRegister(colRegister), fLoopManager(&lm), - fInputColumns(inputColNames), fIsDefine(inputColNames.size()) + : fColNames(colNames), + fVariationNames(variationTags), + fVariationName(variationName), + fType(type), + fLastCheckedEntry(lm.GetNSlots() * CacheLineStep(), -1), + fColumnRegister(colRegister), + fLoopManager(&lm), + fInputColumns(inputColNames), + fIsDefine(inputColNames.size()) { // prepend the variation name to each tag for (auto &tag : fVariationNames) - tag = std::string(variationName) + ':' + tag; + tag = std::string(fVariationName) + ':' + tag; const auto nColumns = fInputColumns.size(); for (auto i = 0u; i < nColumns; ++i) @@ -52,6 +58,11 @@ const std::vector &RVariationBase::GetVariationNames() const return fVariationNames; } +const std::string &RVariationBase::GetVariationName() const +{ + return fVariationName; +} + std::string RVariationBase::GetTypeName() const { return fType;