[RF] Migrate more RooFit classes from evaluateSpan to computeBatch

roofit/histfactory/inc/RooStats/HistFactory/ParamHistFunc.h

@@ -105,7 +105,7 @@ class ParamHistFunc : public RooAbsReal {
   Int_t addParamSet( const RooArgList& params );
   static Int_t GetNumBins( const RooArgSet& vars );
   double evaluate() const override;
-  RooSpan<double> evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t size, RooBatchCompute::DataMap&) const override;
 
 private:
   static NumBins getNumBinsPerDim(RooArgSet const& vars);

roofit/histfactory/inc/RooStats/HistFactory/PiecewiseInterpolation.h

@@ -97,7 +97,7 @@ class PiecewiseInterpolation : public RooAbsReal {
   std::vector<int> _interpCode;
 
   Double_t evaluate() const override;
-  RooSpan<double> evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t size, RooBatchCompute::DataMap&) const override;
 
   ClassDefOverride(PiecewiseInterpolation,4) // Sum of RooAbsReal objects
 };

roofit/histfactory/src/ParamHistFunc.cxx

@@ -591,23 +591,19 @@ Double_t ParamHistFunc::evaluate() const
 /// the associated parameters.
 /// \param[in,out] evalData Input/output data for evaluating the ParamHistFunc.
 /// \param[in] normSet Normalisation set passed on to objects that are serving values to us.
-RooSpan<double> ParamHistFunc::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const {
+void ParamHistFunc::computeBatch(cudaStream_t*, double* output, size_t size, RooBatchCompute::DataMap& dataMap) const {
   std::vector<double> oldValues;
   std::vector<RooSpan<const double>> data;
-  std::size_t batchSize = 0;
 
   // Retrieve data for all variables
   for (auto arg : _dataVars) {
     const auto* var = static_cast<RooRealVar*>(arg);
     oldValues.push_back(var->getVal());
-    data.push_back(var->getValues(evalData, normSet));
-    batchSize = std::max(batchSize, data.back().size());
+    data.push_back(dataMap[var]);
   }
 
   // Run computation for each entry in the dataset
-  RooSpan<double> output = evalData.makeBatch(this, batchSize);
-
-  for (std::size_t i = 0; i < batchSize; ++i) {
+  for (std::size_t i = 0; i < size; ++i) {
     for (unsigned int j = 0; j < _dataVars.size(); ++j) {
       assert(i < data[j].size());
       auto& var = static_cast<RooRealVar&>(_dataVars[j]);
@@ -624,8 +620,6 @@ RooSpan<double> ParamHistFunc::evaluateSpan(RooBatchCompute::RunContext& evalDat
     auto& var = static_cast<RooRealVar&>(_dataVars[j]);
     var.setCachedValue(oldValues[j], /*notifyClients=*/false);
   }
-
-  return output;
 }
 
 ////////////////////////////////////////////////////////////////////////////////

roofit/histfactory/src/PiecewiseInterpolation.cxx

@@ -318,15 +318,16 @@ Double_t PiecewiseInterpolation::evaluate() const
 /// Interpolate between input distributions for all values of the observable in `evalData`.
 /// \param[in,out] evalData Struct holding spans pointing to input data. The results of this function will be stored here.
 /// \param[in] normSet Arguments to normalise over.
-RooSpan<double> PiecewiseInterpolation::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const {
-  auto nominal = _nominal->getValues(evalData, normSet);
-  auto sum = evalData.makeBatch(this, nominal.size());
-  std::copy(nominal.begin(), nominal.end(), sum.begin());
+void PiecewiseInterpolation::computeBatch(cudaStream_t*, double* sum, size_t /*size*/, RooBatchCompute::DataMap& dataMap) const {
+  auto nominal = dataMap[&*_nominal];
+  for(unsigned int j=0; j < nominal.size(); ++j) {
+    sum[j] = nominal[j];
+  }
 
   for (unsigned int i=0; i < _paramSet.size(); ++i) {
     const double param = static_cast<RooAbsReal*>(_paramSet.at(i))->getVal();
-    auto low   = static_cast<RooAbsReal*>(_lowSet.at(i) )->getValues(evalData, normSet);
-    auto high  = static_cast<RooAbsReal*>(_highSet.at(i))->getValues(evalData, normSet);
+    auto low   = dataMap[_lowSet.at(i)];
+    auto high  = dataMap[_highSet.at(i)];
     const int icode = _interpCode[i];
 
     switch(icode) {
@@ -436,13 +437,11 @@ RooSpan<double> PiecewiseInterpolation::evaluateSpan(RooBatchCompute::RunContext
   }
 
   if (_positiveDefinite) {
-    for (double& val : sum) {
-      if (val < 0.)
-        val = 0.;
+    for(unsigned int j=0; j < nominal.size(); ++j) {
+      if (sum[j] < 0.)
+        sum[j] = 0.;
     }
   }
-
-  return sum;
 }
 
 ////////////////////////////////////////////////////////////////////////////////

roofit/roofitcore/inc/RooBinWidthFunction.h

@@ -71,7 +71,7 @@ class RooBinWidthFunction : public RooAbsReal {
   bool divideByBinWidth() const { return _divideByBinWidth; }
   const RooHistFunc& histFunc() const { return (*_histFunc); }
   double evaluate() const override;
-  RooSpan<double> evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t size, RooBatchCompute::DataMap&) const override;
 
 private:
   RooTemplateProxy<const RooHistFunc> _histFunc;

roofit/roofitcore/inc/RooHistFunc.h

@@ -86,15 +86,15 @@ class RooHistFunc : public RooAbsReal {
 
 
   Int_t getBin() const;
-  std::vector<Int_t> getBins(RooBatchCompute::RunContext& evalData) const;
+  std::vector<Int_t> getBins(RooBatchCompute::DataMap& dataMap) const;
 
 protected:
 
   Bool_t importWorkspaceHook(RooWorkspace& ws) override ;
   Bool_t areIdentical(const RooDataHist& dh1, const RooDataHist& dh2) ;
 
   Double_t evaluate() const override;
-  RooSpan<double> evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* /*normSet*/) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t size, RooBatchCompute::DataMap&) const override;
   friend class RooAbsCachedReal ;
 
   void ioStreamerPass2() override ;

roofit/roofitcore/inc/RooProduct.h

@@ -77,7 +77,7 @@ class RooProduct : public RooAbsReal {
 
   Double_t calculate(const RooArgList& partIntList) const;
   Double_t evaluate() const override;
-  RooSpan<double> evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, RooBatchCompute::DataMap&) const override;
 
   const char* makeFPName(const char *pfx,const RooArgSet& terms) const ;
   ProdMap* groupProductTerms(const RooArgSet&) const;

roofit/roofitcore/inc/RooRealSumPdf.h

@@ -69,7 +69,6 @@ class RooRealSumPdf : public RooAbsPdf {
   void setCacheAndTrackHints(RooArgSet&) override ;
 
 protected:
-  RooSpan<double> evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const override;
 
   class CacheElem : public RooAbsCacheElement {
   public:

roofit/roofitcore/src/RooBinWidthFunction.cxx

@@ -46,22 +46,18 @@ double RooBinWidthFunction::evaluate() const {
 /// Compute bin index for all values of the observable(s) in `evalData`, and return their volumes or inverse volumes, depending
 /// on the configuration chosen in the constructor.
 /// If a bin is not valid, return a volume of 1.
-RooSpan<double> RooBinWidthFunction::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* /*normSet*/) const {
+void RooBinWidthFunction::computeBatch(cudaStream_t*, double* output, size_t, RooBatchCompute::DataMap& dataMap) const {
   const RooDataHist& dataHist = _histFunc->dataHist();
-  std::vector<Int_t> bins = _histFunc->getBins(evalData);
+  std::vector<Int_t> bins = _histFunc->getBins(dataMap);
   auto volumes = dataHist.binVolumes(0, dataHist.numEntries());
 
-  auto results = evalData.makeBatch(this, bins.size());
-
   if (_divideByBinWidth) {
     for (std::size_t i=0; i < bins.size(); ++i) {
-      results[i] = bins[i] >= 0 ? 1./volumes[bins[i]] : 1.;
+      output[i] = bins[i] >= 0 ? 1./volumes[bins[i]] : 1.;
     }
   } else {
     for (std::size_t i=0; i < bins.size(); ++i) {
-      results[i] = bins[i] >= 0 ? volumes[bins[i]] : 1.;
+      output[i] = bins[i] >= 0 ? volumes[bins[i]] : 1.;
     }
   }
-
-  return results;
 }

roofit/roofitcore/src/RooFitDriver.cxx

@@ -223,7 +223,7 @@ void RooFitDriver::init()
    // Some checks and logging of used architectures
    {
       auto log = [](std::string_view message) {
-         oocxcoutI(static_cast<RooAbsArg *>(nullptr), FastEvaluations) << message << std::endl;
+         oocxcoutI(static_cast<RooAbsArg *>(nullptr), Fitting) << message << std::endl;
       };
 
       if (_batchMode == RooFit::BatchModeOption::Cuda && !RooBatchCompute::dispatchCUDA) {

roofit/roofitcore/src/RooHistFunc.cxx

@@ -202,26 +202,19 @@ Double_t RooHistFunc::evaluate() const
 }
 
 
-////////////////////////////////////////////////////////////////////////////////
-/// Compute value of the HistFunc for every entry in `evalData`.
-/// \param[in,out] evalData Struct with input data. The computation results will be stored here.
-RooSpan<double> RooHistFunc::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* /*normSet*/) const {
+void RooHistFunc::computeBatch(cudaStream_t*, double* output, size_t size, RooBatchCompute::DataMap& dataMap) const {
   std::vector<RooSpan<const double>> inputValues;
-  std::size_t batchSize = 0;
   for (const auto& obs : _depList) {
     auto realObs = dynamic_cast<const RooAbsReal*>(obs);
     if (realObs) {
-      auto inputs = realObs->getValues(evalData, nullptr);
-      batchSize = std::max(batchSize, inputs.size());
+      auto inputs = dataMap[realObs];
       inputValues.push_back(std::move(inputs));
     } else {
       inputValues.emplace_back();
     }
   }
 
-  auto results = evalData.makeBatch(this, batchSize);
-
-  for (std::size_t i = 0; i < batchSize; ++i) {
+  for (std::size_t i = 0; i < size; ++i) {
     bool skip = false;
 
     for (auto j = 0u; j < _histObsList.size(); ++j) {
@@ -236,10 +229,8 @@ RooSpan<double> RooHistFunc::evaluateSpan(RooBatchCompute::RunContext& evalData,
       }
     }
 
-    results[i] = skip ? 0. : _dataHist->weightFast(_histObsList, _intOrder, false, _cdfBoundaries);
+    output[i] = skip ? 0. : _dataHist->weightFast(_histObsList, _intOrder, false, _cdfBoundaries);
   }
-
-  return results;
 }
 
 
@@ -606,12 +597,12 @@ Int_t RooHistFunc::getBin() const {
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute bin numbers corresponding to all coordinates in `evalData`.
 /// \return Vector of bin numbers. If a bin is not in the current range of the observables, return -1.
-std::vector<Int_t> RooHistFunc::getBins(RooBatchCompute::RunContext& evalData) const {
+std::vector<Int_t> RooHistFunc::getBins(RooBatchCompute::DataMap& dataMap) const {
   std::vector<RooSpan<const double>> depData;
   for (const auto dep : _depList) {
     auto real = dynamic_cast<const RooAbsReal*>(dep);
     if (real) {
-      depData.push_back(real->getValues(evalData, nullptr));
+      depData.push_back(dataMap[real]);
     } else {
       depData.emplace_back(nullptr, 0);
     }

roofit/roofitcore/src/RooProduct.cxx

@@ -386,46 +386,32 @@ Double_t RooProduct::evaluate() const
 }
 
 
-////////////////////////////////////////////////////////////////////////////////
-/// Evaluate product of input functions for all points found in `evalData`.
-RooSpan<double> RooProduct::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* normSet) const {
-  RooSpan<double> prod;
-
-  assert(_compRSet.nset() == normSet);
+void RooProduct::computeBatch(cudaStream_t* /*stream*/, double* output, size_t nEvents, RooBatchCompute::DataMap& dataMap) const
+{
+  for (unsigned int i = 0; i < nEvents; ++i) {
+    output[i] = 1.;
+  }
 
   for (const auto item : _compRSet) {
     auto rcomp = static_cast<const RooAbsReal*>(item);
-    auto componentValues = rcomp->getValues(evalData, normSet);
-
-    if (prod.empty()) {
-      prod = evalData.makeBatch(this, componentValues.size());
-      for (auto& val : prod) val = 1.;
-    } else if (prod.size() == 1 && componentValues.size() > 1) {
-      const double val = prod[0];
-      prod = evalData.makeBatch(this, componentValues.size());
-      std::fill(prod.begin(), prod.end(), val);
-    }
-    assert(prod.size() == componentValues.size() || componentValues.size() == 1);
+    auto componentValues = dataMap[rcomp];
 
-    for (unsigned int i = 0; i < prod.size(); ++i) {
-      prod[i] *= componentValues.size() == 1 ? componentValues[0] : componentValues[i];
+    for (unsigned int i = 0; i < nEvents; ++i) {
+      output[i] *= componentValues.size() == 1 ? componentValues[0] : componentValues[i];
     }
   }
 
   for (const auto item : _compCSet) {
     auto ccomp = static_cast<const RooAbsCategory*>(item);
     const int catIndex = ccomp->getCurrentIndex();
 
-    for (unsigned int i = 0; i < prod.size(); ++i) {
-      prod[i] *= catIndex;
+    for (unsigned int i = 0; i < nEvents; ++i) {
+      output[i] *= catIndex;
     }
   }
-
-  return prod;
 }
 
 
-
 ////////////////////////////////////////////////////////////////////////////////
 /// Forward the plot sampling hint from the p.d.f. that defines the observable obs
 

roofit/roofitcore/src/RooRealSumPdf.cxx

@@ -257,7 +257,7 @@ Double_t RooRealSumPdf::evaluate() const
 }
 
 
-void RooRealSumPdf::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, RooBatchCompute::DataMap& dataMap) const {
+void RooRealSumPdf::computeBatch(cudaStream_t* /*stream*/, double* output, size_t nEvents, RooBatchCompute::DataMap& dataMap) const {
 
   // To evaluate this RooRealSumPdf, we have to undo the normalization of the
   // pdf servers by convention. TODO: find a less hacky solution for this,
@@ -283,33 +283,26 @@ void RooRealSumPdf::computeBatch(cudaStream_t* stream, double* output, size_t nE
       }
   }
 
-  RooAbsPdf::computeBatch(stream, output, nEvents, dataMapCopy);
-}
-
-
-////////////////////////////////////////////////////////////////////////////////
-/// Calculate the value for all values of the observable in `evalData`.
-RooSpan<double> RooRealSumPdf::evaluateSpan(RooBatchCompute::RunContext& evalData, const RooArgSet* /*normSet*/) const {
   // Do running sum of coef/func pairs, calculate lastCoef.
-  RooSpan<double> values;
+  for (unsigned int j = 0; j < nEvents; ++j) {
+    output[j] = 0.0;
+  }
+
   double sumCoeff = 0.;
   for (unsigned int i = 0; i < _funcList.size(); ++i) {
     const auto func = static_cast<RooAbsReal*>(&_funcList[i]);
     const auto coef = static_cast<RooAbsReal*>(i < _coefList.size() ? &_coefList[i] : nullptr);
     const double coefVal = coef != nullptr ? coef->getVal() : (1. - sumCoeff);
 
     if (func->isSelectedComp()) {
-      auto funcValues = func->getValues(evalData, nullptr); // No normSet here, because we are summing functions!
-      if (values.empty() || (values.size() == 1 && funcValues.size() > 1)) {
-        const double init = values.empty() ? 0. : values[0];
-        values = evalData.makeBatch(this, funcValues.size());
-        for (unsigned int j = 0; j < values.size(); ++j) {
-          values[j] = init + funcValues[j] * coefVal;
+      auto funcValues = dataMapCopy[func];
+      if(funcValues.size() == 1) {
+        for (unsigned int j = 0; j < nEvents; ++j) {
+          output[j] += funcValues[0] * coefVal;
         }
       } else {
-        assert(values.size() == funcValues.size());
-        for (unsigned int j = 0; j < values.size(); ++j) {
-          values[j] += funcValues[j] * coefVal;
+        for (unsigned int j = 0; j < nEvents; ++j) {
+          output[j] += funcValues[j] * coefVal;
         }
       }
     }
@@ -323,20 +316,33 @@ RooSpan<double> RooRealSumPdf::evaluateSpan(RooBatchCompute::RunContext& evalDat
         _haveWarned = true;
       }
       // Signal that we are in an undefined region by handing back one NaN.
-      values[0] = RooNaNPacker::packFloatIntoNaN(100.f * (coefVal < 0. ? -coefVal : coefVal - 1.));
+      output[0] = RooNaNPacker::packFloatIntoNaN(100.f * (coefVal < 0. ? -coefVal : coefVal - 1.));
     }
 
     sumCoeff += coefVal;
   }
 
   // Introduce floor if so requested
   if (_doFloor || _doFloorGlobal) {
-    for (unsigned int j = 0; j < values.size(); ++j) {
-      values[j] += std::max(0., values[j]);
+    for (unsigned int j = 0; j < nEvents; ++j) {
+      output[j] += std::max(0., output[j]);
     }
   }
 
-  return values;
+  // normalize
+  auto integralSpan = dataMap.at(_norm);
+
+  if(integralSpan.size() == 1) {
+    double oneOverNorm = 1. / integralSpan[0];
+    for (std::size_t i=0; i < nEvents; ++i) {
+      output[i] *= oneOverNorm;
+    }
+  } else {
+    assert(integralSpan.size() == nEvents);
+    for (std::size_t i=0; i < nEvents; ++i) {
+      output[i] = normalizeWithNaNPacking(output[i], integralSpan[i]);
+    }
+  }
 }