[RF] Architectual improvements in CUDA backend that speed up unbinned fits by up to factor two - #23263
Open
guitargeek wants to merge 8 commits into
Open
[RF] Architectual improvements in CUDA backend that speed up unbinned fits by up to factor two#23263guitargeek wants to merge 8 commits into
guitargeek wants to merge 8 commits into
Conversation
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
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
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
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
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
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
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
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
guitargeek
changed the title
Test Results 23 files 23 suites 3d 18h 24m 7s ⏱️ For more details on these failures, see this check. Results for commit 9faad73. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR overhauls the RooFit CUDA backend around a single CUDA stream per
RooFit::Evaluator, replacing the per-node streams/events and the polling scheduler. On top of that, it removes all allocations and most synchronization from the evaluation hot loop and brings the CUDA NLL reduction to full semantic parity with the CPU path.These are architectural and feature parity improvements that I always knew had to be done at some point, but the CUDA backend was not a priority in the recent years. Now that I can tell AI to implement the architecture vision, the cost function for this work has shifted and it was the time to do it.
All of the unbinned RooFit benchmarks in rootbench (1M events) drastically benefit from the new single-stream architecture. The speedup is up to a factor two:
Note that the benchmarks were only ran after the PR was completed, to not risk biasing the implementation too much by benchmark overfitting.