Add code for ATLAS RooFit AD benchmarks

root/roofit/atlas-benchmarks/CMakeLists.txt

@@ -10,3 +10,6 @@ RB_ADD_GBENCHMARK(roofitAtlasHiggsBenchmark
 
 file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/download_workspaces.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/)
 file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/run_benchmarks.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/)
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/atlasHiggsBackendComparison.C DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/)
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/atlasHiggsBackendComparison_make_plot.py DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/)
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/run_atlasHiggsBackendComparison.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/)

root/roofit/atlas-benchmarks/README.md

@@ -0,0 +1,63 @@
+# ATLAS Higgs benchmarks
+
+Benchmarks running RooFit on the ATLAS Higgs combination workspaces published at
+https://root.cern/files/rootbench/atlas-higgs-workspaces-2021/.
+
+Two things live here:
+
+- `roofitAtlasHiggsBenchmark.cxx` — a Google Benchmark executable, driven by
+  `run_benchmarks.sh`, reporting NLL creation and minimization times.
+- `atlasHiggsBackendComparison.C` — a ROOT macro that fits the VHbb workspace
+  once per evaluation backend, used for the bar plot below.
+
+## Evaluation backend bar plot
+
+Shows where the wall time of a fit goes for each RooFit evaluation backend, as a
+stacked bar per backend: NLL creation, function JIT, gradient generation,
+gradient lowering, seeding, minimization, and the unaccounted remainder.
+
+```bash
+./download_workspaces.sh                          # once; only WS-VHbb-STXS_mu_toy_new.root is needed
+./run_atlasHiggsBackendComparison.sh              # ~1 min per backend, writes atlas/{cpu,codegen}.log
+python3 atlasHiggsBackendComparison_make_plot.py  # writes plot_roofit_ad_atlas_root.pdf
+```
+
+Both scripts use the working directory, so run them from where the workspace was
+downloaded.
+
+Notes:
+
+- The timings are scraped out of ROOT's own log output, so the macro must keep
+  printing at `kInfo`. `atlasHiggsBackendComparison.C` lists the exact lines the
+  plot script looks for.
+- The `time` builtin has to stay wrapped around `root` itself: its `user` line is
+  the total bar height, and the measured steps are subtracted from it to get the
+  "Other" slice.
+- Out of the box this gives two bars, `cpu` and `codegen`. The `legacy` backend
+  no longer exists in recent ROOT; on an older ROOT, add it to the loop in
+  `run_atlasHiggsBackendComparison.sh` to get a third bar. The plot script drops
+  any backend whose log is absent, so no other change is needed.
+- The plot script takes the log directory as an argument, so several runs can be
+  compared: `python3 atlasHiggsBackendComparison_make_plot.py atlas some-other-run`
+  writes one PDF per directory, named after it.
+
+## Machine-readable output
+
+`--json` writes the parsed timings to stdout instead of plotting, as one record
+per (run, backend, step):
+
+```bash
+python3 atlasHiggsBackendComparison_make_plot.py atlas --json > atlas.json
+```
+
+```python
+import pandas as pd
+
+df = pd.read_json("atlas.json")                       # columns: run, backend, step, seconds
+df.groupby(["run", "backend"]).seconds.sum()          # total per bar
+df.pivot_table(index="backend", columns="step", values="seconds")
+```
+
+Pass several directories to get them all in one file, already tagged by `run`,
+ready to concatenate across runs. Progress messages go to stderr, so stdout stays
+valid JSON, and `--json` does not import ROOT at all.

root/roofit/atlas-benchmarks/atlasHiggsBackendComparison.C

@@ -0,0 +1,94 @@
+#include <RooCategory.h>
+#include <RooWorkspace.h>
+#include <RooRealVar.h>
+#include <RooRandom.h>
+#include <RooSimultaneous.h>
+#include <RooUniform.h>
+#include <RooMinimizer.h>
+#include <RooStats/ModelConfig.h>
+#include <Minuit2/GradientCalculator.h>
+
+#include <TFile.h>
+#include <TRandom.h>
+#include <TError.h>
+
+// Driven once per evaluation backend by run_atlasHiggsBackendComparison.sh:
+//
+//     sh -c "{ time root -q -b 'atlasHiggsBackendComparison.C(\"codegen\")'; }" &> atlas/codegen.log
+//
+// The resulting atlas/*.log are what atlasHiggsBackendComparison_make_plot.py
+// parses.
+// It looks for exactly these lines, so anything that changes them changes the plot:
+//
+//     Creation of NLL object took ...           (all backends)
+//     Function JIT time: ...                    (codegen only)
+//     Gradient generation time: ...             (codegen only)
+//     Gradient IR to machine code time: ...     (codegen only)
+//     MnSeedGenerator Evaluated function and gradient in ...
+//     NegativeG2LineSearch Done after ...
+//     VariableMetricBuilder Stop iterating after ...
+//     user<TAB>...                              (from the `time` around root)
+//
+// The "user" total is the bar height, and every step above is subtracted from it
+// to get the "Other" slice. So doHessian and writeDebugMacros default to off:
+// both are codegen-only side studies whose minutes would otherwise land in "Other".
+void atlasHiggsBackendComparison(const char *evalBackend = "cpu", bool doHessian = false, bool writeDebugMacros = false)
+{
+   using namespace RooFit;
+
+   const bool isCodegen = std::string{evalBackend} == "codegen";
+
+   std::cout << "EvalBackend: " << evalBackend << std::endl;
+
+   gErrorIgnoreLevel = kInfo;
+   RooMsgService::instance().getStream(1).removeTopic(RooFit::Minimization);
+   RooMsgService::instance().getStream(1).removeTopic(RooFit::NumIntegration);
+   RooMsgService::instance().getStream(1).removeTopic(RooFit::Eval);
+
+   // Fetched by download_workspaces.sh into the working directory.
+   std::string workspaceFile = "WS-VHbb-STXS_mu_toy_new.root";
+   std::string workspaceName = "combined";
+
+   std::unique_ptr<TFile> tfile{TFile::Open(workspaceFile.c_str())};
+   RooWorkspace *ws = tfile->Get<RooWorkspace>(workspaceName.c_str());
+   auto mc = static_cast<RooStats::ModelConfig *>(ws->obj("ModelConfig"));
+
+   RooAbsPdf *pdf = mc->GetPdf();
+
+   RooArgSet const *globObs = mc->GetGlobalObservables();
+   RooAbsData *data = ws->data("toyData");
+
+   //ROOT::Minuit2::GradientCalculator::DoParallelOMP(false);
+   //ROOT::Minuit2::GradientCalculator::DoParallelOMP(true);
+
+   std::unique_ptr<RooAbsReal> nll{pdf->createNLL(
+         *data, GlobalObservables(*globObs), Offset(true), Optimize(2), EvalBackend(evalBackend))};
+
+   // Only the codegen NLL is a RooEvaluatorWrapper, so this cast is not valid
+   // for the legacy and cpu backends.
+   if (isCodegen && doHessian) {
+      static_cast<RooFit::Experimental::RooEvaluatorWrapper &>(*nll).generateHessian();
+   }
+
+   double val = nll->getVal();
+
+   std::cout << "Initial value: " << val << std::endl;
+
+   if (isCodegen && writeDebugMacros) {
+      RooFit::Experimental::writeCodegenDebugMacro(*nll, "debug_macro");
+   }
+
+   RooMinimizer::Config cfg;
+   cfg.useGradient = true;
+   cfg.useHessian = isCodegen && doHessian;
+
+   RooMinimizer minim{*nll, cfg};
+   minim.setStrategy(0);
+
+   minim.minimize("Minuit2", "MIGRAD");
+
+   if (doHessian) {
+      minim.setStrategy(3);
+      minim.hesse();
+   }
+}