perf: fir post processing in rust

Cargo.toml

@@ -22,6 +22,7 @@ ndarray = { version = "0.15.6", features = ["rayon"] }
 num = "0.4.2"
 numpy = "0.21.0"
 ordered-float = "4.2.0"
+pulp = "0.18.14"
 pyo3 = { version = "0.21.2", features = ["hashbrown", "anyhow"] }
 rayon = "1.10.0"
 thiserror = "1.0.60"

src/lib.rs

@@ -10,7 +10,7 @@ mod shape;
 use std::{borrow::Borrow, fmt::Debug, str::FromStr, sync::Arc};
 
 use hashbrown::HashMap;
-use indoc::indoc;
+use ndarray::ArrayViewMut2;
 use numpy::{prelude::*, AllowTypeChange, PyArray1, PyArray2, PyArrayLike1, PyArrayLike2};
 use pyo3::{
     exceptions::{PyRuntimeError, PyTypeError, PyValueError},
@@ -21,7 +21,10 @@ use rayon::prelude::*;
 
 use crate::{
     executor::Executor,
-    pulse::{PulseList, Sampler},
+    pulse::{
+        apply_fir_inplace, apply_iir_inplace, apply_iq_inplace, apply_offset_inplace, PulseList,
+        Sampler,
+    },
     quant::{Amplitude, ChannelId, Frequency, Phase, ShapeId, Time},
     schedule::{ElementCommonBuilder, ElementRef},
 };
@@ -1959,19 +1962,21 @@ fn generate_waveforms(
         allow_oversize,
     )?;
     let waveforms = sample_waveform(py, &channels, pulse_lists, crosstalk, time_tolerance)?;
-    py.allow_threads(|| {
+    Ok(py.allow_threads(|| {
         waveforms
             .into_par_iter()
             .map(|(n, w)| {
                 Python::with_gil(|py| {
-                    let mut w = w.into_bound(py);
+                    let w = w.bind(py);
+                    let mut w = w.readwrite();
+                    let mut w = w.as_array_mut();
                     let c = &channels[&n];
-                    post_process(py, &mut w, c)?;
-                    Ok((n, w.unbind()))
-                })
+                    post_process(py, &mut w, c);
+                });
+                (n, w)
             })
-            .collect::<PyResult<_>>()
-    })
+            .collect()
+    }))
 }
 
 fn build_pulse_lists(
@@ -1991,9 +1996,12 @@ fn build_pulse_lists(
         let s = s.bind(py);
         executor.add_shape(n.clone(), Shape::get_rust_shape(s)?);
     }
-    executor
-        .execute(&schedule.get().0)
-        .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+    let schedule = &schedule.get().0;
+    py.allow_threads(|| {
+        executor
+            .execute(schedule)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
+    })?;
     Ok(executor.into_result())
 }
 
@@ -2027,80 +2035,45 @@ fn sample_waveform(
     Ok(waveforms)
 }
 
-fn post_process(py: Python, w: &mut Bound<PyArray2<f64>>, c: &Channel) -> PyResult<()> {
-    if let Some(iq_matrix) = &c.iq_matrix {
-        apply_iq_matrix(py, w, iq_matrix.bind(py));
+fn post_process(py: Python, w: &mut ArrayViewMut2<f64>, c: &Channel) {
+    macro_rules! map_as_array {
+        ($n:ident) => {
+            let temp = c.$n.as_ref().map(|x| x.bind(py).readonly());
+            let $n = temp.as_ref().map(|x| x.as_array());
+        };
     }
-    if c.filter_offset {
-        if let Some(offset) = &c.offset {
-            apply_offset(py, w, offset.bind(py));
-        }
-        if let Some(iir) = &c.iir {
-            apply_iir(py, w, iir.bind(py));
-        }
-        if let Some(fir) = &c.fir {
-            apply_fir(py, w, fir.bind(py))?;
-        }
-    } else {
-        if let Some(iir) = &c.iir {
-            apply_iir(py, w, iir.bind(py));
-        }
-        if let Some(fir) = &c.fir {
-            apply_fir(py, w, fir.bind(py))?;
+    map_as_array!(iq_matrix);
+    map_as_array!(offset);
+    map_as_array!(iir);
+    map_as_array!(fir);
+    py.allow_threads(|| {
+        if let Some(iq_matrix) = iq_matrix {
+            apply_iq_inplace(w, iq_matrix);
         }
-        if let Some(offset) = &c.offset {
-            apply_offset(py, w, offset.bind(py));
+        if c.filter_offset {
+            if let Some(offset) = offset {
+                apply_offset_inplace(w, offset);
+            }
+            if let Some(iir) = iir {
+                apply_iir_inplace(w, iir);
+            }
+            if let Some(fir) = fir {
+                apply_fir_inplace(w, fir);
+            }
+        } else {
+            if let Some(iir) = iir {
+                apply_iir_inplace(w, iir);
+            }
+            if let Some(fir) = fir {
+                apply_fir_inplace(w, fir);
+            }
+            if let Some(offset) = offset {
+                apply_offset_inplace(w, offset);
+            }
         }
-    }
-    Ok(())
-}
-
-fn apply_iq_matrix(py: Python, w: &mut Bound<PyArray2<f64>>, iq_matrix: &Bound<PyArray2<f64>>) {
-    let mut w = w.readwrite();
-    let w = w.as_array_mut();
-    let iq_matrix = iq_matrix.readonly();
-    let iq_matrix = iq_matrix.as_array();
-    py.allow_threads(|| {
-        pulse::apply_iq_inplace(w, iq_matrix);
-    });
-}
-
-fn apply_offset(py: Python, w: &Bound<PyArray2<f64>>, offset: &Bound<PyArray1<f64>>) {
-    let mut w = w.readwrite();
-    let w = w.as_array_mut();
-    let offset = offset.readonly();
-    let offset = offset.as_array();
-    py.allow_threads(|| {
-        pulse::apply_offset_inplace(w, offset);
     });
 }
 
-fn apply_iir(py: Python, w: &Bound<PyArray2<f64>>, iir: &Bound<PyArray2<f64>>) {
-    let mut w = w.readwrite();
-    let w = w.as_array_mut();
-    let iir = iir.readonly();
-    let iir = iir.as_array();
-    py.allow_threads(|| {
-        pulse::apply_iir_inplace(w, iir);
-    });
-}
-
-fn apply_fir(py: Python, w: &Bound<PyArray2<f64>>, fir: &Bound<PyArray1<f64>>) -> PyResult<()> {
-    let locals = PyDict::new_bound(py);
-    locals.set_item("w", w)?;
-    locals.set_item("fir", fir)?;
-    py.run_bound(
-        indoc! {"
-            from scipy import signal
-            for wi in w:
-                wi[:] = signal.convolve(wi, fir, mode='full')[:len(wi)]
-        "},
-        None,
-        Some(&locals),
-    )?;
-    Ok(())
-}
-
 /// Generates microwave pulses for superconducting quantum computing
 /// experiments.
 ///

src/pulse.rs

@@ -11,6 +11,7 @@ use hashbrown::HashMap;
 use itertools::{izip, Itertools};
 use ndarray::{azip, s, ArrayView1, ArrayView2, ArrayViewMut2, Axis};
 use numpy::Complex64;
+use pulp::{Arch, Simd, WithSimd};
 use rayon::prelude::*;
 
 use crate::{
@@ -460,7 +461,7 @@ where
     Ok(())
 }
 
-pub(crate) fn apply_iq_inplace(mut waveform: ArrayViewMut2<f64>, iq_matrix: ArrayView2<f64>) {
+pub(crate) fn apply_iq_inplace(waveform: &mut ArrayViewMut2<f64>, iq_matrix: ArrayView2<f64>) {
     assert!(matches!(waveform.shape(), [2, _]));
     assert!(matches!(iq_matrix.shape(), [2, 2]));
     for mut col in waveform.columns_mut() {
@@ -473,12 +474,12 @@ pub(crate) fn apply_iq_inplace(mut waveform: ArrayViewMut2<f64>, iq_matrix: Arra
     }
 }
 
-pub(crate) fn apply_offset_inplace(mut waveform: ArrayViewMut2<f64>, offset: ArrayView1<f64>) {
+pub(crate) fn apply_offset_inplace(waveform: &mut ArrayViewMut2<f64>, offset: ArrayView1<f64>) {
     assert!(waveform.shape()[0] == offset.len());
     azip!((mut row in waveform.axis_iter_mut(Axis(0)), &offset in &offset) row += offset);
 }
 
-pub(crate) fn apply_iir_inplace(mut waveform: ArrayViewMut2<f64>, sos: ArrayView2<f64>) {
+pub(crate) fn apply_iir_inplace(waveform: &mut ArrayViewMut2<f64>, sos: ArrayView2<f64>) {
     let mut biquads: Vec<_> = sos
         .axis_iter(Axis(0))
         .map(|row| {
@@ -508,3 +509,65 @@ fn apply_iir_inplace_1d(waveform: &mut [f64], biquads: &mut [biquad::DirectForm2
         *y = x;
     }
 }
+
+pub(crate) fn apply_fir_inplace(waveform: &mut ArrayViewMut2<f64>, taps: ArrayView1<f64>) {
+    let arch = Arch::new();
+    arch.dispatch(ApplyFirInplace {
+        waveform: waveform.view_mut(),
+        taps,
+    });
+}
+
+struct ApplyFirInplace<'a, 'b> {
+    waveform: ArrayViewMut2<'a, f64>,
+    taps: ArrayView1<'b, f64>,
+}
+
+impl<'a, 'b> WithSimd for ApplyFirInplace<'a, 'b> {
+    type Output = ();
+
+    #[inline(always)]
+    fn with_simd<S: Simd>(mut self, simd: S) -> Self::Output {
+        let lanes = std::mem::size_of::<S::f64s>() / std::mem::size_of::<f64>();
+        let buffer_len = align_ceil(self.taps.len(), lanes);
+        assert!(buffer_len % lanes == 0);
+        let taps_buffer = {
+            let mut buffer = vec![0.0; buffer_len * 2];
+            for (&t, b) in self.taps.iter().zip(buffer[..buffer_len].iter_mut().rev()) {
+                *b = t;
+            }
+            for (&t, b) in self.taps.iter().zip(buffer[buffer_len..].iter_mut().rev()) {
+                *b = t;
+            }
+            buffer
+        };
+        for mut row in self.waveform.axis_iter_mut(Axis(0)) {
+            let mut w_buffer = vec![0.0; buffer_len];
+            for (i, w) in row.iter_mut().enumerate() {
+                let w_buffer_index = i % buffer_len;
+                w_buffer[w_buffer_index] = *w;
+                let tap_buffer_index = buffer_len - w_buffer_index - 1;
+                let (taps_simd, _) = S::f64s_as_simd(&taps_buffer[tap_buffer_index..]);
+                let (w_simd, _) = S::f64s_as_simd(&w_buffer);
+                let sum = taps_simd
+                    .iter()
+                    .zip(w_simd.iter())
+                    .fold(simd.f64s_splat(0.0), |acc, (taps, w)| {
+                        simd.f64s_mul_add_e(*taps, *w, acc)
+                    });
+                let sum = simd.f64s_reduce_sum(sum);
+                *w = sum;
+            }
+        }
+    }
+}
+
+#[inline]
+fn align_ceil(x: usize, n: usize) -> usize {
+    let r = x % n;
+    if r == 0 {
+        x
+    } else {
+        x + n - r
+    }
+}