Snapshot from this documentation build. SHA-256: 0cd68dcfda70cf37d386b1e5edc7befaa7ec788b34bc8e34e97f8d27fcb1b5bc
1 # GPU-resident radix-2 FFT + periodic Poisson root solve (the last CPU step in the
2 # gravity path — FFTW — moved onto the device). Real/imag are kept as SEPARATE
3 # Float32/Float64 arrays (no Complex bitstype in the Metal kernels, for safety).
4 #
5 # 1-D Cooley–Tukey, decimation-in-time: bit-reversal permutation, then log2(N)
6 # butterfly stages (one kernel launch per stage; butterflies within a stage are
7 # independent so they're race-free, and KA queue-orders the stages). 3-D = the
8 # 1-D transform along axis 1, then permutedims to bring axes 2 and 3 to the front.
9 # Only powers of two (the SB root grids are 128³/256³). CPU `fft_poisson_root!`
10 # stays the oracle; this is validated bit-for-(f32)-bit against it.
11
12 using KernelAbstractions: @kernel, @index, @Const
13
14 _ispow2(n) = n > 0 && (n & (n - 1)) == 0
15 _log2i(n) = (m = 0; while (1 << m) < n; m += 1; end; m)
16
17 # bit-reversal index table (0-based) for length N = 2^bits
18 function _bitrev_table(N::Int)
19 bits = _log2i(N); br = Vector{Int32}(undef, N)
20 @inbounds for i in 0:N-1
21 r = 0; x = i
22 for _ in 1:bits; r = (r << 1) | (x & 1); x >>= 1; end
23 br[i+1] = r
24 end
25 br
26 end
27
28 # IN-PLACE bit-reversal along axis 1: swap element i with r=bitrev(i), guarded by
29 # i<r so each pair is touched exactly once (bit-reversal is an involution ⇒ the
30 # partner workitem's guard is false ⇒ no concurrent write, race-free on CPU & GPU).
31 # Avoids the out-of-place scratch write + full copy-back (≈3 array passes/axis).
32 @kernel function _bitrev_inplace_k!(xr, xi, @Const(br))
33 i, j, k = @index(Global, NTuple)
34 @inbounds begin
35 r = br[i] + 1
36 if i < r
37 t = xr[i, j, k]; xr[i, j, k] = xr[r, j, k]; xr[r, j, k] = t
38 t = xi[i, j, k]; xi[i, j, k] = xi[r, j, k]; xi[r, j, k] = t
39 end
40 end
41 end
42
43 # one butterfly stage along axis 1; half = m/2, m = 2^s; sgn=-1 forward, +1 inverse
44 @kernel function _fft_stage_k!(xr, xi, half::Int, m::Int, sgn, twopi)
45 p, j, k = @index(Global, NTuple) # p ∈ 1:N1/2 butterfly index
46 @inbounds begin
47 pp = p - 1; g = pp ÷ half; kk = pp % half
48 i0 = g * m + kk + 1; i1 = i0 + half
49 T = eltype(xr)
50 th = sgn * twopi * T(kk) / T(m)
51 wr = cos(th); wi = sin(th)
52 ar = xr[i0, j, k]; ai = xi[i0, j, k]
53 br_ = xr[i1, j, k]; bi = xi[i1, j, k]
54 tr = wr * br_ - wi * bi; ti = wr * bi + wi * br_
55 xr[i0, j, k] = ar + tr; xi[i0, j, k] = ai + ti
56 xr[i1, j, k] = ar - tr; xi[i1, j, k] = ai - ti
57 end
58 end
59
60 # in-place 1-D FFT along axis 1 of a 3-D (re,im) pair (no scratch — in-place bitrev).
61 function _fft_axis1!(xr, xi, brdev, sgn::Int)
62 be = KA.get_backend(xr); T = eltype(xr); N1, N2, N3 = size(xr)
63 _bitrev_inplace_k!(be)(xr, xi, brdev; ndrange = (N1, N2, N3))
64 twopi = T(2) * T(pi)
65 s = 1
66 while (1 << s) <= N1
67 m = 1 << s; half = m >> 1
68 _fft_stage_k!(be)(xr, xi, half, m, T(sgn), twopi; ndrange = (N1 ÷ 2, N2, N3))
69 s += 1
70 end
71 return xr, xi
72 end
73
74 # Per-(backend,T,N) cached transpose targets tr,ti (the bit-reversal is now in-place,
75 # so no sr,si scratch). Reused across all three axes AND across calls, so a 3-D FFT
76 # allocates NOTHING after the first. Cubic grid ⇒ every buffer shares one N³ shape.
77 const _FFT_SCRATCH = Dict{Any,Any}()
78 _fft_scratch(be, ::Type{T}, N) where {T} =
79 get!(_FFT_SCRATCH, (typeof(be), T, N)) do
80 (KA.zeros(be, T, N...), KA.zeros(be, T, N...))
81 end
82 # Cached imaginary buffer, reused (zeroed) every solve — so the hot gravity FFT allocates NOTHING
83 # per call (the fresh `KA.zeros` xi otherwise collides with the pooled ρpi transient at the ceiling).
84 _fft_imag(be, ::Type{T}, N) where {T} =
85 get!(() -> KA.zeros(be, T, N...), _FFT_SCRATCH, (:imag, typeof(be), T, N))
86
87 """
88 poisson_scratch_i32(be, T, N) -> Int32 view of an N³ FFT scratch buffer (reinterpreted)
89
90 The Stockham FFT's real scratch `sr` (Float32, N³) reinterpreted as an Int32 array of the same shape,
91 for callers that need a large Int32 accumulator BEFORE the FFT runs (e.g. the deterministic CIC
92 deposit). Reusing this buffer avoids a separate N³·4 B allocation at the grid ceiling. The FFT
93 overwrites `sr` on its first stage, so the borrow must be finished before `fft_poisson_root_gpu!`.
94 """
95 poisson_scratch_i32(be, ::Type{T}, N) where {T} = reinterpret(Int32, _fft_scratch(be, T, N)[1])
96
97 # ── axis transposes ───────────────────────────────────────────────────────────
98 # The 3-D FFT brings each axis to the front (axis 1) before the 1-D transform. On
99 # the GPU `permutedims!` is already a tuned device kernel; on the CPU it is Base's
100 # SERIAL transpose — the only non-KA, non-threaded step in the host FFT. So on the
101 # CPU KA backend we transpose with a threaded `@kernel` (one workitem per element,
102 # split over `Threads.nthreads()`), keeping the whole host FFT "purely KA + host
103 # threads"; on every other backend we keep the optimized `permutedims!`.
104 @kernel function _transpose12_k!(dst, @Const(src)) # perm (2,1,3): swap axes 1,2
105 i, j, k = @index(Global, NTuple)
106 @inbounds dst[i, j, k] = src[j, i, k]
107 end
108 @kernel function _transpose13_k!(dst, @Const(src)) # perm (3,2,1): swap axes 1,3
109 i, j, k = @index(Global, NTuple)
110 @inbounds dst[i, j, k] = src[k, j, i]
111 end
112
113 # out-of-place transpose dst ← permutedims(src, perm); KA-threaded on CPU, permutedims! else
114 @inline function _transpose!(dst, src, perm::NTuple{3,Int}, be)
115 if be isa KA.CPU
116 if perm === (2, 1, 3)
117 _transpose12_k!(be)(dst, src; ndrange = size(dst))
118 else
119 _transpose13_k!(be)(dst, src; ndrange = size(dst))
120 end
121 else
122 permutedims!(dst, src, perm)
123 end
124 return dst
125 end
126
127 # 3-D FFT (sgn=-1 fwd, +1 inv) — transform axis 1, then bring axes 2,3 to the front via
128 # out-of-place transposes into cached buffers (no per-call allocation). On the CPU
129 # the transposes are threaded KA kernels (see `_transpose!`); on the GPU, permutedims!.
130 function _fft3d!(xr, xi, brdev, sgn::Int)
131 be = KA.get_backend(xr); T = eltype(xr); N = size(xr)
132 tr, ti = _fft_scratch(be, T, N)
133 _fft_axis1!(xr, xi, brdev, sgn) # axis 1
134 _transpose!(tr, xr, (2, 1, 3), be); _transpose!(ti, xi, (2, 1, 3), be)
135 _fft_axis1!(tr, ti, brdev, sgn) # axis 2
136 _transpose!(xr, tr, (2, 1, 3), be); _transpose!(xi, ti, (2, 1, 3), be)
137 _transpose!(tr, xr, (3, 2, 1), be); _transpose!(ti, xi, (3, 2, 1), be)
138 _fft_axis1!(tr, ti, brdev, sgn) # axis 3
139 _transpose!(xr, tr, (3, 2, 1), be); _transpose!(xi, ti, (3, 2, 1), be)
140 return xr, xi
141 end
142
143 # ── Mixed-radix (2·3·5) Stockham autosort — non-power-of-two cubic grids ─────────
144 # The radix-2 path above is in-place (bit-reversal); mixed radix uses a self-sorting
145 # Stockham (no digit-reversal) that ping-pongs between two buffer pairs. This is what
146 # lets the gravity FFT run at 2ᵃ3ᵇ5ᶜ sizes (768/864/900/960/1024) with a SMALL, poolable
147 # scratch instead of cuFFT's opaque ~18 B/cell work area. Validated bit-close vs FFTW.
148 _factorize235(n::Int) = begin
149 fs = Int[]; m = n
150 for p in (2, 3, 5); while m % p == 0; push!(fs, p); m ÷= p; end; end
151 m == 1 ? fs : nothing # nothing ⇒ has a prime factor > 5
152 end
153
154 # One Stockham stage along axis 1. The twiddle is FUSED into the radix-R DFT (combined phase
155 # r·((jz%Ns)+kk·Ns)/(Ns·R)) so each output reads the R inputs directly — no per-thread local
156 # array (keeps it backend-agnostic; R ≤ 5 ⇒ ≤25 reads/butterfly). Self-sorting scatter (idxD).
157 @kernel function _stockham_stage_k!(outr, outi, @Const(inr), @Const(ini),
158 R::Int, Ns::Int, NR::Int, sgn, twopi)
159 j, a2, a3 = @index(Global, NTuple)
160 @inbounds begin
161 T = eltype(outr); jz = j - 1
162 base = jz % Ns; idxD = (jz ÷ Ns) * Ns * R + base
163 for kk in 0:R-1
164 ph0 = sgn * twopi * T(base + kk*Ns) / T(Ns*R)
165 sr = zero(T); si = zero(T)
166 for r in 0:R-1
167 xr = inr[jz + r*NR + 1, a2, a3]; xi = ini[jz + r*NR + 1, a2, a3]
168 c = cos(T(r)*ph0); s = sin(T(r)*ph0)
169 sr += c*xr - s*xi; si += c*xi + s*xr
170 end
171 outr[idxD + kk*Ns + 1, a2, a3] = sr; outi[idxD + kk*Ns + 1, a2, a3] = si
172 end
173 end
174 end
175
176 # 1-D Stockham along axis 1 (ping-pong cr↔or); result left in (cr,ci), (or,oi) free as scratch.
177 function _stockham_axis1!(cr, ci, or, oi, factors, sgn::Int)
178 be = KA.get_backend(cr); T = eltype(cr); N1, N2, N3 = size(cr); tp = T(2) * T(pi)
179 ar, ai, brr, bii = cr, ci, or, oi; Ns = 1
180 for R in factors
181 _stockham_stage_k!(be)(brr, bii, ar, ai, R, Ns, N1 ÷ R, T(sgn), tp; ndrange = (N1 ÷ R, N2, N3))
182 ar, brr = brr, ar; ai, bii = bii, ai; Ns *= R
183 end
184 ar === cr || (copyto!(cr, ar); copyto!(ci, ai)) # odd #stages ⇒ result in scratch: copy back
185 return nothing
186 end
187
188 # 3-D mixed-radix FFT with ONLY two buffer pairs (xr,xi)+(sr,si): each axis Stockham leaves its
189 # result in the first pair; the out-of-place transpose ping-pongs into the other — no 3rd buffer.
190 function _fft3d_mixed!(xr, xi, sr, si, factors, sgn::Int)
191 be = KA.get_backend(xr)
192 _stockham_axis1!(xr, xi, sr, si, factors, sgn) # axis 1 → (xr,xi)
193 _transpose!(sr, xr, (2,1,3), be); _transpose!(si, xi, (2,1,3), be) # (sr,si) ← transpose
194 _stockham_axis1!(sr, si, xr, xi, factors, sgn) # axis 2 → (sr,si)
195 _transpose!(xr, sr, (2,1,3), be); _transpose!(xi, si, (2,1,3), be) # back → (xr,xi)
196 _transpose!(sr, xr, (3,2,1), be); _transpose!(si, xi, (3,2,1), be) # (sr,si) ← axis1↔3
197 _stockham_axis1!(sr, si, xr, xi, factors, sgn) # axis 3 → (sr,si)
198 _transpose!(xr, sr, (3,2,1), be); _transpose!(xi, si, (3,2,1), be) # back → (xr,xi)
199 return xr, xi
200 end
201
202 # ── Real-input FFT (rfft) Poisson — ~½ the FLOPs, memory, and transpose volume of the c2c path ──
203 # rfft along axis-1: pack the n reals into M=n/2 complex, c2c(M), Hermitian split → the (M+1,n,n)
204 # half-complex grid; then c2c along axes 2,3 on the half-grid, Green's, and the inverse. Needs an
205 # EVEN 2·3·5-smooth n (with M=n/2 also 2·3·5-smooth). Validated bit-close vs the c2c solver / FFTW.
206 @kernel function _rpack_k!(zr, zi, @Const(x)) # (n,n,n) real → (M,n,n) complex
207 m, j, k = @index(Global, NTuple)
208 @inbounds begin zr[m,j,k] = x[2m-1,j,k]; zi[m,j,k] = x[2m,j,k] end
209 end
210 @kernel function _runpack_k!(x, @Const(zr), @Const(zi)) # (M,n,n) complex → (n,n,n) real
211 m, j, k = @index(Global, NTuple)
212 @inbounds begin x[2m-1,j,k] = zr[m,j,k]; x[2m,j,k] = zi[m,j,k] end
213 end
214 # Hermitian split Z (M,n,n) → X (M+1,n,n): E=½(Z[k]+conj Z[M-k]), O=(1/2i)(Z[k]−conj Z[M-k]), X=E+W_N^k O.
215 @kernel function _rsplit_k!(Xr, Xi, @Const(Zr), @Const(Zi), M::Int, tw)
216 kk, j, k = @index(Global, NTuple) # kk∈1:M+1, k1=kk-1
217 @inbounds begin
218 T = eltype(Xr); k1 = kk-1; km = mod(M-k1, M)
219 zkr = Zr[k1%M+1,j,k]; zki = Zi[k1%M+1,j,k]; zcr = Zr[km+1,j,k]; zci = -Zi[km+1,j,k]
220 Er = T(0.5)*(zkr+zcr); Ei = T(0.5)*(zki+zci); Or = T(0.5)*(zki-zci); Oi = T(-0.5)*(zkr-zcr)
221 wv = tw*T(k1); wr = cos(wv); ws = sin(wv) # W_N^{k1}=exp(sgn·2πk1/n)
222 Xr[kk,j,k] = Er + wr*Or - ws*Oi; Xi[kk,j,k] = Ei + wr*Oi + ws*Or
223 end
224 end
225 # inverse split X (M+1,n,n) → Z (M,n,n): E=½(X[k]+conj X[M-k]), O=½W^{-k}(X[k]−conj X[M-k]), Z=E+iO.
226 @kernel function _runsplit_k!(Zr, Zi, @Const(Xr), @Const(Xi), M::Int, tw)
227 kk, j, k = @index(Global, NTuple) # kk∈1:M, k1=kk-1
228 @inbounds begin
229 T = eltype(Zr); k1 = kk-1; mk = M-k1
230 xkr = Xr[kk,j,k]; xki = Xi[kk,j,k]; xcr = Xr[mk+1,j,k]; xci = -Xi[mk+1,j,k]
231 Er = T(0.5)*(xkr+xcr); Ei = T(0.5)*(xki+xci); Dr = xkr-xcr; Di = xki-xci
232 wv = -tw*T(k1); wr = cos(wv); ws = sin(wv) # W_N^{-k1}
233 Or = T(0.5)*(wr*Dr - ws*Di); Oi = T(0.5)*(wr*Di + ws*Dr)
234 Zr[kk,j,k] = Er - Oi; Zi[kk,j,k] = Ei + Or # Z = E + i·O
235 end
236 end
237 # Green's −1/k² on the half-grid (M+1,n,n): kx = 0..M (the non-negative rfft half), ky,kz folded.
238 @kernel function _scale_half_k!(Xr, Xi, coef, n::Int, w1, w2, w3)
239 i, j, k = @index(Global, NTuple)
240 @inbounds begin
241 T = eltype(Xr)
242 kx = w1*T(i-1); ky = w2*T((j-1)<=n÷2 ? (j-1) : (j-1-n)); kz = w3*T((k-1)<=n÷2 ? (k-1) : (k-1-n))
243 ks = kx*kx+ky*ky+kz*kz; c = ks==zero(T) ? zero(T) : coef*(-one(T)/ks)
244 Xr[i,j,k]*=c; Xi[i,j,k]*=c
245 end
246 end
247 # c2c along axis 2 or 3 of a complex (n1,n2,n3): transpose-to-front, Stockham, transpose back.
248 # Only ONE extra pair (tr,ti): after the transpose Xr,Xi→tr,ti the source Xr,Xi is FREE, so it doubles
249 # as the Stockham ping-pong (reshaped to the transposed layout) ⇒ the rfft needs just 2 buffer pairs.
250 function _c2c_axis!(Xr, Xi, tr, ti, axis::Int, factors, sgn::Int, be)
251 dims = size(Xr); perm = axis == 2 ? (2,1,3) : (3,2,1); tsh = ntuple(d -> dims[perm[d]], 3); nel = prod(dims)
252 trv = reshape(view(vec(tr),1:nel),tsh); tiv = reshape(view(vec(ti),1:nel),tsh)
253 _transpose!(trv, Xr, perm, be); _transpose!(tiv, Xi, perm, be) # Xr,Xi → tr,ti; Xr,Xi now free
254 srv = reshape(view(vec(Xr),1:nel),tsh); siv = reshape(view(vec(Xi),1:nel),tsh) # ping-pong ← freed Xr,Xi
255 _stockham_axis1!(trv, tiv, srv, siv, factors, sgn) # length dims[perm[1]] along axis 1
256 _transpose!(Xr, trv, perm, be); _transpose!(Xi, tiv, perm, be) # perm is its own inverse
257 return nothing
258 end
259 # cached rfft buffers: 2 pairs of (M+1,n,n) ≈ 8 B/cell (the c2c path is 12). Allocated once (warm-up).
260 _rfft_bufs(be, ::Type{T}, n) where {T} = get!(_FFT_SCRATCH, (:rfft, typeof(be), T, n)) do
261 M = n÷2
262 (Xr=KA.zeros(be,T,M+1,n,n), Xi=KA.zeros(be,T,M+1,n,n), tr=KA.zeros(be,T,M+1,n,n), ti=KA.zeros(be,T,M+1,n,n))
263 end
264
265 """
266 fft_poisson_rfft_ka!(phi, rho; G, a, boxsize) -> phi
267
268 Periodic Poisson solve with the in-tree REAL FFT (rfft): ~half the work/memory of `fft_poisson_root_gpu!`
269 (the half-complex grid is (n/2+1)×n×n). Even, 2·3·5-smooth cubic `n` with n/2 also 2·3·5-smooth; else
270 it delegates to the c2c solver. Bit-close to `fft_poisson_root_gpu!` / FFTW.
271 """
272 function fft_poisson_rfft_ka!(phi::AbstractArray{T,3}, rho::AbstractArray{T,3};
273 G::Real=1.0, a::Real=1.0, boxsize=1.0) where {T}
274 be = KA.get_backend(rho); N = size(rho); n = N[1]
275 all(==(n), N) || error("fft_poisson_rfft_ka!: cubic grid required, got $N")
276 M = n ÷ 2; facM = iseven(n) ? _factorize235(M) : nothing; facN = _factorize235(n)
277 (facM === nothing || facN === nothing) &&
278 return fft_poisson_root_gpu!(phi, rho; G=G, a=a, boxsize=boxsize) # fall back to c2c
279 L = boxsize isa Number ? ntuple(_->Float64(boxsize),3) : ntuple(d->Float64(boxsize[d]),3)
280 w = ntuple(d->T(2π)/T(L[d]),3); coef = T(G)/T(a); tw = -T(2π)/T(n)
281 B = _rfft_bufs(be, T, n); Mnn = M*n*n # cached (allocated once in the warm-up)
282 zr = reshape(view(vec(B.tr),1:Mnn),M,n,n); zi = reshape(view(vec(B.ti),1:Mnn),M,n,n) # packed/axis-1 ← tr,ti
283 zsr = reshape(view(vec(B.Xr),1:Mnn),M,n,n); zsi = reshape(view(vec(B.Xi),1:Mnn),M,n,n) # axis-1 ping-pong ← Xr,Xi
284 _rpack_k!(be)(zr, zi, rho; ndrange=(M,n,n))
285 _stockham_axis1!(zr, zi, zsr, zsi, facM, -1) # rfft axis-1: c2c(M) on packed
286 _rsplit_k!(be)(B.Xr, B.Xi, zr, zi, M, tw; ndrange=(M+1,n,n)) # → half-complex (Xr,Xi)
287 _c2c_axis!(B.Xr, B.Xi, B.tr, B.ti, 2, facN, -1, be)
288 _c2c_axis!(B.Xr, B.Xi, B.tr, B.ti, 3, facN, -1, be)
289 _scale_half_k!(be)(B.Xr, B.Xi, coef, n, w[1], w[2], w[3]; ndrange=(M+1,n,n))
290 _c2c_axis!(B.Xr, B.Xi, B.tr, B.ti, 3, facN, +1, be)
291 _c2c_axis!(B.Xr, B.Xi, B.tr, B.ti, 2, facN, +1, be)
292 _runsplit_k!(be)(zr, zi, B.Xr, B.Xi, M, tw; ndrange=(M,n,n)) # Xr,Xi → zr,zi (tr,ti)
293 _stockham_axis1!(zr, zi, zsr, zsi, facM, +1) # irfft axis-1
294 _runpack_k!(be)(phi, zr, zi; ndrange=(M,n,n))
295 phi .*= T(2)/T(n)^3 # axis-1 round-trips ×M=×n/2 (inner c2c is length n/2) ⇒ total n³/2
296 KA.synchronize(be)
297 return phi
298 end
299
300 # full (c2c) periodic Green's function G(k) = -1/k², G(0)=0, on the device.
301 function _greens_full(be, ::Type{T}, N::NTuple{3,Int}, L::NTuple{3,Float64}) where {T}
302 Gk = Array{T,3}(undef, N)
303 w = ntuple(d -> T(2π) / T(L[d]), 3)
304 @inbounds for k in 0:N[3]-1
305 kz = w[3] * T(k <= N[3] ÷ 2 ? k : k - N[3])
306 for j in 0:N[2]-1
307 ky = w[2] * T(j <= N[2] ÷ 2 ? j : j - N[2])
308 for i in 0:N[1]-1
309 kx = w[1] * T(i <= N[1] ÷ 2 ? i : i - N[1])
310 ks = kx*kx + ky*ky + kz*kz
311 Gk[i+1, j+1, k+1] = ks == zero(T) ? zero(T) : -one(T) / ks
312 end
313 end
314 end
315 to_device(be, Gk, T)
316 end
317
318 const _GPUFFT_CACHE = Dict{Any,Any}()
319
320 @kernel function _scale_by_k!(xr, xi, @Const(Gk), coef)
321 i, j, k = @index(Global, NTuple)
322 @inbounds begin
323 c = coef * Gk[i, j, k]
324 xr[i, j, k] *= c; xi[i, j, k] *= c
325 end
326 end
327
328 # On-the-fly Green's multiply: φ̂ = coef·(−1/k²)·ρ̂ with k² formed from the indices — no cached
329 # N³ Gk array (−4 B/cell; at the grid ceiling that array is the difference between fitting and OOM).
330 @kernel function _scale_by_k_onfly!(xr, xi, coef, n1::Int, n2::Int, n3::Int, w1, w2, w3)
331 i, j, k = @index(Global, NTuple)
332 @inbounds begin
333 T = eltype(xr)
334 kx = w1 * T((i-1) <= n1÷2 ? (i-1) : (i-1-n1))
335 ky = w2 * T((j-1) <= n2÷2 ? (j-1) : (j-1-n2))
336 kz = w3 * T((k-1) <= n3÷2 ? (k-1) : (k-1-n3))
337 ks = kx*kx + ky*ky + kz*kz
338 c = ks == zero(T) ? zero(T) : coef * (-one(T) / ks)
339 xr[i, j, k] *= c; xi[i, j, k] *= c
340 end
341 end
342
343 """
344 fft_poisson_root_gpu!(phi, rho; G=1.0, a=1.0, boxsize=1.0) -> phi
345
346 GPU-resident periodic Poisson solve (the device counterpart of `fft_poisson_root!`):
347 forward radix-2 FFT of `rho` → multiply by `coef·G(k)` → inverse FFT → real part,
348 all on `rho`'s backend. Dimensions must be powers of two. Bit-identical (to f32) with
349 the FFTW path. `phi`, `rho` are 3-D device arrays of the same shape.
350 """
351 function fft_poisson_root_gpu!(phi::AbstractArray{T,3}, rho::AbstractArray{T,3};
352 G::Real = 1.0, a::Real = 1.0, boxsize = 1.0) where {T}
353 be = KA.get_backend(rho); N = size(rho)
354 all(==(N[1]), N) || error("fft_poisson_root_gpu! needs a cubic grid, got $N")
355 factors = _factorize235(N[1])
356 factors === nothing && error("fft_poisson_root_gpu! needs a 2·3·5-smooth dim, got $(N[1])")
357 L = boxsize isa Number ? ntuple(_ -> Float64(boxsize), 3) : ntuple(d -> Float64(boxsize[d]), 3)
358 pow2 = _ispow2(N[1])
359 brdev = pow2 ? get!(() -> to_device(be, _bitrev_table(N[1]), Int32),
360 _GPUFFT_CACHE, (:brev, typeof(be), T, N[1])) : nothing
361 w = ntuple(d -> T(2π) / T(L[d]), 3); coef = T(G) / T(a) # G(k)=−1/k² formed on the fly (no Gk array)
362 # In-place when φ aliases ρ (the GRAV1BUF shared buffer): use `rho` as the real part directly —
363 # no N³ copy (−4 B/cell), and the FFT working set is ρ + xi + one scratch pair, not two copies.
364 xr = phi === rho ? rho : copyto!(similar(rho), rho)
365 xi = _fft_imag(be, T, N); fill!(xi, zero(T)) # cached (allocated once, in the warmup's clean pool)
366 if pow2 # radix-2, in-place bit-reversal (SB root grids)
367 _fft3d!(xr, xi, brdev, -1)
368 _scale_by_k_onfly!(be)(xr, xi, coef, N[1], N[2], N[3], w[1], w[2], w[3]; ndrange = N)
369 _fft3d!(xr, xi, brdev, +1)
370 else # mixed radix 2·3·5, Stockham (2 buffer pairs)
371 sr, si = _fft_scratch(be, T, N)
372 _fft3d_mixed!(xr, xi, sr, si, factors, -1)
373 _scale_by_k_onfly!(be)(xr, xi, coef, N[1], N[2], N[3], w[1], w[2], w[3]; ndrange = N)
374 _fft3d_mixed!(xr, xi, sr, si, factors, +1)
375 end
376 invN = one(T) / T(prod(N))
377 @inbounds phi .= xr .* invN # real part, normalized
378 KA.synchronize(be)
379 return phi
380 end
381
382 """
383 power_spectrum_gpu(delta; boxsize=1.0, nbins=0) -> (k, P, Nmodes)
384
385 Isotropic power spectrum P(k) of a real overdensity field `delta` (a cubic,
386 power-of-two device array), computed with the GPU radix-2 FFT. Convention: the
387 mean-normalized transform δ̂(k) = (1/N³) Σ_x δ(x) e^{-ik·x} and
388 P(k) = V ⟨|δ̂(k)|²⟩ with V = boxsize³ — so a white-noise field of variance σ²
389 returns P = σ²·V/N³ (the shot/grid floor). Modes are radially binned in |k| in
390 units of the fundamental k_f = 2π/boxsize; `nbins` defaults to N/2 linear bins
391 out to the Nyquist k_f·N/2. Returns bin-centre wavenumbers `k` (physical, =
392 2π/length), the binned power `P`, and the mode count `Nmodes` per bin (bins with
393 no modes are dropped). The FFT runs on `delta`'s backend; only the (small)
394 |δ̂|² radial binning is on the host.
395 """
396 function power_spectrum_gpu(delta::AbstractArray{T,3}; boxsize = 1.0, nbins::Integer = 0) where {T}
397 be = KA.get_backend(delta); N = size(delta)
398 all(_ispow2, N) || error("power_spectrum_gpu needs power-of-two dims, got $N")
399 all(==(N[1]), N) || error("power_spectrum_gpu currently needs a cubic grid, got $N")
400 n = N[1]
401 L = boxsize isa Number ? Float64(boxsize) : Float64(boxsize[1])
402 brdev = get!(_GPUFFT_CACHE, (:brev, typeof(be), T, n)) do
403 to_device(be, _bitrev_table(n), Int32)
404 end
405 xr = copyto!(similar(delta), delta); xi = KA.zeros(be, T, N...)
406 xr, xi = _fft3d!(xr, xi, brdev, -1) # forward DFT (unnormalized)
407 KA.synchronize(be)
408 ar = Array(xr); ai = Array(xi) # bring |δ̂|² home for binning
409 invN3 = 1.0 / Float64(n)^3
410 V = L^3
411 nb = nbins > 0 ? Int(nbins) : n ÷ 2
412 kf = 2π / L # fundamental wavenumber
413 kny_modes = n / 2 # Nyquist in fundamental units
414 Psum = zeros(Float64, nb); ksum = zeros(Float64, nb); cnt = zeros(Int, nb)
415 @inbounds for k in 0:n-1
416 kz = k <= n ÷ 2 ? k : k - n
417 for j in 0:n-1
418 ky = j <= n ÷ 2 ? j : j - n
419 for i in 0:n-1
420 kx = i <= n ÷ 2 ? i : i - n
421 km = sqrt(Float64(kx*kx + ky*ky + kz*kz)) # |k| in fundamental units
422 (km <= 0 || km > kny_modes) && continue
423 b = clamp(1 + floor(Int, (km / kny_modes) * nb), 1, nb)
424 re = Float64(ar[i+1, j+1, k+1]) * invN3
425 im = Float64(ai[i+1, j+1, k+1]) * invN3
426 Psum[b] += V * (re*re + im*im)
427 ksum[b] += km * kf
428 cnt[b] += 1
429 end
430 end
431 end
432 keep = cnt .> 0
433 kc = (ksum[keep] ./ cnt[keep])
434 Pc = (Psum[keep] ./ cnt[keep])
435 return (k = kc, P = Pc, Nmodes = cnt[keep])
436 end
437
438 """
439 power_spectrum_aniso_gpu(field; boxsize=1.0, nmu=4, nbins=0, axis=1, kmax=nothing)
440 -> (k, P, Nmodes)
441
442 Anisotropic P(k,μ) on the device. `field` is a real cubic power-of-two device array
443 (scalar overdensity) OR a 3-tuple of such arrays (a vector field → power Σ_c|f̂_c|², e.g.
444 velocity). μ = |k_axis|/|k| (the v_bc stream is along `axis`; default 1). Returns the
445 radial-bin centre wavenumbers `k` (physical), the binned power `P` of shape `(nbins, nmu)`,
446 and the per-(k,μ) mode count `Nmodes`. Same normalization as [`power_spectrum_gpu`].
447 `kmax` optionally fixes the physical upper edge of the radial bins, allowing spectra from
448 different grid resolutions to use identical bins. It must not exceed the grid Nyquist.
449 The FFT(s) run on `field`'s backend; only the (k,μ) binning of |f̂|² is on the host.
450 """
451 function power_spectrum_aniso_gpu(field; boxsize = 1.0, nmu::Integer = 4,
452 nbins::Integer = 0, axis::Integer = 1,
453 kmax::Union{Nothing,Real} = nothing)
454 fields = field isa Tuple ? field : (field,)
455 f1 = fields[1]; be = KA.get_backend(f1); N = size(f1); T = eltype(f1)
456 all(_ispow2, N) || error("power_spectrum_aniso_gpu needs power-of-two dims, got $N")
457 all(==(N[1]), N) || error("power_spectrum_aniso_gpu needs a cubic grid, got $N")
458 n = N[1]; L = Float64(boxsize isa Number ? boxsize : boxsize[1])
459 brdev = get!(_GPUFFT_CACHE, (:brev, typeof(be), T, n)) do
460 to_device(be, _bitrev_table(n), Int32)
461 end
462 Pacc = nothing # Σ_components |f̂|² (unnormalized), on host
463 for f in fields
464 xr = copyto!(similar(f), f); xi = KA.zeros(be, T, N...)
465 xr, xi = _fft3d!(xr, xi, brdev, -1); KA.synchronize(be) # forward DFT on DEVICE
466 ar = Array(xr); ai = Array(xi)
467 p = Float64.(ar).^2 .+ Float64.(ai).^2
468 Pacc = Pacc === nothing ? p : (Pacc .+ p)
469 end
470 invN3 = 1.0/Float64(n)^3; V = L^3; nb = nbins > 0 ? Int(nbins) : n ÷ 2
471 kf = 2π/L; kny_grid = n/2
472 kny = kmax === nothing ? kny_grid : Float64(kmax) / kf
473 0 < kny <= kny_grid ||
474 error("kmax must satisfy 0 < kmax <= grid Nyquist $(kny_grid*kf), got $kmax")
475 Psum = zeros(Float64, nb, nmu); ksum = zeros(Float64, nb); cnt = zeros(Int, nb, nmu)
476 @inbounds for k in 0:n-1
477 kz = k <= n÷2 ? k : k - n
478 for j in 0:n-1
479 ky = j <= n÷2 ? j : j - n
480 for i in 0:n-1
481 kx = i <= n÷2 ? i : i - n
482 km = sqrt(Float64(kx*kx + ky*ky + kz*kz))
483 (km <= 0 || km > kny) && continue
484 ka = axis == 1 ? kx : axis == 2 ? ky : kz
485 b = clamp(1 + floor(Int, (km/kny)*nb), 1, nb)
486 mb = clamp(1 + floor(Int, (abs(Float64(ka))/km)*nmu), 1, nmu)
487 Psum[b,mb] += V * Pacc[i+1,j+1,k+1] * invN3 * invN3
488 ksum[b] += km*kf; cnt[b,mb] += 1
489 end
490 end
491 end
492 tot = vec(sum(cnt, dims=2))
493 kc = [tot[b] > 0 ? ksum[b]/tot[b] : NaN for b in 1:nb]
494 P = [cnt[b,m] > 0 ? Psum[b,m]/cnt[b,m] : NaN for b in 1:nb, m in 1:nmu]
495 return (k = kc, P = P, Nmodes = cnt)
496 end