From ebd03f9441d19137d615ce4964cdf10bfe2b40a1 Mon Sep 17 00:00:00 2001 From: Oscar Smith Date: Fri, 13 Nov 2026 05:27:24 +0000 Subject: [PATCH] One looped kernel for the prime radices 17 to 73 A stage of a prime radix from 17 on runs prime_stage! (src/primestage.jl): one kernel for all such radices, with the radix and a table of cosines and sines as arguments, compiled once per vector type. It is the symmetric form of the odd-prime codelet as loops, with up to eight output pairs per pass over the inputs, and covers the stages vectorised over the lanes, the last stage, lane blocks and the grouped first stage. The radices up to 16 keep their straight-line codelets. A straight-line codelet of a prime r has about r^2/2 statements and needed 82 kernel methods per radix on AVX-512, which is why the codelets of 17-31 were compiled at first use. That machinery (callable tables, the lazy_radices preference) is removed: a prime now costs a table. MAX_RADIX_PRIME is 73, the largest prime for which the kernel was faster than Bluestein's algorithm in every measured case, the line of the prime alone included. The KernelAbstractions kernels and the generic number types keep their limit of 31 (factorize takes the limit as an argument). The stage kernels have no @generated twins any more: every combination the chain can select has a plain method. One table (STAGE_KERNELS) lists the kernels and their arguments; the methods and the dispatch switches are made from it. stage! and stage_range! share one body builder, and the scalar tail of the grouped first stage is the grouped kernel on a vector of one complex value. The generated bodies of the radices 2-16 are unchanged. The calls of the prime kernel are behind a function of their own, outside the dispatch switches: they take more arguments than fit in registers and cost every case of a switch about 1 ns otherwise. AMD EPYC 9354 (AVX-512), one thread, plan_fft and mul!, rounds in alternating processes, median time over that of the previous code: - lengths with a factor 17-31 (40 cases): geometric mean 0.82, one case above 1.05 (1.06); - long lengths with such a factor (36 cases): 0.94, one above 1.05 (1.08); - the primes 37-73 against Bluestein's algorithm (100 cases): 0.29; - 88 general cases without such factors: 1.00, three above 1.05 (Float64 lines of 32, 48, 60: 1.07-1.15; 1.04, 0.83, 1.05 in a second run of 6 rounds on short lines, whose stage kernels compile to the same code). Error against a 200-bit reference at 841, 6256, 6859 (Float64): 1.4-1.8 eps instead of 3.5-4.6 eps. libfftl.so 26.8 MB instead of 49.5 MB. The package image is 66 MB instead of 59 MB, because the kernel is in the image and the lazy codelets were not. Not verified on aarch64. Tests that used 37, 41 and 43 as examples of primes without a radix use 79, 83 and 89. Co-Authored-By: Claude Fable 5.2 --- AGENTS.md | 7 +- README.md | 11 +- docs/src/explanation/algorithms.md | 12 +- docs/src/explanation/design-decisions.md | 78 +++++++-- docs/src/explanation/image-size.md | 56 +++---- docs/src/explanation/implementation.md | 7 +- docs/src/explanation/performance.md | 2 +- docs/src/reference/platforms.md | 2 +- ext/FFTLKernelAbstractionsExt.jl | 18 +- ext/ka_real.jl | 14 +- libfftl/LocalPreferences.toml | 3 - src/FFTL.jl | 1 + src/factor.jl | 21 ++- src/generic.jl | 9 +- src/kernel_methods.jl | 61 ++----- src/kernels.jl | 203 ++++++----------------- src/plan.jl | 4 +- src/primestage.jl | 175 +++++++++++++++++++ src/rchain.jl | 178 ++++++++++---------- test/inference.jl | 10 +- test/radix_table.jl | 6 +- test/runtests.jl | 83 ++++----- 22 files changed, 532 insertions(+), 429 deletions(-) create mode 100644 src/primestage.jl diff --git a/AGENTS.md b/AGENTS.md index 68bff407..f1b16af0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -26,9 +26,10 @@ Read this file before changing the package. The documentation holds the rest: before it is written down, or it is marked as not verified. A change that claims to be faster comes with an interleaved A/B measurement on one machine (medians, and the spread of the same code against itself); see `docs/src/explanation/performance.md`. -- **Kernels of a fixed set are plain methods, not `@generated`.** A package image keeps the complete - optimised IR of every precompiled `@generated` specialisation. `@generated` methods stay as the - fallback for combinations outside the precompiled set. See `docs/src/explanation/image-size.md`. +- **Kernels are plain methods, not `@generated`.** A package image keeps the complete optimised IR + of every precompiled `@generated` specialisation. The stage kernels are defined with `@eval` for + the fixed set of radices, signs and vector types; there is no `@generated` fallback. See + `docs/src/explanation/image-size.md`. - **Nothing private in the repository.** Code, documentation, data files, commit messages and pull request descriptions contain no host names or accounts of machines, no local paths, no private contact data and no names of people beyond the authors of the package. A machine is named by its processor ("AMD EPYC 7702", diff --git a/README.md b/README.md index e451eac0..faee98a5 100644 --- a/README.md +++ b/README.md @@ -84,7 +84,7 @@ yd = FFTL.fft(MtlArray(randn(ComplexF32, 1024, 1024))) **Fourier and discrete cosine transforms written in Julia, with wide support for number types and CPU and GPU execution.** - Complex, real-input and all eleven real-to-real kinds of FFTW, along any subset of - dimensions, for any length (the part of a length made of prime factors above 31 through + dimensions, for any length (the part of a length made of prime factors above 73 through Bluestein's algorithm). - `Float32` and `Float64` on the fast path (SIMD codelets, threads). `BigFloat`, `Double64`, `Float128` and other floating-point types are transformed in their own precision by a generic @@ -117,19 +117,20 @@ figures there are regenerated from the committed data on every documentation bui ## How it works - **Generated codelets.** A small symbolic generator writes straight-line code for DFTs of - size 2–31 (composite sizes split in registers, odd primes in the symmetric form). Codelets are - generated with `@generated` for the radix, element type and SIMD width actually used. + size 2–16 (composite sizes split in registers, odd primes in the symmetric form), one method + per radix, element type and SIMD width. The prime radices 17–73 share one kernel that loops + over a table of cosines and sines. - **Complex numbers in SIMD registers.** On aarch64 CPUs with the ARMv8.3 complex-number extension (all Apple silicon), additions with a factor ±i and complex multiplications use the `FCADD` and `FCMLA` instructions; elsewhere the same operations are built from shuffles. - **Stockham chains.** A transform is a sequence of radix stages (mostly 4 and 8) without bit - reversal. The sequence is part of the plan's type, so it compiles into one function. The radices + reversal. The sequence is data of the plan; nothing is compiled per length. The radices are chosen by a fixed rule found in offline searches over the decompositions of many lengths; there is no run-time measuring step. - **Multidimensional transforms** run the same chains along every dimension: a dimension with stride S is a chain started with stride S, the S neighbouring columns forming the vectorised index. No transposition. -- **Lengths with prime factors above 31** split off their smooth part in one Cooley–Tukey step +- **Lengths with prime factors above 73** split off their smooth part in one Cooley–Tukey step and use Bluestein's algorithm only for the rest. - **Threads** split each dimension over slices or blocks of columns; long 1D transforms use the four-step algorithm. diff --git a/docs/src/explanation/algorithms.md b/docs/src/explanation/algorithms.md index ac848b16..2df098cc 100644 --- a/docs/src/explanation/algorithms.md +++ b/docs/src/explanation/algorithms.md @@ -18,13 +18,13 @@ which keeps the data in natural order throughout: The radices are chosen by a fixed rule (`src/factor.jl`). The rule is the result of exhaustive searches over the decompositions of many lengths on the development machines, not of a runtime search. It prefers radix 4 and radix 8 for powers of two. On aarch64 it merges a factor 3 with a 4 -into a 12; the x86 rules differ and are in the same file. Every other prime up to 31 is a radix of +into a 12; the x86 rules differ and are in the same file. Every other prime up to 73 is a radix of its own. Larger primes are not radices; the part of the length they make up goes to Bluestein's algorithm (below). ## Codelets -A stage of radix ``r`` is a small DFT applied to ``r`` values at once. For each ``r`` the codelet +A stage of radix ``r`` is a small DFT applied to ``r`` values at once. For each ``r`` up to 16 the codelet is straight-line code without loops, generated by a symbolic generator (`src/codegen.jl`). The generator writes the DFT as a sum of terms and simplifies it: @@ -38,9 +38,15 @@ The generated code runs on complex SIMD vectors (`src/complexvec.jl`), with the the element type. On aarch64 the complex add-and-rotate and the complex multiply are single instructions (`FCADD`, `FCMLA`). +The prime radices from 17 to 73 use the same symmetric form, but as loops: one kernel +(`src/primestage.jl`) takes the radix and a table of the cosines and sines as arguments. Straight-line +code for such a radix has about ``r^2/2`` statements with as many different constants, which no longer +fit in registers; the loop computes up to eight output pairs in one pass over the inputs, so that the +inputs are loaded once for all of them. + ## Bluestein's algorithm -Lengths with a prime factor above 31 go through Bluestein's identity. With ``jk = (j^2 + k^2 - (k-j)^2)/2``, +Lengths with a prime factor above 73 go through Bluestein's identity. With ``jk = (j^2 + k^2 - (k-j)^2)/2``, ```math X_k = c_k \sum_{j} (x_j c_j)\, \overline{c_{k-j}}, \qquad c_j = e^{\sigma \pi i j^2/n}. diff --git a/docs/src/explanation/design-decisions.md b/docs/src/explanation/design-decisions.md index 0babbe08..29ae61c4 100644 --- a/docs/src/explanation/design-decisions.md +++ b/docs/src/explanation/design-decisions.md @@ -58,7 +58,8 @@ showed 100 `Float32` and 64 `Float64` lengths slower than 1.5×, the worst at 3. ## Lengths with large prime factors: split off the smooth part, Bluestein on the rest -**Decision (2026-11-07).** For n = s·q, with q the part made of primes above 31: one Cooley–Tukey +**Decision (2026-11-07).** For n = s·q, with q the part made of primes above `MAX_RADIX_PRIME` (73, at the +time of the decision 31): one Cooley–Tukey step, a chain of length s, twiddles, and s Bluestein transforms of length q (`Composite`). A leftover q stays one Bluestein while its convolution is at most 4 MiB; above that, q is split into its prime powers, recursively. Bluestein convolutions of 8 MiB and more run as a serial four-step algorithm on @@ -115,8 +116,11 @@ for the rest), and before that the idea of removing `Chain` without changing the ## Kernels of a fixed set are plain methods, not `@generated` -**Decision (2026-11-10).** The stage kernels of the codelet set that `RChain` dispatches to are plain -methods defined in `src/kernel_methods.jl` from the bodies the `@generated` methods build. The `@generated` methods remain as the fallback. +**Decision (2026-11-10, fallback removed 2026-11-13).** The stage kernels of the codelet set that `RChain` +dispatches to are plain methods, defined with `@eval` in `src/kernel_methods.jl` from the functions that +build their bodies. One table (`STAGE_KERNELS` in `src/kernels.jl`) lists the kernels and their arguments; +the methods and the dispatch switches of `RChain` are made from it. There are no `@generated` stage +kernels: every combination the chain can select has a plain method, so a fallback would never run. **Why.** A package image keeps the complete optimised IR of every precompiled `@generated` specialisation, 130 MB of the 146 MB image. See [Package size and load time](@ref "Package size and load time"). @@ -149,16 +153,64 @@ radix 16 middle stage 284 → 198, radix 8 100 → 79. One line, `plan_fft`, med rounds: lengths 512–16384 take 0.79–0.94 (`Float32`) and 0.82–0.92 (`Float64`) of the time, lengths from 32768 on are unchanged (the same code against itself differs by 1–5 %). Not verified on aarch64. -## Codelets of the radices 17–31 are compiled on first use - -**Decision (2026-11-10).** The codelets of the radices 17, 19, 23, 29 and 31 are not part of the package -image. The first call compiles them (70–100 ms); later calls go through a table of callable objects at a -cost of about 0.02 µs per transform. - -**Why.** These radices were 54 % of the statements of the stage kernels and are rarely needed. The image -shrinks from 50 MB to 32 MB. Making 11 and 13 lazy as well would save 4 MB more; it was not done, because -lengths with those factors are common. The list is the preference `lazy_radices`; the trimmed C library -sets it to empty, because `juliac` compiles nothing at run time. +## The prime radices from 17 to 73 share one looped kernel + +**Decision (2026-11-13).** A stage of a prime radix from 17 on runs `prime_stage!` (`src/primestage.jl`): +one kernel for all such radices, with the radix and a table of cosines and sines as arguments, compiled +once per vector type. It uses the symmetric form of the odd-prime codelet as loops and computes up to +eight output pairs (four with 16 vector registers) in one pass over the inputs. It covers the stage +kinds that the codelets have separate kernels for: stages vectorised over the lanes, the last stage, +lane blocks, and the grouped first stage. `MAX_RADIX_PRIME` is 73. The radices up to 16 keep their +straight-line codelets. The KernelAbstractions kernels and the generic number types keep their limit +of 31. + +**Why.** A straight-line codelet of a prime r has about r²/2 statements with as many different +constants, which do not fit in registers (525 statements for 31), and every radix needed 82 kernel +methods on AVX-512 (kind × sign × vector type). The codelets of 17–31 were therefore compiled at first +use and not with the package (decision of 2026-11-10, replaced by this one). With the loop, a prime +costs a table. + +Measurements: AMD EPYC 9354 (AVX-512), one thread, `plan_fft` and `mul!`, rounds in alternating +processes, median time over the median time of the previous code. + +- Lengths and shapes with a factor 17–31 against the codelets (40 cases, 4 rounds): geometric mean + 0.82, 36 cases below 0.95, one above 1.05 (`Float64` 100 × 31, 1.06: a single stage in place, which + copies the inputs of a butterfly first). 17 and 31 alone: 0.63–0.84; 1088 = 17·64: 0.96 and 0.89; + 837 = 27·31 and 841 = 29²: 0.72–0.82; 62 × 62: 0.70–0.74. +- Long lengths (36 cases, 3 rounds; 31·2⁶ to 17·2²⁰, and a factor 17–31 times a prime above 1000): + geometric mean 0.94, one case above 1.05 (`Float64` 19·2¹⁷, 1.08). 31·2⁶ to 31·2¹⁴: 0.74–0.91; + 17·2²⁰: 0.99 and 1.01. +- The primes 37–73 against Bluestein's algorithm, which they used before (100 cases, 2 rounds; p, 8p, + 64p, 1024p, 16384p): geometric mean 0.29. The prime alone: 0.43–0.79 (`Float64`) and 0.48–0.90 + (`Float32`); as a factor of a longer length: 0.13–0.53. +- 92 general cases without such factors (the set of the decision above, 4 rounds): geometric mean 1.00 + without the two cases with a dimension of 37 (0.10 and 0.16), 3 cases above 1.05 (`Float64` lines of + 32, 48 and 60: 1.07–1.15). In 6 rounds on 20 short lines alone these three were at 1.04, 0.83 and + 1.05 and the geometric mean 1.00; the stage kernels of these lengths compile to the same machine code + as before. +- Accuracy against a 200-bit reference, `Float64`, 841, 6256 and 6859: 1.4–1.8 ε, with the codelets + 3.5–4.6 ε (the table is built from angles reduced to one turn). + +What made the loop as fast as the unrolled sums: several rows per pass, so that the inputs are loaded +once for all of them and the multiply-add chains of the rows run side by side; no buffer for the sums +and differences of the inputs (storing whole registers was the slowest part: 1.3–1.4× per butterfly +with the buffer); the pass over a block of rows out of line and the outputs of a block written by +recursion over the tuples of accumulators (a loop with an index into them, or a closure over them, put +the accumulators on the stack: up to 1.4× for a whole transform). + +**The limit.** With the limit at 127 for the measurement, a prime from 79 on is still 1.3–4× faster than +Bluestein's algorithm as a factor of a longer length, but the line of that prime alone is slower from +79 on for `Float32` (1.05, 1.5 at 127) and from 97 on for `Float64` (1.01, 1.15 at 127). 73 is the +largest prime for which every measured case is faster. + +**Not done.** The loop for 11 and 13: with one kernel per radix (an earlier form of the loop) it was at +0.98 of the codelets for `Float64` and 1.05–1.06 for `Float32`, and 1.08–1.10 for 7; not measured with +the single kernel. A larger limit for lengths that are not a prime alone. aarch64 is not verified. + +**Costs.** The kernel and its tables are in the package image, which the lazy codelets were not: the +image file is 66 MB instead of 59 MB on this machine (68.5 MB before the decision on the chain form). +The C library contains every kernel in either case and shrinks: `libfftl.so` (x86-64-v3, `--trim=safe`) +is 26.8 MB instead of 49.5 MB, its code 8.8 MB instead of 18.6 MB. ## Run-time pass plans for trimmed programs diff --git a/docs/src/explanation/image-size.md b/docs/src/explanation/image-size.md index 0e4e78b8..ee5ebca5 100644 --- a/docs/src/explanation/image-size.md +++ b/docs/src/explanation/image-size.md @@ -42,35 +42,32 @@ later Julia versions, so re-run the count below after an upgrade. ## What keeps it small -1. **Plain methods for the fixed set.** The stage kernels that `RChain` dispatches to (radix × sign × - vector type) are defined in `src/kernel_methods.jl` with `@eval`, from the same body-building - functions the generators call. The `@generated` methods stay as the fallback for every other combination - (the lazy radices), so behaviour is unchanged. A plain method keeps its (compressed) source - and its machine code, not its optimised IR. -2. **Lazy radices.** The codelets of the radices in `LAZY_RADICES` (17, 19, 23, 29, 31 by default; the - preference `lazy_radices`, read in `src/rchain.jl`) are 54 % of the statements of the stage kernels and are rarely needed. They are not - in the image. A call with such a codelet index reaches `cstage_*_lazy`, which finds a callable object - for the (radix, sign, vector type) in a table and creates it at the first call; Julia compiles its - kernel then. A first call of a length with such a factor takes 70–100 ms; later calls cost about - 0.02 µs more than a precompiled codelet. Calling the stage kernel with run-time `Val` arguments - instead of the table cost 0.22 µs per transform, 10 % at n = 1088. - - `juliac` compiles nothing at run time, so a program that is compiled with `--trim` sets the preference - `lazy_radices = []`: every codelet is then compiled into it. `libfftl/LocalPreferences.toml` does this - for the C library; a project that builds FFTL into its own trimmed library sets it in its - `LocalPreferences.toml` (the project must have FFTL as a direct dependency): - - ```toml - [FFTL] - lazy_radices = [] - ``` - -To add a radix, put it in `RADICES`; `kernel_methods.jl` defines its plain methods. To make a radix lazy, -add it to the default of the preference in `src/rchain.jl`. - -The numbers on this page were measured while the smooth lengths up to 128 had compiled chains of their -own (583 more plain methods). Without them the image on an AMD EPYC 9354 (AVX-512, which has more vector -types than the M1) went from 72 MB to 61 MB; the M1 numbers have not been measured again. +1. **Plain methods for the fixed set.** The stage kernels (`STAGE_KERNELS` in `src/kernels.jl`: radix × + sign × vector type) are defined in `src/kernel_methods.jl` with `@eval`, from the functions that build + their bodies. A plain method keeps its (compressed) source and its machine code, not its optimised IR. + There are no `@generated` stage kernels. +2. **One kernel for the large prime radices.** The radices 17 to 73 have no codelets of their own: one + kernel (`src/primestage.jl`) takes the radix and a table as arguments and is compiled once per vector + type. Straight-line codelets for 17, 19, 23, 29 and 31 were 54 % of the statements of the stage kernels; + they used to be left out of the image and compiled at the first call (70–100 ms, the "lazy radices" + column of the table above). + +To add a radix up to 16, put it in `RADICES`; `kernel_methods.jl` defines its methods. A larger prime +needs no code: raise `MAX_RADIX_PRIME` (see [Design decisions](@ref "Design decisions") for where the +kernel stops being faster than Bluestein's algorithm). + +The numbers on this page were measured on the M1 while the smooth lengths up to 128 had compiled chains +of their own (583 more plain methods) and the radices 17–31 were compiled at first use. On an AMD EPYC +9354 (AVX-512, which has seven vector types where the M1 has three) the image file went from 68.5 MB +with the compiled chains to 58.6 MB without them and to 66.2 MB with the kernel for the large primes +(which is in the image, where the lazy codelets were not). The C library, which contains every kernel +in either case, went from 55.6 MB to 49.5 MB and to 26.8 MB (`libfftl.so`, x86-64-v3, `--trim=safe`). The M1 numbers have not been measured again. + +With the workload, nearly all of the image is the stage kernels: the dispatch switches of `RChain` name +every kernel, so the first transform compiles all of them. On the EPYC 9354 (measured with the prime +limit at 127, which does not change the kernels) the image was 65 MB with the workload and 20.5 MB +without, precompilation took 61 s and 13 s, `using FFTL` 185 ms and 110 ms, and the first transform +16–37 ms with the workload and 23–35 s without. ## What is left @@ -88,7 +85,6 @@ Further steps that were considered: | step | gain | cost | |:--|:--|:--| | methods generic in the element and vector type instead of one per vector type | about 4 MB (stored sources ÷ 3) | the compiler must fold `ncomplex` and `sizeof` of the type parameters; needs a speed check | -| radices 11 and 13 lazy as well | 4 MB (28 MB), a few ms of load time (inside the noise) | 70–100 ms at the first call of lengths with those factors, which are common | | no PrecompileTools | about 25 ms of load time (Dates, TOML, Preferences) | the workload is written with `precompile` calls | | no `Float32` in the workload | 16 MB (at 73 MB) | 6 s at the first `Float32` call | diff --git a/docs/src/explanation/implementation.md b/docs/src/explanation/implementation.md index 65078e2b..40418634 100644 --- a/docs/src/explanation/implementation.md +++ b/docs/src/explanation/implementation.md @@ -68,6 +68,11 @@ is a few nanoseconds of arithmetic, so the loop is written to cost little next t - The switch over the codelet index is inlined into the loop: a stage is one jump and one call. - The rows of a plane run inside one call of the loop function. +A stage of a prime radix from 17 on calls one kernel for all such radices (`prime_stage!` in +`src/primestage.jl`), with the radix and its table as arguments; it is compiled once per vector type. +Its calls sit behind a function of their own, outside the switch, because they take more arguments +than fit in registers and would cost every case of the switch a larger stack frame. + The addresses in the stage records are valid for the chain object they were taken for. A chain remembers its own address at that moment; a copy of a plan, or a plan stored in a package image and loaded into another process, is at another address and links its records again before it uses them. @@ -76,7 +81,7 @@ loaded into another process, is at another address and links its records again b at build time and compiles nothing at run time; the fixed codelet set is all a chain of any length needs. The [libfftl how-to](@ref "C library (libfftl)") describes what the library is built from, and [Design decisions](@ref "Design decisions") gives the measurements behind the single chain form. The stage kernels of -the fixed codelet set are plain methods, not `@generated` ones, to keep the package image small; see +the fixed codelet set are plain methods defined with `@eval`, not `@generated` ones, to keep the package image small; see [Package size and load time](@ref "Package size and load time"). ## Lengths and the factorisation rule diff --git a/docs/src/explanation/performance.md b/docs/src/explanation/performance.md index 9b37eb9d..a23f4373 100644 --- a/docs/src/explanation/performance.md +++ b/docs/src/explanation/performance.md @@ -51,5 +51,5 @@ application they happen at different times. length likewise. The extra length shows in the measurements. - On the GPU, the real-to-real kinds are not tuned yet. - Bluestein transforms a line through two transforms of length ``M ≥ 2q - 1``, where ``q`` is - at most the part of the length made of prime factors above 31 (the smooth part is split off, + at most the part of the length made of prime factors above 73 (the smooth part is split off, see [Algorithms](@ref "Lengths with a smooth part")). A prime length pays this in full. diff --git a/docs/src/reference/platforms.md b/docs/src/reference/platforms.md index 25df9c92..55e25d9f 100644 --- a/docs/src/reference/platforms.md +++ b/docs/src/reference/platforms.md @@ -17,7 +17,7 @@ available. | DHT, `R2HC`, `HC2R` | ✓ | ✓ | ◐ not tuned | ✓ | | batched 1D passes (`plan_pass`, `plan_r2r_pass`) | ✓ | – | – | – | | any rank and any subset of dimensions | ✓ | ✓ | ✓ | ✓ | -| any length (Bluestein above prime factor 31) | ✓ | ✓ | ✓ | ✓ | +| any length (Bluestein above prime factor 73; 31 on the GPU and for generic floats) | ✓ | ✓ | ✓ | ✓ | ## Platforms diff --git a/ext/FFTLKernelAbstractionsExt.jl b/ext/FFTLKernelAbstractionsExt.jl index 48b63a1d..a26bd8cf 100644 --- a/ext/FFTLKernelAbstractionsExt.jl +++ b/ext/FFTLKernelAbstractionsExt.jl @@ -148,14 +148,18 @@ const PKERNELS_LOCK = ReentrantLock() lradix(n) = n >= 512 ? 8 : 16 const TWLEV = 3 # number of factor tables of W_N^t +# Largest prime factor with a codelet in a kernel; lengths with larger ones use Bluestein's +# algorithm. +const KA_MAX_PRIME = 31 + # Radices of a chain of length n: the power-of-two part as radix-R codelets and one remainder -# (first), then the odd part: pairs of 3 as 9, other primes (≤ FFTL.MAX_RADIX_PRIME) alone. +# (first), then the odd part: pairs of 3 as 9, other primes (≤ KA_MAX_PRIME) alone. function local_radices(n, R = lradix(n)) r = Int[]; m = n >> trailing_zeros(n); p2 = n ÷ m while p2 >= R; push!(r, R); p2 ÷= R; end p2 > 1 && push!(r, p2) - f, rest = FFTL.factorize(m) - @assert rest == 1 "prime factor > $(FFTL.MAX_RADIX_PRIME)" + f, rest = FFTL.factorize(m, KA_MAX_PRIME) + @assert rest == 1 "prime factor > $(KA_MAX_PRIME)" e3 = get(f, 3, 0) append!(r, fill(9, e3 ÷ 2)); isodd(e3) && push!(r, 3) for p in sort!(collect(keys(f))); p == 3 || append!(r, fill(p, f[p])); end @@ -770,7 +774,7 @@ passes_G(N, rs, lim = LIM0) = (s = 1; [(G = pass_G(R, s, N ÷ (R * s), 1, 1, 64, # N as few pass lengths ≤ cap as first-fit decreasing packing of the prime factors finds function pack_passes(N, cap) - f, rest = FFTL.factorize(N); @assert rest == 1 + f, rest = FFTL.factorize(N, KA_MAX_PRIME); @assert rest == 1 bins = Int[] for p in sort!(reduce(vcat, [fill(p, k) for (p, k) in f]; init = Int[]); rev = true) i = findfirst(b -> b * p <= cap, bins) @@ -782,7 +786,7 @@ end # N as a product of P pass lengths ≤ cap: prime factors, largest first, each to the pass with # the smallest product (more passes if a length exceeds cap) function balance_passes(N, P, cap = NMAX) - f, rest = FFTL.factorize(N); @assert rest == 1 + f, rest = FFTL.factorize(N, KA_MAX_PRIME); @assert rest == 1 ps = sort!(reduce(vcat, [fill(p, k) for (p, k) in f]; init = Int[]); rev = true) while true bins = ones(Int, P) @@ -792,7 +796,7 @@ function balance_passes(N, P, cap = NMAX) end end -# Bluestein along a dimension of length N (prime factor > FFTL.MAX_RADIX_PRIME): +# Bluestein along a dimension of length N (prime factor > KA_MAX_PRIME): # X_k = c_k Σ_j (x_j c_j) conj(c_{k-j}), c_j = exp(σ iπ j²/N), as a cyclic convolution of length # M ≥ 2N - 1. Forward passes of length M read x times c (load hook, zero padding), the first # backward pass reads times b̂ = FFT_M(conj(c) wrapped) / M, the last one writes the first N @@ -858,7 +862,7 @@ function KAPlan(backend, ::Type{T}, dims::NTuple{D,Int}, σ; region = 1:D, radic scr_next(src) = src === :c ? :d : :c for (i, d) in enumerate(region) N = dims[d]; E = prod(dims[1:d-1]; init = 1); B = prod(dims[d+1:end]; init = 1) - blue = FFTL.factorize(N)[2] > 1 + blue = FFTL.factorize(N, KA_MAX_PRIME)[2] > 1 M, cin, bhat = blue ? bluestein_tables(backend, T, N, σ) : (N, dummy, dummy) rs = radices === nothing ? dim_passes(M, E, E * M * B * sizeof(Complex{T}); padded = pad > 0, first = i == 1, last = i == length(region), prof, gmax = gmax, lim = lim) : radices[i] @assert prod(rs) == M diff --git a/ext/ka_real.jl b/ext/ka_real.jl index e002e309..65e2336d 100644 --- a/ext/ka_real.jl +++ b/ext/ka_real.jl @@ -320,7 +320,7 @@ function KARealPlan(backend, ::Type{T}, sz::NTuple{N,Int}, σ::Int; region = 1:N even = iseven(n) && n > 2 lsz = even ? setdim(sz, d1, n ÷ 2) : sz w = todevice(backend, [Complex{T}((σ < 0 ? complex(0, -0.5) : complex(0, 1)) * cispi(2σ * k / n)) for k in 0:h-1]) - fuse = even && d1 == 1 && FFTL.factorize(n ÷ 2)[2] == 1 + fuse = even && d1 == 1 && FFTL.factorize(n ÷ 2, KA_MAX_PRIME)[2] == 1 line = fuse ? KAPlan(backend, T, lsz, σ; region = d1:d1, rhook = σ, rtab = w) : KAPlan(backend, T, lsz, σ; region = d1:d1) rest = length(reg) > 1 ? KAPlan(backend, T, hsz, σ; region = reg[2:end]) : nothing bufs = Any[devzeros(backend, Complex{T}, lsz), devzeros(backend, Complex{T}, lsz), devzeros(backend, Complex{T}, hsz)] @@ -434,7 +434,7 @@ end function KAFourStep(backend, ::Type{T}, N::Int, P::Int) where {T} S = N ÷ P - @assert S * P == N && FFTL.factorize(S)[2] == 1 + @assert S * P == N && FFTL.factorize(S, KA_MAX_PRIME)[2] == 1 pP = KAPlan(backend, T, (S, P), -1; region = 2:2) pS = KAPlan(backend, T, (S, P), -1; region = 1:1) # W_N^m = thi[m ÷ W + 1] · tlo[m % W + 1] with W ≈ √N: two small tables, computed in double @@ -497,7 +497,7 @@ function KAR2RPlan(backend, ::Type{T}, sz::NTuple{N,Int}, kinds, region) where { for (k, d) in zip(ks, reg) n = sz[d] k == FFTL.REDFT00 && n < 2 && throw(ArgumentError("REDFT00 needs length ≥ 2")) - if k == FFTL.REDFT00 && n >= 3 && FFTL.factorize(n - 1)[2] == 1 + if k == FFTL.REDFT00 && n >= 3 && FFTL.factorize(n - 1, KA_MAX_PRIME)[2] == 1 # DCT-I: r2c of the even extension (length 2(n-1)) by one complex DFT of length n - 1, the # extension built at the first pass's loads, real parts at the last pass (rhook 8) n1 = n - 1 @@ -514,7 +514,7 @@ function KAR2RPlan(backend, ::Type{T}, sz::NTuple{N,Int}, kinds, region) where { # (with Bluestein the pack step runs in the first pass's chirp loads) n1 = n - 1; zsz = setdim(sz, d, n1) u = todevice(backend, [Complex{T}(complex(0, -0.5) * cispi(-j / n1)) for j in 0:n1]) - f1, P1 = FFTL.factorize(n1) + f1, P1 = FFTL.factorize(n1, KA_MAX_PRIME) if nameof(typeof(backend)) === :CUDABackend && length(sz) == 1 && P1 > 1 && n1 >= KA_FOURSTEP_MIN # one large prime factor: four-step DFT (KAFourStep) instead of Bluestein of length 2n push!(dims, (:dct1, false, KAFourStep(backend, T, n1, P1), u, @@ -534,7 +534,7 @@ function KAR2RPlan(backend, ::Type{T}, sz::NTuple{N,Int}, kinds, region) where { devzeros(backend, Complex{T}, zsz), devzeros(backend, Complex{T}, zsz))) continue end - if k in (FFTL.REDFT11, FFTL.RODFT11) && iseven(n) && n >= 4 && FFTL.factorize(n ÷ 2)[2] == 1 + if k in (FFTL.REDFT11, FFTL.RODFT11) && iseven(n) && n >= 4 && FFTL.factorize(n ÷ 2, KA_MAX_PRIME)[2] == 1 # DCT-IV/DST-IV of even length: complex backward DFT of n/2 between element-wise twiddles, # all in the DFT's passes (pass_kernel's rhook 6/7) M = n ÷ 2 @@ -551,7 +551,7 @@ function KAR2RPlan(backend, ::Type{T}, sz::NTuple{N,Int}, kinds, region) where { σ = k in (FFTL.REDFT10, FFTL.RODFT10) ? -1 : 1 w = todevice(backend, [Complex{T}(cispi(-2j / n)) for j in 0:M]) tw = todevice(backend, [Complex{T}(cispi(-j / (2n))) for j in 0:M]) - if σ > 0 && FFTL.factorize(M)[2] == 1 + if σ > 0 && FFTL.factorize(M, KA_MAX_PRIME)[2] == 1 # the whole DCT-III/DST-III in the line plan's passes (pass_kernel's rhook 4/5): # C_K = P_K (y_K - i y_(n-K)) + Q_K (y_(M-K) + i y_(M+K)) with a_k = conj(t_k)/2, # P_K = a_K (2 + 2i conj w_K), Q_K = conj(a_(M-K)) (2 - 2i conj w_K) @@ -564,7 +564,7 @@ function KAR2RPlan(backend, ::Type{T}, sz::NTuple{N,Int}, kinds, region) where { continue end end - if σ < 0 && FFTL.factorize(M)[2] == 1 + if σ < 0 && FFTL.factorize(M, KA_MAX_PRIME)[2] == 1 # the whole DCT-II/DST-II in the line plan's passes (pass_kernel's rhook 2/3) u = todevice(backend, [Complex{T}(complex(0, -0.5) * cispi(-2j / n)) for j in 0:M]) P = KAPlan(backend, T, hsz, -1; region = d:d, rhook = k == FFTL.REDFT10 ? 2 : 3, rtab = (u, tw)) diff --git a/libfftl/LocalPreferences.toml b/libfftl/LocalPreferences.toml index 26d73fa2..548a1ff6 100644 --- a/libfftl/LocalPreferences.toml +++ b/libfftl/LocalPreferences.toml @@ -2,9 +2,6 @@ # would only lengthen the build. [FFTL] precompile_workload = "none" -# juliac compiles nothing at run time, so every codelet is compiled with the package and none on -# first use. -lazy_radices = [] # 256-bit kernels only. The 512-bit complex addsub uses an AVX-512 intrinsic that the x86-64-v3 # target of the trimmed build cannot legalize (LLVM: "Do not know how to split the result"). simd_bytes = 32 diff --git a/src/FFTL.jl b/src/FFTL.jl index 7bdd2ca0..b10bec60 100644 --- a/src/FFTL.jl +++ b/src/FFTL.jl @@ -29,6 +29,7 @@ include("complexvec.jl") include("codegen.jl") include("kernels.jl") include("factor.jl") +include("primestage.jl") include("plan.jl") include("rchain.jl") include("kernel_methods.jl") diff --git a/src/factor.jl b/src/factor.jl index 869fd86f..e3c96a61 100644 --- a/src/factor.jl +++ b/src/factor.jl @@ -26,20 +26,27 @@ """ MAX_RADIX_PRIME -Largest prime factor handled by a codelet; lengths with larger prime factors use Bluestein's -algorithm. +Largest prime factor that is a radix of a chain; lengths with larger prime factors use Bluestein's +algorithm. The primes from `LOOPED_RADIX_MIN` on share one kernel (primestage.jl), so the limit +is where that kernel stops being faster than Bluestein's algorithm for a line of prime length. """ -const MAX_RADIX_PRIME = 31 +const MAX_RADIX_PRIME = 73 -function factorize(n::Int) +# The primes up to MAX_RADIX_PRIME. +const RADIX_PRIMES = Tuple(p for p in 2:MAX_RADIX_PRIME if all(p % d != 0 for d in 2:isqrt(p))) + +# The exponents of the prime factors up to pmax, and the remaining part of n (1 if n has no +# larger prime factor). +function factorize(n::Int, pmax::Int = MAX_RADIX_PRIME) f = Dict{Int,Int}() - for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31) + for p in RADIX_PRIMES + p > pmax && break while n % p == 0 f[p] = get(f, p, 0) + 1 n ÷= p end end - f, n # n > 1: remaining part with prime factors > MAX_RADIX_PRIME + f, n end smooth(n::Int) = factorize(n)[2] == 1 @@ -142,7 +149,7 @@ function heuristic_radices(n::Int, S::Int, esize::Int = 16) push!(r, 5) end end - for p in (7, 11, 13, 17, 19, 23, 29, 31), _ in 1:get(f, p, 0) + for p in RADIX_PRIMES, _ in 1:(p >= 7 ? get(f, p, 0) : 0) push!(r, p) end pow2 = vcat(fill(8, eight), fill(4, fours), fill(2, two)) diff --git a/src/generic.jl b/src/generic.jl index e31896ad..7ecf3a80 100644 --- a/src/generic.jl +++ b/src/generic.jl @@ -1,7 +1,7 @@ # Transforms for element types beyond the hardware floats (BigFloat, Double64, Float128, …): # plain Julia on arrays, no SIMD and no pointers (BigFloat is not an isbits type), single-threaded. # Accuracy over speed: twiddles are computed in T, smooth lengths go through a recursive -# mixed-radix FFT, lengths with a prime factor above MAX_RADIX_PRIME through Bluestein's +# mixed-radix FFT, lengths with a prime factor above GENERIC_MAX_PRIME through Bluestein's # algorithm. Real transforms reduce to complex DFTs: r2c/c2r through the full complex line, the # real-to-real kinds through a complex DFT of length 2n (2(n-1) for DCT-I, 2(n+1) for DST-I) # with pre- and post-twiddles: @@ -36,11 +36,14 @@ function generic_roots(::Type{T}, n::Int, σ::Int) where {T} w end +# Largest prime factor that the mixed-radix transform takes as a radix. +const GENERIC_MAX_PRIME = 31 + """ GLine{T} One-dimensional DFT of length `n`, sign `σ`, for a generic float type: mixed radix, or Bluestein -when a prime factor exceeds MAX_RADIX_PRIME. +when a prime factor exceeds GENERIC_MAX_PRIME. """ struct GLine{T} n::Int @@ -57,7 +60,7 @@ end function GLine{T}(n::Int, σ::Int) where {T} n >= 1 || throw(ArgumentError("length must be positive")) - f, rest = factorize(n) + f, rest = factorize(n, GENERIC_MAX_PRIME) if rest == 1 rs = Int[] for p in reverse!(sort!(collect(keys(f)))), _ in 1:f[p] diff --git a/src/kernel_methods.jl b/src/kernel_methods.jl index 5a1fe9d0..b00228ea 100644 --- a/src/kernel_methods.jl +++ b/src/kernel_methods.jl @@ -1,50 +1,15 @@ -# Compiled-ahead methods of the stage kernels for the fixed codelet set (radix × sign × vector -# type) that `RChain` dispatches to. The -# `@generated` kernels in kernels.jl define every other combination. A package image keeps the -# complete optimised IR of each precompiled `@generated` specialization but none of a plain -# method, so the fixed set is defined as plain methods with the same bodies. The codelets of -# LAZY_RADICES are not in the set. - -kernel_subst(x, d) = x isa Symbol ? get(d, x, x) : - x isa Expr ? Expr(x.head, map(a -> kernel_subst(a, d), x.args)...) : x - -const KERNEL_HEADS = ( - (:stage!, :(stage!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT})), :all), - (:stage_range!, :(stage_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, qlo::Int, qhi::Int, - plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :all), - (:stage_last!, :(stage_last!(py::Ptr{T}, px::Ptr{T}, s::Int, qlo::Int, qhi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT})), :all), - (:stage_lanes!, :(stage_lanes!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s0::Int, Sx::Int, Sy::Int, - B::Int, plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :all), - (:stage_last_lanes!, :(stage_last_lanes!(py::Ptr{T}, px::Ptr{T}, s0::Int, Sx::Int, Sy::Int, B::Int, - qlo::Int, qhi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :all), - # the grouped first stage exists for 2 or 4 complex values per vector only; its scalar tail - # runs on one complex value - (:stage_first_group_range!, :(stage_first_group_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, - plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :group), - (:stage_first_group_1!, :(stage_first_group_1!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, - plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :tail)) - -for T in (Float32, Float64), (name, head, vts) in KERNEL_HEADS - body = getfield(@__MODULE__, Symbol(string(name)[1:end-1], "_body")) - vts === :all ? (types = vec_ladder(T)) : - vts === :group ? (types = filter(V -> ncomplex(V) in (2, 4), vec_ladder(T))) : - (types = any(V -> ncomplex(V) in (2, 4), vec_ladder(T)) ? (Vec{2,T},) : ()) - for σ in (-1, 1), VT in types, R in RADICES - R in LAZY_RADICES && continue - h = kernel_subst(head, Dict{Symbol,Any}(:T => T, :R => R, :σ => σ, :VT => VT)) - # the codelet statements name the type parameters T and VT, which a plain method does not bind - b = kernel_subst(body(T, R, σ, VT), Dict{Symbol,Any}(:T => T, :VT => VT)) - @eval $(Expr(:function, h, b)) +# The methods of the stage kernels: one per kernel of STAGE_KERNELS (kernels.jl), element type, +# sign, vector type of the ladder and radix of RADICES, the set that `RChain` dispatches to. They +# are plain methods: a package image keeps the complete optimised IR of each precompiled +# specialization of a `@generated` method but none of a plain one. +for T in (Float32, Float64), (name, args, body) in STAGE_KERNELS + ladder = vec_ladder(T) + typed = [a === :pt ? :($a::Ptr{$T}) : :($a::Int) for a in args] + for σ in (-1, 1), (i, V) in enumerate(ladder), R in RADICES + kernel_exists(name, ladder, i, R) || continue + @eval function $name(py::Ptr{$T}, px::Ptr{$T}, $(typed...), ::Val{$R}, ::Val{$σ}, ::Type{$V}) + VT = $V # the codelet statements name the vector type + $(body(T, R, σ, V)) + end end end - -# stage_small!: the widest vector type, power-of-two radices. -for T in (Float32, Float64), σ in (-1, 1), R in RADICES - VT = first(vec_ladder(T)) - (ispow2(R) && !(R in LAZY_RADICES)) || continue - h = :(stage_small!(py::Ptr{$T}, px::Ptr{$T}, pt::Ptr{$T}, m::Int, ::Val{$R}, ::Val{$σ}, ::Type{$VT})) - b = kernel_subst(stage_small_body(T, R, σ, VT), Dict{Symbol,Any}(:T => T, :VT => VT)) - @eval $(Expr(:function, h, b)) -end diff --git a/src/kernels.jl b/src/kernels.jl index ef34bede..4cbd876a 100644 --- a/src/kernels.jl +++ b/src/kernels.jl @@ -5,47 +5,42 @@ # (contiguous in p). # Kernels take raw pointers (to the real part of the first complex value); callers keep the # arrays alive. Indices are zero-based offsets in units of complex values. +# This file builds the bodies of the stage kernels; kernel_methods.jl defines one method per +# (radix, sign, vector type) from them, with the arguments listed in STAGE_KERNELS. # # A transform along a dimension with stride S (batch of S interleaved transforms) is the same # chain started with s = S: the S lanes become part of the contiguous index q. @inline cptr(x::Array{Complex{T}}) where {T} = Ptr{T}(pointer(x)) -# General stage, vectorised over q: s must be a multiple of the complex lanes of VT. -stage!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage!(py, px, pt, m, s, 0, m, Val(R), Val(σ), VT) - -function stage_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} +# The statements of one butterfly of radix R in a stage body: the loads of the inputs from the +# addresses `ld(j)`, j = 0…R-1, the DFT, the statements `pre`, and the stores `st(X_k, k)`. +function butterfly(::Type{T}, R::Int, σ::Int, ::Type{VT}, ld, st, pre = ()) where {T,VT} stmts, a, X = codelet(R, σ) + vcat([:($(a[j]) = vload($VT, $(ld(j - 1)))) for j in 1:R], stmts, collect(Any, pre), [st(X[k+1], k) for k in 0:R-1]) +end + +# General stage, vectorised over q: s must be a multiple of the complex lanes of VT. +# The body of stage! (q over 0:s-1) and of stage_range! (q over qlo:qhi-1). +function stage_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, qlo = 0, qhi = :s) where {T,VT} C = ncomplex(VT); B = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, $(stream_addr(:xb, :xr, j - 1)))) for j in 1:R] + st(X, k) = k == 0 ? :(vstore($X, yb0)) : :(vstore(ccmul_at($X, $(Symbol("w", k))), $(stream_addr(:yb, :yr, k)))) twl = [:($(Symbol("w", k)) = pt + ($(R - 1) * p + $(k - 1)) * $(2sizeof(T))) for k in 1:R-1] - stores = Any[:(vstore($(X[1]), yb0))] - for k in 1:R-1 - push!(stores, :(vstore(ccmul_at($(X[k+1]), $(Symbol("w", k))), $(stream_addr(:yb, :yr, k))))) - end quote $(stream_strides(:xr, R, :(s * m * $B))...) $(stream_strides(:yr, R, :(s * $B))...) for p in plo:phi-1 $(twl...) - for q in 0:$C:s-1 + for q in $qlo:$C:$qhi-1 $(stream_bases(:xb, :xr, R, :(px + (q + s * p) * $B))...) - $(loads...) - $(stmts...) - $(stream_bases(:yb, :yr, R, :(py + (q + s * $R * p) * $B))...) - $(stores...) + $(butterfly(T, R, σ, VT, j -> stream_addr(:xb, :xr, j), st, + stream_bases(:yb, :yr, R, :(py + (q + s * $R * p) * $B)))...) end end nothing end end -@generated function stage!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} - stage_body(T, R, σ, VT) -end - # base + i·2^e as one scaled-index address computation (a GEP over an element of 2^e bytes), so # that the scaling stays in the addressing mode instead of becoming a value of its own. for (e, ty) in enumerate(("i8", "i16", "i32", "i64")) @@ -75,87 +70,39 @@ stream_strides(r, R, unit) = # Stage over a q range: a stride that is not a multiple of the complex lanes of VT runs the # first s - s % C lanes in vectors of C complex values and the remainder one complex value at a # time. -function stage_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} - stmts, a, X = codelet(R, σ) - C = ncomplex(VT); B = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, $(stream_addr(:xb, :xr, j - 1)))) for j in 1:R] - twl = [:($(Symbol("w", k)) = pt + ($(R - 1) * p + $(k - 1)) * $(2sizeof(T))) for k in 1:R-1] - stores = Any[:(vstore($(X[1]), yb0))] - for k in 1:R-1 - push!(stores, :(vstore(ccmul_at($(X[k+1]), $(Symbol("w", k))), $(stream_addr(:yb, :yr, k))))) - end - quote - $(stream_strides(:xr, R, :(s * m * $B))...) - $(stream_strides(:yr, R, :(s * $B))...) - for p in plo:phi-1 - $(twl...) - for q in qlo:$C:qhi-1 - $(stream_bases(:xb, :xr, R, :(px + (q + s * p) * $B))...) - $(loads...) - $(stmts...) - $(stream_bases(:yb, :yr, R, :(py + (q + s * $R * p) * $B))...) - $(stores...) - end - end - nothing - end -end - -@generated function stage_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, qlo::Int, qhi::Int, - plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} - stage_range_body(T, R, σ, VT) -end +stage_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} = stage_body(T, R, σ, VT, :qlo, :qhi) # Last stage (m = 1): all twiddles are 1. -stage_last!(py::Ptr{T}, px::Ptr{T}, s::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage_last!(py, px, s, 0, s, Val(R), Val(σ), VT) - function stage_last_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} - stmts, a, X = codelet(R, σ) C = ncomplex(VT); B = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, $(stream_addr(:xb, :r, j - 1)))) for j in 1:R] - stores = [:(vstore($(X[k+1]), $(stream_addr(:yb, :r, k)))) for k in 0:R-1] quote $(stream_strides(:r, R, :(s * $B))...) for q in qlo:$C:qhi-1 $(stream_bases(:xb, :r, R, :(px + q * $B))...) - $(loads...) - $(stmts...) - $(stream_bases(:yb, :r, R, :(py + q * $B))...) - $(stores...) + $(butterfly(T, R, σ, VT, j -> stream_addr(:xb, :r, j), (X, k) -> :(vstore($X, $(stream_addr(:yb, :r, k)))), + stream_bases(:yb, :r, R, :(py + q * $B)))...) end nothing end end -@generated function stage_last!(py::Ptr{T}, px::Ptr{T}, s::Int, qlo::Int, qhi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} - stage_last_body(T, R, σ, VT) -end - -# First stage with s = 1 for vectors of C = 2 or 4 complex values (m a multiple of C): +# First stage with s = 1 for vectors of C = 2 or 4 complex values: # butterflies p, …, p+C-1 share a register. Groups of C outputs X_k…X_(k+C-1) are transposed # (complex value l of each output into one vector) so that each store writes C consecutive # outputs of one butterfly; remaining outputs (R not a multiple of C) are stored one complex -# value at a time. Butterflies beyond a multiple of C (m % C) run one complex value at a time. -stage_first_group!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage_first_group!(py, px, pt, m, 0, m, Val(R), Val(σ), VT) - -@inline function stage_first_group!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, plo::Int, phi::Int, +# value at a time. Butterflies beyond a multiple of C (m % C) run one complex value at a time, with +# the same kernel on a vector of one complex value. +@inline function stage_first_group!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} - C = ncomplex(VT) - rem = (phi - plo) % C - if rem == 0 - stage_first_group_range!(py, px, pt, m, plo, phi, Val(R), Val(σ), VT) - else - stage_first_group_range!(py, px, pt, m, plo, phi - rem, Val(R), Val(σ), VT) - stage_first_group_1!(py, px, pt, m, phi - rem, phi, Val(R), Val(σ), Vec{2,T}) - end + rem = m % ncomplex(VT) + stage_first_group_range!(py, px, pt, m, 0, m - rem, Val(R), Val(σ), VT) + rem == 0 || stage_first_group_range!(py, px, pt, m, m - rem, m, Val(R), Val(σ), Vec{2,T}) + nothing end function stage_first_group_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} C = ncomplex(VT) - C in (2, 4) || error("the grouped first stage needs 2 or 4 complex values per vector") + C in (1, 2, 4) || error("the grouped first stage needs 1, 2 or 4 complex values per vector") stmts, a, X = codelet(R, σ) N = 2C; CB = 2sizeof(T) loads = [:($(a[j]) = vload($VT, px + (p + $(j - 1) * m) * $CB)) for j in 1:R] @@ -172,7 +119,7 @@ function stage_first_group_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) wh k = 0 while k + C <= R for l in 0:C-1 - z = C == 2 ? pick(outs[k+1], outs[k+2], l) : + z = C == 1 ? outs[k+1] : C == 2 ? pick(outs[k+1], outs[k+2], l) : :(shufflevector($(pick(outs[k+1], outs[k+2], l)), $(pick(outs[k+3], outs[k+4], l)), Val((0, 1, 2, 3, 4, 5, 6, 7)))) push!(st, :(vstore($z, py + ($R * (p + $l) + $k) * $CB))) @@ -196,11 +143,6 @@ function stage_first_group_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) wh end end -@generated function stage_first_group_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} - stage_first_group_range_body(T, R, σ, VT) -end - # Complex values of u ++ v (2C) in chunks of S, alternating u_0 v_0 u_1 v_1 …: the lane masks of # the first C and the last C. function chunk_masks(C, S) @@ -269,68 +211,23 @@ function stage_small_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} end end -@generated function stage_small!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, ::Val{R}, ::Val{σ}, - ::Type{VT}) where {T,R,σ,VT} - stage_small_body(T, R, σ, VT) -end - -# Scalar tail of the grouped first stage: one complex value per vector, the twiddle table in -# the grouped layout (k-1)m + p, one store per output. -function stage_first_group_1_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} - ncomplex(VT) == 1 || error("the scalar tail of the grouped first stage needs one complex value per vector") - stmts, a, X = codelet(R, σ) - CB = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, px + (p + $(j - 1) * m) * $CB)) for j in 1:R] - twm = Any[] - stores = Any[:(vstore($(X[1]), py + $R * p * $CB))] - for k in 1:R-1 - b = Symbol("b", k) - push!(twm, :($b = ccmul($(X[k+1]), vload($VT, pt + ($(k - 1) * m + p) * $CB)))) - push!(stores, :(vstore($b, py + ($R * p + $k) * $CB))) - end - quote - for p in plo:phi-1 - $(loads...) - $(stmts...) - $(twm...) - $(stores...) - end - nothing - end -end - -@generated function stage_first_group_1!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} - stage_first_group_1_body(T, R, σ, VT) -end - # Stages on a block of B lanes (cache blocking of a dimension pass, and its unit of work among # threads): element j of lane i # is at i + Sx·j in the input and at i + Sy·j in the output, so a block can be read from the array # (Sx = S) into a compact buffer (Sy = B) and back. The index q of the formulas above is # i + S·q0 with i = 0…B-1, q0 = 0…s0-1. B must be a multiple of the complex lanes of VT. -stage_lanes!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s0::Int, Sx::Int, Sy::Int, B::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage_lanes!(py, px, pt, m, s0, Sx, Sy, B, 0, m, Val(R), Val(σ), VT) - function stage_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} - stmts, a, X = codelet(R, σ) C = ncomplex(VT); Bt = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, px + (i + Sx * (q0 + s0 * (p + $(j - 1) * m))) * $Bt)) for j in 1:R] twl = [:($(Symbol("w", k)) = cbroadcast($VT, pt + ($(R - 1) * p + $(k - 1)) * $(2sizeof(T)))) for k in 1:R-1] - stores = Any[:(vstore($(X[1]), py + (i + Sy * o) * $Bt))] - for k in 1:R-1 - push!(stores, :(vstore(ccmul($(X[k+1]), $(Symbol("w", k))), py + (i + Sy * (o + $k * s0)) * $Bt))) - end + st(X, k) = k == 0 ? :(vstore($X, py + (i + Sy * o) * $Bt)) : + :(vstore(ccmul($X, $(Symbol("w", k))), py + (i + Sy * (o + $k * s0)) * $Bt)) quote for p in plo:phi-1 $(twl...) for q0 in 0:s0-1 o = q0 + s0 * $R * p for i in 0:$C:B-1 - $(loads...) - $(stmts...) - $(stores...) + $(butterfly(T, R, σ, VT, j -> :(px + (i + Sx * (q0 + s0 * (p + $j * m))) * $Bt), st)...) end end end @@ -338,33 +235,33 @@ function stage_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} end end -@generated function stage_lanes!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s0::Int, Sx::Int, Sy::Int, - B::Int, plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} - stage_lanes_body(T, R, σ, VT) -end - -stage_last_lanes!(py::Ptr{T}, px::Ptr{T}, s0::Int, Sx::Int, Sy::Int, B::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage_last_lanes!(py, px, s0, Sx, Sy, B, 0, s0, Val(R), Val(σ), VT) - function stage_last_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} - stmts, a, X = codelet(R, σ) C = ncomplex(VT); Bt = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, px + (i + Sx * (q0 + s0 * $(j - 1))) * $Bt)) for j in 1:R] - stores = [:(vstore($(X[k+1]), py + (i + Sy * (q0 + s0 * $k)) * $Bt)) for k in 0:R-1] quote for q0 in qlo:qhi-1 for i in 0:$C:B-1 - $(loads...) - $(stmts...) - $(stores...) + $(butterfly(T, R, σ, VT, j -> :(px + (i + Sx * (q0 + s0 * $j)) * $Bt), + (X, k) -> :(vstore($X, py + (i + Sy * (q0 + s0 * $k)) * $Bt)))...) end end nothing end end -@generated function stage_last_lanes!(py::Ptr{T}, px::Ptr{T}, s0::Int, Sx::Int, Sy::Int, B::Int, qlo::Int, qhi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} - stage_last_lanes_body(T, R, σ, VT) -end +# The stage kernels: name, the arguments between (py, px) and (Val(R), Val(σ), VT), and the body. +const STAGE_KERNELS = ( + (:stage!, (:pt, :m, :s, :plo, :phi), stage_body), + (:stage_range!, (:pt, :m, :s, :qlo, :qhi, :plo, :phi), stage_range_body), + (:stage_last!, (:s, :qlo, :qhi), stage_last_body), + (:stage_first_group_range!, (:pt, :m, :plo, :phi), stage_first_group_range_body), + (:stage_small!, (:pt, :m), stage_small_body), + (:stage_lanes!, (:pt, :m, :s0, :Sx, :Sy, :B, :plo, :phi), stage_lanes_body), + (:stage_last_lanes!, (:s0, :Sx, :Sy, :B, :qlo, :qhi), stage_last_lanes_body)) + +# Whether a kernel exists for vector type i of the ladder and radix R: the looped radices have +# one kernel of their own (primestage.jl), the grouped first stage takes at most 4 complex values +# per vector, stage_small! the widest vector type and a power-of-two radix. +kernel_exists(name::Symbol, ladder, i::Int, R::Int) = + looped(R) ? false : + name === :stage_first_group_range! ? ncomplex(ladder[i]) <= 4 : + name === :stage_small! ? i == 1 && ispow2(R) : true diff --git a/src/plan.jl b/src/plan.jl index eaa09b85..5bc4de76 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -360,8 +360,8 @@ function composite_split(::Type{T}, n::Int) where {T} _, r = factorize(n) r == 1 && return (n, 1) bluestein_length(T, 2r - 1) * sizeof(Complex{T}) <= COMPOSITE_SPLIT_BYTES && return (n ÷ r, r) - p = r; d = 37 - while d * d <= r # largest prime factor of r (all its factors are ≥ 37) + p = r; d = MAX_RADIX_PRIME + 2 + while d * d <= r # largest prime factor of r (all its factors are above MAX_RADIX_PRIME) while r % d == 0 p = d; r ÷= d end diff --git a/src/primestage.jl b/src/primestage.jl new file mode 100644 index 00000000..fa0a7045 --- /dev/null +++ b/src/primestage.jl @@ -0,0 +1,175 @@ +# One stage kernel for every prime radix from LOOPED_RADIX_MIN on. The radix and its table of +# cosines and sines are arguments, so a further prime costs a table and no code. With h = (R-1)/2, +# S_j = a_j + a_(R-j), D_j = a_j - a_(R-j), +# X_0 = a_0 + Σ_j S_j, +# X_k, X_(R-k) = A_k ∓ i·B_k, A_k = a_0 + Σ_j S_j cos(2π jk/R), B_k = -σ Σ_j D_j sin(2π jk/R), +# the sums of the straight-line codelet of an odd prime (codegen.jl) as loops. + +const LOOPED_RADIX_MIN = 17 +# RADICES holds no composite radix from LOOPED_RADIX_MIN on. +looped(R::Int) = R >= LOOPED_RADIX_MIN + +# Rows of a butterfly that one pass over its inputs computes: every row takes two registers. +const PRIME_BLOCK = (REG_BYTES == 64 || Sys.ARCH === :aarch64) ? 8 : 4 + +# The table of radix R and sign σ: cos(2π jk/R) at (j-1)h + k for j, k = 1…h, and -σ·sin(2π jk/R) +# h² further on, so that the rows k of one input pair j are contiguous. +function build_prime_table(::Type{T}, R::Int, σ::Int) where {T} + h = (R - 1) ÷ 2 + vcat([T(cospi(2 * mod(j * k, R) / R)) for j in 1:h for k in 1:h], + [T(-σ * sinpi(2 * mod(j * k, R) / R)) for j in 1:h for k in 1:h]) +end + +for T in (Float32, Float64) + tables = Symbol(:PRIME_TABLES_, T) + @eval const $tables = Dict{Tuple{Int,Int},Vector{$T}}((R, σ) => build_prime_table($T, R, σ) + for R in RADIX_PRIMES if looped(R) for σ in (-1, 1)) + @eval prime_table(::Type{$T}, R::Int, σ::Int) = $tables[(R, σ)] +end + +# How a stage treats its outputs: no twiddles (a last stage), one twiddle per output k for all +# lanes, or the grouped first stage (kernels.jl): the lanes of a vector are the butterflies +# p, p+1, …, each with its own twiddle, and output k of lane l goes to R·l + k. +const PRIME_LAST = 0 +const PRIME_MID = 1 +const PRIME_GROUP = 2 + +@inline function prime_store!(X::VT, yb::Ptr{T}, ys::Int, k::Int, R::Int, mode::Int) where {T,VT} + if mode == PRIME_GROUP + for l in 1:ncomplex(VT) + vstore(Vec{2,T}((X[2l-1], X[2l])), yb + (R * (l - 1) + k) * 2sizeof(T)) + end + else + vstore(X, yb + k * ys) + end + nothing +end + +# The outputs k and R - k from A_k and B_k. tw is the address of the twiddle of output 1, tws the +# distance to the next one in the grouped table. +@inline function prime_out!(yb::Ptr{T}, ys::Int, tw::Ptr{T}, tws::Int, mode::Int, c::VT, s::VT, k::Int, R::Int) where {T,VT} + Bt = 2sizeof(T) + Xk = fcadd270(c, s); Xr = fcadd90(c, s) + if mode == PRIME_MID + Xk = ccmul_at(Xk, tw + (k - 1) * Bt); Xr = ccmul_at(Xr, tw + (R - k - 1) * Bt) + elseif mode == PRIME_GROUP + Xk = ccmul(Xk, vload(VT, tw + (k - 1) * tws)); Xr = ccmul(Xr, vload(VT, tw + (R - k - 1) * tws)) + end + prime_store!(Xk, yb, ys, k, R, mode); prime_store!(Xr, yb, ys, R - k, R, mode) + nothing +end +# The outputs of the rows k, k+1, … of a block, one per element of the tuples. The recursion over +# the tuples keeps the accumulators in registers: a loop with an index into them, and a closure +# over them, were both measured slower (up to 1.4× for a whole transform). +@inline prime_outs!(yb, ys, tw, tws, mode, ::Tuple{}, ::Tuple{}, k, R) = nothing +@inline function prime_outs!(yb, ys, tw, tws, mode, cs::Tuple, ss::Tuple, k, R) + prime_out!(yb, ys, tw, tws, mode, first(cs), first(ss), k, R) + prime_outs!(yb, ys, tw, tws, mode, Base.tail(cs), Base.tail(ss), k + 1, R) +end + +# acc_b + v·t_b for the NB table values at p +@inline fma_rows(v::VT, p::Ptr{T}, acc::NTuple{NB,VT}) where {T,VT,NB} = + ntuple(b -> muladd(v, unsafe_load(p, b), acc[b]), Val(NB)) + +# Rows k … k+NB-1 of one butterfly (and row 0 with the first block) in one pass over the inputs +# at xb + j·xs. The sums and differences are formed again in every pass: keeping them in a buffer +# costs more than the loads (the stores of whole registers are the slowest part). Not inlined: a +# pass is long enough for a call, and an inlined function keeps its typed code in the package +# image for every block size and vector type. +@noinline function prime_rows!(yb::Ptr{T}, ys::Int, xb::Ptr{T}, xs::Int, tab::Ptr{T}, tw::Ptr{T}, tws::Int, mode::Int, + ::Type{VT}, R::Int, h::Int, k::Int, ::Val{NB}) where {T,VT,NB} + a0 = vload(VT, xb); acc = a0 + cs = ntuple(_ -> a0, Val(NB)); ss = ntuple(_ -> zero(VT), Val(NB)) + psin = tab + h * h * sizeof(T) + for j in 1:h + x = vload(VT, xb + j * xs); y = vload(VT, xb + (R - j) * xs) + sj = x + y + acc += sj + o = ((j - 1) * h + k - 1) * sizeof(T) + cs = fma_rows(sj, tab + o, cs) + ss = fma_rows(x - y, psin + o, ss) + end + k == 1 && prime_store!(acc, yb, ys, 0, R, mode) + prime_outs!(yb, ys, tw, tws, mode, cs, ss, k, R) +end + +# One butterfly in blocks of at most PRIME_BLOCK rows of nearly equal size (q or q + 1, the +# first r blocks the larger ones). +@eval @inline function prime_butterfly!(yb, ys, xb, xs, tab, tw, tws, mode, ::Type{VT}, R, h, nblocks, q, r) where {VT} + k = 1 + for b in 1:nblocks + nb = q + (b <= r) + Base.Cartesian.@nif $PRIME_BLOCK d -> (nb == d) d -> prime_rows!(yb, ys, xb, xs, tab, tw, tws, mode, VT, R, h, k, Val(d)) + k += nb + end + nothing +end + +""" + prime_stage!(py, px, pt, tab, R, m, s0, Sx, Sy, i0, i1, VT) + +A Stockham stage of the prime radix `R` with the table `tab` (`prime_table`), in the form of the +lane stages (kernels.jl): input j of butterfly p at `i + Sx·(q0 + s0·(p + j·m))`, output k at +`i + Sy·(q0 + s0·(R·p + k))`, for q0 < s0 and the lanes i0 ≤ i < i1 in vectors of `VT`. A stage +vectorised over q is the case s0 = 1, Sx = Sy = s; a last stage has m = 1 and no twiddles +(`pt` null). +""" +@noinline function prime_stage!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, tab::Ptr{T}, R::Int, m::Int, s0::Int, Sx::Int, Sy::Int, + i0::Int, i1::Int, ::Type{VT}) where {T,VT} + C = ncomplex(VT); Bt = 2sizeof(T); h = (R - 1) >> 1 + nblocks = cld(h, PRIME_BLOCK); q, r = divrem(h, nblocks) + # In place (a single stage) the inputs of a butterfly must be read before its first store: one + # block does that, several need a copy. + py == px && nblocks > 1 && return prime_stage_inplace!(py, tab, R, s0, Sx, i0, i1, VT) + xs = Sx * s0 * m * Bt; ys = Sy * s0 * Bt + mode = pt == C_NULL ? PRIME_LAST : PRIME_MID + for p in 0:m-1 + tw = pt + (R - 1) * p * Bt + for q0 in 0:s0-1, i in i0:C:i1-1 + xb = px + (i + Sx * (q0 + s0 * p)) * Bt + yb = py + (i + Sy * (q0 + s0 * R * p)) * Bt + prime_butterfly!(yb, ys, xb, xs, tab, tw, 0, mode, VT, R, h, nblocks, q, r) + end + end + nothing +end + +# A single stage in place (m = 1, no twiddles): each butterfly from a copy of its inputs. +@noinline function prime_stage_inplace!(py::Ptr{T}, tab::Ptr{T}, R::Int, s0::Int, S::Int, i0::Int, i1::Int, + ::Type{VT}) where {T,VT} + C = ncomplex(VT); Bt = 2sizeof(T); h = (R - 1) >> 1; W = sizeof(VT) + ys = S * s0 * Bt + nblocks = cld(h, PRIME_BLOCK); q, r = divrem(h, nblocks) + buf = Ref{NTuple{MAX_RADIX_PRIME,VT}}() + GC.@preserve buf begin + pbuf = Ptr{T}(Base.unsafe_convert(Ptr{NTuple{MAX_RADIX_PRIME,VT}}, buf)) + for q0 in 0:s0-1, i in i0:C:i1-1 + yb = py + (i + S * q0) * Bt + for j in 0:R-1 + vstore(vload(VT, yb + j * ys), pbuf + j * W) + end + prime_butterfly!(yb, ys, pbuf, W, tab, Ptr{T}(C_NULL), 0, PRIME_LAST, VT, R, h, nblocks, q, r) + end + end + nothing +end + +# The butterflies p0 … p1-1 of the grouped first stage of a single line (s = 1, the twiddles pt in +# the grouped layout (k-1)m + p), ncomplex(VT) of them per vector. +@noinline function prime_group_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, tab::Ptr{T}, R::Int, m::Int, p0::Int, p1::Int, + ::Type{VT}) where {T,VT} + Bt = 2sizeof(T); h = (R - 1) >> 1 + nblocks = cld(h, PRIME_BLOCK); q, r = divrem(h, nblocks) + for p in p0:ncomplex(VT):p1-1 + prime_butterfly!(py + R * p * Bt, 0, px + p * Bt, m * Bt, tab, pt + p * Bt, m * Bt, PRIME_GROUP, VT, R, h, nblocks, q, r) + end + nothing +end + +# The grouped first stage: m - m % C butterflies in vectors, the rest one complex value at a time. +@inline function prime_group!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, tab::Ptr{T}, R::Int, m::Int, ::Type{VT}) where {T,VT} + rem = m % ncomplex(VT) + prime_group_range!(py, px, pt, tab, R, m, 0, m - rem, VT) + rem == 0 || prime_group_range!(py, px, pt, tab, R, m, m - rem, m, Vec{2,T}) + nothing +end diff --git a/src/rchain.jl b/src/rchain.jl index 3260364d..8cbd2b11 100644 --- a/src/rchain.jl +++ b/src/rchain.jl @@ -10,7 +10,7 @@ The radices `radices` produces for any length and stride (checked by the tests); `RChain` has a codelet for each of them. """ -const RADICES = (2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 23, 29, 31) +const RADICES = (2:16..., filter(>(16), RADIX_PRIMES)...) # complex values per vector for each entry of vec_ladder(T), widest first, as constants (a map # over the tuple of types is not folded by the compiler and would make the schedule dynamic) @@ -34,76 +34,60 @@ end ((σ == 1 ? 1 : 0) * length(ladder_ncomplex(T)) + vi - 1) * length(RADICES) + ri end -# Radices whose codelets are compiled on first use, not with the package: they are 54 % of the -# statements of the stage kernels and rarely needed. The preference `lazy_radices` is a list of -# radices from RADICES; an empty list compiles every codelet with the package, which a trimmed build -# needs because juliac compiles nothing at run time. -using Preferences: @load_preference -const LAZY_RADICES = Tuple(Int.(@load_preference("lazy_radices", [17, 19, 23, 29, 31]))) -all(in(RADICES), LAZY_RADICES) || error("LAZY_RADICES must be radices from RADICES") - # Codelet dispatch: one function per kind of stage call, a switch over the codelet index, -# generated for both element types so that every call has constant type arguments. Codelets of -# LAZY_RADICES are not in the switch; `_lazy` takes their index and calls the stage kernel -# through a callable that is created and compiled at the first call. +# generated for both element types so that every call has constant type arguments. for T in (Float32, Float64) ladder = vec_ladder(T) - LTYPES = Tuple(ladder) NR = length(RADICES); NL = length(ladder) + # name, arguments, the kernel and its call, and the call of prime_stage! for a looped radix + # (which takes its table `tab`; the full-width first stage has none) + null = :(Ptr{$T}(C_NULL)) kinds = ( - (:cstage!, (:py, :px, :pt, :m, :s), (r, σ, V) -> :(stage!(py, px, pt, m, s, Val($r), Val($σ), $V))), - (:cstage_range!, (:py, :px, :pt, :m, :s, :qlo, :qhi), (r, σ, V) -> :(stage_range!(py, px, pt, m, s, qlo, qhi, 0, m, Val($r), Val($σ), $V))), - (:cstage_last!, (:py, :px, :s), (r, σ, V) -> :(stage_last!(py, px, s, Val($r), Val($σ), $V))), - (:cstage_last_range!, (:py, :px, :s, :qlo, :qhi), (r, σ, V) -> :(stage_last!(py, px, s, qlo, qhi, Val($r), Val($σ), $V))), - (:cstage_group!, (:py, :px, :pt, :m), (r, σ, V) -> :(stage_first_group!(py, px, pt, m, Val($r), Val($σ), $V))), - (:cstage_small!, (:py, :px, :pt, :m), (r, σ, V) -> :(stage_small!(py, px, pt, m, Val($r), Val($σ), $V))), - (:cstage_lanes!, (:py, :px, :pt, :m, :s0, :sx, :sy, :B), (r, σ, V) -> :(stage_lanes!(py, px, pt, m, s0, sx, sy, B, Val($r), Val($σ), $V))), - (:cstage_last_lanes!, (:py, :px, :s0, :sx, :sy, :B), (r, σ, V) -> :(stage_last_lanes!(py, px, s0, sx, sy, B, Val($r), Val($σ), $V)))) - for (name, args, call) in kinds + (:cstage!, (:py, :px, :pt, :m, :s), :stage!, + (r, σ, V) -> :(stage!(py, px, pt, m, s, 0, m, Val($r), Val($σ), $V)), + (r, V) -> :(prime_stage!(py, px, pt, tab, $r, m, 1, s, s, 0, s, $V))), + (:cstage_range!, (:py, :px, :pt, :m, :s, :qlo, :qhi), :stage_range!, + (r, σ, V) -> :(stage_range!(py, px, pt, m, s, qlo, qhi, 0, m, Val($r), Val($σ), $V)), + (r, V) -> :(prime_stage!(py, px, pt, tab, $r, m, 1, s, s, qlo, qhi, $V))), + (:cstage_last!, (:py, :px, :s), :stage_last!, + (r, σ, V) -> :(stage_last!(py, px, s, 0, s, Val($r), Val($σ), $V)), + (r, V) -> :(prime_stage!(py, px, $null, tab, $r, 1, 1, s, s, 0, s, $V))), + (:cstage_last_range!, (:py, :px, :s, :qlo, :qhi), :stage_last!, + (r, σ, V) -> :(stage_last!(py, px, s, qlo, qhi, Val($r), Val($σ), $V)), + (r, V) -> :(prime_stage!(py, px, $null, tab, $r, 1, 1, s, s, qlo, qhi, $V))), + (:cstage_group!, (:py, :px, :pt, :m), :stage_first_group_range!, + (r, σ, V) -> :(stage_first_group!(py, px, pt, m, Val($r), Val($σ), $V)), + (r, V) -> :(prime_group!(py, px, pt, tab, $r, m, $V))), + (:cstage_small!, (:py, :px, :pt, :m), :stage_small!, + (r, σ, V) -> :(stage_small!(py, px, pt, m, Val($r), Val($σ), $V)), nothing), + (:cstage_lanes!, (:py, :px, :pt, :m, :s0, :sx, :sy, :B), :stage_lanes!, + (r, σ, V) -> :(stage_lanes!(py, px, pt, m, s0, sx, sy, B, 0, m, Val($r), Val($σ), $V)), + (r, V) -> :(prime_stage!(py, px, pt, tab, $r, m, s0, sx, sy, 0, B, $V))), + (:cstage_last_lanes!, (:py, :px, :s0, :sx, :sy, :B), :stage_last_lanes!, + (r, σ, V) -> :(stage_last_lanes!(py, px, s0, sx, sy, B, 0, s0, Val($r), Val($σ), $V)), + (r, V) -> :(prime_stage!(py, px, $null, tab, $r, 1, s0, sx, sy, 0, B, $V)))) + for (name, args, kernel, call, prime) in kinds typed = [a in (:py, :px, :pt) ? :($a::Ptr{$T}) : :($a::Int) for a in args] - cases = Any[] + cases = Any[]; pcases = Any[] for (si, σ) in enumerate((-1, 1)), (vi, V) in enumerate(ladder), (ri, r) in enumerate(RADICES) - # the grouped first stage exists for 2 or 4 complex values per vector only; the schedule - # never selects another vector type for it, and code compiled ahead of time must not - # contain the call (its generator would fail) - name === :cstage_group! && !(ncomplex(V) in (2, 4)) && continue - # stage_small! exists for the widest vector type and power-of-two radices - name === :cstage_small! && !(vi == 1 && ispow2(r)) && continue - r in LAZY_RADICES && continue c = ((si - 1) * NL + vi - 1) * NR + ri - push!(cases, :(c == $c && return $(call(r, σ, V)))) - end - # Without lazy radices there is no run-time call at all (the trimmed library needs none). - fallback = isempty(LAZY_RADICES) ? :(throw(ArgumentError("no codelet for this index"))) : - :(return $(Symbol(name, :_lazy))($(args...), c)) - @eval function $name($(typed...), c::Int) - $(cases...) - $fallback - end - isempty(LAZY_RADICES) && continue - # One callable singleton per (radix, sign, vector type): the first call of a codelet creates it - # and compiles its kernel, later calls find it in a table by index. A call through the table - # is a dynamic call of a single method, much cheaper than the dispatch of the stage kernel - # on run-time `Val`s. - callable = Symbol(name, :_, T, :_codelet) - table = Symbol(name, :_, T, :_codelets) - @eval struct $callable{R,σ,VT} end - @eval function (::$callable{R,σ,VT})($(typed...)) where {R,σ,VT} - return $(call(:R, :σ, :VT)) + if looped(r) + prime === nothing || (kernel === :stage_first_group_range! && ncomplex(V) > 4) || + push!(pcases, :(c == $c && return $(prime(r, V)))) + elseif kernel_exists(kernel, ladder, vi, r) + push!(cases, :(c == $c && return $(call(r, σ, V)))) + end end - @eval const $table = Any[nothing for _ in 1:2 * $NL * $NR] - # Index c = ((si - 1) * NL + vi - 1) * NR + ri, with ri = 1 … NR into RADICES. - @eval @noinline function $(Symbol(name, :_, T, :_make))(c::Int) - ri = (c - 1) % $NR + 1 - rest = (c - 1) ÷ $NR - vi = rest % $NL + 1 - si = rest ÷ $NL + 1 - return $callable{RADICES[ri], si == 2 ? 1 : -1, $LTYPES[vi]}() + # The looped radices are behind a call of their own: their kernel takes more arguments than + # the registers hold, and a switch that contains such calls pays for them in every case. + looped_name = Symbol(name, :_looped) + @eval @noinline function $looped_name($(typed...), tab::Ptr{$T}, c::Int) + $(pcases...) + throw(ArgumentError("no codelet for this index")) end - @eval @noinline function $(Symbol(name, :_lazy))($(typed...), c::Int) - f = $table[c] - f === nothing && (f = $table[c] = $(Symbol(name, :_, T, :_make))(c)) - return f($(args...)) + @eval function $name($(typed...), tab::Ptr{$T}, c::Int) + $(cases...) + return $looped_name($(args...), tab, c) end end end @@ -117,14 +101,16 @@ const OP_GROUP = 5 # first stage vectorised over the butterflies p const OP_SMALL = 6 # first stage over the butterflies p on the widest vector (stage_small!) struct StageCall{T} - op::Int + op::Int # the kind of call (OP_…) + 8 · the scalar remainder s mod C of the stride c::Int # codelet of the vectorised part c1::Int # codelet of the scalar remainder (one complex value per vector) m::Int s::Int - rem::Int pt::Ptr{T} # twiddle table of the stage (C_NULL for a last stage), see link_calls! + tab::Ptr{T} # table of a looped radix (primestage.jl), else C_NULL end +stage_kind(st::StageCall) = st.op & 7 +stage_rem(st::StageCall) = st.op >> 3 # One stage of a batch of lines or of a lane block: the vector type is chosen per call from the # stride, so the record holds the codelet index of the widest vector type (the index of vector @@ -132,6 +118,7 @@ end struct BatchStage{T} c::Int pt::Ptr{T} # C_NULL for the last stage + tab::Ptr{T} # table of a looped radix, else C_NULL end """ @@ -158,8 +145,8 @@ function build_rchain(::Type{T}, n::Int, σ::Int, rs::AbstractVector{Int}) where all(in(RADICES), rs) || throw(ArgumentError("radix without a codelet in RADICES")) tw, tw1 = chain_tables(T, n, σ, rs) calls = single_schedule(T, n, σ, collect(Int, rs)) - batch = [BatchStage{T}(codelet_index(T, σ, 1, r), Ptr{T}(C_NULL)) for r in rs] - P = RChain{T}(n, σ, collect(Int, rs), tw, tw1, calls, batch, length(calls) == 1 ? calls[1].c : 0, UInt(0)) + batch = [BatchStage{T}(codelet_index(T, σ, 1, r), Ptr{T}(C_NULL), Ptr{T}(C_NULL)) for r in rs] + P = RChain{T}(n, σ, collect(Int, rs), tw, tw1, calls, batch, length(calls) == 1 && !looped(rs[1]) ? calls[1].c : 0, UInt(0)) link_calls!(P) P end @@ -173,16 +160,18 @@ end @inline linked(P::RChain) = P.linked == UInt(pointer_from_objref(P)) @noinline function link_calls!(P::RChain{T}) where {T} + tab(i) = looped(P.rs[i]) ? pointer(prime_table(T, P.rs[i], P.σ)) : Ptr{T}(C_NULL) calls = P.single for i in eachindex(calls) st = calls[i] - pt = st.op == OP_GROUP || st.op == OP_SMALL ? cptr(P.tw1) : - st.op == OP_MID || st.op == OP_MID_TAIL ? cptr(P.tw[i]) : Ptr{T}(C_NULL) - calls[i] = StageCall{T}(st.op, st.c, st.c1, st.m, st.s, st.rem, pt) + op = stage_kind(st) + pt = op == OP_GROUP || op == OP_SMALL ? cptr(P.tw1) : + op == OP_MID || op == OP_MID_TAIL ? cptr(P.tw[i]) : Ptr{T}(C_NULL) + calls[i] = StageCall{T}(st.op, st.c, st.c1, st.m, st.s, pt, tab(i)) end batch = P.batch - for i in 1:length(batch)-1 - batch[i] = BatchStage{T}(batch[i].c, cptr(P.tw[i])) + for i in eachindex(batch) + batch[i] = BatchStage{T}(batch[i].c, i < length(batch) ? cptr(P.tw[i]) : Ptr{T}(C_NULL), tab(i)) end P.linked = UInt(pointer_from_objref(P)) nothing @@ -215,13 +204,13 @@ function single_schedule(::Type{T}, n::Int, σ::Int, rs::Vector{Int}) where {T} c = codelet_index(T, σ, vi, r); c1 = codelet_index(T, σ, nv, r) m = L ÷ r if i == 1 && small_first(ncs, group, r, n) - push!(calls, StageCall{T}(OP_SMALL, codelet_index(T, σ, 1, r), c1, m, s, 0, Ptr{T}(C_NULL))) + push!(calls, StageCall{T}(OP_SMALL, codelet_index(T, σ, 1, r), c1, m, s, Ptr{T}(C_NULL), Ptr{T}(C_NULL))) elseif i == K - push!(calls, StageCall{T}(rem == 0 ? OP_LAST : OP_LAST_TAIL, c, c1, m, s, rem, Ptr{T}(C_NULL))) + push!(calls, StageCall{T}((rem == 0 ? OP_LAST : OP_LAST_TAIL) + 8rem, c, c1, m, s, Ptr{T}(C_NULL), Ptr{T}(C_NULL))) elseif i == 1 && group != 0 && ncs[vi] == 1 - push!(calls, StageCall{T}(OP_GROUP, codelet_index(T, σ, group, r), c1, m, s, 0, Ptr{T}(C_NULL))) + push!(calls, StageCall{T}(OP_GROUP, codelet_index(T, σ, group, r), c1, m, s, Ptr{T}(C_NULL), Ptr{T}(C_NULL))) else - push!(calls, StageCall{T}(rem == 0 ? OP_MID : OP_MID_TAIL, c, c1, m, s, rem, Ptr{T}(C_NULL))) + push!(calls, StageCall{T}((rem == 0 ? OP_MID : OP_MID_TAIL) + 8rem, c, c1, m, s, Ptr{T}(C_NULL), Ptr{T}(C_NULL))) end L = m; s *= r end @@ -280,19 +269,19 @@ end if i == K r = L if rem == 0 - cstage_last!(dst, src, s, c) + cstage_last!(dst, src, s, st.tab, c) else - cstage_last_range!(dst, src, s, 0, s - rem, c) - cstage_last_range!(dst, src, s, s - rem, s, st.c + (nv - 1) * NR) + cstage_last_range!(dst, src, s, 0, s - rem, st.tab, c) + cstage_last_range!(dst, src, s, s - rem, s, st.tab, st.c + (nv - 1) * NR) end else r = RADICES[(st.c - 1) % NR + 1] m = L ÷ r if rem == 0 - cstage!(dst, src, st.pt, m, s, c) + cstage!(dst, src, st.pt, m, s, st.tab, c) else - cstage_range!(dst, src, st.pt, m, s, 0, s - rem, c) - cstage_range!(dst, src, st.pt, m, s, s - rem, s, st.c + (nv - 1) * NR) + cstage_range!(dst, src, st.pt, m, s, 0, s - rem, st.tab, c) + cstage_range!(dst, src, st.pt, m, s, s - rem, s, st.tab, st.c + (nv - 1) * NR) end L = m; s *= r end @@ -306,7 +295,7 @@ end @inline function run_single!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::RChain{T}) where {T} c = P.one if c != 0 - cstage_last!(py, px, 1, c) + cstage_last!(py, px, 1, Ptr{T}(C_NULL), c) return nothing end run_calls!(py, px, pw1, pw2, P) @@ -322,21 +311,21 @@ end for i in 1:K st = unsafe_load(pc, i) dst = stage_dst(py, px, pw1, pw2, i, K, inplace) - op = st.op + op = stage_kind(st); rem = stage_rem(st) if op == OP_LAST - @inline cstage_last!(dst, src, st.s, st.c) + @inline cstage_last!(dst, src, st.s, st.tab, st.c) elseif op == OP_MID - @inline cstage!(dst, src, st.pt, st.m, st.s, st.c) + @inline cstage!(dst, src, st.pt, st.m, st.s, st.tab, st.c) elseif op == OP_SMALL - @inline cstage_small!(dst, src, st.pt, st.m, st.c) + @inline cstage_small!(dst, src, st.pt, st.m, st.tab, st.c) elseif op == OP_GROUP - @inline cstage_group!(dst, src, st.pt, st.m, st.c) + @inline cstage_group!(dst, src, st.pt, st.m, st.tab, st.c) elseif op == OP_MID_TAIL - @inline cstage_range!(dst, src, st.pt, st.m, st.s, 0, st.s - st.rem, st.c) - @inline cstage_range!(dst, src, st.pt, st.m, st.s, st.s - st.rem, st.s, st.c1) + @inline cstage_range!(dst, src, st.pt, st.m, st.s, 0, st.s - rem, st.tab, st.c) + @inline cstage_range!(dst, src, st.pt, st.m, st.s, st.s - rem, st.s, st.tab, st.c1) else - @inline cstage_last_range!(dst, src, st.s, 0, st.s - st.rem, st.c) - @inline cstage_last_range!(dst, src, st.s, st.s - st.rem, st.s, st.c1) + @inline cstage_last_range!(dst, src, st.s, 0, st.s - rem, st.tab, st.c) + @inline cstage_last_range!(dst, src, st.s, st.s - rem, st.s, st.tab, st.c1) end src = dst end @@ -394,7 +383,8 @@ function run_lanes!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::RChain{ NR = length(RADICES) off = (vi - 1) * NR if K == 1 # one stage loads all inputs of a butterfly before it stores: py may equal px - cstage_last_lanes!(py, px, 1, Sx, Sy, B, unsafe_load(pb).c + off) + st = unsafe_load(pb) + cstage_last_lanes!(py, px, 1, Sx, Sy, B, st.tab, st.c + off) return nothing end L = P.n; s0 = 1; src = px; srcS = Sx @@ -402,12 +392,12 @@ function run_lanes!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::RChain{ st = unsafe_load(pb, i) c = st.c + off if i == K - cstage_last_lanes!(py, src, s0, srcS, Sy, B, c) + cstage_last_lanes!(py, src, s0, srcS, Sy, B, st.tab, c) else dst = isodd(i) ? pw1 : pw2 r = @inbounds RADICES[(st.c - 1) % NR + 1] m = L ÷ r - cstage_lanes!(dst, src, st.pt, m, s0, srcS, B, B, c) + cstage_lanes!(dst, src, st.pt, m, s0, srcS, B, B, st.tab, c) L = m; s0 *= r src = dst; srcS = B end @@ -421,7 +411,7 @@ function run_rows!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::RChain{T c = P.one if c != 0 for j in 0:count-1 - @inline cstage_last!(py + j * bytes, px + j * bytes, 1, c) + @inline cstage_last!(py + j * bytes, px + j * bytes, 1, Ptr{T}(C_NULL), c) end else run_calls_rows!(py, px, pw1, pw2, P, count, bytes) diff --git a/test/inference.jl b/test/inference.jl index 14f97e64..125fdd03 100644 --- a/test/inference.jl +++ b/test/inference.jl @@ -19,7 +19,7 @@ @test_broken @inferred FFTL.plan_fft(x, (1, 2)) end # Bluestein length, the plan holds the work buffers - let x = randn(Complex{T}, 37), y = similar(x) + let x = randn(Complex{T}, 79), y = similar(x) p = FFTL.plan_fft(x) mul!(y, p, x) @test (@inferred mul!(y, p, x)) === y @@ -60,7 +60,7 @@ @test_broken @inferred FFTL.plan_pass(v, 2) end # Complex views: independent layouts, multidimensional passes and the large column block. - for (sz, region) in (((10240, 4), (1,)), ((3, 8, 2), (1, 3)), ((4, 37, 2), (2,)), ((3, 1, 2), (2,))) + for (sz, region) in (((10240, 4), (1,)), ((3, 8, 2), (1, 3)), ((4, 79, 2), (2,)), ((3, 1, 2), (2,))) inputs = realpass_layouts(Complex{T}, sz) outputs = realpass_layouts(Complex{T}, sz) for mk in (FFTL.plan_fft, FFTL.plan_bfft, FFTL.plan_ifft) @@ -75,7 +75,7 @@ end end # High-level real plans and their inverses, with independent and changing layouts. - for (sz, region) in (((10240, 4), (1,)), ((3, 8, 2), (2, 3)), ((3, 37, 2), (2,)), ((3, 1, 2), (2,))) + for (sz, region) in (((10240, 4), (1,)), ((3, 8, 2), (2, 3)), ((3, 79, 2), (2,)), ((3, 1, 2), (2,))) d = first(region); n = sz[d]; osz = Base.setindex(sz, n ÷ 2 + 1, d) reals = realpass_layouts(T, sz); spectra = realpass_layouts(Complex{T}, osz) pf = FFTL.plan_rfft(reals[2], region; threads = 1) @@ -92,7 +92,7 @@ end # real transforms # Real passes reuse their offset storage even on the first call with new stride values. - for (sz, d) in (((8,), 1), ((3, 9, 2), 2), ((2, 37, 1), 2), ((3, 1, 2), 2), ((10240, 4), 1)) + for (sz, d) in (((8,), 1), ((3, 9, 2), 2), ((2, 79, 1), 2), ((3, 1, 2), 2), ((10240, 4), 1)) n = sz[d]; osz = Base.setindex(sz, n ÷ 2 + 1, d) for forward in (true, false) inputs = realpass_layouts(forward ? T : Complex{T}, forward ? sz : osz) @@ -159,7 +159,7 @@ exe_r2r!(p, y, x) = GC.@preserve x y FFTL.execute_r2r!(p, pointer(y), pointer(x) @test (@inferred exe!(p, y, x)) === nothing @test (@allocated exe!(p, y, x)) == 0 end - let p = (@inferred FFTL.RPlan{T}([37], [1], 1; threads = 1)), x = randn(Complex{T}, 37), y = similar(x) # Bluestein + let p = (@inferred FFTL.RPlan{T}([79], [1], 1; threads = 1)), x = randn(Complex{T}, 79), y = similar(x) # Bluestein exe!(p, y, x) @test (@allocated exe!(p, y, x)) == 0 end diff --git a/test/radix_table.jl b/test/radix_table.jl index 412ad456..d0ef8862 100644 --- a/test/radix_table.jl +++ b/test/radix_table.jl @@ -9,10 +9,10 @@ const RADIX_TABLE_STRIDES = (1, 1024) const RADIX_TABLE_LARGE = (2048, 4096, 8192, 16384, 32768, 65536, 100000, 120000, 131072, 262144, 786432, 1048576) -# Lengths with prime factors up to MAX_RADIX_PRIME (the lengths a chain handles), from 2 to 1024 -# and the large ones above. +# Lengths with prime factors up to 31 (the radices the rules were derived for), from 2 to 1024 and +# the large ones above. radix_table_lengths() = - vcat(filter(n -> FFTL.factorize(n)[2] == 1, 2:1024), collect(RADIX_TABLE_LARGE)) + vcat(filter(n -> FFTL.factorize(n, 31)[2] == 1, 2:1024), collect(RADIX_TABLE_LARGE)) # x86 has two rule sets: 16 vector registers (AVX2) and 32 (AVX-512) radix_table_path(arch = string(Sys.ARCH)) = diff --git a/test/runtests.jl b/test/runtests.jl index f5d11e2b..86d79881 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -69,8 +69,10 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] @testset "codelets and chains against the direct DFT, $T" for T in (Float64, Float32) for n in pick([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 24, 32, 48, 49, 60, 64, 80, 100, 105, 120, 128, 143, 169, 210, 243, 256, 360, 384, 500, 512, 1000, 1024, - 17, 19, 23, 29, 31, 68, 289, 928, 31 * 27], - thin(T, [1, 2, 3, 4, 5, 7, 8, 11, 13, 16, 31, 60, 128, 243, 1000, 31 * 27])) + 17, 19, 23, 29, 31, 68, 289, 928, 31 * 27, + # the looped prime radices: alone, first, last, twice, with a scalar remainder + 37, 41, 53, 61, 67, 73, 2 * 47, 8 * 59, 64 * 71, 23 * 43, 61^2, 3 * 73 * 5, 17 * 37 * 16], + thin(T, [1, 2, 3, 4, 5, 7, 8, 11, 13, 16, 31, 60, 128, 243, 1000, 31 * 27, 73, 8 * 59, 23 * 43, 3 * 73 * 5])) x = randn(Complex{T}, n) ref = dft(ComplexF64.(x)) @test relerr(FFTL.fft(x), ref) < tol(T) @@ -79,7 +81,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end @testset "Bluestein, $T" for T in (Float64, Float32) - for n in pick([37, 41, 74, 101, 37^2, 1009], [37, 74, 1009]) + for n in pick([79, 83, 158, 101, 79^2, 1009], [79, 158, 1009]) x = randn(Complex{T}, n) @test relerr(FFTL.fft(x), dft(ComplexF64.(x))) < 10tol(T) @test relerr(FFTL.bfft(x), dft(ComplexF64.(x), 1)) < 10tol(T) @@ -88,9 +90,9 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] @testset "smooth part times leftover, $T" for T in (Float64, Float32) # n = s·q: a chain of length s, Bluestein on the leftover q (prime or composite) - for n in pick([2 * 37, 3 * 41, 16 * 43, 29 * 37, 64 * 37, 2 * 37^2, 5 * 13 * 37, 2 * 37 * 41, 6 * 1009, - 2 * 37 * 41 * 43, 6 * 37^2 * 41], - [2 * 37, 29 * 37, 2 * 37 * 41, 6 * 1009]) + for n in pick([2 * 79, 3 * 83, 16 * 89, 29 * 79, 64 * 79, 2 * 79^2, 5 * 13 * 79, 2 * 79 * 83, 6 * 1009, + 17 * 79 * 83, 6 * 79^2], + [2 * 79, 29 * 79, 2 * 79 * 83, 6 * 1009]) x = randn(Complex{T}, n) @test FFTL.plan_fft(x; threads = 1).kernels[1] isa FFTL.Composite @test relerr(FFTL.fft(x), dft(ComplexF64.(x))) < 10tol(T) @@ -98,8 +100,8 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end # a leftover split into prime powers (the planner does this for large leftovers only): first # stages that are composite and Bluestein kernels themselves - for (n, s, q) in pick([(2 * 37 * 41 * 43, 2 * 37 * 41, 43), (37 * 41, 37, 41), (6 * 37^2 * 41, 6 * 37^2, 41)], - [(2 * 37 * 41 * 43, 2 * 37 * 41, 43), (37 * 41, 37, 41)]) + for (n, s, q) in pick([(2 * 79 * 83 * 89, 2 * 79 * 83, 89), (79 * 83, 79, 83), (6 * 79^2 * 83, 6 * 79^2, 83)], + [(2 * 79 * 83 * 89, 2 * 79 * 83, 89), (79 * 83, 79, 83)]) x = randn(Complex{T}, n) for σ in (-1, 1) k = FFTL.Composite{T}(n, s, q, σ, 1) @@ -110,7 +112,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end end # lanes of a 2D array: along dimension 1 (contiguous lines) and dimension 2 (strided lanes) - for (sz, region) in ((((2 * 37, 5)), (1,)), ((3, 2 * 37), (2,)), ((5, 3 * 41), (1, 2))) + for (sz, region) in ((((2 * 79, 5)), (1,)), ((3, 2 * 79), (2,)), ((5, 3 * 83), (1, 2))) x = randn(Complex{T}, sz) @test relerr(FFTL.fft(x, region), dftn(x, region)) < 10tol(T) z = copy(x); p = FFTL.plan_fft(x, region); mul!(z, p, z) @@ -129,11 +131,11 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end @testset "multidimensional, $T" for T in (Float64, Float32) - cases = pick([((8, 12), (1, 2)), ((16, 16), (1, 2)), ((17, 6), (1, 2)), ((5, 37), (1, 2)), + cases = pick([((8, 12), (1, 2)), ((16, 16), (1, 2)), ((17, 6), (1, 2)), ((5, 79), (1, 2)), ((3, 8), (2,)), ((3, 8), (1, 2)), ((12, 7), (1,)), ((1, 9), (1, 2)), ((4, 6, 10), (1, 2, 3)), ((3, 5, 7), (1, 2, 3)), ((6, 4, 5), (1, 3)), ((6, 4, 5), (2,)), ((16, 9, 2, 3), (1, 2, 3, 4))], - thin(T, [((8, 12), (1, 2)), ((5, 37), (1, 2)), ((3, 8), (2,)), ((1, 9), (1, 2)), + thin(T, [((8, 12), (1, 2)), ((5, 79), (1, 2)), ((3, 8), (2,)), ((1, 9), (1, 2)), ((4, 6, 10), (1, 2, 3)), ((6, 4, 5), (1, 3)), ((16, 9, 2, 3), (1, 2, 3, 4))])) for (sz, region) in cases x = randn(Complex{T}, sz) @@ -150,9 +152,9 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end @testset "complex strided plans, $T" for T in (Float64, Float32) - for (sz, region) in (((8,), (1,)), ((37, 3), (1,)), ((3, 74, 2), (2,)), + for (sz, region) in (((8,), (1,)), ((79, 3), (1,)), ((3, 158, 2), (2,)), ((8, 6), (1, 2)), ((3, 8, 2), (2,)), ((4, 3, 2), (1, 3)), - ((19, 8, 2), (2,)), ((16, 37, 2), (2,)), ((1, 9), (1, 2)), + ((19, 8, 2), (2,)), ((16, 79, 2), (2,)), ((1, 9), (1, 2)), ((3, 1, 2), (2,)), ((2, 3, 4, 2), (3,))) x = randn(Complex{T}, sz) inputs = realpass_layouts(Complex{T}, sz) @@ -205,7 +207,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] @test_throws DimensionMismatch mul!(zeros(Complex{T}, 2, 4), p, v) if Threads.nthreads() > 1 for (sz, region) in (((128, 512), (1,)), ((33, 1024), (2,)), - ((8192,), (1,)), ((20011,), (1,)), ((2 * 37 * 41 * 43,), (1,))) + ((8192,), (1,)), ((20011,), (1,)), ((2 * 79 * 83 * 89,), (1,))) x = randn(Complex{T}, sz) v = realpass_layouts(Complex{T}, sz)[4]; v .= x w = realpass_layouts(Complex{T}, sz)[5] @@ -221,7 +223,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] # slices, interior-axis lane blocks (including a partial block), and multiple passes. # The three-lane case has gather workspace with two planned tasks, but none with three. for (sz, region) in (((1009, 64), (1,)), ((65, 1009), (2,)), ((2018, 33), (1,)), - ((65, 2018), (2,)), ((19, 129, 32), (2, 3)), ((3, 6054, 2), (2,))) + ((65, 2018), (2,)), ((19, 237, 32), (2, 3)), ((3, 6054, 2), (2,))) x = randn(Complex{T}, sz) inputs = realpass_layouts(Complex{T}, sz) third = view(fill(Complex{T}(7), 3 .* sz .+ 4), ntuple(i -> 2:3:3sz[i]-1, length(sz))...) @@ -255,8 +257,8 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] @testset "threads, $T" for T in (Float64, Float32) # sizes above the threading threshold: slices (dimension 1 of a 2D array), lane blocks # (last dimension), odd strides (last block with fewer lanes than the vector width) - for sz in pick([(256, 256), (64, 64, 64), (33, 1024), (100, 100, 10), (8, 4096), (1024, 37), (2 * 3^11,), (5^7,)], - [(64, 64, 64), (33, 1024), (1024, 37), (2 * 3^11,)]) + for sz in pick([(256, 256), (64, 64, 64), (33, 1024), (100, 100, 10), (8, 4096), (1024, 79), (2 * 3^11,), (5^7,)], + [(64, 64, 64), (33, 1024), (1024, 79), (2 * 3^11,)]) x = randn(Complex{T}, sz) ref = FFTL.plan_fft(x; threads = 1) * x for nt in pick((2, 3, 4), (3,)) @@ -302,10 +304,10 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end end # many lines with a prime factor above MAX_RADIX_PRIME: slices and lane blocks over the tasks, - # Bluestein kernels (127, 1009, 37) and composite kernels (129 = 3·43, 2018 = 2·1009) + # Bluestein kernels (127, 1009, 79) and composite kernels (237 = 3·79, 2018 = 2·1009) for (sz, region) in pick([((256, 127, 16), 1:3), ((16, 256, 127), (3,)), ((1009, 64), (1,)), ((64, 1009), (2,)), - ((37, 37, 37), 1:3), ((129, 300), 1:2), ((300, 129), 1:2), ((2018, 32), (1,)), ((32, 2018), (2,))], - [((256, 127, 16), 1:3), ((64, 1009), (2,)), ((1009, 64), (1,)), ((300, 129), 1:2)]) + ((79, 79, 79), 1:3), ((237, 300), 1:2), ((300, 237), 1:2), ((2018, 32), (1,)), ((32, 2018), (2,))], + [((256, 127, 16), 1:3), ((64, 1009), (2,)), ((1009, 64), (1,)), ((300, 237), 1:2)]) x = randn(Complex{T}, sz...) for mk in (FFTL.plan_fft, FFTL.plan_bfft) p1 = mk(x, region; threads = 1) @@ -326,7 +328,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end @testset "real strided plans, $T" for T in (Float64, Float32) - for (sz, region) in (((1,), (1,)), ((8,), (1,)), ((9, 4), (1,)), ((3, 37, 2), (2,)), + for (sz, region) in (((1,), (1,)), ((8,), (1,)), ((9, 4), (1,)), ((3, 79, 2), (2,)), ((4, 6), (1, 2)), ((4, 5), (2, 1)), ((3, 8, 2), (2, 3)), ((8, 3, 4), (1, 3)), ((3, 1, 2), (2,))) # Preserve the existing sorted region and choice of the halved dimension. @@ -394,7 +396,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end # The complex passes of real plans also select LinesMT, including inverse passes into # the plan's dense scratch. Use independent dense references for the inverse inputs. - for (sz, region) in (((4, 1009, 32), (1, 2)), ((8, 65, 129), (1, 3)), ((65, 8, 129), (2, 3))) + for (sz, region) in (((4, 1009, 32), (1, 2)), ((8, 65, 237), (1, 3)), ((65, 8, 237), (2, 3))) d = first(region); n = sz[d]; osz = Base.setindex(sz, n ÷ 2 + 1, d) x = randn(T, sz) ref = FFTL.plan_rfft(x, region; threads = 1) * x @@ -458,7 +460,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] @test relerr(yv ./ n, x) < 10tol(T) end # Reuse each plan while changing input strides, output strides, both, then back to dense. - for (sz, d) in (((8,), 1), ((9, 4), 1), ((3, 8, 2), 2), ((2, 37, 1), 2), ((3, 1, 2), 2)) + for (sz, d) in (((8,), 1), ((9, 4), 1), ((3, 8, 2), 2), ((2, 79, 1), 2), ((3, 1, 2), 2)) n = sz[d] x = randn(T, sz) osz = Base.setindex(sz, n ÷ 2 + 1, d) @@ -550,12 +552,12 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] @testset "run-time pass plans, $T" for T in (Float64, Float32) # the plan type depends on the element type and the dimension count only - for (sz, d) in (((100, 3), 1), ((1009, 3), 1), ((64, 5), 1), ((37, 6), 2), ((6, 130), 2)) + for (sz, d) in (((100, 3), 1), ((1009, 3), 1), ((64, 5), 1), ((79, 6), 2), ((6, 130), 2)) p = FFTL.plan_pass_runtime(randn(Complex{T}, sz), d) @test p isa FFTL.PassPlan{T,2,FFTL.RPassKern{T}} end # against the length-specialised plan, every kernel kind: smooth, Bluestein, composite, four-step - for (n, sz_other, nt) in ((64, 5, 1), (100, 7, 1), (1000, 3, 2), (1009, 4, 1), (2 * 37 * 41, 2, 1), + for (n, sz_other, nt) in ((64, 5, 1), (100, 7, 1), (1000, 3, 2), (1009, 4, 1), (2 * 79 * 83, 2, 1), (2^17, 1, 3)) x = randn(Complex{T}, n, sz_other) for sign in (-1, 1) @@ -591,11 +593,11 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end @testset "real transforms, $T" for T in (Float64, Float32) - for (sz, region) in pick([((8,), 1), ((9,), 1), ((10,), 1), ((15,), 1), ((16,), 1), ((37,), 1), ((100,), 1), + for (sz, region) in pick([((8,), 1), ((9,), 1), ((10,), 1), ((15,), 1), ((16,), 1), ((79,), 1), ((100,), 1), ((101,), 1), ((1024,), 1), ((1023,), 1), ((3 * 2^11,), 1), ((8, 6), (1, 2)), ((9, 6), (1, 2)), ((7, 10), (1, 2)), ((8, 6), (2,)), ((12, 7), (2,)), ((6, 8, 5), (1, 2, 3)), ((6, 8, 5), (2, 3)), ((6, 9, 5), (2, 3)), ((5, 4, 6), (3,)), ((16, 9), (1,))], - thin(T, [((9,), 1), ((16,), 1), ((37,), 1), ((1024,), 1), ((3 * 2^11,), 1), + thin(T, [((9,), 1), ((16,), 1), ((79,), 1), ((1024,), 1), ((3 * 2^11,), 1), ((9, 6), (1, 2)), ((12, 7), (2,)), ((6, 8, 5), (1, 2, 3)), ((6, 9, 5), (2, 3))])) x = randn(T, sz) reg = sort([region...]); d1 = reg[1]; n = sz[d1]; h = n ÷ 2 + 1 @@ -759,22 +761,23 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end # a chain whose stage records are not linked to it (a copy, or one loaded from a package # image) links them again at the first run - rs = FFTL.radices(60, 1; esize = sizeof(Complex{T})) - ref = FFTL.RChain{T}(60, -1, rs); rc = FFTL.build_rchain(T, 60, -1, rs) + # (4·3·17: the table of the looped radix is linked as well) + rs = FFTL.radices(204, 1; esize = sizeof(Complex{T})) + ref = FFTL.RChain{T}(204, -1, rs); rc = FFTL.build_rchain(T, 204, -1, rs) unlink!(c) = (for i in eachindex(c.single) st = c.single[i] - c.single[i] = FFTL.StageCall{T}(st.op, st.c, st.c1, st.m, st.s, st.rem, C_NULL) - c.batch[i] = FFTL.BatchStage{T}(c.batch[i].c, C_NULL) + c.single[i] = FFTL.StageCall{T}(st.op, st.c, st.c1, st.m, st.s, C_NULL, C_NULL) + c.batch[i] = FFTL.BatchStage{T}(c.batch[i].c, C_NULL, C_NULL) end; c.linked = 0) - x = randn(Complex{T}, 4 * 60); y1 = similar(x); y2 = similar(x); w1 = similar(x); w2 = similar(x) + x = randn(Complex{T}, 4 * 204); y1 = similar(x); y2 = similar(x); w1 = similar(x); w2 = similar(x) GC.@preserve x y1 y2 w1 w2 begin unlink!(rc); FFTL.run!(p(y1), p(x), p(w1), p(w2), ref, 1); FFTL.run!(p(y2), p(x), p(w1), p(w2), rc, 1) - @test y1[1:60] == y2[1:60] + @test y1[1:204] == y2[1:204] unlink!(rc); FFTL.run!(p(y1), p(x), p(w1), p(w2), ref, 4); FFTL.run!(p(y2), p(x), p(w1), p(w2), rc, 4) @test y1 == y2 unlink!(rc); FFTL.run_lanes!(p(y1), p(x), p(w1), p(w2), ref, 4, 4, 4); FFTL.run_lanes!(p(y2), p(x), p(w1), p(w2), rc, 4, 4, 4) @test y1 == y2 - unlink!(rc); FFTL.run_rows!(p(y1), p(x), p(w1), p(w2), ref, 4, 60); FFTL.run_rows!(p(y2), p(x), p(w1), p(w2), rc, 4, 60) + unlink!(rc); FFTL.run_rows!(p(y1), p(x), p(w1), p(w2), ref, 4, 204); FFTL.run_rows!(p(y2), p(x), p(w1), p(w2), rc, 4, 204) @test y1 == y2 end # a chain of length 1 copies @@ -794,7 +797,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] rs = [r, 2C] n = prod(rs) rc = FFTL.RChain{T}(n, σ, rs) - @test (rc.single[1].op == FFTL.OP_SMALL) == (C > 4) + @test (FFTL.stage_kind(rc.single[1]) == FFTL.OP_SMALL) == (C > 4) x = randn(Complex{T}, n); w1 = similar(x); w2 = similar(x); y = similar(x) GC.@preserve x y w1 w2 FFTL.run_single!(p(y), p(x), p(w1), p(w2), rc) @test relerr(y, dft(ComplexF64.(x), σ)) < tol(T) @@ -804,8 +807,8 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] # run-time plans (RChain kernels): same results as the reference, four-step and plane fusion @testset "RPlan, $T" for T in pick((Float64, Float32), (Float64,)) for (sz, reg, nt) in pick([((1000,), [1], 1), ((2 * 3^11,), [1], 3), ((1009,), [1], 1), ((64, 48), [1, 2], 1), - ((37, 100), [1, 2], 3), ((16, 9, 5), [1, 2, 3], 1), ((8, 37, 6), [2], 3), ((40, 36, 10), [1, 3], 3)], - [((2 * 3^11,), [1], 3), ((1009,), [1], 1), ((37, 100), [1, 2], 3), ((16, 9, 5), [1, 2, 3], 1), ((8, 37, 6), [2], 3)]) + ((79, 100), [1, 2], 3), ((16, 9, 5), [1, 2, 3], 1), ((8, 79, 6), [2], 3), ((40, 36, 10), [1, 3], 3)], + [((2 * 3^11,), [1], 3), ((1009,), [1], 1), ((79, 100), [1, 2], 3), ((16, 9, 5), [1, 2, 3], 1), ((8, 79, 6), [2], 3)]) x = randn(Complex{T}, sz) p = FFTL.RPlan{T}(collect(sz), reg, -1; threads = nt) ref = FFTL.plan_fft(x, Tuple(reg); threads = 1) * x @@ -855,7 +858,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] setprecision(BigFloat, 128) do B = BigFloat; btol = 1e-35 dftB(x, σ) = [sum(x[j+1] * cispi(B(2σ) * mod(j * k, length(x)) / length(x)) for j in 0:length(x)-1) for k in 0:length(x)-1] - for n in pick((1, 2, 7, 12, 16, 37, 97, 100), (7, 16, 37)) + for n in pick((1, 2, 7, 12, 16, 79, 97, 100), (7, 16, 79)) x = Complex{B}.(randn(ComplexF64, n)) @test eltype(FFTL.fft(x)) == Complex{B} @test relerr(FFTL.fft(x), dftB(x, -1)) < btol @@ -865,7 +868,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] x = Complex{B}.(randn(ComplexF64, 6, 5, 4)) @test relerr(FFTL.fft(x, (1, 3)), FFTL.fft(FFTL.fft(x, 1), 3)) < btol @test relerr(FFTL.ifft(FFTL.fft(x)), x) < btol - for (sz, reg) in (((16,), 1:1), ((37,), 1:1), ((6, 5), 1:2), ((6, 5), 2:2), ((4, 5, 3), (1, 3))) + for (sz, reg) in (((16,), 1:1), ((79,), 1:1), ((6, 5), 1:2), ((6, 5), 2:2), ((4, 5, 3), (1, 3))) xr = B.(randn(sz...)) X = FFTL.rfft(xr, reg) d1 = minimum(reg) @@ -897,7 +900,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] @test relerr(FFTL.plan_brfft(X, n; threads = 3) * X, n .* x) < 10tol(T) end # non-smooth lengths, even and odd: the half-length (even) or the full-length (odd) DFT on threads - for (T, n) in ((Float64, 37888), (Float64, 92935), (Float64, 2 * 1009 * 37), (Float32, 33490), (Float32, 92935), (Float32, 100003)) + for (T, n) in ((Float64, 80896), (Float64, 92935), (Float64, 2 * 1009 * 79), (Float32, 33490), (Float32, 92935), (Float32, 100003)) x = randn(T, n) p = FFTL.plan_rfft(x; threads = 3) @test p.four isa Union{FFTL.CompositeMT,FFTL.BluesteinMT}