[RF] Architectual improvements in CUDA backend that speed up unbinned fits by up to factor two

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<double> 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<AbsBufferManager> 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;
 };
 
 /**

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

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));

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;

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))

roofit/batchcompute/src/CudaInterface.cu

@@ -12,20 +12,6 @@
 
 #include "CudaInterface.h"
 
-#include <stdexcept>
-#include <sstream>
-#include <string>
-
-#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<PinnedHostMemory>::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<CudaEvent>::operator()(void *ptr)
-{
-   auto event = reinterpret_cast<cudaEvent_t *>(ptr);
-   ERRCHECK(cudaEventDestroy(*event));
-   delete event;
-   ptr = nullptr;
-}
-
 template <>
 void Deleter<CudaStream>::operator()(void *ptr)
 {
@@ -87,17 +50,6 @@ void Deleter<CudaStream>::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)

roofit/batchcompute/src/CudaInterface.h

@@ -16,6 +16,29 @@
 #include <cstddef>
 #include <memory>
 
+#ifdef __CUDACC__
+#include <sstream>
+#include <stdexcept>
+#include <string>
+
+#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,31 +60,13 @@ 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<cudaEvent_t *>(_ptr.get()); }
-#endif
-private:
-   std::unique_ptr<void, Deleter<CudaEvent>> _ptr;
-};
-
 /*
  * Wrapper around cudaStream_t.
  */
 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<cudaStream_t *&>(_ptr); }
@@ -71,9 +76,6 @@ class CudaStream {
    std::unique_ptr<void, Deleter<CudaStream>> _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 <class T>
-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 <class T>
-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 <class T>
-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