[RF] New RooFit batch mode with GPU support

cmake/modules/RootConfiguration.cmake

@@ -378,6 +378,11 @@ if(CMAKE_USE_PTHREADS_INIT)
 else()
   set(haspthread undef)
 endif()
+if(cuda)
+  set(hascuda define)
+else()
+  set(hascuda undef)
+endif()
 if(x11)
   set(hasxft define)
 else()

config/RConfigure.in

@@ -25,6 +25,7 @@
 #@setresuid@ R__HAS_SETRESUID   /**/
 #@hasmathmore@ R__HAS_MATHMORE   /**/
 #@haspthread@ R__HAS_PTHREAD    /**/
+#@hascuda@ R__HAS_CUDA    /**/
 #@hasxft@ R__HAS_XFT    /**/
 #@hascocoa@ R__HAS_COCOA    /**/
 #@hasvc@ R__HAS_VC    /**/

roofit/batchcompute/CMakeLists.txt

@@ -1,6 +1,7 @@
 # Library which powers fast batch computations in Roofit.
 ROOT_LINKER_LIBRARY(RooBatchCompute
     src/Initialisation.cxx
+    src/RooMath.cxx
     src/RunContext.cxx
   DEPENDENCIES
     Core
@@ -24,32 +25,38 @@ endif()
 # Instantiations of the shared objects which provide the actual computation functions.
 
 # Generic implementation for CPUs that don't support vector instruction sets.
-ROOT_LINKER_LIBRARY(RooBatchCompute_GENERIC src/RooBatchCompute.cxx TYPE SHARED DEPENDENCIES RooFitCore RooBatchCompute)
+ROOT_LINKER_LIBRARY(RooBatchCompute_GENERIC src/RooBatchCompute.cxx src/ComputeFunctions.cxx TYPE SHARED DEPENDENCIES RooBatchCompute)
+target_compile_options(RooBatchCompute_GENERIC  PRIVATE ${common-flags} -DRF_ARCH=GENERIC)
 
 # Windows platform and ICC compiler need special code and testing, thus the feature has not been implemented yet for these.
 if (ROOT_PLATFORM MATCHES "linux|macosx" AND CMAKE_SYSTEM_PROCESSOR MATCHES x86_64 AND CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
 
   target_compile_options(RooBatchCompute PRIVATE -DR__RF_ARCHITECTURE_SPECIFIC_LIBS)
 
-  ROOT_LINKER_LIBRARY(RooBatchCompute_SSE4.1  src/RooBatchCompute.cxx TYPE SHARED DEPENDENCIES RooFitCore RooBatchCompute)
-  ROOT_LINKER_LIBRARY(RooBatchCompute_AVX     src/RooBatchCompute.cxx TYPE SHARED DEPENDENCIES RooFitCore RooBatchCompute)
-  ROOT_LINKER_LIBRARY(RooBatchCompute_AVX2    src/RooBatchCompute.cxx TYPE SHARED DEPENDENCIES RooFitCore RooBatchCompute)
+  ROOT_LINKER_LIBRARY(RooBatchCompute_SSE4.1  src/RooBatchCompute.cxx src/ComputeFunctions.cxx TYPE SHARED DEPENDENCIES RooBatchCompute)
+  ROOT_LINKER_LIBRARY(RooBatchCompute_AVX     src/RooBatchCompute.cxx src/ComputeFunctions.cxx TYPE SHARED DEPENDENCIES RooBatchCompute)
+  ROOT_LINKER_LIBRARY(RooBatchCompute_AVX2    src/RooBatchCompute.cxx src/ComputeFunctions.cxx TYPE SHARED DEPENDENCIES RooBatchCompute)
 
   # Flags -fno-signaling-nans, -fno-trapping-math and -O3 are necessary to enable autovectorization (especially for GCC).
   set(common-flags $<$<CXX_COMPILER_ID:GNU>:-fno-signaling-nans>)
   list(APPEND common-flags $<$<OR:$<CONFIG:Release>,$<CONFIG:RelWithDebInfo>>: -fno-trapping-math -O3>)
 
   target_compile_options(RooBatchCompute_SSE4.1  PRIVATE ${common-flags} -msse4    -DRF_ARCH=SSE4)
-  target_compile_options(RooBatchCompute_AVX     PRIVATE ${common-flags} -mavx     -DRF_ARCH=AVX )
+  target_compile_options(RooBatchCompute_AVX     PRIVATE ${common-flags} -mavx     -DRF_ARCH=AVX)
   target_compile_options(RooBatchCompute_AVX2    PRIVATE ${common-flags} -mavx2    -DRF_ARCH=AVX2)
 
   # AVX512 is only supported in gcc 6+
   # We focus on AVX512 capable processors that support at least the skylake-avx512 instruction sets.
   if(NOT (CMAKE_CXX_COMPILER_ID STREQUAL "GNU") OR CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 6)
-    ROOT_LINKER_LIBRARY(RooBatchCompute_AVX512  src/RooBatchCompute.cxx TYPE SHARED DEPENDENCIES RooFitCore RooBatchCompute)
+    ROOT_LINKER_LIBRARY(RooBatchCompute_AVX512  src/RooBatchCompute.cxx src/ComputeFunctions.cxx TYPE SHARED DEPENDENCIES RooBatchCompute)
     target_compile_options(RooBatchCompute_AVX512  PRIVATE ${common-flags} -march=skylake-avx512 -DRF_ARCH=AVX512)
   endif()
 
+endif() # vector versions of library
+
+if (cuda)
+  ROOT_LINKER_LIBRARY(RooBatchCompute_CUDA  src/RooBatchCompute.cu src/ComputeFunctions.cu TYPE SHARED DEPENDENCIES RooBatchCompute)
+  target_compile_options(RooBatchCompute_CUDA  PRIVATE -DRF_ARCH=CUDA -lineinfo --expt-relaxed-constexpr)
 endif()
 
 ROOT_INSTALL_HEADERS()

roofit/batchcompute/README.md

@@ -5,63 +5,54 @@ _Contains optimized computation functions for PDFs that enable significantly fas
 ### Purpose
 While fitting, a significant amount of time and processing power is spent on computing the probability function for every event and PDF involved in the fitting model. To speed up this process, roofit can use the computation functions provided in this library. The functions provided here process whole data arrays (batches) instead of a single event at a time, as in the legacy evaluate() function in roofit. In addition, the code is written in a manner that allows for compiler optimizations, notably auto-vectorization. This library is compiled multiple times for different [vector instuction set architectures](https://en.wikipedia.org/wiki/SIMD) and the optimal code is executed during runtime, as a result of an automatic hardware detection mechanism that this library contains. **As a result, fits can benefit by a speedup of 3x-16x.**
 
+As of ROOT v6.26, RooBatchComputes also provides multithread and [CUDA](https://en.wikipedia.org/wiki/CUDA) instances of the computation functions, resulting in even greater improvements for fitting times.
+
 ### How to use
+This library is an internal component of RooFit, as a results users are not supposed to actively interact with it. Instead, they can benefit from significantly faster times for fitting by calling `fitTo()` and providing a `BatchMode(RooBatchCompute::Cpu)` or a `BatchMode(RooBatchCompute::Cuda)` option. 
+```c++
+  // fit using the most efficient library that the computer's cpu can support
+  RooMyPDF.fitTo(data, BatchMode("cpu")); 
+
+  // fit using the cuda library along with the most efficient library that the computer's cpu can support
+  RooMyPDF.fitTo(data, BatchMode("cuda")); 
+```
+**Note: In case the system does not support vector instructions, the `RooBatchCompute::Cpu` option is guaranteed to work properly by using a generic cpu library. In contrast, users must first make sure that their system supports cuda in order to use the `RooBatchCompute::Cuda` option. If this is not the case, an exemption will be thrown.**
+
+If `RooBatchCompute::Cuda` is selected, RooFit will launch CUDA kernels for computing possibilities and potentially other intense computations. At the same time, the most efficent cpu library loaded will also handle parts of the computations in parallel with the GPU (or potentially, if it's faster all of them), thus gaining full advantage of the available hardware. For this purpose `RooFitDriver`, a newly created RooFit class (in roofitcore) takes over the task of analyzing the computations and assigning each to the correct piece of hardware, taking into consideration the performance boost or penalty that may arise with every method of computing.
+
+#### Multithread computations
+The CPU instance of the computing library can furthermore execute multithread computations. This also applies for computations handled by the CPU in the `RooBatchCompute::Cuda` mode. To use them, one needs to set the desired number of parallel tasks before calling `fitTo()` as shown below:
+```c++
+  ROOT::EnableImplicitMT(nThreads);
+  RooMyPDF.fitTo(data, BatchMode(RooBatchCompute::Cpu)); // can also use RooBatchCompute::Cuda
+```
+
+### User-made PDFs
 The easiest and most efficient way of accelerating your PDFs is to request their addition to the official RooFit by submiting a ticket [here](https://github.com/root-project/root/issues/new). The ROOT team will gladly assist you and take care of the details.
 
-The above process might take some time and the users will be required to update ROOT to use the newly introduced PDFs. In the meantime, you are able to significantly improve the speed of fitting (but not take full advantage of the RooBatchCompute library), at least by using the batch evaluation feature.
-To make use of it, one should override [`RooAbsReal::evaluateSpan()`](https://root.cern.ch/doc/master/classRooAbsReal.html#a1e5129ffbc63bfd04c01511fd354b1b8)
+While your code is integrated, you are able to significantly improve the speed of fitting (but not take full advantage of the RooBatchCompute library), at least by using the batch evaluation feature.
+To make use of it, one should override `RooAbsReal::computeBatch()`
 ```c++
-  RooSpan<double> RooMyPDF::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const
+  void RooMyPDF::computeBatch(RooBatchCompute::RooBatchComputeInterface*, double* output, size_t nEvents, RooBatchCompute::DataMap& dataMap) const
 ```
-The evalData is a simple struct that holds the vector data for the fitting in the form of `RooSpan<double>`.
-The normSet (normalization set) is used for invoking the computation and retrieving the values of the variables of the PDF.
-You don't need to worry about these arguments as they will be provided by the RooFit internal functions that will call `evaluateSpan()`.
+This method must be implemented so that it fills the `output` array with the **normalized** probabilities computed for `nEvents` events, the data of which can be retrieved from `dataMap`. `dataMap` is a simple `std::map<RooRealVar*, RooSpan<const double>>`. Note that it is not necessary to evaluate any of the objects that the PDF relies to, because they have already been evaluated by the RooFitDriver, so that their updated results are always present in `dataMap`. The `RooBatchCompute::RooBatchComputeInterface` pointer should be ignored.
+
 ```c++
-RooSpan<double> RooMyPDF::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const
+void RooMyPDF::computeBatch(RooBatchCompute::RooBatchComputeInterface*, double* output, size_t nEvents, RooBatchCompute::DataMap& dataMap) const
 {
   // Retrieve `RooSpan`s for each parameter of the PDF
-  RooSpan<const double> span1 = var1->getValues(evalData, normSet);
-  // or: auto span1 = var1->getValues(evalData, normSet);
-  RooSpan<const double> span2 = var2->getValues(evalData, normSet);
+  RooSpan<const double> span1 = dataMap.at(&*proxyVar1);
+  // or: auto span1 = dataMap.at(&*proxyVar1);
+  RooSpan<const double> span2 = dataMap.at(&*proxyVar2);
 
-  // let's assume c is a scalar parameter of the PDF. In this case getValues will return a RooSpan with only one value.
-  RooSpan<const double> scalar = c->getValues(evalData, normset);
-
-  // Get the number of nEvents
-  size_t nEvents=0;
-  for (auto& i:{xData,meanData,sigmaData})
-    nEvents = std::max(nEvents,i.size());
-
-  // Allocate the output array
-  evalData.makeBatch(this, nEvents);
+  // let's assume c is a scalar parameter of the PDF. In this case the dataMap contains a RooSpan with only one value.
+  RooSpan<const double> scalar = dataMap.at(&*c);
 
   // Perform computations in a for-loop
   // Use VDT if possible to facilitate auto-vectorization
   for (size_t i=0; i<nEvents; ++i) {
     output[i] = RooBatchCompute::fast_log(span1[i]+span2[i]) + scalar[0]; //scalar is a RooSpan of length 1
   }
-  return output;
 }
 ```
-Make sure to add the `evaluateSpan()` function signature in the header `RooMyPDF.h` and mark it as `override` to ensure that you have successfully overriden the method. In case the data types (scalar or vector) for the variables can not be predicted when writing the source code, you can use [BracketAdapterWithMask](https://github.com/root-project/root/blob/2b84398d4f52462a120083b3c5d1e0b952cc5221/roofit/batchcompute/inc/BracketAdapter.h#L55). This class overloads the `operator[]` and is constructed by a RooSpan. In case the RooSpan used for construction has a length of 1, ie represents a scalar variable, then `BracketAdapterWithMask::operator[]` always returns the scalar value, regrdless of the index used. This allows us to write:
-
-```c++
-RooSpan<double> RooMyPDF::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const
-{
-  // Construct BracketAdapterWithMasks for each variable if we're not sure whether they are scalar of vectors.
-  BracketAdapterWithMask adapter1(var1->getValues(evalData, normSet));
-  BracketAdapterWithMask adapter2(var2->getValues(evalData, normSet));
-  BracketAdapterWithMask scalar(c->getValues(evalData, normSet));
-
-  // prepare the computations as above
-  ...
-
-  // by calling adapter[i] we either get the i-th or 0-th element, if the variable is a vector or a scalar respectively.
-  for (size_t i=0; i<nEvents; ++i) {
-    output[i] = RooBatchCompute::fast_log(adapter1[i]+adapter2[i]) + scalar[i]; 
-  }
-  return output;
-}
-  ```
-
-  As a final note, always remember to append `RooBatchCompute::` to the classes defined in the RooBatchCompute library, or write `using namespace RooBatchCompute`.
+Make sure to add the `computeBatch()` function signature in the header `RooMyPDF.h` and mark it as `override` to ensure that you have successfully overriden the method. As a final note, always remember to append `RooBatchCompute::` to the classes defined in the RooBatchCompute library, or write `using namespace RooBatchCompute`.