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
use crate::FrameInvariants;
use crate::FrameState;
use crate::partition::*;
Thomas Daede's avatar
Thomas Daede committed
19
use crate::partition::RefType::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
20
use crate::plane::*;
21
use crate::util::Pixel;
22

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

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

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

Luca Barbato's avatar
Luca Barbato committed
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 63
  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;
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 93

    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
94
  }
Kyle Siefring's avatar
Kyle Siefring committed
95

Luca Barbato's avatar
Luca Barbato committed
96
  #[target_feature(enable = "ssse3")]
97
  unsafe fn sad_hbd_ssse3(
98
    plane_org: &PlaneSlice<'_, u16>, plane_ref: &PlaneSlice<'_, u16>, blk_h: usize,
Luca Barbato's avatar
Luca Barbato committed
99 100 101
    blk_w: usize, bit_depth: usize
  ) -> u32 {
    let mut sum = 0 as u32;
102 103
    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
104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119
    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
120 121
        let org_ptr = org_slice.as_ptr();
        let ref_ptr = ref_slice.as_ptr();
122 123 124
        // 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
125 126
        sum += func(org_ptr, org_stride, ref_ptr, ref_stride);
      }
Kyle Siefring's avatar
Kyle Siefring committed
127
    }
128
    sum
Kyle Siefring's avatar
Kyle Siefring committed
129 130
  }

131
  #[target_feature(enable = "sse2")]
132 133
  unsafe fn sad_sse2(
    plane_org: &PlaneSlice<'_, u8>, plane_ref: &PlaneSlice<'_, u8>, blk_h: usize,
134 135
    blk_w: usize
  ) -> u32 {
136 137
    let org_ptr = plane_org.as_ptr();
    let ref_ptr = plane_ref.as_ptr();
138 139
    let org_stride = plane_org.plane.cfg.stride as libc::ptrdiff_t;
    let ref_stride = plane_ref.plane.cfg.stride as libc::ptrdiff_t;
140 141 142 143 144 145 146
    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);
147
    let step_size = blk_h.min(blk_w).min(ptr_align);
148 149 150 151 152 153 154 155 156
    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
    };
157 158 159 160 161
    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);
162 163 164 165 166 167
        sum += func(org_ptr, org_stride, ref_ptr, ref_stride);
      }
    }
    sum
  }

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

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

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

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

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

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

Romain Vimont's avatar
Romain Vimont committed
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 251
#[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]
  }
}

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

  (mvx_min, mvx_max, mvy_min, mvy_max)
}

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

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

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

279 280
  // EPZS subset A and B predictors.

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

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

299
  if !median_preds.is_empty() {
Vladimir Kazakov's avatar
Vladimir Kazakov committed
300
    let mut median_mv = MotionVector::default();
301
    for mv in median_preds.iter() {
302 303
      median_mv = median_mv + *mv;
    }
304 305 306
    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); }
307 308 309 310 311
  }

  // EPZS subset C predictors.

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

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

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

  predictors
}

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

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

356 357
  fn motion_estimation<T: Pixel> (
    fi: &FrameInvariants<T>, fs: &FrameState<T>, bsize: BlockSize,
Thomas Daede's avatar
Thomas Daede committed
358
    bo: BlockOffset, ref_frame: RefType, cmv: MotionVector,
359
    pmv: [MotionVector; 2]
360
  ) -> MotionVector {
Thomas Daede's avatar
Thomas Daede committed
361
    match fi.rec_buffer.frames[fi.ref_frames[ref_frame.to_index()] as usize]
362 363 364 365 366 367 368 369 370 371 372 373 374 375 376
    {
      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();

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

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

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

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

  fn estimate_motion_ss2<T: Pixel>(
    fi: &FrameInvariants<T>, fs: &FrameState<T>, bsize: BlockSize, ref_idx: usize,
Romain Vimont's avatar
Romain Vimont committed
394
    bo: BlockOffset, pmvs: &[Option<MotionVector>; 3], ref_frame: usize
395 396 397 398 399
  ) -> 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);
Romain Vimont's avatar
Romain Vimont committed
400
      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);
401 402 403

      let global_mv = [MotionVector{row: 0, col: 0}; 2];
      let frame_mvs = &fs.frame_mvs[ref_frame];
404
      let frame_ref_opt = fi.rec_buffer.frames[fi.ref_frames[0] as usize].as_ref().map(Arc::as_ref);
405 406 407 408 409 410 411 412

      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(
413 414
        fi, fs, pmvs, bo_adj,
        frame_mvs, frame_ref_opt, rec, global_mv, lambda,
415 416 417 418 419 420 421 422 423 424 425 426 427
        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,
428
    frame_mvs: &FrameMotionVectors, frame_ref_opt: Option<&ReferenceFrame<T>>,
429
    rec: &ReferenceFrame<T>, global_mv: [MotionVector; 2], lambda: u32,
430 431 432 433
    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
  );
434
}
435

436 437 438
pub struct DiamondSearch {}
pub struct FullSearch {}

439 440
impl MotionEstimation for DiamondSearch {
  fn full_pixel_me<T: Pixel>(
441
    fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &ReferenceFrame<T>,
442
    bo: BlockOffset, lambda: u32,
443 444
    cmv: MotionVector, pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
    mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
Thomas Daede's avatar
Thomas Daede committed
445
    best_mv: &mut MotionVector, lowest_cost: &mut u64, ref_frame: RefType
446
  ) {
Thomas Daede's avatar
Thomas Daede committed
447
    let frame_mvs = &fs.frame_mvs[ref_frame.to_index()];
448
    let frame_ref = fi.rec_buffer.frames[fi.ref_frames[0] as usize].as_ref().map(Arc::as_ref);
449
    let predictors =
Thomas Daede's avatar
Thomas Daede committed
450
      get_subset_predictors(bo, cmv, fi.w_in_b, fi.h_in_b, frame_mvs, frame_ref, ref_frame.to_index());
451 452 453

    diamond_me_search(
      fi,
454
      bo.to_luma_plane_offset(),
455 456 457 458 459 460 461 462 463 464 465 466 467 468
      &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,
469
      false,
470 471 472
      ref_frame
    );
  }
473 474

  fn sub_pixel_me<T: Pixel>(
475
    fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &ReferenceFrame<T>,
476
    bo: BlockOffset, lambda: u32,
477 478
    pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
    mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
Thomas Daede's avatar
Thomas Daede committed
479
    best_mv: &mut MotionVector, lowest_cost: &mut u64, ref_frame: RefType,
480 481 482 483 484
  )
  {
    let predictors = vec![*best_mv];
    diamond_me_search(
      fi,
485
      bo.to_luma_plane_offset(),
486 487 488 489 490 491 492 493 494 495 496 497 498 499
      &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,
500
      true,
501 502 503
      ref_frame
    );
  }
504 505 506

  fn me_ss2<T: Pixel>(
    fi: &FrameInvariants<T>, fs: &FrameState<T>,
507
    pmvs: &[Option<MotionVector>; 3], bo_adj: BlockOffset,
508
    frame_mvs: &FrameMotionVectors, frame_ref_opt: Option<&ReferenceFrame<T>>,
509
    rec: &ReferenceFrame<T>, global_mv: [MotionVector; 2], lambda: u32,
510 511 512 513
    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
  ) {
514 515 516 517
    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,
    };
518 519 520
    for omv in pmvs.iter() {
      if let Some(pmv) = omv {
        let mut predictors = get_subset_predictors::<T>(
521
          bo_adj,
522 523 524 525 526 527 528 529 530 531 532
          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
533
          fi, po,
534 535 536 537 538 539
          &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,
Thomas Daede's avatar
Thomas Daede committed
540
          false, LAST_FRAME
541 542 543 544
        );
      }
    }
  }
545 546 547 548
}

impl MotionEstimation for FullSearch {
  fn full_pixel_me<T: Pixel>(
549
    fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &ReferenceFrame<T>,
550
    bo: BlockOffset, lambda: u32,
551 552
    cmv: MotionVector, pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
    mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
Thomas Daede's avatar
Thomas Daede committed
553
    best_mv: &mut MotionVector, lowest_cost: &mut u64, _ref_frame: RefType
554
  ) {
555
    let po = bo.to_luma_plane_offset();
556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576
    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
577
      po,
578 579 580 581 582 583 584
      2,
      fi.sequence.bit_depth,
      lambda,
      pmv,
      fi.allow_high_precision_mv
    );
  }
585 586

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

  fn me_ss2<T: Pixel>(
    fi: &FrameInvariants<T>, fs: &FrameState<T>,
614
    pmvs: &[Option<MotionVector>; 3], bo_adj: BlockOffset,
615
    _frame_mvs: &FrameMotionVectors, _frame_ref_opt: Option<&ReferenceFrame<T>>,
616
    rec: &ReferenceFrame<T>, _global_mv: [MotionVector; 2], lambda: u32,
617 618 619 620
    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
  ) {
621 622 623 624
    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,
    };
625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642
    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
643
          po,
644 645 646 647 648 649 650 651 652
          1,
          fi.sequence.bit_depth,
          lambda,
          [MotionVector::default(); 2],
          fi.allow_high_precision_mv
        );
      }
    }
  }
653
}
654

655 656
fn get_best_predictor<T: Pixel>(
  fi: &FrameInvariants<T>,
Romain Vimont's avatar
Romain Vimont committed
657
  po: PlaneOffset, p_org: &Plane<T>, p_ref: &Plane<T>,
658
  predictors: &[MotionVector],
659
  bit_depth: usize, pmv: [MotionVector; 2], lambda: u32,
660 661
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
  blk_w: usize, blk_h: usize,
662
  center_mv: &mut MotionVector, center_mv_cost: &mut u64,
Thomas Daede's avatar
Thomas Daede committed
663
  tmp_plane_opt: &mut Option<Plane<T>>, ref_frame: RefType) {
Vladimir Kazakov's avatar
Vladimir Kazakov committed
664
  *center_mv = MotionVector::default();
665 666 667 668 669 670
  *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,
671
      blk_w, blk_h, init_mv, tmp_plane_opt, ref_frame);
672 673 674 675 676 677 678 679

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

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

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

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

    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
723
          fi, po, p_org, p_ref, bit_depth,
724
          pmv, lambda, mvx_min, mvx_max, mvy_min, mvy_max,
725
          blk_w, blk_h, cand_mv, &mut tmp_plane_opt, ref_frame);
726 727 728 729 730 731 732 733

        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 {
734
      if diamond_radius == diamond_radius_end {
735 736 737 738 739 740 741 742 743 744 745 746 747 748
        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);
}

749 750
fn get_mv_rd_cost<T: Pixel>(
  fi: &FrameInvariants<T>,
Romain Vimont's avatar
Romain Vimont committed
751
  po: PlaneOffset, p_org: &Plane<T>, p_ref: &Plane<T>, bit_depth: usize,
752
  pmv: [MotionVector; 2], lambda: u32,
753 754
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
  blk_w: usize, blk_h: usize,
755
  cand_mv: MotionVector, tmp_plane_opt: &mut Option<Plane<T>>,
Thomas Daede's avatar
Thomas Daede committed
756
  ref_frame: RefType) -> u64
757 758 759 760 761 762 763 764 765 766
{
  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);

767
  if let Some(ref mut tmp_plane) = tmp_plane_opt {
Romain Vimont's avatar
Romain Vimont committed
768
    let mut tmp_slice = &mut tmp_plane.mut_slice(PlaneOffset { x: 0, y: 0 });
769 770 771
    PredictionMode::NEWMV.predict_inter(
      fi,
      0,
Romain Vimont's avatar
Romain Vimont committed
772
      po,
773 774 775 776 777 778
      &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
779
    let plane_ref = tmp_plane.slice(PlaneOffset { x: 0, y: 0 });
780 781 782 783 784 785
    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
786
    let plane_ref = p_ref.slice(PlaneOffset {
787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803
      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
{
804 805 806 807 808 809 810 811 812
  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
}

813
fn telescopic_subpel_search<T: Pixel>(
814
  fi: &FrameInvariants<T>, fs: &FrameState<T>, po: PlaneOffset,
Thomas Daede's avatar
Thomas Daede committed
815
  lambda: u32, ref_frame: RefType, pmv: [MotionVector; 2],
816
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
817
  blk_w: usize, blk_h: usize,
Romain Vimont's avatar
Romain Vimont committed
818
  best_mv: &mut MotionVector, lowest_cost: &mut u64
819 820 821 822 823 824 825 826
) {
  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
827 828
  let mut tmp_plane = Plane::new(blk_w, blk_h, 0, 0, 0, 0);

829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851
  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
852
            &mut tmp_plane.mut_slice(PlaneOffset { x: 0, y: 0 });
853 854 855 856

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

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

        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;
        }
      }
    }
  }
}

885
fn full_search<T: Pixel>(
Kyle Siefring's avatar
Kyle Siefring committed
886
  x_lo: isize, x_hi: isize, y_lo: isize, y_hi: isize, blk_h: usize,
887
  blk_w: usize, p_org: &Plane<T>, p_ref: &Plane<T>, best_mv: &mut MotionVector,
Romain Vimont's avatar
Romain Vimont committed
888
  lowest_cost: &mut u64, po: PlaneOffset, step: usize, bit_depth: usize,
889
  lambda: u32, pmv: [MotionVector; 2], allow_high_precision_mv: bool
Kyle Siefring's avatar
Kyle Siefring committed
890
) {
Luca Barbato's avatar
Luca Barbato committed
891 892 893 894 895
    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)| {
896
      let plane_org = p_org.slice(po);
Romain Vimont's avatar
Romain Vimont committed
897
      let plane_ref = p_ref.slice(PlaneOffset { x, y });
898

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

Frank Bossen's avatar
Frank Bossen committed
901 902 903 904 905 906 907 908
      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);
909
      let cost = 256 * sad as u64 + rate as u64 * lambda as u64;
Frank Bossen's avatar
Frank Bossen committed
910

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

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

// Adjust block offset such that entire block lies within frame boundaries
Romain Vimont's avatar
Romain Vimont committed
919
fn adjust_bo<T: Pixel>(bo: BlockOffset, fi: &FrameInvariants<T>, blk_w: usize, blk_h: usize) -> BlockOffset {
920 921 922 923 924 925
  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
  }
}

926
#[inline(always)]
Frank Bossen's avatar
Frank Bossen committed
927
fn get_mv_rate(a: MotionVector, b: MotionVector, allow_high_precision_mv: bool) -> u32 {
928
  #[inline(always)]
Frank Bossen's avatar
Frank Bossen committed
929 930 931 932 933 934 935 936 937 938 939 940
  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)
}

941 942
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
943
  bo: BlockOffset
944 945 946 947 948 949 950 951 952
) -> 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
    };
953

954 955
    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
956
    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);
957 958 959 960
    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);
961

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

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

968
    full_search(
Kyle Siefring's avatar
Kyle Siefring committed
969 970 971 972 973 974 975 976 977
      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
978
      &mut lowest_cost,
Romain Vimont's avatar
Romain Vimont committed
979
      po,
Kyle Siefring's avatar
Kyle Siefring committed
980
      1,
981
      fi.sequence.bit_depth,
Frank Bossen's avatar
Frank Bossen committed
982
      lambda,
Vladimir Kazakov's avatar
Vladimir Kazakov committed
983
      [MotionVector::default(); 2],
Frank Bossen's avatar
Frank Bossen committed
984
      fi.allow_high_precision_mv
985 986 987 988 989 990 991 992
    );

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

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

  // Generate plane data for get_sad_same()
1000
  fn setup_sad<T: Pixel>() -> (Plane<T>, Plane<T>) {
1001 1002
    let mut input_plane = Plane::new(640, 480, 0, 0, 128 + 8, 128 + 8);
    let mut rec_plane = input_plane.clone();
1003 1004
    // 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
1005

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

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

    (input_plane, rec_plane)
  }

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

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

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

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

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

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

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