gpu_compute.hpp

Go to the documentation of this file.

#pragma once
 
#if defined(PULP_GPU_RUNTIME)
// This build's device layer is gpu-runtime's (tools/cmake/PulpGpuRuntime.cmake):
// these are its types and functions under their Pulp names, so a Pulp object
// and a gpu-runtime object are one and the same.
#include <gpu_runtime/gpu_compute.hpp>
// The Pulp names of the types it uses, as the header always provided.
#include <pulp/render/gpu_surface.hpp>
#ifdef PULP_BENCHMARK
#include <pulp/render/bench/perf_counters.hpp>
#endif
 
namespace pulp::render {
using gpu_runtime::GpuCompute;
} // namespace pulp::render
#else // !PULP_GPU_RUNTIME
 
#include <pulp/render/gpu_surface.hpp>
 
#include <atomic>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <vector>
 
#ifdef PULP_BENCHMARK
#include <pulp/render/bench/perf_counters.hpp>
#endif
 
namespace pulp::render {

class GpuCompute {
public:
    virtual ~GpuCompute() = default;

    virtual bool initialize_from_surface(GpuSurface& surface) = 0;

    virtual bool initialize_standalone() = 0;

    virtual bool initialize_from_device(void* wgpu_device) = 0;
 
    bool is_initialized() const {
        return initialized_.load(std::memory_order_acquire);
    }

    virtual bool device_lost() const = 0;

    virtual std::uint64_t device_generation() const {
        return 0;
    }

    virtual void notify_device_lost() = 0;
 
    virtual void notify_device_lost(std::uint64_t generation) {
        if (generation == device_generation())
            notify_device_lost();
    }

    virtual void post_device_lost(std::uint64_t generation) {
        notify_device_lost(generation);
    }
 
    // ── Compute operations ──────────────────────────────────────────────

    virtual bool compute_magnitude(const float* complex_pairs, float* magnitudes,
                                   uint32_t num_bins) = 0;

    virtual bool complex_multiply(const float* a, const float* b, float* result,
                                  uint32_t count) = 0;

    virtual bool batch_magnitude(const float* complex_frames, float* magnitude_frames,
                                 uint32_t bins_per_frame, uint32_t num_frames) = 0;
 
    // ── FFT ──────────────────────────────────────────────────────────────
    //
    // GPU radix-2 FFT (Stockham auto-sort, no bit-reversal pass). Operates on
    // interleaved complex data [re0, im0, re1, im1, ...] of length `n` complex
    // values (2*n floats in and out). `n` must be a power of two and small
    // enough that 2*n*sizeof(float) fits the device storage-buffer limit.
    // Not for the audio callback — these build/cache a per-`n` plan and block
    // on readback. Validated against a direct DFT and the CPU FFT
    // (magnitude parity + round-trip identity); see test/test_gpu_compute.cpp.
    //
    // Thread-safety: like the rest of GpuCompute, a single instance is NOT
    // safe for concurrent use — call all methods from one thread (e.g. the
    // GPU worker). The plan/pipeline caches are unsynchronized by design.

    virtual bool fft_forward(const float* complex_in, float* complex_out, uint32_t n) = 0;

    virtual bool fft_inverse(const float* complex_in, float* complex_out, uint32_t n) = 0;

    virtual bool fft_forward_timed(const float* complex_in, float* complex_out,
                                   uint32_t n, double* gpu_compute_us) = 0;
 
    // ── GPU-resident convolution ───────────────────────────────────────────
    //
    // Fused FFT convolution that keeps intermediates GPU-resident: forward FFT
    // → complex-multiply by a resident IR spectrum → inverse FFT (1/N
    // normalized), all encoded into ONE command buffer with ONE readback
    // (vs three readbacks for fft_forward + complex_multiply + fft_inverse).
    // Not real-time-safe (still blocks on the single readback) but the
    // amortizable building block toward a real-time GPU convolver.

    virtual bool prepare_convolution(uint32_t n, const float* ir_spec) = 0;

    virtual bool convolve(const float* in_complex, float* out_complex, uint32_t n) = 0;

    virtual bool prepare_convolution_batch(uint32_t n, const float* ir_spec,
                                           uint32_t batch) = 0;
    virtual bool convolve_batch(const float* in_complex, float* out_complex,
                                uint32_t n, uint32_t batch) = 0;
 
    // ── Multi-IR convolution (one input → many distinct IRs → stereo) ───────
    //
    // The mass-parallel regime: convolve ONE input block against `num_ir`
    // DISTINCT IR spectra ("rooms"/"taps") in ONE submit, then reduce the
    // num_ir results to a stereo pair on the GPU with per-room pan weights —
    // so only one stereo block (2*n floats) is read back regardless of num_ir.
    // The forward FFT runs once and is shared across all rooms; each room gets
    // its own complex-multiply and inverse FFT, all batched. On the CPU this is
    // num_ir independent convolutions; on the GPU it is one batched job — the
    // structural win that scales past real time on the CPU while the GPU holds.
    // Not real-time-safe (blocks on the single readback); the amortizable
    // building block for a GPU multi-room convolution reverb.

    virtual bool prepare_multi_convolution(uint32_t n, const float* ir_specs,
                                           uint32_t num_ir) = 0;

    virtual bool multi_convolve(const float* in_complex, const float* pan_l,
                                const float* pan_r, float* out_lr, uint32_t n,
                                uint32_t num_ir) = 0;

    virtual bool multi_convolve_timed(const float* in_complex, const float* pan_l,
                                      const float* pan_r, float* out_lr, uint32_t n,
                                      uint32_t num_ir, double* gpu_compute_us) = 0;
    // ── Partitioned-FDL multi-convolution (roofline #1) ──────────────────────
    //
    // The same one-input → many-IRs → stereo result, but with the IR split into
    // `num_part` block-size partitions so the FFT stays at the small fixed block
    // size (n = 2*block) instead of growing to next_pow2(block + IR_len). The
    // convolution becomes a sum of spectral products over a delay line of recent
    // input spectra — the algorithm signal::PartitionedConvolver uses on the CPU.
    // This trades one giant FFT for P small-FFT-equivalent MAC rows, which is the
    // removable waste multi_convolve_timed measured. Not real-time-safe (blocks
    // on the readback); the amortizable building block for a real-time GPU
    // multi-room convolver.

    virtual bool prepare_multi_fdl(uint32_t n, const float* ir_part_specs,
                                   uint32_t num_ir, uint32_t num_part) = 0;

    virtual bool multi_fdl_convolve(const float* in_block, const float* pan_l,
                                    const float* pan_r, float* out_block,
                                    uint32_t n, uint32_t num_ir) = 0;
 
    // ── Multi-layer spectral stack (N resident layer-spectra → one frame) ───
    //
    // The spectral analogue of the multi-IR regime: keep `num_layers` frozen
    // spectral frames (each a magnitude array + a persistent phase array)
    // RESIDENT on the GPU, and per hop process EVERY layer's bins in ONE batched
    // submit — advance each layer's phase by its nominal per-hop frequency, blur
    // (smear) each layer's magnitude across frequency, then weighted-sum all
    // layers into one combined spectrum and inverse-FFT it to a single real
    // frame. Only that one frame (2n floats) is read back per hop, regardless of
    // how many layers there are. On the CPU this is num_layers × (per-bin smear
    // + transcendentals), serial; on the GPU it is one batch parallel across all
    // n bins — the structural win that scales past real time on the CPU at high
    // num_layers while the GPU holds. Not real-time-safe (blocks on the single
    // readback); driven from the GPU audio worker / offline.

    virtual bool prepare_spectral_stack(uint32_t n, uint32_t hop,
                                        uint32_t num_layers) = 0;

    virtual bool spectral_stack_set_layer(uint32_t layer, const float* mag,
                                          const float* phase, uint32_t n,
                                          uint32_t num_layers) = 0;

    virtual bool spectral_stack_render(const float* layer_weights,
                                       uint32_t num_layers, float smear,
                                       float jitter, uint32_t rng_seed,
                                       float* frame_out, uint32_t n) = 0;
 
    // ── Linear algebra ─────────────────────────────────────────────────────

    virtual bool matmul(const float* a, const float* b, float* c,
                        uint32_t m, uint32_t k, uint32_t n) = 0;
 
    // ── Synthesis ──────────────────────────────────────────────────────────

    virtual bool additive_synth(const float* partials, float* out,
                                uint32_t num_partials, uint32_t num_samples,
                                float sample_rate, float t0_samples) = 0;

    virtual bool modal_strike(const float* modes, float* out, uint32_t num_modes,
                              uint32_t num_samples, float sample_rate, float t0_samples) = 0;

    virtual bool granular_cloud(const float* grains, const float* source, float* out,
                                uint32_t num_grains, uint32_t source_len,
                                uint32_t num_samples) = 0;
 
    // ── Neural inference ────────────────────────────────────────────────────

    virtual bool dense_tanh(const float* input, const float* weights, const float* bias,
                            float* output, uint32_t in_dim, uint32_t out_dim) = 0;
 
    // ── Fused conv-stack inference (WaveNet-style neural amp) ───────────────
    //
    // A stack of gated, dilated, causal 1-D conv layers — the architecture
    // behind neural-amp captures. The whole block's samples are computed in
    // PARALLEL (the network is feedforward: output[t] depends on past inputs,
    // not past outputs) and the entire forward pass runs as successive compute
    // passes in ONE submit with device-resident activations — so the CPU<->GPU
    // round-trip is paid ONCE per block, not per layer or per sample. That is
    // the regime where GPU neural inference beats the CPU, and the win grows
    // with model size (big captures the CPU can't run in real time stay smooth
    // on the GPU). Per-sample/per-layer dispatch loses badly and is never used.
    //
    // Per-layer weight layout (row-major f32, concatenated for all layers, then
    // the two global blocks): conv W [2C*C*K] + bias [2C], residual W [C*C] +
    // bias [C], skip W [C*C] + bias [C]; then once: input W [C] + bias [C], head
    // W [C] + bias [1]. `channels` = C, `kernel` = K, `dilations` has one entry
    // per layer. Not real-time-safe (blocks on readback); run on the worker.

    virtual bool prepare_conv_stack(uint32_t channels, uint32_t kernel,
                                    const uint32_t* dilations, uint32_t num_layers,
                                    const float* weights, uint32_t weights_len,
                                    uint32_t block_size, float head_scale) = 0;

    virtual bool conv_stack_forward(const float* in_block, float* out_block,
                                    uint32_t block_size) = 0;
 
    // ── Fused conditioned WaveNet inference (gated dilated conv sequence) ────
    //
    // A conditioned WaveNet forward — a sequence of layer-arrays of gated,
    // dilated, causal 1-D conv layers (van den Oord et al.) — reproduced
    // GPU-resident and block-parallel. Unlike conv_stack (a single gated layer
    // stack), this models the full shape: a sequence of layer-arrays, each with
    // an input rechannel (1x1), dilated causal conv layers that mix in a
    // separate condition signal (here the raw mono input), Tanh (or gated
    // tanh·sigmoid) activation, a 1x1 residual, a per-array head accumulator,
    // and a head rechannel — with the layer output and the head accumulator
    // chained array→array. The whole block's samples compute in PARALLEL (the
    // net is feedforward), every pass encodes into ONE submit over
    // device-resident activations, and only the final mono block is read back —
    // the round-trip is paid once per block. Dilation history carries across
    // calls so streaming blocks are continuous. Not real-time-safe (blocks on
    // the readback); run on the GPU worker / offline. It is the GPU counterpart
    // to a streaming CPU WaveNet and drives neural waveshaping / amp inference.
    //
    // Weight layout (flat f32 blob). Per layer-array, walking the blob:
    // rechannel W [channels*input_size] (no bias); then per layer: conv W
    // [Z*channels*kernel] + bias [Z], input_mixin W [Z*condition_size] (no bias),
    // layer1x1 W [channels*channels] + bias [channels]; then head_rechannel W
    // [head_size*channels] (+ bias [head_size] if head_bias). After all arrays:
    // one trailing head_scale scalar. Z = gated ? 2*channels : channels.
    // Activation is Tanh; the gate is sigmoid.

    struct WavenetLayerArraySpec {
        uint32_t input_size = 1;
        uint32_t condition_size = 1;
        uint32_t channels = 1;
        uint32_t kernel = 1;
        uint32_t head_size = 1;
        uint32_t gated = 0;
        uint32_t head_bias = 0;
        const uint32_t* dilations = nullptr;
        uint32_t num_layers = 0;
    };

    virtual bool prepare_wavenet(const WavenetLayerArraySpec* arrays, uint32_t num_arrays,
                                 const float* weights, uint32_t weights_len,
                                 uint32_t block_size, float head_scale,
                                 uint32_t instance = 0) = 0;

    virtual bool wavenet_forward(const float* in_block, float* out_block,
                                 uint32_t block_size, uint32_t instance = 0) = 0;

    virtual bool prepare_wavenet_group(const WavenetLayerArraySpec* /*arrays*/,
                                       uint32_t /*num_arrays*/, const float* const* /*weights*/,
                                       uint32_t /*weights_len*/, const float* /*head_scales*/,
                                       uint32_t /*count*/, uint32_t /*block_size*/,
                                       uint32_t /*group*/) {
        return false;
    }

    virtual bool wavenet_forward_group(const float* const* /*in_blocks*/,
                                       float* const* /*out_blocks*/, uint32_t /*count*/,
                                       uint32_t /*block_size*/, uint32_t /*group*/,
                                       double* gpu_busy_us = nullptr) {
        if (gpu_busy_us)
            *gpu_busy_us = -1.0;
        return false;
    }

    virtual bool wavenet_forward_batch(const float* const* in_blocks, float* const* out_blocks,
                                       const uint32_t* instances, uint32_t count,
                                       uint32_t block_size, double* gpu_busy_us = nullptr) {
        if (gpu_busy_us)
            *gpu_busy_us = -1.0;
        if (!in_blocks || !out_blocks || !instances || count == 0)
            return false;
        for (uint32_t i = 0; i < count; ++i)
            for (uint32_t j = i + 1; j < count; ++j)
                if (instances[i] == instances[j])
                    return false;
        for (uint32_t i = 0; i < count; ++i)
            if (!wavenet_forward(in_blocks[i], out_blocks[i], block_size, instances[i]))
                return false;
        return true;
    }
 
    // ── Asynchronous readback (compute_magnitude + convolve_batch) ──────────
    //
    // SCOPE FIRST, so nobody reads the contract below as a general promise:
    // exactly TWO operations have an async variant — compute_magnitude_async()
    // and convolve_batch_async(). fft_forward, convolve, multi_convolve,
    // multi_fdl_convolve, spectral_stack_render, matmul, conv_stack_forward and
    // wavenet_forward ALL still block. This is not a non-blocking mode for
    // GpuCompute; it is two non-blocking operations plus the completion
    // machinery the rest of the ops reuse when they are converted.
    //
    // Every blocking operation pumps the device event queue until the GPU map
    // resolves. That is only viable where the caller owns a thread it may stall
    // (the GPU worker, an offline render). On a single-threaded event loop it is
    // fatal — spinning starves the very loop that would resolve the map. The
    // async path submits the dispatch, returns immediately, and hands results
    // back from poll_readbacks(), so the caller decides when completions run.
    //
    // In the browser (Emscripten/emdawnwebgpu) the blocking ops are compiled out
    // or hard-fail: the async path is the ONLY way to get bytes back from the
    // GPU there.
    //
    // Contract (of the async path):
    // - An accepted request (non-zero id) completes EXACTLY ONCE, from
    //   poll_readbacks() or from ~GpuCompute(). It never hangs: a request whose
    //   map has not resolved by its deadline completes as `Expired`, which the
    //   caller routes to its miss policy (gpu_audio::MissPolicy) exactly as it
    //   would route a late block.
    // - Completions fire in submission order.
    // - `dest` must stay valid until the completion fires. It is written only
    //   on `Success`; a map that lands after an `Expired` completion is
    //   discarded without touching `dest`.
    // - Not real-time-safe: submit and poll from the GPU worker, never from the
    //   audio callback. The audio thread's contract (never wait on, allocate
    //   for, or synchronize with the GPU) is upheld by the transport, not here.
 
    enum class ReadbackStatus {
        Success,   
        Expired,   
        Failed,    
    };
 
    struct ReadbackResult {
        uint64_t id = 0;                                  
        ReadbackStatus status = ReadbackStatus::Failed;
        uint32_t bytes = 0;                               
    };
 
    using ReadbackCallback = std::function<void(const ReadbackResult&)>;

    virtual uint64_t compute_magnitude_async(const float* complex_pairs,
                                             float* magnitudes, uint32_t num_bins,
                                             std::chrono::microseconds deadline,
                                             ReadbackCallback on_complete) = 0;

    virtual uint64_t convolve_batch_async(const float* in_complex, float* out_complex,
                                          uint32_t n, uint32_t batch,
                                          std::chrono::microseconds deadline,
                                          ReadbackCallback on_complete) = 0;

    virtual std::size_t poll_readbacks() = 0;

    virtual std::size_t readbacks_in_flight() const = 0;

    virtual void set_max_readbacks_in_flight(std::size_t max_in_flight) = 0;

    struct AsyncStats {
        uint64_t submits = 0;
        uint64_t map_resolves = 0;
        uint64_t expired = 0;
        uint64_t failed = 0;
        uint64_t write_buffer_calls = 0;
        uint64_t write_buffer_bytes = 0;
        uint64_t output_copy_calls = 0;
        uint64_t output_copy_bytes = 0;
        uint64_t map_async_calls = 0;
        uint64_t mapped_readback_memcpy_calls = 0;
        uint64_t mapped_readback_memcpy_bytes = 0;
    };
 
    virtual AsyncStats async_stats() const = 0;

    virtual void set_async_timing_enabled(bool /*enabled*/) {}

    virtual double last_gpu_busy_ns() const { return 0.0; }
 
    // ── Kernel-source seam (diagnostic) ──────────────────────────────────
    //
    // A named compute pipeline's WGSL is readable, and replaceable before the
    // pipeline is built. This exists for ONE purpose: to demonstrate that the
    // WGSL text is causally upstream of the audio samples — mutate the kernel,
    // hear/measure the mutation in the output. It is what makes the browser
    // GPU-audio claim checkable rather than asserted. It is a diagnostic seam,
    // not a plugin/effect API: there is no validation of the replacement beyond
    // WGSL compilation, and a bad override simply fails pipeline creation.

    virtual const char* kernel_source(const char* label) const = 0;

    virtual bool override_kernel_source(const char* label, const char* wgsl) = 0;
 
    // ── Capabilities ─────────────────────────────────────────────────────

    struct CapabilityReport {
        bool available = false;
        bool adapter_info_authentic = false;
        std::string backend;
        std::string adapter_type;
        std::string name;
        std::string vendor;
        std::string architecture;
        std::string description;
        uint32_t vendor_id = 0;
        uint32_t device_id = 0;
        bool timestamp_query = false;   // compute-pass GPU timing available
        bool shader_f16 = false;        // half-precision in WGSL
        uint64_t max_storage_buffer_binding_size = 0;
        uint64_t max_buffer_size = 0;
        uint32_t max_compute_invocations_per_workgroup = 0;
        uint32_t max_compute_workgroup_size_x = 0;
        uint32_t max_compute_workgroup_storage_size = 0;
        // Largest power-of-two complex FFT this device's storage-buffer limit
        // admits, capped at the implementation maximum (0 if unavailable).
        uint32_t max_fft_size = 0;
    };

    virtual CapabilityReport capabilities() const = 0;
 
    // ── Device sharing verification ──────────────────────────────────────
 
    struct DeviceSharingReport {
        bool device_obtained = false;
        bool second_consumer_works = false;
        bool concurrent_submission_ok = false;
        bool memory_pressure_ok = false;
        std::string backend_name;     // "Metal", "D3D12", "Vulkan"
        std::string notes;
    };

    virtual DeviceSharingReport verify_device_sharing(GpuSurface& surface) = 0;
 
    // ── Benchmarking ────────────────────────────────────────────────────
 
    struct BenchmarkResult {
        uint32_t num_elements = 0;
        double upload_us = 0;      // microseconds: CPU → GPU buffer
        double dispatch_us = 0;    // microseconds: compute shader execution
        double readback_us = 0;    // microseconds: GPU buffer → CPU
        double total_us = 0;       // upload + dispatch + readback
        double cpu_baseline_us = 0; // equivalent CPU operation for comparison
        bool gpu_faster = false;
    };

    virtual std::vector<BenchmarkResult> benchmark_magnitude(
        const std::vector<uint32_t>& sizes, int iterations = 10) = 0;

    virtual std::vector<BenchmarkResult> benchmark_complex_multiply(
        const std::vector<uint32_t>& sizes, int iterations = 10) = 0;
 
    static std::unique_ptr<GpuCompute> create();
 
#ifdef PULP_BENCHMARK
    virtual void set_bench_counters(bench::PerfCounters* counters) = 0;
#endif
 
  protected:
    std::atomic<bool> initialized_{false};
};
 
} // namespace pulp::render
 
#endif // PULP_GPU_RUNTIME