Real plans over any number of dimensions - #138
Open
pankgeorg wants to merge 11 commits into
Open
Conversation
With FFTW.jl loaded alongside FFTA, plan_rfft(::Vector{Float64}, ::Int) was
ambiguous between FFTW's StridedArray method and FFTA's method annotated
with region::RegionTypes, turning rfft(x) into a MethodError. Leave region
unannotated on the AbstractFFTs entry points (as plan_fft already does) and
normalise it in an internal function, so FFTW's methods are strictly more
specific and take over as AbstractFFTs intends.
A coexistence test runs in a subprocess (loading FFTW in the test process
would make every other test exercise FFTW).
Real-input/real-output plans only implemented *, so every rfft/irfft allocated, mul!(y, p, x) with a preallocated output was a MethodError, rfft along one dimension of an N-d array went through mapslices, and 2D real plans ran a full complex transform and discarded half of it. FFTAPlan_re now carries a scratch buffer and two pencil kernels (_rfft_pencil!/_brfft_pencil!) implement the even-length half-size trick and the odd-length transform on AbstractVectors, so views work as input and output. mul! is defined for 1D plans on 1D and N-d arrays (looping over pencils along the region) and for 2D plans on N-d arrays (real transform along the first region dimension, complex along the second), and * allocates the output and calls mul!. The 2D plan's first call graph is built for the half length like the 1D plan's.
inv(p) threw a TypeError for every FFTA plan: no plan_inv method existed
and the dummy pinv::FFTAInvPlan field made AbstractFFTs' pinv_type
resolve to Union{}. The plan structs are now mutable with an initially
undefined pinv field, as in FFTW.jl, so AbstractFFTs.inv caches the
result of the new plan_inv methods for complex and real plans; p \ x,
ldiv!(y, p, x), plan_ifft and plan_irfft work through them.
plan_fft!/plan_bfft! return an FFTAPlan_inplace wrapping an ordinary plan
plus a buffer: when input and output alias, the input is copied to the
buffer and transformed out of place (FFTA's kernels are out of place, and
the 1D pencil path is not alias-safe); otherwise the wrapped plan is used
directly. fft!, bfft! and ifft! from AbstractFFTs now work — the internal
kernel that shadowed AbstractFFTs.fft! is renamed fft_kernel!.
…llow AbstractFFTs backend hooks in the ExplicitImports check
…e the other AbstractFFTs entry points
… across pencils Multidimensional mul! allocated two pencil buffers per call and copied every pencil in and out of them; a 1D plan applied along one dimension of an array transformed its pencils one after the other; and the single call-graph workspace of a plan made it impossible to use threads. A plan now owns one Worker per thread it may use (num_threads keyword of the plan_* functions, default Threads.nthreads(), kept by inv): call graphs sharing the plan's nodes but with their own workspace, plus the pencil buffers. One driver, _foreach_pencil, runs every pencil loop (1D plans on N-d arrays, N-d plans, real pencil kernels): serially and allocation-free below THREAD_THRESHOLD or with one worker, otherwise as one Threads.@Spawn task per worker over a contiguous chunk of pencils. Each chunk uses its own worker, so results are bit-identical whatever the thread count. Kernels read the strided pencil views directly; only the copy-out from the worker's output buffer remains.
plan_rfft/plan_brfft threw for regions of more than two dimensions. The 2D real path (real transform along the first region dimension, complex transform along the second) generalises directly: forward runs the real pencil kernel along the first region dimension and then the complex pencil pass along each further one; backward does the reverse. DSP.jl's conv on 3-D float arrays needs exactly this.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #138 +/- ##
==========================================
- Coverage 98.80% 98.17% -0.63%
==========================================
Files 5 4 -1
Lines 585 658 +73
==========================================
+ Hits 578 646 +68
- Misses 7 12 +5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…nels Replacing mapslices with strided views made real transforms along dims=2 of wide matrices (64 x N) 1.2-1.3x slower on x86-64: the mapslices copy had been an unlabelled copy-in that turned a stride-of-a-cache-line gather into one contiguous pass before the kernel. Pencils whose parent arrays are unit-stride along the transform dimension still go to the kernels directly (the dims=1 gain stays); any other pencil is copied to two plan-owned contiguous buffers first and copied back after, so execution stays allocation-free and the dims=2 result is now identical to the mapslices one.
…d-buffers-threads The copy buffers live in each Worker so that pencils transformed on different tasks never share them.
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.
Stacked on #137 (the diff includes #132, #133 and #137 until they land).
plan_rfft(x, 1:3)/plan_brfft(y, n, 1:3)threwArgumentError("only supports 1D and 2D FFTs"). The 2D real path introduced in #132 — a real transform along the first region dimension (half length when even), then a complex transform along the second — generalises directly to any number of region dimensions: forward runs the real pencil kernel along the first region dimension and then the complex pencil pass (fft_along_dim!, so threaded across pencils with #137) along each further one; backward does the complex passes in reverse order and then the complex-to-real pencil kernel. Argument checks report the expected transform sizes as before.Motivation: DSP.jl's
convfor float arrays of three or more dimensions plans an N-drfft/brffton its padded buffers (src/dspbase.jl), so this was the one hard blocker for running DSP.jl on FFTA.Tests:
test/ndim/real.jlchecks 3-D and 4-Drfft/irfft/brfftagainst the complex transform for several sizes and regions (including non-leading and non-contiguous regions),mul!,inv,\,Float32, and that results are identical with several workers; the argument-checking test that asserted the old error now asserts the size checks instead.