tract_gpu/ops/
stft.rs

1use crate::tensor::{DeviceTensor, DeviceTensorExt, IntoDevice};
2use tract_core::internal::*;
3
4/// The power-of-two FFT frame lengths the GPU STFT kernels support (the radix-2 kernel
5/// generalizes to any pow2; these cover Kaldi-style featurizers at 8/16/32/48 kHz).
6pub const SUPPORTED_FRAMES: [usize; 4] = [256, 512, 1024, 2048];
7
8/// Whether the GPU STFT kernels handle this frame length (a supported power of two).
9pub fn is_supported_frame(frame: usize) -> bool {
10    SUPPORTED_FRAMES.contains(&frame)
11}
12
13/// Per-backend STFT kernel launcher: `(stride, input, window, output)`. `input` is
14/// interleaved-complex f32 `[lead.., T, 2]`, `window` the pre-padded real window
15/// `[frame]`, `output` `[lead.., frames, frame, 2]`. The kernel reads the frame length
16/// from the output shape (`output[axis + 1]`).
17pub type DispatchStftFn = fn(usize, &DeviceTensor, &DeviceTensor, &DeviceTensor) -> TractResult<()>;
18
19/// Backend-agnostic fused STFT (frame + window + forward FFT). `frame` is a supported
20/// power of two ([`SUPPORTED_FRAMES`]); the window is pre-padded to `[frame]` (all-ones
21/// when the source had none), matching `core::ops::fft::Stft`'s symmetric padding. The
22/// time axis must sit just before the trailing complex pair (`axis == rank - 2`). Each
23/// backend supplies its own `dispatch` kernel; everything else (facts, output allocation,
24/// window upload) is shared.
25#[derive(Clone)]
26pub struct GpuStft {
27    pub axis: usize,
28    pub frame: usize,
29    pub stride: usize,
30    pub window: Arc<Tensor>,
31    pub backend_name: &'static str,
32    pub dispatch: DispatchStftFn,
33}
34
35impl GpuStft {
36    fn output_shape<D: DimLike>(&self, input: &[D]) -> TVec<D> {
37        let mut shape: TVec<D> = input.into();
38        let frames = (input[self.axis].clone() - self.frame) / self.stride + 1;
39        shape[self.axis] = frames;
40        shape.insert(self.axis + 1, self.frame.into());
41        shape
42    }
43}
44
45impl std::fmt::Debug for GpuStft {
46    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
47        write!(f, "{}Stft(frame={}, stride={})", self.backend_name, self.frame, self.stride)
48    }
49}
50
51impl PartialEq for GpuStft {
52    fn eq(&self, other: &Self) -> bool {
53        self.backend_name == other.backend_name
54            && self.axis == other.axis
55            && self.frame == other.frame
56            && self.stride == other.stride
57            && self.window == other.window
58    }
59}
60
61impl Eq for GpuStft {}
62
63impl std::hash::Hash for GpuStft {
64    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
65        self.backend_name.hash(state);
66        self.axis.hash(state);
67        self.frame.hash(state);
68        self.stride.hash(state);
69        self.window.hash(state);
70    }
71}
72
73impl Op for GpuStft {
74    fn name(&self) -> StaticName {
75        format!("{}Stft", self.backend_name).into()
76    }
77
78    op_as_typed_op!();
79}
80
81impl EvalOp for GpuStft {
82    op_out_of_plan!();
83
84    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
85        let input = inputs[0].to_device_tensor()?;
86        let window = (*self.window).clone().into_device()?;
87        let output = crate::turn_handler::make_tensor_for_node(
88            ctx,
89            input.datum_type(),
90            &self.output_shape(input.shape()),
91        )?;
92        (self.dispatch)(self.stride, input, &window, &output)
93            .with_context(|| format!("Error while dispatching eval for {}", self.name()))?;
94        Ok(tvec!(output.into_tensor().into_tvalue()))
95    }
96}
97
98impl TypedOp for GpuStft {
99    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
100        crate::utils::facts_to_device_facts(inputs, |facts| {
101            let input = facts[0];
102            ensure!(
103                input.rank() >= 2 && input.shape[input.rank() - 1] == 2.to_dim(),
104                "{} expects a complex input [.., T, 2]",
105                self.name()
106            );
107            Ok(tvec!(input.datum_type.fact(self.output_shape(&input.shape.to_tvec()))))
108        })
109        .with_context(|| format!("Error while computing facts for {:?}", self.name()))
110    }
111
112    as_op!();
113}
114
115/// Per-backend FFT kernel launcher: `(inverse, input, output)`, both interleaved-complex
116/// f32 `[lead.., N, 2]` (N a supported power of two, transformed axis at `rank-2`). The
117/// inverse is UNNORMALIZED, matching core `Fft` (rustfft). The kernel reads N from the
118/// input shape.
119pub type DispatchFftFn = fn(bool, &DeviceTensor, &DeviceTensor) -> TractResult<()>;
120
121/// Backend-agnostic complex FFT over the innermost-but-one axis (the trailing dim is the
122/// `[re, im]` pair); forward or `inverse`, shape-preserving. Mirrors `core::ops::fft::Fft`
123/// for a supported power-of-two length with `axis == rank - 2`.
124#[derive(Clone)]
125pub struct GpuFft {
126    pub axis: usize,
127    pub inverse: bool,
128    pub backend_name: &'static str,
129    pub dispatch: DispatchFftFn,
130}
131
132impl std::fmt::Debug for GpuFft {
133    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
134        write!(f, "{}Fft({})", self.backend_name, if self.inverse { "inverse" } else { "forward" })
135    }
136}
137
138impl PartialEq for GpuFft {
139    fn eq(&self, other: &Self) -> bool {
140        self.backend_name == other.backend_name
141            && self.axis == other.axis
142            && self.inverse == other.inverse
143    }
144}
145
146impl Eq for GpuFft {}
147
148impl std::hash::Hash for GpuFft {
149    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
150        self.backend_name.hash(state);
151        self.axis.hash(state);
152        self.inverse.hash(state);
153    }
154}
155
156impl Op for GpuFft {
157    fn name(&self) -> StaticName {
158        format!("{}Fft", self.backend_name).into()
159    }
160
161    op_as_typed_op!();
162}
163
164impl EvalOp for GpuFft {
165    op_out_of_plan!();
166
167    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
168        let input = inputs[0].to_device_tensor()?;
169        let output =
170            crate::turn_handler::make_tensor_for_node(ctx, input.datum_type(), input.shape())?;
171        (self.dispatch)(self.inverse, input, &output)
172            .with_context(|| format!("Error while dispatching eval for {}", self.name()))?;
173        Ok(tvec!(output.into_tensor().into_tvalue()))
174    }
175}
176
177impl TypedOp for GpuFft {
178    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
179        crate::utils::facts_to_device_facts(inputs, |facts| {
180            let input = facts[0];
181            ensure!(
182                input.rank() >= 2 && input.shape[input.rank() - 1] == 2.to_dim(),
183                "{} expects a complex input [.., N, 2]",
184                self.name()
185            );
186            Ok(tvec!(input.datum_type.fact(input.shape.clone())))
187        })
188        .with_context(|| format!("Error while computing facts for {:?}", self.name()))
189    }
190
191    as_op!();
192}
193
194/// Pre-pad `window` (or all-ones) to `[frame]` exactly as core `Stft` does: symmetric
195/// padding when shorter than the frame. Shared by every backend's lowering rule.
196pub fn padded_window(window: Option<&Arc<Tensor>>, frame: usize) -> TractResult<Arc<Tensor>> {
197    let mut win = vec![0f32; frame];
198    match window {
199        Some(w) => {
200            let w = w.cast_to::<f32>()?;
201            let w = w.try_as_plain()?;
202            let w = w.as_slice::<f32>()?;
203            ensure!(w.len() <= frame, "STFT window longer than frame");
204            let pad_left = (frame - w.len()) / 2;
205            win[pad_left..pad_left + w.len()].copy_from_slice(w);
206        }
207        None => win.fill(1.0),
208    }
209    Ok(Arc::new(tensor1(&win)))
210}
tract_gpu/ops/stft.rs

tract_gpu/ops/
stft.rs
