[RF] Change AddPdf BatchMode coef updating to use values from dataMap

guitargeek · guitargeek · commit 2f75c418f600 · 2022-08-11T12:53:26.000+02:00
So far, the raw coefficient values for the RooAddPdf in BatchMode were
still obstained with the `getVal()` interface. This is redundant,
because as value servers the coefficients are also evaluated by the
RooFitDriver. This caused slowdowns in more complicated fits, like the
ATLAS Higgs combination.

This commit solves that problem by not evaluating the `_coefList`
directly with `getVal()` in `updateCoefficients()`. Instead, the raw
coeficients will be passed to the function via the already existing
`_coefCache`, which is also used as an output parameter.

In the BatchMode, we can then just fill the initial `_coefCache` with
the values from the data map.
diff --git a/roofit/roofitcore/inc/RooAddPdf.h b/roofit/roofitcore/inc/RooAddPdf.h
@@ -106,7 +106,7 @@ class RooAddPdf : public RooAbsPdf {
 
   mutable RooObjCacheManager _projCacheMgr ;  //! Manager of cache with coefficient projections and transformations
   AddCacheElem* getProjCache(const RooArgSet* nset, const RooArgSet* iset=nullptr, const char* rangeName=nullptr) const ;
-  void updateCoefficients(AddCacheElem& cache, const RooArgSet* nset) const ;
+  void updateCoefficients(AddCacheElem& cache, const RooArgSet* nset, bool syncCoefValues=true) const ;
 
 
   friend class RooAddGenContext ;
diff --git a/roofit/roofitcore/src/RooAddHelpers.cxx b/roofit/roofitcore/src/RooAddHelpers.cxx
@@ -17,6 +17,10 @@
 #include <RooRealIntegral.h>
 #include <RooRealVar.h>
 
+////////////////////////////////////////////////////////////////////////////////
+/// Create a RooAddPdf cache element for a given normalization set and
+/// projection configuration.
+
 AddCacheElem::AddCacheElem(RooAbsPdf const &addPdf, RooArgList const &pdfList, RooArgList const &coefList,
                            const RooArgSet *nset, const RooArgSet *iset, const char *rangeName, bool projectCoefs,
                            RooArgSet const &refCoefNorm, TNamed const *refCoefRangeName, int verboseEval)
@@ -110,12 +114,12 @@ AddCacheElem::AddCacheElem(RooAbsPdf const &addPdf, RooArgList const &pdfList, R
                                   << (refCoefRangeName ? RooNameReg::str(refCoefRangeName) : "") << std::endl;
 
       // Recalculate projection integrals of PDFs
-      for (auto *thePdf : static_range_cast<const RooAbsPdf *>(pdfList)) {
+      for (auto *pdf : static_range_cast<const RooAbsPdf *>(pdfList)) {
 
          // Calculate projection integral
          std::unique_ptr<RooAbsReal> pdfProj;
          if (!nset2.equals(refCoefNorm)) {
-            pdfProj = std::unique_ptr<RooAbsReal>{thePdf->createIntegral(nset2, refCoefNorm, addPdf.normRange())};
+            pdfProj = std::unique_ptr<RooAbsReal>{pdf->createIntegral(nset2, refCoefNorm, addPdf.normRange())};
             pdfProj->setOperMode(addPdf.operMode());
             oocxcoutD(&addPdf, Caching) << addPdf.ClassName() << "(" << addPdf.GetName() << ")::getPC nset2(" << nset2
                                         << ")!=_refCoefNorm(" << refCoefNorm << ") --> pdfProj = " << pdfProj->GetName()
@@ -128,11 +132,11 @@ AddCacheElem::AddCacheElem(RooAbsPdf const &addPdf, RooArgList const &pdfList, R
 
          // Calculation optional supplemental normalization term
          RooArgSet supNormSet(refCoefNorm);
-         auto deps = std::unique_ptr<RooArgSet>{thePdf->getParameters(RooArgSet())};
+         auto deps = std::unique_ptr<RooArgSet>{pdf->getParameters(RooArgSet())};
          supNormSet.remove(*deps, true, true);
 
          std::unique_ptr<RooAbsReal> snorm;
-         auto name = std::string(addPdf.GetName()) + "_" + thePdf->GetName() + "_ProjSupNorm";
+         auto name = std::string(addPdf.GetName()) + "_" + pdf->GetName() + "_ProjSupNorm";
          if (!supNormSet.empty() && !nset2.equals(refCoefNorm)) {
             snorm = std::make_unique<RooRealIntegral>(name.c_str(), "Projection Supplemental normalization integral",
                                                       RooRealConstant::value(1.0), supNormSet);
@@ -155,7 +159,7 @@ AddCacheElem::AddCacheElem(RooAbsPdf const &addPdf, RooArgList const &pdfList, R
 
          // ----------
          RooArgSet tmpObs;
-         thePdf->getObservables(&refCoefNorm, tmpObs);
+         pdf->getObservables(&refCoefNorm, tmpObs);
          bool allIdent = true;
          for (auto *rvarg : dynamic_range_cast<RooRealVar *>(tmpObs)) {
             if (rvarg) {
@@ -170,9 +174,8 @@ AddCacheElem::AddCacheElem(RooAbsPdf const &addPdf, RooArgList const &pdfList, R
          if (refCoefRangeName && !refCoefNorm.empty() && !allIdent) {
 
             RooArgSet tmp;
-            thePdf->getObservables(&refCoefNorm, tmp);
-            rangeProj1 =
-               std::unique_ptr<RooAbsReal>{thePdf->createIntegral(tmp, tmp, RooNameReg::str(refCoefRangeName))};
+            pdf->getObservables(&refCoefNorm, tmp);
+            rangeProj1 = std::unique_ptr<RooAbsReal>{pdf->createIntegral(tmp, tmp, RooNameReg::str(refCoefRangeName))};
 
             // rangeProj1->setOperMode(operMode()) ;
             oocxcoutD(&addPdf, Caching) << " " << addPdf.ClassName() << "::syncCoefProjList(" << addPdf.GetName()
@@ -187,14 +190,14 @@ AddCacheElem::AddCacheElem(RooAbsPdf const &addPdf, RooArgList const &pdfList, R
                                      << " nset = " << (nset ? *nset : RooArgSet()) << std::endl;
          if (rangeName && !refCoefNorm.empty()) {
 
-            rangeProj2 = std::unique_ptr<RooAbsReal>{thePdf->createIntegral(refCoefNorm, refCoefNorm, rangeName)};
+            rangeProj2 = std::unique_ptr<RooAbsReal>{pdf->createIntegral(refCoefNorm, refCoefNorm, rangeName)};
             // rangeProj2->setOperMode(operMode()) ;
 
          } else if (addPdf.normRange()) {
 
             RooArgSet tmp;
-            thePdf->getObservables(&refCoefNorm, tmp);
-            rangeProj2 = std::unique_ptr<RooAbsReal>{thePdf->createIntegral(tmp, tmp, addPdf.normRange())};
+            pdf->getObservables(&refCoefNorm, tmp);
+            rangeProj2 = std::unique_ptr<RooAbsReal>{pdf->createIntegral(tmp, tmp, addPdf.normRange())};
             oocxcoutD(&addPdf, Caching) << " " << addPdf.ClassName() << "::syncCoefProjList(" << addPdf.GetName()
                                         << ") R2 = " << rangeProj2->GetName() << std::endl;
          }
@@ -230,15 +233,20 @@ RooArgList AddCacheElem::containedArgs(Action)
    return allNodes;
 }
 
-void RooAddHelpers::updateCoefficients(RooAbsPdf const &addPdf, RooArgList const &pdfList, RooArgList const &coefList,
-                                       AddCacheElem &cache, const RooArgSet *nset, bool projectCoefs,
-                                       RooArgSet const &refCoefNorm, bool allExtendable, std::vector<double> &coefCache,
-                                       int &coefErrCount)
+////////////////////////////////////////////////////////////////////////////////
+/// Update the RooAddPdf coefficients for a given normalization set and
+/// projection configuration. The `coefCache` argument should have the same
+/// size as `pdfList`. It needs to be initialized with the raw values of the
+/// coefficients, as obtained from the `_coefList` proxy in the RooAddPdf. If
+/// the last coefficient is not given, the initial value of the last element of
+/// `_coefCache` does not matter. After this function, the `_coefCache` will be
+/// filled with the correctly scaled coefficients for each pdf.
+
+void RooAddHelpers::updateCoefficients(RooAbsPdf const &addPdf, std::vector<double> &coefCache,
+                                       RooArgList const &pdfList, bool haveLastCoef, AddCacheElem &cache,
+                                       const RooArgSet *nset, bool projectCoefs, RooArgSet const &refCoefNorm,
+                                       bool allExtendable, int &coefErrCount)
 {
-   bool haveLastCoef = pdfList.size() == coefList.size();
-
-   coefCache.resize(haveLastCoef ? coefList.size() : pdfList.size(), 0.);
-
    // Straight coefficients
    if (allExtendable) {
 
@@ -265,38 +273,26 @@ void RooAddHelpers::updateCoefficients(RooAbsPdf const &addPdf, RooArgList const
       if (haveLastCoef) {
 
          // coef[i] = coef[i] / SUM(coef)
-         double coefSum(0);
-         std::size_t i = 0;
-         for (auto coefArg : coefList) {
-            auto coef = static_cast<RooAbsReal *>(coefArg);
-            coefCache[i] = coef->getVal(nset);
-            coefSum += coefCache[i++];
-         }
+         double coefSum = std::accumulate(coefCache.begin(), coefCache.end(), 0.0);
          if (coefSum == 0.) {
             oocoutW(&addPdf, Eval) << addPdf.ClassName() << "::updateCoefCache(" << addPdf.GetName()
                                    << ") WARNING: sum of coefficients is zero 0" << std::endl;
          } else {
             const double invCoefSum = 1. / coefSum;
-            for (std::size_t j = 0; j < coefList.size(); j++) {
+            for (std::size_t j = 0; j < coefCache.size(); j++) {
                coefCache[j] *= invCoefSum;
             }
          }
       } else {
 
          // coef[i] = coef[i] ; coef[n] = 1-SUM(coef[0...n-1])
-         double lastCoef(1);
-         std::size_t i = 0;
-         for (auto coefArg : coefList) {
-            auto coef = static_cast<RooAbsReal *>(coefArg);
-            coefCache[i] = coef->getVal(nset);
-            lastCoef -= coefCache[i++];
-         }
-         coefCache[coefList.size()] = lastCoef;
+         double lastCoef = 1.0 - std::accumulate(coefCache.begin(), coefCache.end() - 1, 0.0);
+         coefCache.back() = lastCoef;
 
          // Treat coefficient degeneration
          const float coefDegen = lastCoef < 0. ? -lastCoef : (lastCoef > 1. ? lastCoef - 1. : 0.);
          if (coefDegen > 1.E-5) {
-            coefCache[coefList.size()] = RooNaNPacker::packFloatIntoNaN(100.f * coefDegen);
+            coefCache.back() = RooNaNPacker::packFloatIntoNaN(100.f * coefDegen);
 
             std::stringstream msg;
             if (coefErrCount-- > 0) {
diff --git a/roofit/roofitcore/src/RooAddHelpers.h b/roofit/roofitcore/src/RooAddHelpers.h
@@ -61,10 +61,9 @@ class AddCacheElem : public RooAbsCacheElement {
 
 class RooAddHelpers {
 public:
-   static void updateCoefficients(RooAbsPdf const &addPdf, RooArgList const &pdfList, RooArgList const &coefList,
-                                  AddCacheElem &cache, const RooArgSet *nset, bool projectCoefs,
-                                  RooArgSet const &refCoefNorm, bool allExtendable, std::vector<double> &coefCache,
-                                  int &coefErrCount);
+   static void updateCoefficients(RooAbsPdf const &addPdf, std::vector<double> &coefCache, RooArgList const &pdfList,
+                                  bool haveLastCoef, AddCacheElem &cache, const RooArgSet *nset, bool projectCoefs,
+                                  RooArgSet const &refCoefNorm, bool allExtendable, int &coefErrCount);
 };
 
 #endif
diff --git a/roofit/roofitcore/src/RooAddModel.cxx b/roofit/roofitcore/src/RooAddModel.cxx
@@ -64,7 +64,6 @@ RooAddModel::RooAddModel() :
   _projCacheMgr(this,10),
   _intCacheMgr(this,10)
 {
-  _coefCache.resize(10);
   _coefErrCount = _errorCount ;
 }
 
@@ -149,7 +148,6 @@ RooAddModel::RooAddModel(const char *name, const char *title, const RooArgList&
       _haveLastCoef = true;
    }
 
-   _coefCache.resize(_pdfList.getSize());
    _coefErrCount = _errorCount;
 
    if (ownPdfList) {
@@ -173,7 +171,6 @@ RooAddModel::RooAddModel(const RooAddModel& other, const char* name) :
   _haveLastCoef(other._haveLastCoef),
   _allExtendable(other._allExtendable)
 {
-  _coefCache.resize(_pdfList.getSize());
   _coefErrCount = _errorCount ;
 }
 
@@ -340,8 +337,12 @@ AddCacheElem* RooAddModel::getProjCache(const RooArgSet* nset, const RooArgSet*
 
 void RooAddModel::updateCoefficients(AddCacheElem& cache, const RooArgSet* nset) const
 {
-  RooAddHelpers::updateCoefficients(*this, _pdfList, _coefList, cache, nset, _projectCoefs,
-                                    _refCoefNorm, _allExtendable, _coefCache, _coefErrCount);
+  _coefCache.resize(_pdfList.getSize());
+  for(std::size_t i = 0; i < _coefList.size(); ++i) {
+    _coefCache[i] = static_cast<RooAbsReal const&>(_coefList[i]).getVal(nset);
+  }
+  RooAddHelpers::updateCoefficients(*this, _coefCache, _pdfList, _haveLastCoef, cache, nset, _projectCoefs,
+                                    _refCoefNorm, _allExtendable, _coefErrCount);
 }
 
 
diff --git a/roofit/roofitcore/src/RooAddPdf.cxx b/roofit/roofitcore/src/RooAddPdf.cxx
@@ -121,8 +121,6 @@ void RooAddPdf::finalizeConstruction() {
     }
     seen.insert(elem);
   }
-
-  _coefCache.resize(_pdfList.size());
 }
 
 
@@ -378,11 +376,25 @@ AddCacheElem* RooAddPdf::getProjCache(const RooArgSet* nset, const RooArgSet* is
 /// fraction, normalize fractions obtained from extended ML terms to unity, and
 /// multiply the various range and dimensional corrections needed in the
 /// current use context.
-
-void RooAddPdf::updateCoefficients(AddCacheElem& cache, const RooArgSet* nset) const
+///
+/// param[in] cache The cache element for the given normalization set that
+///                 stores the supplementary normalization values and
+///                 projection-related objects.
+/// param[in] nset The set of variables to normalize over.
+/// param[in] syncCoefValues If the initial values of the coefficients still
+///                          need to be copied from the `_coefList` elements to
+///                          the `_coefCache`. True by default.
+
+void RooAddPdf::updateCoefficients(AddCacheElem& cache, const RooArgSet* nset, bool syncCoefValues) const
 {
-  RooAddHelpers::updateCoefficients(*this, _pdfList, _coefList, cache, nset, _projectCoefs,
-                                    _refCoefNorm, _allExtendable, _coefCache, _coefErrCount);
+  _coefCache.resize(_pdfList.size());
+  if(syncCoefValues) {
+    for(std::size_t i = 0; i < _coefList.size(); ++i) {
+      _coefCache[i] = static_cast<RooAbsReal const&>(_coefList[i]).getVal(nset);
+    }
+  }
+  RooAddHelpers::updateCoefficients(*this, _coefCache, _pdfList, _haveLastCoef, cache, nset, _projectCoefs,
+                                    _refCoefNorm, _allExtendable, _coefErrCount);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -435,7 +447,6 @@ std::pair<const RooArgSet*, AddCacheElem*> RooAddPdf::getNormAndCache(const RooA
 
 
   AddCacheElem* cache = getProjCache(nset) ;
-  updateCoefficients(*cache, nset);
 
   return {nset, cache};
 }
@@ -449,6 +460,7 @@ double RooAddPdf::getValV(const RooArgSet* normSet) const
   auto normAndCache = getNormAndCache(normSet);
   const RooArgSet* nset = normAndCache.first;
   AddCacheElem* cache = normAndCache.second;
+  updateCoefficients(*cache, nset);
 
   // Process change in last data set used
   bool nsetChanged(false) ;
@@ -481,9 +493,32 @@ double RooAddPdf::getValV(const RooArgSet* normSet) const
 /// Compute addition of PDFs in batches.
 void RooAddPdf::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, RooFit::Detail::DataMap const& dataMap) const
 {
+  _coefCache.resize(_pdfList.size());
+  for(std::size_t i = 0; i < _coefList.size(); ++i) {
+    auto coefVals = dataMap.at(&_coefList[i]);
+    // We don't support per-event coefficients in this function. If the CPU
+    // mode is used, we can just fall back to the RooAbsReal implementation.
+    // With CUDA, we can't do that because the inputs might be on the device.
+    // That's why we throw an exception then.
+    if(coefVals.size() > 1) {
+      if(stream) {
+        throw std::runtime_error("The RooAddPdf doesn't support per-event coefficients in CUDA mode yet!");
+      }
+      RooAbsReal::computeBatch(stream, output, nEvents, dataMap);
+      return;
+    }
+    _coefCache[i] = coefVals[0];
+  }
+
   RooBatchCompute::VarVector pdfs;
   RooBatchCompute::ArgVector coefs;
-  AddCacheElem* cache = getNormAndCache(nullptr).second;
+  auto normAndCache = getNormAndCache(nullptr);
+  const RooArgSet* nset = normAndCache.first;
+  AddCacheElem* cache = normAndCache.second;
+  // We don't sync the coefficient values from the _coefList to the _coefCache
+  // because we have already done it using the dataMap.
+  updateCoefficients(*cache, nset, /*syncCoefValues=*/false);
+
   for (unsigned int pdfNo = 0; pdfNo < _pdfList.size(); ++pdfNo)
   {
     auto pdf = static_cast<RooAbsPdf*>(&_pdfList[pdfNo]);

Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,6 @@ RooAddModel::RooAddModel() :`
`64`	`64`	`_projCacheMgr(this,10),`
`65`	`65`	`_intCacheMgr(this,10)`
`66`	`66`	`{`
`67`		`- _coefCache.resize(10);`
`68`	`67`	`_coefErrCount = _errorCount ;`
`69`	`68`	`}`
`70`	`69`
`@@ -149,7 +148,6 @@ RooAddModel::RooAddModel(const char name, const char title, const RooArgList&`
`149`	`148`	`_haveLastCoef = true;`
`150`	`149`	`}`
`151`	`150`
`152`		`- _coefCache.resize(_pdfList.getSize());`
`153`	`151`	`_coefErrCount = _errorCount;`
`154`	`152`
`155`	`153`	`if (ownPdfList) {`
`@@ -173,7 +171,6 @@ RooAddModel::RooAddModel(const RooAddModel& other, const char* name) :`
`173`	`171`	`_haveLastCoef(other._haveLastCoef),`
`174`	`172`	`_allExtendable(other._allExtendable)`
`175`	`173`	`{`
`176`		`- _coefCache.resize(_pdfList.getSize());`
`177`	`174`	`_coefErrCount = _errorCount ;`
`178`	`175`	`}`
`179`	`176`
`@@ -340,8 +337,12 @@ AddCacheElem* RooAddModel::getProjCache(const RooArgSet* nset, const RooArgSet*`
`340`	`337`
`341`	`338`	`void RooAddModel::updateCoefficients(AddCacheElem& cache, const RooArgSet* nset) const`
`342`	`339`	`{`
`343`		`- RooAddHelpers::updateCoefficients(*this, _pdfList, _coefList, cache, nset, _projectCoefs,`
`344`		`- _refCoefNorm, _allExtendable, _coefCache, _coefErrCount);`
	`340`	`+ _coefCache.resize(_pdfList.getSize());`
	`341`	`+ for(std::size_t i = 0; i < _coefList.size(); ++i) {`
	`342`	`+ _coefCache[i] = static_cast<RooAbsReal const&>(_coefList[i]).getVal(nset);`
	`343`	`+ }`
	`344`	`+ RooAddHelpers::updateCoefficients(*this, _coefCache, _pdfList, _haveLastCoef, cache, nset, _projectCoefs,`
	`345`	`+ _refCoefNorm, _allExtendable, _coefErrCount);`
`345`	`346`	`}`
`346`	`347`
`347`	`348`