Use a RwLock in frame_me_stats

src/api/internal.rs

@@ -743,7 +743,7 @@ impl<T: Pixel> ContextInner<T> {
       // backwards lower reference (so the closest previous frame).
       let index = if second_ref_frame.to_index() != 0 { 0 } else { 1 };
 
-      let me_stats = &fs.frame_me_stats[index];
+      let me_stats = &fs.frame_me_stats.read().expect("poisoned lock")[index];
       use byteorder::{NativeEndian, WriteBytesExt};
       // dynamic allocation: debugging only
       let mut buf = vec![];
@@ -1084,68 +1084,70 @@ impl<T: Pixel> ContextInner<T> {
       // is more performant because it avoids a very expensive clone.
       let output_frame_data =
         self.frame_data.remove(&output_frameno).unwrap().unwrap();
-      let fi = &output_frame_data.fi;
-      let fs = &output_frame_data.fs;
-
-      let frame = self.frame_q[&fi.input_frameno].as_ref().unwrap();
+      {
+        let fi = &output_frame_data.fi;
+        let fs = &output_frame_data.fs;
 
-      // There can be at most 3 of these.
-      let mut unique_indices = ArrayVec::<_, 3>::new();
+        let frame = self.frame_q[&fi.input_frameno].as_ref().unwrap();
 
-      for (mv_index, &rec_index) in fi.ref_frames.iter().enumerate() {
-        if !unique_indices.iter().any(|&(_, r)| r == rec_index) {
-          unique_indices.push((mv_index, rec_index));
-        }
-      }
+        // There can be at most 3 of these.
+        let mut unique_indices = ArrayVec::<_, 3>::new();
 
-      let bit_depth = self.config.bit_depth;
-      let frame_data = &mut self.frame_data;
-      let len = unique_indices.len();
-
-      let lookahead_me_stats = &fs.frame_me_stats;
-
-      // Compute and propagate the importance, split evenly between the
-      // referenced frames.
-      unique_indices.iter().for_each(|&(mv_index, rec_index)| {
-        // Use rec_buffer here rather than lookahead_rec_buffer because
-        // rec_buffer still contains the reference frames for the current frame
-        // (it's only overwritten when the frame is encoded), while
-        // lookahead_rec_buffer already contains reference frames for the next
-        // frame (for the reference propagation to work correctly).
-        let reference =
-          fi.rec_buffer.frames[rec_index as usize].as_ref().unwrap();
-        let reference_frame = &reference.frame;
-        let reference_output_frameno = reference.output_frameno;
-        let me_stats = &lookahead_me_stats[mv_index];
-
-        // We should never use frame as its own reference.
-        assert_ne!(reference_output_frameno, output_frameno);
-
-        if let Some(reference_frame_block_importances) =
-          frame_data.get_mut(&reference_output_frameno).map(|data| {
-            &mut data
-              .as_mut()
-              .unwrap()
-              .fi
-              .coded_frame_data
-              .as_mut()
-              .unwrap()
-              .block_importances
-          })
-        {
-          Self::update_block_importances(
-            fi,
-            me_stats,
-            frame,
-            reference_frame,
-            bit_depth,
-            bsize,
-            len,
-            reference_frame_block_importances,
-          );
+        for (mv_index, &rec_index) in fi.ref_frames.iter().enumerate() {
+          if !unique_indices.iter().any(|&(_, r)| r == rec_index) {
+            unique_indices.push((mv_index, rec_index));
+          }
         }
-      });
 
+        let bit_depth = self.config.bit_depth;
+        let frame_data = &mut self.frame_data;
+        let len = unique_indices.len();
+
+        let lookahead_me_stats =
+          fs.frame_me_stats.read().expect("poisoned lock");
+
+        // Compute and propagate the importance, split evenly between the
+        // referenced frames.
+        unique_indices.iter().for_each(|&(mv_index, rec_index)| {
+          // Use rec_buffer here rather than lookahead_rec_buffer because
+          // rec_buffer still contains the reference frames for the current frame
+          // (it's only overwritten when the frame is encoded), while
+          // lookahead_rec_buffer already contains reference frames for the next
+          // frame (for the reference propagation to work correctly).
+          let reference =
+            fi.rec_buffer.frames[rec_index as usize].as_ref().unwrap();
+          let reference_frame = &reference.frame;
+          let reference_output_frameno = reference.output_frameno;
+          let me_stats = &lookahead_me_stats[mv_index];
+
+          // We should never use frame as its own reference.
+          assert_ne!(reference_output_frameno, output_frameno);
+
+          if let Some(reference_frame_block_importances) =
+            frame_data.get_mut(&reference_output_frameno).map(|data| {
+              &mut data
+                .as_mut()
+                .unwrap()
+                .fi
+                .coded_frame_data
+                .as_mut()
+                .unwrap()
+                .block_importances
+            })
+          {
+            Self::update_block_importances(
+              fi,
+              me_stats,
+              frame,
+              reference_frame,
+              bit_depth,
+              bsize,
+              len,
+              reference_frame_block_importances,
+            );
+          }
+        });
+      }
       self.frame_data.insert(output_frameno, Some(output_frame_data));
     }
 

src/api/lookahead.rs

@@ -7,8 +7,8 @@ use crate::encoder::{
   FrameInvariants, FrameState, Sequence, IMPORTANCE_BLOCK_SIZE,
 };
 use crate::frame::{AsRegion, PlaneOffset};
-use crate::me::{estimate_tile_motion, FrameMEStats};
-use crate::partition::{get_intra_edges, BlockSize, REF_FRAMES};
+use crate::me::{estimate_tile_motion, RefMEStats};
+use crate::partition::{get_intra_edges, BlockSize};
 use crate::predict::{IntraParam, PredictionMode};
 use crate::rayon::iter::*;
 use crate::tiling::{Area, PlaneRegion, TileRect};
@@ -178,8 +178,7 @@ pub(crate) fn estimate_importance_block_difference<T: Pixel>(
 #[hawktracer(estimate_inter_costs)]
 pub(crate) fn estimate_inter_costs<T: Pixel>(
   frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>, bit_depth: usize,
-  mut config: EncoderConfig, sequence: Arc<Sequence>,
-  buffer: Arc<[FrameMEStats; REF_FRAMES]>,
+  mut config: EncoderConfig, sequence: Arc<Sequence>, buffer: RefMEStats,
 ) -> f64 {
   config.low_latency = true;
   config.speed_settings.multiref = false;
@@ -210,7 +209,7 @@ pub(crate) fn estimate_inter_costs<T: Pixel>(
   let plane_ref = &ref_frame.planes[0];
   let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
   let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
-  let stats = &fs.frame_me_stats[0];
+  let stats = &fs.frame_me_stats.read().expect("poisoned lock")[0];
   let bsize = BlockSize::from_width_and_height(
     IMPORTANCE_BLOCK_SIZE,
     IMPORTANCE_BLOCK_SIZE,

src/encoder.rs

@@ -87,7 +87,7 @@ pub struct ReferenceFrame<T: Pixel> {
   pub input_hres: Arc<Plane<T>>,
   pub input_qres: Arc<Plane<T>>,
   pub cdfs: CDFContext,
-  pub frame_me_stats: Arc<[FrameMEStats; REF_FRAMES as usize]>,
+  pub frame_me_stats: RefMEStats,
   pub output_frameno: u64,
   pub segmentation: SegmentationState,
 }
@@ -436,7 +436,7 @@ pub struct FrameState<T: Pixel> {
   pub restoration: RestorationState,
   // Because we only reference these within a tile context,
   // these are stored per-tile for easier access.
-  pub frame_me_stats: Arc<[FrameMEStats; REF_FRAMES as usize]>,
+  pub frame_me_stats: RefMEStats,
   pub enc_stats: EncoderStats,
 }
 
@@ -457,8 +457,8 @@ impl<T: Pixel> FrameState<T> {
   /// it does not create hres or qres versions of `frame` as downscaling is
   /// somewhat expensive and are not needed for [`estimate_inter_costs`].
   pub fn new_with_frame_and_me_stats_and_rec(
-    fi: &FrameInvariants<T>, frame: Arc<Frame<T>>,
-    me_stats: Arc<[FrameMEStats; REF_FRAMES]>, rec: Arc<Frame<T>>,
+    fi: &FrameInvariants<T>, frame: Arc<Frame<T>>, me_stats: RefMEStats,
+    rec: Arc<Frame<T>>,
   ) -> Self {
     let rs = RestorationState::new(fi, &frame);
 
@@ -519,8 +519,16 @@ impl<T: Pixel> FrameState<T> {
   {
     let PlaneConfig { width, height, .. } = self.rec.planes[0].cfg;
     let sbo_0 = PlaneSuperBlockOffset(SuperBlockOffset { x: 0, y: 0 });
-    let ts =
-      &mut TileStateMut::new(self, sbo_0, self.sb_size_log2, width, height);
+    let frame_me_stats = self.frame_me_stats.clone();
+    let frame_me_stats = &mut *frame_me_stats.write().expect("poisoned lock");
+    let ts = &mut TileStateMut::new(
+      self,
+      sbo_0,
+      self.sb_size_log2,
+      width,
+      height,
+      frame_me_stats,
+    );
 
     f(ts)
   }

src/me.rs

@@ -13,7 +13,6 @@ use crate::context::{
   MI_SIZE_LOG2, SB_SIZE,
 };
 use crate::dist::*;
-use crate::encoder::ReferenceFrame;
 use crate::frame::*;
 use crate::mc::MotionVector;
 use crate::partition::*;
@@ -27,7 +26,7 @@ use arrayvec::*;
 use crate::api::InterConfig;
 use crate::util::ILog;
 use std::ops::{Index, IndexMut};
-use std::sync::Arc;
+use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 use rust_hawktracer::*;
 
@@ -45,6 +44,15 @@ pub struct FrameMEStats {
   pub rows: usize,
 }
 
+/// cbindgen:ignore
+pub type RefMEStats = Arc<RwLock<[FrameMEStats; REF_FRAMES]>>;
+/// cbindgen:ignore
+pub type ReadGuardMEStats<'a> =
+  RwLockReadGuard<'a, [FrameMEStats; REF_FRAMES]>;
+/// cbindgen:ignore
+pub type WriteGuardMEStats<'a> =
+  RwLockWriteGuard<'a, [FrameMEStats; REF_FRAMES]>;
+
 impl FrameMEStats {
   pub fn new(cols: usize, rows: usize) -> Self {
     Self {
@@ -54,8 +62,8 @@ impl FrameMEStats {
       rows,
     }
   }
-  pub fn new_arc_array(cols: usize, rows: usize) -> Arc<[Self; REF_FRAMES]> {
-    Arc::new([
+  pub fn new_arc_array(cols: usize, rows: usize) -> RefMEStats {
+    Arc::new(RwLock::new([
       FrameMEStats::new(cols, rows),
       FrameMEStats::new(cols, rows),
       FrameMEStats::new(cols, rows),
@@ -64,7 +72,7 @@ impl FrameMEStats {
       FrameMEStats::new(cols, rows),
       FrameMEStats::new(cols, rows),
       FrameMEStats::new(cols, rows),
-    ])
+    ]))
   }
 }
 
@@ -373,9 +381,9 @@ impl MotionEstimationSubsets {
   }
 }
 
-fn get_subset_predictors<T: Pixel>(
+fn get_subset_predictors(
   tile_bo: TileBlockOffset, tile_me_stats: &TileMEStats<'_>,
-  frame_ref_opt: Option<&ReferenceFrame<T>>, ref_frame_id: usize,
+  frame_ref_opt: Option<ReadGuardMEStats<'_>>, ref_frame_id: usize,
   pix_w: usize, pix_h: usize, mvx_min: isize, mvx_max: isize, mvy_min: isize,
   mvy_max: isize, corner: MVSamplingMode, ssdec: u8,
 ) -> MotionEstimationSubsets {
@@ -464,8 +472,8 @@ fn get_subset_predictors<T: Pixel>(
   // blocks of the previous frame.
   // Sample the middle of this block in the previous frame.
 
-  if let Some(frame_ref) = frame_ref_opt {
-    let prev_frame = &frame_ref.frame_me_stats[ref_frame_id];
+  if let Some(frame_me_stats) = frame_ref_opt {
+    let prev_frame = &frame_me_stats[ref_frame_id];
 
     let frame_bo = PlaneBlockOffset(BlockOffset {
       x: tile_me_stats.x() + tile_bo.0.x,
@@ -689,8 +697,9 @@ fn full_pixel_me<T: Pixel>(
 ) -> MotionSearchResult {
   let ref_frame_id = ref_frame.to_index();
   let tile_me_stats = &ts.me_stats[ref_frame_id].as_const();
-  let frame_ref =
-    fi.rec_buffer.frames[fi.ref_frames[0] as usize].as_ref().map(Arc::as_ref);
+  let frame_ref = fi.rec_buffer.frames[fi.ref_frames[0] as usize]
+    .as_ref()
+    .map(|frame_ref| frame_ref.frame_me_stats.read().expect("poisoned lock"));
   let subsets = get_subset_predictors(
     tile_bo,
     tile_me_stats,

src/scenechange/mod.rs

@@ -14,8 +14,7 @@ use crate::api::{EncoderConfig, SceneDetectionSpeed};
 use crate::cpu_features::CpuFeatureLevel;
 use crate::encoder::Sequence;
 use crate::frame::*;
-use crate::me::FrameMEStats;
-use crate::partition::REF_FRAMES;
+use crate::me::RefMEStats;
 use crate::util::Pixel;
 use debug_unreachable::debug_unreachable;
 use rust_hawktracer::*;
@@ -73,7 +72,7 @@ pub struct SceneChangeDetector<T: Pixel> {
     bool,
   )>,
   /// Buffer for FrameMEStats for cost scenecut
-  frame_me_stats_buffer: Option<Arc<[FrameMEStats; REF_FRAMES]>>,
+  frame_me_stats_buffer: Option<RefMEStats>,
   /// Frame buffer for holding references to source frames.
   ///
   /// Useful for not copying data into the downscaled frame buffer

src/tiling/tile_state.rs

@@ -14,6 +14,8 @@ use crate::encoder::*;
 use crate::frame::*;
 use crate::lrf::{IntegralImageBuffer, SOLVE_IMAGE_SIZE};
 use crate::mc::MotionVector;
+use crate::me::FrameMEStats;
+use crate::me::WriteGuardMEStats;
 use crate::partition::{RefType, REF_FRAMES};
 use crate::predict::{InterCompoundBuffers, PredictionMode};
 use crate::quantize::*;
@@ -125,6 +127,7 @@ impl<'a, T: Pixel> TileStateMut<'a, T> {
   pub fn new(
     fs: &'a mut FrameState<T>, sbo: PlaneSuperBlockOffset,
     sb_size_log2: usize, width: usize, height: usize,
+    frame_me_stats: &'a mut [FrameMEStats],
   ) -> Self {
     debug_assert!(
       width % MI_SIZE == 0,
@@ -170,7 +173,7 @@ impl<'a, T: Pixel> TileStateMut<'a, T> {
         sb_width,
         sb_height,
       ),
-      me_stats: Arc::make_mut(&mut fs.frame_me_stats)
+      me_stats: frame_me_stats
         .iter_mut()
         .map(|fmvs| {
           TileMEStatsMut::new(