Changes to memory ownership of jitted nodes

Author: vepadulano

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<RJittedFilter> BookFilterJit(std::shared_ptr<RNodeBase> *prevNodeOnHeap, std::string_view name,
+std::shared_ptr<RJittedFilter> BookFilterJit(std::shared_ptr<RNodeBase> prevNode, std::string_view name,
                                              std::string_view expression, const RColumnRegister &colRegister,
                                              TTree *tree, RDataSource *ds);
 
 std::shared_ptr<RJittedDefine> BookDefineJit(std::string_view name, std::string_view expression, RLoopManager &lm,
-                                             RDataSource *ds, const RColumnRegister &colRegister,
-                                             std::shared_ptr<RNodeBase> *prevNodeOnHeap);
+                                             RDataSource *ds, const RColumnRegister &colRegister);
 
 std::shared_ptr<RJittedDefine> BookDefinePerSampleJit(std::string_view name, std::string_view expression,
-                                                      RLoopManager &lm, const RColumnRegister &colRegister,
-                                                      std::shared_ptr<RNodeBase> *upcastNodeOnHeap);
+                                                      RLoopManager &lm, const RColumnRegister &colRegister);
 
 std::shared_ptr<RJittedVariation>
 BookVariationJit(const std::vector<std::string> &colNames, std::string_view variationName,
                  const std::vector<std::string> &variationTags, std::string_view expression, RLoopManager &lm,
-                 RDataSource *ds, const RColumnRegister &colRegister, std::shared_ptr<RNodeBase> *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<std::string> &requiredCols, ROOT::Detail::RD
                   ROOT::Internal::RDF::RColumnRegister &colRegister);
 
 // this function is meant to be called by the jitted code generated by BookFilterJit
-template <typename F, typename PrevNode>
-void JitFilterHelper(F &&f, const ColumnNames_t &cols, std::string_view name,
-                     std::weak_ptr<RJittedFilter> *wkJittedFilter, std::shared_ptr<PrevNode> *prevNodeOnHeap,
-                     RColumnRegister *colRegister) noexcept
+template <typename F>
+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<F>;
-   using F_t = RFilter<Callable_t, PrevNode>;
+   auto prevNode = jittedFilter->MoveOutPrevNode();
+   using PrevNode_t = typename decltype(prevNode)::element_type;
+   using F_t = RFilter<Callable_t, PrevNode_t>;
    using ColTypes_t = typename TTraits::CallableTraits<Callable_t>::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<RFilterBase>(new F_t(std::forward<F>(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<RFilterBase>(new F_t(std::forward<F>(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 <typename RDefineTypeTag, typename F>
-void JitDefineHelper(F &&f, const ColumnNames_t &cols, std::string_view name, RLoopManager *lm,
-                     std::weak_ptr<RJittedDefine> *wkJittedDefine, RColumnRegister *colRegister,
-                     std::shared_ptr<RNodeBase> *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<F>;
    using ColTypes_t = typename TTraits::CallableTraits<Callable_t>::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<RDefineBase> instead of make_unique<NewCol_t> to reduce jit/compile-times
    std::unique_ptr<RDefineBase> newCol{
-      MakeDefineNode(RDefineTypeTag{}, name, dummyType, std::forward<F>(f), cols, *colRegister, *lm)};
+      MakeDefineNode(RDefineTypeTag{}, name, dummyType, std::forward<F>(f), cols, colRegister, lm)};
    jittedDefine->SetDefine(std::move(newCol));
-
-   doDeletes();
 }
 
 template <bool IsSingleColumn, typename F>
-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<RJittedVariation> *wkJittedVariation,
-                        RColumnRegister *colRegister, std::shared_ptr<RNodeBase> *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<std::string> tags(variationTags, variationTags + variationTagsSize);
-
-   auto jittedVariation = wkJittedVariation->lock();
-
    using Callable_t = std::decay_t<F>;
    using ColTypes_t = typename TTraits::CallableTraits<Callable_t>::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<RDefineBase> instead of make_unique<NewCol_t> to reduce jit/compile-times
-   std::unique_ptr<RVariationBase> newVariation{new RVariation<std::decay_t<F>, IsSingleColumn>(
-      std::move(variedColNames), variationName, std::forward<F>(f), std::move(tags), jittedVariation->GetTypeName(),
-      *colRegister, *lm, inputColNames)};
+   std::unique_ptr<RVariationBase> newVariation{
+      new RVariation<std::decay_t<F>, IsSingleColumn>(variedColNames, variationName, std::forward<F>(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 <typename ActionTag, typename... ColTypes, typename PrevNodeType, typename HelperArgType>
-void CallBuildAction(std::shared_ptr<PrevNodeType> *prevNodeOnHeap, const ColumnNames_t &cols,
-                     const unsigned int nSlots, std::shared_ptr<HelperArgType> *helperArgOnHeap,
-                     std::weak_ptr<RJittedAction> *wkJittedActionOnHeap, RColumnRegister *colRegister) noexcept
+template <typename ActionTag, typename... ColTypes, typename HelperArgType>
+void CallBuildAction(const ColumnNames_t &cols, RColumnRegister &colRegister, ROOT::Detail::RDF::RLoopManager &lm,
+                     RJittedAction *jittedAction, unsigned int nSlots,
+                     std::shared_ptr<HelperArgType> *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<ColTypes...>;
    constexpr auto nColumns = ColTypes_t::list_size;
-   auto ds = loopManager.GetDataSource();
-   if (ds != nullptr)
-      AddDSColumns(cols, loopManager, *ds, ColTypes_t(), *colRegister);
-
-   auto actionPtr = BuildAction<ColTypes...>(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<ColTypes...>(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<U>`, `T` otherwise.

tree/dataframe/inc/ROOT/RDF/RInterface.hxx

@@ -293,12 +293,8 @@ public:
    /// ~~~
    RInterface<RDFDetail::RJittedFilter, DS_t> 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<decltype(upcastNodeOnHeap)>::element_type;
-      RInterface<BaseNodeType_t> 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<RDFDetail::RJittedFilter, DS_t>(std::move(jittedFilter), *fLoopManager, fColRegister);
    }
@@ -540,9 +536,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));
@@ -630,9 +624,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));
@@ -807,9 +799,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));
@@ -3420,10 +3410,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));

tree/dataframe/inc/ROOT/RDF/RInterfaceBase.hxx

@@ -196,19 +196,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<RDFInternal::RJittedAction>(
+         *fLoopManager, validColumnNames, fColRegister, proxiedPtr->GetVariations(), proxiedPtr);
 
-      auto upcastNodeOnHeap = RDFInternal::MakeSharedOnHeap(RDFInternal::UpcastNode(proxiedPtr));
-
-      const auto jittedAction = std::make_shared<RDFInternal::RJittedAction>(*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<ROOT::Internal::RDF::RColumnRegister>(fColRegister),
+                                          validColumnNames, jittedAction, helperArg);
       return MakeResultPtr(r, *fLoopManager, std::move(jittedAction));
    }
 

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<RActionBase> fConcreteAction;
+   std::shared_ptr<ROOT::Detail::RDF::RNodeBase> fPrevNode;
 
 public:
    RJittedAction(RLoopManager &lm, const ROOT::RDF::ColumnNames_t &columns, const RColumnRegister &colRegister,
-                 const std::vector<std::string> &prevVariations);
+                 const std::vector<std::string> &prevVariations,
+                 std::shared_ptr<ROOT::Detail::RDF::RNodeBase> prevNode = nullptr);
    ~RJittedAction();
 
    void SetAction(std::unique_ptr<RActionBase> a) { fConcreteAction = std::move(a); }
@@ -67,6 +70,7 @@ public:
 
    std::unique_ptr<RActionBase> MakeVariedAction(std::vector<void *> &&results) final;
    std::unique_ptr<ROOT::Internal::RDF::RActionBase> CloneAction(void *newResult) final;
+   std::shared_ptr<ROOT::Detail::RDF::RNodeBase> MoveOutPrevNode();
 };
 
 } // ns RDF

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<RFilterBase> fConcreteFilter = nullptr;
+   std::shared_ptr<RNodeBase> fPrevNode;
 
 public:
-   RJittedFilter(RLoopManager *lm, std::string_view name, const std::vector<std::string> &variations);
+   RJittedFilter(RLoopManager *lm, std::string_view name, const std::vector<std::string> &variations,
+                 std::shared_ptr<ROOT::Detail::RDF::RNodeBase> prevNode = nullptr);
 
    // Rule of five
 
@@ -68,6 +70,7 @@ public:
    std::shared_ptr<RDFGraphDrawing::GraphNode>
    GetGraph(std::unordered_map<void *, std::shared_ptr<RDFGraphDrawing::GraphNode>> &visitedMap) final;
    std::shared_ptr<RNodeBase> GetVariedFilter(const std::string &variationName) final;
+   std::shared_ptr<RNodeBase> MoveOutPrevNode();
 };
 
 } // ns RDF