Fix ~13x force-evaluation regression in classic ACEModel

.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

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`.

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"

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.

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:  "<natoms> <ms>".
+#
+# 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

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)")