Skip to content

Commit da045e2

Browse files
SorBaldaclaude
SorBalda
and
claude
committed
memopt(kron): never materialize Lambda in get_odd_straight_with_v4
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>
1 parent aa4da39 commit da045e2

1 file changed

Lines changed: 207 additions & 116 deletions

File tree

‎SCHAModules/get_odd_straight_with_v4.f90‎

Lines changed: 207 additions & 116 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,89 @@
11

2-
! This subroutine calculates the L mat needed to get the average of the
3-
! third order derivatives. It is formed by four polarization vectors
4-
! times the mass^1/2 divided by the normal length.
2+
! This subroutine computes the odd (third order) SCHA correction to the
3+
! free energy Hessian, including the fourth order term v4:
4+
!
5+
! phi_sc_odd = v3^T . Lambda . (I - v4.Lambda)^-1 . v3
6+
!
7+
! (R. Bianco et al., PRB 96, 014111 (2017), Eq. 27), with nl = n_mode^2.
8+
!
9+
! MEMORY-OPTIMIZED "KRON" VERSION: the nl x nl Lambda matrix is NEVER
10+
! materialized. The old path called get_cmat, which allocated TWO extra
11+
! nl x nl temporaries (mat_e, mat_et) plus the nl x nl output cmat and did
12+
! an O(N^6) dgemm; together with v4 that was a 4-array (4 "d4 units") peak.
13+
! Instead we exploit the exact factorized structure that get_cmat builds
14+
! (see get_cmat.f90, loops at its lines ~78-90):
15+
!
16+
! mat_e (ka,ja) = e(nu,x) * e(mu,y)
17+
! mat_et(ja,ka) = mat_e(ka,ja) * mat_w(mu,nu) * 0.5
18+
! Lambda = mat_e . mat_et
19+
!
20+
! with ka = (x-1)*N + y (y fast) and ja = (mu-1)*N + nu (nu fast), i.e.
21+
!
22+
! Lambda(ka,ka') = sum_{mu,nu} e(nu,x) e(mu,y) D(mu,nu) e(nu,x') e(mu,y')
23+
!
24+
! where e(N,N) comes from get_emat, D(mu,nu) = mat_w(mu,nu)/2 with mat_w
25+
! from get_g (both small, O(N^2)). Applying Lambda therefore reduces to
26+
! contractions with the SMALL matrix e (each an O(N^5) dgemm) plus a
27+
! diagonal scaling over the mode pair (O(N^4)). The only O(N^6) operation
28+
! left is the inversion of (I - v4.Lambda).
29+
!
30+
! EXACT REPRODUCTION OF THE ORIGINAL PRODUCT (index bookkeeping).
31+
! The original code built maux = I - v42 . Lambda with the reordered copy
32+
! v42(ja,ka) = v4(w,z,x,y), ja = (w-1)*N + z (z fast),
33+
! ka = (x-1)*N + y (y fast).
34+
! Expanding, the matrix subtracted from the identity is, as a 4-tensor,
35+
!
36+
! P(w,z,x,y) = sum_{x',y',mu,nu} v4(w,z,x',y') e(nu,x') e(mu,y')
37+
! * D(mu,nu) * e(nu,x) e(mu,y)
38+
!
39+
! placed at maux( (w-1)*N + z , (x-1)*N + y ). NO permutation symmetry of
40+
! v4 is assumed anywhere: this is an identity in the indices, valid for
41+
! arbitrary v4 (verified numerically against v1.5 with a NON-symmetrized
42+
! random v4 as well; the old optional flag use_v4_symmetry is gone since
43+
! the v42 copy it avoided no longer exists).
44+
!
45+
! P is evaluated with partial (Kronecker-factor) contractions that
46+
! ping-pong between exactly TWO nl x nl buffers: maux and v4 ITSELF used
47+
! as scratch. Layouts below are column-major, leftmost index fastest:
48+
!
49+
! A: B(w,z,x',mu) = sum_y' v4(w,z,x',y') e(mu,y') v4 -> maux
50+
! one dgemm: (N^3 x N) . (N x N)^T, O(N^5)
51+
! B: T(w,z,nu,mu) = sum_x' B(w,z,x',mu) e(nu,x') maux-> v4
52+
! N slice dgemms over mu: (N^2 x N) . (N x N)^T, O(N^5)
53+
! C: T(w,z,nu,mu) *= D(mu,nu) v4 in place, O(N^4)
54+
! D: C(w,z,nu,y) = sum_mu T(w,z,nu,mu) e(mu,y) v4 -> maux
55+
! one dgemm: (N^3 x N) . (N x N), O(N^5)
56+
! E: P(w,z,x,y) = sum_nu C(w,z,nu,y) e(nu,x) maux-> v4
57+
! N slice dgemms over y: (N^2 x N) . (N x N), O(N^5)
58+
! F: maux((w-1)N+z,(x-1)N+y) = delta - P(w,z,x,y) v4 -> maux, O(N^4)
59+
! explicit permuted copy: the natural buffer layout of P has (w fast,
60+
! z slow) in the row pair and (x fast, y slow) in the column pair,
61+
! while the original v42.Lambda uses (z fast, w slow) rows and
62+
! (y fast, x slow) columns, so BOTH index pairs are swapped here.
63+
!
64+
! After F: LU inversion in place in maux (dgetrf/dgetri, unchanged), then
65+
! tmp = maux . v3^T (nl x ns, O(N^5))
66+
! cf = Lambda . tmp via the same e/D/e^T factorization
67+
! applied to skinny matrices, using two
68+
! O(N^3) buffers (N x N x ns)
69+
! phi_sc_odd = v32 . cf (as before)
70+
!
71+
! !!! WARNING: v4 IS INTENT(INOUT) AND ITS CONTENT IS DESTROYED !!!
72+
! v4 is deliberately used as one of the two nl x nl scratch buffers (its
73+
! original content is consumed by step A before step B overwrites it).
74+
! This is safe for the only caller, Ensemble.get_free_energy_hessian
75+
! (Modules/Ensemble.py): it creates d4 with SCHAModules.get_v4 (hence
76+
! F-contiguous, so f2py's intent(inout) wraps it in place without a copy),
77+
! optionally symmetrizes it in place, calls this routine, and never uses
78+
! d4 again. Any NEW caller must pass a throw-away, F-contiguous v4.
79+
!
80+
! Peak large-memory budget: exactly 2 nl x nl arrays alive (v4 + maux),
81+
! plus O(N^3) skinny buffers (v32, tmp, cf, b1, b2) and O(N^2)/O(N) work
82+
! arrays. For N = 192 (Au 4x4x4): 2 * 10.87 GB = 21.7 GB, vs ~43.5 GB
83+
! inside the old get_cmat call (v4 + mat_e + mat_et + cmat).
584

685
subroutine get_odd_straight_with_v4 ( a, wr, er, transmode, amass, ityp_sc, T, v3, v4, phi_sc_odd, &
7-
n_mode, nat_sc, ntyp, use_v4_symmetry)
86+
n_mode, nat_sc, ntyp)
887

988
implicit none
1089

@@ -15,163 +94,175 @@ subroutine get_odd_straight_with_v4 ( a, wr, er, transmode, amass, ityp_sc, T, v
1594
integer, dimension(nat_sc), intent(in) :: ityp_sc
1695
double precision, intent(in) :: T
1796
double precision, dimension(n_mode,n_mode,n_mode), intent(in) :: v3
18-
double precision, dimension(n_mode,n_mode,n_mode, n_mode), intent(in) :: v4
97+
! v4 is used as scratch and DESTROYED on exit -- see warning above.
98+
double precision, dimension(n_mode,n_mode,n_mode,n_mode), intent(inout) :: v4
1999
double precision, dimension(n_mode, n_mode), intent(out) :: phi_sc_odd
20-
! Optional: if present and .true., assume v4 is permutation-symmetric
21-
! (standard case, after ApplySymmetryToTensor4 / use_symmetries=True) and
22-
! skip building the explicit reordered copy v42 -- see note below.
23-
logical, intent(in), optional :: use_v4_symmetry
24-
25100

26101
integer :: nat_sc, n_mode, nl, ns, ntyp
27-
double precision, dimension(:,:), allocatable :: l, g, phi_aux, v1, v2, v32
28-
double precision :: lsum
29-
double precision, dimension(:), allocatable :: laux1, lres1, veclong
30-
double precision, dimension(:), allocatable :: laux2, lres2
31102

32-
double precision, dimension(:,:), allocatable :: lamat, v42, maux
103+
! The ONLY allocated nl x nl array (v4, the other big buffer, is the
104+
! caller's own array).
105+
double precision, dimension(:,:), allocatable :: maux
106+
107+
! Small O(N^2) factors of Lambda.
108+
double precision, dimension(:,:), allocatable :: e ! from get_emat
109+
double precision, dimension(:,:), allocatable :: dmat ! D = mat_w/2, from get_g
110+
111+
! Skinny O(N^3) buffers.
112+
double precision, dimension(:,:), allocatable :: v32 ! ns x nl
113+
double precision, dimension(:,:), allocatable :: tmp, cf ! nl x ns
114+
double precision, dimension(:,:,:), allocatable :: b1, b2 ! ns x ns x ns
115+
33116
double precision, dimension(:), allocatable :: work
34117
integer, dimension(:), allocatable :: ipiv
35118
integer :: info
36119

37-
double precision, dimension(:,:), allocatable :: cf
38-
double precision, dimension(:,:), allocatable :: tmp
39-
40-
integer :: mu, nu, alpha
41-
integer :: ka, ja
42-
integer :: i, j, x, y, z, w
43-
44-
real :: t1, t2
120+
integer :: mu, nu
121+
integer :: ka
122+
integer :: i, s, x, y, z, w
45123

46124
logical, parameter :: debug = .true.
47-
logical :: sym
48-
49-
! Get integers
50125

51126
if (debug) then
52-
print *, "=== DEBUG ODD STRAIGHT ==="
127+
print *, "=== DEBUG ODD STRAIGHT (kron, Lambda never materialized) ==="
53128
print *, "N_MODE:", n_mode
54129
print *, "NTYP:", ntyp
55130
print *, "NAT_SC:", nat_sc
56131
call flush()
57132
end if
58133

59-
!nat_sc = size(er(:,1,1))
60-
!n_mode = 3*nat_sc
61-
62134
ns = n_mode
63135
nl = n_mode*n_mode
64136

65-
sym = .false.
66-
if (present(use_v4_symmetry)) sym = use_v4_symmetry
67-
68-
! Allocate stuff
69-
70-
allocate(lamat(nl,nl))
71137
allocate(maux(nl,nl))
138+
allocate(e(ns,ns))
139+
allocate(dmat(ns,ns))
140+
allocate(v32(ns,nl))
141+
allocate(tmp(nl,ns))
142+
allocate(cf(nl,ns))
143+
allocate(b1(ns,ns,ns))
144+
allocate(b2(ns,ns,ns))
72145
allocate(ipiv(nl))
73146
allocate(work(nl))
74-
allocate(v32(n_mode,n_mode*n_mode))
75-
76-
allocate(cf(nl,ns))
77-
allocate(tmp(nl,ns))
78-
79-
! v42 (explicit reordered copy of v4) is only needed when we cannot rely on
80-
! the permutation symmetry of v4. It is one full nl x nl array.
81-
if (.not. sym) allocate(v42(nl,nl))
82147

83-
! Get lambda matrix
84-
85-
call get_cmat ( a, wr, er, transmode, amass, ityp_sc, T, .true., lamat,n_mode, nat_sc, ntyp )
86-
87-
!print *, "AFTER CMAT"
88-
!call flush()
89-
90-
! Write third and fourth order force constants as rank 2
148+
! Small factors of Lambda: exactly what get_cmat used internally.
149+
! (v3_log = .true., as in the old get_cmat call from this routine.)
150+
call get_emat ( er, a, amass, ityp_sc, .true., transmode, e, n_mode, nat_sc, ntyp)
151+
call get_g (a, wr, transmode, T, dmat, n_mode)
152+
dmat = 0.5d0 * dmat ! D(mu,nu) = mat_w(mu,nu)/2 (the 0.5 of mat_et)
91153

154+
! Third order force constants as rank 2, v32(:,ka) with ka=(x-1)*N+y.
92155
ka = 0
93-
94-
do x = 1, n_mode
95-
do y = 1, n_mode
156+
do x = 1, ns
157+
do y = 1, ns
96158
ka = ka + 1
97159
v32(:,ka) = v3(:,x,y)
98-
if (.not. sym) then
99-
ja = 0
100-
do w = 1, n_mode
101-
do z = 1, n_mode
102-
ja = ja + 1
103-
v42(ja,ka) = v4(w,z,x,y)
104-
end do
105-
end do
106-
end if
107160
end do
108161
end do
109162

110-
! Prepare identity matrix directly in maux (was: iden then maux = iden;
111-
! bitwise identical, avoids one nl x nl temporary)
112-
113-
maux = 0.0d0
114-
115-
do x = 1, nl
116-
maux(x,x) = 1.0d0
163+
! ---------------------------------------------------------------------
164+
! Build maux = I - v42.Lambda without forming Lambda (steps A-F above).
165+
! ---------------------------------------------------------------------
166+
167+
! A: B(w,z,x',mu) = sum_y' v4(w,z,x',y') e(mu,y') [v4 -> maux]
168+
! v4 viewed as (N^3 x N) with columns y'; op(B)=e^T has entry
169+
! (y',mu) = e(mu,y'). Fills maux completely (N^3 * N = nl*nl).
170+
call dgemm('N','T', nl*ns, ns, ns, 1.0d0, v4(1,1,1,1), nl*ns, &
171+
e(1,1), ns, 0.0d0, maux(1,1), nl*ns)
172+
173+
! B: T(w,z,nu,mu) = sum_x' B(w,z,x',mu) e(nu,x') [maux -> v4]
174+
! For each mu, the slice B(:,:,:,mu) is the (N^2 x N) block starting
175+
! at maux(1,(mu-1)*ns+1) (linear offset (mu-1)*N^3), columns x';
176+
! result panel written at v4(:,:,:,mu) with layout (w,z,nu).
177+
! From here on the original content of v4 is destroyed.
178+
do mu = 1, ns
179+
call dgemm('N','T', nl, ns, ns, 1.0d0, maux(1,(mu-1)*ns+1), nl, &
180+
e(1,1), ns, 0.0d0, v4(1,1,1,mu), nl)
117181
end do
118182

119-
! Calculate ** iden - v4 lamat ** matrix
120-
121-
!print *, "BEFORE I - V4Lambda"
122-
!call flush()
123-
124-
if (.not. sym) then
125-
call dgemm('N','N',nl,nl,nl,-1.0d0,v42,nl,lamat,nl,1.0d0,maux,nl)
126-
else
127-
! v4 fully permutation-symmetric => the reordered copy v42(ja,ka)=v4(w,z,x,y)
128-
! equals the plain column-major reshape of v4 into an (nl,nl) matrix, so we
129-
! can feed v4 directly to dgemm with leading dimension nl and skip v42.
130-
! (Not valid with use_symmetries=False -> guarded by use_v4_symmetry.)
131-
call dgemm('N','N',nl,nl,nl,-1.0d0,v4,nl,lamat,nl,1.0d0,maux,nl)
132-
end if
133-
134-
if (allocated(v42)) deallocate(v42)
183+
! C: T(w,z,nu,mu) *= D(mu,nu) [v4 in place]
184+
! NOTE the argument order: 4th dim of T is mu, 3rd is nu.
185+
do mu = 1, ns
186+
do nu = 1, ns
187+
v4(:,:,nu,mu) = v4(:,:,nu,mu) * dmat(mu,nu)
188+
end do
189+
end do
135190

136-
! Invert ** iden - lamat v4 **
191+
! D: C(w,z,nu,y) = sum_mu T(w,z,nu,mu) e(mu,y) [v4 -> maux]
192+
call dgemm('N','N', nl*ns, ns, ns, 1.0d0, v4(1,1,1,1), nl*ns, &
193+
e(1,1), ns, 0.0d0, maux(1,1), nl*ns)
137194

138-
!print *, "BEFORE (I - V4Lambda)^-1"
139-
!call flush()
195+
! E: P(w,z,x,y) = sum_nu C(w,z,nu,y) e(nu,x) [maux -> v4]
196+
do y = 1, ns
197+
call dgemm('N','N', nl, ns, ns, 1.0d0, maux(1,(y-1)*ns+1), nl, &
198+
e(1,1), ns, 0.0d0, v4(1,1,1,y), nl)
199+
end do
140200

201+
! F: maux = I - P with BOTH index pairs swapped to the original
202+
! v42.Lambda convention: rows (z fast, w slow), cols (y fast, x slow).
203+
do y = 1, ns
204+
do x = 1, ns
205+
ka = (x-1)*ns + y
206+
do z = 1, ns
207+
do w = 1, ns
208+
maux((w-1)*ns + z, ka) = -v4(w,z,x,y)
209+
end do
210+
end do
211+
end do
212+
end do
213+
do i = 1, nl
214+
maux(i,i) = maux(i,i) + 1.0d0
215+
end do
141216

217+
! Invert ** iden - v4 lamat ** in place (unchanged from v1.5).
142218
call dgetrf ( nl, nl, maux, nl, ipiv, info )
143219
call dgetri ( nl, maux, nl, ipiv, work, nl, info )
144220

145-
! Take product between lamat and the inverted matrix, contracted with v3.
146-
!
147-
! REASSOCIATION (numerically equivalent to ~1e-15 relative, NOT bitwise):
148-
! the original code formed the nl x nl product v42 := lamat * maux and then
149-
! cf := v42 * v3^T. We instead evaluate cf = lamat * (maux * v3^T) using an
150-
! nl x ns buffer tmp, which is mathematically identical by associativity
151-
! ( Lambda * M^-1 * v3^T = Lambda * (M^-1 * v3^T) ) and reduces both the flop
152-
! count (O(N^6) -> O(N^5)) and the memory (no nl x nl buffer needed).
153-
154-
! tmp = (I - v4 lamat)^-1 * v3^T (nl x ns)
155-
call dgemm('N','T',nl,ns,nl,1.0d0,maux,nl,&
156-
v32,ns,0.0d0,tmp,nl)
157-
158-
! cf = lamat * tmp = lamat (I - v4 lamat)^-1 v3^T (nl x ns)
159-
call dgemm('N','N',nl,ns,nl,1.0d0,lamat,nl,&
160-
tmp,nl,0.0d0,cf,nl)
161-
162-
! Now get:
163-
! v3 * ( 1 - lamat*v4)^-1 lamat * v3
221+
! tmp = (I - v4 lamat)^-1 . v3^T (nl x ns)
222+
call dgemm('N','T', nl, ns, nl, 1.0d0, maux(1,1), nl, &
223+
v32(1,1), ns, 0.0d0, tmp(1,1), nl)
224+
225+
! ---------------------------------------------------------------------
226+
! cf = Lambda . tmp, again without forming Lambda. With tmp's row index
227+
! ka' = (x'-1)*N + y' (y' fast) read as tmp3(y',x',s):
228+
! u(mu,nu,s) = sum_{x',y'} e(mu,y') e(nu,x') tmp3(y',x',s)
229+
! cf3(y,x,s) = sum_{mu,nu} e(mu,y) e(nu,x) D(mu,nu) u(mu,nu,s)
230+
! All buffers are O(N^3).
231+
! ---------------------------------------------------------------------
232+
233+
! C1: b1(mu,x',s) = sum_y' e(mu,y') tmp3(y',x',s)
234+
! tmp reshaped as (N x N*ns), one dgemm.
235+
call dgemm('N','N', ns, nl, ns, 1.0d0, e(1,1), ns, &
236+
tmp(1,1), ns, 0.0d0, b1(1,1,1), ns)
237+
238+
! C2: b2(mu,nu,s) = sum_x' b1(mu,x',s) e(nu,x')
239+
do s = 1, ns
240+
call dgemm('N','T', ns, ns, ns, 1.0d0, b1(1,1,s), ns, &
241+
e(1,1), ns, 0.0d0, b2(1,1,s), ns)
242+
end do
164243

165-
!print *, "BEFORE V3 Lambda(I - V4Lambda)^-1 V3"
166-
!call flush()
167-
call dgemm('N','N',ns,ns,nl,1.0d0,v32,ns,&
168-
cf,nl,0.0d0,phi_sc_odd,ns)
244+
! C3: b2(mu,nu,s) *= D(mu,nu)
245+
do s = 1, ns
246+
b2(:,:,s) = b2(:,:,s) * dmat(:,:)
247+
end do
169248

249+
! C4: b1(y,nu,s) = sum_mu e(mu,y) b2(mu,nu,s) (e^T . b2_s)
250+
do s = 1, ns
251+
call dgemm('T','N', ns, ns, ns, 1.0d0, e(1,1), ns, &
252+
b2(1,1,s), ns, 0.0d0, b1(1,1,s), ns)
253+
end do
170254

171-
!call get_odd_from_cmat_fu2 (v42, v32, phi_sc_odd)
255+
! C5: cf3(y,x,s) = sum_nu b1(y,nu,s) e(nu,x); column s of cf viewed as
256+
! an (N x N) block (y fast, x slow), matching ka = (x-1)*N + y.
257+
do s = 1, ns
258+
call dgemm('N','N', ns, ns, ns, 1.0d0, b1(1,1,s), ns, &
259+
e(1,1), ns, 0.0d0, cf(1,s), ns)
260+
end do
172261

173-
! Deallocate stuff
262+
! phi_sc_odd = v3 . lamat (I - v4 lamat)^-1 . v3
263+
call dgemm('N','N', ns, ns, nl, 1.0d0, v32(1,1), ns, &
264+
cf(1,1), nl, 0.0d0, phi_sc_odd(1,1), ns)
174265

175-
deallocate(lamat,v32,maux,ipiv,work, cf, tmp)
266+
deallocate(maux, e, dmat, v32, tmp, cf, b1, b2, ipiv, work)
176267

177268
end subroutine get_odd_straight_with_v4

0 commit comments

Comments
 (0)