[RF] Optimizations to avoid taking logarithms of Gaussian constraints

Large likelihood fits (e.g. in CMS and ATLAS) often contain many
Gaussian constraint terms. In the standard evaluation, these incur
unnecessary overhead by computing the PDF value via `exp(...)` only to
immediately take the logarithm in the constraint term.

This commit introduces a log-space optimization that avoids the
redundant exp -> log roundtrip.

We need this optimization because it makes a difference for CMS, and
they have also implemented something like this in the custom CMS Combine
likelihood classes. If we want them to gradually move to "vanilla"
RooFit, we need at least the same optimizations.

Author: guitargeek

roofit/batchcompute/res/RooBatchCompute.h

@@ -58,8 +58,12 @@ class Config {
    void setCudaStream(CudaInterface::CudaStream *cudaStream) { _cudaStream = cudaStream; }
    CudaInterface::CudaStream *cudaStream() const { return _cudaStream; }
 
+   bool takeLog() const { return _takeLog; }
+   void setTakeLog(bool takeLog) { _takeLog = takeLog; }
+
 private:
    CudaInterface::CudaStream *_cudaStream = nullptr;
+   bool _takeLog = false;
 };
 
 enum class Architecture {
@@ -90,6 +94,7 @@ enum Computer {
    Gamma,
    GaussModelExpBasis,
    Gaussian,
+   LogGaussian,
    Identity,
    Johnson,
    Landau,

roofit/batchcompute/src/ComputeFunctions.cxx

@@ -470,6 +470,18 @@ __rooglobal__ void computeGaussian(Batches &batches)
    }
 }
 
+__rooglobal__ void computeLogGaussian(Batches &batches)
+{
+   auto x = batches.args[0];
+   auto mean = batches.args[1];
+   auto sigma = batches.args[2];
+   for (size_t i = BEGIN; i < batches.nEvents; i += STEP) {
+      const double arg = x[i] - mean[i];
+      const double halfBySigmaSq = -0.5 / (sigma[i] * sigma[i]);
+      batches.output[i] = arg * arg * halfBySigmaSq;
+   }
+}
+
 __rooglobal__ void computeIdentity(Batches &batches)
 {
    for (size_t i = BEGIN; i < batches.nEvents; i += STEP) {
@@ -938,6 +950,7 @@ std::vector<void (*)(Batches &)> getFunctions()
            computeGamma,
            computeGaussModelExpBasis,
            computeGaussian,
+           computeLogGaussian,
            computeIdentity,
            computeJohnson,
            computeLandau,

roofit/roofit/inc/RooGaussian.h

@@ -58,6 +58,7 @@ class RooGaussian : public RooAbsPdf {
 
   double evaluate() const override;
   void doEval(RooFit::EvalContext &) const override;
+  bool canOptimizeLogarithm() const override { return true; }
   inline bool canComputeBatchWithCuda() const override { return true; }
 
 private:

roofit/roofit/src/RooGaussian.cxx

@@ -27,8 +27,6 @@ Plain Gaussian p.d.f
 
 #include <RooFit/Detail/MathFuncs.h>
 
-#include <vector>
-
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -61,10 +59,22 @@ double RooGaussian::evaluate() const
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of Gaussian distribution.
-void RooGaussian::doEval(RooFit::EvalContext & ctx) const
+void RooGaussian::doEval(RooFit::EvalContext &ctx) const
 {
-  RooBatchCompute::compute(ctx.config(this), RooBatchCompute::Gaussian, ctx.output(),
-          {ctx.at(x), ctx.at(mean), ctx.at(sigma)});
+   if (ctx.config(this).takeLog()) {
+      auto output = ctx.output();
+      if (output.size() == 1) {
+         // If the ouput size is just one, which is common for constraints,
+         // calling into RooBatchCompute is not worth its overhead.
+         output[0] = RooFit::Detail::MathFuncs::logGaussian(ctx.at(x)[0], ctx.at(mean)[0], ctx.at(sigma)[0]);
+      } else {
+         RooBatchCompute::compute(ctx.config(this), RooBatchCompute::LogGaussian, output,
+                                  {ctx.at(x), ctx.at(mean), ctx.at(sigma)});
+      }
+      return;
+   }
+   RooBatchCompute::compute(ctx.config(this), RooBatchCompute::Gaussian, ctx.output(),
+                            {ctx.at(x), ctx.at(mean), ctx.at(sigma)});
 }
 
 ////////////////////////////////////////////////////////////////////////////////

roofit/roofitcore/inc/RooAbsArg.h

@@ -510,8 +510,10 @@ class RooAbsArg : public TNamed, public RooPrintable {
       return false;
    };
 
-   virtual bool canComputeBatchWithCuda() const { return false; }
+   /// Information to expose to the RooFit::Evaluator for optimized evaluation:
+   virtual bool canOptimizeLogarithm() const { return false; }
    virtual bool isReducerNode() const { return false; }
+   virtual bool canComputeBatchWithCuda() const { return false; }
 
    virtual void applyWeightSquared(bool flag);
 

roofit/roofitcore/inc/RooFit/Detail/MathFuncs.h

@@ -83,12 +83,18 @@ double bernstein(double x, double xmin, double xmax, DoubleArray coefs, int nCoe
    return result;
 }
 
-/// @brief Function to evaluate an un-normalized RooGaussian.
-inline double gaussian(double x, double mean, double sigma)
+/// Evaluate the logarithm of an un-normalized Gaussian.
+inline double logGaussian(double x, double mean, double sigma)
 {
    const double arg = x - mean;
    const double sig = sigma;
-   return std::exp(-0.5 * arg * arg / (sig * sig));
+   return -0.5 * arg * arg / (sig * sig);
+}
+
+/// @brief Function to evaluate an un-normalized Gaussian.
+inline double gaussian(double x, double mean, double sigma)
+{
+   return std::exp(logGaussian(x, mean, sigma));
 }
 
 template <typename DoubleArray>

roofit/roofitcore/inc/RooFit/Detail/RooNLLVarNew.h

@@ -53,7 +53,11 @@ class RooNLLVarNew : public RooAbsReal {
 
    void enableBinOffsetting(bool on = true) { _doBinOffset = on; }
 
-   void setSimCount(int simCount) { _simCount = simCount; }
+   void setSimCount(int simCount)
+   {
+      _simCount = simCount;
+      _logSimCount = std::log(static_cast<double>(simCount));
+   }
 
    RooAbsPdf const &pdf() const { return *_pdf; }
    RooAbsReal const &weightVar() const { return *_weightVar; }
@@ -78,6 +82,7 @@ class RooNLLVarNew : public RooAbsReal {
    bool _doOffset = false;
    bool _doBinOffset = false;
    int _simCount = 1;
+   double _logSimCount = 0.;
    std::string _prefix;
    std::vector<double> _binw;
    mutable ROOT::Math::KahanSum<double> _offset{0.}; ///<! Offset as KahanSum to avoid loss of precision

roofit/roofitcore/inc/RooFit/Detail/RooNormalizedPdf.h

@@ -69,6 +69,7 @@ class RooNormalizedPdf : public RooAbsPdf {
       return _pdf->createExpectedEventsFunc(&_normSet);
    }
 
+   bool canOptimizeLogarithm() const override { return true; }
    bool canComputeBatchWithCuda() const override { return true; }
 
    RooAbsPdf const &pdf() const { return *_pdf; }

roofit/roofitcore/inc/RooRealIntegral.h

@@ -103,6 +103,9 @@ class RooRealIntegral : public RooAbsReal {
   double evaluate() const override ;
   bool isValidReal(double value, bool printError=false) const override ;
 
+  void doEval(RooFit::EvalContext &) const override;
+  bool canOptimizeLogarithm() const override { return true; }
+
   bool redirectServersHook(const RooAbsCollection& newServerList,
                  bool mustReplaceAll, bool nameChange, bool isRecursive) override ;
 

roofit/roofitcore/src/RooConstraintSum.cxx

@@ -28,6 +28,7 @@ arguments.
 
 
 #include "RooConstraintSum.h"
+#include "RooBatchCompute.h"
 #include "RooAbsData.h"
 #include "RooAbsReal.h"
 #include "RooAbsPdf.h"
@@ -78,12 +79,17 @@ double RooConstraintSum::evaluate() const
   return sum;
 }
 
+/// Evaluate with the vectorizing CPU backend.
 void RooConstraintSum::doEval(RooFit::EvalContext &ctx) const
 {
    double sum(0);
 
    for (const auto comp : _set1) {
-      sum -= std::log(ctx.at(comp)[0]);
+      // We only need to take the logarithm if the server didn't do it already:
+      if (!ctx.config(comp).takeLog())
+         sum -= std::log(ctx.at(comp)[0]);
+      else
+         sum -= ctx.at(comp)[0];
    }
 
    ctx.output()[0] = sum;