diff --git a/benches/me.rs b/benches/me.rs index 0beac7b176bbacf50a08913edd66c317b21b5d6a..84f9fcc2afe8dfb5f35f52ab5c4bec457c442482 100644 --- a/benches/me.rs +++ b/benches/me.rs @@ -51,6 +51,26 @@ fn bench_get_sad(b: &mut Bencher, bs: &BlockSize) { }) } +fn bench_get_sad_iter(b: &mut Bencher, bs: &BlockSize) { + let mut ra = ChaChaRng::from_seed([0; 32]); + let bsw = bs.width(); + let bsh = bs.height(); + let w = 640; + let h = 480; + let input_plane = new_plane(&mut ra, w, h); + let rec_plane = new_plane(&mut ra, w, h); + let po = PlaneOffset { x: 0, y: 0 }; + + let mut plane_org = input_plane.slice(&po); + let mut plane_ref = rec_plane.slice(&po); + + b.iter(|| { + let _ = me::get_sad_iter(&mut plane_org, &mut plane_ref, bsw, bsh); + plane_org.y = 0; + plane_ref.y = 0; + }) +} + pub fn get_sad(c: &mut Criterion) { use partition::BlockSize::*; let blocks = vec![ @@ -78,5 +98,6 @@ pub fn get_sad(c: &mut Criterion) { BLOCK_64X16, ]; - c.bench_function_over_inputs("get_sad", bench_get_sad, blocks); + c.bench_function_over_inputs("get_sad", bench_get_sad, blocks.clone()); + c.bench_function_over_inputs("get_sad_iter", bench_get_sad_iter, blocks); } diff --git a/src/me.rs b/src/me.rs index bd81b9844f5f63afe1d860db7905dd19df4b0694..60e0a45f15a5b3b52806a45a467e901123a17491 100644 --- a/src/me.rs +++ b/src/me.rs @@ -39,6 +39,26 @@ pub fn get_sad( sum } +#[inline(always)] +pub fn get_sad_iter( + plane_org: &mut PlaneSlice, plane_ref: &mut PlaneSlice, blk_h: usize, + blk_w: usize +) -> u32 { + let mut sum = 0 as u32; + + let org_iter = plane_org.iter_width(blk_w); + let ref_iter = plane_ref.iter_width(blk_w); + + for (slice_org, slice_ref) in org_iter.take(blk_h).zip(ref_iter) { + sum += slice_org + .iter() + .zip(slice_ref) + .map(|(&a, &b)| (a as i32 - b as i32).abs() as u32) + .sum::<u32>(); + } + + sum +} pub fn motion_estimation( fi: &FrameInvariants, fs: &FrameState, bsize: BlockSize, bo: &BlockOffset, ref_frame: usize, pmv: &MotionVector diff --git a/src/plane.rs b/src/plane.rs index 42023b82b17372221166861c3188da18d632d4be..9c38ac7ee570e17bc00a49004acfac68d4e0364e 100644 --- a/src/plane.rs +++ b/src/plane.rs @@ -220,12 +220,38 @@ impl Plane { } } +#[derive(Clone, Copy)] pub struct PlaneSlice<'a> { pub plane: &'a Plane, pub x: isize, pub y: isize } +pub struct IterWidth<'a> { + ps: PlaneSlice<'a>, + width: usize, +} + +// TODO: Implement more methods +impl<'a> Iterator for IterWidth<'a> { + type Item = &'a [u16]; + + #[inline] + fn next(&mut self) -> Option<&'a [u16]> { + let x = self.ps.plane.cfg.xorigin + self.ps.x; + let y = self.ps.plane.cfg.yorigin + self.ps.y; + let stride = self.ps.plane.cfg.stride; + let base = y as usize * stride + x as usize; + + if self.ps.plane.data.len() < base + self.width { + None + } else { + self.ps.y += 1; + Some(&self.ps.plane.data[base..base + self.width]) + } + } +} + impl<'a> PlaneSlice<'a> { pub fn as_slice(&'a self) -> &'a [u16] { let stride = self.plane.cfg.stride; @@ -252,6 +278,10 @@ impl<'a> PlaneSlice<'a> { &self.plane.data[base..base + width] } + pub fn iter_width(&self, width: usize) -> IterWidth<'a> { + IterWidth { ps: *self, width } + } + pub fn subslice(&'a self, xo: usize, yo: usize) -> PlaneSlice<'a> { PlaneSlice { plane: self.plane,