fft/engine.hpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include <concepts> | ||
| 4 | #include <span> | ||
| 5 | #include <type_traits> | ||
| 6 | #include <utility> | ||
| 7 | |||
| 8 | #include "fft/common.hpp" | ||
| 9 | |||
| 10 | namespace wala::fft { | ||
| 11 | |||
| 12 | // ==== engine concept ==== | ||
| 13 | |||
| 14 | // Output operations for the finish step to express arbitrary fusion into the output buffer. | ||
| 15 | 706779328 | struct assign_op { template <typename T> void operator()(T& d, T v) const { d = v; } }; | |
| 16 | 56367 | struct add_op { template <typename T> void operator()(T& d, T v) const { d += v; } }; | |
| 17 | ✗ | struct sub_op { template <typename T> void operator()(T& d, T v) const { d -= v; } }; | |
| 18 | 4760 | struct add_twice_op { template <typename T> void operator()(T& d, T v) const { d += v + v; } }; | |
| 19 | |||
| 20 | // `engine` contract | ||
| 21 | // engine represents a way of packing/unpacking sequences over an arbitrary ring into FFT-style transforms. | ||
| 22 | // We expect transforms/products of transforms to be linear but potentially lossy/imprecise, so we'll track precision | ||
| 23 | // at compile-time as a template parameter. | ||
| 24 | // | ||
| 25 | // E::value_type The ring we operate over | ||
| 26 | // E::unit_scale 0 or 1 depending on whether there's error that can accumulate | ||
| 27 | // E::commutative A marker for whether the ring is commutative | ||
| 28 | // | ||
| 29 | // transformed_t<A> The transform of a sequence. This object owns its data buffer. | ||
| 30 | // product_t<A> The product of 2 transforms. May equal transformed_t, particularly when unit_scale = 0. | ||
| 31 | // | ||
| 32 | // transformed alias for transformed_t<unit_scale> | ||
| 33 | // product alias for product_t<unit_scale> | ||
| 34 | // | ||
| 35 | // The basic multiplication API is | ||
| 36 | // transform(span<const value_type> in, int n) -> transformed_t<unit_scale> | ||
| 37 | // mul(transformed_t<A>, transformed_t<B>, int n) -> product_t<A*B> | ||
| 38 | // mul2(a1, b1, a2, b2, int n) -> product_t<A1*B1 + A2*B2>, computing a1*b1 + a2*b2 in one pass | ||
| 39 | // finish(product_t<A>, span<value_type>& out, Op) -> void | ||
| 40 | // | ||
| 41 | // Input span can be length up to 2n. | ||
| 42 | // Output spans can be length up to n; only the prefix that exists is filled. | ||
| 43 | // finish applies Op exactly once per out element, in index order, to | ||
| 44 | // value_type targets (so ops may be stateful). | ||
| 45 | // Transforms can be longer than necessary, and only the relevant prefix is used. | ||
| 46 | // | ||
| 47 | // For non-exact engines, there's some subtlety in whether we wrap before or after packing. | ||
| 48 | // We will choose to wrap *after* packing, which hurts error bounds but makes the prefix condition more uniform. | ||
| 49 | // | ||
| 50 | // Additionally, we have APIs to take advantage of linearity in both transformed and product space: | ||
| 51 | // add(transformed_t<A>, transformed_t<B>) -> transformed_t<A+B> | ||
| 52 | // add(product_t<K1>, product_t<K2>) -> product_t<K1+K2> | ||
| 53 | // | ||
| 54 | // Finally, we expose some additional fast-transform optimization paths. | ||
| 55 | // extend_to only operates on transformed_t<unit_scale>; the others are scale-generic. | ||
| 56 | // downsample is also defined on product_t (halving before finish saves inverse-transform work). | ||
| 57 | // extend_to build (if empty) or grow a transform to size m by repeated doubling; feed the same coefficients (sz <= 2m) every time | ||
| 58 | // (each doubling step reads only the coefficients that fit, so zero-padded buffers are fine: | ||
| 59 | // coefficients past twice the existing transform's size must be zero, or it couldn't be a prefix) | ||
| 60 | // downsample compute the half-sized transform/product of just the even (odd = false) or odd terms of the input | ||
| 61 | // upsample size n (n >= 2) transform/product of the input spread as evens (odd = false) or odds (odd = true); inverse of downsample | ||
| 62 | // negate_arg size n transform of A(-x) | ||
| 63 | // of_reverse size n transform of A(1/x) mod x^n - 1, i.e. of the cyclically reversed coefficients a[-k mod n] | ||
| 64 | // | ||
| 65 | // Optional (exact-ish engines only, currently ntt and real): read a product out without an inverse transform. | ||
| 66 | // dot(product_t<K>, transformed_t<A>, int n) -> value_type | ||
| 67 | // sum_k finish(p)[k] * b[k] over k < n, where b is the coefficient sequence (mod x^n - 1) of the | ||
| 68 | // transform; Parseval, so it costs O(n) instead of an inverse transform. | ||
| 69 | // <P, B> with B = of_reverse(T) is the constant term of P * T. | ||
| 70 | // Not offered by split/crt: their pointwise data isn't exact enough to sum before reduction. | ||
| 71 | template <typename E> | ||
| 72 | concept engine = requires( | ||
| 73 | std::span<const typename E::value_type> in, | ||
| 74 | std::span<typename E::value_type> out, | ||
| 75 | typename E::transformed& t, | ||
| 76 | const typename E::transformed& ct, | ||
| 77 | typename E::product& p, | ||
| 78 | const typename E::product& cp, | ||
| 79 | int n | ||
| 80 | ) { | ||
| 81 | typename E::value_type; | ||
| 82 | { E::transform(in, n) } -> std::same_as<typename E::transformed>; | ||
| 83 | { ct.size() } -> std::same_as<int>; | ||
| 84 | E::extend_to(t, n, in); | ||
| 85 | { E::downsample(ct, n, false) } -> std::same_as<typename E::transformed>; | ||
| 86 | { E::downsample(cp, n, false) } -> std::same_as<typename E::product>; | ||
| 87 | { E::upsample(ct, n, false) } -> std::same_as<typename E::transformed>; | ||
| 88 | { E::upsample(cp, n, false) } -> std::same_as<typename E::product>; | ||
| 89 | { E::negate_arg(ct, n) } -> std::same_as<typename E::transformed>; | ||
| 90 | { E::of_reverse(ct, n) } -> std::same_as<typename E::transformed>; | ||
| 91 | { E::mul(ct, ct, n) } -> std::same_as<typename E::product>; | ||
| 92 | { E::sq(ct, n) } -> std::same_as<typename E::product>; | ||
| 93 | { E::mul2(ct, ct, ct, ct, n) } -> std::same_as<typename E::template product_t<2 * E::unit_scale>>; | ||
| 94 | E::finish(std::move(p), out); | ||
| 95 | E::finish(std::move(p), out, add_op{}); | ||
| 96 | E::finish(E::add(std::move(p), std::move(p)), out); | ||
| 97 | { E::add(E::transform(in, n), ct) } -> std::same_as<typename E::template transformed_t<2 * E::unit_scale>>; | ||
| 98 | { E::add(std::move(p), std::move(p)) } -> std::same_as<typename E::template product_t<2 * E::unit_scale>>; | ||
| 99 | requires std::same_as<std::remove_cvref_t<decltype(E::commutative)>, bool>; | ||
| 100 | requires std::same_as<std::remove_cvref_t<decltype(E::unit_scale)>, int>; | ||
| 101 | }; | ||
| 102 | |||
| 103 | // Constrains two engine-parameterized value types to share the same engine. | ||
| 104 | template <typename A, typename B> | ||
| 105 | concept same_engine = std::same_as<typename A::engine_t, typename B::engine_t>; | ||
| 106 | |||
| 107 | // short spelling for E::transformed at use sites | ||
| 108 | template <engine E> using transformed = typename E::transformed; | ||
| 109 | |||
| 110 | } // namespace wala::fft | ||
| 111 |