diff --git a/.gitignore b/.gitignore index 13b72b45..a6a3c1ad 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,6 @@ scratch .vscode .ipynb_checkpoints .CondaPkg + +# PkgBenchmark tuning artifact +benchmark/tune.json diff --git a/benchmark/bench_forces_regression.jl b/benchmark/bench_forces_regression.jl index 429775b7..160ba52d 100644 --- a/benchmark/bench_forces_regression.jl +++ b/benchmark/bench_forces_regression.jl @@ -84,13 +84,13 @@ function run_table() 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. +# Isolation: time the ETACE gradient step `site_grads` vs the full basis jacobian +# `site_basis_jacobian`, with the graph precomputed — isolates the gradient 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("| Atoms | site_grads (ms) | site_basis_jacobian (full, 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) diff --git a/benchmark/profile_et_forces.jl b/benchmark/profile_et_forces.jl new file mode 100644 index 00000000..b8f625c3 --- /dev/null +++ b/benchmark/profile_et_forces.jl @@ -0,0 +1,43 @@ +# Profile the ET ACE force path to find optimization scope. +# Breaks the stacked-calculator forces into components, and the ETACE component +# into graph-build / site_grads(Zygote) / forces_from_edge_grads, with allocations. +include("common.jl") +using Printf +import EquivariantTensors as ET + +s = build_model_and_calcs() +et_calc = s.et_calc # StackedCalculator(onebody, pair, ace) +rcut = s.rcut +onebody, pair, ace = et_calc.calcs + +bench(f) = (f(); minimum(@benchmark $f() samples=10 evals=3).time/1e6) # ms +allocOf(f) = (f(); @allocated f()) + +for cfg in SIZE_CONFIGS[[1, 3, 5]] + sys = make_system(cfg) + nat = length(sys) + G = ET.Atoms.interaction_graph(sys, rcut * u"Å") + ne = length(G.edge_data) + + t_full = bench(() -> AtomsCalculators.forces(sys, et_calc)) + t_1b = bench(() -> AtomsCalculators.forces(sys, onebody)) + t_pair = bench(() -> AtomsCalculators.forces(sys, pair)) + t_ace = bench(() -> AtomsCalculators.forces(sys, ace)) + + m, ps, st = ace.model, ace.ps, ace.st + t_graph = bench(() -> ET.Atoms.interaction_graph(sys, rcut * u"Å")) + t_sg = bench(() -> ETM.site_grads(m, G, ps, st)) + ∂G = ETM.site_grads(m, G, ps, st) + t_ffe = bench(() -> ET.Atoms.forces_from_edge_grads(sys, G, ∂G.edge_data)) + a_sg = allocOf(() -> ETM.site_grads(m, G, ps, st)) + + @printf("\n=== %d atoms, %d edges ===\n", nat, ne) + @printf("full stacked forces : %8.3f ms\n", t_full) + @printf(" onebody forces : %8.3f ms\n", t_1b) + @printf(" pair forces : %8.3f ms\n", t_pair) + @printf(" ETACE forces : %8.3f ms\n", t_ace) + @printf(" -- ETACE breakdown --\n") + @printf(" graph build : %8.3f ms\n", t_graph) + @printf(" site_grads : %8.3f ms alloc=%d B\n", t_sg, a_sg) + @printf(" forces_from_edge_grads: %7.3f ms\n", t_ffe) +end diff --git a/src/et_models/et_ace.jl b/src/et_models/et_ace.jl index df244645..b8d822a9 100644 --- a/src/et_models/et_ace.jl +++ b/src/et_models/et_ace.jl @@ -37,16 +37,22 @@ function _apply_etace(l::ETACE, X::ET.ETGraph, ps, st) return φ end -# ----------------------------------------------------------- +# ----------------------------------------------------------- -import Zygote +import Zygote # -# At first glance this looks like we are computing ∂E / ∂ri but this is not -# actually true. Because E = ∑ Ei and by interpreting G as a list of edges -# we are differentiating E w.r.t. 𝐫ij which is the same is Ei w.r.t. 𝐫ij. +# At first glance this looks like we are computing ∂E / ∂ri but this is not +# actually true. Because E = ∑ Ei and by interpreting G as a list of edges +# we are differentiating E w.r.t. 𝐫ij which is the same as Ei w.r.t. 𝐫ij. +# +# NB: an analytic VJP (via ET._ka_pullback) was prototyped but only ~10-15% +# faster here — ET's many-body kernel intermediates dominate the cost, not +# Zygote's overhead — and it coupled ACEpotentials tightly to ET internals. +# We therefore keep the Zygote gradient for the many-body model. The pair model +# (et_pair.jl) does use the analytic gradient, since there it is a trivial, +# low-coupling contraction. # - function site_grads(l::ETACE, X::ET.ETGraph, ps, st) ∂X = Zygote.gradient( X -> sum(_apply_etace(l, X, ps, st)), X)[1] return ∂X diff --git a/src/et_models/et_calculators.jl b/src/et_models/et_calculators.jl index ac1c3674..4cc77538 100644 --- a/src/et_models/et_calculators.jl +++ b/src/et_models/et_calculators.jl @@ -14,7 +14,7 @@ # See also: stackedcalc.jl for StackedCalculator (combines multiple calculators) import AtomsCalculators -import AtomsBase: AbstractSystem, ChemicalSpecies +import AtomsBase: AbstractSystem, ChemicalSpecies, position, species import EquivariantTensors as ET using DecoratedParticles: PState using StaticArrays @@ -25,6 +25,58 @@ using LinearAlgebra: norm import ..Models: length_basis, energy_forces_virial_basis, potential_energy_basis +# ============================================================================ +# Analytic site gradients (used by ETPairModel) +# ============================================================================ +# +# Forces come from ∂E/∂𝐫ij. Since E = ∑_i ∑_k W[1,k,iZ[i]] · 𝔹[i,k] (a linear +# readout of the basis), the edge gradient is the basis jacobian contracted once +# with the readout weights: +# ∇E[j,i] = ∑_k W[1,k,iZ[i]] · ∂𝔹[j,i,k] +# The same contraction is already used per-basis-function in +# `energy_forces_virial_basis`; here it is done once with the actual weights. +# +# This is used for the PAIR model, where `site_basis_jacobian` returns ∂𝔹 == ∂R +# directly (the pair basis is a linear sum over neighbours), so the jacobian is +# exactly the per-edge gradient — no blow-up. The many-body ETACE model keeps the +# Zygote `site_grads` (see et_ace.jl): there `site_basis_jacobian` would compute +# the full `nbasis`× larger jacobian, and the leaner VJP alternative coupled too +# tightly to EquivariantTensors internals for the modest gain. + +# The contraction runs in its own function so it acts as a FUNCTION BARRIER: +# `site_basis_jacobian` → `ET._jacobian_X` is not type-stable (returns Any-typed +# ∂𝔹), and doing the inner loop inline would dispatch every product dynamically +# (cf. the classic-model fix in src/models/ace.jl and EquivariantTensors.jl#135). +function _contract_readout!(∇E, ∂𝔹, W, iZ) + (isempty(∂𝔹) || isempty(iZ)) && return ∇E + z = zero(∂𝔹[1, 1, 1]) + @inbounds for i in axes(∂𝔹, 2) + iz = iZ[i] + for j in axes(∂𝔹, 1) + s = z + for k in axes(∂𝔹, 3) + s = s + W[1, k, iz] * ∂𝔹[j, i, k] + end + ∇E[j, i] = s + end + end + return ∇E +end + +# Analytic `site_grads` for any ETACE-pattern model with a SelectLinL readout +# (ETACE, ETPairModel). Returns edge gradients in the same `(; edge_data = …)` +# form as the former Zygote implementation, consumed by `_wrapped_forces` / +# `_compute_virial`. +function _site_grads_analytic(l, X::ET.ETGraph, ps, st) + _, ∂𝔹 = site_basis_jacobian(l, X, ps, st) # (maxneigs, nnodes, nbasis) + iZ = l.readout.selector.(X.node_data) + ∇E = similar(∂𝔹, size(∂𝔹, 1), size(∂𝔹, 2)) # concrete eltype of ∂𝔹 + _contract_readout!(∇E, ∂𝔹, ps.readout.W, iZ) + ∇E3 = reshape(∇E, size(∇E, 1), size(∇E, 2), 1) + return (; edge_data = ET.rev_reshape_embedding(∇E3, X)[:]) +end + + # ============================================================================ # WrappedSiteCalculator - Unified wrapper for ETACE-pattern models # ============================================================================ @@ -77,14 +129,26 @@ end cutoff_radius(calc::WrappedSiteCalculator) = calc.rcut * u"Å" -function _wrapped_energy(calc::WrappedSiteCalculator, sys::AbstractSystem) - G = ET.Atoms.interaction_graph(sys, calc.rcut * u"Å") +# Build the interaction graph for a calculator's own cutoff. +_wrapped_graph(calc::WrappedSiteCalculator, sys::AbstractSystem) = + ET.Atoms.interaction_graph(sys, calc.rcut * u"Å") + +# Each `_wrapped_*` has a method that accepts a prebuilt graph `G`, so the graph +# can be shared across stacked components that use the same cutoff (see +# stackedcalc.jl). The no-`G` methods just build the graph and delegate. + +_wrapped_energy(calc::WrappedSiteCalculator, sys::AbstractSystem) = + _wrapped_energy(calc, sys, _wrapped_graph(calc, sys)) + +function _wrapped_energy(calc::WrappedSiteCalculator, sys::AbstractSystem, G) Ei, _ = calc.model(G, calc.ps, calc.st) return sum(Ei) end -function _wrapped_forces(calc::WrappedSiteCalculator, sys::AbstractSystem) - G = ET.Atoms.interaction_graph(sys, calc.rcut * u"Å") +_wrapped_forces(calc::WrappedSiteCalculator, sys::AbstractSystem) = + _wrapped_forces(calc, sys, _wrapped_graph(calc, sys)) + +function _wrapped_forces(calc::WrappedSiteCalculator, sys::AbstractSystem, G) ∂G = site_grads(calc.model, G, calc.ps, calc.st) # Handle empty edge case (e.g., ETOneBody with small cutoff) if isempty(∂G.edge_data) @@ -104,8 +168,10 @@ function _compute_virial(G::ET.ETGraph, ∂G) return V end -function _wrapped_virial(calc::WrappedSiteCalculator, sys::AbstractSystem) - G = ET.Atoms.interaction_graph(sys, calc.rcut * u"Å") +_wrapped_virial(calc::WrappedSiteCalculator, sys::AbstractSystem) = + _wrapped_virial(calc, sys, _wrapped_graph(calc, sys)) + +function _wrapped_virial(calc::WrappedSiteCalculator, sys::AbstractSystem, G) ∂G = site_grads(calc.model, G, calc.ps, calc.st) # Handle empty edge case if isempty(∂G.edge_data) @@ -114,9 +180,10 @@ function _wrapped_virial(calc::WrappedSiteCalculator, sys::AbstractSystem) return _compute_virial(G, ∂G) end -function _wrapped_energy_forces_virial(calc::WrappedSiteCalculator, sys::AbstractSystem) - G = ET.Atoms.interaction_graph(sys, calc.rcut * u"Å") +_wrapped_energy_forces_virial(calc::WrappedSiteCalculator, sys::AbstractSystem) = + _wrapped_energy_forces_virial(calc, sys, _wrapped_graph(calc, sys)) +function _wrapped_energy_forces_virial(calc::WrappedSiteCalculator, sys::AbstractSystem, G) # Energy from site energies (call model directly - ETACE interface) Ei, _ = calc.model(G, calc.ps, calc.st) E = sum(Ei) @@ -136,6 +203,35 @@ function _wrapped_energy_forces_virial(calc::WrappedSiteCalculator, sys::Abstrac return (energy=E, forces=F, virial=V) end +# ---------------------------------------------------------------------------- +# Graph-cached dispatch used by StackedCalculator (see stackedcalc.jl). +# Components that share a cutoff (e.g. pair + many-body) reuse one interaction +# graph per force/energy call instead of each rebuilding its own. The cache is +# keyed on each calculator's own `rcut`, so per-component cutoffs are preserved. +# Non-WrappedSiteCalculator components fall back to the plain AtomsCalculators +# interface (cache ignored), keeping StackedCalculator generic. +# ---------------------------------------------------------------------------- +_cached_graph!(gcache, calc::WrappedSiteCalculator, sys) = + get!(() -> ET.Atoms.interaction_graph(sys, calc.rcut * u"Å"), gcache, calc.rcut) + +_cached_energy(c, sys, gcache) = AtomsCalculators.potential_energy(sys, c) +_cached_forces(c, sys, gcache) = AtomsCalculators.forces(sys, c) +_cached_virial(c, sys, gcache) = AtomsCalculators.virial(sys, c) +_cached_efv(c, sys, gcache) = AtomsCalculators.energy_forces_virial(sys, c) + +_cached_energy(c::WrappedSiteCalculator, sys, gcache) = + _wrapped_energy(c, sys, _cached_graph!(gcache, c, sys)) * u"eV" +_cached_forces(c::WrappedSiteCalculator, sys, gcache) = + _wrapped_forces(c, sys, _cached_graph!(gcache, c, sys)) .* u"eV/Å" +_cached_virial(c::WrappedSiteCalculator, sys, gcache) = + _wrapped_virial(c, sys, _cached_graph!(gcache, c, sys)) * u"eV" +function _cached_efv(c::WrappedSiteCalculator, sys, gcache) + efv = _wrapped_energy_forces_virial(c, sys, _cached_graph!(gcache, c, sys)) + return (energy = efv.energy * u"eV", + forces = efv.forces .* u"eV/Å", + virial = efv.virial * u"eV") +end + # AtomsCalculators interface for WrappedSiteCalculator AtomsCalculators.@generate_interface function AtomsCalculators.potential_energy( sys::AbstractSystem, calc::WrappedSiteCalculator; kwargs...) @@ -603,6 +699,41 @@ function ETOneBodyPotential(model::ETOneBody, ps, st, rcut::Real) return WrappedSiteCalculator(model, ps, st, Float64(rcut)) end +# ---------------------------------------------------------------------------- +# ETOneBody is a pure one-body model: energy depends only on atom species +# (node_data), never on neighbours, and forces/virial are identically zero. +# We therefore SKIP the interaction-graph / neighbour search entirely and build +# the node states directly. `node_data` is rcut-independent (it loops over all +# atoms regardless of edges), so this matches the graph path exactly — and these +# methods also bypass the StackedCalculator graph cache (no graph is needed). +# ---------------------------------------------------------------------------- +_onebody_nodes(sys::AbstractSystem) = + [ PState(𝐫 = ustrip.(position(sys, i)), z = species(sys, i)) for i in 1:length(sys) ] + +function _wrapped_energy(calc::ETOneBodyPotential, sys::AbstractSystem) + Ei, _ = calc.model(_onebody_nodes(sys), calc.ps, calc.st) + return sum(Ei) +end +_wrapped_forces(calc::ETOneBodyPotential, sys::AbstractSystem) = + zeros(SVector{3, Float64}, length(sys)) +_wrapped_virial(calc::ETOneBodyPotential, sys::AbstractSystem) = + zero(SMatrix{3, 3, Float64, 9}) +_wrapped_energy_forces_virial(calc::ETOneBodyPotential, sys::AbstractSystem) = + (energy = _wrapped_energy(calc, sys), + forces = zeros(SVector{3, Float64}, length(sys)), + virial = zero(SMatrix{3, 3, Float64, 9})) + +# StackedCalculator dispatch: onebody needs no graph, so ignore the cache. +_cached_energy(c::ETOneBodyPotential, sys, gcache) = _wrapped_energy(c, sys) * u"eV" +_cached_forces(c::ETOneBodyPotential, sys, gcache) = _wrapped_forces(c, sys) .* u"eV/Å" +_cached_virial(c::ETOneBodyPotential, sys, gcache) = _wrapped_virial(c, sys) * u"eV" +function _cached_efv(c::ETOneBodyPotential, sys, gcache) + efv = _wrapped_energy_forces_virial(c, sys) + return (energy = efv.energy * u"eV", + forces = efv.forces .* u"eV/Å", + virial = efv.virial * u"eV") +end + # ============================================================================ # ETOneBodyPotential Training Assembly (empty - no learnable parameters) # ============================================================================ @@ -706,8 +837,9 @@ function convert2et_full(model, ps, st; rng::AbstractRNG=default_rng()) E0_dict = Dict(z => E0s[z.atomic_number] for z in zlist) et_onebody = one_body(E0_dict, x -> x.z) _, onebody_st = setup(rng, et_onebody) - # Use minimum cutoff for graph construction (ETOneBody needs no neighbors) - onebody_calc = WrappedSiteCalculator(et_onebody, nothing, onebody_st, 3.0) + # ETOneBody needs no neighbour graph (energy depends only on species), so its + # cutoff is irrelevant — pass 0.0 to make that explicit. + onebody_calc = WrappedSiteCalculator(et_onebody, nothing, onebody_st, 0.0) # 2. Convert pair potential to ETPairModel et_pair = convertpair(model) diff --git a/src/et_models/et_pair.jl b/src/et_models/et_pair.jl index 419321d9..d9fa33d5 100644 --- a/src/et_models/et_pair.jl +++ b/src/et_models/et_pair.jl @@ -35,10 +35,9 @@ end # ----------------------------------------------------------- -function site_grads(l::ETPairModel, X::ET.ETGraph, ps, st) - ∂X = Zygote.gradient( X -> sum(_apply_etpairmodel(l, X, ps, st)), X)[1] - return ∂X -end +# analytic edge gradients (see _site_grads_analytic in et_calculators.jl); +# replaces the previous allocation-heavy Zygote.gradient implementation. +site_grads(l::ETPairModel, X::ET.ETGraph, ps, st) = _site_grads_analytic(l, X, ps, st) # ----------------------------------------------------------- diff --git a/src/et_models/stackedcalc.jl b/src/et_models/stackedcalc.jl index cdcb737b..edd853c2 100644 --- a/src/et_models/stackedcalc.jl +++ b/src/et_models/stackedcalc.jl @@ -60,30 +60,38 @@ end # Efficient implementations using @generated for compile-time unrolling # ============================================================================ +# `gcache` caches one interaction graph per distinct cutoff within a single call, +# so stacked components that share a cutoff (e.g. pair + many-body) build the +# graph once instead of once each. See `_cached_*` in et_calculators.jl. + @generated function _stacked_energy(sys::AbstractSystem, calc::StackedCalculator{N}) where {N} quote - @nexprs $N i -> E_i = AtomsCalculators.potential_energy(sys, calc.calcs[i]) + gcache = Dict{Float64, Any}() + @nexprs $N i -> E_i = _cached_energy(calc.calcs[i], sys, gcache) return sum(@ntuple $N E) end end @generated function _stacked_forces(sys::AbstractSystem, calc::StackedCalculator{N}) where {N} quote - @nexprs $N i -> F_i = AtomsCalculators.forces(sys, calc.calcs[i]) + gcache = Dict{Float64, Any}() + @nexprs $N i -> F_i = _cached_forces(calc.calcs[i], sys, gcache) return reduce(.+, @ntuple $N F) end end @generated function _stacked_virial(sys::AbstractSystem, calc::StackedCalculator{N}) where {N} quote - @nexprs $N i -> V_i = AtomsCalculators.virial(sys, calc.calcs[i]) + gcache = Dict{Float64, Any}() + @nexprs $N i -> V_i = _cached_virial(calc.calcs[i], sys, gcache) return sum(@ntuple $N V) end end @generated function _stacked_efv(sys::AbstractSystem, calc::StackedCalculator{N}) where {N} quote - @nexprs $N i -> efv_i = AtomsCalculators.energy_forces_virial(sys, calc.calcs[i]) + gcache = Dict{Float64, Any}() + @nexprs $N i -> efv_i = _cached_efv(calc.calcs[i], sys, gcache) return ( energy = sum(@ntuple $N i -> efv_i.energy), forces = reduce(.+, @ntuple $N i -> efv_i.forces),