50#include <unordered_set>
58void assignSpan(std::span<T> &to, std::span<T>
const &from)
63void logArchitectureInfo(
bool useGPU)
75 static std::pair<bool, bool> lastUseGPU;
76 if (lastUseGPU.second && lastUseGPU.first == useGPU)
78 lastUseGPU = {useGPU,
true};
81 auto log = [](std::string_view message) {
86 log(
"using generic CPU library compiled with no vectorizations");
91 log(
"using CUDA computation library");
101void logMissingCudaSupport(
RooAbsArg const &arg)
106 static std::set<std::string> warnedClasses;
107 static std::mutex warnedClassesMutex;
108 std::scoped_lock guard{warnedClassesMutex};
109 if (warnedClasses.insert(arg.
ClassName()).second) {
111 <<
" could not be evaluated on the GPU because it doesn't support it."
112 <<
" Consider requesting or implementing it to benefit from a speed up."
128 std::shared_ptr<RooBatchCompute::AbsBuffer>
buffer;
170 : _topNode{const_cast<
RooAbsReal &>(absReal)}, _useGPU{useGPU}
174 throw std::runtime_error(
"Can't create Evaluator in CUDA mode because RooBatchCompute CUDA could not be loaded!");
190 std::map<RooFit::Detail::DataKey, NodeInfo *> nodeInfos;
194 std::size_t iNode = 0;
198 auto &nodeInfo =
_nodes.back();
201 nodeInfo.absArg = arg;
202 nodeInfo.originalOperMode = arg->
operMode();
203 nodeInfo.iNode = iNode;
204 nodeInfos[arg] = &nodeInfo;
207 nodeInfo.isVariable =
true;
212 nodeInfo.isCategory =
true;
219 info.serverInfos.reserve(info.absArg->servers().size());
220 for (
RooAbsArg *server : info.absArg->servers()) {
221 if (server->isValueServer(*info.absArg)) {
222 auto *serverInfo = nodeInfos.at(server);
223 info.serverInfos.emplace_back(serverInfo);
224 serverInfo->clientInfos.emplace_back(&info);
230 _nodes.back().isValueServer =
true;
231 for (
auto iter =
_nodes.rbegin(); iter !=
_nodes.rend(); ++iter) {
232 if (!iter->isValueServer)
234 for (
auto &serverInfo : iter->serverInfos) {
235 serverInfo->isValueServer =
true;
249 for (
auto &info :
_nodes) {
261 std::size_t iValueServer = 0;
262 for (
RooAbsArg *server : info.absArg->servers()) {
263 if (server->isValueServer(*info.absArg)) {
264 auto *knownServer = info.serverInfos[iValueServer]->absArg;
265 if (knownServer->hasDataToken()) {
266 server->setDataToken(knownServer->dataToken());
277 throw std::runtime_error(
"Evaluator can only take device array as input in CUDA mode!");
294 static std::atomic<std::size_t> nextInputGeneration{1};
295 const std::size_t gen = ++nextInputGeneration;
326 info.
buffer->assignFromDevice(gpuSpan);
332 info.
buffer->assignFromHost(cpuSpan);
339 std::map<RooFit::Detail::DataKey, std::size_t> sizeMap;
340 for (
auto &info :
_nodes) {
341 if (info.fromArrayInput) {
342 sizeMap[info.absArg] = info.outputSize;
351 auto found = sizeMap.find(key);
352 return found != sizeMap.end() ? found->second : -1;
355 for (
auto &info :
_nodes) {
356 info.outputSize = outputSizeMap.at(info.absArg);
369 for (
auto &info :
_nodes) {
370 if (!info.isVariable) {
371 info.absArg->resetDataToken();
381 using namespace Detail;
385 double *buffer =
nullptr;
395 logMissingCudaSupport(*info.
absArg);
402 buffer = info.
buffer->hostWritePtr();
412 buffer[0] = nodeAbsCategory->getCurrentIndex();
414 throw std::runtime_error(
"RooFit::Evaluator - non-scalar category values are not supported!");
417 auto nodeAbsReal =
static_cast<RooAbsReal const *
>(node);
435 auto *var =
static_cast<RooRealVar const *
>(node);
439 clientInfo->isDirty =
true;
452 if (nodeInfo.
lastCatVal != cat->getCurrentIndex()) {
455 clientInfo->isDirty =
true;
467 clientInfo->isDirty =
true;
488 for (
auto &nodeInfo :
_nodes) {
489 if (!nodeInfo.fromArrayInput) {
490 if (nodeInfo.isVariable) {
492 }
else if (nodeInfo.isCategory) {
495 if (nodeInfo.isDirty) {
498 nodeInfo.isDirty =
false;
516 for (
auto &info :
_nodes) {
517 info.remClients = info.clientInfos.size();
518 if (info.buffer && !info.fromArrayInput) {
531 for (
auto &info :
_nodes) {
532 if (!info.fromArrayInput) {
533 if (info.computeInGPU) {
546 for (
auto *serverInfo : info.serverInfos) {
547 serverInfo->decrementRemainingClients();
580 for (
auto &action : ctx->_deferredActions) {
583 ctx->_deferredActions.clear();
593 using namespace Detail;
599 double *buffer =
nullptr;
606 buffer = info.
buffer->deviceWritePtr();
623 for (
auto &info :
_nodes) {
624 info.computeInGPU =
false;
625 if (!info.absArg->canComputeBatchWithCuda()) {
628 for (
NodeInfo const *serverInfo : info.serverInfos) {
629 if (serverInfo->outputSize > 1) {
630 info.computeInGPU =
true;
637 for (
auto &info :
_nodes) {
638 info.copyAfterEvaluation =
false;
640 if (!info.isScalar()) {
641 for (
auto *clientInfo : info.clientInfos) {
642 if (info.computeInGPU != clientInfo->computeInGPU) {
643 info.copyAfterEvaluation =
true;
660 for (
auto &info :
_nodes) {
661 if (info.isVariable) {
694 auto out = std::make_unique<ChangeOperModeRAII>();
695 std::unordered_set<RooAbsArg *> visited;
697 std::vector<RooAbsArg *> queue;
698 queue.reserve(
_nodes.size());
699 for (
auto &info :
_nodes) {
700 queue.push_back(info.absArg);
703 while (!queue.empty()) {
706 if (!visited.insert(node).second)
709 out->change(node, opMode);
715 queue.push_back(client);
724 std::cout <<
"--- RooFit BatchMode evaluation ---\n";
726 std::vector<int> widths{9, 37, 20, 9, 10, 20};
728 auto printElement = [&](
int iCol,
auto const &t) {
729 const char separator =
' ';
730 os << separator << std::left << std::setw(widths[iCol]) << std::setfill(separator) << t;
734 auto printHorizontalRow = [&]() {
736 for (
int w : widths) {
739 for (
int i = 0; i <
n; i++) {
745 printHorizontalRow();
748 printElement(0,
"Index");
749 printElement(1,
"Name");
750 printElement(2,
"Class");
751 printElement(3,
"Size");
752 printElement(4,
"From Data");
753 printElement(5,
"1st value");
756 printHorizontalRow();
758 for (std::size_t iNode = 0; iNode <
_nodes.size(); ++iNode) {
759 auto &nodeInfo =
_nodes[iNode];
765 printElement(0, iNode);
766 printElement(1, node->
GetName());
768 printElement(3, nodeInfo.outputSize);
769 printElement(4, nodeInfo.fromArrayInput);
770 printElement(5, span[0]);
775 printHorizontalRow();
787 for (
auto &nodeInfo :
_nodes) {
788 if (nodeInfo.isValueServer && nodeInfo.absArg->isFundamental()) {
789 parameters.
add(*nodeInfo.absArg);
814 for (
auto &nodeInfo :
_nodes) {
815 if (nodeInfo.absArg->isReducerNode()) {
816 nodeInfo.isDirty =
true;
Option_t Option_t TPoint TPoint const char mode
Common abstract base class for objects that represent a value and a "shape" in RooFit.
const TNamed * namePtr() const
De-duplicated pointer to this object's name.
void setDataToken(std::size_t index)
Sets the token for retrieving results in the BatchMode. For internal use only.
const RefCountList_t & valueClients() const
List of all value clients of this object. Value clients receive value updates.
OperMode operMode() const
Query the operation mode of this node.
A space to attach TBranches.
virtual bool add(const RooAbsArg &var, bool silent=false)
Add the specified argument to list.
Storage_t::size_type size() const
void sort(bool reverse=false)
Sort collection using std::sort and name comparison.
Abstract base class for objects that represent a real value and implements functionality common to al...
virtual void doEval(RooFit::EvalContext &) const
Base function for computing multiple values of a RooAbsReal.
RooArgSet is a container object that can hold multiple RooAbsArg objects.
Minimal configuration struct to steer the evaluation of a single node with the RooBatchCompute librar...
void setCudaStream(CudaInterface::CudaStream *cudaStream)
void setNThreads(int nThreads)
Number of threads to use for CPU batch computations and reductions.
virtual void synchronizeCudaStream(CudaInterface::CudaStream *) const =0
Wait until all work that was enqueued on the stream has completed.
virtual std::unique_ptr< AbsBufferManager > createBufferManager() const =0
virtual CudaInterface::CudaStream * newCudaStream() const =0
virtual void deleteCudaStream(CudaInterface::CudaStream *) const =0
std::size_t _inputGeneration
std::vector< std::function< void()> > _deferredActions
void set(RooAbsArg const *arg, std::span< const double > const &span)
std::span< const double > at(RooAbsArg const *arg, RooAbsArg const *caller=nullptr)
void resetVectorBuffers()
void enableVectorBuffers(bool enable)
RooBatchCompute::Config config(RooAbsArg const *arg) const
void setConfig(RooAbsArg const *arg, RooBatchCompute::Config const &config)
std::span< double > _currentOutput
void resize(std::size_t n)
void print(std::ostream &os)
void setClientsDirty(NodeInfo &nodeInfo)
Flags all the clients of a given node dirty.
std::unique_ptr< ChangeOperModeRAII > setOperModes(RooAbsArg::OperMode opMode)
RooArgSet getParameters() const
Gets all the parameters of the RooAbsReal.
void setOffsetMode(RooFit::EvalContext::OffsetMode)
Sets the offset mode for evaluation.
void syncDataTokens()
If there are servers with the same name that got de-duplicated in the _nodes list,...
std::unordered_map< TNamed const *, NodeInfo * > _nodesMap
std::unique_ptr< ChangeOperModeRAII > _operModeChanges
std::vector< NodeInfo > _nodes
bool _needToUpdateOutputSizes
std::span< const double > getValHeterogeneous()
Returns the value of the top node in the computation graph.
std::span< const double > run()
Returns the value of the top node in the computation graph.
Evaluator(const RooAbsReal &absReal, bool useGPU=false)
Construct a new Evaluator.
void setNThreads(int nThreads)
Sets the number of threads to use for the evaluation of a single node.
void processVariable(NodeInfo &nodeInfo)
Process a variable in the computation graph.
void processCategory(NodeInfo &nodeInfo)
Process a category in the computation graph.
RooBatchCompute::CudaInterface::CudaStream * _cudaStream
std::unique_ptr< RooBatchCompute::AbsBufferManager > _bufferManager
void markGPUNodes()
Decides which nodes are assigned to the GPU in a CUDA fit.
void assignToGPU(NodeInfo &info)
Enqueue the computation of a node on the GPU.
void setInput(std::string const &name, std::span< const double > inputArray, bool isOnDevice)
RooFit::EvalContext _evalContextCUDA
RooFit::EvalContext _evalContextCPU
void computeCPUNode(const RooAbsArg *node, NodeInfo &info)
void setOperMode(RooAbsArg *arg, RooAbsArg::OperMode opMode)
Temporarily change the operation mode of a RooAbsArg until the Evaluator gets deleted.
static RooMsgService & instance()
Return reference to singleton instance.
static const TNamed * ptr(const char *stringPtr)
Return a unique TNamed pointer for given C++ string.
Variable that can be changed from the outside.
const char * GetName() const override
Returns name of object.
virtual const char * ClassName() const
Returns name of class to which the object belongs.
RVec< PromoteType< T > > log(const RVec< T > &v)
R__EXTERN RooBatchComputeInterface * dispatchCUDA
std::string cpuArchitectureName()
R__EXTERN RooBatchComputeInterface * dispatchCPU
This dispatch pointer points to an implementation of the compute library, provided one has been loade...
Architecture cpuArchitecture()
int initCPU()
Inspect hardware capabilities, and load the optimal library for RooFit computations.
The namespace RooFit contains mostly switches that change the behaviour of functions of PDFs (or othe...
void getSortedComputationGraph(RooAbsArg const &func, RooArgSet &out)
A struct used by the Evaluator to store information on the RooAbsArgs in the computation graph.
std::size_t lastSetValCount
std::vector< NodeInfo * > serverInfos
RooAbsArg::OperMode originalOperMode
std::vector< NodeInfo * > clientInfos
std::shared_ptr< RooBatchCompute::AbsBuffer > buffer
void decrementRemainingClients()
Check the servers of a node that has been computed and release its resources if they are no longer ne...