speakrs 0.6.0

//! Inspection helpers for direct CUDA tests

#[cfg(feature = "_cuda-libraries")]
use super::{ActivationForm, BlockForm, dispatch::Plan, trunk::ConvLayer};
use super::{EmbeddingBatch, ResNetEmbedding, SharedEmbeddingActivations};
#[cfg(feature = "_cuda-libraries")]
use super::{EmbeddingTapFn, SPEAKERS_PER_CHUNK};
#[cfg(feature = "_cuda-libraries")]
use crate::inference::cuda::candidate::{WideconvAlgorithm, WideconvFp16Tiles};
use crate::inference::cuda::{CudaError, CudaRuntime};
use cudarc::driver::CudaSlice;
use std::sync::Arc;

#[cfg(feature = "_cuda-libraries")]
impl EmbeddingBatch {
    /// Runs the forward pass eagerly and calls `tap` with every intermediate
    /// activation listed in [`EmbeddingTap`], in execution order
    ///
    /// The view is only valid during the call; download it there if needed. This
    /// never replays a captured graph
    pub fn forward_with_taps(
        &mut self,
        runtime: &CudaRuntime,
        tap: &mut EmbeddingTapFn<'_>,
    ) -> Result<(), CudaError> {
        let activations = self.activations.clone();
        let mut storage = activations.lock()?;
        self.run_with_activations(
            runtime,
            super::PlanSet::Selected,
            super::ActivationForm::Fp32,
            tap,
            storage.buffers_mut(),
        )
    }

    /// Copies host fbank and masks into the batch, runs the forward pass and
    /// downloads the embeddings
    ///
    /// `fbank` is `[chunks, 998, 80]` and `masks` is `[chunks * 3, 589]`; the result
    /// is `[chunks * 3, 256]`, all row-major
    pub fn embed(
        &mut self,
        runtime: &CudaRuntime,
        fbank: &[f32],
        masks: &[f32],
    ) -> Result<Vec<f32>, CudaError> {
        let stream = runtime.stream();
        self.fbank.copy_from_host(stream, fbank)?;
        self.masks.copy_from_host(stream, masks)?;
        self.forward(runtime)?;
        self.download_output(runtime)
    }

    /// Embeddings per forward pass, `chunks * 3`
    pub fn rows(&self) -> usize {
        self.chunks * SPEAKERS_PER_CHUNK
    }

    /// Distinct convolution shapes, one cuDNN plan each
    pub fn plan_count(&self) -> usize {
        self.plans.len()
    }

    /// Residual blocks whose two convolutions run FP16 tiles, wide on 32 or 64
    /// channels or either tile on 128 or 256, and how many blocks [`Self::forward`]
    /// passes a half hidden activation
    pub fn half_hidden_blocks(&self) -> (usize, usize) {
        let (fp16_pairs, _, half_pairs, _) = self.half_blocks();
        (fp16_pairs, half_pairs)
    }

    /// Residual blocks whose output feeds a next block without a shortcut, where the
    /// block's second convolution and both of the next block's run such FP16 tiles,
    /// and how many blocks [`Self::forward`] writes as halves
    pub fn half_output_blocks(&self) -> (usize, usize) {
        let (_, fp16_outputs, _, half_outputs) = self.half_blocks();
        (fp16_outputs, half_outputs)
    }

    /// The FP16 pairs and links read from the selected plans' configs, then the
    /// half hidden activations and outputs the forward pass's own [`BlockForm::plan`]
    /// gives, so a forward pass that stops passing halves fails the comparison
    fn half_blocks(&self) -> (usize, usize, usize, usize) {
        let fp16 = |layer: &ConvLayer| {
            let conv = layer.conv(self.chunks, self.model.math);
            let tiles: &[WideconvFp16Tiles] = match conv.in_channels {
                32 | 64 => &[WideconvFp16Tiles::Wide],
                128 | 256 => &[WideconvFp16Tiles::Wide, WideconvFp16Tiles::Narrow],
                _ => return false,
            };
            conv.out_channels == conv.in_channels
                && conv.stride == [1, 1]
                && self.plans.iter().any(|(name, plan)| {
                    name == layer.name()
                        && matches!(plan, Plan::Wideconv(plan)
                        if tiles.iter().any(|&tiles| {
                            plan.config().algorithm == WideconvAlgorithm::Fp16(tiles)
                        }))
                })
        };
        let blocks = &self.model.trunk.blocks;
        let fp16_pairs = blocks
            .iter()
            .filter(|block| fp16(&block.conv1) && fp16(&block.conv2))
            .count();
        // a shortcut reads the block input as FP32, so its block's input stays FP32
        let fp16_outputs = blocks
            .windows(2)
            .filter(|pair| {
                let [block, next] = pair else { return false };
                next.shortcut.is_none()
                    && fp16(&block.conv2)
                    && fp16(&next.conv1)
                    && fp16(&next.conv2)
            })
            .count();

        let forms = BlockForm::plan(blocks, &self.plans, ActivationForm::Half);
        let half_pairs = forms.iter().filter(|form| form.hidden).count();
        let half_outputs = forms.iter().filter(|form| form.output).count();
        (fp16_pairs, fp16_outputs, half_pairs, half_outputs)
    }

    /// Whether [`Self::forward`] replays a captured CUDA graph
    pub fn has_graph(&self) -> bool {
        self.graph.is_some()
    }

    /// Bytes of the shared cuDNN workspace
    pub fn workspace_bytes(&self) -> usize {
        self.workspace.len()
    }

    /// Bytes held by this batch's activation, input and output buffers, excluding
    /// the cuDNN workspace
    pub fn buffer_bytes(&self) -> usize {
        let floats = self.fbank.len()
            + self.masks.len()
            + self.stem_input.len()
            + self.pooled.len()
            + self.output.len();
        floats * size_of::<f32>() + self.activations.retained_bytes()
    }
}

impl ResNetEmbedding {
    /// Allocates the device buffers and plans the convolutions for `chunks` fbank
    /// chunks per forward pass
    ///
    /// Do this once per batch class and reuse the batch for every forward pass of
    /// that size. The batch keeps a handle to this model and must run on
    /// `runtime`, whose stream owns its buffers
    pub fn batch(&self, runtime: &CudaRuntime, chunks: usize) -> Result<EmbeddingBatch, CudaError> {
        let activations = self.activations(runtime, chunks)?;
        self.batch_with_activations(runtime, chunks, activations)
    }
}

impl SharedEmbeddingActivations {
    pub(crate) fn retained_bytes(&self) -> usize {
        let storage = self.lock().unwrap();
        let buffers = storage.buffers();
        (buffers.trunk.iter().map(CudaSlice::len).sum::<usize>()
            + buffers.hidden.len()
            + buffers.shortcut.len())
            * size_of::<f32>()
    }
}

impl EmbeddingBatch {
    pub(crate) fn shares_activations(&self, other: &SharedEmbeddingActivations) -> bool {
        Arc::ptr_eq(&self.activations.0, &other.0)
    }
}