[RF] Bugfix (cuda computations always run in default stream)

root-project · Oct 24, 2021 · eb20234 · eb20234
1 parent 977be07
commit eb20234
Show file tree

Hide file tree

Showing 57 changed files with 124 additions and 100 deletions.
diff --git a/roofit/batchcompute/inc/rbc.h b/roofit/batchcompute/inc/rbc.h
@@ -64,8 +64,8 @@ class RbcInterface {
   public:
     virtual ~RbcInterface() = default;
     virtual void   init() { throw std::bad_function_call(); }
-    virtual void   compute(Computer, RestrictArr, size_t, const DataMap&, const VarVector&, const ArgVector& ={}) = 0;
-    virtual double sumReduce(InputArr, size_t) = 0;
+    virtual void   compute(cudaStream_t*, Computer, RestrictArr, size_t, const DataMap&, const VarVector&, const ArgVector& ={}) = 0;
+    virtual double sumReduce(cudaStream_t*, InputArr, size_t) = 0;
 
     //cuda functions that need to be interfaced
     virtual void* cudaMalloc(size_t)                { throw std::bad_function_call(); }

diff --git a/roofit/batchcompute/src/RooBatchCompute.cu b/roofit/batchcompute/src/RooBatchCompute.cu
@@ -51,14 +51,14 @@ class RbcClass : public RbcInterface {
     \param varData A std::map containing the values of the variables involved in the computation.
     \param vars A std::vector containing pointers to the variables involved in the computation.
     \param extraArgs An optional std::vector containing extra double values that may participate in the computation. **/ 
-    void compute(Computer computer, RestrictArr output, size_t nEvents, const DataMap& varData, const VarVector& vars, const ArgVector& extraArgs) override  
+    void compute(cudaStream_t* stream, Computer computer, RestrictArr output, size_t nEvents, const DataMap& varData, const VarVector& vars, const ArgVector& extraArgs) override
     {
       Batches batches(output, nEvents, varData, vars, extraArgs);
-      _computeFunctions[computer]<<<128,512>>>(batches);
+      _computeFunctions[computer]<<<128,512,0,*stream>>>(batches);
     }
     /// Return the sum of an input array
-    double sumReduce(InputArr input, size_t n) override {
-      return thrust::reduce(thrust::device, input, input+n, 0.0);
+    double sumReduce(cudaStream_t* stream, InputArr input, size_t n) override {
+      return thrust::reduce(thrust::cuda::par.on(*stream), input, input+n, 0.0);
     }
 
     //cuda functions

diff --git a/roofit/batchcompute/src/RooBatchCompute.cxx b/roofit/batchcompute/src/RooBatchCompute.cxx
@@ -39,7 +39,7 @@ class RbcClass : public RbcInterface {
     \param varData A std::map containing the values of the variables involved in the computation.
     \param vars A std::vector containing pointers to the variables involved in the computation.
     \param extraArgs An optional std::vector containing extra double values that may participate in the computation. **/ 
-    void compute(Computer computer, RestrictArr output, size_t nEvents, const DataMap& varData, const VarVector& vars, const ArgVector& extraArgs) override  
+    void compute(cudaStream_t*, Computer computer, RestrictArr output, size_t nEvents, const DataMap& varData, const VarVector& vars, const ArgVector& extraArgs) override
     {
       double buffer[maxParams][bufferSize];
       ROOT::Internal::TExecutor ex;
@@ -72,7 +72,7 @@ class RbcClass : public RbcInterface {
       ex.Map(task, batchesArr);
     }
     /// Return the sum of an input array
-    double sumReduce(InputArr input, size_t n) override
+    double sumReduce(cudaStream_t*, InputArr input, size_t n) override
     {
       long double sum=0.0;
       for (size_t i=0; i<n; i++) sum += input[i];

diff --git a/roofit/roofit/inc/RooArgusBG.h b/roofit/roofit/inc/RooArgusBG.h
@@ -40,7 +40,7 @@ class RooArgusBG : public RooAbsPdf {
   RooRealProxy p ;
 
   Double_t evaluate() const ;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t size, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t size, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 

diff --git a/roofit/roofit/inc/RooBernstein.h b/roofit/roofit/inc/RooBernstein.h
@@ -45,7 +45,7 @@ class RooBernstein : public RooAbsPdf {
   std::string _refRangeName ;
 
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
   ClassDef(RooBernstein,2) // Bernstein polynomial PDF

diff --git a/roofit/roofit/inc/RooBifurGauss.h b/roofit/roofit/inc/RooBifurGauss.h
@@ -43,7 +43,7 @@ class RooBifurGauss : public RooAbsPdf {
   RooRealProxy sigmaR;
 
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooBreitWigner.h b/roofit/roofit/inc/RooBreitWigner.h
@@ -41,7 +41,7 @@ class RooBreitWigner : public RooAbsPdf {
   RooRealProxy width ;
 
   Double_t evaluate() const ;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 //   void initGenerator();

diff --git a/roofit/roofit/inc/RooBukinPdf.h b/roofit/roofit/inc/RooBukinPdf.h
@@ -48,7 +48,7 @@ class RooBukinPdf : public RooAbsPdf {
   RooRealProxy rho1;
   RooRealProxy rho2;
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooCBShape.h b/roofit/roofit/inc/RooCBShape.h
@@ -51,7 +51,7 @@ class RooCBShape : public RooAbsPdf {
   RooRealProxy n;
 
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 

diff --git a/roofit/roofit/inc/RooChebychev.h b/roofit/roofit/inc/RooChebychev.h
@@ -45,7 +45,7 @@ class RooChebychev : public RooAbsPdf {
   mutable TNamed* _refRangeName ; 
 
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
   Double_t evalAnaInt(const Double_t a, const Double_t b) const;

diff --git a/roofit/roofit/inc/RooChiSquarePdf.h b/roofit/roofit/inc/RooChiSquarePdf.h
@@ -41,7 +41,7 @@ class RooChiSquarePdf : public RooAbsPdf {
   RooRealProxy _ndof;
 
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
   ClassDef(RooChiSquarePdf,1) // Chi Square distribution (eg. the PDF )

diff --git a/roofit/roofit/inc/RooDstD0BG.h b/roofit/roofit/inc/RooDstD0BG.h
@@ -45,7 +45,7 @@ class RooDstD0BG : public RooAbsPdf {
   RooRealProxy C,A,B ;
 
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooExponential.h b/roofit/roofit/inc/RooExponential.h
@@ -39,7 +39,7 @@ class RooExponential : public RooAbsPdf {
   RooRealProxy c;
 
   Double_t evaluate() const override;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const override;
   inline bool canComputeBatchWithCuda() const override { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooGamma.h b/roofit/roofit/inc/RooGamma.h
@@ -40,7 +40,7 @@ class RooGamma : public RooAbsPdf {
   RooRealProxy mu ;
 
   Double_t evaluate() const ;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooGaussian.h b/roofit/roofit/inc/RooGaussian.h
@@ -45,7 +45,7 @@ class RooGaussian : public RooAbsPdf {
   RooRealProxy sigma ;
 
   Double_t evaluate() const override;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t size, rbc::DataMap& dataMap) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t size, rbc::DataMap&) const override;
   inline bool canComputeBatchWithCuda() const override { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooJohnson.h b/roofit/roofit/inc/RooJohnson.h
@@ -57,7 +57,7 @@ class RooJohnson final : public RooAbsPdf {
   double _massThreshold{-1.E300};
 
   Double_t evaluate() const override;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const override;
   inline bool canComputeBatchWithCuda() const override { return true; }
 
   ClassDefOverride(RooJohnson,1)

diff --git a/roofit/roofit/inc/RooLandau.h b/roofit/roofit/inc/RooLandau.h
@@ -39,7 +39,7 @@ class RooLandau : public RooAbsPdf {
   RooRealProxy sigma ;
 
   Double_t evaluate() const ;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooLognormal.h b/roofit/roofit/inc/RooLognormal.h
@@ -38,7 +38,7 @@ class RooLognormal : public RooAbsPdf {
   RooRealProxy k ;
 
   Double_t evaluate() const ;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooNovosibirsk.h b/roofit/roofit/inc/RooNovosibirsk.h
@@ -45,7 +45,7 @@ class RooNovosibirsk : public RooAbsPdf {
 
 protected:
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 private:

diff --git a/roofit/roofit/inc/RooPoisson.h b/roofit/roofit/inc/RooPoisson.h
@@ -43,7 +43,7 @@ class RooPoisson : public RooAbsPdf {
   Bool_t  _protectNegative{true};
 
   Double_t evaluate() const override;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const override;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const override;
   inline bool canComputeBatchWithCuda() const override { return true; }
 
   ClassDefOverride(RooPoisson,3) // A Poisson PDF

diff --git a/roofit/roofit/inc/RooPolynomial.h b/roofit/roofit/inc/RooPolynomial.h
@@ -50,7 +50,7 @@ class RooPolynomial : public RooAbsPdf {
 
   /// Evaluation
   Double_t evaluate() const;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
   ClassDef(RooPolynomial,1) // Polynomial PDF

diff --git a/roofit/roofit/inc/RooVoigtian.h b/roofit/roofit/inc/RooVoigtian.h
@@ -47,7 +47,7 @@ class RooVoigtian : public RooAbsPdf {
   RooRealProxy sigma ;
 
   Double_t evaluate() const ;
-  void computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const;
+  void computeBatch(cudaStream_t*, double* output, size_t nEvents, rbc::DataMap&) const;
   inline bool canComputeBatchWithCuda() const { return true; }
 
 private:

diff --git a/roofit/roofit/src/RooArgusBG.cxx b/roofit/roofit/src/RooArgusBG.cxx
@@ -88,9 +88,10 @@ Double_t RooArgusBG::evaluate() const {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-void RooArgusBG::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooArgusBG::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::ArgusBG, output, nEvents, dataMap, {&*m,&*m0,&*c,&*p,&*_norm});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::ArgusBG, output, nEvents, dataMap, {&*m,&*m0,&*c,&*p,&*_norm});
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/roofit/roofit/src/RooBernstein.cxx b/roofit/roofit/src/RooBernstein.cxx
@@ -142,7 +142,7 @@ Double_t RooBernstein::evaluate() const
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of Bernstein distribution.  
-void RooBernstein::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooBernstein::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
   const int nCoef = _coefList.size();
   std::vector<double> extraArgs(nCoef+2);
@@ -151,7 +151,8 @@ void RooBernstein::computeBatch(rbc::RbcInterface* dispatch, double* output, siz
   extraArgs[nCoef] = _x.min();
   extraArgs[nCoef+1] = _x.max();
 
-  dispatch->compute(rbc::Bernstein, output, nEvents, dataMap, {&*_x, &*_norm}, extraArgs);
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::Bernstein, output, nEvents, dataMap, {&*_x, &*_norm}, extraArgs);
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/roofit/roofit/src/RooBifurGauss.cxx b/roofit/roofit/src/RooBifurGauss.cxx
@@ -80,9 +80,10 @@ Double_t RooBifurGauss::evaluate() const {
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of BifurGauss distribution.  
-void RooBifurGauss::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooBifurGauss::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::BifurGauss, output, nEvents, dataMap, {&*x,&*mean,&*sigmaL,&*sigmaR,&*_norm});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::BifurGauss, output, nEvents, dataMap, {&*x,&*mean,&*sigmaL,&*sigmaR,&*_norm});
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/roofit/roofit/src/RooBreitWigner.cxx b/roofit/roofit/src/RooBreitWigner.cxx
@@ -66,9 +66,10 @@ Double_t RooBreitWigner::evaluate() const
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of BreitWigner distribution.  
-void RooBreitWigner::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooBreitWigner::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::BreitWigner, output, nEvents, dataMap, {&*x,&*mean,&*width,&*_norm});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::BreitWigner, output, nEvents, dataMap, {&*x,&*mean,&*width,&*_norm});
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/roofit/roofit/src/RooBukinPdf.cxx b/roofit/roofit/src/RooBukinPdf.cxx
@@ -142,7 +142,8 @@ Double_t RooBukinPdf::evaluate() const
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of Bukin distribution.  
-void RooBukinPdf::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooBukinPdf::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::Bukin, output, nEvents, dataMap, {&*x,&*Xp,&*sigp,&*xi,&*rho1,&*rho2,&*_norm});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::Bukin, output, nEvents, dataMap, {&*x,&*Xp,&*sigp,&*xi,&*rho1,&*rho2,&*_norm});
 }
diff --git a/roofit/roofit/src/RooCBShape.cxx b/roofit/roofit/src/RooCBShape.cxx
@@ -92,9 +92,10 @@ Double_t RooCBShape::evaluate() const {
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of Crystal ball Shape distribution.
-void RooCBShape::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooCBShape::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::CBShape, output, nEvents, dataMap, {&*m,&*m0,&*sigma,&*alpha,&*n,&*_norm});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::CBShape, output, nEvents, dataMap, {&*m,&*m0,&*sigma,&*alpha,&*n,&*_norm});
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/roofit/roofit/src/RooChebychev.cxx b/roofit/roofit/src/RooChebychev.cxx
@@ -184,14 +184,15 @@ Double_t RooChebychev::evaluate() const
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of Chebychev.  
-void RooChebychev::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooChebychev::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
   rbc::ArgVector extraArgs;
   for (auto* coef:_coefList)
     extraArgs.push_back( static_cast<const RooAbsReal*>(coef)->getVal() );
   extraArgs.push_back( _x.min(_refRangeName?_refRangeName->GetName() : nullptr) );
   extraArgs.push_back( _x.max(_refRangeName?_refRangeName->GetName() : nullptr) );
-  dispatch->compute(rbc::Chebychev, output, nEvents, dataMap, {&*_x,&*_norm}, extraArgs);
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::Chebychev, output, nEvents, dataMap, {&*_x,&*_norm}, extraArgs);
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/roofit/roofit/src/RooChiSquarePdf.cxx b/roofit/roofit/src/RooChiSquarePdf.cxx
@@ -64,9 +64,10 @@ Double_t RooChiSquarePdf::evaluate() const
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of ChiSquare distribution.  
-void RooChiSquarePdf::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooChiSquarePdf::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::ChiSquare, output, nEvents, dataMap, {&*_x,&*_norm}, {_ndof});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::ChiSquare, output, nEvents, dataMap, {&*_x,&*_norm}, {_ndof});
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/roofit/roofit/src/RooDstD0BG.cxx b/roofit/roofit/src/RooDstD0BG.cxx
@@ -81,9 +81,10 @@ Double_t RooDstD0BG::evaluate() const
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of D*-D0 mass difference distribution. 
-void RooDstD0BG::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooDstD0BG::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::DstD0BG, output, nEvents, dataMap, {&*dm,&*dm0,&*C,&*A,&*B,&*_norm});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::DstD0BG, output, nEvents, dataMap, {&*dm,&*dm0,&*C,&*A,&*B,&*_norm});
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/roofit/roofit/src/RooExponential.cxx b/roofit/roofit/src/RooExponential.cxx
@@ -62,9 +62,10 @@ Double_t RooExponential::evaluate() const{
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of Exponential distribution.
-void RooExponential::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooExponential::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::Exponential, output, nEvents, dataMap, {&*x,&*c,&*_norm});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::Exponential, output, nEvents, dataMap, {&*x,&*c,&*_norm});
 }
 
 

diff --git a/roofit/roofit/src/RooGamma.cxx b/roofit/roofit/src/RooGamma.cxx
@@ -87,9 +87,10 @@ Double_t RooGamma::evaluate() const
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Compute multiple values of Gamma PDF.  
-void RooGamma::computeBatch(rbc::RbcInterface* dispatch, double* output, size_t nEvents, rbc::DataMap& dataMap) const
+void RooGamma::computeBatch(cudaStream_t* stream, double* output, size_t nEvents, rbc::DataMap& dataMap) const
 {
-  dispatch->compute(rbc::Gamma, output, nEvents, dataMap, {&*x,&*gamma,&*beta,&*mu,&*_norm});
+  auto dispatch = stream ? rbc::dispatchCUDA : rbc::dispatchCPU;
+  dispatch->compute(stream, rbc::Gamma, output, nEvents, dataMap, {&*x,&*gamma,&*beta,&*mu,&*_norm});
 }
 
 ////////////////////////////////////////////////////////////////////////////////