include/root/RooONNXFunc.h

0001 /*
0002  * Project: RooFit
0003  * Authors:
0004  *   Jonas Rembser, CERN  04/2026
0005  *
0006  * Copyright (c) 2026, CERN
0007  *
0008  * Redistribution and use in source and binary forms,
0009  * with or without modification, are permitted according to the terms
0010  * listed in LICENSE (http://roofit.sourceforge.net/license.txt)
0011  */
0012
0013 #ifndef RooFit_RooONNXFunc_h
0014 #define RooFit_RooONNXFunc_h
0015
0016 #include <RooAbsReal.h>
0017 #include <RooListProxy.h>
0018
0019 #include <any>
0020
0021 class RooONNXFunc : public RooAbsReal {
0022 public:
0023    RooONNXFunc() = default;
0024
0025    RooONNXFunc(const char *name, const char *title, const std::vector<RooArgList> &inputTensors,
0026                    const std::string &onnxFile, const std::vector<std::string> &inputNames = {},
0027                    const std::vector<std::vector<int>> &inputShapes = {});
0028
0029    RooONNXFunc(const RooONNXFunc &other, const char *newName = nullptr);
0030
0031    TObject *clone(const char *newName) const override { return new RooONNXFunc(*this, newName); }
0032
0033    std::size_t nInputTensors() const { return _inputTensors.size(); }
0034    RooArgList const &inputTensorList(int iTensor) const { return *(_inputTensors[iTensor]); }
0035
0036    std::string funcName() const { return _funcName; }
0037    std::string outerWrapperName() const { return "TMVA_SOFIE_" + funcName() + "::roo_outer_wrapper"; }
0038
0039 protected:
0040    double evaluate() const override;
0041
0042 private:
0043    /// Build transient runtime backend on first use.
0044    void initialize();
0045
0046    /// Gather current RooFit inputs into a contiguous feature buffer.
0047    void fillInputBuffer() const;
0048
0049    struct RuntimeCache;
0050
0051    std::vector<std::unique_ptr<RooListProxy>> _inputTensors; ///< Inputs mapping to flattened input tensors.
0052    std::vector<std::uint8_t> _onnxBytes;                     ///< Persisted ONNX model bytes.
0053    std::shared_ptr<RuntimeCache> _runtime;                   ///<! Transient runtime information.
0054    mutable std::vector<float> _inputBuffer;                  ///<!
0055    std::string _funcName;                                    ///<!
0056
0057    ClassDefOverride(RooONNXFunc, 1)
0058 };
0059
0060 namespace RooFit::Detail {
0061
0062 struct AnyWithVoidPtr {
0063    std::any any;
0064    void *ptr = nullptr;
0065
0066    template <class T>
0067    void emplace()
0068    {
0069       any = std::make_any<T>();
0070       ptr = std::any_cast<T>(&any);
0071    }
0072
0073    void emplace(std::string const &typeName);
0074 };
0075
0076 template <class Session_t, class... Inputs>
0077 void doInferWithSessionVoidPtr(void *session, float *out, Inputs const *...inputs)
0078 {
0079    doInfer(*reinterpret_cast<Session_t *>(session), inputs..., out);
0080 }
0081
0082 } // namespace RooFit::Detail
0083
0084 #endif