diff --git a/.github/workflows/Benchmark.yml b/.github/workflows/Benchmark.yml new file mode 100644 index 00000000..0003cf19 --- /dev/null +++ b/.github/workflows/Benchmark.yml @@ -0,0 +1,62 @@ +name: Benchmark +# Performance guard for force/energy evaluation. +# +# Runs the PkgBenchmark suite in benchmark/benchmarks.jl on the PR head and posts +# the results to the job summary. Non-blocking (continue-on-error) so noisy shared +# runners never block a merge. +# +# NOTE: comparison against the base branch (PkgBenchmark `judge`) is intentionally +# NOT enabled yet — the benchmark suite does not exist on `main` until this lands. +# Once merged, a follow-up can switch the run step to `judge` against the base ref +# (e.g. via BenchmarkCI) to get automatic regression detection on the same runner. +on: + pull_request: + workflow_dispatch: +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +jobs: + benchmark: + name: Force/energy benchmark (non-blocking) + runs-on: ubuntu-latest + timeout-minutes: 60 + continue-on-error: true # keep non-blocking + permissions: + actions: write + contents: read + env: + JULIA_NUM_THREADS: '1' # classic path is threaded via Folds; pin for stable timings + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: '1.12' + arch: x64 + - name: Add registries + run: | + using Pkg + Pkg.Registry.add("General") + Pkg.Registry.add(Pkg.RegistrySpec(url = "https://github.com/ACEsuit/ACEregistry")) + shell: bash -c "julia --color=yes {0}" + - uses: julia-actions/cache@v2 + - name: Instantiate benchmark env + run: julia --color=yes --project=benchmark -e 'using Pkg; Pkg.instantiate()' + - name: Run benchmark suite + # Must run from the benchmark project: PkgBenchmark propagates the active + # project to the benchmark subprocess, so benchmark/Project.toml (which has + # ACEpotentials, BenchmarkTools, PkgBenchmark, …) must be active here. + run: | + using PkgBenchmark + results = benchmarkpkg(".") # uses benchmark/benchmarks.jl + benchmark/Project.toml + export_markdown("benchmark_results.md", results) + println(read("benchmark_results.md", String)) + shell: bash -c "julia --color=yes --project=benchmark {0}" + - name: Post results to job summary + if: always() + run: | + echo "## Force/energy benchmark results" >> "$GITHUB_STEP_SUMMARY" + if [ -f benchmark_results.md ]; then + cat benchmark_results.md >> "$GITHUB_STEP_SUMMARY" + else + echo "_Benchmark run did not produce results (see job log)._" >> "$GITHUB_STEP_SUMMARY" + fi diff --git a/benchmark/FORCE_REGRESSION_FINDINGS.md b/benchmark/FORCE_REGRESSION_FINDINGS.md new file mode 100644 index 00000000..b2c8115a --- /dev/null +++ b/benchmark/FORCE_REGRESSION_FINDINGS.md @@ -0,0 +1,69 @@ +# Force-evaluation performance regression — investigation & fix + +## TL;DR + +The v0.10 (EquivariantTensors) series had a **~13× force-evaluation slowdown** vs +the previous release (0.9.1) for the **default `ace1_model` / classic `ACEModel`**. + +- The original hypothesis — that the slowdown came from the ET backend computing + forces via **autograd** — is **refuted**. The ET autograd path is actually + *faster* than the classic analytic path was. +- The real cause is a **type instability** in the classic analytic `evaluate_ed` + (`src/models/ace.jl`): `EquivariantTensors.pullback` returns `Tuple{Any,Any}`, + so the inline gradient-assembly loops ran with per-element dynamic dispatch. +- **Fix:** a one-function **function barrier** (`_assemble_grad_ed!`). Restores + performance to ~0.9.1 levels; forces are bit-identical. + +## How it was found + +Benchmarks (`benchmark/`), single-threaded, parameter-matched models: + +1. **In-repo, classic vs ET** (`bench_forces_regression.jl`): ET autograd forces + were ~4× *faster* than classic analytic forces — opposite to the hypothesis. + Classic forces were ~12.7× their own energy cost (ET only ~2.1×). +2. **Cross-version** (`bench_crossversion.jl`, 0.9.1 vs dev): classic forces were + **11–14× slower** in dev than in 0.9.1. +3. **Profiling** `evaluate_ed`: numerical kernels (radial, Ylm, tensor fwd/bwd, + pair) summed to ~30 µs, but the full call took ~400 µs and allocated 764 KB. + Per-block timing pinned ~225 µs (80%) on the gradient-**assembly loops**. +4. **Isolation:** the same loops with concrete-typed inputs ran in **1.5 µs / 0 B**. + `Base.return_types(EquivariantTensors.pullback, …) == Tuple{Any,Any}` → + `∂Rnl`/`∂Ylm` reach the loop as `Any` → dynamic dispatch per element. + +## The fix + +`evaluate_ed` built `∇Ei` with the products `∂Rnl[j,t]*dRnl[j,t]*∇rs[j]` and +`∂Ylm[j,t]*dYlm[j,t]` **inline**. Because `∂Rnl`/`∂Ylm` are `Any`, every product +was dynamically dispatched. Moving the loops into `_assemble_grad_ed!` lets Julia +re-specialise on the concrete runtime types (a standard *function barrier*). + +``` +evaluate_ed, 40 neighbours: 386 µs → 30 µs (ΔE = 0, max|Δg| = 0) +``` + +## Results (single-thread) + +Classic `forces`, before vs after the fix (Si/O, order 2): + +| atoms | before | after | speedup | +|------:|-------:|------:|--------:| +| 64 | 22.0 ms | 2.69 ms | 8.2× | +| 512 | 190.8 ms | 18.3 ms | 10.4× | +| 800 | 298.3 ms | 29.2 ms | 10.2× | + +Cross-version (Si, `ace1_model`), dev/0.9.1 ratio: **~13× → ~1.15×**. + +All ACE model + calculator tests pass (finite-difference force checks included). + +## Notes / follow-ups + +- The residual ~15% vs 0.9.1 is minor (ET `pullback`/`evaluate` overhead vs the + old EquivariantModels backend) — not pursued here. +- The same `Tuple{Any,Any}` return from `EquivariantTensors.pullback` could be + fixed upstream in EquivariantTensors.jl; the function barrier is the robust + in-repo workaround. `grad_params` (fitting path) already passes the pullback + result through a function (`pb_Rnl`), so it is largely insulated. +- Regression guard: `benchmark/benchmarks.jl` (PkgBenchmark) run by the + non-blocking `.github/workflows/Benchmark.yml` job (results posted to the job + summary). Base-branch comparison (`judge`) can be enabled once the suite is on + `main`; see `benchmark/README.md`. diff --git a/benchmark/Project.toml b/benchmark/Project.toml new file mode 100644 index 00000000..9994e881 --- /dev/null +++ b/benchmark/Project.toml @@ -0,0 +1,16 @@ +[deps] +ACEpotentials = "3b96b61c-0fcc-4693-95ed-1ef9f35fcc53" +AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a" +AtomsBuilder = "f5cc8831-eeb7-4288-8d9f-d6c1ddb77004" +AtomsCalculators = "a3e0e189-c65a-42c1-833c-339540406eb1" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +EquivariantTensors = "5e107534-7145-4f8f-b06f-47a52840c895" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" +PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" diff --git a/benchmark/README.md b/benchmark/README.md new file mode 100644 index 00000000..9986be1d --- /dev/null +++ b/benchmark/README.md @@ -0,0 +1,72 @@ +# Benchmarks + +Performance benchmarks for ACEpotentials, with a focus on force evaluation. + +The benchmarks run in their own environment (`benchmark/Project.toml`, which +`develop`s the package at `..`), keeping perf-only dependencies out of the main +test path. They require the ACE registry: + +```julia +using Pkg +Pkg.Registry.add(Pkg.RegistrySpec(url="https://github.com/ACEsuit/ACEregistry")) +Pkg.activate("benchmark"); Pkg.instantiate() +``` + +## Files + +| File | Purpose | +|------|---------| +| `common.jl` | Shared model/system setup + `measure` helper. Single source of truth. | +| `bench_forces_regression.jl` | Reproduce & quantify the autograd force regression: classic (analytic) vs ET, energy & forces, with allocations + a `site_grads` isolation timing. | +| `bench_crossversion.jl` | Compare a pinned **previous release** of ACEpotentials against the current checkout (release-vs-release drift). | +| `benchmarks.jl` | PkgBenchmark `SUITE` used by CI to guard against regressions. | +| `benchmark_full_model.jl` | Pre-existing full-model ACE vs ETACE comparison (CPU/GPU). | + +## Reproduce the force regression locally + +```bash +JULIA_NUM_THREADS=1 julia --project=benchmark benchmark/bench_forces_regression.jl +``` + +The hypothesis (autograd `site_grads` dominates force cost) is confirmed when the +**force** ET/classic ratio is far larger than the **energy** ratio and ET force +allocations scale much worse than ET energy allocations. + +## Cross-version comparison + +```bash +julia --project=. benchmark/bench_crossversion.jl +``` + +This installs a previous release into a temporary, isolated environment and +benchmarks `forces` on the same systems, then compares to the current dev +checkout. Note the previous release's default `ace1_model` path was already +analytic, so this measures total release drift; `bench_forces_regression.jl` +isolates the autograd contribution within the current codebase. + +## Regression testing in CI + +`.github/workflows/Benchmark.yml` runs the PkgBenchmark suite +(`benchmark/benchmarks.jl`) on the PR head and posts the results to the job +summary. It is **non-blocking** (`continue-on-error: true`) and pins +`JULIA_NUM_THREADS=1` for stable timings. + +Comparison against the base branch (PkgBenchmark `judge`) is **not** enabled yet: +the suite does not exist on `main` until this work lands. Once merged, switch the +run step to judge the PR against the base ref on the same runner (which cancels +most hardware noise) — e.g. with +[BenchmarkCI.jl](https://github.com/tkf/BenchmarkCI.jl) — for automatic +regression detection, and promote it to a required check once the threshold is +calibrated. + +Run the suite locally: + +```julia +using PkgBenchmark +r = benchmarkpkg(".") # uses benchmark/benchmarks.jl + benchmark/Project.toml +export_markdown(stdout, r) +``` + +Sanity-check the detector by temporarily restoring the Zygote-based `site_grads` +in `src/et_models/et_ace.jl` / `et_pair.jl`; the force benchmarks should regress +sharply. diff --git a/benchmark/_crossversion_worker.jl b/benchmark/_crossversion_worker.jl new file mode 100644 index 00000000..cc44ea7d --- /dev/null +++ b/benchmark/_crossversion_worker.jl @@ -0,0 +1,22 @@ +# Worker for the cross-version force benchmark. Run inside an environment that +# has a specific ACEpotentials version installed. Prints lines: " ". +# +# Uses ONLY the stable public API (ace1_model -> ACEPotential calculator, +# AtomsCalculators.forces) and single-element Si so it works unchanged across +# the 0.9.x and 0.10.x series. +using ACEpotentials, AtomsBuilder, AtomsCalculators, Unitful, BenchmarkTools, Printf + +label = isempty(ARGS) ? "?" : ARGS[1] + +# ace1_model returns a ready-to-use ACEPotential calculator (random params). +calc = ACEpotentials.ace1_model(elements = [:Si], order = 3, + totaldegree = 10, rcut = 5.5) + +println("# version=", pkgversion(ACEpotentials), " label=", label) +for n in (2, 3, 4) # 16, 54, 128 atoms (cubic Si has 8/cell) + sys = AtomsBuilder.bulk(:Si, cubic = true) * n + rattle!(sys, 0.1u"Å") + AtomsCalculators.forces(sys, calc) # warmup / compile + t = minimum(@benchmark AtomsCalculators.forces($sys, $calc) samples = 5 evals = 3).time + @printf("%d %.6f\n", length(sys), t / 1e6) # ms +end diff --git a/benchmark/bench_crossversion.jl b/benchmark/bench_crossversion.jl new file mode 100644 index 00000000..264f5fbf --- /dev/null +++ b/benchmark/bench_crossversion.jl @@ -0,0 +1,69 @@ +# Cross-version force benchmark driver. +# +# Compares classic `ace1_model` force evaluation between a pinned PREVIOUS +# release of ACEpotentials and the CURRENT checkout, each in its own isolated +# environment, on identical Si systems. This measures release-vs-release drift +# in the default (analytic) force path. +# +# Run: julia --project=. benchmark/bench_crossversion.jl [old_version] +# (default old_version = 0.9.1, the last release before the ET-backend series) + +using Pkg, Printf + +const OLD_VERSION = isempty(ARGS) ? "0.9.1" : ARGS[1] +const REPO = dirname(@__DIR__) +const WORKER = joinpath(@__DIR__, "_crossversion_worker.jl") +const SCRATCH = mktempdir() + +worker_deps() = ["AtomsBuilder", "AtomsCalculators", "Unitful", "BenchmarkTools"] + +"Set up an isolated env with ACEpotentials pinned to `version` (or `path` for dev)." +function setup_env(dir; version = nothing, path = nothing) + Pkg.activate(dir) + if path !== nothing + Pkg.develop(path = path) + else + Pkg.add(Pkg.PackageSpec(name = "ACEpotentials", version = version)) + end + Pkg.add(worker_deps()) + Pkg.instantiate() +end + +"Run the worker in `env` and parse `natoms => ms`." +function run_worker(env, label) + out = read(`$(Base.julia_cmd()) --project=$env --threads=1 $WORKER $label`, String) + print(out) + res = Dict{Int,Float64}() + for ln in split(out, '\n') + (isempty(ln) || startswith(ln, "#")) && continue + parts = split(strip(ln)) + length(parts) == 2 || continue + res[parse(Int, parts[1])] = parse(Float64, parts[2]) + end + return res +end + +println("Cross-version force benchmark: old=$OLD_VERSION vs current dev\n") + +old_env = joinpath(SCRATCH, "old") +new_env = joinpath(SCRATCH, "new") +mkpath(old_env); mkpath(new_env) + +@info "Setting up OLD env (ACEpotentials@$OLD_VERSION)…" +setup_env(old_env; version = OLD_VERSION) +@info "Setting up NEW env (current dev checkout)…" +setup_env(new_env; path = REPO) + +@info "Benchmarking OLD…" +old = run_worker(old_env, "old-$OLD_VERSION") +@info "Benchmarking NEW…" +new = run_worker(new_env, "new-dev") + +println("\n", "="^60) +println("| Atoms | old $OLD_VERSION (ms) | new dev (ms) | new/old |") +println("|-------|--------------|--------------|---------|") +for nat in sort(collect(keys(old))) + haskey(new, nat) || continue + @printf("| %5d | %12.3f | %12.3f | %7.2f |\n", nat, old[nat], new[nat], new[nat]/old[nat]) +end +println("\n(new/old > 1 ⇒ current release is SLOWER than $OLD_VERSION for default analytic forces)") diff --git a/benchmark/bench_forces_regression.jl b/benchmark/bench_forces_regression.jl new file mode 100644 index 00000000..429775b7 --- /dev/null +++ b/benchmark/bench_forces_regression.jl @@ -0,0 +1,119 @@ +# Force-evaluation performance benchmark: classic (analytic) vs ET backend. +# +# Background: a ~13x force-evaluation regression was reported in the v0.10 (ET) +# series vs the previous release. The original hypothesis was that the ET +# backend's Zygote autograd forces (`site_grads`) were the cause. This benchmark +# (together with benchmark/bench_crossversion.jl) REFUTED that: the ET autograd +# path is actually faster than the classic path was. The real cause was a TYPE +# INSTABILITY in the classic analytic `evaluate_ed` (src/models/ace.jl): +# `EquivariantTensors.pullback` returns `Tuple{Any,Any}`, so the inline gradient +# assembly ran with per-element dynamic dispatch (~10x slower). Fixed with a +# function barrier (`_assemble_grad_ed!`). +# +# This script measures, for parameter-matched classic vs ET calculators on the +# SAME systems: potential_energy and forces (time + allocations), plus an +# isolation timing of the gradient step. It is kept as a backend comparison and +# a manual regression check (the automated guard is benchmark/benchmarks.jl). +# +# Run: julia --project=benchmark benchmark/bench_forces_regression.jl + +include("common.jl") +using Printf + +const setup = build_model_and_calcs() +const ace_calc = setup.ace_calc +const et_calc = setup.et_calc +const rcut = setup.rcut + +# --------------------------------------------------------------------------- +# 1c. Correctness gate: ET forces must match classic forces before timing. +# --------------------------------------------------------------------------- +function check_consistency() + sys = make_system(SIZE_CONFIGS[1]) + Ec = AtomsCalculators.potential_energy(sys, ace_calc) + Ee = AtomsCalculators.potential_energy(sys, et_calc) + Fc = AtomsCalculators.forces(sys, ace_calc) + Fe = AtomsCalculators.forces(sys, et_calc) + eE = abs(ustrip(Ee - Ec)) / max(abs(ustrip(Ec)), 1) + maxF = maximum(norm.(ustrip.(Fc))) + eF = maximum(norm.(ustrip.(Fe .- Fc))) / max(maxF, 1) + @printf("Consistency check (%d atoms): ΔE/|E| = %.2e max|ΔF|/max|F| = %.2e\n", + length(sys), eE, eF) + if eE > 1e-6 || eF > 1e-6 + @warn "Classic and ET results disagree beyond 1e-6 — timings of an " * + "incorrect path are meaningless." eE eF + end + return eE, eF +end + +# --------------------------------------------------------------------------- +# 1a. Head-to-head energy & forces benchmark. +# --------------------------------------------------------------------------- +function run_table() + println("\n", "="^104) + println("FORCE REGRESSION BENCHMARK: classic (analytic) vs ET (Zygote autograd)") + println("="^104, "\n") + println("| Atoms | Edges | E classic (ms) | E ET (ms) | E ratio | F classic (ms) | F ET (ms) | F ratio | F ET allocs |") + println("|-------|-------|----------------|-----------|---------|----------------|-----------|---------|-------------|") + + results = Vector{Any}() + for cfg in SIZE_CONFIGS + sys = make_system(cfg) + nat = length(sys) + ne = nedges(sys, rcut) + + eC = measure(() -> AtomsCalculators.potential_energy(sys, ace_calc)) + eE = measure(() -> AtomsCalculators.potential_energy(sys, et_calc)) + fC = measure(() -> AtomsCalculators.forces(sys, ace_calc)) + fE = measure(() -> AtomsCalculators.forces(sys, et_calc)) + + e_ratio = eE.time_ns / eC.time_ns + f_ratio = fE.time_ns / fC.time_ns + + @printf("| %5d | %5d | %14.3f | %9.3f | %7.2f | %14.3f | %9.3f | %7.2f | %11d |\n", + nat, ne, eC.time_ns/1e6, eE.time_ns/1e6, e_ratio, + fC.time_ns/1e6, fE.time_ns/1e6, f_ratio, fE.allocs) + + push!(results, (; natoms = nat, nedges = ne, + e_classic_ns = eC.time_ns, e_et_ns = eE.time_ns, e_ratio, + f_classic_ns = fC.time_ns, f_et_ns = fE.time_ns, f_ratio, + f_et_allocs = fE.allocs, f_et_mem = fE.memory, + e_et_allocs = eE.allocs, e_et_mem = eE.memory)) + end + return results +end + +# --------------------------------------------------------------------------- +# Isolation: time the gradient step itself (Zygote `site_grads`) vs the +# analytic `site_basis_jacobian`, with the graph precomputed — isolates the AD +# cost from graph construction and force assembly. +# --------------------------------------------------------------------------- +function run_isolation() + println("\nIsolation: gradient step only (graph precomputed)") + println("| Atoms | site_grads (Zygote, ms) | site_basis_jacobian (analytic, ms) | ratio |") + println("|-------|-------------------------|------------------------------------|-------|") + # Reach into the StackedCalculator to grab the ETACE many-body component. + ace_sub = et_calc.calcs[end] # ETACE WrappedSiteCalculator (last in the stack) + m, ps, st = ace_sub.model, ace_sub.ps, ace_sub.st + for cfg in SIZE_CONFIGS[1:min(3, end)] + sys = make_system(cfg) + G = ET.Atoms.interaction_graph(sys, rcut * u"Å") + tg = measure(() -> ETM.site_grads(m, G, ps, st)) + tj = measure(() -> ETM.site_basis_jacobian(m, G, ps, st)) + @printf("| %5d | %23.3f | %34.3f | %5.2f |\n", + length(sys), tg.time_ns/1e6, tj.time_ns/1e6, tg.time_ns/tj.time_ns) + end +end + +check_consistency() +results = run_table() +run_isolation() + +# --- summary --- +mean_e = sum(r.e_ratio for r in results) / length(results) +mean_f = sum(r.f_ratio for r in results) / length(results) +mean_fe_classic = sum(r.f_classic_ns / r.e_classic_ns for r in results) / length(results) +println("\nMean energy ratio (ET/classic): ", round(mean_e, digits=2)) +println("Mean force ratio (ET/classic): ", round(mean_f, digits=2)) +println("Mean classic force/energy ratio: ", round(mean_fe_classic, digits=2), + " (was ~12-13x before the evaluate_ed type-stability fix; ~1.5x after)") diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl new file mode 100644 index 00000000..ae16202f --- /dev/null +++ b/benchmark/benchmarks.jl @@ -0,0 +1,48 @@ +# PkgBenchmark suite for ACEpotentials force/energy evaluation. +# +# Used by PkgBenchmark.jl / BenchmarkCI.jl to compare a PR against its base +# branch on the same runner (the comparison cancels most hardware noise). +# +# Local run: +# julia --project=benchmark -e 'using PkgBenchmark; \ +# r = benchmarkpkg("ACEpotentials"; script="benchmark/benchmarks.jl"); \ +# export_markdown(stdout, r)' +# +# The suite is intentionally small (2 system sizes, both backends) to keep CI +# runtime bounded. The `et` (autograd→analytic) force entries are the ones that +# guard against re-introducing the autograd regression. + +using BenchmarkTools + +include("common.jl") + +const SUITE = BenchmarkGroup() + +# Build the parameter-matched calculators once. +const _setup = build_model_and_calcs() +const _ace = _setup.ace_calc +const _et = _setup.et_calc + +SUITE["energy"] = BenchmarkGroup() +SUITE["forces"] = BenchmarkGroup() + +for cfg in CI_SIZE_CONFIGS + sys = make_system(cfg) + nat = length(sys) + key = "$(nat)atoms" + + # warm up so the recorded samples exclude first-call compilation + AtomsCalculators.potential_energy(sys, _ace) + AtomsCalculators.potential_energy(sys, _et) + AtomsCalculators.forces(sys, _ace) + AtomsCalculators.forces(sys, _et) + + SUITE["energy"]["classic_$key"] = + @benchmarkable AtomsCalculators.potential_energy($sys, $_ace) + SUITE["energy"]["et_$key"] = + @benchmarkable AtomsCalculators.potential_energy($sys, $_et) + SUITE["forces"]["classic_$key"] = + @benchmarkable AtomsCalculators.forces($sys, $_ace) + SUITE["forces"]["et_$key"] = + @benchmarkable AtomsCalculators.forces($sys, $_et) +end diff --git a/benchmark/common.jl b/benchmark/common.jl new file mode 100644 index 00000000..2674c057 --- /dev/null +++ b/benchmark/common.jl @@ -0,0 +1,103 @@ +# Shared setup for the force-evaluation performance benchmarks. +# +# Single source of truth for: +# - model construction (classic ACEModel + ET conversion) +# - test systems of varying size +# - the measurement helper used by both the reproduction script +# (bench_forces_regression.jl) and the PkgBenchmark suite (benchmarks.jl). +# +# This deliberately mirrors benchmark/benchmark_full_model.jl so the numbers are +# comparable to the historical benchmark, while adding allocation capture and a +# parameter-matched classic-vs-ET pair built from the SAME (ps, st). + +using ACEpotentials +const M = ACEpotentials.Models +const ETM = ACEpotentials.ETModels + +import EquivariantTensors as ET +import AtomsCalculators +using StaticArrays, Lux, Random, LuxCore, LinearAlgebra +using AtomsBase, AtomsBuilder, Unitful +using BenchmarkTools + +# --- Model definition (kept small enough for CI, large enough to be meaningful) --- +const ELEMENTS = (:Si, :O) +const ORDER = 2 +const MAXLEVEL = 8 +const MAXL = 4 + +""" + build_model_and_calcs(; rng) + +Build a classic ACE model and its parameter-matched ET conversion from the SAME +`(model, ps, st)`, so the two calculators represent identical physics and any +timing difference is purely backend/implementation. + +Returns `(; model, ps, st, ace_calc, et_calc, rcut)`. +""" +function build_model_and_calcs(; rng = Random.MersenneTwister(1234)) + rin0cuts = M._default_rin0cuts(ELEMENTS) + rin0cuts = (x -> (rin = x.rin, r0 = x.r0, rcut = 5.5)).(rin0cuts) + E0s = Dict(:Si => -158.54496821, :O => -2042.0330099956639) + + model = M.ace_model(; elements = ELEMENTS, order = ORDER, + Ytype = :solid, level = M.TotalDegree(), + max_level = MAXLEVEL, maxl = MAXL, pair_maxn = MAXLEVEL, + rin0cuts = rin0cuts, + init_WB = :glorot_normal, init_Wpair = :glorot_normal, + pair_learnable = true, # required for ET conversion + E0s = E0s) + + ps, st = Lux.setup(rng, model) + + ace_calc = M.ACEPotential(model, ps, st) # classic analytic backend + et_calc = ETM.convert2et_full(model, ps, st) # ET (autograd) backend + rcut = maximum(a.rcut for a in model.pairbasis.rin0cuts) + + return (; model, ps, st, ace_calc, et_calc, rcut) +end + +# --- Systems --- +# (nx, ny, nz) supercell repeats of cubic Si (8 atoms/cell), randomly half-Si/half-O. +const SIZE_CONFIGS = [ + (2, 2, 2), # 64 atoms + (3, 3, 2), # 144 atoms + (4, 4, 2), # 256 atoms + (4, 4, 4), # 512 atoms + (5, 5, 4), # 800 atoms +] + +# Smaller subset used by the CI PkgBenchmark suite to bound runtime. +const CI_SIZE_CONFIGS = [(2, 2, 2), (4, 4, 2)] # 64, 256 atoms + +""" + make_system(cfg; rng) + +Deterministic rattled Si/O supercell for the given `(nx,ny,nz)` config. +""" +function make_system(cfg; rng = Random.MersenneTwister(1234)) + sys = AtomsBuilder.bulk(:Si, cubic = true) * cfg + rattle!(sys, 0.1u"Å") + AtomsBuilder.randz!(sys, [:Si => 0.5, :O => 0.5]) + return sys +end + +""" + nedges(sys, rcut) + +Number of directed edges in the interaction graph (for reporting). +""" +nedges(sys, rcut) = length(ET.Atoms.interaction_graph(sys, rcut * u"Å").edge_data) + +""" + measure(f) + +Run `f()` once to warm up, then benchmark it. Returns `(; time_ns, memory, allocs)` +taken from the minimum sample (noise-resistant). +""" +function measure(f) + f() # warmup / compile + b = @benchmark $f() samples = 5 evals = 3 + m = minimum(b) + return (; time_ns = m.time, memory = m.memory, allocs = m.allocs) +end diff --git a/benchmark/profile_classic_ed.jl b/benchmark/profile_classic_ed.jl new file mode 100644 index 00000000..e87f850a --- /dev/null +++ b/benchmark/profile_classic_ed.jl @@ -0,0 +1,52 @@ +# Profile the classic analytic evaluate_ed to locate the force-path hotspot. +# Times each sub-step on a representative single-site environment. +include("common.jl") +using Printf +const P4ML = ACEpotentials.Models.P4ML +import EquivariantTensors + +const s = build_model_and_calcs() +const model = s.model +const ps = s.ps +const st = s.st + +# representative atomic environment (~ number of neighbours within rcut) +Nat = 40 +Rs, Zs, z0 = M.rand_atenv(model, Nat) + +bench(f) = (f(); minimum(@benchmark $f() samples=10 evals=5).time / 1e3) # µs + +i_z0 = M._z2i(model.rbasis, z0) +rs = norm.(Rs) + +# component timings +t_full_e = bench(() -> M.evaluate(model, Rs, Zs, z0, ps, st)) +t_full_ed = bench(() -> M.evaluate_ed(model, Rs, Zs, z0, ps, st)) + +Rnl, dRnl = M.evaluate_ed_batched(model.rbasis, rs, z0, Zs, ps.rbasis, st.rbasis) +Ylm, dYlm = P4ML.evaluate_ed(model.ybasis, Rs) +TA = promote_type(eltype(Rnl), eltype(Ylm)) +A = zeros(TA, length(model.tensor.abasis)) +EquivariantTensors.evaluate!(A, model.tensor.abasis, (Rnl, Ylm)) +∂B = [ps.WB[:, i_z0]] + +t_rnl_ed = bench(() -> M.evaluate_ed_batched(model.rbasis, rs, z0, Zs, ps.rbasis, st.rbasis)) +t_ylm_ed = bench(() -> P4ML.evaluate_ed(model.ybasis, Rs)) +t_tensor_eval = bench(() -> EquivariantTensors.evaluate(model.tensor, Rnl, Ylm, NamedTuple(), NamedTuple())) +t_tensor_pb = bench(() -> EquivariantTensors.pullback(∂B, model.tensor, Rnl, Ylm, A)) +t_pair_ed = model.pairbasis === nothing ? 0.0 : + bench(() -> M.evaluate_ed_batched(model.pairbasis, rs, z0, Zs, ps.pairbasis, st.pairbasis)) +alloc_ed = (M.evaluate_ed(model, Rs, Zs, z0, ps, st); @allocated M.evaluate_ed(model, Rs, Zs, z0, ps, st)) + +@printf("\n=== classic evaluate_ed component profile (Nat=%d neighbours) ===\n", Nat) +@printf("evaluate (energy only) : %8.2f µs\n", t_full_e) +@printf("evaluate_ed (energy+forces) : %8.2f µs (%.1fx energy)\n", t_full_ed, t_full_ed/t_full_e) +@printf(" -- components --\n") +@printf(" radial evaluate_ed_batched: %8.2f µs\n", t_rnl_ed) +@printf(" Ylm evaluate_ed : %8.2f µs\n", t_ylm_ed) +@printf(" tensor evaluate (fwd) : %8.2f µs\n", t_tensor_eval) +@printf(" tensor pullback (bwd) : %8.2f µs <-- ET migration code\n", t_tensor_pb) +@printf(" pair basis evaluate_ed : %8.2f µs <-- pair path\n", t_pair_ed) +@printf(" (sum of components) : %8.2f µs\n", + t_rnl_ed + t_ylm_ed + t_tensor_eval + t_tensor_pb + t_pair_ed) +@printf("evaluate_ed allocations : %8d bytes\n", alloc_ed) diff --git a/src/models/ace.jl b/src/models/ace.jl index b7ebe0c6..4e4965bc 100644 --- a/src/models/ace.jl +++ b/src/models/ace.jl @@ -297,8 +297,28 @@ end -function evaluate_ed(model::ACEModel, - Rs::AbstractVector{SVector{3, T}}, Zs, Z0, +# Function barrier for the gradient assembly in `evaluate_ed`. Because +# `EquivariantTensors.pullback` is not type-stable, ∂Rnl / ∂Ylm reach the caller as +# `Any`; passing them through this function call forces Julia to re-specialise on +# their concrete runtime types, turning the inner products back into statically +# dispatched (fast, non-allocating) operations. +function _assemble_grad_ed!(∇Ei, ∂Rnl, dRnl, ∂Ylm, dYlm, ∇rs) + @inbounds for t = 1:size(∂Rnl, 2) + for j = 1:size(∂Rnl, 1) + ∇Ei[j] += (∂Rnl[j, t] * dRnl[j, t]) * ∇rs[j] + end + end + @inbounds for t = 1:size(∂Ylm, 2) + for j = 1:size(∂Ylm, 1) + ∇Ei[j] += ∂Ylm[j, t] * dYlm[j, t] + end + end + return ∇Ei +end + + +function evaluate_ed(model::ACEModel, + Rs::AbstractVector{SVector{3, T}}, Zs, Z0, ps, st) where {T} i_z0 = _z2i(model.rbasis, Z0) @@ -345,21 +365,18 @@ function evaluate_ed(model::ACEModel, ∂Rnl, ∂Ylm = EquivariantTensors.pullback([∂B], model.tensor, Rnl, Ylm, A) # ---------- ASSEMBLE DERIVATIVES ------------ - # The ∂Ei / ∂𝐫ⱼ can now be obtained from the ∂Ei / ∂Rnl, ∂Ei / ∂Ylm - # as follows: - # ∂Ei / ∂𝐫ⱼ = ∑_nl ∂Ei / ∂Rnl[j] * ∂Rnl[j] / ∂𝐫ⱼ + # The ∂Ei / ∂𝐫ⱼ can now be obtained from the ∂Ei / ∂Rnl, ∂Ei / ∂Ylm + # as follows: + # ∂Ei / ∂𝐫ⱼ = ∑_nl ∂Ei / ∂Rnl[j] * ∂Rnl[j] / ∂𝐫ⱼ # + ∑_lm ∂Ei / ∂Ylm[j] * ∂Ylm[j] / ∂𝐫ⱼ + # NB: `EquivariantTensors.pullback` is not type-stable (returns Tuple{Any,Any}), + # so ∂Rnl / ∂Ylm are inferred as `Any`. The assembly is therefore done in a + # separate function (_assemble_grad_ed!) which acts as a function barrier: + # it re-specialises on the concrete runtime types. Doing it inline makes the + # element-wise products dynamically dispatched and is ~10x slower (see + # benchmark/bench_forces_regression.jl). ∇Ei = zeros(SVector{3, T}, length(Rs)) - for t = 1:size(∂Rnl, 2) - for j = 1:size(∂Rnl, 1) - ∇Ei[j] += (∂Rnl[j, t] * dRnl[j, t]) * ∇rs[j] - end - end - for t = 1:size(∂Ylm, 2) - for j = 1:size(∂Ylm, 1) - ∇Ei[j] += ∂Ylm[j, t] * dYlm[j, t] - end - end + _assemble_grad_ed!(∇Ei, ∂Rnl, dRnl, ∂Ylm, dYlm, ∇rs) # ------------------- # pair potential