root/RooFit/Evaluator.h

0001 /*
0002  * Project: RooFit
0003  * Authors:
0004  *   Jonas Rembser, CERN 2021
0005  *   Emmanouil Michalainas, CERN 2021
0006  *
0007  * Copyright (c) 2023, CERN
0008  *
0009  * Redistribution and use in source and binary forms,
0010  * with or without modification, are permitted according to the terms
0011  * listed in LICENSE (http://roofit.sourceforge.net/license.txt)
0012  */
0013
0014 #ifndef RooFit_Evaluator_h
0015 #define RooFit_Evaluator_h
0016
0017 #include <RooAbsReal.h>
0018 #include <RooFit/EvalContext.h>
0019
0020 #include <RConfig.h>
0021
0022 #include <memory>
0023 #include <stack>
0024
0025 class ChangeOperModeRAII;
0026 class RooAbsArg;
0027
0028 namespace RooBatchCompute {
0029 class AbsBufferManager;
0030 }
0031
0032 namespace RooFit {
0033
0034 struct NodeInfo;
0035
0036 class Evaluator {
0037 public:
0038    Evaluator(const RooAbsReal &absReal, bool useGPU = false);
0039    ~Evaluator();
0040
0041    std::span<const double> run();
0042    void setInput(std::string const &name, std::span<const double> inputArray, bool isOnDevice);
0043    RooArgSet getParameters() const;
0044    void print(std::ostream &os);
0045
0046    void setOffsetMode(RooFit::EvalContext::OffsetMode);
0047
0048 private:
0049    void processVariable(NodeInfo &nodeInfo);
0050    void setClientsDirty(NodeInfo &nodeInfo);
0051    std::span<const double> getValHeterogeneous();
0052    void markGPUNodes();
0053    void assignToGPU(NodeInfo &info);
0054    void computeCPUNode(const RooAbsArg *node, NodeInfo &info);
0055    void setOperMode(RooAbsArg *arg, RooAbsArg::OperMode opMode);
0056    void syncDataTokens();
0057    void updateOutputSizes();
0058
0059    std::unique_ptr<RooBatchCompute::AbsBufferManager> _bufferManager;
0060    RooAbsReal &_topNode;
0061    const bool _useGPU = false;
0062    int _nEvaluations = 0;
0063    bool _needToUpdateOutputSizes = false;
0064    RooFit::EvalContext _evalContextCPU;
0065    RooFit::EvalContext _evalContextCUDA;
0066    std::vector<NodeInfo> _nodes; // the ordered computation graph
0067    std::stack<std::unique_ptr<ChangeOperModeRAII>> _changeOperModeRAIIs;
0068 };
0069
0070 } // end namespace RooFit
0071
0072 #endif