From 578a6fe47d187923a22e503962a23f491e1d60e7 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 4 Sep 2026 16:31:01 +0000 Subject: [PATCH 1/8] [RF] Use a single CUDA stream in the RooFit::Evaluator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the per-node CUDA streams and events plus the polling scheduler in the RooFit::Evaluator with a single stream per Evaluator instance on which all GPU work is enqueued in topological order. Rationale: * The per-node multi-stream design never delivered actual concurrency: all data transfers go through synchronous cudaMemcpy calls on the legacy default stream, which serializes across streams. Independent measurements show the single-stream version is never slower and up to 7 % faster on multi-channel simultaneous fits, the workload the multi-stream design targeted. * The elementwise kernels saturate the GPU by themselves at realistic event counts, so overlapping them across streams cannot win anything. In the small-kernel regime, launch and allocation overhead dominates, which is better addressed with pooled allocations and CUDA graph capture. Both become much simpler with a fixed single-stream launch sequence. * CPU/GPU overlap is preserved: kernel launches are asynchronous, so the topological loop keeps evaluating independent CPU nodes while enqueued GPU work is in flight. Host-side reads of GPU results synchronize in the buffer implementation. This removes the remServers state machine, the 1 ms sleep polling loop, and the per-node stream and event bookkeeping (net -47 lines). Benchmarks on an RTX A4500 (Gaussian+Exponential fits, wall time, identical minNll in all cases): single channel: 1M events: 0.220 s -> 0.217 s, 10M: 1.31 -> 1.30 s 8-channel simultaneous, 125k/chan: 1.63 s -> 1.57 s 8-channel simultaneous, 12.5k/chan: 0.90 s -> 0.84 s test-stressroofit-cuda and test-stressroostats-cuda pass. 🤖 Done with the help of AI --- roofit/roofitcore/inc/RooFit/Evaluator.h | 5 + roofit/roofitcore/src/RooFit/Evaluator.cxx | 120 ++++++--------------- 2 files changed, 39 insertions(+), 86 deletions(-) diff --git a/roofit/roofitcore/inc/RooFit/Evaluator.h b/roofit/roofitcore/inc/RooFit/Evaluator.h index f50424c634abf..8c325a0b154ec 100644 --- a/roofit/roofitcore/inc/RooFit/Evaluator.h +++ b/roofit/roofitcore/inc/RooFit/Evaluator.h @@ -26,7 +26,10 @@ class RooAbsArg; namespace RooBatchCompute { class AbsBufferManager; +namespace CudaInterface { +class CudaStream; } +} // namespace RooBatchCompute namespace RooFit { @@ -68,6 +71,8 @@ class Evaluator { std::vector _nodes; // the ordered computation graph std::unordered_map _nodesMap; // for quick lookup of nodes std::unique_ptr _operModeChanges; + // the single CUDA stream on which all GPU work of this Evaluator is enqueued + RooBatchCompute::CudaInterface::CudaStream *_cudaStream = nullptr; }; } // end namespace RooFit diff --git a/roofit/roofitcore/src/RooFit/Evaluator.cxx b/roofit/roofitcore/src/RooFit/Evaluator.cxx index 06ced37a724a4..9abefb47055c8 100644 --- a/roofit/roofitcore/src/RooFit/Evaluator.cxx +++ b/roofit/roofitcore/src/RooFit/Evaluator.cxx @@ -42,10 +42,8 @@ RooAbsPdf::fitTo() is called and gets destroyed when the fitting ends. #include "BatchModeDataHelpers.h" #include "RooFitImplHelpers.h" -#include #include #include -#include #include namespace RooFit { @@ -105,7 +103,6 @@ struct NodeInfo { std::shared_ptr buffer; std::size_t iNode = 0; int remClients = 0; - int remServers = 0; bool copyAfterEvaluation = false; bool fromArrayInput = false; bool isVariable = false; @@ -121,9 +118,6 @@ struct NodeInfo { std::vector serverInfos; std::vector clientInfos; - RooBatchCompute::CudaInterface::CudaEvent *event = nullptr; - RooBatchCompute::CudaInterface::CudaStream *stream = nullptr; - /// Check the servers of a node that has been computed and release its /// resources if they are no longer needed. void decrementRemainingClients() @@ -132,14 +126,6 @@ struct NodeInfo { buffer.reset(); } } - - ~NodeInfo() - { - if (event) - RooBatchCompute::dispatchCUDA->deleteCudaEvent(event); - if (stream) - RooBatchCompute::dispatchCUDA->deleteCudaStream(stream); - } }; /// Construct a new Evaluator. The constructor analyzes and saves metadata about the graph, @@ -222,12 +208,14 @@ Evaluator::Evaluator(const RooAbsReal &absReal, bool useGPU) syncDataTokens(); if (_useGPU) { - // create events and streams for every node + // Create the single CUDA stream on which all GPU computations and data + // transfers of this Evaluator are enqueued. The graph is evaluated in + // topological order, so ordering the operations by the stream is enough + // to guarantee correct results. + _cudaStream = RooBatchCompute::dispatchCUDA->newCudaStream(); + RooBatchCompute::Config cfg; + cfg.setCudaStream(_cudaStream); for (auto &info : _nodes) { - info.event = RooBatchCompute::dispatchCUDA->newCudaEvent(false); - info.stream = RooBatchCompute::dispatchCUDA->newCudaStream(); - RooBatchCompute::Config cfg; - cfg.setCudaStream(info.stream); _evalContextCUDA.setConfig(info.absArg, cfg); } } @@ -342,6 +330,9 @@ Evaluator::~Evaluator() info.absArg->resetDataToken(); } } + if (_cudaStream) { + RooBatchCompute::dispatchCUDA->deleteCudaStream(_cudaStream); + } } void Evaluator::computeCPUNode(const RooAbsArg *node, NodeInfo &info) @@ -366,7 +357,7 @@ void Evaluator::computeCPUNode(const RooAbsArg *node, NodeInfo &info) info.hasLogged = true; } if (!info.buffer) { - info.buffer = info.copyAfterEvaluation ? _bufferManager->makePinnedBuffer(nOut, info.stream) + info.buffer = info.copyAfterEvaluation ? _bufferManager->makePinnedBuffer(nOut, _cudaStream) : _bufferManager->makeCpuBuffer(nOut); } buffer = info.buffer->hostWritePtr(); @@ -390,10 +381,10 @@ void Evaluator::computeCPUNode(const RooAbsArg *node, NodeInfo &info) _evalContextCPU.resetVectorBuffers(); _evalContextCPU.enableVectorBuffers(false); if (info.copyAfterEvaluation) { + // The deviceReadPtr() call triggers the copy of the result to the GPU. + // The copy is ordered by the CUDA stream, so GPU clients enqueued later + // will see the result without any further synchronization. _evalContextCUDA.set(node, {info.buffer->deviceReadPtr(), nOut}); - if (info.event) { - RooBatchCompute::dispatchCUDA->cudaEventRecord(info.event, info.stream); - } } } @@ -475,66 +466,31 @@ std::span Evaluator::getValHeterogeneous() { for (auto &info : _nodes) { info.remClients = info.clientInfos.size(); - info.remServers = info.serverInfos.size(); if (info.buffer && !info.fromArrayInput) { info.buffer.reset(); } } - // find initial GPU nodes and assign them to GPU + // 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. for (auto &info : _nodes) { - if (info.remServers == 0 && info.computeInGPU) { - assignToGPU(info); - } - } - - NodeInfo const &topNodeInfo = _nodes.back(); - while (topNodeInfo.remServers != -2) { - // find finished GPU nodes - for (auto &info : _nodes) { - if (info.remServers == -1 && !RooBatchCompute::dispatchCUDA->cudaStreamIsActive(info.stream)) { - info.remServers = -2; - // Decrement number of remaining servers for clients and start GPU computations - for (auto *infoClient : info.clientInfos) { - --infoClient->remServers; - if (infoClient->computeInGPU && infoClient->remServers == 0) { - assignToGPU(*infoClient); - } - } - for (auto *serverInfo : info.serverInfos) { - serverInfo->decrementRemainingClients(); - } - } - } - - // find next CPU node - auto it = _nodes.begin(); - for (; it != _nodes.end(); it++) { - if (it->remServers == 0 && !it->computeInGPU) - break; - } - - // if no CPU node available sleep for a while to save CPU usage - if (it == _nodes.end()) { - std::this_thread::sleep_for(std::chrono::milliseconds(1)); - continue; - } - - // compute next CPU node - NodeInfo &info = *it; - RooAbsArg const *node = info.absArg; - info.remServers = -2; // so that it doesn't get picked again - if (!info.fromArrayInput) { - computeCPUNode(node, info); - } - - // Assign the clients that are computed on the GPU - for (auto *infoClient : info.clientInfos) { - if (--infoClient->remServers == 0 && infoClient->computeInGPU) { - assignToGPU(*infoClient); + if (info.computeInGPU) { + assignToGPU(info); + } else { + computeCPUNode(info.absArg, info); } } + + // Release the buffers of server nodes that are no longer needed. 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. for (auto *serverInfo : info.serverInfos) { serverInfo->decrementRemainingClients(); } @@ -544,22 +500,13 @@ std::span Evaluator::getValHeterogeneous() return _evalContextCUDA.at(&_topNode); } -/// Assign a node to be computed in the GPU. Scan it's clients and also assign them -/// in case they only depend on GPU nodes. +/// Enqueue the computation of a node on the GPU. void Evaluator::assignToGPU(NodeInfo &info) { using namespace Detail; - info.remServers = -1; - auto node = static_cast(info.absArg); - // wait for every server to finish - for (auto *infoServer : info.serverInfos) { - if (infoServer->event) - RooBatchCompute::dispatchCUDA->cudaStreamWaitForEvent(info.stream, infoServer->event); - } - const std::size_t nOut = info.outputSize; double *buffer = nullptr; @@ -567,15 +514,16 @@ void Evaluator::assignToGPU(NodeInfo &info) buffer = &info.scalarBuffer; _evalContextCPU.set(node, {buffer, nOut}); } else { - info.buffer = info.copyAfterEvaluation ? _bufferManager->makePinnedBuffer(nOut, info.stream) + info.buffer = info.copyAfterEvaluation ? _bufferManager->makePinnedBuffer(nOut, _cudaStream) : _bufferManager->makeGpuBuffer(nOut); buffer = info.buffer->deviceWritePtr(); } assignSpan(_evalContextCUDA._currentOutput, {buffer, nOut}); _evalContextCUDA.set(node, {buffer, nOut}); node->doEval(_evalContextCUDA); - RooBatchCompute::dispatchCUDA->cudaEventRecord(info.event, info.stream); if (info.copyAfterEvaluation) { + // The hostReadPtr() call triggers the copy of the result to the host, + // which waits for the enqueued computation via the CUDA stream. _evalContextCPU.set(node, {info.buffer->hostReadPtr(), nOut}); } } From f478bade23fc14988d027b1a337ecd7424ca23f9 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 4 Sep 2026 16:31:01 +0000 Subject: [PATCH 2/8] [RF] Remove unused CUDA event API from RooBatchComputeInterface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After the single-stream refactoring of the RooFit::Evaluator, the per-node CUDA event and stream-query methods on the RooBatchComputeInterface have no users left. Remove them, together with the CudaEvent wrapper and the event-related functions in the CudaInterface, and the throwing stubs in the CPU implementation. 🤖 Done with the help of AI --- roofit/batchcompute/res/RooBatchCompute.h | 6 -- roofit/batchcompute/src/CudaInterface.cu | 74 --------------------- roofit/batchcompute/src/CudaInterface.h | 21 ------ roofit/batchcompute/src/RooBatchCompute.cu | 15 ----- roofit/batchcompute/src/RooBatchCompute.cxx | 11 --- 5 files changed, 127 deletions(-) diff --git a/roofit/batchcompute/res/RooBatchCompute.h b/roofit/batchcompute/res/RooBatchCompute.h index bd7a7301b1d7a..d520efc5e1bbd 100644 --- a/roofit/batchcompute/res/RooBatchCompute.h +++ b/roofit/batchcompute/res/RooBatchCompute.h @@ -37,7 +37,6 @@ namespace RooBatchCompute { namespace CudaInterface { -class CudaEvent; class CudaStream; } // namespace CudaInterface @@ -180,13 +179,8 @@ class RooBatchComputeInterface { virtual std::unique_ptr createBufferManager() const = 0; - virtual CudaInterface::CudaEvent *newCudaEvent(bool forTiming) const = 0; virtual CudaInterface::CudaStream *newCudaStream() const = 0; - virtual void deleteCudaEvent(CudaInterface::CudaEvent *) const = 0; virtual void deleteCudaStream(CudaInterface::CudaStream *) const = 0; - virtual void cudaEventRecord(CudaInterface::CudaEvent *, CudaInterface::CudaStream *) const = 0; - virtual void cudaStreamWaitForEvent(CudaInterface::CudaStream *, CudaInterface::CudaEvent *) const = 0; - virtual bool cudaStreamIsActive(CudaInterface::CudaStream *) const = 0; }; /** diff --git a/roofit/batchcompute/src/CudaInterface.cu b/roofit/batchcompute/src/CudaInterface.cu index 34a8c1c042643..1d402ce4a6d6c 100644 --- a/roofit/batchcompute/src/CudaInterface.cu +++ b/roofit/batchcompute/src/CudaInterface.cu @@ -55,29 +55,6 @@ void Deleter::operator()(void *ptr) ptr = nullptr; } -/** - * Creates a new CUDA event. - * - * @param[in] forTiming Set to true if the event is intended for timing purposes. - * If `false`, the `cudaEventDisableTiming` is passed to CUDA. - * @return CudaEvent object representing the new event. - */ -CudaEvent::CudaEvent(bool forTiming) -{ - auto event = new cudaEvent_t; - ERRCHECK(cudaEventCreateWithFlags(event, forTiming ? 0 : cudaEventDisableTiming)); - _ptr.reset(event); -} - -template <> -void Deleter::operator()(void *ptr) -{ - auto event = reinterpret_cast(ptr); - ERRCHECK(cudaEventDestroy(*event)); - delete event; - ptr = nullptr; -} - template <> void Deleter::operator()(void *ptr) { @@ -87,17 +64,6 @@ void Deleter::operator()(void *ptr) ptr = nullptr; } -/** - * Records a CUDA event. - * - * @param[in] event CudaEvent object representing the event to be recorded. - * @param[in] stream CudaStream in which to record the event. - */ -void cudaEventRecord(CudaEvent &event, CudaStream &stream) -{ - ERRCHECK(::cudaEventRecord(event, stream)); -} - /** * Creates a new CUDA stream. * @@ -110,46 +76,6 @@ CudaStream::CudaStream() _ptr.reset(stream); } -/** - * Checks if a CUDA stream is currently active. - * - * @return True if the stream is active, false otherwise. - */ -bool CudaStream::isActive() -{ - cudaError_t err = cudaStreamQuery(*this); - if (err == cudaErrorNotReady) - return true; - else if (err == cudaSuccess) - return false; - ERRCHECK(err); - return false; -} - -/** - * Makes a CUDA stream wait for a CUDA event. - * - * @param[in] event CudaEvent object representing the event to wait for. - */ -void CudaStream::waitForEvent(CudaEvent &event) -{ - ERRCHECK(::cudaStreamWaitEvent(*this, event, 0)); -} - -/** - * Calculates the elapsed time between two CUDA events. - * - * @param[in] begin CudaEvent representing the start event. - * @param[in] end CudaEvent representing the end event. - * @return Elapsed time in milliseconds. - */ -float cudaEventElapsedTime(CudaEvent &begin, CudaEvent &end) -{ - float ret; - ERRCHECK(::cudaEventElapsedTime(&ret, begin, end)); - return ret; -} - /// \cond ROOFIT_INTERNAL void copyHostToDeviceImpl(const void *src, void *dest, size_t nBytes, CudaStream *stream) diff --git a/roofit/batchcompute/src/CudaInterface.h b/roofit/batchcompute/src/CudaInterface.h index cf4a688210988..df969d96589a4 100644 --- a/roofit/batchcompute/src/CudaInterface.h +++ b/roofit/batchcompute/src/CudaInterface.h @@ -37,21 +37,6 @@ struct Deleter { /// \endcond -/* - * Wrapper around cudaEvent_t. - */ -class CudaEvent { -public: - CudaEvent(bool forTiming); - -// When compiling with NVCC, we allow setting and getting the actual CUDA objects from the wrapper. -#ifdef __CUDACC__ - inline operator cudaEvent_t() { return *reinterpret_cast(_ptr.get()); } -#endif -private: - std::unique_ptr> _ptr; -}; - /* * Wrapper around cudaStream_t. */ @@ -59,9 +44,6 @@ class CudaStream { public: CudaStream(); - bool isActive(); - void waitForEvent(CudaEvent &); - // When compiling with NVCC, we allow setting and getting the actual CUDA objects from the wrapper. #ifdef __CUDACC__ inline cudaStream_t *get() { return reinterpret_cast(_ptr); } @@ -71,9 +53,6 @@ class CudaStream { std::unique_ptr> _ptr; }; -void cudaEventRecord(CudaEvent &, CudaStream &); -float cudaEventElapsedTime(CudaEvent &, CudaEvent &); - /// \cond ROOFIT_INTERNAL void copyHostToDeviceImpl(const void *src, void *dest, std::size_t n, CudaStream * = nullptr); void copyDeviceToHostImpl(const void *src, void *dest, std::size_t n, CudaStream * = nullptr); diff --git a/roofit/batchcompute/src/RooBatchCompute.cu b/roofit/batchcompute/src/RooBatchCompute.cu index 5397589ad5511..ce284ba2f0c69 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cu +++ b/roofit/batchcompute/src/RooBatchCompute.cu @@ -155,24 +155,9 @@ public: std::unique_ptr createBufferManager() const override; - CudaInterface::CudaEvent *newCudaEvent(bool forTiming) const override - { - return new CudaInterface::CudaEvent{forTiming}; - } CudaInterface::CudaStream *newCudaStream() const override { return new CudaInterface::CudaStream{}; } - void deleteCudaEvent(CudaInterface::CudaEvent *event) const override { delete event; } void deleteCudaStream(CudaInterface::CudaStream *stream) const override { delete stream; } - void cudaEventRecord(CudaInterface::CudaEvent *event, CudaInterface::CudaStream *stream) const override - { - CudaInterface::cudaEventRecord(*event, *stream); - } - void cudaStreamWaitForEvent(CudaInterface::CudaStream *stream, CudaInterface::CudaEvent *event) const override - { - stream->waitForEvent(*event); - } - bool cudaStreamIsActive(CudaInterface::CudaStream *stream) const override { return stream->isActive(); } - private: const std::vector _computeFunctions; diff --git a/roofit/batchcompute/src/RooBatchCompute.cxx b/roofit/batchcompute/src/RooBatchCompute.cxx index d98e0c828d6b5..9971ba81b4c05 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cxx +++ b/roofit/batchcompute/src/RooBatchCompute.cxx @@ -104,19 +104,8 @@ class RooBatchComputeClass : public RooBatchComputeInterface { std::unique_ptr createBufferManager() const override; - CudaInterface::CudaEvent *newCudaEvent(bool) const override { throw std::bad_function_call(); } CudaInterface::CudaStream *newCudaStream() const override { throw std::bad_function_call(); } - void deleteCudaEvent(CudaInterface::CudaEvent *) const override { throw std::bad_function_call(); } void deleteCudaStream(CudaInterface::CudaStream *) const override { throw std::bad_function_call(); } - void cudaEventRecord(CudaInterface::CudaEvent *, CudaInterface::CudaStream *) const override - { - throw std::bad_function_call(); - } - void cudaStreamWaitForEvent(CudaInterface::CudaStream *, CudaInterface::CudaEvent *) const override - { - throw std::bad_function_call(); - } - bool cudaStreamIsActive(CudaInterface::CudaStream *) const override { throw std::bad_function_call(); } private: #ifdef ROOBATCHCOMPUTE_USE_IMT From 330d5b31ffc4ca430ae769c32dba5844f2ebadd8 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 4 Sep 2026 16:31:01 +0000 Subject: [PATCH 3/8] [RF] Don't mutate the extra args in the Bernstein compute function MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Bernstein compute function applied the binomial coefficients to the extra args in-place and restored them at the end of the function. In the CUDA case, the extra args live in global device memory shared by all threads, so this read-modify-write pattern was a cross-block data race: a thread block that finished early could restore the coefficients while other blocks were still using them. Fold the binomial coefficient recurrence into the evaluation loops instead, so the extra args are never written to. The results are bit-for-bit identical because the same products are computed in the same order. This also makes Bernstein safe for skipping the device-to-host copy-back of the extra args after kernel launches. 🤖 Done with the help of AI --- roofit/batchcompute/src/ComputeFunctions.cxx | 26 +++++++++----------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/roofit/batchcompute/src/ComputeFunctions.cxx b/roofit/batchcompute/src/ComputeFunctions.cxx index 3f67a736c7eb2..f6e8508f3a899 100644 --- a/roofit/batchcompute/src/ComputeFunctions.cxx +++ b/roofit/batchcompute/src/ComputeFunctions.cxx @@ -100,12 +100,10 @@ __rooglobal__ void computeBernstein(Batches &batches) const double xmax = batches.extra[nCoef + 1]; Batch xData = batches.args[0]; - // apply binomial coefficient in-place so we don't have to allocate new memory - double binomial = 1.0; - for (int k = 0; k < nCoef; k++) { - batches.extra[k] = batches.extra[k] * binomial; - binomial = (binomial * (degree - k)) / (k + 1); - } + // The binomial coefficients are applied on the fly in the evaluation loops + // below. Note for the CUDA case: the coefficients must not be applied to + // batches.extra in-place, because the extra arguments live in global device + // memory that is shared by all threads. if (STEP == 1) { double X[bufferSize]; @@ -134,9 +132,12 @@ __rooglobal__ void computeBernstein(Batches &batches) for (size_t i = BEGIN; i < batches.nEvents; i += STEP) _1_X[i] = 1 / _1_X[i]; + double binomial = 1.0; for (int k = 0; k < nCoef; k++) { + const double coef = batches.extra[k] * binomial; + binomial = (binomial * (degree - k)) / (k + 1); for (size_t i = BEGIN; i < batches.nEvents; i += STEP) { - batches.output[i] += batches.extra[k] * powX[i] * pow_1_X[i]; + batches.output[i] += coef * powX[i] * pow_1_X[i]; // calculating next power for x and 1-x powX[i] *= X[i]; @@ -152,20 +153,15 @@ __rooglobal__ void computeBernstein(Batches &batches) for (int k = 1; k <= degree; k++) pow_1_X *= 1 - X; const double _1_X = 1 / (1 - X); + double binomial = 1.0; for (int k = 0; k < nCoef; k++) { - batches.output[i] += batches.extra[k] * powX * pow_1_X; + batches.output[i] += batches.extra[k] * binomial * powX * pow_1_X; + binomial = (binomial * (degree - k)) / (k + 1); powX *= X; pow_1_X *= _1_X; } } } - - // reset extraArgs values so we don't mutate the Batches object - binomial = 1.0; - for (int k = 0; k < nCoef; k++) { - batches.extra[k] = batches.extra[k] / binomial; - binomial = (binomial * (degree - k)) / (k + 1); - } } __rooglobal__ void computeBifurGauss(Batches &batches) From 3c65fe0525ed09b882f52281c606ad4cdef7939a Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 4 Sep 2026 16:31:01 +0000 Subject: [PATCH 4/8] [RF] Avoid allocations and synchronization in CUDA evaluation hot loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Overhaul the memory and transfer handling of the RooBatchCompute CUDA backend: * The typed CudaInterface copy helpers silently dropped their stream argument, so every copy was a synchronous cudaMemcpy on the legacy default stream that serialized with the compute stream. The stream is now forwarded, and all callers were audited for the new asynchronous semantics: PinnedBufferContainer::hostReadPtr() and the reduction results now synchronize the stream explicitly, and the buffer copies in the Evaluator::setInput() path keep their synchronous behavior via the nullptr-stream default. * Every compute() call allocated its staging area with cudaMalloc and freed it with cudaFree, and staged through a pageable host vector. The reductions likewise allocated their output buffers per call. All of this is replaced by StreamScratch, a per-stream ring of pinned host + device staging slots guarded by CUDA events, so the hot loop performs no CUDA allocations at all and staging uploads are asynchronous. * The extra args were unconditionally copied back after every kernel launch, which forced a round trip per node. Only the NormalizedPdf computer actually mutates its extra args (evaluation error counters), so only that case copies back and synchronizes. * Evaluator::run() now synchronizes the stream before returning, which makes the eager buffer recycling at the start of the next evaluation safe and restores the completion guarantee that the old polling scheduler provided. Also initialize Batches::extra and relax the debug assertion in reduceNLL, which wrongly required scalar probability spans to be device pointers (an observable-independent pdf legitimately provides a host scalar). Benchmarks on an RTX A4500 (wall time, identical minNll everywhere, CUDA baseline before the single-stream refactoring in parentheses): 1 channel, 1M events: 0.13 s (was 0.22 s) 1 channel, 10M events: 1.21 s (was 1.31 s) 8-channel simultaneous, 125k/chan: 1.30 s (was 1.63 s) 8-channel simultaneous, 12.5k/chan: 0.56 s (was 0.90 s) test-stressroofit-cuda and test-stressroostats-cuda pass. 🤖 Done with the help of AI --- roofit/batchcompute/res/RooBatchCompute.h | 2 + roofit/batchcompute/src/Batches.h | 2 +- roofit/batchcompute/src/CudaInterface.cu | 14 -- roofit/batchcompute/src/CudaInterface.h | 35 +++- roofit/batchcompute/src/RooBatchCompute.cu | 170 ++++++++++++++++---- roofit/batchcompute/src/RooBatchCompute.cxx | 1 + roofit/roofitcore/src/RooFit/Evaluator.cxx | 6 + 7 files changed, 181 insertions(+), 49 deletions(-) diff --git a/roofit/batchcompute/res/RooBatchCompute.h b/roofit/batchcompute/res/RooBatchCompute.h index d520efc5e1bbd..a0bbbf5dec73b 100644 --- a/roofit/batchcompute/res/RooBatchCompute.h +++ b/roofit/batchcompute/res/RooBatchCompute.h @@ -181,6 +181,8 @@ class RooBatchComputeInterface { virtual CudaInterface::CudaStream *newCudaStream() const = 0; virtual void deleteCudaStream(CudaInterface::CudaStream *) const = 0; + /// Wait until all work that was enqueued on the stream has completed. + virtual void synchronizeCudaStream(CudaInterface::CudaStream *) const = 0; }; /** diff --git a/roofit/batchcompute/src/Batches.h b/roofit/batchcompute/src/Batches.h index c63d689f821ba..c9189cfe7a14c 100644 --- a/roofit/batchcompute/src/Batches.h +++ b/roofit/batchcompute/src/Batches.h @@ -42,7 +42,7 @@ class Batch { class Batches { public: Batch *args = nullptr; - double *extra; + double *extra = nullptr; std::size_t nEvents = 0; std::size_t nBatches = 0; std::size_t nExtra = 0; diff --git a/roofit/batchcompute/src/CudaInterface.cu b/roofit/batchcompute/src/CudaInterface.cu index 1d402ce4a6d6c..13e56d1b3b75d 100644 --- a/roofit/batchcompute/src/CudaInterface.cu +++ b/roofit/batchcompute/src/CudaInterface.cu @@ -12,20 +12,6 @@ #include "CudaInterface.h" -#include -#include -#include - -#define ERRCHECK(err) __checkCudaErrors((err), __func__, __FILE__, __LINE__) -inline static void __checkCudaErrors(cudaError_t error, std::string func, std::string file, int line) -{ - if (error != cudaSuccess) { - std::stringstream errMsg; - errMsg << func << "(), " << file + ":" << std::to_string(line) << " : " << cudaGetErrorString(error); - throw std::runtime_error(errMsg.str()); - } -} - namespace RooBatchCompute { namespace CudaInterface { diff --git a/roofit/batchcompute/src/CudaInterface.h b/roofit/batchcompute/src/CudaInterface.h index df969d96589a4..ef8f241e522dc 100644 --- a/roofit/batchcompute/src/CudaInterface.h +++ b/roofit/batchcompute/src/CudaInterface.h @@ -16,6 +16,29 @@ #include #include +#ifdef __CUDACC__ +#include +#include +#include + +#define ERRCHECK(err) RooBatchCompute::CudaInterface::checkCudaErrors((err), __func__, __FILE__, __LINE__) + +namespace RooBatchCompute { +namespace CudaInterface { + +inline void checkCudaErrors(cudaError_t error, std::string const &func, std::string const &file, int line) +{ + if (error != cudaSuccess) { + std::stringstream errMsg; + errMsg << func << "(), " << file << ":" << std::to_string(line) << " : " << cudaGetErrorString(error); + throw std::runtime_error(errMsg.str()); + } +} + +} // namespace CudaInterface +} // namespace RooBatchCompute +#endif // __CUDACC__ + namespace RooBatchCompute { /* @@ -68,9 +91,9 @@ void copyDeviceToDeviceImpl(const void *src, void *dest, std::size_t n, CudaStre * @param[in] stream CudaStream for asynchronous memory transfer (optional). */ template -void copyHostToDevice(const T *src, T *dest, std::size_t n, CudaStream * = nullptr) +void copyHostToDevice(const T *src, T *dest, std::size_t n, CudaStream *stream = nullptr) { - copyHostToDeviceImpl(src, dest, sizeof(T) * n); + copyHostToDeviceImpl(src, dest, sizeof(T) * n, stream); } /** @@ -82,9 +105,9 @@ void copyHostToDevice(const T *src, T *dest, std::size_t n, CudaStream * = nullp * @param[in] stream CudaStream for asynchronous memory transfer (optional). */ template -void copyDeviceToHost(const T *src, T *dest, std::size_t n, CudaStream * = nullptr) +void copyDeviceToHost(const T *src, T *dest, std::size_t n, CudaStream *stream = nullptr) { - copyDeviceToHostImpl(src, dest, sizeof(T) * n); + copyDeviceToHostImpl(src, dest, sizeof(T) * n, stream); } /** @@ -96,9 +119,9 @@ void copyDeviceToHost(const T *src, T *dest, std::size_t n, CudaStream * = nullp * @param[in] stream CudaStream for asynchronous memory transfer (optional). */ template -void copyDeviceToDevice(const T *src, T *dest, std::size_t n, CudaStream * = nullptr) +void copyDeviceToDevice(const T *src, T *dest, std::size_t n, CudaStream *stream = nullptr) { - copyDeviceToDeviceImpl(src, dest, sizeof(T) * n); + copyDeviceToDeviceImpl(src, dest, sizeof(T) * n, stream); } /// \cond ROOFIT_INTERNAL diff --git a/roofit/batchcompute/src/RooBatchCompute.cu b/roofit/batchcompute/src/RooBatchCompute.cu index ce284ba2f0c69..0b8a060747f2b 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cu +++ b/roofit/batchcompute/src/RooBatchCompute.cu @@ -23,10 +23,12 @@ This file contains the code for cuda computations using the RooBatchCompute libr #include "CudaInterface.h" #include +#include #include #include #include #include +#include #include namespace RooBatchCompute { @@ -81,6 +83,84 @@ int getGridSize(std::size_t n) return std::min(int(std::ceil(double(n) / blockSize)), maxGridSize); } +/// Scratch memory attached to a CUDA stream, used for staging small +/// per-kernel-launch data like the Batches descriptor and reduction results. +/// +/// The slots form a ring: acquire() returns the next slot, waiting for the +/// completion of the work that was previously enqueued from that slot if it +/// is still in flight (which is rare, given the depth of the ring). Each slot +/// pairs a pinned host buffer with a device buffer of the same capacity, so +/// staging copies are truly asynchronous and no cudaMalloc()/cudaFree() calls +/// happen in the evaluation hot loop. +/// +/// Like the rest of the RooBatchCompute library, this class is not +/// thread-safe: RooFit evaluates on a single thread per process. +class StreamScratch { +public: + struct Slot { + char *host = nullptr; // pinned host memory + char *device = nullptr; + std::size_t capacity = 0; + cudaEvent_t event = nullptr; // recorded after the last enqueued use + bool inFlight = false; + }; + + StreamScratch() = default; + StreamScratch(StreamScratch const &) = delete; + StreamScratch &operator=(StreamScratch const &) = delete; + + Slot &acquire(std::size_t n) + { + Slot &slot = _slots[_next]; + _next = (_next + 1) % _slots.size(); + if (slot.inFlight) { + ERRCHECK(cudaEventSynchronize(slot.event)); + slot.inFlight = false; + } + if (slot.capacity < n) { + if (slot.host) + ERRCHECK(cudaFreeHost(slot.host)); + if (slot.device) + ERRCHECK(cudaFree(slot.device)); + const std::size_t newCapacity = std::max(n, 1024); + ERRCHECK(cudaMallocHost(reinterpret_cast(&slot.host), newCapacity)); + ERRCHECK(cudaMalloc(reinterpret_cast(&slot.device), newCapacity)); + slot.capacity = newCapacity; + } + if (slot.event == nullptr) { + ERRCHECK(cudaEventCreateWithFlags(&slot.event, cudaEventDisableTiming)); + } + return slot; + } + + /// Mark the last enqueued use of the slot on the stream. The slot will not + /// be handed out again before that work has completed. + void release(Slot &slot, cudaStream_t stream) + { + ERRCHECK(cudaEventRecord(slot.event, stream)); + slot.inFlight = true; + } + + ~StreamScratch() + { + // Don't use ERRCHECK here: throwing from a destructor would terminate. + for (Slot &slot : _slots) { + if (slot.inFlight) + cudaEventSynchronize(slot.event); + if (slot.event) + cudaEventDestroy(slot.event); + if (slot.host) + cudaFreeHost(slot.host); + if (slot.device) + cudaFree(slot.device); + } + } + +private: + std::array _slots; + std::size_t _next = 0; +}; + } // namespace std::vector getFunctions(); @@ -115,14 +195,17 @@ public: const std::size_t memSize = sizeof(Batches) + vars.size() * sizeof(Batch) + vars.size() * sizeof(double) + extraArgs.size() * sizeof(double); - std::vector hostMem(memSize); - auto batches = reinterpret_cast(hostMem.data()); + cudaStream_t stream = *cfg.cudaStream(); + StreamScratch::Slot &slot = scratch(cfg.cudaStream()).acquire(memSize); + + // The staging area has the same layout in the pinned host buffer and in + // the device buffer, so it can be uploaded with a single copy. + auto batches = reinterpret_cast(slot.host); auto arrays = reinterpret_cast(batches + 1); auto scalarBuffer = reinterpret_cast(arrays + vars.size()); auto extraArgsHost = reinterpret_cast(scalarBuffer + vars.size()); - DeviceArray deviceMem(memSize); - auto batchesDevice = reinterpret_cast(deviceMem.data()); + auto batchesDevice = reinterpret_cast(slot.device); auto arraysDevice = reinterpret_cast(batchesDevice + 1); auto scalarBufferDevice = reinterpret_cast(arraysDevice + vars.size()); auto extraArgsDevice = reinterpret_cast(scalarBufferDevice + vars.size()); @@ -136,17 +219,22 @@ public: batches->extra = extraArgsDevice; } - copyHostToDevice(hostMem.data(), deviceMem.data(), hostMem.size(), cfg.cudaStream()); + copyHostToDevice(slot.host, slot.device, memSize, cfg.cudaStream()); const int gridSize = getGridSize(nEvents); - _computeFunctions[computer]<<>>(*batchesDevice); - - // The compute might have modified the mutable extra args, so we need to - // copy them back. This can be optimized if necessary in the future by - // flagging if the extra args were actually changed. - if (!extraArgs.empty()) { + _computeFunctions[computer]<<>>(*batchesDevice); + + // Only the NormalizedPdf computer mutates its extra args: it uses them + // as output parameters for the evaluation error counts. Only then the + // extra args need to be copied back, and the stream needs to be + // synchronized because the caller inspects the counters right after + // this function returns. + if (computer == NormalizedPdf && !extraArgs.empty()) { copyDeviceToHost(extraArgsDevice, extraArgs.data(), extraArgs.size(), cfg.cudaStream()); + ERRCHECK(cudaStreamSynchronize(stream)); } + + scratch(cfg.cudaStream()).release(slot, stream); } /// Return the sum of an input array double reduceSum(RooBatchCompute::Config const &cfg, InputArr input, size_t n) override; @@ -156,10 +244,21 @@ public: std::unique_ptr createBufferManager() const override; CudaInterface::CudaStream *newCudaStream() const override { return new CudaInterface::CudaStream{}; } - void deleteCudaStream(CudaInterface::CudaStream *stream) const override { delete stream; } + void deleteCudaStream(CudaInterface::CudaStream *stream) const override + { + _scratchMap.erase(stream); + delete stream; + } + void synchronizeCudaStream(CudaInterface::CudaStream *stream) const override + { + ERRCHECK(::cudaStreamSynchronize(*stream)); + } private: + StreamScratch &scratch(CudaInterface::CudaStream *stream) { return _scratchMap[stream]; } + const std::vector _computeFunctions; + mutable std::unordered_map _scratchMap; }; // End class RooBatchComputeClass @@ -264,13 +363,17 @@ double RooBatchComputeClass::reduceSum(RooBatchCompute::Config const &cfg, Input return 0.0; const int gridSize = getGridSize(n); cudaStream_t stream = *cfg.cudaStream(); - CudaInterface::DeviceArray devOut(2 * gridSize); + StreamScratch::Slot &slot = scratch(cfg.cudaStream()).acquire(2 * gridSize * sizeof(double)); + auto devOut = reinterpret_cast(slot.device); + auto hostOut = reinterpret_cast(slot.host); constexpr int shMemSize = 2 * blockSize * sizeof(double); - kahanSum<<>>(input, nullptr, n, devOut.data(), 0); - kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut.data(), devOut.data() + gridSize, gridSize, devOut.data(), 0); - double tmp = 0.0; - CudaInterface::copyDeviceToHost(devOut.data(), &tmp, 1, cfg.cudaStream()); - return tmp; + kahanSum<<>>(input, nullptr, n, devOut, 0); + kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut, devOut + gridSize, gridSize, devOut, 0); + CudaInterface::copyDeviceToHost(devOut, hostOut, 1, cfg.cudaStream()); + ERRCHECK(cudaStreamSynchronize(stream)); + const double result = hostOut[0]; + scratch(cfg.cudaStream()).release(slot, stream); + return result; } ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &cfg, std::span probas, @@ -281,31 +384,37 @@ ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &c return out; } const int gridSize = getGridSize(weights.size()); - CudaInterface::DeviceArray devOut(2 * gridSize); cudaStream_t stream = *cfg.cudaStream(); + StreamScratch::Slot &slot = scratch(cfg.cudaStream()).acquire(2 * gridSize * sizeof(double)); + auto devOut = reinterpret_cast(slot.device); + auto hostOut = reinterpret_cast(slot.host); constexpr int shMemSize = 2 * blockSize * sizeof(double); #ifndef NDEBUG for (auto span : {probas, weights, offsetProbas}) { + // Scalar spans can point to host memory (e.g. the scalar buffer of an + // observable-independent pdf), so only spans with more than one element + // are required to be on the device. cudaPointerAttributes attr; - assert(span.size() == 0 || span.data() == nullptr || + assert(span.size() <= 1 || span.data() == nullptr || (cudaPointerGetAttributes(&attr, span.data()) == cudaSuccess && attr.type == cudaMemoryTypeDevice)); } #endif nllSumKernel<<>>( probas.data(), weights.data(), offsetProbas.empty() ? nullptr : offsetProbas.data(), probas.size(), - probas.size() == 1 ? probas[0] : 0.0, weights.size(), devOut.data()); + probas.size() == 1 ? probas[0] : 0.0, weights.size(), devOut); - kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut.data(), devOut.data() + gridSize, gridSize, devOut.data(), 0); + kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut, devOut + gridSize, gridSize, devOut, 0); - double tmpSum = 0.0; - double tmpCarry = 0.0; - CudaInterface::copyDeviceToHost(devOut.data(), &tmpSum, 1, cfg.cudaStream()); - CudaInterface::copyDeviceToHost(devOut.data() + 1, &tmpCarry, 1, cfg.cudaStream()); + // The sum and its Kahan carry are adjacent in the output buffer, so they + // can be read back with a single copy. + CudaInterface::copyDeviceToHost(devOut, hostOut, 2, cfg.cudaStream()); + ERRCHECK(cudaStreamSynchronize(stream)); - out.nllSum = tmpSum; - out.nllSumCarry = tmpCarry; + out.nllSum = hostOut[0]; + out.nllSumCarry = hostOut[1]; + scratch(cfg.cudaStream()).release(slot, stream); return out; } @@ -408,6 +517,11 @@ public: if (_lastAccess == LastAccessType::GPU_WRITE) { CudaInterface::copyDeviceToHost(_gpuBuffer.deviceReadPtr(), const_cast(_arr.data()), size(), _cudaStream); + // The copy is asynchronous, and the caller reads the host memory + // right away, so the stream needs to be synchronized here. + if (_cudaStream) { + ERRCHECK(cudaStreamSynchronize(*_cudaStream)); + } } _lastAccess = LastAccessType::CPU_READ; diff --git a/roofit/batchcompute/src/RooBatchCompute.cxx b/roofit/batchcompute/src/RooBatchCompute.cxx index 9971ba81b4c05..4dcf6bbf1d1eb 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cxx +++ b/roofit/batchcompute/src/RooBatchCompute.cxx @@ -106,6 +106,7 @@ class RooBatchComputeClass : public RooBatchComputeInterface { CudaInterface::CudaStream *newCudaStream() const override { throw std::bad_function_call(); } void deleteCudaStream(CudaInterface::CudaStream *) const override { throw std::bad_function_call(); } + void synchronizeCudaStream(CudaInterface::CudaStream *) const override { throw std::bad_function_call(); } private: #ifdef ROOBATCHCOMPUTE_USE_IMT diff --git a/roofit/roofitcore/src/RooFit/Evaluator.cxx b/roofit/roofitcore/src/RooFit/Evaluator.cxx index 9abefb47055c8..26fd73ea74f98 100644 --- a/roofit/roofitcore/src/RooFit/Evaluator.cxx +++ b/roofit/roofitcore/src/RooFit/Evaluator.cxx @@ -496,6 +496,12 @@ std::span Evaluator::getValHeterogeneous() } } + // Ensure that all enqueued GPU work has completed when run() returns. For + // the usual likelihood evaluations this is 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. + RooBatchCompute::dispatchCUDA->synchronizeCudaStream(_cudaStream); + // return the final value return _evalContextCUDA.at(&_topNode); } From bca55467661b48c4d330978b2db0c3a1a00316a6 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 4 Sep 2026 16:45:05 +0000 Subject: [PATCH 5/8] [RF] Report evaluation errors from the CUDA NLL reduction like on CPU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CUDA implementation of RooBatchComputeInterface::reduceNLL() diverged from the CPU implementation in several user-visible ways: * Non-positive, NaN, and infinite probabilities were not counted, so RooNLLVarNew silently dropped all evaluation error logging in CUDA fits. * The accumulated "badness" of problematic events was not packed into the returned NaN with RooNaNPacker, so the error recovery in the minimizer (RecoverFromUndefinedRegions) did not work on the GPU: the minimizer only saw a plain NaN without the recovery information. * Zero-weight events were not skipped, so a zero-weight event with zero probability turned the NLL into NaN via 0 * inf, while the CPU implementation skips such events. The nllSumKernel now mirrors the CPU getLog() semantics per event and accumulates the badness and the three error counters with atomicAdd on the rare path. The statistics are stored next to the Kahan sum in the stream scratch slot, so they are read back in the same single copy, and the host packs the badness into the NaN like the CPU implementation. For pre-sm_60 devices, the canonical atomicCAS-based double atomicAdd fallback is provided in RooBatchComputeTypes.h. RooNaNPacker::unpackNaN and isNaNWithPayload are marked __roodevice__ so they can be used in kernels. Also fix the evaluation error counters of computeNormalizedPdf: the thresholds for reporting type-1 and type-2 errors were "> 1" and "> 2" instead of "> 0" (so single errors were never reported, also on CPU), and the counter accumulation in global device memory now uses atomicAdd to avoid losing counts to the read-modify-write race between threads. The new RooNaNPacker.CudaEvalErrorParity test verifies that CPU and CUDA backends produce the same packed badness payload and that zero-weight events in undefined pdf regions are skipped identically. 🤖 Done with the help of AI --- .../batchcompute/res/RooBatchComputeTypes.h | 17 ++++ roofit/batchcompute/res/RooNaNPacker.h | 4 +- roofit/batchcompute/src/ComputeFunctions.cxx | 18 ++-- roofit/batchcompute/src/RooBatchCompute.cu | 87 ++++++++++++++++--- roofit/roofitcore/test/testNaNPacker.cxx | 60 +++++++++++++ 5 files changed, 165 insertions(+), 21 deletions(-) diff --git a/roofit/batchcompute/res/RooBatchComputeTypes.h b/roofit/batchcompute/res/RooBatchComputeTypes.h index ecb4421c742f7..913d7a7c4a757 100644 --- a/roofit/batchcompute/res/RooBatchComputeTypes.h +++ b/roofit/batchcompute/res/RooBatchComputeTypes.h @@ -24,4 +24,21 @@ #define __rooglobal__ #endif // #indef __CUDACC__ +// Double-precision atomicAdd() is only provided by the CUDA runtime for +// compute capability 6.0 and higher. This is the canonical fallback +// implementation from the CUDA C++ Programming Guide for older devices. +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 600 +static __inline__ __device__ double atomicAdd(double *address, double val) +{ + unsigned long long int *address_as_ull = (unsigned long long int *)address; + unsigned long long int old = *address_as_ull; + unsigned long long int assumed; + do { + assumed = old; + old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed))); + } while (assumed != old); + return __longlong_as_double(old); +} +#endif + #endif diff --git a/roofit/batchcompute/res/RooNaNPacker.h b/roofit/batchcompute/res/RooNaNPacker.h index 1bc4586b1780f..f0cb0d6222e44 100644 --- a/roofit/batchcompute/res/RooNaNPacker.h +++ b/roofit/batchcompute/res/RooNaNPacker.h @@ -96,7 +96,7 @@ struct RooNaNPacker { bool isNaNWithPayload() const { return isNaNWithPayload(_payload); } /// Test if `val` has a float packed into its mantissa. - static bool isNaNWithPayload(double val) + __roodevice__ __roohost__ static bool isNaNWithPayload(double val) { uint64_t tmp; std::memcpy(&tmp, &val, sizeof(uint64_t)); @@ -120,7 +120,7 @@ struct RooNaNPacker { /// If `val` is NaN and a this NaN has been tagged as containing /// a payload, unpack the float from the mantissa. /// Return 0 otherwise. - static float unpackNaN(double val) + __roodevice__ __roohost__ static float unpackNaN(double val) { float tmp; std::memcpy(&tmp, &val, sizeof(float)); diff --git a/roofit/batchcompute/src/ComputeFunctions.cxx b/roofit/batchcompute/src/ComputeFunctions.cxx index f6e8508f3a899..f218bca6db459 100644 --- a/roofit/batchcompute/src/ComputeFunctions.cxx +++ b/roofit/batchcompute/src/ComputeFunctions.cxx @@ -660,12 +660,20 @@ __rooglobal__ void computeNormalizedPdf(Batches &batches) batches.output[i] = out; } + // The counters live in memory that is shared between all threads in the + // CUDA case, so they need to be accumulated atomically there. +#ifdef __CUDACC__ if (nEvalErrorsType0 > 0) - batches.extra[0] = batches.extra[0] + nEvalErrorsType0; - if (nEvalErrorsType1 > 1) - batches.extra[1] = batches.extra[1] + nEvalErrorsType1; - if (nEvalErrorsType2 > 2) - batches.extra[2] = batches.extra[2] + nEvalErrorsType2; + atomicAdd(&batches.extra[0], double(nEvalErrorsType0)); + if (nEvalErrorsType1 > 0) + atomicAdd(&batches.extra[1], double(nEvalErrorsType1)); + if (nEvalErrorsType2 > 0) + atomicAdd(&batches.extra[2], double(nEvalErrorsType2)); +#else + batches.extra[0] = batches.extra[0] + nEvalErrorsType0; + batches.extra[1] = batches.extra[1] + nEvalErrorsType1; + batches.extra[2] = batches.extra[2] + nEvalErrorsType2; +#endif } /* TMath::ASinH(x) needs to be replaced with ln( x + sqrt(x^2+1)) diff --git a/roofit/batchcompute/src/RooBatchCompute.cu b/roofit/batchcompute/src/RooBatchCompute.cu index 0b8a060747f2b..b5f0dc73ecff1 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cu +++ b/roofit/batchcompute/src/RooBatchCompute.cu @@ -19,6 +19,7 @@ This file contains the code for cuda computations using the RooBatchCompute libr **/ #include "RooBatchCompute.h" +#include "RooNaNPacker.h" #include "Batches.h" #include "CudaInterface.h" @@ -323,9 +324,15 @@ __global__ void kahanSum(const double *__restrict__ input, const double *__restr kahanSumReduction(shared, n, result, carry_index); } +/// Computes the negative log likelihood sum with the same semantics as the +/// CPU implementation of RooBatchComputeInterface::reduceNLL(): zero-weight +/// events are skipped, and evaluation problems are counted and accumulated +/// into a "badness" value that the host can pack into a NaN for the error +/// recovery in the minimizer. The `stats` output has the layout +/// [badness, nNonPositive, nNaN, nInfinite] and must be zero-initialized. __global__ void nllSumKernel(const double *__restrict__ probas, const double *__restrict__ weights, const double *__restrict__ offsetProbas, size_t nProbas, double scalarProba, - size_t nWeights, double *__restrict__ result) + size_t nWeights, double *__restrict__ result, double *__restrict__ stats) { int thIdx = threadIdx.x; int gthIdx = thIdx + blockIdx.x * blockSize; @@ -337,16 +344,51 @@ __global__ void nllSumKernel(const double *__restrict__ probas, const double *__ double sum = 0.0; double carry = 0.0; + double badness = 0.0; + unsigned int nNonPositive = 0; + unsigned int nNaN = 0; + unsigned int nInfinite = 0; for (int i = gthIdx; i < nWeights; i += nThreadsTotal) { - // Note: it does not make sense to use the nll option and provide at the - // same time external carries. - double val = -std::log(nProbas == 1 ? scalarProba : probas[i]); + const double weight = weights[i]; + // Zero-weight events don't contribute to the likelihood. Skipping them + // also avoids 0 * inf = NaN for zero probabilities. + if (weight == 0.0) { + continue; + } + const double proba = nProbas == 1 ? scalarProba : probas[i]; + double term; + if (proba <= 0.0) { + ++nNonPositive; + badness += -proba; + term = std::log(proba); + } else if (std::isnan(proba)) { + ++nNaN; + badness += RooNaNPacker::unpackNaN(proba); + term = proba; + } else { + if (std::isinf(proba)) { + ++nInfinite; + } + term = std::log(proba); + } if (offsetProbas) - val += std::log(offsetProbas[i]); - val = weights[i] * val; - kahanSumUpdate(sum, carry, val, 0.0); - } + term -= std::log(offsetProbas[i]); + term *= -weight; + kahanSumUpdate(sum, carry, term, 0.0); + } + + // Accumulate the evaluation error statistics over the whole grid. These + // atomics are on the rare path: they are only executed by threads that + // actually encountered problematic values. + if (badness != 0.0) + atomicAdd(&stats[0], badness); + if (nNonPositive != 0) + atomicAdd(&stats[1], double(nNonPositive)); + if (nNaN != 0) + atomicAdd(&stats[2], double(nNaN)); + if (nInfinite != 0) + atomicAdd(&stats[3], double(nInfinite)); shared[thIdx] = sum; shared[carry_index] = carry; @@ -385,7 +427,9 @@ ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &c } const int gridSize = getGridSize(weights.size()); cudaStream_t stream = *cfg.cudaStream(); - StreamScratch::Slot &slot = scratch(cfg.cudaStream()).acquire(2 * gridSize * sizeof(double)); + // Layout of the scratch buffer: [sum, carry, badness, nNonPositive, nNaN, + // nInfinite, partial sums (gridSize), partial carries (gridSize)]. + StreamScratch::Slot &slot = scratch(cfg.cudaStream()).acquire((6 + 2 * gridSize) * sizeof(double)); auto devOut = reinterpret_cast(slot.device); auto hostOut = reinterpret_cast(slot.host); constexpr int shMemSize = 2 * blockSize * sizeof(double); @@ -401,19 +445,34 @@ ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &c } #endif + // Zero-initialize the evaluation error statistics for the atomic updates. + ERRCHECK(cudaMemsetAsync(devOut + 2, 0, 4 * sizeof(double), stream)); + nllSumKernel<<>>( probas.data(), weights.data(), offsetProbas.empty() ? nullptr : offsetProbas.data(), probas.size(), - probas.size() == 1 ? probas[0] : 0.0, weights.size(), devOut); + probas.size() == 1 ? probas[0] : 0.0, weights.size(), devOut + 6, devOut + 2); - kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut, devOut + gridSize, gridSize, devOut, 0); + kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut + 6, devOut + 6 + gridSize, gridSize, devOut, 0); - // The sum and its Kahan carry are adjacent in the output buffer, so they - // can be read back with a single copy. - CudaInterface::copyDeviceToHost(devOut, hostOut, 2, cfg.cudaStream()); + // The sum, its Kahan carry, and the evaluation error statistics are + // adjacent in the output buffer, so they can be read back in a single copy. + CudaInterface::copyDeviceToHost(devOut, hostOut, 6, cfg.cudaStream()); ERRCHECK(cudaStreamSynchronize(stream)); out.nllSum = hostOut[0]; out.nllSumCarry = hostOut[1]; + out.nNonPositiveValues = hostOut[3]; + out.nNaNValues = hostOut[4]; + out.nInfiniteValues = hostOut[5]; + + if (hostOut[2] != 0.0) { + // Some events had evaluation errors: return the accumulated "badness" + // of the errors packed into a NaN, like the CPU implementation, so the + // minimizer can use it to recover. + out.nllSum = RooNaNPacker::packFloatIntoNaN(hostOut[2]); + out.nllSumCarry = 0.0; + } + scratch(cfg.cudaStream()).release(slot, stream); return out; } diff --git a/roofit/roofitcore/test/testNaNPacker.cxx b/roofit/roofitcore/test/testNaNPacker.cxx index c1f307ab60277..fb023c3c888e0 100644 --- a/roofit/roofitcore/test/testNaNPacker.cxx +++ b/roofit/roofitcore/test/testNaNPacker.cxx @@ -261,6 +261,66 @@ INSTANTIATE_TEST_SUITE_P(RooNaNPacker, TestForDifferentBackends, testing::Values #undef BATCH_MODE_VALS +#ifdef ROOFIT_CUDA +/// Verify that the CUDA implementation of the NLL reduction reports +/// evaluation errors exactly like the CPU implementation: the accumulated +/// "badness" of non-positive pdf values is packed into the returned NaN so +/// the minimizer can recover, and zero-weight events are skipped. +TEST(RooNaNPacker, CudaEvalErrorParity) +{ + using namespace RooFit; + + RooRealVar x("x", "x", -10, 10); + RooRealVar a1("a1", "a1", 12., -10., 20.); + RooRealVar a2("a2", "a2", 1.1, -10., 20.); + RooGenericPdf pdf("pdf", "a1 + x + a2*x*x", {x, a1, a2}); + std::unique_ptr data{pdf.generate(x, 10000)}; + + // We provoke a lot of evaluation errors in this test: don't log them. + const auto prevErrorLoggingMode = RooAbsReal::evalErrorLoggingMode(); + RooAbsReal::setEvalErrorLoggingMode(RooAbsReal::Ignore); + + // Move the parameters into a region where the pdf is negative for many + // events: both backends must return a NaN with the same packed "badness". + a1.setVal(-9.); + a2.setVal(-1.); + std::unique_ptr nllCpu{pdf.createNLL(*data, EvalBackend::Cpu())}; + const double valCpu = nllCpu->getVal(); + std::unique_ptr nllCuda{pdf.createNLL(*data, EvalBackend::Cuda())}; + const double valCuda = nllCuda->getVal(); + + EXPECT_TRUE(RooNaNPacker::isNaNWithPayload(valCpu)); + EXPECT_TRUE(RooNaNPacker::isNaNWithPayload(valCuda)); + // The accumulation order differs between the backends, so the float + // payloads only agree within a relative tolerance. + EXPECT_THAT(RooNaNPacker::unpackNaN(valCuda), RelativeNear(RooNaNPacker::unpackNaN(valCpu), 1e-5)); + + // Zero-weight events must be skipped, even where the pdf is not positive. + a1.setVal(-9.); + a2.setVal(1.1); // pdf negative around x=0, positive for large x + RooRealVar w("w", "w", 0, 1); + RooDataSet wdata("wdata", "wdata", RooArgSet(x, w), WeightVar("w")); + for (int i = 0; i < 1000; ++i) { + x.setVal(5.0 + 4.0 * i / 1000.0); // pdf positive here + wdata.add(RooArgSet(x), 1.0); + } + x.setVal(0.0); // pdf negative here, but the events are weightless + for (int i = 0; i < 10; ++i) { + wdata.add(RooArgSet(x), 0.0); + } + std::unique_ptr nllCpuW{pdf.createNLL(wdata, EvalBackend::Cpu())}; + const double valCpuW = nllCpuW->getVal(); + std::unique_ptr nllCudaW{pdf.createNLL(wdata, EvalBackend::Cuda())}; + const double valCudaW = nllCudaW->getVal(); + + EXPECT_TRUE(std::isfinite(valCpuW)); + EXPECT_TRUE(std::isfinite(valCudaW)); + EXPECT_THAT(valCudaW, RelativeNear(valCpuW, 1e-10)); + + RooAbsReal::setEvalErrorLoggingMode(prevErrorLoggingMode); +} +#endif // ROOFIT_CUDA + /// Make coefficients of RooAddPdf sum to more than 1. Fitter should recover from this. TEST(RooNaNPacker, FitAddPdf_DegenerateCoeff) { From 95960cc91aea3bd3d99946b3b69b2c29071056a7 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 4 Sep 2026 17:02:46 +0000 Subject: [PATCH 6/8] [RF] Cache the sums of event weights in RooNLLVarNew MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sums of the event weights and of the squared event weights only depend on the dataset, but they were recomputed with a full reduction in every evaluation of the RooNLLVarNew (and for the weight sum, in every chi2 evaluation). In CUDA fits, each of these reductions also implied kernel launches and a stream synchronization per likelihood component per minimizer evaluation. The sums are now cached in the RooNLLVarNew, keyed on a new input data generation counter in the EvalContext that the Evaluator bumps whenever new input arrays are loaded via setInput(). This covers the dataset swapping in RooEvaluatorWrapper::setData(), e.g. for toy studies. The counter values are globally unique in the process, so a cached sum can never be wrongly validated by the context of a different Evaluator. The two sums carry separate generation stamps, because the squared weight sum is only computed on demand when applyWeightSquared() is enabled during the fit. This speeds up both backends, since the CPU reduction over the weights was equally redundant. Wall times on an RTX A4500 (before -> after, identical minNll): 1 channel, 10M events: cpu 14.3 -> 13.0 s, cuda 1.21 -> 1.19 s 8-channel simultaneous, 125k/chan: cpu 3.35 -> 3.03 s, cuda 1.30 -> 1.22 s 8-channel simultaneous, 12.5k/chan: cpu 0.32 -> 0.29 s, cuda 0.56 -> 0.47 s Verified that NLL values after setData() are bit-identical to freshly created NLL objects on both backends, that SumW2Error fits agree between CPU and CUDA, and that the full RooFit test suite including the CUDA stress tests passes. 🤖 Done with the help of AI --- .../inc/RooFit/Detail/RooNLLVarNew.h | 5 ++++ roofit/roofitcore/inc/RooFit/EvalContext.h | 10 ++++++++ roofit/roofitcore/src/RooFit/Evaluator.cxx | 11 +++++++++ roofit/roofitcore/src/RooNLLVarNew.cxx | 23 +++++++++++++++---- 4 files changed, 45 insertions(+), 4 deletions(-) diff --git a/roofit/roofitcore/inc/RooFit/Detail/RooNLLVarNew.h b/roofit/roofitcore/inc/RooFit/Detail/RooNLLVarNew.h index c55222c0c35b2..98c31a66d1e96 100644 --- a/roofit/roofitcore/inc/RooFit/Detail/RooNLLVarNew.h +++ b/roofit/roofitcore/inc/RooFit/Detail/RooNLLVarNew.h @@ -95,6 +95,7 @@ class RooNLLVarNew : public RooAbsReal { private: double evaluate() const override { return _value; } void resetWeightVarNames(); + double sumOfWeights(RooFit::EvalContext &, std::span weights, bool squared) const; void finalizeResult(RooFit::EvalContext &, ROOT::Math::KahanSum result, double weightSum) const; void fillBinWidthsFromPdfBoundaries(RooAbsReal const &pdf, RooArgSet const &observables); void doEvalBinnedL(RooFit::EvalContext &, std::span preds, std::span weights) const; @@ -120,6 +121,10 @@ class RooNLLVarNew : public RooAbsReal { std::string _prefix; std::vector _binw; mutable ROOT::Math::KahanSum _offset{0.}; /// output() { return _currentOutput; } @@ -118,6 +127,7 @@ class EvalContext { friend class Evaluator; OffsetMode _offsetMode = OffsetMode::WithoutOffset; + std::size_t _inputGeneration = 1; std::span _currentOutput; std::vector> _ctx; bool _enableVectorBuffers = false; diff --git a/roofit/roofitcore/src/RooFit/Evaluator.cxx b/roofit/roofitcore/src/RooFit/Evaluator.cxx index 26fd73ea74f98..c3ac9b97d9c66 100644 --- a/roofit/roofitcore/src/RooFit/Evaluator.cxx +++ b/roofit/roofitcore/src/RooFit/Evaluator.cxx @@ -42,6 +42,7 @@ RooAbsPdf::fitTo() is called and gets destroyed when the fitting ends. #include "BatchModeDataHelpers.h" #include "RooFitImplHelpers.h" +#include #include #include #include @@ -256,6 +257,16 @@ void Evaluator::setInput(std::string const &name, std::span inputA _needToUpdateOutputSizes = true; + // Invalidate the caches that reducer nodes key on the input data, like the + // cached sum of event weights in RooNLLVarNew. The counter is global so + // that generation values can never alias between different Evaluators. + { + static std::atomic nextInputGeneration{1}; + const std::size_t gen = ++nextInputGeneration; + _evalContextCPU._inputGeneration = gen; + _evalContextCUDA._inputGeneration = gen; + } + NodeInfo &info = *found->second; info.fromArrayInput = true; diff --git a/roofit/roofitcore/src/RooNLLVarNew.cxx b/roofit/roofitcore/src/RooNLLVarNew.cxx index 46ca910a42d30..3b9221e5a9369 100644 --- a/roofit/roofitcore/src/RooNLLVarNew.cxx +++ b/roofit/roofitcore/src/RooNLLVarNew.cxx @@ -292,6 +292,22 @@ void RooNLLVarNew::doEvalBinnedL(RooFit::EvalContext &ctx, std::span weights, bool squared) const +{ + std::size_t &generation = squared ? _sumWeight2Gen : _sumWeightGen; + double &cache = squared ? _sumWeight2Cache : _sumWeightCache; + if (generation != ctx.inputGeneration()) { + cache = RooBatchCompute::reduceSum(ctx.config(this), weights.data(), weights.size()); + generation = ctx.inputGeneration(); + } + return cache; +} + void RooNLLVarNew::doEvalChi2(RooFit::EvalContext &ctx, std::span preds, std::span weights, std::span weightsSumW2) const { @@ -302,12 +318,11 @@ void RooNLLVarNew::doEvalChi2(RooFit::EvalContext &ctx, std::span return; } - auto config = ctx.config(this); std::span binVol = ctx.at(*_binVolumes); std::span errLo = _weightErrLo ? ctx.at(*_weightErrLo) : std::span{}; std::span errHi = _weightErrHi ? ctx.at(*_weightErrHi) : std::span{}; - const double sumWeight = RooBatchCompute::reduceSum(config, weights.data(), weights.size()); + const double sumWeight = sumOfWeights(ctx, weights, false); double normFactor = 1.0; switch (_funcMode) { @@ -370,10 +385,10 @@ void RooNLLVarNew::doEval(RooFit::EvalContext &ctx) const auto probas = ctx.at(_func); - double sumWeight = RooBatchCompute::reduceSum(config, weights.data(), weights.size()); + double sumWeight = sumOfWeights(ctx, weights, false); double sumWeight2 = 0.; if (_expectedEvents && _weightSquared) { - sumWeight2 = RooBatchCompute::reduceSum(config, weightsSumW2.data(), weightsSumW2.size()); + sumWeight2 = sumOfWeights(ctx, weightsSumW2, true); } auto nllOut = RooBatchCompute::reduceNLL(config, probas, _weightSquared ? weightsSumW2 : weights, From 761f5829200517776882ee1964a52b7a88ca9948 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 4 Sep 2026 17:33:19 +0000 Subject: [PATCH 7/8] [RF] Defer the readback of CUDA evaluation error counters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The NormalizedPdf compute function uses its extra args as output parameters for the evaluation error counters, and the caller inspected them right after the compute() call. In CUDA mode, this forced a full stream synchronization after every NormalizedPdf kernel launch: one pipeline stall per likelihood component per minimizer evaluation. The readback is now deferred: compute() enqueues an asynchronous copy of the counters into a persistent pinned slot of the per-stream scratch memory, and synchronizeCudaStream() forwards them to the caller's memory after the synchronization at the end of the evaluation of the computation graph. Correspondingly, RooNormalizedPdf keeps its counters in a member instead of on the stack, and defers the logging of the evaluation errors with the new EvalContext::deferAction() mechanism: the RooFit::Evaluator runs the deferred actions right after the final stream synchronization of each evaluation. The deferral stays within one evaluation, so the evaluation errors are still logged at the same point in the fit and with the same parameter value snapshots as before, and the number of logged errors is identical between the CPU and CUDA backends (verified: 1802 errors on both for a pdf that goes negative in part of its range). The CPU code path is unchanged apart from the counters living in the member. Wall times on an RTX A4500 (identical minNll): 1 channel, 1M events: cuda 0.13 -> 0.12 s 8-channel simultaneous, 125k/chan: cuda 1.22 -> 1.18 s 8-channel simultaneous, 12.5k/chan: cuda 0.47 -> 0.42 s The full RooFit test suite including the CUDA stress tests passes. 🤖 Done with the help of AI --- roofit/batchcompute/src/RooBatchCompute.cu | 79 +++++++++++++++++-- .../inc/RooFit/Detail/RooNormalizedPdf.h | 10 +++ roofit/roofitcore/inc/RooFit/EvalContext.h | 9 +++ roofit/roofitcore/src/RooFit/Evaluator.cxx | 25 +++++- roofit/roofitcore/src/RooNormalizedPdf.cxx | 26 ++++-- 5 files changed, 132 insertions(+), 17 deletions(-) diff --git a/roofit/batchcompute/src/RooBatchCompute.cu b/roofit/batchcompute/src/RooBatchCompute.cu index b5f0dc73ecff1..8615f8aaf6031 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cu +++ b/roofit/batchcompute/src/RooBatchCompute.cu @@ -26,6 +26,7 @@ This file contains the code for cuda computations using the RooBatchCompute libr #include #include #include +#include #include #include #include @@ -142,6 +143,50 @@ public: slot.inFlight = true; } + /// A persistent slot for a deferred device-to-host readback: an + /// asynchronous copy delivers device results (e.g. evaluation error + /// counters) into the pinned host buffer, and flushDeferred() forwards + /// them to the destination in the caller's memory once the stream was + /// synchronized. Slots stay valid from acquireDeferred() until the flush. + struct DeferredSlot { + char *host = nullptr; // pinned host memory + std::size_t capacity = 0; + double *dst = nullptr; + std::size_t nPending = 0; + }; + + DeferredSlot &acquireDeferred(std::size_t n) + { + if (_deferredCursor == _deferredSlots.size()) { + _deferredSlots.emplace_back(); + } + DeferredSlot &slot = _deferredSlots[_deferredCursor++]; + if (slot.capacity < n) { + // The slot is idle here: its previous use ended with the flush after + // a stream synchronization. + if (slot.host) + ERRCHECK(cudaFreeHost(slot.host)); + ERRCHECK(cudaMallocHost(reinterpret_cast(&slot.host), n)); + slot.capacity = n; + } + return slot; + } + + /// Copy the completed readbacks to their destinations. Must only be + /// called after the stream was synchronized. + void flushDeferred() + { + for (std::size_t i = 0; i < _deferredCursor; ++i) { + DeferredSlot &slot = _deferredSlots[i]; + if (slot.dst) { + std::memcpy(slot.dst, slot.host, slot.nPending * sizeof(double)); + slot.dst = nullptr; + slot.nPending = 0; + } + } + _deferredCursor = 0; + } + ~StreamScratch() { // Don't use ERRCHECK here: throwing from a destructor would terminate. @@ -155,11 +200,17 @@ public: if (slot.device) cudaFree(slot.device); } + for (DeferredSlot &slot : _deferredSlots) { + if (slot.host) + cudaFreeHost(slot.host); + } } private: std::array _slots; std::size_t _next = 0; + std::vector _deferredSlots; + std::size_t _deferredCursor = 0; }; } // namespace @@ -197,7 +248,8 @@ public: extraArgs.size() * sizeof(double); cudaStream_t stream = *cfg.cudaStream(); - StreamScratch::Slot &slot = scratch(cfg.cudaStream()).acquire(memSize); + StreamScratch &streamScratch = scratch(cfg.cudaStream()); + StreamScratch::Slot &slot = streamScratch.acquire(memSize); // The staging area has the same layout in the pinned host buffer and in // the device buffer, so it can be uploaded with a single copy. @@ -226,16 +278,21 @@ public: _computeFunctions[computer]<<>>(*batchesDevice); // Only the NormalizedPdf computer mutates its extra args: it uses them - // as output parameters for the evaluation error counts. Only then the - // extra args need to be copied back, and the stream needs to be - // synchronized because the caller inspects the counters right after - // this function returns. + // as output parameters for the evaluation error counts. Instead of + // synchronizing the stream to read the counters back immediately, the + // readback is deferred to avoid stalling the pipeline: an asynchronous + // copy delivers them into a persistent pinned buffer, and the next + // synchronizeCudaStream() call forwards them to the caller's span. The + // caller's memory therefore has to stay valid until then. if (computer == NormalizedPdf && !extraArgs.empty()) { - copyDeviceToHost(extraArgsDevice, extraArgs.data(), extraArgs.size(), cfg.cudaStream()); - ERRCHECK(cudaStreamSynchronize(stream)); + const std::size_t nBytes = extraArgs.size() * sizeof(double); + StreamScratch::DeferredSlot &deferredSlot = streamScratch.acquireDeferred(nBytes); + ERRCHECK(cudaMemcpyAsync(deferredSlot.host, extraArgsDevice, nBytes, cudaMemcpyDeviceToHost, stream)); + deferredSlot.dst = extraArgs.data(); + deferredSlot.nPending = extraArgs.size(); } - scratch(cfg.cudaStream()).release(slot, stream); + streamScratch.release(slot, stream); } /// Return the sum of an input array double reduceSum(RooBatchCompute::Config const &cfg, InputArr input, size_t n) override; @@ -253,6 +310,12 @@ public: void synchronizeCudaStream(CudaInterface::CudaStream *stream) const override { ERRCHECK(::cudaStreamSynchronize(*stream)); + // Deliver deferred readbacks (e.g. the evaluation error counters from + // compute()) that have completed with the synchronization. + auto found = _scratchMap.find(stream); + if (found != _scratchMap.end()) { + found->second.flushDeferred(); + } } private: diff --git a/roofit/roofitcore/inc/RooFit/Detail/RooNormalizedPdf.h b/roofit/roofitcore/inc/RooFit/Detail/RooNormalizedPdf.h index ff70d3a96edf6..f2424b1cb3051 100644 --- a/roofit/roofitcore/inc/RooFit/Detail/RooNormalizedPdf.h +++ b/roofit/roofitcore/inc/RooFit/Detail/RooNormalizedPdf.h @@ -16,6 +16,8 @@ #include #include +#include + namespace RooFit::Detail { class RooNormalizedPdf : public RooAbsPdf { @@ -85,10 +87,18 @@ class RooNormalizedPdf : public RooAbsPdf { double getValV(const RooArgSet * normSet) const override; private: + void logEvalErrorCounts() const; + RooTemplateProxy _pdf; RooRealProxy _normIntegral; RooArgSet _normSet; + /// Evaluation error counters, filled by the compute function. In CUDA + /// mode, they are read back from the GPU asynchronously and only arrive + /// after the evaluation of the computation graph, so they have to live in + /// a member and not on the stack of doEval(). Transient and not copied. + mutable std::array _evalErrorCounts{}; // +#include #include #include #include @@ -123,6 +124,13 @@ class EvalContext { void setOutputWithOffset(RooAbsArg const *arg, ROOT::Math::KahanSum val, ROOT::Math::KahanSum const &offset); + /// Register an action to be run after the evaluation of the full + /// computation graph, when all potentially asynchronous computations and + /// data transfers have completed. Used to defer work that depends on + /// results that are read back from the GPU without synchronization, like + /// the logging of evaluation error counts. + void deferAction(std::function action) { _deferredActions.emplace_back(std::move(action)); } + private: friend class Evaluator; @@ -134,6 +142,7 @@ class EvalContext { std::vector> _buffers; std::size_t _bufferIdx = 0; std::vector _cfgs; + std::vector> _deferredActions; }; } // namespace RooFit diff --git a/roofit/roofitcore/src/RooFit/Evaluator.cxx b/roofit/roofitcore/src/RooFit/Evaluator.cxx index c3ac9b97d9c66..7de4163112566 100644 --- a/roofit/roofitcore/src/RooFit/Evaluator.cxx +++ b/roofit/roofitcore/src/RooFit/Evaluator.cxx @@ -448,6 +448,11 @@ std::span Evaluator::run() ++_nEvaluations; + // Discard leftover deferred actions in case a previous evaluation was + // aborted by an exception. + _evalContextCPU._deferredActions.clear(); + _evalContextCUDA._deferredActions.clear(); + if (_useGPU) { return getValHeterogeneous(); } @@ -468,6 +473,11 @@ std::span Evaluator::run() } } + for (auto &action : _evalContextCPU._deferredActions) { + action(); + } + _evalContextCPU._deferredActions.clear(); + // return the final output return _evalContextCPU.at(&_topNode); } @@ -508,11 +518,20 @@ std::span Evaluator::getValHeterogeneous() } // Ensure that all enqueued GPU work has completed when run() returns. For - // the usual likelihood evaluations this is 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. + // 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. RooBatchCompute::dispatchCUDA->synchronizeCudaStream(_cudaStream); + // Run the deferred actions now that all results have arrived on the host, + // e.g. the logging of evaluation errors that were counted on the GPU. + for (auto &action : _evalContextCUDA._deferredActions) { + action(); + } + _evalContextCUDA._deferredActions.clear(); + // return the final value return _evalContextCUDA.at(&_topNode); } diff --git a/roofit/roofitcore/src/RooNormalizedPdf.cxx b/roofit/roofitcore/src/RooNormalizedPdf.cxx index b6eabfc20662b..1505c0e811548 100644 --- a/roofit/roofitcore/src/RooNormalizedPdf.cxx +++ b/roofit/roofitcore/src/RooNormalizedPdf.cxx @@ -32,14 +32,28 @@ void RooNormalizedPdf::doEval(RooFit::EvalContext &ctx) const auto integralSpan = ctx.at(_normIntegral); // We use the extraArgs as output parameter to count evaluation errors. - std::array extraArgs{0.0, 0.0, 0.0}; + _evalErrorCounts = {}; - RooBatchCompute::compute(ctx.config(this), RooBatchCompute::NormalizedPdf, ctx.output(), {nums, integralSpan}, - extraArgs); + auto config = ctx.config(this); + RooBatchCompute::compute(config, RooBatchCompute::NormalizedPdf, ctx.output(), {nums, integralSpan}, + _evalErrorCounts); - std::size_t nEvalErrorsType0 = extraArgs[0]; - std::size_t nEvalErrorsType1 = extraArgs[1]; - std::size_t nEvalErrorsType2 = extraArgs[2]; + if (config.useCuda()) { + // In CUDA mode, the counters are read back from the GPU without + // synchronizing the stream: they only arrive in _evalErrorCounts with + // the synchronization at the end of the evaluation of the full + // computation graph, so the logging has to be deferred until then. + ctx.deferAction([this] { logEvalErrorCounts(); }); + } else { + logEvalErrorCounts(); + } +} + +void RooNormalizedPdf::logEvalErrorCounts() const +{ + const std::size_t nEvalErrorsType0 = _evalErrorCounts[0]; + const std::size_t nEvalErrorsType1 = _evalErrorCounts[1]; + const std::size_t nEvalErrorsType2 = _evalErrorCounts[2]; for (std::size_t i = 0; i < nEvalErrorsType0; ++i) { logEvalError("p.d.f normalization integral is zero or negative"); From 9faad7377c14bb186ec4002e2af40da71c4c6780 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 4 Sep 2026 18:18:55 +0000 Subject: [PATCH 8/8] [RF] Harden the CUDA evaluation against races and aborted evaluations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the confirmed findings of a code review of the CUDA evaluation refactoring: * Don't recycle pinned host buffers eagerly during an evaluation. Since the host-to-device uploads are genuinely asynchronous now, a released pinned buffer could be handed to another CPU-evaluated node that overwrites the pinned host array with a plain CPU write - which is not ordered by the CUDA stream - while the previous owner's upload is still pending, silently corrupting GPU inputs. Buffers of nodes with copyAfterEvaluation are now released only at the beginning of the next evaluation, after the stream synchronization at the end of the current one. * If an evaluation is aborted by an exception, synchronize the stream and deliver the deferred readbacks while the destination memory in the computation graph is guaranteed to be alive, so that no armed readback survives into a later evaluation where it would flush stale data (or, in the worst case, write to freed memory). * Make the reallocation paths of the stream scratch memory exception safe: reset the slot state before reallocating, so a throwing cudaMalloc can't leave dangling pointers with a stale capacity behind that would later lead to a double free. * In the reductions, record the slot guard event right after the last enqueued use instead of after the (throwing) stream synchronization, so the slot stays protected on the exception path. * Drain the deferred actions of both evaluation contexts in the heterogeneous evaluation, so actions registered by CPU-evaluated nodes can not be silently discarded. * Document the narrowed extra-args contract of compute() in the interface header: only NormalizedPdf receives output through the extra args, and in CUDA mode the caller's memory must stay valid until the next stream synchronization. Verified with a heterogeneous model (two CUDA-unsupported pdfs in a RooAddPdf, exercising multiple same-size pinned boundary buffers per evaluation), all previous parity checks, and the full RooFit test suite including the CUDA stress tests. Benchmark timings are unchanged. 🤖 Done with the help of AI --- roofit/batchcompute/res/RooBatchCompute.h | 11 ++++ roofit/batchcompute/src/ComputeFunctions.cxx | 5 +- roofit/batchcompute/src/RooBatchCompute.cu | 37 ++++++++--- roofit/roofitcore/src/RooFit/Evaluator.cxx | 65 ++++++++++++++------ 4 files changed, 89 insertions(+), 29 deletions(-) diff --git a/roofit/batchcompute/res/RooBatchCompute.h b/roofit/batchcompute/res/RooBatchCompute.h index a0bbbf5dec73b..db57d1f7adc88 100644 --- a/roofit/batchcompute/res/RooBatchCompute.h +++ b/roofit/batchcompute/res/RooBatchCompute.h @@ -168,6 +168,17 @@ class AbsBufferManager { class RooBatchComputeInterface { public: virtual ~RooBatchComputeInterface() = default; + + /// Compute the values for a batch of events. + /// + /// The extra args (the last parameter) are read-only inputs for all + /// computers except `NormalizedPdf`, which uses them as output parameters + /// for its evaluation error counters. In the CUDA implementation, these + /// outputs are read back from the device *asynchronously*: they only + /// arrive in the caller's span with the next synchronizeCudaStream() call + /// on the stream of the passed config. The memory backing the extra args + /// of a `NormalizedPdf` call must therefore stay valid until that + /// synchronization, so it must not live on the caller's stack. virtual void compute(Config const &cfg, Computer, std::span output, VarSpan, ArgSpan) = 0; virtual double reduceSum(Config const &cfg, InputArr input, size_t n) = 0; diff --git a/roofit/batchcompute/src/ComputeFunctions.cxx b/roofit/batchcompute/src/ComputeFunctions.cxx index f218bca6db459..e1ae37eb1137a 100644 --- a/roofit/batchcompute/src/ComputeFunctions.cxx +++ b/roofit/batchcompute/src/ComputeFunctions.cxx @@ -661,7 +661,10 @@ __rooglobal__ void computeNormalizedPdf(Batches &batches) } // The counters live in memory that is shared between all threads in the - // CUDA case, so they need to be accumulated atomically there. + // CUDA case, so they need to be accumulated atomically there. Note that + // the CPU branch below is only safe because the CPU implementation runs + // single-threaded: with implicit multi-threading, the workers would share + // this memory as well and would also need atomic accumulation. #ifdef __CUDACC__ if (nEvalErrorsType0 > 0) atomicAdd(&batches.extra[0], double(nEvalErrorsType0)); diff --git a/roofit/batchcompute/src/RooBatchCompute.cu b/roofit/batchcompute/src/RooBatchCompute.cu index 8615f8aaf6031..29c9617b57bb0 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cu +++ b/roofit/batchcompute/src/RooBatchCompute.cu @@ -120,10 +120,18 @@ public: slot.inFlight = false; } if (slot.capacity < n) { - if (slot.host) + // Reset the slot state before reallocating, so that a throwing + // allocation can't leave dangling pointers with a stale capacity + // behind (which would lead to a double free later). + if (slot.host) { ERRCHECK(cudaFreeHost(slot.host)); - if (slot.device) + slot.host = nullptr; + } + if (slot.device) { ERRCHECK(cudaFree(slot.device)); + slot.device = nullptr; + } + slot.capacity = 0; const std::size_t newCapacity = std::max(n, 1024); ERRCHECK(cudaMallocHost(reinterpret_cast(&slot.host), newCapacity)); ERRCHECK(cudaMalloc(reinterpret_cast(&slot.device), newCapacity)); @@ -163,9 +171,13 @@ public: DeferredSlot &slot = _deferredSlots[_deferredCursor++]; if (slot.capacity < n) { // The slot is idle here: its previous use ended with the flush after - // a stream synchronization. - if (slot.host) + // a stream synchronization. Reset the state before reallocating for + // exception safety, like in acquire(). + if (slot.host) { ERRCHECK(cudaFreeHost(slot.host)); + slot.host = nullptr; + } + slot.capacity = 0; ERRCHECK(cudaMallocHost(reinterpret_cast(&slot.host), n)); slot.capacity = n; } @@ -468,17 +480,19 @@ double RooBatchComputeClass::reduceSum(RooBatchCompute::Config const &cfg, Input return 0.0; const int gridSize = getGridSize(n); cudaStream_t stream = *cfg.cudaStream(); - StreamScratch::Slot &slot = scratch(cfg.cudaStream()).acquire(2 * gridSize * sizeof(double)); + StreamScratch &streamScratch = scratch(cfg.cudaStream()); + StreamScratch::Slot &slot = streamScratch.acquire(2 * gridSize * sizeof(double)); auto devOut = reinterpret_cast(slot.device); auto hostOut = reinterpret_cast(slot.host); constexpr int shMemSize = 2 * blockSize * sizeof(double); kahanSum<<>>(input, nullptr, n, devOut, 0); kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut, devOut + gridSize, gridSize, devOut, 0); CudaInterface::copyDeviceToHost(devOut, hostOut, 1, cfg.cudaStream()); + // Release right after the last enqueued use of the slot, so that the slot + // is protected by its event even if the synchronization below throws. + streamScratch.release(slot, stream); ERRCHECK(cudaStreamSynchronize(stream)); - const double result = hostOut[0]; - scratch(cfg.cudaStream()).release(slot, stream); - return result; + return hostOut[0]; } ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &cfg, std::span probas, @@ -492,7 +506,8 @@ ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &c cudaStream_t stream = *cfg.cudaStream(); // Layout of the scratch buffer: [sum, carry, badness, nNonPositive, nNaN, // nInfinite, partial sums (gridSize), partial carries (gridSize)]. - StreamScratch::Slot &slot = scratch(cfg.cudaStream()).acquire((6 + 2 * gridSize) * sizeof(double)); + StreamScratch &streamScratch = scratch(cfg.cudaStream()); + StreamScratch::Slot &slot = streamScratch.acquire((6 + 2 * gridSize) * sizeof(double)); auto devOut = reinterpret_cast(slot.device); auto hostOut = reinterpret_cast(slot.host); constexpr int shMemSize = 2 * blockSize * sizeof(double); @@ -520,6 +535,9 @@ ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &c // The sum, its Kahan carry, and the evaluation error statistics are // adjacent in the output buffer, so they can be read back in a single copy. CudaInterface::copyDeviceToHost(devOut, hostOut, 6, cfg.cudaStream()); + // Release right after the last enqueued use of the slot, so that the slot + // is protected by its event even if the synchronization below throws. + streamScratch.release(slot, stream); ERRCHECK(cudaStreamSynchronize(stream)); out.nllSum = hostOut[0]; @@ -536,7 +554,6 @@ ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &c out.nllSumCarry = 0.0; } - scratch(cfg.cudaStream()).release(slot, stream); return out; } diff --git a/roofit/roofitcore/src/RooFit/Evaluator.cxx b/roofit/roofitcore/src/RooFit/Evaluator.cxx index 7de4163112566..5a10c36ed039a 100644 --- a/roofit/roofitcore/src/RooFit/Evaluator.cxx +++ b/roofit/roofitcore/src/RooFit/Evaluator.cxx @@ -120,10 +120,16 @@ struct NodeInfo { std::vector clientInfos; /// Check the servers of a node that has been computed and release its - /// resources if they are no longer needed. + /// 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) { + if (--remClients == 0 && !fromArrayInput && !copyAfterEvaluation) { buffer.reset(); } } @@ -499,22 +505,41 @@ std::span Evaluator::getValHeterogeneous() // 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. - for (auto &info : _nodes) { - if (!info.fromArrayInput) { - if (info.computeInGPU) { - assignToGPU(info); - } else { - computeCPUNode(info.absArg, info); + try { + for (auto &info : _nodes) { + if (!info.fromArrayInput) { + if (info.computeInGPU) { + assignToGPU(info); + } else { + computeCPUNode(info.absArg, info); + } } - } - // Release the buffers of server nodes that are no longer needed. 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. - for (auto *serverInfo : info.serverInfos) { - serverInfo->decrementRemainingClients(); + // 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(); + } + } + } catch (...) { + // The evaluation was aborted, but readbacks that compute() calls + // 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. + try { + RooBatchCompute::dispatchCUDA->synchronizeCudaStream(_cudaStream); + } catch (...) { + // The stream is in an unrecoverable error state. The deferred + // readbacks are dropped together with the scratch memory when the + // stream gets deleted. } + _evalContextCUDA._deferredActions.clear(); + _evalContextCPU._deferredActions.clear(); + throw; } // Ensure that all enqueued GPU work has completed when run() returns. For @@ -527,10 +552,14 @@ std::span Evaluator::getValHeterogeneous() // Run the deferred actions now that all results have arrived on the host, // e.g. the logging of evaluation errors that were counted on the GPU. - for (auto &action : _evalContextCUDA._deferredActions) { - action(); + // Nodes evaluated on the CPU register their actions in the CPU context, + // so both contexts are drained. + for (auto *ctx : {&_evalContextCUDA, &_evalContextCPU}) { + for (auto &action : ctx->_deferredActions) { + action(); + } + ctx->_deferredActions.clear(); } - _evalContextCUDA._deferredActions.clear(); // return the final value return _evalContextCUDA.at(&_topNode);