From 7ab0de1fe907b16e3ce6c7ff0e15d27f71e5d5b4 Mon Sep 17 00:00:00 2001
From: Josh Holmer <jholmer.in@gmail.com>
Date: Tue, 11 Apr 2023 18:28:09 -0400
Subject: [PATCH] Add benchmarks for motion_compensate function

This was identified as a hot function, so a benchmark will help us to
optimize it.
---
 benches/bench.rs |  6 ++--
 benches/mc.rs    | 83 +++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/benches/bench.rs b/benches/bench.rs
index fd767b9584..6e48e395f4 100644
--- a/benches/bench.rs
+++ b/benches/bench.rs
@@ -49,7 +49,7 @@ fn write_b_bench(b: &mut Bencher, tx_size: TxSize, qindex: usize) {
     speed_settings: SpeedSettings::from_preset(10),
     ..Default::default()
   });
-  let sequence = Arc::new(Sequence::new(&Default::default()));
+  let sequence = Arc::new(Sequence::new(&config));
   let fi = FrameInvariants::<u16>::new(config, sequence);
   let mut w = WriterEncoder::new();
   let mut fc = CDFContext::new(fi.base_q_idx);
@@ -138,7 +138,7 @@ fn cdef_frame_bench(b: &mut Bencher, width: usize, height: usize) {
     speed_settings: SpeedSettings::from_preset(10),
     ..Default::default()
   });
-  let sequence = Arc::new(Sequence::new(&Default::default()));
+  let sequence = Arc::new(Sequence::new(&config));
   let fi = FrameInvariants::<u16>::new(config, sequence);
   let fb = FrameBlocks::new(fi.sb_width * 16, fi.sb_height * 16);
   let mut fs = FrameState::new(&fi);
@@ -170,7 +170,7 @@ fn cfl_rdo_bench(b: &mut Bencher, bsize: BlockSize) {
     speed_settings: SpeedSettings::from_preset(10),
     ..Default::default()
   });
-  let sequence = Arc::new(Sequence::new(&Default::default()));
+  let sequence = Arc::new(Sequence::new(&config));
   let fi = FrameInvariants::<u16>::new(config, sequence);
   let mut fs = FrameState::new(&fi);
   fs.apply_tile_state_mut(|ts| {
diff --git a/benches/mc.rs b/benches/mc.rs
index 88072d2abb..38c6fdbb3a 100644
--- a/benches/mc.rs
+++ b/benches/mc.rs
@@ -3,11 +3,20 @@
 use criterion::*;
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaChaRng;
+use rav1e::bench::context::BlockContext;
+use rav1e::bench::context::CDFContext;
+use rav1e::bench::context::ContextWriter;
+use rav1e::bench::context::FrameBlocks;
 use rav1e::bench::cpu_features::*;
+use rav1e::bench::encoder::FrameState;
+use rav1e::bench::encoder::{motion_compensate, FrameInvariants};
 use rav1e::bench::frame::{AsRegion, PlaneOffset, PlaneSlice};
 use rav1e::bench::mc::*;
+use rav1e::bench::partition::RefType;
 use rav1e::bench::util::Aligned;
+use rav1e::context::{BlockOffset, TileBlockOffset};
 use rav1e::prelude::*;
+use std::sync::Arc;
 
 fn bench_put_8tap_top_left_lbd(c: &mut Criterion) {
   let mut ra = ChaChaRng::from_seed([0; 32]);
@@ -505,6 +514,76 @@ fn bench_prep_8tap_center_hbd(c: &mut Criterion) {
   });
 }
 
+fn bench_motion_compensate(c: &mut Criterion) {
+  let config = Arc::new(EncoderConfig {
+    width: 640,
+    height: 480,
+    bit_depth: 8,
+    speed_settings: SpeedSettings::from_preset(2),
+    ..Default::default()
+  });
+  let sequence = Arc::new(Sequence::new(&config));
+  let fi = FrameInvariants::<u8>::new(config, sequence);
+  let mut fc = CDFContext::new(fi.base_q_idx);
+  let mut fb = FrameBlocks::new(fi.sb_width * 16, fi.sb_height * 16);
+  let mut tb = fb.as_tile_blocks_mut();
+  let bc = BlockContext::new(&mut tb);
+  let mut fs = FrameState::new(&fi);
+  let mut cw = ContextWriter::new(&mut fc, bc);
+  fs.apply_tile_state_mut(|ts| {
+    c.bench_function("motion_compensate", |b| {
+      b.iter(|| {
+        let _ = black_box(motion_compensate(
+          &fi,
+          ts,
+          &mut cw,
+          PredictionMode::NEAREST_NEARESTMV,
+          [RefType::LAST_FRAME, RefType::ALTREF_FRAME],
+          [MotionVector { row: 0, col: 0 }, MotionVector { row: 0, col: 0 }],
+          BlockSize::BLOCK_32X16,
+          TileBlockOffset(BlockOffset { x: 152, y: 88 }),
+          false,
+        ));
+      })
+    });
+  });
+}
+
+fn bench_motion_compensate_hbd(c: &mut Criterion) {
+  let config = Arc::new(EncoderConfig {
+    width: 640,
+    height: 480,
+    bit_depth: 10,
+    speed_settings: SpeedSettings::from_preset(2),
+    ..Default::default()
+  });
+  let sequence = Arc::new(Sequence::new(&config));
+  let fi = FrameInvariants::<u16>::new(config, sequence);
+  let mut fc = CDFContext::new(fi.base_q_idx);
+  let mut fb = FrameBlocks::new(fi.sb_width * 16, fi.sb_height * 16);
+  let mut tb = fb.as_tile_blocks_mut();
+  let bc = BlockContext::new(&mut tb);
+  let mut fs = FrameState::new(&fi);
+  let mut cw = ContextWriter::new(&mut fc, bc);
+  fs.apply_tile_state_mut(|ts| {
+    c.bench_function("motion_compensate_hbd", |b| {
+      b.iter(|| {
+        let _ = black_box(motion_compensate(
+          &fi,
+          ts,
+          &mut cw,
+          PredictionMode::NEAREST_NEARESTMV,
+          [RefType::LAST_FRAME, RefType::ALTREF_FRAME],
+          [MotionVector { row: 0, col: 0 }, MotionVector { row: 0, col: 0 }],
+          BlockSize::BLOCK_32X16,
+          TileBlockOffset(BlockOffset { x: 152, y: 88 }),
+          false,
+        ));
+      })
+    });
+  });
+}
+
 criterion_group!(
   mc,
   bench_put_8tap_top_left_lbd,
@@ -522,7 +601,9 @@ criterion_group!(
   bench_prep_8tap_top_left_hbd,
   bench_prep_8tap_top_hbd,
   bench_prep_8tap_left_hbd,
-  bench_prep_8tap_center_hbd
+  bench_prep_8tap_center_hbd,
+  bench_motion_compensate,
+  bench_motion_compensate_hbd,
 );
 
 fn fill_plane<T: Pixel>(ra: &mut ChaChaRng, plane: &mut Plane<T>) {