Speed up FFT demag: parallelise glue loops + FFT-friendly padding - #174
Merged
Conversation
The demag field eval is ~66% of LLG runtime. The FFTs were already r2c/c2r, FFTW_MEASURE, threaded and aligned, but two things were leaving CPU on the table: 1. Parallelise the O(N) glue loops around the FFTs. The spectral tensor multiply (H = N.M; 9 complex-mul + 6 add per point) ran single-threaded while the FFTs used all cores; also it iterated over total_length instead of the lenz*leny*(lenx/2+1) valid r2c spectrum, doing ~2x the needed work. Now bounded to nfreq and OpenMP-parallelised, along with the pack/zero/unpack loops (collapse(2) so 2D meshes parallelise too). 2. Pad to the next even 7-smooth (2,3,5,7) size >= 2n instead of exactly 2n, so awkward mesh sizes avoid FFTW's slow large-prime path (e.g. 2*127 = 254 = 2x127). The tensor Nyquist-zeroing is generalised from i==nx to i==lenx/2, which is identical for the 2n case, so good sizes are unchanged. Measured (2D, 4 threads), demag field per call: n=32 0.078 -> 0.037 ms (2.1x) n=64 0.301 -> 0.121 ms (2.5x) n=128 1.052 -> 0.468 ms (2.2x) n=127 2.184 -> 0.464 ms (4.7x, was large-prime bound) Field magnitudes are bit-for-bit unchanged; test_demag_libraries (FFT vs brute-force DemagFull, cuboid + hexagonal) passes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Unrelated to the FFT-speedup change, but surfaced by its CI: fill_demag_tensors_c mapped padded indices to the physical grid via x = abs(i-nx+1), which reaches x == nx at the last padded index (i = lenx-1). Reading tensors[id + c*nx*ny*nz] there indexes past the caller's 6*nx*ny*nz buffer, returning heap garbage - an intermittent NaN, most visible in tensor_yz (largest component offset, so it reads furthest past the end). Flaky by heap luck, which is why only the 3.14 CI job caught it. Zero the padded wrap edge (x>=nx || y>=ny || z>=nz), matching the Nyquist zeroing in compute_demag_tensors; those points have no physical counterpart. Strengthen test_demag_2d_pbc: assert np.isfinite (catches inf too, not just NaN) and that the padded wrap planes are exactly zero. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
davidcortesortuno
merged commit Aug 13, 2026
23ee98a
into
port-to-uv-scikit-core-build
8 of 9 checks passed
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
1 participant
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.
Speed up the FFT demag (CPU-only)
Profiling an LLG run shows the demag field evaluation is ~66% of runtime (the FFT convolution + the spectral tensor multiply), so any win here compounds across every CPU user. The FFTs were already well set up (r2c/c2r,
FFTW_MEASURE, threaded,fftw_malloc-aligned); this PR fixes two things around them.1. Parallelise the O(N) glue loops
H = N·M, 9 complex-mul + 6 add per point) ran single-threaded while the FFTs used all cores. It also iterated overtotal_lengthinstead of thelenz·leny·(lenx/2+1)valid r2c spectrum — ~2× the necessary work. Now bounded tonfreqand OpenMP-parallelised.collapse(2)so 2D meshes,nz=1, also thread).2. FFT-friendly padding
≥ 2ninstead of exactly2n, so awkward mesh sizes avoid FFTW's slow large-prime path (e.g.2·127 = 254 = 2×127).i==nx → i==lenx/2— identical for the2ncase, so good sizes are byte-for-byte unchanged.Measured (2D, 4 threads) — demag field per call
≈ 1.4× overall for demag-bound simulations. No GPU required — this is the CPU path everyone uses.
Correctness
n=127).tests/test_demag_libraries.py(FFT demag vs brute-forceDemagFull, cuboid and hexagonal meshes) passes.Not included (measured, rejected)
fftw_plan_many: measured; only helped tinyn=32(−18%, already negligible in absolute terms) and was within noise for the sizes that matter, since each FFT already saturates the cores. Not worth the added allocation/stride complexity.Notes
port-to-uv-scikit-core-build(one commit) — intended for review after the uv branch is accepted.demag.cis shared withmaster, so this applies cleanly to either line.🤖 Generated with Claude Code