convolver_messages.hpp

Go to the documentation of this file.

#pragma once
 
#include <pulp/runtime/slot.hpp>
#include <pulp/signal/fft.hpp>
 
#include <algorithm>
#include <atomic>
#include <complex>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <span>
#include <vector>
 
namespace pulp::signal {

template <typename SampleType = float> class ConvolverSpectrumArenaT {
  public:
    using value_type = std::complex<SampleType>;

    static constexpr std::size_t stride_for(std::size_t bins) noexcept {
        constexpr std::size_t per_line = 64 / sizeof(value_type);
        return (bins + per_line - 1) / per_line * per_line;
    }
 
    ConvolverSpectrumArenaT() = default;
    ConvolverSpectrumArenaT(std::size_t count, std::size_t bins) {
        assign(count, bins);
    }
 
    // The alignment offset is relative to the storage address, which a copy
    // would change; a move keeps the buffer, so moving is safe.
    ConvolverSpectrumArenaT(const ConvolverSpectrumArenaT&) = delete;
    ConvolverSpectrumArenaT& operator=(const ConvolverSpectrumArenaT&) = delete;
    ConvolverSpectrumArenaT(ConvolverSpectrumArenaT&&) noexcept = default;
    ConvolverSpectrumArenaT& operator=(ConvolverSpectrumArenaT&&) noexcept = default;

    void assign(std::size_t count, std::size_t bins) {
        constexpr std::size_t per_line = 64 / sizeof(value_type);
        count_ = count;
        bins_ = bins;
        stride_ = stride_for(bins);
        storage_.assign(count_ * stride_ + per_line, value_type{});
        const auto address = reinterpret_cast<std::uintptr_t>(storage_.data());
        offset_ = ((64u - address % 64u) % 64u) / sizeof(value_type);
    }
 
    std::size_t size() const noexcept {
        return count_;
    }
    bool empty() const noexcept {
        return count_ == 0;
    }
    std::size_t bins() const noexcept {
        return bins_;
    }
    std::size_t stride() const noexcept {
        return stride_;
    }

    std::size_t footprint_bytes() const noexcept {
        return storage_.size() * sizeof(value_type);
    }
 
    std::span<value_type> operator[](std::size_t i) noexcept {
        return {storage_.data() + offset_ + i * stride_, bins_};
    }
    std::span<const value_type> operator[](std::size_t i) const noexcept {
        return {storage_.data() + offset_ + i * stride_, bins_};
    }

    void fill_zero() noexcept {
        std::fill(storage_.begin(), storage_.end(), value_type{});
    }
 
  private:
    std::vector<value_type> storage_;
    std::size_t offset_ = 0, count_ = 0, bins_ = 0, stride_ = 0;
};

template <typename SampleType = float>
struct ConvolverInputHistoryT {
    int block_size = 0;
    int fft_size = 0;

    ConvolverSpectrumArenaT<SampleType> spectra;
    std::vector<std::complex<SampleType>> overlap;
    std::size_t write_pos = 0;
};
 
using ConvolverInputHistory = ConvolverInputHistoryT<float>;
using ConvolverInputHistory64 = ConvolverInputHistoryT<double>;

template <typename SampleType = float>
struct ConvolverIrStateT {
    int block_size = 0;
    int fft_size = 0;
    std::size_t num_partitions = 0;
 
    std::unique_ptr<FftT<SampleType>> fft;
    ConvolverSpectrumArenaT<SampleType> ir_spectra;
    std::vector<std::complex<SampleType>> accum;
    std::unique_ptr<ConvolverInputHistoryT<SampleType>> history;
};
 
using ConvolverIrState = ConvolverIrStateT<float>;
using ConvolverIrState64 = ConvolverIrStateT<double>;
 
namespace detail {

constexpr std::size_t convolver_block_size(std::size_t block_size) noexcept {
    std::size_t pot = 1;
    while (pot < block_size && pot <= (static_cast<std::size_t>(-1) >> 1))
        pot <<= 1;
    return pot;
}

template <typename SampleType = float>
inline std::unique_ptr<ConvolverIrStateT<SampleType>>
build_convolver_ir_state(const SampleType* ir,
                         std::size_t ir_length,
                         std::size_t block_size) {
    if (ir == nullptr || ir_length == 0)
        return nullptr;
 
    block_size = convolver_block_size(block_size);
 
    auto state = std::make_unique<ConvolverIrStateT<SampleType>>();
    state->block_size = static_cast<int>(block_size);
    state->fft_size = state->block_size * 2;
    state->fft = std::make_unique<FftT<SampleType>>(state->fft_size);
    state->num_partitions = (ir_length + block_size - 1) / block_size;
 
    const auto bins = static_cast<std::size_t>(state->fft_size / 2 + 1);
    state->ir_spectra.assign(state->num_partitions, bins);
    std::vector<std::complex<SampleType>> padded(
        state->fft_size, {SampleType{0.0f}, SampleType{0.0f}});
 
    for (std::size_t p = 0; p < state->num_partitions; ++p) {
        const std::size_t offset = p * block_size;
        const std::size_t count = std::min(block_size, ir_length - offset);
 
        std::fill(padded.begin(), padded.end(),
                  std::complex<SampleType>{SampleType{0.0f}, SampleType{0.0f}});
        for (std::size_t i = 0; i < count; ++i)
            padded[i] = {ir[offset + i], SampleType{0.0f}};
 
        state->fft->forward(padded.data());
        std::copy_n(padded.begin(), bins, state->ir_spectra[p].begin());
    }
 
    state->accum.assign(state->fft_size,
                        {SampleType{0.0f}, SampleType{0.0f}});
    state->history = std::make_unique<ConvolverInputHistoryT<SampleType>>();
    state->history->block_size = state->block_size;
    state->history->fft_size = state->fft_size;
    state->history->overlap.assign(state->fft_size, {0.0f, 0.0f});
    state->history->spectra.assign(state->num_partitions, bins);
 
    return state;
}

template <typename SampleType = float>
inline void migrate_input_history(ConvolverInputHistoryT<SampleType>& from,
                                  ConvolverInputHistoryT<SampleType>& to) {
    if (from.fft_size != to.fft_size || from.block_size != to.block_size)
        return;
    const std::size_t old_capacity = from.spectra.size();
    const std::size_t new_capacity = to.spectra.size();
    if (old_capacity == 0 || new_capacity == 0) return;
    const std::size_t carried = std::min(old_capacity, new_capacity);
    // Both rings share the FFT size, so their spectra have the same length.
    // Each `a` maps to a distinct slot in each of two distinct rings, so a copy
    // never overwrites a slot a later iteration reads.
    for (std::size_t a = 0; a < carried; ++a) {
        const std::size_t oldi = (from.write_pos + old_capacity - 1 - a) % old_capacity;
        const std::size_t newi = (new_capacity - 1 - a) % new_capacity;
        const auto source = from.spectra[oldi];
        std::copy(source.begin(), source.end(), to.spectra[newi].begin());
    }
    // The overlap buffer's lower half holds the previous block's samples (the
    // overlap the next block's FFT needs); it is IR-independent, so carry it too.
    to.overlap.swap(from.overlap);
    to.write_pos = 0;
}
 
} // namespace detail

template <typename SampleType = float>
class ConvolverIrSwapperT {
public:
    ConvolverIrSwapperT() = default;
 
    ConvolverIrSwapperT(const ConvolverIrSwapperT&) = delete;
    ConvolverIrSwapperT& operator=(const ConvolverIrSwapperT&) = delete;

    bool stage_ir(const SampleType* ir,
                  std::size_t ir_length,
                  std::size_t block_size) {
        auto next = detail::build_convolver_ir_state(ir, ir_length, block_size);
        if (!next)
            return false;
        return stage_ir(std::move(next));
    }

    bool stage_ir(std::unique_ptr<ConvolverIrStateT<SampleType>> next) {
        return handoff_.publish(std::move(next));
    }

    std::unique_ptr<ConvolverIrStateT<SampleType>> try_consume() {
        return handoff_.try_consume();
    }

    [[nodiscard]] bool retire(std::unique_ptr<ConvolverIrStateT<SampleType>>& displaced) {
        return handoff_.retire(displaced);
    }

    bool has_retire_capacity() const { return handoff_.has_retire_capacity(); }

    bool has_retire_capacity(std::size_t n) const { return handoff_.has_retire_capacity(n); }

    std::size_t drain_old() { return handoff_.drain_retired(); }

    std::unique_ptr<ConvolverIrStateT<SampleType>> drain_old_one() {
        return handoff_.drain_retired_one();
    }

    bool has_pending() const { return handoff_.has_pending(); }

    bool has_retired() const { return handoff_.has_retired(); }

    static constexpr std::size_t retire_capacity() { return kRetireRingCapacity; }
 
private:
    static constexpr std::size_t kRetireRingCapacity = 8;
 
    // The publish/consume/retire machinery is `runtime::Handoff` — the shared
    // real-time ownership-transfer primitive. It owns the pending slot, the
    // retire ring, and teardown reclamation, so this class only adds the
    // IR-specific `stage_ir` build step and the `drain_old` naming its callers
    // already use.
    pulp::runtime::Handoff<ConvolverIrStateT<SampleType>, kRetireRingCapacity> handoff_;
};
 
using ConvolverIrSwapper = ConvolverIrSwapperT<float>;
using ConvolverIrSwapper64 = ConvolverIrSwapperT<double>;
 
} // namespace pulp::signal