SCHAModules: cut peak memory of get_odd_straight_with_v4 (~2.4× lighter, ~1.45× faster)#423
Open
SorBalda wants to merge 4 commits into
Open
SCHAModules: cut peak memory of get_odd_straight_with_v4 (~2.4× lighter, ~1.45× faster)#423SorBalda wants to merge 4 commits into
SorBalda wants to merge 4 commits into
Conversation
…with_v4
Bitwise-identical memory optimization of get_odd_straight_with_v4.
- Remove zz(nl,nl), vv(nl*(nl+1)/2) and ww(nl): declared and allocated
but never read/written anywhere in the subroutine (verified by grep on
the v1.5 source). They were also never deallocated -> a memory leak.
zz alone is one full nl x nl array (= 1 "v4 unit").
- Remove the iden(nl,nl) scratch matrix. Instead of building iden and
copying "maux = iden", initialize maux directly:
maux = 0.0d0 ; do x=1,nl: maux(x,x)=1.0d0
Same value, same bit pattern; saves one nl x nl allocation.
No BLAS call, no operation order, no numerics changed.
Standalone Fortran test (n_mode=6, symmetrized v3/v4): max|phi_new-phi_orig| = 0.0 exactly.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
NOT bitwise identical: this commit reassociates floating-point products,
so results differ from the original at the ~1e-15 relative level (verified
below). The algorithm and the physics are unchanged.
(d) Replace the nl x nl buffer product with an nl x ns contraction.
Original:
v42 := lamat * maux (nl x nl dgemm, reuses v42 as buffer)
cf := v42 * v3^T (nl x ns)
New (associativity Lambda*M^-1*v3^T = Lambda*(M^-1*v3^T)):
tmp := maux * v3^T (nl x ns)
cf := lamat * tmp (nl x ns)
Removes the nl x nl buffer and lowers that stage from O(N^6) to O(N^5).
tmp is nl x ns = O(N^3), negligible.
(e) Optional argument use_v4_symmetry (default .false. -> old safe behavior).
When .true. (standard case: v4 permutation-symmetric after
ApplySymmetryToTensor4 / use_symmetries=True), the reordered copy
v42(ja,ka)=v4(w,z,x,y) equals the plain column-major reshape of v4, so
v4 is fed directly to the first dgemm (lda=nl) and the whole nl x nl
v42 array is never allocated. INVALID for use_symmetries=False, hence
it is opt-in only, never the silent default. The optional argument is
interface-compatible: existing callers that omit it are unaffected.
Peak internal nl x nl arrays: default path lamat+v42+maux (3 units, v42
freed right after the first dgemm); symmetry path lamat+maux (2 units).
Standalone test (n_mode=6, symmetrized v3/v4), vs original:
default path : max abs diff 2.1e-16, max elementwise rel diff 7.6e-15
symmetry path: max abs diff 1.1e-16, max elementwise rel diff 1.8e-14
Both far below the 1e-12 tolerance.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The nl x nl Lambda matrix (nl = n_mode^2) was built by get_cmat, which
internally allocates mat_e + mat_et (two MORE nl x nl temporaries) and
does cmat = mat_e . mat_et with an O(N^6) dgemm: true peak inside that
call was v4 + mat_e + mat_et + cmat = 4 units = 43.5 GB for N = 192 --
the actual OOM site on a 31 GiB machine, untouched by commits 1-2.
get_cmat is no longer called (nobody calls it now). We use its exact
factorized structure Lambda(ka,ka') = sum_{mu,nu} e(nu,x) e(mu,y)
D(mu,nu) e(nu,x') e(mu,y') (e from get_emat, D = mat_w/2 from get_g,
both N x N) and build M = I - v4.Lambda with partial Kronecker
contractions against the small e (four O(N^5) dgemm stages + O(N^4)
diagonal scaling), ping-ponging between exactly TWO nl x nl buffers:
maux and v4 itself used as scratch. The final assembly swaps both index
pairs to reproduce the original v42(ja,ka)=v4(w,z,x,y) ordering EXACTLY;
no permutation symmetry of v4 is assumed anywhere. cf = Lambda.tmp is
factorized the same way on skinny nl x ns matrices (O(N^3) buffers).
The only O(N^6) operation left is the dgetrf/dgetri inversion.
CONTRACT CHANGE: v4 is now intent(inout) and its content is DESTROYED.
Safe for the only caller, Ensemble.get_free_energy_hessian: d4 comes
from SCHAModules.get_v4 (F-contiguous, so f2py intent(inout) does not
copy), is symmetrized in place, and is never used after this call.
The commit-2 optional flag use_v4_symmetry is removed: it only existed
to skip the v42 copy, which no longer exists; the kron path is exact
for arbitrary v4 and no Python code ever passed the flag.
Memory: true whole-call peak 2 units + O(N^3) = 21.7 GB for N = 192
(was 43.5 GB inside get_cmat). Verified vs pristine v1.5 with the
standalone driver (n_mode = 6): max elementwise relative diff 1.46e-14
with permutation-symmetrized v3/v4 and 1.86e-14 with raw
non-symmetrized v4 (tolerance 1e-12; not bitwise, different but
mathematically identical summation order).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Cuts the peak memory of
SCHAModules/get_odd_straight_with_v4by ~2.4× (and makes it ~1.45× faster at equal BLAS), with no change to the physics. The routine computes the odd (3rd-order) correctionphi_sc_odd = v3ᵀ·Λ·(I − v4·Λ)⁻¹·v3(Bianco, PRB 96, 014111, Eq. 27) with densenl × nlmatrices,nl = n_mode². For a gold 4×4×4 cell (N = n_mode = 192) onenl × nldouble array is 10.87 GB, so the old peak of ~4–5.5 such arrays (43–60 GB) OOMs on a 31 GiB machine. This is what motivated the work.The change is fully localized to
get_odd_straight_with_v4.f90— no other file is touched,get_cmat.f90is left intact (it is simply no longer called), and the f2py interface is unchanged except for the contract note below.What changed (3 incremental commits)
memopt(safe)— bitwise identical. Removezz,vv,ww(declared + allocated but never read/written — also never deallocated, i.e. a leak) and theidenscratch identity (initializemauxdirectly). No BLAS call or operation order changed.memopt(reassoc)— equal to ~1e-15. ReassociateΛ·M⁻¹·v3ᵀ = Λ·(M⁻¹·v3ᵀ)so thenl × nlbuffer product becomes annl × nscontraction (drops an O(N⁶) stage to O(N⁵) and frees a fullnl × nlbuffer).memopt(kron)— Λ never materialized. The real OOM site was the oldget_cmatcall, which internally allocated two morenl × nlarrays (mat_e,mat_et) and did an O(N⁶)cmat = mat_e·mat_et— a true peak of 4 units (43.5 GB) never counted before. Usingget_cmat's own factorized structureΛ(ka,ka') = Σ_{μν} e(ν,x)e(μ,y)·D(μ,ν)·e(ν,x')e(μ,y')(with the smallN×Ne,Dfromget_emat/get_g), the matrixI − v4·Λis built by partial Kronecker contractions ping-ponging between exactly twonl × nlbuffers (maux, andv4reused as scratch): four O(N⁵) dgemm stages + O(N⁴) diagonal scaling. The only O(N⁶) op left is the unavoidabledgetrf/dgetriinversion. No permutation symmetry of v4 is assumed — it is a pure index-bookkeeping identity, valid for arbitrary v4.v4(thed4argument) is nowintent(inout)and its content is destroyed on exit (it is one of the two scratch buffers). This is safe for the only caller,Ensemble.get_free_energy_hessian, which buildsd4fresh viaSCHAModules.get_v4(...), (optionally) symmetrizes it in place, calls this routine, and never readsd4again. f2py wrapsintent(inout)without copying only for an F-contiguous float64 array (and raises loudly otherwise), so a misusing caller fails fast rather than silently. Any custom code callingget_odd_straight_with_v4directly must pass a throw-away, F-contiguousv4.Correctness
Validated against the pristine v1.5 subroutine with a standalone Fortran driver (
n_mode = 6, synthetic consistent inputs), required tolerance max elementwise rel diff < 1e-12:Not bitwise for commits 2–3 (different but mathematically identical summation order), all ≈ 2e-14 ≪ 1e-12.
Full test suite on the final build (
pytest --ignore=aiida_ensemble, cellconstructor 1.5.2): 23 passed, 1 skipped, 0 failed — includingtest_hessian(exercises this routine) and the gold vc-relax workflows. The 1 skip istest_parallel_calculation(optional mpi4py + F3Calc).Memory & performance
Peak concurrent
nl × nlarrays (Au 4×4×4, N = 192, 1 unit = 10.87 GB):get_cmatinternals)Laptop benchmark, Au 3×3×3 (nconf 50, T = 100,
OMP_NUM_THREADS=1, same OpenBLAS):Same physics throughout (max |Δfreq| vs historical hessian = 7.7e-14 cm⁻¹).
Notes
The standalone Fortran validation driver (original vs patched subroutine, two random-input cases including non-symmetrized v4) and a detailed derivation are available on request / in my fork — kept out of this PR to keep it to the single-file optimization. Happy to add a regression test or adjust anything.