me.rs 33.1 KB
Newer Older
1 2 3 4 5 6 7 8 9
// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

10
#[cfg(all(target_arch = "x86_64", feature = "nasm"))]
Raphaël Zumer's avatar
Raphaël Zumer committed
11
pub use self::nasm::get_sad;
12
#[cfg(any(not(target_arch = "x86_64"), not(feature = "nasm")))]
Raphaël Zumer's avatar
Raphaël Zumer committed
13
pub use self::native::get_sad;
Raphaël Zumer's avatar
Raphaël Zumer committed
14
use crate::context::{BlockOffset, BLOCK_TO_PLANE_SHIFT, MI_SIZE};
15
use crate::encoder::ReferenceFrame;
Raphaël Zumer's avatar
Raphaël Zumer committed
16 17 18 19
use crate::FrameInvariants;
use crate::FrameState;
use crate::partition::*;
use crate::plane::*;
20
use crate::util::Pixel;
21

Romain Vimont's avatar
Romain Vimont committed
22
use std::ops::{Index, IndexMut};
23
use std::sync::Arc;
24

25
#[cfg(all(target_arch = "x86_64", feature = "nasm"))]
26
mod nasm {
Raphaël Zumer's avatar
Raphaël Zumer committed
27 28
  use crate::plane::*;
  use crate::util::*;
29
  use std::mem;
Luca Barbato's avatar
Luca Barbato committed
30

Raphaël Zumer's avatar
Raphaël Zumer committed
31 32
  use libc;

Luca Barbato's avatar
Luca Barbato committed
33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
  extern {
    fn rav1e_sad_4x4_hbd_ssse3(
      src: *const u16, src_stride: libc::ptrdiff_t, dst: *const u16,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad_8x8_hbd10_ssse3(
      src: *const u16, src_stride: libc::ptrdiff_t, dst: *const u16,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad_16x16_hbd_ssse3(
      src: *const u16, src_stride: libc::ptrdiff_t, dst: *const u16,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad_32x32_hbd10_ssse3(
      src: *const u16, src_stride: libc::ptrdiff_t, dst: *const u16,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad_64x64_hbd10_ssse3(
      src: *const u16, src_stride: libc::ptrdiff_t, dst: *const u16,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad_128x128_hbd10_ssse3(
      src: *const u16, src_stride: libc::ptrdiff_t, dst: *const u16,
      dst_stride: libc::ptrdiff_t
    ) -> u32;
63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92

    fn rav1e_sad4x4_sse2(
      src: *const u8, src_stride: libc::ptrdiff_t, dst: *const u8,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad8x8_sse2(
      src: *const u8, src_stride: libc::ptrdiff_t, dst: *const u8,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad16x16_sse2(
      src: *const u8, src_stride: libc::ptrdiff_t, dst: *const u8,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad32x32_sse2(
      src: *const u8, src_stride: libc::ptrdiff_t, dst: *const u8,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad64x64_sse2(
      src: *const u8, src_stride: libc::ptrdiff_t, dst: *const u8,
      dst_stride: libc::ptrdiff_t
    ) -> u32;

    fn rav1e_sad128x128_sse2(
      src: *const u8, src_stride: libc::ptrdiff_t, dst: *const u8,
      dst_stride: libc::ptrdiff_t
    ) -> u32;
Luca Barbato's avatar
Luca Barbato committed
93
  }
Kyle Siefring's avatar
Kyle Siefring committed
94

Luca Barbato's avatar
Luca Barbato committed
95
  #[target_feature(enable = "ssse3")]
96
  unsafe fn sad_hbd_ssse3(
97
    plane_org: &PlaneSlice<'_, u16>, plane_ref: &PlaneSlice<'_, u16>, blk_h: usize,
Luca Barbato's avatar
Luca Barbato committed
98 99 100
    blk_w: usize, bit_depth: usize
  ) -> u32 {
    let mut sum = 0 as u32;
101 102
    let org_stride = (plane_org.plane.cfg.stride * 2) as libc::ptrdiff_t;
    let ref_stride = (plane_ref.plane.cfg.stride * 2) as libc::ptrdiff_t;
Luca Barbato's avatar
Luca Barbato committed
103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118
    assert!(blk_h >= 4 && blk_w >= 4);
    let step_size =
      blk_h.min(blk_w).min(if bit_depth <= 10 { 128 } else { 4 });
    let func = match step_size.ilog() {
      3 => rav1e_sad_4x4_hbd_ssse3,
      4 => rav1e_sad_8x8_hbd10_ssse3,
      5 => rav1e_sad_16x16_hbd_ssse3,
      6 => rav1e_sad_32x32_hbd10_ssse3,
      7 => rav1e_sad_64x64_hbd10_ssse3,
      8 => rav1e_sad_128x128_hbd10_ssse3,
      _ => rav1e_sad_128x128_hbd10_ssse3
    };
    for r in (0..blk_h).step_by(step_size) {
      for c in (0..blk_w).step_by(step_size) {
        let org_slice = plane_org.subslice(c, r);
        let ref_slice = plane_ref.subslice(c, r);
Romain Vimont's avatar
Romain Vimont committed
119 120
        let org_ptr = org_slice.as_ptr();
        let ref_ptr = ref_slice.as_ptr();
121 122 123
        // FIXME for now, T == u16
        let org_ptr = org_ptr as *const u16;
        let ref_ptr = ref_ptr as *const u16;
Luca Barbato's avatar
Luca Barbato committed
124 125
        sum += func(org_ptr, org_stride, ref_ptr, ref_stride);
      }
Kyle Siefring's avatar
Kyle Siefring committed
126
    }
127
    sum
Kyle Siefring's avatar
Kyle Siefring committed
128 129
  }

130
  #[target_feature(enable = "sse2")]
131 132
  unsafe fn sad_sse2(
    plane_org: &PlaneSlice<'_, u8>, plane_ref: &PlaneSlice<'_, u8>, blk_h: usize,
133 134
    blk_w: usize
  ) -> u32 {
135 136
    let org_ptr = plane_org.as_ptr();
    let ref_ptr = plane_ref.as_ptr();
137 138
    let org_stride = plane_org.plane.cfg.stride as libc::ptrdiff_t;
    let ref_stride = plane_ref.plane.cfg.stride as libc::ptrdiff_t;
139 140 141 142 143 144 145
    if blk_w == 16 && blk_h == 16 && (org_ptr as usize & 15) == 0 {
      return rav1e_sad16x16_sse2(org_ptr, org_stride, ref_ptr, ref_stride);
    }
    // Note: unaligned blocks come from hres/qres ME search
    let ptr_align_log2 = (org_ptr as usize).trailing_zeros() as usize;
    // The largest unaligned-safe function is for 8x8
    let ptr_align = 1 << ptr_align_log2.max(3);
146
    let step_size = blk_h.min(blk_w).min(ptr_align);
147 148 149 150 151 152 153 154 155
    let func = match step_size.ilog() {
      3 => rav1e_sad4x4_sse2,
      4 => rav1e_sad8x8_sse2,
      5 => rav1e_sad16x16_sse2,
      6 => rav1e_sad32x32_sse2,
      7 => rav1e_sad64x64_sse2,
      8 => rav1e_sad128x128_sse2,
      _ => rav1e_sad128x128_sse2
    };
156 157 158 159 160
    let mut sum = 0 as u32;
    for r in (0..blk_h as isize).step_by(step_size) {
      for c in (0..blk_w as isize).step_by(step_size) {
        let org_ptr = org_ptr.offset(r * org_stride + c);
        let ref_ptr = ref_ptr.offset(r * ref_stride + c);
161 162 163 164 165 166
        sum += func(org_ptr, org_stride, ref_ptr, ref_stride);
      }
    }
    sum
  }

Luca Barbato's avatar
Luca Barbato committed
167
  #[inline(always)]
168 169
  pub fn get_sad<T: Pixel>(
    plane_org: &PlaneSlice<'_, T>, plane_ref: &PlaneSlice<'_, T>, blk_h: usize,
Luca Barbato's avatar
Luca Barbato committed
170 171
    blk_w: usize, bit_depth: usize
  ) -> u32 {
172
    #[cfg(all(target_arch = "x86_64", feature = "nasm"))]
Luca Barbato's avatar
Luca Barbato committed
173
    {
174
      if mem::size_of::<T>() == 2 && is_x86_feature_detected!("ssse3") && blk_h >= 4 && blk_w >= 4 {
Luca Barbato's avatar
Luca Barbato committed
175
        return unsafe {
176 177
          let plane_org = &*(plane_org as *const _ as *const PlaneSlice<'_, u16>);
          let plane_ref = &*(plane_ref as *const _ as *const PlaneSlice<'_, u16>);
178
          sad_hbd_ssse3(plane_org, plane_ref, blk_h, blk_w, bit_depth)
Luca Barbato's avatar
Luca Barbato committed
179 180
        };
      }
181 182
      if mem::size_of::<T>() == 1 && is_x86_feature_detected!("sse2") && blk_h >= 4 && blk_w >= 4 {
        return unsafe {
183 184
          let plane_org = &*(plane_org as *const _ as *const PlaneSlice<'_, u8>);
          let plane_ref = &*(plane_ref as *const _ as *const PlaneSlice<'_, u8>);
185 186 187
          sad_sse2(plane_org, plane_ref, blk_h, blk_w)
        };
      }
Kyle Siefring's avatar
Kyle Siefring committed
188
    }
Luca Barbato's avatar
Luca Barbato committed
189
    super::native::get_sad(plane_org, plane_ref, blk_h, blk_w, bit_depth)
Kyle Siefring's avatar
Kyle Siefring committed
190
  }
191 192 193
}

mod native {
Raphaël Zumer's avatar
Raphaël Zumer committed
194
  use crate::plane::*;
195
  use crate::util::*;
196

Luca Barbato's avatar
Luca Barbato committed
197
  #[inline(always)]
198 199
  pub fn get_sad<T: Pixel>(
    plane_org: &PlaneSlice<'_, T>, plane_ref: &PlaneSlice<'_, T>, blk_h: usize,
Luca Barbato's avatar
Luca Barbato committed
200 201 202
    blk_w: usize, _bit_depth: usize
  ) -> u32 {
    let mut sum = 0 as u32;
203

Luca Barbato's avatar
Luca Barbato committed
204 205
    let org_iter = plane_org.iter_width(blk_w);
    let ref_iter = plane_ref.iter_width(blk_w);
206

Luca Barbato's avatar
Luca Barbato committed
207
    for (slice_org, slice_ref) in org_iter.take(blk_h).zip(ref_iter) {
208 209 210
      sum += slice_org
        .iter()
        .zip(slice_ref)
211
        .map(|(&a, &b)| (i32::cast_from(a) - i32::cast_from(b)).abs() as u32)
212
        .sum::<u32>();
Luca Barbato's avatar
Luca Barbato committed
213
    }
214

Luca Barbato's avatar
Luca Barbato committed
215 216
    sum
  }
217 218
}

Romain Vimont's avatar
Romain Vimont committed
219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250
#[derive(Debug, Clone)]
pub struct FrameMotionVectors {
  mvs: Box<[MotionVector]>,
  pub cols: usize,
  pub rows: usize,
}

impl FrameMotionVectors {
  pub fn new(cols: usize, rows: usize) -> Self {
    Self {
      mvs: vec![MotionVector::default(); cols * rows].into_boxed_slice(),
      cols,
      rows,
    }
  }
}

impl Index<usize> for FrameMotionVectors {
  type Output = [MotionVector];
  #[inline]
  fn index(&self, index: usize) -> &Self::Output {
    &self.mvs[index * self.cols..(index + 1) * self.cols]
  }
}

impl IndexMut<usize> for FrameMotionVectors {
  #[inline]
  fn index_mut(&mut self, index: usize) -> &mut Self::Output {
    &mut self.mvs[index * self.cols..(index + 1) * self.cols]
  }
}

251
fn get_mv_range(
Romain Vimont's avatar
Romain Vimont committed
252
  w_in_b: usize, h_in_b: usize, bo: BlockOffset, blk_w: usize, blk_h: usize
253
) -> (isize, isize, isize, isize) {
254 255 256
  let border_w = 128 + blk_w as isize * 8;
  let border_h = 128 + blk_h as isize * 8;
  let mvx_min = -(bo.x as isize) * (8 * MI_SIZE) as isize - border_w;
257
  let mvx_max = (w_in_b - bo.x - blk_w / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_w;
258
  let mvy_min = -(bo.y as isize) * (8 * MI_SIZE) as isize - border_h;
259
  let mvy_max = (h_in_b - bo.y - blk_h / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_h;
260 261 262 263

  (mvx_min, mvx_max, mvy_min, mvy_max)
}

264
pub fn get_subset_predictors<T: Pixel>(
Romain Vimont's avatar
Romain Vimont committed
265
  bo: BlockOffset, cmv: MotionVector,
266
  w_in_b: usize, h_in_b: usize,
267
  frame_mvs: &FrameMotionVectors, frame_ref_opt: Option<&ReferenceFrame<T>>,
268
  ref_frame_id: usize
269 270 271
) -> (Vec<MotionVector>) {
  let mut predictors = Vec::new();

272 273 274 275 276 277
  // Zero motion vector
  predictors.push(MotionVector::default());

  // Coarse motion estimation.
  predictors.push(cmv.quantize_to_fullpel());

278 279
  // EPZS subset A and B predictors.

280
  let mut median_preds = Vec::new();
281
  if bo.x > 0 {
Romain Vimont's avatar
Romain Vimont committed
282
    let left = frame_mvs[bo.y][bo.x - 1];
283 284
    median_preds.push(left);
    if !left.is_zero() { predictors.push(left); }
285 286
  }
  if bo.y > 0 {
Romain Vimont's avatar
Romain Vimont committed
287
    let top = frame_mvs[bo.y - 1][bo.x];
288 289
    median_preds.push(top);
    if !top.is_zero() { predictors.push(top); }
290

291
    if bo.x < w_in_b - 1 {
Romain Vimont's avatar
Romain Vimont committed
292
      let top_right = frame_mvs[bo.y - 1][bo.x + 1];
293 294
      median_preds.push(top_right);
      if !top_right.is_zero() { predictors.push(top_right); }
295 296 297
    }
  }

298
  if !median_preds.is_empty() {
Vladimir Kazakov's avatar
Vladimir Kazakov committed
299
    let mut median_mv = MotionVector::default();
300
    for mv in median_preds.iter() {
301 302
      median_mv = median_mv + *mv;
    }
303 304 305
    median_mv = median_mv / (median_preds.len() as i16);
    let median_mv_quant = median_mv.quantize_to_fullpel();
    if !median_mv_quant.is_zero() { predictors.push(median_mv_quant); }
306 307 308 309 310
  }

  // EPZS subset C predictors.

  if let Some(ref frame_ref) = frame_ref_opt {
311
    let prev_frame_mvs = &frame_ref.frame_mvs[ref_frame_id];
312 313

    if bo.x > 0 {
Romain Vimont's avatar
Romain Vimont committed
314
      let left = prev_frame_mvs[bo.y][bo.x - 1];
315
      if !left.is_zero() { predictors.push(left); }
316 317
    }
    if bo.y > 0 {
Romain Vimont's avatar
Romain Vimont committed
318
      let top = prev_frame_mvs[bo.y - 1][bo.x];
319
      if !top.is_zero() { predictors.push(top); }
320
    }
321
    if bo.x < w_in_b - 1 {
Romain Vimont's avatar
Romain Vimont committed
322
      let right = prev_frame_mvs[bo.y][bo.x + 1];
323
      if !right.is_zero() { predictors.push(right); }
324
    }
325
    if bo.y < h_in_b - 1 {
Romain Vimont's avatar
Romain Vimont committed
326
      let bottom = prev_frame_mvs[bo.y + 1][bo.x];
327
      if !bottom.is_zero() { predictors.push(bottom); }
328 329
    }

330 331
    let previous = prev_frame_mvs[bo.y][bo.x];
    if !previous.is_zero() { predictors.push(previous); }
332 333 334 335 336
  }

  predictors
}

337
pub trait MotionEstimation {
338
  fn full_pixel_me<T: Pixel>(
339
    fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &ReferenceFrame<T>,
Romain Vimont's avatar
Romain Vimont committed
340
    bo: BlockOffset, lambda: u32,
341
    cmv: MotionVector, pmv: [MotionVector; 2],
342 343 344
    mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
    blk_w: usize, blk_h: usize, best_mv: &mut MotionVector,
    lowest_cost: &mut u64, ref_frame: usize
345
  );
346

347
  fn sub_pixel_me<T: Pixel>(
348
    fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &ReferenceFrame<T>,
Romain Vimont's avatar
Romain Vimont committed
349
    bo: BlockOffset, lambda: u32, pmv: [MotionVector; 2],
350 351
    mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
    blk_w: usize, blk_h: usize, best_mv: &mut MotionVector,
352
    lowest_cost: &mut u64, ref_frame: usize
353 354
  );

355 356
  fn motion_estimation<T: Pixel> (
    fi: &FrameInvariants<T>, fs: &FrameState<T>, bsize: BlockSize,
Romain Vimont's avatar
Romain Vimont committed
357
    bo: BlockOffset, ref_frame: usize, cmv: MotionVector,
358
    pmv: [MotionVector; 2]
359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375
  ) -> MotionVector {
    match fi.rec_buffer.frames[fi.ref_frames[ref_frame - LAST_FRAME] as usize]
    {
      Some(ref rec) => {
        let blk_w = bsize.width();
        let blk_h = bsize.height();
        let (mvx_min, mvx_max, mvy_min, mvy_max) =
          get_mv_range(fi.w_in_b, fi.h_in_b, bo, blk_w, blk_h);

        // 0.5 is a fudge factor
        let lambda = (fi.me_lambda * 256.0 * 0.5) as u32;

        // Full-pixel motion estimation

        let mut lowest_cost = std::u64::MAX;
        let mut best_mv = MotionVector::default();

376
        Self::full_pixel_me(fi, fs, rec, bo, lambda, cmv, pmv,
377 378
                           mvx_min, mvx_max, mvy_min, mvy_max, blk_w, blk_h,
                           &mut best_mv, &mut lowest_cost, ref_frame);
379

380
        Self::sub_pixel_me(fi, fs, rec, bo, lambda, pmv,
381
                           mvx_min, mvx_max, mvy_min, mvy_max, blk_w, blk_h,
382
                           &mut best_mv, &mut lowest_cost, ref_frame);
383 384

        best_mv
Frank Bossen's avatar
Frank Bossen committed
385 386
      }

387
      None => MotionVector::default()
388
    }
389
  }
390 391 392

  fn estimate_motion_ss2<T: Pixel>(
    fi: &FrameInvariants<T>, fs: &FrameState<T>, bsize: BlockSize, ref_idx: usize,
Romain Vimont's avatar
Romain Vimont committed
393
    bo: BlockOffset, pmvs: &[Option<MotionVector>; 3], ref_frame: usize
394 395 396 397 398 399 400 401 402 403
  ) -> Option<MotionVector> {
    if let Some(ref rec) = fi.rec_buffer.frames[ref_idx] {
      let blk_w = bsize.width();
      let blk_h = bsize.height();
      let bo_adj = adjust_bo(bo, fi, blk_w, blk_h);
      let bo_adj_h = BlockOffset{x: bo_adj.x >> 1, y: bo_adj.y >> 1};
      let po = PlaneOffset {
        x: (bo_adj.x as isize) << BLOCK_TO_PLANE_SHIFT >> 1,
        y: (bo_adj.y as isize) << BLOCK_TO_PLANE_SHIFT >> 1
      };
Romain Vimont's avatar
Romain Vimont committed
404
      let (mvx_min, mvx_max, mvy_min, mvy_max) = get_mv_range(fi.w_in_b, fi.h_in_b, bo_adj, blk_w, blk_h);
405 406 407

      let global_mv = [MotionVector{row: 0, col: 0}; 2];
      let frame_mvs = &fs.frame_mvs[ref_frame];
408
      let frame_ref_opt = fi.rec_buffer.frames[fi.ref_frames[0] as usize].as_ref().map(Arc::as_ref);
409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431

      let mut lowest_cost = std::u64::MAX;
      let mut best_mv = MotionVector::default();

      // Divide by 4 to account for subsampling, 0.125 is a fudge factor
      let lambda = (fi.me_lambda * 256.0 / 4.0 * 0.125) as u32;

      Self::me_ss2(
        fi, fs, pmvs, bo_adj_h,
        frame_mvs, frame_ref_opt, po, rec, global_mv, lambda,
        mvx_min, mvx_max, mvy_min, mvy_max, blk_w, blk_h,
        &mut best_mv, &mut lowest_cost
      );

      Some(MotionVector { row: best_mv.row * 2, col: best_mv.col * 2 })
    } else {
      None
    }
  }

  fn me_ss2<T: Pixel>(
    fi: &FrameInvariants<T>, fs: &FrameState<T>,
    pmvs: &[Option<MotionVector>; 3], bo_adj_h: BlockOffset,
432
    frame_mvs: &FrameMotionVectors, frame_ref_opt: Option<&ReferenceFrame<T>>,
433
    po: PlaneOffset, rec: &ReferenceFrame<T>,
434 435 436 437 438
    global_mv: [MotionVector; 2], lambda: u32,
    mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
    blk_w: usize, blk_h: usize,
    best_mv: &mut MotionVector, lowest_cost: &mut u64
  );
439
}
440

441 442 443
pub struct DiamondSearch {}
pub struct FullSearch {}

444 445
impl MotionEstimation for DiamondSearch {
  fn full_pixel_me<T: Pixel>(
446
    fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &ReferenceFrame<T>,
447
    bo: BlockOffset, lambda: u32,
448 449 450 451
    cmv: MotionVector, pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
    mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
    best_mv: &mut MotionVector, lowest_cost: &mut u64, ref_frame: usize
  ) {
452
    let frame_mvs = &fs.frame_mvs[ref_frame - LAST_FRAME];
453
    let frame_ref = fi.rec_buffer.frames[fi.ref_frames[0] as usize].as_ref().map(Arc::as_ref);
454
    let predictors =
455
      get_subset_predictors(bo, cmv, fi.w_in_b, fi.h_in_b, frame_mvs, frame_ref, ref_frame - LAST_FRAME);
456 457 458

    diamond_me_search(
      fi,
459
      bo.to_luma_plane_offset(),
460 461 462 463 464 465 466 467 468 469 470 471 472 473
      &fs.input.planes[0],
      &rec.frame.planes[0],
      &predictors,
      fi.sequence.bit_depth,
      pmv,
      lambda,
      mvx_min,
      mvx_max,
      mvy_min,
      mvy_max,
      blk_w,
      blk_h,
      best_mv,
      lowest_cost,
474
      false,
475 476 477
      ref_frame
    );
  }
478 479

  fn sub_pixel_me<T: Pixel>(
480
    fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &ReferenceFrame<T>,
481
    bo: BlockOffset, lambda: u32,
482 483 484 485 486 487 488 489
    pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
    mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
    best_mv: &mut MotionVector, lowest_cost: &mut u64, ref_frame: usize,
  )
  {
    let predictors = vec![*best_mv];
    diamond_me_search(
      fi,
490
      bo.to_luma_plane_offset(),
491 492 493 494 495 496 497 498 499 500 501 502 503 504
      &fs.input.planes[0],
      &rec.frame.planes[0],
      &predictors,
      fi.sequence.bit_depth,
      pmv,
      lambda,
      mvx_min,
      mvx_max,
      mvy_min,
      mvy_max,
      blk_w,
      blk_h,
      best_mv,
      lowest_cost,
505
      true,
506 507 508
      ref_frame
    );
  }
509 510 511 512

  fn me_ss2<T: Pixel>(
    fi: &FrameInvariants<T>, fs: &FrameState<T>,
    pmvs: &[Option<MotionVector>; 3], bo_adj_h: BlockOffset,
513
    frame_mvs: &FrameMotionVectors, frame_ref_opt: Option<&ReferenceFrame<T>>,
514
    po: PlaneOffset, rec: &ReferenceFrame<T>,
515 516 517 518 519 520 521 522
    global_mv: [MotionVector; 2], lambda: u32,
    mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
    blk_w: usize, blk_h: usize,
    best_mv: &mut MotionVector, lowest_cost: &mut u64
  ) {
    for omv in pmvs.iter() {
      if let Some(pmv) = omv {
        let mut predictors = get_subset_predictors::<T>(
Romain Vimont's avatar
Romain Vimont committed
523
          bo_adj_h,
524 525 526 527 528 529 530 531 532 533 534
          MotionVector{row: pmv.row, col: pmv.col},
          fi.w_in_b, fi.h_in_b,
          &frame_mvs, frame_ref_opt, 0
        );

        for predictor in &mut predictors {
          predictor.row >>= 1;
          predictor.col >>= 1;
        }

        diamond_me_search(
Romain Vimont's avatar
Romain Vimont committed
535
          fi, po,
536 537 538 539 540 541
          &fs.input_hres, &rec.input_hres,
          &predictors, fi.sequence.bit_depth,
          global_mv, lambda,
          mvx_min >> 1, mvx_max >> 1, mvy_min >> 1, mvy_max >> 1,
          blk_w >> 1, blk_h >> 1,
          best_mv, lowest_cost,
542
          false, 0
543 544 545 546
        );
      }
    }
  }
547 548 549 550
}

impl MotionEstimation for FullSearch {
  fn full_pixel_me<T: Pixel>(
551
    fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &ReferenceFrame<T>,
552
    bo: BlockOffset, lambda: u32,
553 554 555 556
    cmv: MotionVector, pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
    mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
    best_mv: &mut MotionVector, lowest_cost: &mut u64, _ref_frame: usize
  ) {
557
    let po = bo.to_luma_plane_offset();
558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578
    let range = 16;
    let x_lo = po.x
      + ((-range + (cmv.col / 8) as isize).max(mvx_min / 8).min(mvx_max / 8));
    let x_hi = po.x
      + ((range + (cmv.col / 8) as isize).max(mvx_min / 8).min(mvx_max / 8));
    let y_lo = po.y
      + ((-range + (cmv.row / 8) as isize).max(mvy_min / 8).min(mvy_max / 8));
    let y_hi = po.y
      + ((range + (cmv.row / 8) as isize).max(mvy_min / 8).min(mvy_max / 8));

    full_search(
      x_lo,
      x_hi,
      y_lo,
      y_hi,
      blk_h,
      blk_w,
      &fs.input.planes[0],
      &rec.frame.planes[0],
      best_mv,
      lowest_cost,
Romain Vimont's avatar
Romain Vimont committed
579
      po,
580 581 582 583 584 585 586
      2,
      fi.sequence.bit_depth,
      lambda,
      pmv,
      fi.allow_high_precision_mv
    );
  }
587 588

  fn sub_pixel_me<T: Pixel>(
589
    fi: &FrameInvariants<T>, fs: &FrameState<T>, _rec: &ReferenceFrame<T>,
590
    bo: BlockOffset, lambda: u32,
591
    pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
592
    mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
593 594 595 596 597 598
    best_mv: &mut MotionVector, lowest_cost: &mut u64, ref_frame: usize,
  )
  {
    telescopic_subpel_search(
      fi,
      fs,
599
      bo.to_luma_plane_offset(),
600 601 602 603 604 605 606
      lambda,
      ref_frame,
      pmv,
      mvx_min,
      mvx_max,
      mvy_min,
      mvy_max,
607 608
      blk_w,
      blk_h,
609 610 611 612
      best_mv,
      lowest_cost
    );
  }
613 614 615 616

  fn me_ss2<T: Pixel>(
    fi: &FrameInvariants<T>, fs: &FrameState<T>,
    pmvs: &[Option<MotionVector>; 3], _bo_adj_h: BlockOffset,
617
    _frame_mvs: &FrameMotionVectors, _frame_ref_opt: Option<&ReferenceFrame<T>>,
618
    po: PlaneOffset, rec: &ReferenceFrame<T>,
619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641
    _global_mv: [MotionVector; 2], lambda: u32,
    mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
    blk_w: usize, blk_h: usize,
    best_mv: &mut MotionVector, lowest_cost: &mut u64
  ) {
    let range = 16;
    for omv in pmvs.iter() {
      if let Some(pmv) = omv {
        let x_lo = po.x + (((pmv.col as isize / 8 - range).max(mvx_min / 8).min(mvx_max / 8)) >> 1);
        let x_hi = po.x + (((pmv.col as isize / 8 + range).max(mvx_min / 8).min(mvx_max / 8)) >> 1);
        let y_lo = po.y + (((pmv.row as isize / 8 - range).max(mvy_min / 8).min(mvy_max / 8)) >> 1);
        let y_hi = po.y + (((pmv.row as isize / 8 + range).max(mvy_min / 8).min(mvy_max / 8)) >> 1);
        full_search(
          x_lo,
          x_hi,
          y_lo,
          y_hi,
          blk_h >> 1,
          blk_w >> 1,
          &fs.input_hres,
          &rec.input_hres,
          best_mv,
          lowest_cost,
Romain Vimont's avatar
Romain Vimont committed
642
          po,
643 644 645 646 647 648 649 650 651
          1,
          fi.sequence.bit_depth,
          lambda,
          [MotionVector::default(); 2],
          fi.allow_high_precision_mv
        );
      }
    }
  }
652
}
653

654 655
fn get_best_predictor<T: Pixel>(
  fi: &FrameInvariants<T>,
Romain Vimont's avatar
Romain Vimont committed
656
  po: PlaneOffset, p_org: &Plane<T>, p_ref: &Plane<T>,
657
  predictors: &[MotionVector],
658
  bit_depth: usize, pmv: [MotionVector; 2], lambda: u32,
659 660
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
  blk_w: usize, blk_h: usize,
661 662
  center_mv: &mut MotionVector, center_mv_cost: &mut u64,
  tmp_plane_opt: &mut Option<Plane<T>>, ref_frame: usize) {
Vladimir Kazakov's avatar
Vladimir Kazakov committed
663
  *center_mv = MotionVector::default();
664 665 666 667 668 669
  *center_mv_cost = std::u64::MAX;

  for &init_mv in predictors.iter() {
    let cost = get_mv_rd_cost(
      fi, po, p_org, p_ref, bit_depth,
      pmv, lambda, mvx_min, mvx_max, mvy_min, mvy_max,
670
      blk_w, blk_h, init_mv, tmp_plane_opt, ref_frame);
671 672 673 674 675 676 677 678

    if cost < *center_mv_cost {
      *center_mv = init_mv;
      *center_mv_cost = cost;
    }
  }
}

679 680
fn diamond_me_search<T: Pixel>(
  fi: &FrameInvariants<T>,
Romain Vimont's avatar
Romain Vimont committed
681
  po: PlaneOffset, p_org: &Plane<T>, p_ref: &Plane<T>,
682
  predictors: &[MotionVector],
683
  bit_depth: usize, pmv: [MotionVector; 2], lambda: u32,
684 685
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
  blk_w: usize, blk_h: usize,
686
  center_mv: &mut MotionVector, center_mv_cost: &mut u64,
687
  subpixel: bool, ref_frame: usize)
688 689
{
  let diamond_pattern = [(1i16, 0i16), (0, 1), (-1, 0), (0, -1)];
690 691
  let (mut diamond_radius, diamond_radius_end, mut tmp_plane_opt) = {
    if subpixel {
692
      // Sub-pixel motion estimation
693 694 695 696 697
      (
        4i16,
        if fi.allow_high_precision_mv {1i16} else {2i16},
        Some(Plane::new(blk_w, blk_h, 0, 0, 0, 0)),
      )
698 699
    } else {
      // Full pixel motion estimation
700
      (16i16, 8i16, None)
701 702
    }
  };
703 704 705 706

  get_best_predictor(
    fi, po, p_org, p_ref, &predictors,
    bit_depth, pmv, lambda, mvx_min, mvx_max, mvy_min, mvy_max,
707
    blk_w, blk_h, center_mv, center_mv_cost,
708
    &mut tmp_plane_opt, ref_frame);
709 710 711

  loop {
    let mut best_diamond_rd_cost = std::u64::MAX;
Vladimir Kazakov's avatar
Vladimir Kazakov committed
712
    let mut best_diamond_mv = MotionVector::default();
713 714 715 716 717 718 719 720 721

    for p in diamond_pattern.iter() {

        let cand_mv = MotionVector {
          row: center_mv.row + diamond_radius * p.0,
          col: center_mv.col + diamond_radius * p.1
        };

        let rd_cost = get_mv_rd_cost(
Romain Vimont's avatar
Romain Vimont committed
722
          fi, po, p_org, p_ref, bit_depth,
723
          pmv, lambda, mvx_min, mvx_max, mvy_min, mvy_max,
724
          blk_w, blk_h, cand_mv, &mut tmp_plane_opt, ref_frame);
725 726 727 728 729 730 731 732

        if rd_cost < best_diamond_rd_cost {
          best_diamond_rd_cost = rd_cost;
          best_diamond_mv = cand_mv;
        }
    }

    if *center_mv_cost <= best_diamond_rd_cost {
733
      if diamond_radius == diamond_radius_end {
734 735 736 737 738 739 740 741 742 743 744 745 746 747
        break;
      } else {
        diamond_radius /= 2;
      }
    }
    else {
      *center_mv = best_diamond_mv;
      *center_mv_cost = best_diamond_rd_cost;
    }
  }

  assert!(*center_mv_cost < std::u64::MAX);
}

748 749
fn get_mv_rd_cost<T: Pixel>(
  fi: &FrameInvariants<T>,
Romain Vimont's avatar
Romain Vimont committed
750
  po: PlaneOffset, p_org: &Plane<T>, p_ref: &Plane<T>, bit_depth: usize,
751
  pmv: [MotionVector; 2], lambda: u32,
752 753
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
  blk_w: usize, blk_h: usize,
754 755
  cand_mv: MotionVector, tmp_plane_opt: &mut Option<Plane<T>>,
  ref_frame: usize) -> u64
756 757 758 759 760 761 762 763 764 765
{
  if (cand_mv.col as isize) < mvx_min || (cand_mv.col as isize) > mvx_max {
    return std::u64::MAX;
  }
  if (cand_mv.row as isize) < mvy_min || (cand_mv.row as isize) > mvy_max {
    return std::u64::MAX;
  }

  let plane_org = p_org.slice(po);

766
  if let Some(ref mut tmp_plane) = tmp_plane_opt {
Romain Vimont's avatar
Romain Vimont committed
767
    let mut tmp_slice = &mut tmp_plane.mut_slice(PlaneOffset { x: 0, y: 0 });
768 769 770
    PredictionMode::NEWMV.predict_inter(
      fi,
      0,
Romain Vimont's avatar
Romain Vimont committed
771
      po,
772 773 774 775 776 777
      &mut tmp_slice,
      blk_w,
      blk_h,
      [ref_frame, NONE_FRAME],
      [cand_mv, MotionVector { row: 0, col: 0 }]
    );
Romain Vimont's avatar
Romain Vimont committed
778
    let plane_ref = tmp_plane.slice(PlaneOffset { x: 0, y: 0 });
779 780 781 782 783 784
    compute_mv_rd_cost(
      fi, pmv, lambda, bit_depth, blk_w, blk_h, cand_mv,
      &plane_org, &plane_ref
    )
  } else {
    // Full pixel motion vector
Romain Vimont's avatar
Romain Vimont committed
785
    let plane_ref = p_ref.slice(PlaneOffset {
786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802
      x: po.x + (cand_mv.col / 8) as isize,
      y: po.y + (cand_mv.row / 8) as isize
    });
    compute_mv_rd_cost(
      fi, pmv, lambda, bit_depth, blk_w, blk_h, cand_mv,
      &plane_org, &plane_ref
    )
  }
}

fn compute_mv_rd_cost<T: Pixel>(
  fi: &FrameInvariants<T>,
  pmv: [MotionVector; 2], lambda: u32,
  bit_depth: usize, blk_w: usize, blk_h: usize, cand_mv: MotionVector,
  plane_org: &PlaneSlice<T>, plane_ref: &PlaneSlice<T>
) -> u64
{
803 804 805 806 807 808 809 810 811
  let sad = get_sad(&plane_org, &plane_ref, blk_h, blk_w, bit_depth);

  let rate1 = get_mv_rate(cand_mv, pmv[0], fi.allow_high_precision_mv);
  let rate2 = get_mv_rate(cand_mv, pmv[1], fi.allow_high_precision_mv);
  let rate = rate1.min(rate2 + 1);

  256 * sad as u64 + rate as u64 * lambda as u64
}

812
fn telescopic_subpel_search<T: Pixel>(
813
  fi: &FrameInvariants<T>, fs: &FrameState<T>, po: PlaneOffset,
814 815
  lambda: u32, ref_frame: usize, pmv: [MotionVector; 2],
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
816
  blk_w: usize, blk_h: usize,
Romain Vimont's avatar
Romain Vimont committed
817
  best_mv: &mut MotionVector, lowest_cost: &mut u64
818 819 820 821 822 823 824 825
) {
  let mode = PredictionMode::NEWMV;

  let mut steps = vec![8, 4, 2];
  if fi.allow_high_precision_mv {
    steps.push(1);
  }

Romain Vimont's avatar
Romain Vimont committed
826 827
  let mut tmp_plane = Plane::new(blk_w, blk_h, 0, 0, 0, 0);

828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850
  for step in steps {
    let center_mv_h = *best_mv;
    for i in 0..3 {
      for j in 0..3 {
        // Skip the center point that was already tested
        if i == 1 && j == 1 {
          continue;
        }

        let cand_mv = MotionVector {
          row: center_mv_h.row + step * (i as i16 - 1),
          col: center_mv_h.col + step * (j as i16 - 1)
        };

        if (cand_mv.col as isize) < mvx_min || (cand_mv.col as isize) > mvx_max {
          continue;
        }
        if (cand_mv.row as isize) < mvy_min || (cand_mv.row as isize) > mvy_max {
          continue;
        }

        {
          let tmp_slice =
Romain Vimont's avatar
Romain Vimont committed
851
            &mut tmp_plane.mut_slice(PlaneOffset { x: 0, y: 0 });
852 853 854 855

          mode.predict_inter(
            fi,
            0,
Romain Vimont's avatar
Romain Vimont committed
856
            po,
857 858 859 860 861 862 863 864
            tmp_slice,
            blk_w,
            blk_h,
            [ref_frame, NONE_FRAME],
            [cand_mv, MotionVector { row: 0, col: 0 }]
          );
        }

Romain Vimont's avatar
Romain Vimont committed
865 866
        let plane_org = fs.input.planes[0].slice(po);
        let plane_ref = tmp_plane.slice(PlaneOffset { x: 0, y: 0 });
867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883

        let sad = get_sad(&plane_org, &plane_ref, blk_h, blk_w, fi.sequence.bit_depth);

        let rate1 = get_mv_rate(cand_mv, pmv[0], fi.allow_high_precision_mv);
        let rate2 = get_mv_rate(cand_mv, pmv[1], fi.allow_high_precision_mv);
        let rate = rate1.min(rate2 + 1);
        let cost = 256 * sad as u64 + rate as u64 * lambda as u64;

        if cost < *lowest_cost {
          *lowest_cost = cost;
          *best_mv = cand_mv;
        }
      }
    }
  }
}

884
fn full_search<T: Pixel>(
Kyle Siefring's avatar
Kyle Siefring committed
885
  x_lo: isize, x_hi: isize, y_lo: isize, y_hi: isize, blk_h: usize,
886
  blk_w: usize, p_org: &Plane<T>, p_ref: &Plane<T>, best_mv: &mut MotionVector,
Romain Vimont's avatar
Romain Vimont committed
887
  lowest_cost: &mut u64, po: PlaneOffset, step: usize, bit_depth: usize,
888
  lambda: u32, pmv: [MotionVector; 2], allow_high_precision_mv: bool
Kyle Siefring's avatar
Kyle Siefring committed
889
) {
Luca Barbato's avatar
Luca Barbato committed
890 891 892 893 894
    let search_range_y = (y_lo..=y_hi).step_by(step);
    let search_range_x = (x_lo..=x_hi).step_by(step);
    let search_area = search_range_y.flat_map(|y| { search_range_x.clone().map(move |x| (y, x)) });

    let (cost, mv) = search_area.map(|(y, x)| {
895
      let plane_org = p_org.slice(po);
Romain Vimont's avatar
Romain Vimont committed
896
      let plane_ref = p_ref.slice(PlaneOffset { x, y });
897

Kyle Siefring's avatar
Kyle Siefring committed
898
      let sad = get_sad(&plane_org, &plane_ref, blk_h, blk_w, bit_depth);
899

Frank Bossen's avatar
Frank Bossen committed
900 901 902 903 904 905 906 907
      let mv = MotionVector {
        row: 8 * (y as i16 - po.y as i16),
        col: 8 * (x as i16 - po.x as i16)
      };

      let rate1 = get_mv_rate(mv, pmv[0], allow_high_precision_mv);
      let rate2 = get_mv_rate(mv, pmv[1], allow_high_precision_mv);
      let rate = rate1.min(rate2 + 1);
908
      let cost = 256 * sad as u64 + rate as u64 * lambda as u64;
Frank Bossen's avatar
Frank Bossen committed
909

Luca Barbato's avatar
Luca Barbato committed
910 911 912 913 914
      (cost, mv)
  }).min_by_key(|(c, _)| *c).unwrap();

    *lowest_cost = cost;
    *best_mv = mv;
915 916 917
}

// Adjust block offset such that entire block lies within frame boundaries
Romain Vimont's avatar
Romain Vimont committed
918
fn adjust_bo<T: Pixel>(bo: BlockOffset, fi: &FrameInvariants<T>, blk_w: usize, blk_h: usize) -> BlockOffset {
919 920 921 922 923 924
  BlockOffset {
    x: (bo.x as isize).min(fi.w_in_b as isize - blk_w as isize / 4).max(0) as usize,
    y: (bo.y as isize).min(fi.h_in_b as isize - blk_h as isize / 4).max(0) as usize
  }
}

925
#[inline(always)]
Frank Bossen's avatar
Frank Bossen committed
926
fn get_mv_rate(a: MotionVector, b: MotionVector, allow_high_precision_mv: bool) -> u32 {
927
  #[inline(always)]
Frank Bossen's avatar
Frank Bossen committed
928 929 930 931 932 933 934 935 936 937 938 939
  fn diff_to_rate(diff: i16, allow_high_precision_mv: bool) -> u32 {
    let d = if allow_high_precision_mv { diff } else { diff >> 1 };
    if d == 0 {
      0
    } else {
      2 * (16 - d.abs().leading_zeros())
    }
  }

  diff_to_rate(a.row - b.row, allow_high_precision_mv) + diff_to_rate(a.col - b.col, allow_high_precision_mv)
}

940 941
pub fn estimate_motion_ss4<T: Pixel>(
  fi: &FrameInvariants<T>, fs: &FrameState<T>, bsize: BlockSize, ref_idx: usize,
Romain Vimont's avatar
Romain Vimont committed
942
  bo: BlockOffset
943 944 945 946 947 948 949 950 951
) -> Option<MotionVector> {
  if let Some(ref rec) = fi.rec_buffer.frames[ref_idx] {
    let blk_w = bsize.width();
    let blk_h = bsize.height();
    let bo_adj = adjust_bo(bo, fi, blk_w, blk_h);
    let po = PlaneOffset {
      x: (bo_adj.x as isize) << BLOCK_TO_PLANE_SHIFT >> 2,
      y: (bo_adj.y as isize) << BLOCK_TO_PLANE_SHIFT >> 2
    };
952

953 954
    let range_x = 192 * fi.me_range_scale as isize;
    let range_y = 64 * fi.me_range_scale as isize;
Romain Vimont's avatar
Romain Vimont committed
955
    let (mvx_min, mvx_max, mvy_min, mvy_max) = get_mv_range(fi.w_in_b, fi.h_in_b, bo_adj, blk_w, blk_h);
956 957 958 959
    let x_lo = po.x + (((-range_x).max(mvx_min / 8)) >> 2);
    let x_hi = po.x + (((range_x).min(mvx_max / 8)) >> 2);
    let y_lo = po.y + (((-range_y).max(mvy_min / 8)) >> 2);
    let y_hi = po.y + (((range_y).min(mvy_max / 8)) >> 2);
960

961
    let mut lowest_cost = std::u64::MAX;
Vladimir Kazakov's avatar
Vladimir Kazakov committed
962
    let mut best_mv = MotionVector::default();
963

Frank Bossen's avatar
Frank Bossen committed
964
    // Divide by 16 to account for subsampling, 0.125 is a fudge factor
965
    let lambda = (fi.me_lambda * 256.0 / 16.0 * 0.125) as u32;
Frank Bossen's avatar
Frank Bossen committed
966

967
    full_search(
Kyle Siefring's avatar
Kyle Siefring committed
968 969 970 971 972 973 974 975 976
      x_lo,
      x_hi,
      y_lo,
      y_hi,
      blk_h >> 2,
      blk_w >> 2,
      &fs.input_qres,
      &rec.input_qres,
      &mut best_mv,
Frank Bossen's avatar
Frank Bossen committed
977
      &mut lowest_cost,
Romain Vimont's avatar
Romain Vimont committed
978
      po,
Kyle Siefring's avatar
Kyle Siefring committed
979
      1,
980
      fi.sequence.bit_depth,
Frank Bossen's avatar
Frank Bossen committed
981
      lambda,
Vladimir Kazakov's avatar
Vladimir Kazakov committed
982
      [MotionVector::default(); 2],
Frank Bossen's avatar
Frank Bossen committed
983
      fi.allow_high_precision_mv
984 985 986 987 988 989 990 991
    );

    Some(MotionVector { row: best_mv.row * 4, col: best_mv.col * 4 })
  } else {
    None
  }
}

992 993 994
#[cfg(test)]
pub mod test {
  use super::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
995 996
  use crate::partition::BlockSize;
  use crate::partition::BlockSize::*;
997 998

  // Generate plane data for get_sad_same()
999
  fn setup_sad<T: Pixel>() -> (Plane<T>, Plane<T>) {
1000 1001
    let mut input_plane = Plane::new(640, 480, 0, 0, 128 + 8, 128 + 8);
    let mut rec_plane = input_plane.clone();
1002 1003
    // Make the test pattern robust to data alignment
    let xpad_off = (input_plane.cfg.xorigin - input_plane.cfg.xpad) as i32 - 8i32;
Luca Barbato's avatar
Luca Barbato committed
1004

1005
    for (i, row) in input_plane.data.chunks_mut(input_plane.cfg.stride).enumerate() {
1006
      for (j, pixel) in row.into_iter().enumerate() {
1007
        let val = (j + i) as i32 - xpad_off & 255i32;
Luca Barbato's avatar
Luca Barbato committed
1008
        assert!(val >= u8::min_value().into() &&
1009
            val <= u8::max_value().into());
1010
        *pixel = T::cast_from(val);
1011 1012 1013 1014
      }
    }

    for (i, row) in rec_plane.data.chunks_mut(rec_plane.cfg.stride).enumerate() {
1015
      for (j, pixel) in row.into_iter().enumerate() {
1016
        let val = j as i32 - i as i32 - xpad_off & 255i32;
Luca Barbato's avatar
Luca Barbato committed
1017
        assert!(val >= u8::min_value().into() &&
1018
            val <= u8::max_value().into());
1019
        *pixel = T::cast_from(val);
1020 1021 1022 1023 1024 1025 1026
      }
    }

    (input_plane, rec_plane)
  }

  // Regression and validation test for SAD computation
1027
  fn get_sad_same_inner<T: Pixel>() {
1028
    let blocks: Vec<(BlockSize, u32)> = vec![
1029
      (BLOCK_4X4, 1912),
1030 1031
      (BLOCK_4X8, 4296),
      (BLOCK_8X4, 3496),
1032
      (BLOCK_8X8, 7824),
1033 1034
      (BLOCK_8X16, 16592),
      (BLOCK_16X8, 14416),
1035
      (BLOCK_16X16, 31136),
1036 1037
      (BLOCK_16X32, 60064),
      (BLOCK_32X16, 59552),
1038
      (BLOCK_32X32, 120128),
1039 1040
      (BLOCK_32X64, 186688),
      (BLOCK_64X32, 250176),
1041
      (BLOCK_64X64, 438912),
1042 1043
      (BLOCK_64X128, 654272),
      (BLOCK_128X64, 1016768),
1044
      (BLOCK_128X128, 1689792),
1045 1046 1047 1048 1049 1050
      (BLOCK_4X16, 8680),
      (BLOCK_16X4, 6664),
      (BLOCK_8X32, 31056),
      (BLOCK_32X8, 27600),
      (BLOCK_16X64, 93344),
      (BLOCK_64X16, 116384),
1051 1052
    ];

Kyle Siefring's avatar
Kyle Siefring committed
1053
    let bit_depth: usize = 8;
1054
    let (input_plane, rec_plane) = setup_sad::<T>();
1055 1056

    for block in blocks {
Kyle Siefring's avatar
Kyle Siefring committed
1057 1058
      let bsw = block.0.width();
      let bsh = block.0.height();
1059
      let po = PlaneOffset { x: 32, y: 40 };
1060

Romain Vimont's avatar
Romain Vimont committed
1061 1062
      let mut input_slice = input_plane.slice(po);
      let mut rec_slice = rec_plane.slice(po);
1063

Kyle Siefring's avatar
Kyle Siefring committed
1064 1065
      assert_eq!(
        block.1,
1066
        get_sad(&mut input_slice, &mut rec_slice, bsh, bsw, bit_depth)
Kyle Siefring's avatar
Kyle Siefring committed
1067
      );
1068 1069
    }
  }
1070 1071 1072 1073 1074 1075 1076 1077 1078 1079

  #[test]
  fn get_sad_same_u8() {
    get_sad_same_inner::<u8>();
  }

  #[test]
  fn get_sad_same_u16() {
    get_sad_same_inner::<u16>();
  }
1080
}