use std::collections::HashMap;
use std::marker::PhantomData;
use std::sync::{Arc, Mutex};
use cubecl::prelude::*;
use cubecl::server::Handle;
use cubecl::tune::{LocalTuner, Tunable, TunableSet, local_tuner};
use dspk_core::compute::LaunchGeometry;
use dspk_core::compute::tune::{size_class, tune_id};
use super::kernels::sos::{SVF_COEFFS, coeffs_len, sos_block_kernel, sos_block_scan_kernel, sos_block_vector_kernel};
use crate::core::{buffer, cast, cast_all, layout, to_f64, DspFloat, Scratch};
use crate::filter::design::{OpsError, FilterMode, FilterSpec, FilterStart, Sos};
const BLOCK_COUNT_CANDIDATES: [usize; 5] = [1, 4, 16, 64, 256];
pub const PINNED_BLOCK_COUNT: usize = 16;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PassLayout {
ChannelMajor,
TimeMajor,
TimeMajorVectorized { width: usize },
}
#[derive(Debug, Default)]
struct PassWorkspace {
slots: Scratch,
input_t: Scratch,
output_t: Scratch,
state_padded: Scratch,
}
#[derive(Debug, Clone)]
pub struct DeviceFilter<F: DspFloat = f32> {
sos: Sos,
mode: FilterMode,
start: FilterStart,
settling: usize,
coeffs_host: Vec<f64>,
coeffs: Handle,
transitions: Arc<Mutex<HashMap<usize, Handle>>>,
block_len: Option<usize>,
layout: Option<PassLayout>,
workspace: Arc<Mutex<HashMap<std::thread::ThreadId, PassWorkspace>>>,
_float: PhantomData<F>,
}
impl<F: DspFloat> DeviceFilter<F> {
pub fn new(client: &Client, spec: &FilterSpec, sample_rate: f64) -> Result<Self, OpsError> {
let sos = spec.design(sample_rate)?;
Ok(Self::from_sos(client, sos, spec.mode).with_start(spec.start))
}
pub fn from_sos(client: &Client, sos: Sos, mode: FilterMode) -> Self {
let settling = sos.settling_samples(crate::filter::design::DEFAULT_SETTLING_TOLERANCE);
let coeffs_host = kernel_coeffs::<F>(&sos);
let coeffs = buffer::upload(client, &cast_all::<F>(&coeffs_host));
Self {
sos,
mode,
start: FilterStart::default(),
settling,
coeffs_host,
coeffs,
transitions: Arc::default(),
block_len: None,
layout: None,
workspace: Arc::default(),
_float: PhantomData,
}
}
pub fn with_start(mut self, start: FilterStart) -> Self {
self.start = start;
self
}
pub fn with_block_len(mut self, block_len: usize) -> Self {
self.block_len = Some(block_len.max(1));
self
}
pub fn with_layout(mut self, layout: PassLayout) -> Self {
self.layout = Some(layout);
self
}
pub fn sos(&self) -> &Sos {
&self.sos
}
pub fn mode(&self) -> FilterMode {
self.mode
}
pub fn n_sections(&self) -> usize {
self.sos.len()
}
pub fn settling_samples(&self) -> usize {
self.settling
}
pub fn state_len(&self) -> usize {
self.n_sections() * 2 + 1
}
pub fn scratch_len(&self, channels: usize, samples: usize) -> usize {
match self.mode {
FilterMode::Forward => 0,
FilterMode::ForwardBackward => channels * (samples + 2 * self.edge_pad(samples)),
}
}
fn edge_pad(&self, samples: usize) -> usize {
self.settling.min(samples.saturating_sub(1))
}
#[allow(clippy::too_many_arguments)]
pub fn apply(
&self,
client: &Client,
input: &Handle,
output: &Handle,
scratch: &Handle,
state: &Handle,
channels: usize,
samples: usize,
) {
if channels == 0 || samples == 0 {
return;
}
self.check_buffer("state", state, channels * self.state_len());
self.check_buffer("scratch", scratch, self.scratch_len(channels, samples));
let pass = |input: &Handle, output: &Handle, in_len, pad, out_start, out_len, reverse, rest| Pass {
input: input.clone(),
output: output.clone(),
state: state.clone(),
channels,
in_len,
pad,
out_start,
out_len,
reverse,
carry: false,
rest,
};
let rest = self.start == FilterStart::Rest;
match self.mode {
FilterMode::Forward => self.run(client, pass(input, output, samples, 0, 0, samples, false, rest)),
FilterMode::ForwardBackward => {
let pad = self.edge_pad(samples);
let ext = samples + 2 * pad;
self.run(client, pass(input, scratch, samples, pad, 0, ext, false, false));
self.run(client, pass(scratch, output, ext, 0, pad, samples, true, false));
}
}
}
#[allow(clippy::too_many_arguments)]
pub fn apply_stateful(
&self,
client: &Client,
input: &Handle,
output: &Handle,
state: &Handle,
channels: usize,
samples: usize,
first: bool,
) -> Result<(), OpsError> {
if self.mode != FilterMode::Forward {
return Err(OpsError::ForwardBackwardOnLiveStream);
}
if channels > 0 && samples > 0 {
self.check_buffer("state", state, channels * self.state_len());
let pass = Pass {
input: input.clone(),
output: output.clone(),
state: state.clone(),
channels,
in_len: samples,
pad: 0,
out_start: 0,
out_len: samples,
reverse: false,
carry: !first,
rest: first && self.start == FilterStart::Rest,
};
self.run(client, pass);
}
Ok(())
}
fn check_buffer(&self, name: &str, handle: &Handle, len: usize) {
let need = (len * size_of::<F>()) as u64;
assert!(handle.size_in_used() >= need, "{name} buffer holds {} bytes, the filter needs {need}", handle.size_in_used());
}
fn run(&self, client: &Client, pass: Pass) {
let inputs = PassInputs { client: client.clone(), filter: self.clone(), pass };
match (self.block_len, self.layout) {
(None, None) if dspk_core::compute::tuned_choices_pinned() => {
inputs.run(inputs.steps().div_ceil(PINNED_BLOCK_COUNT), PassLayout::ChannelMajor)
}
(None, None) => tuned_pass(inputs),
(len, layout) => inputs.run(len.unwrap_or(usize::MAX), layout.unwrap_or(PassLayout::ChannelMajor)),
}
}
fn transition(&self, client: &Client, block_len: usize) -> Handle {
let mut cache = self.transitions.lock().expect("transition cache");
cache
.entry(block_len)
.or_insert_with(|| {
let power = transition_power(&self.coeffs_host, self.n_sections(), block_len);
buffer::upload(client, &cast_all::<F>(&power))
})
.clone()
}
}
fn kernel_coeffs<F: DspFloat>(sos: &Sos) -> Vec<f64> {
let round = |v: f64| to_f64(cast::<F>(v));
let n = sos.len();
let mut c: Vec<f64> = sos.sections.iter().flat_map(|s| s.svf()).map(round).collect();
c.push(round(sos.dc_gain()));
let rest = unit_step_states(&c, n);
c.extend(rest.into_iter().map(round));
debug_assert_eq!(c.len(), coeffs_len(n));
c
}
fn svf_step(coeffs: &[f64], n_sections: usize, z: &mut [f64], x: f64) -> f64 {
let mut v = x;
for s in 0..n_sections {
let c = &coeffs[s * SVF_COEFFS..(s + 1) * SVF_COEFFS];
let (ic1, ic2) = (z[2 * s], z[2 * s + 1]);
let v3 = v - ic2;
let v1 = c[0] * ic1 + c[1] * v3;
let v2 = ic2 + c[1] * ic1 + c[2] * v3;
z[2 * s] = 2.0 * v1 - ic1;
z[2 * s + 1] = 2.0 * v2 - ic2;
v = c[3] * v + c[4] * v1 + c[5] * v2;
}
v
}
fn transition_matrix(coeffs: &[f64], n_sections: usize) -> Vec<f64> {
let dim = 2 * n_sections;
let mut a = vec![0.0f64; dim * dim];
for k in 0..dim {
let mut z = vec![0.0f64; dim];
z[k] = 1.0;
svf_step(coeffs, n_sections, &mut z, 0.0);
for r in 0..dim {
a[r * dim + k] = z[r];
}
}
a
}
fn unit_step_states(coeffs: &[f64], n_sections: usize) -> Vec<f64> {
let dim = 2 * n_sections;
let a = transition_matrix(coeffs, n_sections);
let mut b = vec![0.0f64; dim];
svf_step(coeffs, n_sections, &mut b, 1.0);
let i_minus_a: Vec<f64> = (0..dim * dim).map(|i| if i % (dim + 1) == 0 { 1.0 } else { 0.0 } - a[i]).collect();
solve(i_minus_a, b, dim)
}
fn solve(mut m: Vec<f64>, mut rhs: Vec<f64>, dim: usize) -> Vec<f64> {
for col in 0..dim {
let pivot = (col..dim).max_by(|&a, &b| m[a * dim + col].abs().total_cmp(&m[b * dim + col].abs())).unwrap_or(col);
if pivot != col {
for k in 0..dim {
m.swap(col * dim + k, pivot * dim + k);
}
rhs.swap(col, pivot);
}
let p = m[col * dim + col];
for row in col + 1..dim {
let f = m[row * dim + col] / p;
for k in col..dim {
m[row * dim + k] -= f * m[col * dim + k];
}
rhs[row] -= f * rhs[col];
}
}
let mut x = vec![0.0f64; dim];
for row in (0..dim).rev() {
let s: f64 = (row + 1..dim).map(|k| m[row * dim + k] * x[k]).sum();
x[row] = (rhs[row] - s) / m[row * dim + row];
}
x
}
fn transition_power(coeffs: &[f64], n_sections: usize, power: usize) -> Vec<f64> {
let dim = 2 * n_sections;
let mul = |x: &[f64], y: &[f64]| {
let mut out = vec![0.0f64; dim * dim];
for r in 0..dim {
for c in 0..dim {
out[r * dim + c] = (0..dim).map(|k| x[r * dim + k] * y[k * dim + c]).sum();
}
}
out
};
let mut result: Vec<f64> = (0..dim * dim).map(|i| if i % (dim + 1) == 0 { 1.0 } else { 0.0 }).collect();
let (mut base, mut e) = (transition_matrix(coeffs, n_sections), power);
while e > 0 {
if e & 1 == 1 {
result = mul(&result, &base);
}
base = mul(&base, &base);
e >>= 1;
}
for d in 0..dim {
result[d * dim + d] -= 1.0;
}
result
}
#[derive(Clone)]
struct Pass {
input: Handle,
output: Handle,
state: Handle,
channels: usize,
in_len: usize,
pad: usize,
out_start: usize,
out_len: usize,
reverse: bool,
carry: bool,
rest: bool,
}
#[derive(Clone)]
struct PassInputs<F: DspFloat> {
client: Client,
filter: DeviceFilter<F>,
pass: Pass,
}
impl<F: DspFloat> PassInputs<F> {
fn steps(&self) -> usize {
self.pass.in_len + 2 * self.pass.pad
}
fn run(&self, block_len: usize, layout: PassLayout) {
let (client, f, p) = (&self.client, &self.filter, &self.pass);
let n = f.n_sections();
let vector = match layout {
PassLayout::TimeMajorVectorized { width } => {
client.io_optimized_vector_sizes(size_of::<F>()).filter(|&w| w > 1 && w <= width).max()
}
_ => None,
};
let time_major = layout != PassLayout::ChannelMajor;
let steps = self.steps();
let block_len = block_len.clamp(1, steps.max(1));
let blocks = steps.div_ceil(block_len).max(1);
let mut workspaces = f.workspace.lock().expect("filter workspace");
let ws = workspaces.entry(std::thread::current().id()).or_default();
let lanes = vector.map_or(p.channels, |w| p.channels.next_multiple_of(w));
let padded = lanes > p.channels;
let state_len = f.state_len();
let slots = ws.slots.get::<F>(client, lanes * blocks * state_len);
let (input, output, in_strides, out_strides) = match time_major {
false => (p.input.clone(), p.output.clone(), (p.in_len, 1), (p.out_len, 1)),
true => {
let input_t = ws.input_t.get::<F>(client, lanes * p.in_len);
if padded {
buffer::fill::<F>(client, &input_t, lanes * p.in_len, cast::<F>(0.0));
}
layout::transpose_strided::<F>(client, &p.input, &input_t, p.channels, p.in_len, p.in_len, lanes);
let output_t = ws.output_t.get::<F>(client, lanes * p.out_len);
(input_t, output_t, (1, lanes), (1, lanes))
}
};
let state = match padded {
false => p.state.clone(),
true => {
let state = ws.state_padded.get::<F>(client, lanes * state_len);
buffer::fill::<F>(client, &state, lanes * state_len, cast::<F>(0.0));
buffer::copy::<F>(client, &p.state, &state, p.channels * state_len);
state
}
};
let launch = |phase: u32, active: usize| unsafe {
if let Some(width) = vector {
let groups = lanes / width;
let geom = LaunchGeometry::elementwise(client, groups * active);
return sos_block_vector_kernel::launch::<F>(
client,
geom.cube_count,
geom.cube_dim,
width,
BufferArg::from_raw_parts(input.clone(), lanes * p.in_len),
BufferArg::from_raw_parts(output.clone(), lanes * p.out_len),
BufferArg::from_raw_parts(f.coeffs.clone(), coeffs_len(n)),
BufferArg::from_raw_parts(state.clone(), lanes * state_len),
BufferArg::from_raw_parts(slots.clone(), (lanes * blocks * state_len).max(1)),
groups as u32,
p.in_len as u32,
p.pad as u32,
p.out_start as u32,
p.out_len as u32,
p.reverse as u32,
p.carry as u32,
p.rest as u32,
block_len as u32,
blocks as u32,
phase,
n,
);
}
let geom = LaunchGeometry::elementwise(client, p.channels * active);
sos_block_kernel::launch::<F>(
client,
geom.cube_count,
geom.cube_dim,
BufferArg::from_raw_parts(input.clone(), p.channels * p.in_len),
BufferArg::from_raw_parts(output.clone(), p.channels * p.out_len),
BufferArg::from_raw_parts(f.coeffs.clone(), coeffs_len(n)),
BufferArg::from_raw_parts(p.state.clone(), p.channels * f.state_len()),
BufferArg::from_raw_parts(slots.clone(), (p.channels * blocks * f.state_len()).max(1)),
p.channels as u32,
p.in_len as u32,
p.pad as u32,
p.out_start as u32,
p.out_len as u32,
in_strides.0 as u32,
in_strides.1 as u32,
out_strides.0 as u32,
out_strides.1 as u32,
p.reverse as u32,
p.carry as u32,
p.rest as u32,
block_len as u32,
blocks as u32,
phase,
n,
);
};
if blocks > 1 {
launch(0, blocks - 1);
let transition = f.transition(client, block_len);
let geom = LaunchGeometry::per_channel(client, lanes);
unsafe {
sos_block_scan_kernel::launch::<F>(
client,
geom.cube_count,
geom.cube_dim,
BufferArg::from_raw_parts(slots.clone(), lanes * blocks * state_len),
BufferArg::from_raw_parts(transition, (2 * n) * (2 * n)),
lanes as u32,
blocks as u32,
n,
);
}
}
launch(1, blocks);
if time_major {
layout::transpose_strided::<F>(client, &output, &p.output, p.out_len, p.channels, lanes, p.out_len);
}
if padded {
buffer::copy::<F>(client, &state, &p.state, p.channels * state_len);
}
}
}
fn tuned_pass<F: DspFloat>(inputs: PassInputs<F>) {
static TUNER: LocalTuner<String, String> = local_tuner!("sos-blocks-v3");
let client = inputs.client.clone();
let id = tune_id(&client);
let widths: Vec<usize> = client.io_optimized_vector_sizes(size_of::<F>()).filter(|&w| w > 1).collect();
let set = TUNER.init(&id, move || {
let key = |p: &PassInputs<F>| {
format!(
"{}-n{}-c{}-t{}-a{}",
F::type_name(),
p.filter.n_sections(),
size_class(p.pass.channels),
size_class(p.steps()),
p.pass.channels.trailing_zeros().min(4)
)
};
let scratch = |_: &String, p: &PassInputs<F>| {
let mut pass = p.pass.clone();
pass.output = buffer::empty::<F>(&p.client, pass.channels * pass.out_len);
pass.state = buffer::zeros::<F>(&p.client, pass.channels * p.filter.state_len());
PassInputs { client: p.client.clone(), filter: p.filter.clone(), pass }
};
let set: TunableSet<String, PassInputs<F>, ()> = TunableSet::new(key, scratch);
let layouts = [PassLayout::ChannelMajor, PassLayout::TimeMajor].into_iter().chain(widths.iter().map(|&width| PassLayout::TimeMajorVectorized { width }));
let candidates = layouts.flat_map(|layout| BLOCK_COUNT_CANDIDATES.into_iter().map(move |blocks| (layout, blocks)));
candidates.fold(set, |set, (layout, blocks)| {
set.with(Tunable::new(&format!("{layout:?}-blocks{blocks}"), move |p: PassInputs<F>| {
p.run(p.steps().div_ceil(blocks), layout);
Ok::<_, String>(())
}))
})
});
TUNER.execute(&id, &client, set, inputs)
}
#[allow(clippy::too_many_arguments)]
pub fn execute_filter<F: DspFloat>(
client: &Client,
input: &Handle,
output: &Handle,
channels: usize,
samples: usize,
sample_rate: f64,
spec: &FilterSpec,
) -> Result<(), OpsError> {
let filter = DeviceFilter::<F>::new(client, spec, sample_rate)?;
let scratch = buffer::empty::<F>(client, filter.scratch_len(channels, samples));
let state = buffer::empty::<F>(client, channels * filter.state_len());
filter.apply(client, input, output, &scratch, &state, channels, samples);
Ok(())
}