From d99cff16586950a83e5aa2271b48d98214580592 Mon Sep 17 00:00:00 2001
From: Jonas Rembser <jonas.rembser@cern.ch>
Date: Wed, 9 Sep 2026 08:19:32 +0000
Subject: [PATCH] [RF] Speed up CodegenContext::beginLoop() with single graph
 traversal
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

To figure out which vector observables to loop over, beginLoop() called
RooAbsArg::dependsOn() once per vector observable. These dependsOn()
walks don't deduplicate visited nodes, so their cost scales with the
number of paths in the computation graph instead of the number of
nodes, which gets very expensive for large models with many shared
nodes.

Replace the per-observable walks with a single depth-first traversal
that visits every node only once and collects the reachable vector
observables, which is equivalent.

For the ATLAS VHbb benchmark workspace from rootbench, this reduces the
"Function JIT time" of the codegen backend (which includes the code
generation itself) from 9.4 s to 5.3 s.

🤖 Done with the help of AI
---
 .../roofitcore/src/RooFit/CodegenContext.cxx  | 25 ++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/roofit/roofitcore/src/RooFit/CodegenContext.cxx b/roofit/roofitcore/src/RooFit/CodegenContext.cxx
index 72338604602da..a4c7abb464a76 100644
--- a/roofit/roofitcore/src/RooFit/CodegenContext.cxx
+++ b/roofit/roofitcore/src/RooFit/CodegenContext.cxx
@@ -25,6 +25,7 @@
 #include <locale>
 #include <type_traits>
 #include <unordered_map>
+#include <unordered_set>
 
 namespace {
 
@@ -166,13 +167,35 @@ std::unique_ptr<CodegenContext::LoopScope> CodegenContext::beginLoop(RooAbsArg c
 
    std::vector<TNamed const *> vars;
 
+   // Figure out which vector observables are in the server tree of "in" with
+   // a single depth-first traversal that visits each node only once. This is
+   // equivalent to calling RooAbsArg::dependsOn() for each vector observable,
+   // but much faster for large computation graphs: dependsOn() doesn't
+   // deduplicate the visited nodes, so its cost scales with the number of
+   // paths in the graph instead of the number of nodes.
+   std::unordered_set<TNamed const *> reachableVecObs;
+   {
+      std::unordered_set<RooAbsArg const *> visited;
+      std::vector<RooAbsArg const *> stack{in};
+      while (!stack.empty()) {
+         RooAbsArg const *arg = stack.back();
+         stack.pop_back();
+         if (!visited.insert(arg).second)
+            continue;
+         if (_vecObsIndices.find(arg->namePtr()) != _vecObsIndices.end())
+            reachableVecObs.insert(arg->namePtr());
+         for (RooAbsArg const *server : arg->servers())
+            stack.push_back(server);
+      }
+   }
+
    // Set the results of the vector observables.
    // TODO: we are using the size of the first loop variable to the the number
    // of iterations, but it should be made sure that all loop vars are either
    // scalar or have the same size.
    int firstObsIdx = -1;
    for (auto const &it : _vecObsIndices) {
-      if (!in->dependsOn(it.first))
+      if (reachableVecObs.find(it.first) == reachableVecObs.end())
          continue;
 
       vars.push_back(it.first);