fft.cpp

#include "fft.hpp"
#ifdef _WIN32
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#endif
#include <memory>
#include <mutex>

#include <cmath>
#include <algorithm>
#include <limits>
#include <stdexcept>

namespace lvm {

namespace {
constexpr double kPi = 3.14159265358979323846;
constexpr std::size_t kButterflyChunk = 16384;
constexpr std::size_t kPlanCacheBytes = 64 * 1024 * 1024;
struct BluesteinPlan {
    std::size_t n = 0;
    std::vector<std::complex<double>> chirp, kernel;
};
std::atomic<std::size_t> retained_plan_bytes{0};
struct BluesteinCache;
struct PlanRegistry {std::mutex mutex;BluesteinCache* first=nullptr;};
// Process lifetime: native FLS/TLS callbacks can run during static teardown.
PlanRegistry& plan_registry(){static auto* registry=new PlanRegistry;return *registry;}
struct BluesteinCache {
    BluesteinPlan plan;
    std::size_t accounted=0;
    BluesteinCache* previous=nullptr;
    BluesteinCache* next=nullptr;
    bool active=false;
    BluesteinCache() {
        auto& registry=plan_registry();std::lock_guard<std::mutex> lock(registry.mutex);
        next=registry.first;if(next)next->previous=this;registry.first=this;
    }
    void release() {
        retained_plan_bytes.fetch_sub(accounted,std::memory_order_relaxed);
        accounted=0;plan={};
    }
    ~BluesteinCache() {
        auto& registry=plan_registry();std::lock_guard<std::mutex> lock(registry.mutex);
        if(previous)previous->next=next;else registry.first=next;
        if(next)next->previous=previous;
        retained_plan_bytes.fetch_sub(accounted,std::memory_order_relaxed);
    }
};
struct PlanUse {
    BluesteinCache& cache;
    explicit PlanUse(BluesteinCache& c):cache(c){std::lock_guard<std::mutex> lock(plan_registry().mutex);cache.active=true;}
    ~PlanUse(){std::lock_guard<std::mutex> lock(plan_registry().mutex);cache.active=false;}
};
BluesteinCache& thread_plan_cache() {
#ifdef _WIN32
    // Native thread/fiber cleanup avoids the MinGW emulated TLS destruction
    // failure reproduced when short-lived FFT threads terminate concurrently.
    struct Slot {
        DWORD index=FlsAlloc([](void* value){delete static_cast<BluesteinCache*>(value);});
        Slot() {if(index==FLS_OUT_OF_INDEXES)throw std::bad_alloc();}
        ~Slot() {FlsFree(index);}
    };
    static Slot slot;
    auto* cache=static_cast<BluesteinCache*>(FlsGetValue(slot.index));
    if(!cache) {
        auto owned=std::make_unique<BluesteinCache>();
        if(!FlsSetValue(slot.index,owned.get()))throw std::bad_alloc();
        cache=owned.release();
    }
    return *cache;
#else
    static thread_local BluesteinCache cache;
    return cache;
#endif
}
void check_cancel(const std::atomic<bool>* cancel) {
    if (cancel && cancel->load(std::memory_order_relaxed)) throw std::runtime_error("Operation cancelled.");
}
inline void butterfly_range(std::vector<std::complex<double>>& a, std::size_t i,
                            std::size_t half, std::size_t first, std::size_t last,
                            std::complex<double>& w, const std::complex<double>& wlen) {
    for (std::size_t k = first; k < last; ++k) {
        const std::complex<double> u = a[i + k];
        const std::complex<double> v = a[i + k + half] * w;
        a[i + k] = u + v;
        a[i + k + half] = u - v;
        w *= wlen;
    }
}
}

void fft_radix2(std::vector<std::complex<double>>& a, bool inverse, const std::atomic<bool>* cancel) {
    check_cancel(cancel);
    const std::size_t n = a.size();
    if (n <= 1) return;
    if ((n & (n - 1)) != 0) throw std::invalid_argument("Radix-2 FFT requires a power-of-two size.");

    // Bit-reversal permutation.
    for (std::size_t i = 1, j = 0; i < n; ++i) {
        if ((i & 0xFFFF) == 0) check_cancel(cancel);
        std::size_t bit = n >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j ^= bit;
        if (i < j) std::swap(a[i], a[j]);
    }

    for (std::size_t len = 2; len <= n; len <<= 1) {
        check_cancel(cancel);
        const double ang = 2.0 * kPi / static_cast<double>(len) * (inverse ? 1.0 : -1.0);
        const std::complex<double> wlen(std::cos(ang), std::sin(ang));
        const std::size_t half = len / 2;
        if (half <= kButterflyChunk) {
            // Group short blocks: k=0 in each block must not cause a cancel
            // call (and impede optimisation) in the arithmetic loop.
            const std::size_t span = std::min(n, 2 * kButterflyChunk);
            for (std::size_t first = 0; first < n; first += span) {
                check_cancel(cancel);
                const std::size_t last = first + span;
                for (std::size_t i = first; i < last; i += len) {
                    std::complex<double> w(1.0, 0.0);
                    butterfly_range(a, i, half, 0, half, w, wlen);
                }
            }
        } else {
            for (std::size_t i = 0; i < n; i += len) {
                std::complex<double> w(1.0, 0.0);
                for (std::size_t first = 0; first < half; first += kButterflyChunk) {
                    check_cancel(cancel);
                    // Keep w across chunks: restarting it changes rounding.
                    butterfly_range(a, i, half, first, first + kButterflyChunk, w, wlen);
                }
            }
        }
    }

    if (inverse) {
        for (std::size_t first = 0; first < n; first += kButterflyChunk) {
            check_cancel(cancel);
            const std::size_t last = std::min(n, first + kButterflyChunk);
            for (std::size_t i = first; i < last; ++i) a[i] /= static_cast<double>(n);
        }
    }
    check_cancel(cancel);
}

std::size_t fft_cached_plan_bytes() {return retained_plan_bytes.load(std::memory_order_relaxed);}
void trim_fft_cached_plans(std::size_t bytes) {
    auto& registry=plan_registry();std::lock_guard<std::mutex> lock(registry.mutex);
    for(auto* cache=registry.first;cache && retained_plan_bytes.load(std::memory_order_relaxed)>bytes;cache=cache->next)
        if(!cache->active)cache->release();
}

std::vector<std::complex<double>> dft(const std::vector<std::complex<double>>& in, const std::atomic<bool>* cancel) {
    const std::size_t n = in.size();
    check_cancel(cancel);
    if (n == 0) return {};

    // Power-of-two fast path.
    if ((n & (n - 1)) == 0) {
        std::vector<std::complex<double>> a = in;
        fft_radix2(a, false, cancel);
        return a;
    }

    // Bluestein's algorithm for arbitrary length.
    if (n > std::numeric_limits<std::size_t>::max() / 2) throw std::length_error("FFT is too large.");
    std::size_t m = 1;
    while (m < 2 * n - 1) {
        if (m > std::numeric_limits<std::size_t>::max() / 2) throw std::length_error("FFT is too large.");
        m <<= 1;
    }

    // One bounded plan per worker/thread. Publish only a complete plan; a
    // cancelled or failed rebuild must never leave reusable partial data.
    auto& cache_owner=thread_plan_cache();
    PlanUse use(cache_owner);
    auto& cached=cache_owner.plan;
    BluesteinPlan temporary;
    const BluesteinPlan* plan = &cached;
    if (cached.n != n) {
        cache_owner.release();
        temporary.chirp.resize(n);
        for (std::size_t k = 0; k < n; ++k) {
            if ((k & 0x3FFF) == 0) check_cancel(cancel);
            // angle = -pi * k^2 / n, reduced mod 2n to keep precision for large k.
            const unsigned long long sq = (static_cast<unsigned long long>(k) * k) % (2ULL * n);
            const double ang = -kPi * static_cast<double>(sq) / static_cast<double>(n);
            temporary.chirp[k] = std::complex<double>(std::cos(ang), std::sin(ang));
        }
        temporary.kernel.assign(m, {0.0, 0.0});
        for (std::size_t k = 0; k < n; ++k) {
            if ((k & 0x3FFF) == 0) check_cancel(cancel);
            temporary.kernel[k] = std::conj(temporary.chirp[k]);
            if (k != 0) temporary.kernel[m - k] = std::conj(temporary.chirp[k]);
        }
        fft_radix2(temporary.kernel, false, cancel);
        temporary.n = n;
        check_cancel(cancel);
        if (m <= kPlanCacheBytes / sizeof(std::complex<double>) &&
            n <= kPlanCacheBytes / sizeof(std::complex<double>) - m) {
            cached = std::move(temporary);
            cache_owner.accounted=(cached.chirp.capacity()+cached.kernel.capacity())*sizeof(std::complex<double>);
            retained_plan_bytes.fetch_add(cache_owner.accounted,std::memory_order_relaxed);
        } else {
            plan = &temporary; // Large plans live only for this call.
        }
    }
    const auto& chirp = plan->chirp;
    const auto& b = plan->kernel;

    std::vector<std::complex<double>> a(m, {0.0, 0.0});
    for (std::size_t k = 0; k < n; ++k) {
        if ((k & 0x3FFF) == 0) check_cancel(cancel);
        a[k] = in[k] * chirp[k];
    }

    fft_radix2(a, false, cancel);
    for (std::size_t first = 0; first < m; first += kButterflyChunk) {
        check_cancel(cancel);
        const std::size_t last = std::min(m, first + kButterflyChunk);
        for (std::size_t i = first; i < last; ++i) a[i] *= b[i];
    }
    fft_radix2(a, true, cancel);

    std::vector<std::complex<double>> out(n);
    for (std::size_t first = 0; first < n; first += kButterflyChunk) {
        check_cancel(cancel);
        const std::size_t last = std::min(n, first + kButterflyChunk);
        for (std::size_t k = first; k < last; ++k) out[k] = a[k] * chirp[k];
    }
    check_cancel(cancel);
    return out;
}

}  // namespace lvm