Skip to content

Commit d0368f7

Browse files
mesonepigrecomesonepigreco
authored
Merge branch 'master' into remove-unused-q-start-current-q-arrays
2 parents 68b3dee + 0b39036 commit d0368f7

1 file changed

Lines changed: 212 additions & 103 deletions

File tree

‎SCHAModules/get_odd_straight_with_v4.f90‎

Lines changed: 212 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,86 @@
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, &
786
n_mode, nat_sc, ntyp)
@@ -15,145 +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
20100

21-
22101
integer :: nat_sc, n_mode, nl, ns, ntyp
23-
double precision, dimension(:,:), allocatable :: l, g, phi_aux, v1, v2, v32, iden
24-
double precision :: lsum
25-
double precision, dimension(:), allocatable :: laux1, lres1, veclong
26-
double precision, dimension(:), allocatable :: laux2, lres2
27-
28-
double precision, dimension(:,:), allocatable :: lamat, v42, maux
102+
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+
29116
double precision, dimension(:), allocatable :: work
30117
integer, dimension(:), allocatable :: ipiv
31118
integer :: info
32119

33-
double precision, dimension(:), allocatable :: vv
34-
double precision, dimension(:), allocatable :: ww
35-
double precision, dimension(:,:), allocatable :: zz
36-
double precision, dimension(:,:), allocatable :: cf
37-
38-
integer :: mu, nu, alpha
39-
integer :: ka, ja
40-
integer :: i, j, x, y, z, w
41-
42-
real :: t1, t2
120+
integer :: mu, nu
121+
integer :: ka
122+
integer :: i, s, x, y, z, w
43123

44124
logical, parameter :: debug = .true.
45125

46-
! Get integers
47-
48126
if (debug) then
49-
print *, "=== DEBUG ODD STRAIGHT ==="
127+
print *, "=== DEBUG ODD STRAIGHT (kron, Lambda never materialized) ==="
50128
print *, "N_MODE:", n_mode
51-
print *, "NTYP:", ntyp
129+
print *, "NTYP:", ntyp
52130
print *, "NAT_SC:", nat_sc
53131
call flush()
54132
end if
55133

56-
!nat_sc = size(er(:,1,1))
57-
!n_mode = 3*nat_sc
58-
59134
ns = n_mode
60135
nl = n_mode*n_mode
61136

62-
! Allocate stuff
63-
64-
allocate(lamat(nl,nl))
65-
allocate(v42(nl,nl))
66137
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))
67145
allocate(ipiv(nl))
68146
allocate(work(nl))
69-
allocate(v32(n_mode,n_mode*n_mode))
70-
allocate(iden(nl,nl))
71-
72-
allocate(cf(nl,ns))
73-
74-
allocate(vv(nl*(nl+1)/2))
75-
allocate(ww(nl))
76-
allocate(zz(nl,nl))
77-
78-
! Get lambda matrix
79147

80-
call get_cmat ( a, wr, er, transmode, amass, ityp_sc, T, .true., lamat,n_mode, nat_sc, ntyp )
81-
82-
!print *, "AFTER CMAT"
83-
!call flush()
84-
85-
! 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)
86153

154+
! Third order force constants as rank 2, v32(:,ka) with ka=(x-1)*N+y.
87155
ka = 0
88-
89-
do x = 1, n_mode
90-
do y = 1, n_mode
156+
do x = 1, ns
157+
do y = 1, ns
91158
ka = ka + 1
92159
v32(:,ka) = v3(:,x,y)
93-
ja = 0
94-
do w = 1, n_mode
95-
do z = 1, n_mode
96-
ja = ja + 1
97-
v42(ja,ka) = v4(w,z,x,y)
98-
end do
99-
end do
100160
end do
101161
end do
102162

103-
! Prepare identity matrix
104-
105-
iden = 0.0d0
106-
107-
do x = 1, nl
108-
iden(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)
109181
end do
110182

111-
! Calculate ** iden - v4 lamat ** matrix
112-
113-
!print *, "BEFORE I - V4Lambda"
114-
!call flush()
115-
116-
maux = iden
117-
118-
call dgemm('N','N',nl,nl,nl,-1.0d0,v42,nl,lamat,nl,1.0d0,maux,nl)
119-
120-
! Invert ** iden - lamat v4 **
121-
122-
!print *, "BEFORE (I - V4Lambda)^-1"
123-
!call flush()
124-
125-
126-
call dgetrf ( nl, nl, maux, nl, ipiv, info )
127-
call dgetri ( nl, maux, nl, ipiv, work, nl, info )
128-
129-
! Take product between lamat and the inverted matrix
130-
131-
!print *, "BEFORE Lambda(I - V4Lambda)^-1"
132-
!call flush()
133-
call dgemm('N','N',nl,nl,nl,1.0d0,lamat,nl,maux,nl,0.0d0,v42,nl)
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
134190

135-
! Calculate final matrix products and assign the correction matrix
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)
136194

137-
! Calculate cf = ( 1 - lamat*v4)^-1 lamat * v3
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
138200

139-
!print *, "BEFORE Lambda(I - V4Lambda)^-1 V3"
140-
!call flush()
141-
call dgemm('N','T',nl,ns,nl,1.0d0,v42,nl,&
142-
v32,ns,0.0d0,cf,nl)
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
143216

144-
! Now get:
145-
! v3 * ( 1 - lamat*v4)^-1 lamat * v3
217+
! Invert ** iden - v4 lamat ** in place (unchanged from v1.5).
218+
call dgetrf ( nl, nl, maux, nl, ipiv, info )
219+
call dgetri ( nl, maux, nl, ipiv, work, nl, info )
220+
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
146243

147-
!print *, "BEFORE V3 Lambda(I - V4Lambda)^-1 V3"
148-
!call flush()
149-
call dgemm('N','N',ns,ns,nl,1.0d0,v32,ns,&
150-
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
151248

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
152254

153-
!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
154261

155-
! 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)
156265

157-
deallocate(lamat,v32,v42,maux,ipiv,work, cf)
266+
deallocate(maux, e, dmat, v32, tmp, cf, b1, b2, ipiv, work)
158267

159-
end subroutine get_odd_straight_with_v4
268+
end subroutine get_odd_straight_with_v4

0 commit comments

Comments
 (0)