diff --git a/roofit/batchcompute/res/RooBatchCompute.h b/roofit/batchcompute/res/RooBatchCompute.h index bd7a7301b1d7a..db57d1f7adc88 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 @@ -169,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; @@ -180,13 +190,10 @@ 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; + /// Wait until all work that was enqueued on the stream has completed. + virtual void synchronizeCudaStream(CudaInterface::CudaStream *) const = 0; }; /** 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/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/ComputeFunctions.cxx b/roofit/batchcompute/src/ComputeFunctions.cxx index 3f67a736c7eb2..e1ae37eb1137a 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) @@ -664,12 +660,23 @@ __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. 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) - 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/CudaInterface.cu b/roofit/batchcompute/src/CudaInterface.cu index 34a8c1c042643..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 { @@ -55,29 +41,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 +50,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 +62,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..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 { /* @@ -37,21 +60,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 +67,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 +76,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); @@ -89,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); } /** @@ -103,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); } /** @@ -117,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 5397589ad5511..29c9617b57bb0 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cu +++ b/roofit/batchcompute/src/RooBatchCompute.cu @@ -19,14 +19,18 @@ This file contains the code for cuda computations using the RooBatchCompute libr **/ #include "RooBatchCompute.h" +#include "RooNaNPacker.h" #include "Batches.h" #include "CudaInterface.h" #include +#include #include +#include #include #include #include +#include #include namespace RooBatchCompute { @@ -81,6 +85,146 @@ 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) { + // 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)); + 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)); + 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; + } + + /// 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. 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; + } + 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. + 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); + } + 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 std::vector getFunctions(); @@ -115,14 +259,18 @@ 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 &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. + 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 +284,27 @@ 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()) { - copyDeviceToHost(extraArgsDevice, extraArgs.data(), extraArgs.size(), cfg.cudaStream()); + _computeFunctions[computer]<<>>(*batchesDevice); + + // Only the NormalizedPdf computer mutates its extra args: it uses them + // 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()) { + 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(); } + + streamScratch.release(slot, stream); } /// Return the sum of an input array double reduceSum(RooBatchCompute::Config const &cfg, InputArr input, size_t n) override; @@ -155,26 +313,28 @@ 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 + void deleteCudaStream(CudaInterface::CudaStream *stream) const override { - CudaInterface::cudaEventRecord(*event, *stream); + _scratchMap.erase(stream); + delete stream; } - void cudaStreamWaitForEvent(CudaInterface::CudaStream *stream, CudaInterface::CudaEvent *event) const override + void synchronizeCudaStream(CudaInterface::CudaStream *stream) const override { - stream->waitForEvent(*event); + 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(); + } } - bool cudaStreamIsActive(CudaInterface::CudaStream *stream) const override { return stream->isActive(); } private: + StreamScratch &scratch(CudaInterface::CudaStream *stream) { return _scratchMap[stream]; } + const std::vector _computeFunctions; + mutable std::unordered_map _scratchMap; }; // End class RooBatchComputeClass @@ -239,9 +399,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; @@ -253,17 +419,52 @@ __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; @@ -279,13 +480,19 @@ 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 &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.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()); + // 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)); + return hostOut[0]; } ReduceNLLOutput RooBatchComputeClass::reduceNLL(RooBatchCompute::Config const &cfg, std::span probas, @@ -296,31 +503,57 @@ 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(); + // Layout of the scratch buffer: [sum, carry, badness, nNonPositive, nNaN, + // nInfinite, partial sums (gridSize), partial carries (gridSize)]. + 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); #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 + // 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.data()); - - kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut.data(), devOut.data() + gridSize, gridSize, devOut.data(), 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()); + probas.size() == 1 ? probas[0] : 0.0, weights.size(), devOut + 6, devOut + 2); + + kahanSum<<<1, blockSize, shMemSize, stream>>>(devOut + 6, devOut + 6 + gridSize, gridSize, devOut, 0); + + // 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]; + 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; + } - out.nllSum = tmpSum; - out.nllSumCarry = tmpCarry; return out; } @@ -423,6 +656,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 d98e0c828d6b5..4dcf6bbf1d1eb 100644 --- a/roofit/batchcompute/src/RooBatchCompute.cxx +++ b/roofit/batchcompute/src/RooBatchCompute.cxx @@ -104,19 +104,9 @@ 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(); } + void synchronizeCudaStream(CudaInterface::CudaStream *) const override { throw std::bad_function_call(); } private: #ifdef ROOBATCHCOMPUTE_USE_IMT 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.}; /// #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 @@ -107,6 +108,15 @@ class EvalContext { } RooBatchCompute::Config config(RooAbsArg const *arg) const; + + /// A counter that is incremented every time new input data is loaded into + /// the evaluation context. Reducer nodes can use it as a cache + /// invalidation key for quantities that only depend on the input data, + /// like the sum of event weights. The counter values are unique across + /// all evaluation contexts in the process, so cached values can not be + /// wrongly validated by an unrelated context. + std::size_t inputGeneration() const { return _inputGeneration; } + void enableVectorBuffers(bool enable) { _enableVectorBuffers = enable; } void resetVectorBuffers() { _bufferIdx = 0; } std::span output() { return _currentOutput; } @@ -114,16 +124,25 @@ 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; OffsetMode _offsetMode = OffsetMode::WithoutOffset; + std::size_t _inputGeneration = 1; std::span _currentOutput; std::vector> _ctx; bool _enableVectorBuffers = false; std::vector> _buffers; std::size_t _bufferIdx = 0; std::vector _cfgs; + std::vector> _deferredActions; }; } // namespace RooFit 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..5a10c36ed039a 100644 --- a/roofit/roofitcore/src/RooFit/Evaluator.cxx +++ b/roofit/roofitcore/src/RooFit/Evaluator.cxx @@ -42,10 +42,9 @@ RooAbsPdf::fitTo() is called and gets destroyed when the fitting ends. #include "BatchModeDataHelpers.h" #include "RooFitImplHelpers.h" -#include +#include #include #include -#include #include namespace RooFit { @@ -105,7 +104,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,25 +119,20 @@ 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. + /// 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(); } } - - ~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 +215,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); } } @@ -268,6 +263,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; @@ -342,6 +347,9 @@ Evaluator::~Evaluator() info.absArg->resetDataToken(); } } + if (_cudaStream) { + RooBatchCompute::dispatchCUDA->deleteCudaStream(_cudaStream); + } } void Evaluator::computeCPUNode(const RooAbsArg *node, NodeInfo &info) @@ -366,7 +374,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 +398,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); - } } } @@ -446,6 +454,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(); } @@ -466,6 +479,11 @@ std::span Evaluator::run() } } + for (auto &action : _evalContextCPU._deferredActions) { + action(); + } + _evalContextCPU._deferredActions.clear(); + // return the final output return _evalContextCPU.at(&_topNode); } @@ -475,91 +493,85 @@ 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 - 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 + // 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. + try { 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(); + if (!info.fromArrayInput) { + if (info.computeInGPU) { + assignToGPU(info); + } else { + computeCPUNode(info.absArg, info); } } - } - - // 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); + // 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(); } } - 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 + // 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. + // 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(); } // return the final value 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 +579,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}); } } 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, 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"); 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) {