From a14b0979d50d28fcaf99648ce50b53117d784f3f Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Sun, 18 Oct 2026 08:27:41 +0000 Subject: [PATCH] [RF] Skip clean nodes in the CUDA evaluation path of RooFit::Evaluator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CPU path in RooFit::Evaluator::run() only recomputes nodes that are dirty, i.e. that depend on a variable whose value changed since the last evaluation. The CUDA path in getValHeterogeneous() re-enqueued the whole computation graph on every call, because the buffers of intermediate results were released as soon as the clients of a node were computed, so nothing could be reused. This commit uses the same dirty tracking in the CUDA path: the output buffers of the nodes now persist across evaluations, and a node is only recomputed if one of its value servers was recomputed in the same evaluation or if a variable it depends on changed. Since no buffer is recycled within an evaluation anymore, the special handling of pinned buffers for the eager release is not needed either. During a minimization, most likelihood evaluations change only one or two parameters: the numeric gradient varies one parameter at a time, and Hesse two. For a RooAddPdf with several components, only the component that depends on the varied parameter needs to be evaluated again, and for a varied yield no component kernel has to run at all. Measured with the FitModel case from rootbench (5-component extended RooAddPdf, 15 parameters, Minuit2) on an RTX A4500, with identical results, number of function calls and parameter errors before and after: events before after 10k 0.29 s 0.12 s 100k 0.77 s 0.31 s 1M 5.05 s 2.12 s At 1M events the kernel time per likelihood evaluation goes from 2.2 ms to 0.94 ms, with 5.4 instead of 13 kernel launches and 5.6 instead of 17 memcpys per evaluation on average. The price is that the device memory for the intermediate results of all vector nodes stays allocated for the lifetime of the Evaluator instead of being recycled between nodes of the same size. For the usual computation graphs this is a few arrays of the size of the dataset. 🤖 Done with the help of AI --- roofit/roofitcore/src/RooFit/Evaluator.cxx | 87 ++++++++++------------ 1 file changed, 40 insertions(+), 47 deletions(-) diff --git a/roofit/roofitcore/src/RooFit/Evaluator.cxx b/roofit/roofitcore/src/RooFit/Evaluator.cxx index 4c647249d9d2f..6a64e94aa9016 100644 --- a/roofit/roofitcore/src/RooFit/Evaluator.cxx +++ b/roofit/roofitcore/src/RooFit/Evaluator.cxx @@ -127,7 +127,6 @@ struct NodeInfo { std::shared_ptr buffer; std::size_t iNode = 0; - int remClients = 0; bool copyAfterEvaluation = false; bool fromArrayInput = false; bool isVariable = false; @@ -142,21 +141,6 @@ struct NodeInfo { double scalarBuffer = 0.0; std::vector serverInfos; std::vector clientInfos; - - /// Check the servers of a node that has been computed and release its - /// resources if they are no longer needed. Buffers of nodes whose results - /// are copied between host and device (copyAfterEvaluation) must not be - /// released eagerly: their pinned host memory can still be the source of - /// an asynchronous copy that was enqueued on the CUDA stream, and a new - /// owner would overwrite it from the CPU without any stream ordering. - /// Those buffers are released at the beginning of the next evaluation - /// instead, after the stream was synchronized at the end of this one. - void decrementRemainingClients() - { - if (--remClients == 0 && !fromArrayInput && !copyAfterEvaluation) { - buffer.reset(); - } - } }; /// Construct a new Evaluator. The constructor analyzes and saves metadata about the graph, @@ -513,38 +497,45 @@ std::span Evaluator::run() /// Returns the value of the top node in the computation graph std::span Evaluator::getValHeterogeneous() { - for (auto &info : _nodes) { - info.remClients = info.clientInfos.size(); - if (info.buffer && !info.fromArrayInput) { - info.buffer.reset(); - } - } - - // Iterate over the nodes in topological order. Nodes that are computed on - // the GPU only enqueue their computation on the single CUDA stream and - // return immediately, so independent CPU nodes that come later in the - // ordering naturally overlap with the GPU computations. Ordering by the - // stream guarantees that GPU nodes see the results of their GPU servers, - // and host-side reads of GPU results synchronize on the stream in the - // buffer implementation. + // Iterate over the nodes in topological order. Like in the CPU-only path + // in run(), only dirty nodes are recomputed: a node is dirty if one of its + // value servers was recomputed in this evaluation, or if a variable + // changed its value since the last evaluation. Clean nodes keep the output + // buffers from their last evaluation, which is why the buffers persist + // across evaluations in this mode instead of being released as soon as + // the clients of a node are computed. During a minimization, most + // likelihood evaluations change only one or two parameters (e.g. for the + // numeric gradient and the Hessian), so skipping the clean nodes saves + // most of the GPU work and the kernel launches. The price is the device + // memory for the intermediate results of the whole computation graph, + // which is small compared to the input data for the usual graphs where + // most nodes are scalars. + // + // Nodes that are computed on the GPU only enqueue their computation on the + // single CUDA stream and return immediately, so independent CPU nodes that + // come later in the ordering naturally overlap with the GPU computations. + // Ordering by the stream guarantees that GPU nodes see the results of + // their GPU servers, and host-side reads of GPU results synchronize on the + // stream in the buffer implementation. Since no buffer is recycled within + // an evaluation, the results of clean nodes can't be overwritten by + // enqueued work either. try { for (auto &info : _nodes) { - if (!info.fromArrayInput) { + if (info.fromArrayInput) { + continue; + } + if (info.isVariable) { + processVariable(info); + } else if (info.isCategory) { + processCategory(info); + } else if (info.isDirty) { + setClientsDirty(info); if (info.computeInGPU) { assignToGPU(info); } else { computeCPUNode(info.absArg, info); } - } - - // Release the buffers of server nodes that are no longer needed. For - // device-only buffers this is safe to do right away even if GPU work - // is still in flight, because any reuse of a released device buffer - // happens through operations that are enqueued later on the same - // stream. Pinned buffers are exempted from the eager release, see - // the comment in NodeInfo::decrementRemainingClients(). - for (auto *serverInfo : info.serverInfos) { - serverInfo->decrementRemainingClients(); + info.isDirty = false; } } } catch (...) { @@ -552,6 +543,8 @@ std::span Evaluator::getValHeterogeneous() // deferred may still be armed. Deliver them now, while the destination // memory in the nodes of the computation graph is guaranteed to be // alive, so that no armed readback survives into a later evaluation. + // The node that threw stays dirty and gets recomputed in the next + // evaluation, together with its clients that were already flagged. try { RooBatchCompute::dispatchCUDA->synchronizeCudaStream(_cudaStream); } catch (...) { @@ -566,10 +559,8 @@ std::span Evaluator::getValHeterogeneous() // Ensure that all enqueued GPU work has completed when run() returns. For // the usual likelihood evaluations this is mostly a no-op, because the - // final reduction has synchronized the stream already. It also guarantees - // that recycling the buffers at the beginning of the next evaluation is - // safe, and it delivers the deferred readbacks like the evaluation error - // counters. + // final reduction has synchronized the stream already. It also delivers + // the deferred readbacks like the evaluation error counters. RooBatchCompute::dispatchCUDA->synchronizeCudaStream(_cudaStream); // Run the deferred actions now that all results have arrived on the host, @@ -601,8 +592,10 @@ void Evaluator::assignToGPU(NodeInfo &info) buffer = &info.scalarBuffer; _evalContextCPU.set(node, {buffer, nOut}); } else { - info.buffer = info.copyAfterEvaluation ? _bufferManager->makePinnedBuffer(nOut, _cudaStream) - : _bufferManager->makeGpuBuffer(nOut); + if (!info.buffer) { + info.buffer = info.copyAfterEvaluation ? _bufferManager->makePinnedBuffer(nOut, _cudaStream) + : _bufferManager->makeGpuBuffer(nOut); + } buffer = info.buffer->deviceWritePtr(); } assignSpan(_evalContextCUDA._currentOutput, {buffer, nOut});