perf: optimize MlpgAdjust

.github/workflows/mlsafir.yml

@@ -3,6 +3,7 @@ on:
   pull_request:
     paths:
       - src/vocoder/mlsa.rs
+      - src/mlpg_adjust/**
 
 jobs:
   collect:

src/lib.rs

@@ -87,7 +87,7 @@ mod tests {
 
         assert_eq!(speech.len(), 74880);
         approx::assert_abs_diff_eq!(speech[2000], 2.3158134981607754e-5, epsilon = 1.0e-10);
-        approx::assert_abs_diff_eq!(speech[30000], 6459.375032316974, epsilon = 1.0e-10);
+        approx::assert_abs_diff_eq!(speech[30000], 6459.375032318177, epsilon = 1.0e-10);
     }
 
     // これ,名詞,代名詞,一般,*,*,*,これ,コレ,コレ,0/2,C3,-1
@@ -127,7 +127,7 @@ mod tests {
 
         assert_eq!(speech.len(), 100800);
         approx::assert_abs_diff_eq!(speech[2000], 17.15977345625943, epsilon = 1.0e-10);
-        approx::assert_abs_diff_eq!(speech[30000], 2566.2058730889985, epsilon = 1.0e-10);
+        approx::assert_abs_diff_eq!(speech[30000], 2566.205873089126, epsilon = 1.0e-10);
         approx::assert_abs_diff_eq!(speech[70000], -1898.2890228814217, epsilon = 1.0e-10);
         approx::assert_abs_diff_eq!(speech[100799], -13.514971382534956, epsilon = 1.0e-10);
     }

src/mlpg_adjust/mlpg.rs

@@ -15,7 +15,7 @@ pub struct MlpgMatrix {
     win_size: usize,
     length: usize,
     width: usize,
-    wuw: Vec<Vec<f64>>,
+    wuw: Vec<f64>,
     wum: Vec<f64>,
 }
 
@@ -25,36 +25,31 @@ impl MlpgMatrix {
     pub fn calc_wuw_and_wum(windows: &Windows, parameters: Vec<Vec<MeanVari>>) -> Self {
         let length = parameters[0].len();
         let width = windows.max_width() * 2 + 1;
-        let mut wum = Vec::with_capacity(length);
-        let mut wuw = Vec::with_capacity(length);
+        let mut wum = vec![0.0; length];
+        let mut wuw = vec![0.0; length * width];
 
-        for t in 0..length {
-            wuw.push(vec![0.0; width]);
-            wum.push(0.0);
+        for (window, parameter) in std::iter::zip(windows, &parameters) {
+            let parameter = &parameter[..length];
 
-            for (i, window) in windows.iter().enumerate() {
-                for (index, coef) in window.iter_rev(0) {
+            for t in 0..length {
+                for (index, coef) in window.iter(window.left_width()) {
                     if coef == 0.0 {
                         continue;
                     }
 
-                    let idx = (t as isize) - index.position();
+                    let idx = (t as isize) - index;
                     if idx < 0 || idx >= length as isize {
                         continue;
                     }
-                    let wu = coef * parameters[i][idx as usize].1;
-                    wum[t] += wu * parameters[i][idx as usize].0;
+                    let MeanVari(mean, vari) = parameter[idx as usize];
+                    wum[t] += coef * vari * mean;
 
-                    for (inner_index, coef) in window.iter_rev(index.index()) {
-                        if coef == 0.0 {
+                    for (inner_index, inner_coef) in window.iter(index) {
+                        if inner_coef == 0.0 {
                             continue;
                         }
-                        let j = inner_index.index() - index.index();
-                        if t + j >= length {
-                            break;
-                        }
-
-                        wuw[t][j] += wu * coef;
+                        let j = (inner_index - index) as usize;
+                        wuw[t * width + j] += coef * inner_coef * vari;
                     }
                 }
             }
@@ -77,43 +72,59 @@ impl MlpgMatrix {
 
     /// Perform Cholesky decomposition.
     fn ldl_factorization(&mut self) {
-        for t in 0..self.length {
-            for i in 1..self.width.min(t + 1) {
-                self.wuw[t][0] -= self.wuw[t - i][i] * self.wuw[t - i][i] * self.wuw[t - i][0];
+        let Self { width, length, .. } = *self;
+        let wuw = &mut self.wuw[..length * width];
+        for t in 0..length {
+            for i in 1..width.min(t + 1) {
+                wuw[width * t] -=
+                    wuw[(t - i) * width + i] * wuw[(t - i) * width + i] * wuw[(t - i) * width];
             }
-            for i in 1..self.width {
-                for j in 1..(self.width - i).min(t + 1) {
-                    self.wuw[t][i] -=
-                        self.wuw[t - j][j] * self.wuw[t - j][i + j] * self.wuw[t - j][0];
+            for i in 1..width {
+                for j in 1..(width - i).min(t + 1) {
+                    wuw[width * t + i] -= wuw[(t - j) * width + j]
+                        * wuw[(t - j) * width + i + j]
+                        * wuw[(t - j) * width];
                 }
-                self.wuw[t][i] /= self.wuw[t][0];
+                wuw[width * t + i] /= wuw[width * t];
             }
         }
     }
 
     /// Forward & backward substitution.
     fn substitutions(&self) -> Vec<f64> {
+        let Self { width, length, .. } = *self;
+        let wum = &self.wum[..length];
+        let wuw = &self.wuw[..length * width];
         let mut g = vec![0.0; self.length];
         // forward
-        for t in 0..self.length {
-            g[t] = self.wum[t];
-            for i in 1..self.width.min(t + 1) {
-                g[t] -= self.wuw[t - i][i] * g[t - i];
+        for t in 0..length {
+            g[t] = wum[t];
+            for i in 1..width.min(t + 1) {
+                g[t] -= wuw[(t - i) * width + i] * g[t - i];
             }
         }
 
         let mut par = vec![0.0; self.length];
         // backward
         for t in (0..self.length).rev() {
-            par[t] = g[t] / self.wuw[t][0];
-            for i in 1..self.width.min(self.length - t) {
-                par[t] -= self.wuw[t][i] * par[t + i];
+            par[t] = g[t] / wuw[t * width];
+            for i in 1..width.min(length - t) {
+                par[t] -= wuw[t * width + i] * par[t + i];
             }
         }
 
         par
     }
 
+    fn calculate_gv_switch(gv_switch: &[bool], durations: &[usize], mask: &[bool]) -> Vec<bool> {
+        gv_switch
+            .iter()
+            .copied()
+            .duration(durations)
+            .filter_by(mask)
+            .collect()
+    }
+
     /// Solve the equasion, and if necessary, applies GV (global variance).
     pub fn par(
         &mut self,
@@ -126,12 +137,8 @@ impl MlpgMatrix {
         if let Some((gv_param, gv_switch)) = gv {
             let mtx_before = self.clone();
             let par = self.solve();
-            let gv_switch: Vec<_> = gv_switch
-                .iter()
-                .copied()
-                .duration(durations)
-                .filter_by(msd_flag.mask())
-                .collect();
+            let gv_switch: Vec<_> =
+                Self::calculate_gv_switch(gv_switch, durations, msd_flag.mask());
             let mgv = MlpgGlobalVariance::new(mtx_before, par, &gv_switch);
 
             let MeanVari(gv_mean, gv_vari) = gv_param[vector_index];
@@ -171,89 +178,100 @@ impl<'a> MlpgGlobalVariance<'a> {
     }
 
     fn calc_gv(&self) -> (f64, f64) {
-        let mean = self
-            .par
-            .iter()
-            .zip(self.gv_switch.iter())
-            .filter(|(_, sw)| **sw)
-            .map(|(p, _)| *p)
-            .sum::<f64>()
-            / self.gv_length as f64;
-        let vari = self
-            .par
-            .iter()
-            .zip(self.gv_switch.iter())
-            .filter(|(_, sw)| **sw)
-            .map(|(p, _)| (*p - mean) * (*p - mean))
-            .sum::<f64>()
-            / self.gv_length as f64;
+        let mut sum = 0.0;
+        let mut sum_quad = 0.0;
 
+        for (par, sw) in std::iter::zip(&self.par, self.gv_switch) {
+            if *sw {
+                sum += *par;
+                sum_quad += *par * *par;
+            }
+        }
+
+        let mean = sum / self.gv_length as f64;
+        let vari = (sum_quad / self.gv_length as f64) - (mean * mean);
         (mean, vari)
     }
 
     /// Adjust parameter's deviation from mean value using gv_mean
     fn conv_gv(&mut self, gv_mean: f64) {
         let (mean, vari) = self.calc_gv();
         let ratio = (gv_mean / vari).sqrt();
-        self.par
-            .iter_mut()
-            .zip(self.gv_switch.iter())
-            .filter(|(_, sw)| **sw)
-            .for_each(|(p, _)| *p = ratio * (*p - mean) + mean);
+
+        for (par, sw) in std::iter::zip(&mut self.par, self.gv_switch) {
+            if *sw {
+                *par = ratio * (*par - mean) + mean;
+            }
+        }
     }
     fn calc_hmmobj_derivative(&self) -> (f64, Vec<f64>) {
-        let mut g = vec![0.0; self.mtx.length];
-
-        #[allow(clippy::needless_range_loop)]
-        for t in 0..self.mtx.length {
-            g[t] = self.mtx.wuw[t][0] * self.par[t];
-            for i in 1..self.mtx.width {
-                if t + i < self.mtx.length {
-                    g[t] += self.mtx.wuw[t][i] * self.par[t + i];
-                }
-                if t + 1 > i {
-                    g[t] += self.mtx.wuw[t - i][i] * self.par[t - i];
+        let MlpgMatrix {
+            width,
+            length,
+            win_size,
+            ..
+        } = self.mtx;
+        assert!(width >= 1); // required for `wuw[0]` access
+        let wuw = self.mtx.wuw.chunks_exact(width);
+        let wum = &self.mtx.wum[..length];
+        let par = &self.par[..length];
+
+        let mut g = vec![0.0; length];
+
+        // .zip(0..length) to help optimizer recognize t < length
+        for (wuw, t) in wuw.zip(0..length) {
+            g[t] += wuw[0] * par[t];
+            for i in 1..width {
+                if t + i < length {
+                    g[t] += wuw[i] * par[t + i];
+                    g[t + i] += wuw[i] * par[t];
                 }
             }
         }
 
-        let w = 1.0 / ((self.mtx.win_size * self.mtx.length) as f64);
+        let w = 1.0 / ((win_size * length) as f64);
         let mut hmmobj = 0.0;
 
-        #[allow(clippy::needless_range_loop)]
-        for t in 0..self.mtx.length {
-            hmmobj += W1 * w * self.par[t] * (self.mtx.wum[t] - 0.5 * g[t]);
+        for t in 0..length {
+            hmmobj += W1 * w * par[t] * (wum[t] - 0.5 * g[t]);
         }
 
         (hmmobj, g)
     }
     fn next_step(
         &mut self,
-        g: Vec<f64>,
+        g: &[f64],
         step: f64,
         mean: f64,
         vari: f64,
         gv_mean: f64,
         gv_vari: f64,
     ) {
-        let length = self.mtx.length;
+        let MlpgMatrix { width, length, .. } = self.mtx;
 
         let w = 1.0 / ((self.mtx.win_size * length) as f64);
         let dv = -2.0 * gv_vari * (vari - gv_mean) / self.mtx.length as f64;
 
-        #[allow(clippy::needless_range_loop)]
-        for t in 0..length {
-            let h = -W1 * w * self.mtx.wuw[t][0]
+        assert!(width >= 1); // required for `wuw[0]` access
+        let wuw = self.mtx.wuw.chunks_exact(width);
+        let wum = &self.mtx.wum[..length];
+        let par = &mut self.par[..length];
+        let gv_switch = &self.gv_switch[..length];
+        let g = &g[..length];
+
+        // .zip(0..length) to help optimizer recognize t < length
+        for (wuw, t) in wuw.zip(0..length) {
+            let h = -W1 * w * wuw[0]
                 - W2 * 2.0 / (length * length) as f64
                     * ((length - 1) as f64 * gv_vari * (vari - gv_mean)
-                        + 2.0 * gv_vari * (self.par[t] - mean) * (self.par[t] - mean));
-            let next_g = if self.gv_switch[t] {
-                1.0 / h * (W1 * w * (-g[t] + self.mtx.wum[t]) + W2 * dv * (self.par[t] - mean))
+                        + 2.0 * gv_vari * (par[t] - mean) * (par[t] - mean));
+            let next_g = if gv_switch[t] {
+                1.0 / h * (W1 * w * (-g[t] + wum[t]) + W2 * dv * (par[t] - mean))
             } else {
-                1.0 / h * (W1 * w * (-g[t] + self.mtx.wum[t]))
+                1.0 / h * (W1 * w * (-g[t] + wum[t]))
             };
 
-            self.par[t] += step * next_g;
+            par[t] += step * next_g;
         }
     }
 
@@ -285,7 +303,7 @@ impl<'a> MlpgGlobalVariance<'a> {
                 }
             }
 
-            self.next_step(g, step, mean, vari, gv_mean, gv_vari);
+            self.next_step(&g, step, mean, vari, gv_mean, gv_vari);
 
             prev = obj;
         }