From 17735c91440fb9f19bc266b2ee33deaa9c8d3c6c Mon Sep 17 00:00:00 2001
From: Jonas Rembser <jonas.rembser@cern.ch>
Date: Tue, 10 Aug 2021 19:13:46 +0200
Subject: [PATCH] [RF] Bugfix in RooFit driver: fix node computation order

The RooFit driver got the order of the nodes wrong because the correct
order information got lost in the `std::unordered_map`. A new
`std::vector` is introduced to keep track of the evaluation order.
---
 roofit/roofitcore/inc/RooFitDriver.h   |  2 ++
 roofit/roofitcore/src/RooFitDriver.cxx | 19 +++++++++++--------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/roofit/roofitcore/inc/RooFitDriver.h b/roofit/roofitcore/inc/RooFitDriver.h
index 4df96a1beab2a..a40c2c4e49971 100644
--- a/roofit/roofitcore/inc/RooFitDriver.h
+++ b/roofit/roofitcore/inc/RooFitDriver.h
@@ -61,6 +61,8 @@ class RooFitDriver {
     const RooNLLVarNew& _topNode;
     const RooAbsData* const _data = nullptr;
     const size_t _nEvents;
+
+    std::vector<const RooAbsReal*> _nodes;
     std::unordered_map<const RooAbsReal*, NodeInfo> _nodeInfos;
 
     //used for preserving resources
diff --git a/roofit/roofitcore/src/RooFitDriver.cxx b/roofit/roofitcore/src/RooFitDriver.cxx
index cf66b9307077a..688c5a7d7fb80 100644
--- a/roofit/roofitcore/src/RooFitDriver.cxx
+++ b/roofit/roofitcore/src/RooFitDriver.cxx
@@ -64,6 +64,8 @@ RooFitDriver::RooFitDriver(const RooAbsData& data, const RooNLLVarNew& topNode,
     }
     else //this node needs evaluation, mark it's clients
     {
+      _nodes.push_back(pAbsReal);
+
       // If the node doesn't depend on any observables, there is no need to
       // loop over events and we don't need to use the batched evaluation.
       RooArgSet observablesForNode;
@@ -155,7 +157,8 @@ double RooFitDriver::getVal()
   for (const auto& it:_nodeInfos)
     if (it.second.remServers==0 && it.second.computeInGPU)
       assignToGPU(it.first);
-  
+
+
   int nNodes = _nodeInfos.size();
   while (nNodes)
   {
@@ -171,22 +174,22 @@ double RooFitDriver::getVal()
         }
     
     // find next cpu node
-    auto it=_nodeInfos.begin();
-    for ( ; it!=_nodeInfos.end(); it++)
-      if (it->second.remServers==0 && !it->second.computeInGPU) break;
+    auto it = std::find_if(_nodes.begin(), _nodes.end(), [&](const RooAbsReal* a){
+            auto const& info = _nodeInfos[a]; return info.remServers==0 && !info.computeInGPU; });
 
     // if no cpu node available sleep for a while to save cpu usage
-    if (it==_nodeInfos.end())
+    if (it==_nodes.end())
     {
       std::this_thread::sleep_for(std::chrono::milliseconds(1));
       continue;
     }
-    
+
     // compute next cpu node
-    const RooAbsReal* node = it->first;
-    NodeInfo& info = it->second;
+    const RooAbsReal* node = *it;
+    NodeInfo& info = _nodeInfos[*it];
     info.remServers=-2; //so that it doesn't get picked again
     nNodes--;
+
     if (info.computeInScalarMode)
     {
       _nonDerivedValues.push_back(node->getVal(_data->get()));