Skip to main content

rufft_host/
dispatch.rs

1use ruda_core::tensor::{DType, host::HostTensor};
2
3pub fn rfft(
4    signal: HostTensor,
5    dim: usize,
6    n: Option<usize>,
7) -> (HostTensor, HostTensor) {
8    match signal.dtype() {
9        DType::F32 => crate::rfft_f32(signal, dim, n),
10        DType::F64 => crate::rfft_f64(signal, dim, n),
11        DType::F16 => crate::rfft_f16(signal, dim, n),
12        DType::BF16 => crate::rfft_bf16(signal, dim, n),
13        dtype => panic!("rfft: unsupported dtype {:?}", dtype),
14    }
15}
16
17pub fn irfft(
18    spectrum_re: HostTensor,
19    spectrum_im: HostTensor,
20    dim: usize,
21    n: Option<usize>,
22) -> HostTensor {
23    match spectrum_re.dtype() {
24        DType::F32 => crate::irfft_f32(spectrum_re, spectrum_im, dim, n),
25        DType::F64 => crate::irfft_f64(spectrum_re, spectrum_im, dim, n),
26        DType::F16 => crate::irfft_f16(spectrum_re, spectrum_im, dim, n),
27        DType::BF16 => crate::irfft_bf16(spectrum_re, spectrum_im, dim, n),
28        dtype => panic!("irfft: unsupported dtype {:?}", dtype),
29    }
30}
31