Evaluator.h

Go to the documentation of this file.

/*
 * Project: RooFit
 * Authors:
 *   Jonas Rembser, CERN 2021
 *   Emmanouil Michalainas, CERN 2021
 *
 * Copyright (c) 2023, CERN
 *
 * Redistribution and use in source and binary forms,
 * with or without modification, are permitted according to the terms
 * listed in LICENSE (http://roofit.sourceforge.net/license.txt)
 */
 
#ifndef RooFit_Evaluator_h
#define RooFit_Evaluator_h
 
#include <RooAbsReal.h>
#include <RooFit/EvalContext.h>
 
#include <ROOT/RConfig.hxx>
 
#include <memory>
 
class ChangeOperModeRAII;
class RooAbsArg;
 
namespace RooBatchCompute {
class AbsBufferManager;
namespace CudaInterface {
class CudaStream;
}
} // namespace RooBatchCompute
 
namespace RooFit {
 
struct NodeInfo;
 
class Evaluator {
public:
   Evaluator(const RooAbsReal &absReal, bool useGPU = false);
   ~Evaluator();
 
   std::span<const double> run();
   void setInput(std::string const &name, std::span<const double> inputArray, bool isOnDevice);
   RooArgSet getParameters() const;
   void print(std::ostream &os);
 
   void setNThreads(int nThreads);
 
   void setOffsetMode(RooFit::EvalContext::OffsetMode);
 
   std::unique_ptr<ChangeOperModeRAII> setOperModes(RooAbsArg::OperMode opMode);
 
private:
   void processVariable(NodeInfo &nodeInfo);
   void processCategory(NodeInfo &nodeInfo);
   void setClientsDirty(NodeInfo &nodeInfo);
   std::span<const double> getValHeterogeneous();
   void markGPUNodes();
   void assignToGPU(NodeInfo &info);
   void computeCPUNode(const RooAbsArg *node, NodeInfo &info);
   void setOperMode(RooAbsArg *arg, RooAbsArg::OperMode opMode);
   void syncDataTokens();
   void updateOutputSizes();
 
   std::unique_ptr<RooBatchCompute::AbsBufferManager> _bufferManager;
   RooAbsReal &_topNode;
   const bool _useGPU = false;
   int _nEvaluations = 0;
   bool _needToUpdateOutputSizes = false;
   RooFit::EvalContext _evalContextCPU;
   RooFit::EvalContext _evalContextCUDA;
   std::vector<NodeInfo> _nodes;                             // the ordered computation graph
   std::unordered_map<TNamed const *, NodeInfo *> _nodesMap; // for quick lookup of nodes
   std::unique_ptr<ChangeOperModeRAII> _operModeChanges;
   // the single CUDA stream on which all GPU work of this Evaluator is enqueued
   RooBatchCompute::CudaInterface::CudaStream *_cudaStream = nullptr;
};
 
} // end namespace RooFit
 
#endif