From d034ccd0cc89e92fe81658420ede8b92232a962c Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Wed, 8 Oct 2025 20:44:06 +0200 Subject: [PATCH 1/6] [RF] Use global workspace array in the generated code This will make it easier to manage intermediate results. --- roofit/codegen/src/CodegenImpl.cxx | 9 ++++++++ roofit/roofitcore/inc/RooFit/CodegenContext.h | 8 +++++-- roofit/roofitcore/src/RooEvaluatorWrapper.cxx | 2 +- .../roofitcore/src/RooFit/CodegenContext.cxx | 21 +++++++++++++------ 4 files changed, 31 insertions(+), 9 deletions(-) diff --git a/roofit/codegen/src/CodegenImpl.cxx b/roofit/codegen/src/CodegenImpl.cxx index a583983eabc09..33e9c7ff809ff 100644 --- a/roofit/codegen/src/CodegenImpl.cxx +++ b/roofit/codegen/src/CodegenImpl.cxx @@ -655,6 +655,15 @@ void codegenImpl(RooRealSumPdf &arg, CodegenContext &ctx) void codegenImpl(RooRealVar &arg, CodegenContext &ctx) { + auto found = ctx._paramIndices.find(&arg); + if (found != ctx._paramIndices.end()) { + auto savedName = "wksp[" + std::to_string(ctx._nWksp++) + "]"; + std::string outVarDecl = savedName + " = params[" + std::to_string(found->second) + "];\n"; + ctx.addToCodeBody(outVarDecl, ctx.isScopeIndependent(&arg)); + ctx.addResult(&arg, savedName); + return; + } + if (!arg.isConstant()) { ctx.addResult(&arg, arg.GetName()); } diff --git a/roofit/roofitcore/inc/RooFit/CodegenContext.h b/roofit/roofitcore/inc/RooFit/CodegenContext.h index 961b8176d4758..65c2426692f73 100644 --- a/roofit/roofitcore/inc/RooFit/CodegenContext.h +++ b/roofit/roofitcore/inc/RooFit/CodegenContext.h @@ -70,6 +70,7 @@ class CodegenContext { void addToGlobalScope(std::string const &str); void addVecObs(const char *key, int idx); + void addParam(const RooAbsArg *key, int idx); int observableIndexOf(const RooAbsArg &arg) const; void addToCodeBody(RooAbsArg const *klass, std::string const &in); @@ -135,14 +136,17 @@ class CodegenContext { }; ScopeRAII OutputScopeRangeComment(RooAbsArg const *arg) { return {arg, *this}; } + bool isScopeIndependent(RooAbsArg const *in) const; + + std::size_t _nWksp = 0; + std::unordered_map _paramIndices; + private: void pushScope(); void popScope(); template std::string buildArgSpanImpl(std::span arr); - bool isScopeIndependent(RooAbsArg const *in) const; - void endLoop(LoopScope const &scope); void addResult(TNamed const *key, std::string const &value); diff --git a/roofit/roofitcore/src/RooEvaluatorWrapper.cxx b/roofit/roofitcore/src/RooEvaluatorWrapper.cxx index e975a3002dd50..70fef04357aa2 100644 --- a/roofit/roofitcore/src/RooEvaluatorWrapper.cxx +++ b/roofit/roofitcore/src/RooEvaluatorWrapper.cxx @@ -239,7 +239,7 @@ RooFuncWrapper::RooFuncWrapper(RooAbsReal &obj, const RooAbsData *data, RooSimul // First update the result variable of params in the compute graph to in[]. int idx = 0; for (RooAbsArg *param : _params) { - ctx.addResult(param, "params[" + std::to_string(idx) + "]"); + ctx.addParam(param, idx); idx++; } diff --git a/roofit/roofitcore/src/RooFit/CodegenContext.cxx b/roofit/roofitcore/src/RooFit/CodegenContext.cxx index 2bf3497b685d4..17529cbaf2873 100644 --- a/roofit/roofitcore/src/RooFit/CodegenContext.cxx +++ b/roofit/roofitcore/src/RooFit/CodegenContext.cxx @@ -105,6 +105,11 @@ void CodegenContext::addVecObs(const char *key, int idx) _vecObsIndices[namePtr] = idx; } +void CodegenContext::addParam(RooAbsArg const *arg, int idx) +{ + _paramIndices[arg] = idx; +} + int CodegenContext::observableIndexOf(RooAbsArg const &arg) const { auto it = _vecObsIndices.find(arg.namePtr()); @@ -214,8 +219,7 @@ std::string CodegenContext::getTmpVarName() const /// @param valueToSave The actual string value to save as a temporary. void CodegenContext::addResult(RooAbsArg const *in, std::string const &valueToSave) { - // std::string savedName = RooFit::Detail::makeValidVarName(in->GetName()); - std::string savedName = getTmpVarName(); + std::string savedName; // Only save values if they contain operations or they are numerals. Otherwise, we can use them directly. @@ -229,7 +233,8 @@ void CodegenContext::addResult(RooAbsArg const *in, std::string const &valueToSa // If the name is not empty and this value is worth saving, save it to the correct scope. // otherwise, just return the actual value itself if (hasOperations || isNumeric) { - std::string outVarDecl = "const double " + savedName + " = " + valueToSave + ";\n"; + savedName = "wksp[" + std::to_string(_nWksp++) + "]"; + std::string outVarDecl = savedName + " = " + valueToSave + ";\n"; addToCodeBody(in, outVarDecl); } else { savedName = valueToSave; @@ -332,9 +337,10 @@ CodegenContext::buildFunction(RooAbsArg const &arg, std::map::infinity();\n" - << funcBody << "\n}"; + << "constexpr double inf = std::numeric_limits::infinity();\n"; + if (ctx._nWksp > 0) { + bodyWithSigStrm << "double wksp [" << ctx._nWksp << "]{};\n"; + } + bodyWithSigStrm << funcBody << "\n}"; ctx._collectedFunctions.emplace_back(funcName); if (!gInterpreter->Declare(bodyWithSigStrm.str().c_str())) { std::stringstream errorMsg; From fc2428ed5b3318d25af559b305fdcb39125691f4 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Wed, 8 Oct 2025 22:53:47 +0200 Subject: [PATCH 2/6] continue --- roofit/codegen/src/CodegenImpl.cxx | 12 ------ roofit/roofitcore/inc/RooFit/CodegenContext.h | 6 +-- roofit/roofitcore/src/RooEvaluatorWrapper.cxx | 1 - .../roofitcore/src/RooFit/CodegenContext.cxx | 40 +++++++++++++++---- 4 files changed, 34 insertions(+), 25 deletions(-) diff --git a/roofit/codegen/src/CodegenImpl.cxx b/roofit/codegen/src/CodegenImpl.cxx index 33e9c7ff809ff..ca192c0eaf6ff 100644 --- a/roofit/codegen/src/CodegenImpl.cxx +++ b/roofit/codegen/src/CodegenImpl.cxx @@ -655,18 +655,6 @@ void codegenImpl(RooRealSumPdf &arg, CodegenContext &ctx) void codegenImpl(RooRealVar &arg, CodegenContext &ctx) { - auto found = ctx._paramIndices.find(&arg); - if (found != ctx._paramIndices.end()) { - auto savedName = "wksp[" + std::to_string(ctx._nWksp++) + "]"; - std::string outVarDecl = savedName + " = params[" + std::to_string(found->second) + "];\n"; - ctx.addToCodeBody(outVarDecl, ctx.isScopeIndependent(&arg)); - ctx.addResult(&arg, savedName); - return; - } - - if (!arg.isConstant()) { - ctx.addResult(&arg, arg.GetName()); - } ctx.addResult(&arg, doubleToString(arg.getVal())); } diff --git a/roofit/roofitcore/inc/RooFit/CodegenContext.h b/roofit/roofitcore/inc/RooFit/CodegenContext.h index 65c2426692f73..8158b1cb4af1f 100644 --- a/roofit/roofitcore/inc/RooFit/CodegenContext.h +++ b/roofit/roofitcore/inc/RooFit/CodegenContext.h @@ -140,6 +140,8 @@ class CodegenContext { std::size_t _nWksp = 0; std::unordered_map _paramIndices; + /// @brief A map to keep track of the observable indices if they are non scalar. + std::unordered_map _vecObsIndices; private: void pushScope(); @@ -149,8 +151,6 @@ class CodegenContext { void endLoop(LoopScope const &scope); - void addResult(TNamed const *key, std::string const &value); - template {}, bool>::type = true> std::string buildArg(T x) { @@ -197,8 +197,6 @@ class CodegenContext { /// @brief Map of node names to their result strings. std::unordered_map _nodeNames; - /// @brief A map to keep track of the observable indices if they are non scalar. - std::unordered_map _vecObsIndices; /// @brief Map of node output sizes. std::map _nodeOutputSizes; /// @brief The code layered by lexical scopes used as a stack. diff --git a/roofit/roofitcore/src/RooEvaluatorWrapper.cxx b/roofit/roofitcore/src/RooEvaluatorWrapper.cxx index 70fef04357aa2..cbfd93242a3bf 100644 --- a/roofit/roofitcore/src/RooEvaluatorWrapper.cxx +++ b/roofit/roofitcore/src/RooEvaluatorWrapper.cxx @@ -251,7 +251,6 @@ RooFuncWrapper::RooFuncWrapper(RooAbsReal &obj, const RooAbsData *data, RooSimul if (item.second.size == 1) { ctx.addResult(obsName, "obs[" + std::to_string(item.second.idx) + "]"); } else { - ctx.addResult(obsName, "obs"); ctx.addVecObs(obsName, item.second.idx); } } diff --git a/roofit/roofitcore/src/RooFit/CodegenContext.cxx b/roofit/roofitcore/src/RooFit/CodegenContext.cxx index 17529cbaf2873..dbb6372222bd3 100644 --- a/roofit/roofitcore/src/RooFit/CodegenContext.cxx +++ b/roofit/roofitcore/src/RooFit/CodegenContext.cxx @@ -44,12 +44,7 @@ void CodegenContext::addResult(const char *key, std::string const &value) { const TNamed *namePtr = RooNameReg::known(key); if (namePtr) - addResult(namePtr, value); -} - -void CodegenContext::addResult(TNamed const *key, std::string const &value) -{ - _nodeNames[key] = value; + _nodeNames[namePtr] = value; } /// @brief Gets the result for the given node using the node name. This node also performs the necessary @@ -175,7 +170,6 @@ std::unique_ptr CodegenContext::beginLoop(RooAbsArg c continue; vars.push_back(it.first); - _nodeNames[it.first] = "obs[" + std::to_string(it.second) + " + " + idx + "]"; } // TODO: we are using the size of the first loop variable to the the number @@ -193,6 +187,16 @@ std::unique_ptr CodegenContext::beginLoop(RooAbsArg c // Make sure that the name of this variable doesn't clash with other stuff addToCodeBody(in, "for(int " + idx + " = 0; " + idx + " < " + std::to_string(numEntries) + "; " + idx + "++) {\n"); + // set the results of the vector observables + for (auto const &it : _vecObsIndices) { + if (!in->dependsOn(it.first)) + continue; + + auto iWksp = std::to_string(_nWksp++); + addToCodeBody(in, "wksp[" + iWksp + "] = obs[" + std::to_string(it.second) + " + " + idx + "];\n"); + _nodeNames[it.first] = "wksp[" + iWksp + "]"; + } + return std::make_unique(*this, std::move(vars)); } @@ -240,7 +244,7 @@ void CodegenContext::addResult(RooAbsArg const *in, std::string const &valueToSa savedName = valueToSave; } - addResult(in->namePtr(), savedName); + _nodeNames[in->namePtr()] = savedName; } /// @brief Function to save a RooListProxy as an array in the squashed code. @@ -422,6 +426,26 @@ struct Caller_FUNC_NAME { void codegen(RooAbsArg &arg, CodegenContext &ctx) { + // parameters + auto foundParam = ctx._paramIndices.find(&arg); + if (foundParam != ctx._paramIndices.end()) { + auto savedName = "wksp[" + std::to_string(ctx._nWksp++) + "]"; + std::string outVarDecl = savedName + " = params[" + std::to_string(foundParam->second) + "];\n"; + ctx.addToCodeBody(outVarDecl, ctx.isScopeIndependent(&arg)); + ctx.addResult(&arg, savedName); + return; + } + + // observables + auto foundObs = ctx._vecObsIndices.find(&arg); + if (foundObs != ctx._vecObsIndices.end()) { + auto savedName = "wksp[" + std::to_string(ctx._nWksp++) + "]"; + std::string outVarDecl = savedName + " = obs[" + std::to_string(foundObs->second) + "];\n"; + ctx.addToCodeBody(outVarDecl, ctx.isScopeIndependent(&arg)); + ctx.addResult(&arg, savedName); + return; + } + static bool codeDeclared = false; if (!codeDeclared) { declareDispatcherCode("codegenImpl"); From c42160f978b839a55766b8e17bb74fc2b47e7bb3 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Wed, 8 Oct 2025 23:53:57 +0200 Subject: [PATCH 3/6] continue --- roofit/codegen/src/CodegenImpl.cxx | 9 ++- roofit/roofitcore/inc/RooFit/CodegenContext.h | 9 ++- roofit/roofitcore/src/RooEvaluatorWrapper.cxx | 9 +-- .../roofitcore/src/RooFit/CodegenContext.cxx | 58 +++++-------------- 4 files changed, 25 insertions(+), 60 deletions(-) diff --git a/roofit/codegen/src/CodegenImpl.cxx b/roofit/codegen/src/CodegenImpl.cxx index ca192c0eaf6ff..0d3a270d631cb 100644 --- a/roofit/codegen/src/CodegenImpl.cxx +++ b/roofit/codegen/src/CodegenImpl.cxx @@ -294,7 +294,7 @@ void codegenImpl(RooCategory &arg, CodegenContext &ctx) if (idx < 0) { idx = 1; - ctx.addVecObs(arg.GetName(), idx); + ctx.addVecObs(arg.GetName(), idx, 1); } std::string result = std::to_string(arg.getCurrentIndex()); @@ -609,10 +609,11 @@ void codegenImpl(RooRealIntegral &arg, CodegenContext &ctx) auto &intVar = static_cast(*arg.numIntRealVars()[0]); std::string obsName = ctx.getTmpVarName(); - std::string oldIntVarResult = ctx.getResult(intVar); - ctx.addResult(&intVar, "obs[0]"); + auto oldVecObsInfo = ctx._vecObsIndices[intVar.namePtr()]; + ctx.addVecObs(intVar.GetName(), 0, 1); std::string funcName = ctx.buildFunction(arg.integrand(), {}); + ctx._vecObsIndices[intVar.namePtr()] = oldVecObsInfo; std::stringstream ss; @@ -639,8 +640,6 @@ void codegenImpl(RooRealIntegral &arg, CodegenContext &ctx) << "}\n"; ctx.addToGlobalScope(ss.str()); - - ctx.addResult(&intVar, oldIntVarResult); } void codegenImpl(RooRealSumFunc &arg, CodegenContext &ctx) diff --git a/roofit/roofitcore/inc/RooFit/CodegenContext.h b/roofit/roofitcore/inc/RooFit/CodegenContext.h index 8158b1cb4af1f..fbac603da8c53 100644 --- a/roofit/roofitcore/inc/RooFit/CodegenContext.h +++ b/roofit/roofitcore/inc/RooFit/CodegenContext.h @@ -46,7 +46,6 @@ using PrioLowest = Prio<10>; class CodegenContext { public: void addResult(RooAbsArg const *key, std::string const &value); - void addResult(const char *key, std::string const &value); std::string const &getResult(RooAbsArg const &arg); @@ -69,7 +68,7 @@ class CodegenContext { } void addToGlobalScope(std::string const &str); - void addVecObs(const char *key, int idx); + void addVecObs(const char *key, int idx, std::size_t size); void addParam(const RooAbsArg *key, int idx); int observableIndexOf(const RooAbsArg &arg) const; @@ -138,10 +137,12 @@ class CodegenContext { bool isScopeIndependent(RooAbsArg const *in) const; + /// @brief Map of node names to their result strings. + std::unordered_map _nodeNames; std::size_t _nWksp = 0; std::unordered_map _paramIndices; /// @brief A map to keep track of the observable indices if they are non scalar. - std::unordered_map _vecObsIndices; + std::unordered_map> _vecObsIndices; private: void pushScope(); @@ -195,8 +196,6 @@ class CodegenContext { template std::string typeName() const; - /// @brief Map of node names to their result strings. - std::unordered_map _nodeNames; /// @brief Map of node output sizes. std::map _nodeOutputSizes; /// @brief The code layered by lexical scopes used as a stack. diff --git a/roofit/roofitcore/src/RooEvaluatorWrapper.cxx b/roofit/roofitcore/src/RooEvaluatorWrapper.cxx index cbfd93242a3bf..8fac387f0e6c3 100644 --- a/roofit/roofitcore/src/RooEvaluatorWrapper.cxx +++ b/roofit/roofitcore/src/RooEvaluatorWrapper.cxx @@ -245,14 +245,7 @@ RooFuncWrapper::RooFuncWrapper(RooAbsReal &obj, const RooAbsData *data, RooSimul for (auto const &item : _obsInfos) { const char *obsName = item.first->GetName(); - // If the observable is scalar, set name to the start idx. else, store - // the start idx and later set the the name to obs[start_idx + curr_idx], - // here curr_idx is defined by a loop producing parent node. - if (item.second.size == 1) { - ctx.addResult(obsName, "obs[" + std::to_string(item.second.idx) + "]"); - } else { - ctx.addVecObs(obsName, item.second.idx); - } + ctx.addVecObs(obsName, item.second.idx, item.second.size); } gInterpreter->Declare("#pragma cling optimize(2)"); diff --git a/roofit/roofitcore/src/RooFit/CodegenContext.cxx b/roofit/roofitcore/src/RooFit/CodegenContext.cxx index dbb6372222bd3..7ffae9f3b01c8 100644 --- a/roofit/roofitcore/src/RooFit/CodegenContext.cxx +++ b/roofit/roofitcore/src/RooFit/CodegenContext.cxx @@ -25,28 +25,9 @@ #include #include -namespace { - -bool startsWith(std::string_view str, std::string_view prefix) -{ - return str.size() >= prefix.size() && 0 == str.compare(0, prefix.size(), prefix); -} - -} // namespace - namespace RooFit { namespace Experimental { -/// @brief Adds (or overwrites) the string representing the result of a node. -/// @param key The name of the node to add the result for. -/// @param value The new name to assign/overwrite. -void CodegenContext::addResult(const char *key, std::string const &value) -{ - const TNamed *namePtr = RooNameReg::known(key); - if (namePtr) - _nodeNames[namePtr] = value; -} - /// @brief Gets the result for the given node using the node name. This node also performs the necessary /// code generation through recursive calls to 'translate'. A call to this function modifies the already /// existing code body. @@ -63,14 +44,6 @@ std::string const &CodegenContext::getResult(RooAbsArg const &arg) if (found != _nodeNames.end()) return found->second; - // The result for vector observables should already be in the map if you - // opened the loop scope. This is just to check if we did not request the - // result of a vector-valued observable outside of the scope of a loop. - auto foundVecObs = _vecObsIndices.find(arg.namePtr()); - if (foundVecObs != _vecObsIndices.end()) { - throw std::runtime_error("You requested the result of a vector observable outside a loop scope for it!"); - } - auto RAII(OutputScopeRangeComment(&arg)); // Now, recursively call translate into the current argument to load the correct result. @@ -93,11 +66,11 @@ void CodegenContext::addToGlobalScope(std::string const &str) /// element. For example, a vector valued variable x with 10 entries will be squashed to obs[start_idx + i]. /// @param key The name of the node representing the vector valued observable. /// @param idx The start index (or relative position of the observable in the set of all observables). -void CodegenContext::addVecObs(const char *key, int idx) +void CodegenContext::addVecObs(const char *key, int idx, std::size_t size) { const TNamed *namePtr = RooNameReg::known(key); if (namePtr) - _vecObsIndices[namePtr] = idx; + _vecObsIndices[namePtr] = {idx, size}; } void CodegenContext::addParam(RooAbsArg const *arg, int idx) @@ -109,7 +82,7 @@ int CodegenContext::observableIndexOf(RooAbsArg const &arg) const { auto it = _vecObsIndices.find(arg.namePtr()); if (it != _vecObsIndices.end()) { - return it->second; + return it->second.first; } return -1; // Not found @@ -192,9 +165,16 @@ std::unique_ptr CodegenContext::beginLoop(RooAbsArg c if (!in->dependsOn(it.first)) continue; - auto iWksp = std::to_string(_nWksp++); - addToCodeBody(in, "wksp[" + iWksp + "] = obs[" + std::to_string(it.second) + " + " + idx + "];\n"); - _nodeNames[it.first] = "wksp[" + iWksp + "]"; + std::size_t n = it.second.second; + + auto savedName = "wksp[" + std::to_string(_nWksp++) + "]"; + std::string outVarDecl; + if (n == 1) + outVarDecl = savedName + " = obs[" + std::to_string(it.second.first) + "];\n"; + else + outVarDecl = savedName + " = obs[" + std::to_string(it.second.first) + " + " + idx + "];\n"; + addToCodeBody(outVarDecl, n == 1); + _nodeNames[it.first] = savedName; } return std::make_unique(*this, std::move(vars)); @@ -342,12 +322,6 @@ CodegenContext::buildFunction(RooAbsArg const &arg, std::mapsecond) + "];\n"; + std::string outVarDecl = savedName + " = obs[" + std::to_string(foundObs->second.first) + "];\n"; ctx.addToCodeBody(outVarDecl, ctx.isScopeIndependent(&arg)); ctx.addResult(&arg, savedName); return; From a8195f4459c9db0777f337d9f9f1496b4ed80ae3 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Thu, 9 Oct 2025 01:09:13 +0200 Subject: [PATCH 4/6] Use only workspace array --- roofit/codegen/src/CodegenImpl.cxx | 68 ++++++++++--------- roofit/roofitcore/inc/RooFit/CodegenContext.h | 6 +- .../roofitcore/src/RooFit/CodegenContext.cxx | 23 +------ 3 files changed, 40 insertions(+), 57 deletions(-) diff --git a/roofit/codegen/src/CodegenImpl.cxx b/roofit/codegen/src/CodegenImpl.cxx index 0d3a270d631cb..7c47e8fcf9ec2 100644 --- a/roofit/codegen/src/CodegenImpl.cxx +++ b/roofit/codegen/src/CodegenImpl.cxx @@ -62,8 +62,7 @@ #include -namespace RooFit { -namespace Experimental { +namespace RooFit::Experimental { namespace { @@ -103,7 +102,7 @@ void rooHistTranslateImpl(RooAbsArg const &arg, CodegenContext &ctx, int intOrde } std::string realSumPdfTranslateImpl(CodegenContext &ctx, RooAbsArg const &arg, RooArgList const &funcList, - RooArgList const &coefList, bool normalize) + RooArgList const &coefList, bool normalize, bool forceScopeIndependent) { bool noLastCoeff = funcList.size() != coefList.size(); @@ -113,7 +112,12 @@ std::string realSumPdfTranslateImpl(CodegenContext &ctx, RooAbsArg const &arg, R std::string sum = ctx.getTmpVarName(); std::string coeffSum = ctx.getTmpVarName(); - ctx.addToCodeBody(&arg, "double " + sum + " = 0;\ndouble " + coeffSum + "= 0;\n"); + std::string code1 = "double " + sum + " = 0;\ndouble " + coeffSum + "= 0;\n"; + + if (forceScopeIndependent) + ctx.addToCodeBody(code1, true); + else + ctx.addToCodeBody(&arg, code1); std::string iterator = "i_" + ctx.getTmpVarName(); std::string subscriptExpr = "[" + iterator + "]"; @@ -128,7 +132,10 @@ std::string realSumPdfTranslateImpl(CodegenContext &ctx, RooAbsArg const &arg, R } else if (normalize) { code += sum + " /= " + coeffSum + ";\n"; } - ctx.addToCodeBody(&arg, code); + if (forceScopeIndependent) + ctx.addToCodeBody(code, true); + else + ctx.addToCodeBody(&arg, code); return sum; } @@ -240,7 +247,7 @@ void codegenImpl(RooAbsArg &arg, CodegenContext &ctx) void codegenImpl(RooAddPdf &arg, CodegenContext &ctx) { - ctx.addResult(&arg, realSumPdfTranslateImpl(ctx, arg, arg.pdfList(), arg.coefList(), true)); + ctx.addResult(&arg, realSumPdfTranslateImpl(ctx, arg, arg.pdfList(), arg.coefList(), true, false)); } void codegenImpl(RooMultiVarGaussian &arg, CodegenContext &ctx) @@ -261,30 +268,25 @@ void codegenImpl(RooMultiPdf &arg, CodegenContext &ctx) // indices MathFunc call becomes more efficient. if (numPdfs > 2) { ctx.addResult(&arg, ctx.buildCall(mathFunc("multipdf"), arg.indexCategory(), arg.getPdfList())); + return; + } + // Ternary nested expression + std::string indexExpr = ctx.getResult(arg.indexCategory()); - std::cout << "MathFunc call used\n"; - - } else { - - // Ternary nested expression - std::string indexExpr = ctx.getResult(arg.indexCategory()); - - // int numPdfs = arg.getNumPdfs(); - std::string expr; + // int numPdfs = arg.getNumPdfs(); + std::string expr; - for (int i = 0; i < numPdfs; ++i) { - RooAbsPdf *pdf = arg.getPdf(i); - std::string pdfExpr = ctx.getResult(*pdf); + for (int i = 0; i < numPdfs; ++i) { + RooAbsPdf *pdf = arg.getPdf(i); + std::string pdfExpr = ctx.getResult(*pdf); - expr += "(" + indexExpr + " == " + std::to_string(i) + " ? (" + pdfExpr + ") : "; - } + expr += "(" + indexExpr + " == " + std::to_string(i) + " ? (" + pdfExpr + ") : "; + } - expr += "0.0"; - expr += std::string(numPdfs, ')'); // Close all ternary operators + expr += "0.0"; + expr += std::string(numPdfs, ')'); // Close all ternary operators - ctx.addResult(&arg, expr); - std::cout << "Ternary expression call used \n"; - } + ctx.addResult(&arg, expr); } // RooCategory index added. @@ -305,6 +307,7 @@ void codegenImpl(RooAddition &arg, CodegenContext &ctx) { if (arg.list().empty()) { ctx.addResult(&arg, "0.0"); + return; } std::string result; if (arg.list().size() > 1) @@ -469,7 +472,6 @@ void codegenImpl(RooFit::Detail::RooNLLVarNew &arg, CodegenContext &ctx) std::string weightSumName = RooFit::Detail::makeValidVarName(arg.GetName()) + "WeightSum"; std::string resName = RooFit::Detail::makeValidVarName(arg.GetName()) + "Result"; - ctx.addResult(&arg, resName); ctx.addToGlobalScope("double " + weightSumName + " = 0.0;\n"); ctx.addToGlobalScope("double " + resName + " = 0.0;\n"); @@ -496,6 +498,8 @@ void codegenImpl(RooFit::Detail::RooNLLVarNew &arg, CodegenContext &ctx) std::string expected = ctx.getResult(*arg.expectedEvents()); ctx.addToCodeBody(resName + " += " + expected + " - " + weightSumName + " * std::log(" + expected + ");\n"); } + + ctx.addResult(&arg, resName); } void codegenImpl(RooFit::Detail::RooNormalizedPdf &arg, CodegenContext &ctx) @@ -620,7 +624,6 @@ void codegenImpl(RooRealIntegral &arg, CodegenContext &ctx) ss << "double " << obsName << "[1];\n"; std::string resName = RooFit::Detail::makeValidVarName(arg.GetName()) + "Result"; - ctx.addResult(&arg, resName); ctx.addToGlobalScope("double " + resName + " = 0.0;\n"); // TODO: once Clad has support for higher-order functions (follow also the @@ -640,16 +643,18 @@ void codegenImpl(RooRealIntegral &arg, CodegenContext &ctx) << "}\n"; ctx.addToGlobalScope(ss.str()); + + ctx.addResult(&arg, resName); } void codegenImpl(RooRealSumFunc &arg, CodegenContext &ctx) { - ctx.addResult(&arg, realSumPdfTranslateImpl(ctx, arg, arg.funcList(), arg.coefList(), false)); + ctx.addResult(&arg, realSumPdfTranslateImpl(ctx, arg, arg.funcList(), arg.coefList(), false, false)); } void codegenImpl(RooRealSumPdf &arg, CodegenContext &ctx) { - ctx.addResult(&arg, realSumPdfTranslateImpl(ctx, arg, arg.funcList(), arg.coefList(), false)); + ctx.addResult(&arg, realSumPdfTranslateImpl(ctx, arg, arg.funcList(), arg.coefList(), false, false)); } void codegenImpl(RooRealVar &arg, CodegenContext &ctx) @@ -894,7 +899,7 @@ std::string codegenIntegralImpl(RooPolynomial &arg, int, const char *rangeName, std::string codegenIntegralImpl(RooRealSumPdf &arg, int code, const char *rangeName, CodegenContext &ctx) { // Re-use translate, since integration is linear. - return realSumPdfTranslateImpl(ctx, arg, arg.funcIntListFromCache(code, rangeName), arg.coefList(), false); + return realSumPdfTranslateImpl(ctx, arg, arg.funcIntListFromCache(code, rangeName), arg.coefList(), false, true); } std::string codegenIntegralImpl(RooUniform &arg, int code, const char *rangeName, CodegenContext &) @@ -904,5 +909,4 @@ std::string codegenIntegralImpl(RooUniform &arg, int code, const char *rangeName return doubleToString(arg.analyticalIntegral(code, rangeName)); } -} // namespace Experimental -} // namespace RooFit +} // namespace RooFit::Experimental diff --git a/roofit/roofitcore/inc/RooFit/CodegenContext.h b/roofit/roofitcore/inc/RooFit/CodegenContext.h index fbac603da8c53..3f8a1b54d06c6 100644 --- a/roofit/roofitcore/inc/RooFit/CodegenContext.h +++ b/roofit/roofitcore/inc/RooFit/CodegenContext.h @@ -30,8 +30,7 @@ template class RooTemplateProxy; -namespace RooFit { -namespace Experimental { +namespace RooFit::Experimental { template struct Prio { @@ -243,7 +242,6 @@ void declareDispatcherCode(std::string const &funcName); void codegen(RooAbsArg &arg, CodegenContext &ctx); -} // namespace Experimental -} // namespace RooFit +} // namespace RooFit::Experimental #endif diff --git a/roofit/roofitcore/src/RooFit/CodegenContext.cxx b/roofit/roofitcore/src/RooFit/CodegenContext.cxx index 7ffae9f3b01c8..df3664cb63988 100644 --- a/roofit/roofitcore/src/RooFit/CodegenContext.cxx +++ b/roofit/roofitcore/src/RooFit/CodegenContext.cxx @@ -203,27 +203,8 @@ std::string CodegenContext::getTmpVarName() const /// @param valueToSave The actual string value to save as a temporary. void CodegenContext::addResult(RooAbsArg const *in, std::string const &valueToSave) { - std::string savedName; - - // Only save values if they contain operations or they are numerals. Otherwise, we can use them directly. - - // Check if string is numeric. - char *end; - std::strtod(valueToSave.c_str(), &end); - bool isNumeric = (*end == '\0'); - - const bool hasOperations = valueToSave.find_first_of(":-+/*") != std::string::npos; - - // If the name is not empty and this value is worth saving, save it to the correct scope. - // otherwise, just return the actual value itself - if (hasOperations || isNumeric) { - savedName = "wksp[" + std::to_string(_nWksp++) + "]"; - std::string outVarDecl = savedName + " = " + valueToSave + ";\n"; - addToCodeBody(in, outVarDecl); - } else { - savedName = valueToSave; - } - + std::string savedName = "wksp[" + std::to_string(_nWksp++) + "]"; + addToCodeBody(in, savedName + " = " + valueToSave + ";\n"); _nodeNames[in->namePtr()] = savedName; } From f3050ef4c7d02687dae4a2529bab5e0812b03b40 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Thu, 9 Oct 2025 01:21:51 +0200 Subject: [PATCH 5/6] [RF] Only cache by index in codegen --- roofit/roofitcore/inc/RooFit/CodegenContext.h | 10 ++--- .../roofitcore/src/RooFit/CodegenContext.cxx | 37 +++++++++---------- roofit/roofitcore/test/testRooFuncWrapper.cxx | 2 + 3 files changed, 24 insertions(+), 25 deletions(-) diff --git a/roofit/roofitcore/inc/RooFit/CodegenContext.h b/roofit/roofitcore/inc/RooFit/CodegenContext.h index 3f8a1b54d06c6..7d3642ea66387 100644 --- a/roofit/roofitcore/inc/RooFit/CodegenContext.h +++ b/roofit/roofitcore/inc/RooFit/CodegenContext.h @@ -46,10 +46,10 @@ class CodegenContext { public: void addResult(RooAbsArg const *key, std::string const &value); - std::string const &getResult(RooAbsArg const &arg); + std::string getResult(RooAbsArg const &arg); template - std::string const &getResult(RooTemplateProxy const &key) + std::string getResult(RooTemplateProxy const &key) { return getResult(key.arg()); } @@ -134,10 +134,8 @@ class CodegenContext { }; ScopeRAII OutputScopeRangeComment(RooAbsArg const *arg) { return {arg, *this}; } - bool isScopeIndependent(RooAbsArg const *in) const; - /// @brief Map of node names to their result strings. - std::unordered_map _nodeNames; + std::unordered_map _nodeNames; std::size_t _nWksp = 0; std::unordered_map _paramIndices; /// @brief A map to keep track of the observable indices if they are non scalar. @@ -149,6 +147,8 @@ class CodegenContext { template std::string buildArgSpanImpl(std::span arr); + bool isScopeIndependent(RooAbsArg const *in) const; + void endLoop(LoopScope const &scope); template {}, bool>::type = true> diff --git a/roofit/roofitcore/src/RooFit/CodegenContext.cxx b/roofit/roofitcore/src/RooFit/CodegenContext.cxx index df3664cb63988..10d15adaa08a8 100644 --- a/roofit/roofitcore/src/RooFit/CodegenContext.cxx +++ b/roofit/roofitcore/src/RooFit/CodegenContext.cxx @@ -33,7 +33,7 @@ namespace Experimental { /// existing code body. /// @param key The node to get the result string for. /// @return String representing the result of this node. -std::string const &CodegenContext::getResult(RooAbsArg const &arg) +std::string CodegenContext::getResult(RooAbsArg const &arg) { // If the result has already been recorded, just return the result. // It is usually the responsibility of each translate function to assign @@ -41,15 +41,19 @@ std::string const &CodegenContext::getResult(RooAbsArg const &arg) // for a particular node, it means the node has already been 'translate'd and we // dont need to visit it again. auto found = _nodeNames.find(arg.namePtr()); - if (found != _nodeNames.end()) - return found->second; + std::size_t idx = 0; + if (found != _nodeNames.end()) { + idx = found->second; + } else { - auto RAII(OutputScopeRangeComment(&arg)); + auto RAII(OutputScopeRangeComment(&arg)); - // Now, recursively call translate into the current argument to load the correct result. - codegen(const_cast(arg), *this); + // Now, recursively call translate into the current argument to load the correct result. + codegen(const_cast(arg), *this); - return _nodeNames.at(arg.namePtr()); + idx = _nodeNames.at(arg.namePtr()); + } + return "wksp[" + std::to_string(idx) + "]"; } /// @brief Adds the given string to the string block that will be emitted at the top of the squashed function. Useful @@ -167,14 +171,14 @@ std::unique_ptr CodegenContext::beginLoop(RooAbsArg c std::size_t n = it.second.second; - auto savedName = "wksp[" + std::to_string(_nWksp++) + "]"; + auto savedName = "wksp[" + std::to_string(_nWksp) + "]"; std::string outVarDecl; if (n == 1) outVarDecl = savedName + " = obs[" + std::to_string(it.second.first) + "];\n"; else outVarDecl = savedName + " = obs[" + std::to_string(it.second.first) + " + " + idx + "];\n"; addToCodeBody(outVarDecl, n == 1); - _nodeNames[it.first] = savedName; + _nodeNames[it.first] = _nWksp++; } return std::make_unique(*this, std::move(vars)); @@ -203,9 +207,8 @@ std::string CodegenContext::getTmpVarName() const /// @param valueToSave The actual string value to save as a temporary. void CodegenContext::addResult(RooAbsArg const *in, std::string const &valueToSave) { - std::string savedName = "wksp[" + std::to_string(_nWksp++) + "]"; - addToCodeBody(in, savedName + " = " + valueToSave + ";\n"); - _nodeNames[in->namePtr()] = savedName; + addToCodeBody(in, "wksp[" + std::to_string(_nWksp) + "]" + " = " + valueToSave + ";\n"); + _nodeNames[in->namePtr()] = _nWksp++; } /// @brief Function to save a RooListProxy as an array in the squashed code. @@ -384,20 +387,14 @@ void codegen(RooAbsArg &arg, CodegenContext &ctx) // parameters auto foundParam = ctx._paramIndices.find(&arg); if (foundParam != ctx._paramIndices.end()) { - auto savedName = "wksp[" + std::to_string(ctx._nWksp++) + "]"; - std::string outVarDecl = savedName + " = params[" + std::to_string(foundParam->second) + "];\n"; - ctx.addToCodeBody(outVarDecl, ctx.isScopeIndependent(&arg)); - ctx.addResult(&arg, savedName); + ctx.addResult(&arg, "params[" + std::to_string(foundParam->second) + "]"); return; } // observables auto foundObs = ctx._vecObsIndices.find(arg.namePtr()); if (foundObs != ctx._vecObsIndices.end()) { - auto savedName = "wksp[" + std::to_string(ctx._nWksp++) + "]"; - std::string outVarDecl = savedName + " = obs[" + std::to_string(foundObs->second.first) + "];\n"; - ctx.addToCodeBody(outVarDecl, ctx.isScopeIndependent(&arg)); - ctx.addResult(&arg, savedName); + ctx.addResult(&arg, "obs[" + std::to_string(foundObs->second.first) + "]"); return; } diff --git a/roofit/roofitcore/test/testRooFuncWrapper.cxx b/roofit/roofitcore/test/testRooFuncWrapper.cxx index 095828c98741d..77ae003d18f69 100644 --- a/roofit/roofitcore/test/testRooFuncWrapper.cxx +++ b/roofit/roofitcore/test/testRooFuncWrapper.cxx @@ -163,6 +163,8 @@ TEST_P(FactoryTest, NLLFit) // this, we make sure to validate also the NLL values of the generated code. static_cast(*nllFunc).setUseGeneratedFunctionCode(true); + RooFit::Experimental::writeCodegenDebugMacro(*nllFunc, _params._name); + double tol = _params._fitResultTolerance; EXPECT_NEAR(nllRef->getVal(observables), nllFunc->getVal(), tol); From 4ac487649944ffcd77967b5c1a16cdc0fe357ce2 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Thu, 9 Oct 2025 02:27:39 +0200 Subject: [PATCH 6/6] finish --- roofit/roofitcore/inc/RooFit/CodegenContext.h | 2 -- roofit/roofitcore/inc/RooFit/Detail/MathFuncs.h | 17 +++++++++++++++++ roofit/roofitcore/src/RooFit/CodegenContext.cxx | 17 ++++++++--------- 3 files changed, 25 insertions(+), 11 deletions(-) diff --git a/roofit/roofitcore/inc/RooFit/CodegenContext.h b/roofit/roofitcore/inc/RooFit/CodegenContext.h index 7d3642ea66387..ea299ee7510c0 100644 --- a/roofit/roofitcore/inc/RooFit/CodegenContext.h +++ b/roofit/roofitcore/inc/RooFit/CodegenContext.h @@ -203,8 +203,6 @@ class CodegenContext { unsigned _indent = 0; /// @brief Index to get unique names for temporary variables. mutable int _tmpVarIdx = 0; - /// @brief A map to keep track of list names as assigned by addResult. - std::unordered_map::Value_t, std::string> _listNames; std::vector _xlArr; std::vector _collectedFunctions; }; diff --git a/roofit/roofitcore/inc/RooFit/Detail/MathFuncs.h b/roofit/roofitcore/inc/RooFit/Detail/MathFuncs.h index adf2169500654..28a900e88a778 100644 --- a/roofit/roofitcore/inc/RooFit/Detail/MathFuncs.h +++ b/roofit/roofitcore/inc/RooFit/Detail/MathFuncs.h @@ -812,7 +812,23 @@ double stepFunctionIntegral(double xmin, double xmax, std::size_t nBins, DoubleA } // namespace RooFit::Detail::MathFuncs +inline void fillFromWorkspace(double *out, std::size_t n, double const *wksp, double const *idx) +{ + for (std::size_t i = 0; i < n; ++i) { + out[i] += wksp[static_cast(idx[i])]; + } +} + namespace clad::custom_derivatives { + +inline void fillFromWorkspace_pullback(double *, std::size_t n, double const *, double const *idx, double *d_out, + std::size_t *, double *d_wksp, double *) +{ + for (std::size_t i = 0; i < n; ++i) { + d_wksp[static_cast(idx[i])] += d_out[i]; + } +} + namespace RooFit::Detail::MathFuncs { // Clad can't generate the pullback for binNumber because of the @@ -826,6 +842,7 @@ void binNumber_pullback(Types...) } } // namespace RooFit::Detail::MathFuncs + } // namespace clad::custom_derivatives #endif diff --git a/roofit/roofitcore/src/RooFit/CodegenContext.cxx b/roofit/roofitcore/src/RooFit/CodegenContext.cxx index 10d15adaa08a8..7d7af3bdb7610 100644 --- a/roofit/roofitcore/src/RooFit/CodegenContext.cxx +++ b/roofit/roofitcore/src/RooFit/CodegenContext.cxx @@ -220,25 +220,24 @@ std::string CodegenContext::buildArg(RooAbsCollection const &in, std::string con return "nullptr"; } - auto it = _listNames.find(in.uniqueId().value()); - if (it != _listNames.end()) - return it->second; - std::string savedName = getTmpVarName(); bool canSaveOutside = true; + std::vector indices; + indices.reserve(in.size()); + std::stringstream declStrm; - declStrm << arrayType << " " << savedName << "[]{"; + declStrm << arrayType << " " << savedName << "[" << in.size() << "]{};\n"; for (const auto arg : in) { - declStrm << getResult(*arg) << ","; + getResult(*arg); // fill the result cache + indices.push_back(_nodeNames.at(arg->namePtr())); canSaveOutside = canSaveOutside && isScopeIndependent(arg); } - declStrm.seekp(-1, declStrm.cur); - declStrm << "};\n"; + + declStrm << "fillFromWorkspace(" << savedName << ", " << in.size() << ", wksp, " << buildArg(indices) << ");\n"; addToCodeBody(declStrm.str(), canSaveOutside); - _listNames.insert({in.uniqueId().value(), savedName}); return savedName; }