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
0024 class ChangeOperModeRAII;
0025 class RooAbsArg;
0026
0027 namespace RooBatchCompute {
0028 class AbsBufferManager;
0029 }
0030
0031 namespace RooFit {
0032
0033 struct NodeInfo;
0034
0035 class Evaluator {
0036 public:
0037    Evaluator(const RooAbsReal &absReal, bool useGPU = false);
0038    ~Evaluator();
0039
0040    std::span<const double> run();
0041    void setInput(std::string const &name, std::span<const double> inputArray, bool isOnDevice);
0042    RooArgSet getParameters() const;
0043    void print(std::ostream &os);
0044
0045    void setOffsetMode(RooFit::EvalContext::OffsetMode);
0046
0047    std::unique_ptr<ChangeOperModeRAII> setOperModes(RooAbsArg::OperMode opMode);
0048
0049 private:
0050    void processVariable(NodeInfo &nodeInfo);
0051    void processCategory(NodeInfo &nodeInfo);
0052    void setClientsDirty(NodeInfo &nodeInfo);
0053    std::span<const double> getValHeterogeneous();
0054    void markGPUNodes();
0055    void assignToGPU(NodeInfo &info);
0056    void computeCPUNode(const RooAbsArg *node, NodeInfo &info);
0057    void setOperMode(RooAbsArg *arg, RooAbsArg::OperMode opMode);
0058    void syncDataTokens();
0059    void updateOutputSizes();
0060
0061    std::unique_ptr<RooBatchCompute::AbsBufferManager> _bufferManager;
0062    RooAbsReal &_topNode;
0063    const bool _useGPU = false;
0064    int _nEvaluations = 0;
0065    bool _needToUpdateOutputSizes = false;
0066    RooFit::EvalContext _evalContextCPU;
0067    RooFit::EvalContext _evalContextCUDA;
0068    std::vector<NodeInfo> _nodes;                             // the ordered computation graph
0069    std::unordered_map<TNamed const *, NodeInfo *> _nodesMap; // for quick lookup of nodes
0070    std::unique_ptr<ChangeOperModeRAII> _operModeChanges;
0071 };
0072
0073 } // end namespace RooFit
0074
0075 #endif