speakrs 0.6.0

Fast Rust speaker diarization with pyannote-level accuracy and native CPU/CoreML/CUDA inference
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

//! Complete dense plans shared by segmentation and embedding

use super::candidate::{DenseCandidate, DenseOxide, DenseSpec, Phases};
use super::implementation::{AreaTarget, LibraryNeed, Selected, plan_selection};
use super::{CudaError, CudaLibrary, CudaRuntime, KernelModule};
use cudarc::driver::{CudaSlice, CudaView, CudaViewMut};

/// One implementation owner for a dense operation and its epilogue
#[derive(Debug)]
pub(super) enum DensePlan {
    #[cfg(feature = "_cuda-libraries")]
    Library,
    Oxide(Box<DenseOxide>),
}

impl DensePlan {
    /// Library plans must finish lazy setup before graph capture
    pub(super) fn requires_warmup(&self) -> bool {
        match self {
            #[cfg(feature = "_cuda-libraries")]
            Self::Library => true,
            Self::Oxide(_) => false,
        }
    }

    pub(super) fn new(
        runtime: &CudaRuntime,
        spec: DenseSpec,
        rows: usize,
        weight: &CudaSlice<f32>,
        bias: &CudaSlice<f32>,
    ) -> Result<Self, CudaError> {
        let boundary = spec.site().boundary();
        let selected = plan_selection(
            runtime,
            boundary,
            spec.batch(),
            spec.math(),
            #[cfg(all(test, feature = "_cuda-libraries"))]
            None,
        )?;
        if let Selected::Oxide(token) = selected
            && let Some(plan) = token.dense(runtime, spec, rows, weight, bias)?
        {
            return Ok(Self::Oxide(Box::new(plan)));
        }
        LibraryNeed::new(
            boundary,
            spec.batch(),
            spec.math(),
            AreaTarget::for_area(runtime, KernelModule::Segdense)?,
            CudaLibrary::Cublas,
        )
        .prepare(runtime)?;
        #[cfg(feature = "_cuda-libraries")]
        return Ok(Self::Library);
        #[cfg(not(feature = "_cuda-libraries"))]
        Err(CudaError::LibraryUnavailable {
            library: CudaLibrary::Cublas,
        })
    }

    pub(super) fn enqueue_slice(
        &self,
        runtime: &CudaRuntime,
        input: &CudaSlice<f32>,
        output: &mut CudaSlice<f32>,
        library: impl FnOnce(&mut CudaSlice<f32>) -> Result<(), CudaError>,
    ) -> Result<(), CudaError> {
        #[cfg(not(feature = "_cuda-libraries"))]
        let _ = library;
        match self {
            #[cfg(feature = "_cuda-libraries")]
            Self::Library => library(output),
            Self::Oxide(plan) => plan.enqueue(
                &input.as_view(),
                &mut output.as_view_mut(),
                &Phases::new(),
                runtime,
            ),
        }
    }

    pub(super) fn enqueue(
        &self,
        runtime: &CudaRuntime,
        input: &CudaView<'_, f32>,
        output: &mut CudaViewMut<'_, f32>,
        library: impl FnOnce(&mut CudaViewMut<'_, f32>) -> Result<(), CudaError>,
    ) -> Result<(), CudaError> {
        #[cfg(not(feature = "_cuda-libraries"))]
        let _ = library;
        match self {
            #[cfg(feature = "_cuda-libraries")]
            Self::Library => library(output),
            Self::Oxide(plan) => plan.enqueue(input, output, &Phases::new(), runtime),
        }
    }
}