From 7f0a73e9e9b3082faab97177af2857748707733e Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Sun, 30 Aug 2026 11:22:30 +0200 Subject: [PATCH] Add code for ATLAS RooFit AD benchmarks --- root/roofit/atlas-benchmarks/CMakeLists.txt | 3 + root/roofit/atlas-benchmarks/README.md | 63 ++++ .../atlasHiggsBackendComparison.C | 94 ++++++ .../atlasHiggsBackendComparison_make_plot.py | 273 ++++++++++++++++++ .../atlas-benchmarks/download_workspaces.sh | 2 +- .../run_atlasHiggsBackendComparison.sh | 19 ++ .../roofit/atlas-benchmarks/run_benchmarks.sh | 2 +- 7 files changed, 454 insertions(+), 2 deletions(-) create mode 100644 root/roofit/atlas-benchmarks/README.md create mode 100644 root/roofit/atlas-benchmarks/atlasHiggsBackendComparison.C create mode 100644 root/roofit/atlas-benchmarks/atlasHiggsBackendComparison_make_plot.py create mode 100755 root/roofit/atlas-benchmarks/run_atlasHiggsBackendComparison.sh diff --git a/root/roofit/atlas-benchmarks/CMakeLists.txt b/root/roofit/atlas-benchmarks/CMakeLists.txt index 25b17178..efad276c 100644 --- a/root/roofit/atlas-benchmarks/CMakeLists.txt +++ b/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}/) diff --git a/root/roofit/atlas-benchmarks/README.md b/root/roofit/atlas-benchmarks/README.md new file mode 100644 index 00000000..2b08cd92 --- /dev/null +++ b/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. diff --git a/root/roofit/atlas-benchmarks/atlasHiggsBackendComparison.C b/root/roofit/atlas-benchmarks/atlasHiggsBackendComparison.C new file mode 100644 index 00000000..cc6c9dfe --- /dev/null +++ b/root/roofit/atlas-benchmarks/atlasHiggsBackendComparison.C @@ -0,0 +1,94 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +// 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... (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::Open(workspaceFile.c_str())}; + RooWorkspace *ws = tfile->Get(workspaceName.c_str()); + auto mc = static_cast(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 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(*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(); + } +} diff --git a/root/roofit/atlas-benchmarks/atlasHiggsBackendComparison_make_plot.py b/root/roofit/atlas-benchmarks/atlasHiggsBackendComparison_make_plot.py new file mode 100644 index 00000000..b39c8d2d --- /dev/null +++ b/root/roofit/atlas-benchmarks/atlasHiggsBackendComparison_make_plot.py @@ -0,0 +1,273 @@ +#!/usr/bin/env python +# coding: utf-8 + +# https://github.com/root-project/root/pull/17926 + +import argparse +import json +import os +import sys + +import numpy as np + +parser = argparse.ArgumentParser( + description="Plot the timing breakdown of the RooFit evaluation backends." +) +parser.add_argument( + "experiments", + nargs="*", + default=["atlas"], + metavar="DIR", + help="directory holding {legacy,cpu,codegen}.log (default: atlas)", +) +parser.add_argument( + "--json", + action="store_true", + help="write the parsed timings to stdout as JSON instead of plotting; the " + "result is a tidy table that pandas reads directly with pd.read_json()", +) +# Parsed before importing ROOT, so that --help stays instant and PyROOT never +# gets a look at sys.argv. +args = parser.parse_args() + +# Only the plot needs ROOT, so --json also works without PyROOT available. +if not args.json: + import ROOT + + # ROOT style + ROOT.gROOT.SetBatch(True) + ROOT.gStyle.SetOptStat(0) + + +def info(*args_): + """Progress output, on stderr so that --json keeps stdout to itself.""" + print(*args_, file=sys.stderr) + + +def time_string_to_seconds(s): + s = s.replace("ms", "*1e-3") + s = s.replace("s", "") + s = s.replace("min", "*60+") + s = s.replace("m", "*60+") + s = s.replace("μ", "*1e-6") + + return float(eval(s)) + + +steps = { + "NLL creation": "Creation of NLL object took ", + "Function JIT": "Function JIT time: ", + "Gradient generation": "Gradient generation time: ", + "Gradient to machine code": "Gradient IR to machine code time: ", + "Seeding step": "MnSeedGenerator Evaluated function and gradient in ", + "NegativeG2LineSearch": "NegativeG2LineSearch Done after ", + "Minimization": "VariableMetricBuilder Stop iterating after ", + "Hesse": "MnHesse Done after ", + "Total": "user ", +} + + +def parse_output(filename): + """Parse output generated like + sh -c { time root -q -b 'atlasHiggsBackendComparison.C("cpu")'; } &> filename + """ + + info(f"--- {filename} ---") + + with open(filename, "r") as f: + lines = f.read().split("\n") + + relevant_lines = {} + + for line in lines: + for key, val in steps.items(): + if val in line: + time_str = line.split(val)[-1] + relevant_lines[key] = time_string_to_seconds(time_str) + + # Without the "user" line from the surrounding `time` there is no bar height + # to subtract the measured steps from, and "Other" would come out negative. + if "Total" not in relevant_lines: + info(f"{filename}: no 'user' line from `time`, dropping this backend") + return None + + # Replace Total with other: + measured = 0.0 + for key, val in relevant_lines.items(): + if key != "Total": + measured = measured + val + + # Fill irrelevant steps with zeros: + for key in steps: + if key not in relevant_lines: + relevant_lines[key] = 0.0 + + # merge NegativeG2LineSearch into Seeding step + relevant_lines["Seeding step"] += relevant_lines["NegativeG2LineSearch"] + del relevant_lines["NegativeG2LineSearch"] + + relevant_lines["Other"] = relevant_lines["Total"] - measured + del relevant_lines["Total"] + + # We don't care of measuring Hesse (for now) + del relevant_lines["Hesse"] + + for ( + k, + v, + ) in relevant_lines.items(): + info(f"{k} : {v}") + + return relevant_lines + + +# Bars appear left to right in this order. A backend whose log is absent is +# dropped rather than being an error, so a run without the legacy backend simply +# yields a two-bar plot. +backend_names = {"RooFit legacy": "legacy", "RooFit": "cpu", "RooFit AD": "codegen"} + +# The legacy backend no longer exists in recent ROOT, so out of the box this +# produces a two-bar plot. To compare ROOT versions instead of backends, relabel: +# backend_names = {"ROOT 6.32": "legacy", "ROOT 6.40": "cpu", "ROOT 6.40 AD": "codegen"} + +# One row per (run, backend, step) for --json: a tidy table stays easy to +# concatenate across runs and to pivot downstream. +records = [] + +for experiment in args.experiments: + + # The directory may be given as a path, but the plot is named after it. + tag = os.path.basename(experiment.rstrip("/")) or experiment + + parsed = {} + for backend, name in backend_names.items(): + filename = os.path.join(experiment, f"{name}.log") + if not os.path.exists(filename): + info(f"{filename}: missing, dropping the '{backend}' bar") + continue + result = parse_output(filename) + if result is not None: + parsed[backend] = result + + # Only the backends that actually produced a log, in the order above. + backends = list(parsed) + + if not backends: + info(f"{experiment}: no usable logs, skipping") + continue + + if args.json: + # Every backend carries the same steps (the ones it does not have are + # zero-filled), so the first one fixes a stable step order. + for backend in backends: + for step in parsed[backends[0]]: + records.append( + { + "run": tag, + "backend": backend, + "step": step, + "seconds": parsed[backend][step], + } + ) + continue + + data = {} + for key in parsed[backends[0]]: + data[key] = np.array([parsed[backend][key] for backend in backends]) + + jit_color = np.array([0.06, 0.21, 0.70]) + + colors = { + "NLL creation": [0.30, 0.70, 0.50], + "Function JIT": (1.0 - 0.3 * (1.0 - jit_color)), + "Gradient generation": (1.0 - 0.6 * (1.0 - jit_color)), + "Gradient to machine code": (1.0 - (1.0 - jit_color)), + "Seeding step": [0.70, 0.05, 0.70], + "Minimization": [1.0, 0.2, 0.2], + "Hesse": [1.0, 0.9, 0.3], + "Other": [0.8, 0.8, 0.8], + } + + # Convert color float to ROOT TColor index + root_colors = {} + for label, (r, g, b) in colors.items(): + root_colors[label] = ROOT.TColor.GetColor(int(r * 255), int(g * 255), int(b * 255)) + + # ROOT histogram stack + stack = ROOT.THStack("stack", "") + histos = {} + + nbins = len(backends) + + for label, values in data.items(): + h = ROOT.TH1F(label, label, nbins*2 + 1, 0, nbins*2 + 1) + + for i, v in enumerate(values): + h.SetBinContent(2*i+1 + 1, v) # fill every second bin + + h.SetFillColor(root_colors[label]) + + # invisible outline + h.SetLineColor(0) + h.SetLineWidth(0) + + histos[label] = h + + stack.Add(h) + + # Legend geometry in NDC, also used below to keep bars out from under it. + leg_x1, leg_y1, leg_x2, leg_y2 = 0.40, 0.55, 0.90, 0.88 + + # Canvas + c = ROOT.TCanvas(f"c_{tag}", "c", 2 * 800, 2 * 500) + + # Draw stack + stack.Draw("hist") + + # The legend is filled opaque, so a bar drawn underneath it is simply hidden + # above its lower edge. That does not show up while the tallest backend sits + # to the left of the legend, but as soon as it is dropped the axis rescales + # and the remaining bars run into the legend. Give the y axis enough headroom + # that every bar the legend spans horizontally stays below it. + margin_lo, margin_hi = 0.1, 0.9 # default pad margins, same on both axes + totals = np.array(list(data.values())).sum(axis=0) + bar_x = [ + margin_lo + (2 * i + 1.5) / (2 * nbins + 1) * (margin_hi - margin_lo) + for i in range(nbins) + ] + under_legend = [t for t, x in zip(totals, bar_x) if leg_x1 <= x <= leg_x2] + if under_legend: + needed = max(under_legend) * (margin_hi - margin_lo) / (leg_y1 - margin_lo) + if needed > 1.05 * totals.max(): # otherwise ROOT's own scaling is fine + stack.SetMaximum(needed) + + stack.GetYaxis().SetTitle("Time [s]") + # stack.GetXaxis().SetTitle("Evaluation backend") + xaxis = stack.GetXaxis() + xaxis.SetLabelSize(0.05) + # xaxis.RotateTitle(False) + xaxis.SetTitleOffset(1.5) # closer to the axis + + # Set backend labels on x-axis + axis = stack.GetXaxis() + for i, backend in enumerate(backends): + axis.SetBinLabel(2*i+1 + 1, backend) + + # Legend + legend = ROOT.TLegend(leg_x1, leg_y1, leg_x2, leg_y2) + legend.SetNColumns(2) + legend.SetBorderSize(0) + for label, h in histos.items(): + legend.AddEntry(h, label, "f") + legend.SetBorderSize(0) + legend.Draw() + + # Optional grid + c.SetGrid() + + # c.SaveAs(f"plot_{tag}_root.png") + c.SaveAs(f"plot_roofit_ad_{tag}_root.pdf") + +if args.json: + json.dump(records, sys.stdout, indent=2) + sys.stdout.write("\n") diff --git a/root/roofit/atlas-benchmarks/download_workspaces.sh b/root/roofit/atlas-benchmarks/download_workspaces.sh index e5a6d326..d1b3c30f 100755 --- a/root/roofit/atlas-benchmarks/download_workspaces.sh +++ b/root/roofit/atlas-benchmarks/download_workspaces.sh @@ -1,4 +1,4 @@ -#!/bin/bash +#!/usr/bin/env bash BASE_URL=https://root.cern/files/rootbench diff --git a/root/roofit/atlas-benchmarks/run_atlasHiggsBackendComparison.sh b/root/roofit/atlas-benchmarks/run_atlasHiggsBackendComparison.sh new file mode 100755 index 00000000..c7727b33 --- /dev/null +++ b/root/roofit/atlas-benchmarks/run_atlasHiggsBackendComparison.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +# Times the RooFit evaluation backends on the ATLAS Higgs VHbb workspace and +# writes one log per backend into atlas/, the input of +# atlasHiggsBackendComparison_make_plot.py. +# +# Run download_workspaces.sh first: the macro opens the workspace from the +# working directory. +# +# The `time` builtin has to wrap root itself, because the plot script reads its +# "user" line as the total bar height. +# +# The legacy backend is not run here, as it no longer exists in recent ROOT. +# Add it to the loop below on an older ROOT to get a third bar. + +mkdir -p atlas + +for backend in cpu codegen; do + sh -c "{ time root -q -b 'atlasHiggsBackendComparison.C(\"$backend\")'; }" &> "atlas/$backend.log" +done diff --git a/root/roofit/atlas-benchmarks/run_benchmarks.sh b/root/roofit/atlas-benchmarks/run_benchmarks.sh index 840e5853..46d1c2a5 100755 --- a/root/roofit/atlas-benchmarks/run_benchmarks.sh +++ b/root/roofit/atlas-benchmarks/run_benchmarks.sh @@ -1,4 +1,4 @@ -#!/bin/bash +#!/usr/bin/env bash ./roofitAtlasHiggsBenchmark 0 ./roofitAtlasHiggsBenchmark 1