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

23 24
use arrayvec::*;

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

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

Raphaël Zumer's avatar
Raphaël Zumer committed
34 35
  use libc;

36 37 38 39 40 41 42 43 44
  macro_rules! declare_asm_sad {
    ($(($name: ident, $T: ident)),+) => (
      $(
        extern { fn $name (
          src: *const $T, src_stride: libc::ptrdiff_t, dst: *const $T,
          dst_stride: libc::ptrdiff_t
        ) -> u32; }
      )+
    )
Luca Barbato's avatar
Luca Barbato committed
45
  }
Kyle Siefring's avatar
Kyle Siefring committed
46

47
  declare_asm_sad![
48
    // SSSE3
49 50 51 52 53 54 55
    (rav1e_sad_4x4_hbd_ssse3, u16),
    (rav1e_sad_8x8_hbd10_ssse3, u16),
    (rav1e_sad_16x16_hbd_ssse3, u16),
    (rav1e_sad_32x32_hbd10_ssse3, u16),
    (rav1e_sad_64x64_hbd10_ssse3, u16),
    (rav1e_sad_128x128_hbd10_ssse3, u16),

56
    // SSE2
57
    (rav1e_sad4x4_sse2, u8),
58 59 60 61
    (rav1e_sad4x8_sse2, u8),
    (rav1e_sad4x16_sse2, u8),

    (rav1e_sad8x4_sse2, u8),
62
    (rav1e_sad8x8_sse2, u8),
63 64 65
    (rav1e_sad8x16_sse2, u8),
    (rav1e_sad8x32_sse2, u8),

66
    (rav1e_sad16x16_sse2, u8),
67

68
    (rav1e_sad32x32_sse2, u8),
69

70
    (rav1e_sad64x64_sse2, u8),
71

72 73
    (rav1e_sad128x128_sse2, u8),

74 75 76
    // AVX
    (rav1e_sad16x4_avx2, u8),
    (rav1e_sad16x8_avx2, u8),
77
    (rav1e_sad16x16_avx2, u8),
78 79 80 81 82
    (rav1e_sad16x32_avx2, u8),
    (rav1e_sad16x64_avx2, u8),

    (rav1e_sad32x8_avx2, u8),
    (rav1e_sad32x16_avx2, u8),
83
    (rav1e_sad32x32_avx2, u8),
84 85 86 87
    (rav1e_sad32x64_avx2, u8),

    (rav1e_sad64x16_avx2, u8),
    (rav1e_sad64x32_avx2, u8),
88
    (rav1e_sad64x64_avx2, u8),
89 90 91
    (rav1e_sad64x128_avx2, u8),

    (rav1e_sad128x64_avx2, u8),
92 93 94
    (rav1e_sad128x128_avx2, u8)
  ];

Luca Barbato's avatar
Luca Barbato committed
95
  #[target_feature(enable = "ssse3")]
96
  unsafe fn sad_hbd_ssse3(
Romain Vimont's avatar
Romain Vimont committed
97 98 99 100 101
    plane_org: &PlaneRegion<'_, u16>,
    plane_ref: &PlaneRegion<'_, u16>,
    blk_w: usize,
    blk_h: usize,
    bit_depth: usize,
Luca Barbato's avatar
Luca Barbato committed
102 103
  ) -> u32 {
    let mut sum = 0 as u32;
Romain Vimont's avatar
Romain Vimont committed
104 105
    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
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) {
120
        // FIXME for now, T == u16
Romain Vimont's avatar
Romain Vimont committed
121 122
        let org_ptr = &plane_org[r][c] as *const u16;
        let ref_ptr = &plane_ref[r][c] as *const u16;
Luca Barbato's avatar
Luca Barbato committed
123 124
        sum += func(org_ptr, org_stride, ref_ptr, ref_stride);
      }
Kyle Siefring's avatar
Kyle Siefring committed
125
    }
126
    sum
Kyle Siefring's avatar
Kyle Siefring committed
127 128
  }

129
  #[target_feature(enable = "sse2")]
130
  unsafe fn sad_sse2(
Romain Vimont's avatar
Romain Vimont committed
131 132 133 134
    plane_org: &PlaneRegion<'_, u8>,
    plane_ref: &PlaneRegion<'_, u8>,
    blk_w: usize,
    blk_h: usize,
135
  ) -> u32 {
Romain Vimont's avatar
Romain Vimont committed
136 137 138 139
    let org_ptr = plane_org.data_ptr();
    let ref_ptr = plane_ref.data_ptr();
    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
    let mut sum = 0 as u32;
Romain Vimont's avatar
Romain Vimont committed
158 159 160 161
    for r in (0..blk_h).step_by(step_size) {
      for c in (0..blk_w).step_by(step_size) {
        let org_ptr = &plane_org[r][c] as *const u8;
        let ref_ptr = &plane_ref[r][c] as *const u8;
162 163 164 165 166 167
        sum += func(org_ptr, org_stride, ref_ptr, ref_stride);
      }
    }
    sum
  }

168 169
  #[target_feature(enable = "avx2")]
  unsafe fn sad_avx2(
Romain Vimont's avatar
Romain Vimont committed
170 171 172 173
    plane_org: &PlaneRegion<'_, u8>,
    plane_ref: &PlaneRegion<'_, u8>,
    blk_w: usize,
    blk_h: usize,
174
  ) -> u32 {
Romain Vimont's avatar
Romain Vimont committed
175 176 177 178
    let org_ptr = plane_org.data_ptr();
    let ref_ptr = plane_ref.data_ptr();
    let org_stride = plane_org.plane_cfg.stride as libc::ptrdiff_t;
    let ref_stride = plane_ref.plane_cfg.stride as libc::ptrdiff_t;
179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209

    let func = match (blk_w, blk_h) {
      (4, 4) => rav1e_sad4x4_sse2,
      (4, 8) => rav1e_sad4x8_sse2,
      (4, 16) => rav1e_sad4x16_sse2,

      (8, 4) => rav1e_sad8x4_sse2,
      (8, 8) => rav1e_sad8x8_sse2,
      (8, 16) => rav1e_sad8x16_sse2,
      (8, 32) => rav1e_sad8x32_sse2,

      (16, 4) => rav1e_sad16x4_avx2,
      (16, 8) => rav1e_sad16x8_avx2,
      (16, 16) => rav1e_sad16x16_avx2,
      (16, 32) => rav1e_sad16x32_avx2,
      (16, 64) => rav1e_sad16x64_avx2,

      (32, 8) => rav1e_sad32x8_avx2,
      (32, 16) => rav1e_sad32x16_avx2,
      (32, 32) => rav1e_sad32x32_avx2,
      (32, 64) => rav1e_sad32x64_avx2,

      (64, 16) => rav1e_sad64x16_avx2,
      (64, 32) => rav1e_sad64x32_avx2,
      (64, 64) => rav1e_sad64x64_avx2,
      (64, 128) => rav1e_sad64x128_avx2,

      (128, 64) => rav1e_sad128x64_avx2,
      (128, 128) => rav1e_sad128x128_avx2,

      _ => unreachable!()
210
    };
211 212
    func(org_ptr, org_stride, ref_ptr, ref_stride)

213 214
  }

Luca Barbato's avatar
Luca Barbato committed
215
  #[inline(always)]
216
  pub fn get_sad<T: Pixel>(
Romain Vimont's avatar
Romain Vimont committed
217 218 219 220 221
    plane_org: &PlaneRegion<'_, T>,
    plane_ref: &PlaneRegion<'_, T>,
    blk_w: usize,
    blk_h: usize,
    bit_depth: usize,
Luca Barbato's avatar
Luca Barbato committed
222
  ) -> u32 {
223
    #[cfg(all(target_arch = "x86_64", feature = "nasm"))]
Luca Barbato's avatar
Luca Barbato committed
224
    {
225
      if mem::size_of::<T>() == 2 && is_x86_feature_detected!("ssse3") && blk_h >= 4 && blk_w >= 4 {
Luca Barbato's avatar
Luca Barbato committed
226
        return unsafe {
Romain Vimont's avatar
Romain Vimont committed
227 228
          let plane_org = &*(plane_org as *const _ as *const PlaneRegion<'_, u16>);
          let plane_ref = &*(plane_ref as *const _ as *const PlaneRegion<'_, u16>);
229
          sad_hbd_ssse3(plane_org, plane_ref, blk_w, blk_h, bit_depth)
Luca Barbato's avatar
Luca Barbato committed
230 231
        };
      }
232
      if mem::size_of::<T>() == 1 && is_x86_feature_detected!("avx2") && blk_h >= 4 && blk_w >= 4 {
233
        return unsafe {
Romain Vimont's avatar
Romain Vimont committed
234 235
          let plane_org = &*(plane_org as *const _ as *const PlaneRegion<'_, u8>);
          let plane_ref = &*(plane_ref as *const _ as *const PlaneRegion<'_, u8>);
236
          sad_avx2(plane_org, plane_ref, blk_w, blk_h)
237 238
        };
      }
239 240
      if mem::size_of::<T>() == 1 && is_x86_feature_detected!("sse2") && blk_h >= 4 && blk_w >= 4 {
        return unsafe {
Romain Vimont's avatar
Romain Vimont committed
241 242
          let plane_org = &*(plane_org as *const _ as *const PlaneRegion<'_, u8>);
          let plane_ref = &*(plane_ref as *const _ as *const PlaneRegion<'_, u8>);
243
          sad_sse2(plane_org, plane_ref, blk_w, blk_h)
244 245
        };
      }
Kyle Siefring's avatar
Kyle Siefring committed
246
    }
247
    super::native::get_sad(plane_org, plane_ref, blk_w, blk_h, bit_depth)
Kyle Siefring's avatar
Kyle Siefring committed
248
  }
249 250 251
}

mod native {
Romain Vimont's avatar
Romain Vimont committed
252
  use crate::tiling::*;
253
  use crate::util::*;
254

Luca Barbato's avatar
Luca Barbato committed
255
  #[inline(always)]
256
  pub fn get_sad<T: Pixel>(
Romain Vimont's avatar
Romain Vimont committed
257 258 259 260 261
    plane_org: &PlaneRegion<'_, T>,
    plane_ref: &PlaneRegion<'_, T>,
    blk_w: usize,
    blk_h: usize,
    _bit_depth: usize,
Luca Barbato's avatar
Luca Barbato committed
262 263
  ) -> u32 {
    let mut sum = 0 as u32;
264

Romain Vimont's avatar
Romain Vimont committed
265
    for (slice_org, slice_ref) in plane_org.rows_iter().take(blk_h).zip(plane_ref.rows_iter()) {
266 267
      sum += slice_org
        .iter()
Romain Vimont's avatar
Romain Vimont committed
268
        .take(blk_w)
269
        .zip(slice_ref)
270
        .map(|(&a, &b)| (i32::cast_from(a) - i32::cast_from(b)).abs() as u32)
271
        .sum::<u32>();
Luca Barbato's avatar
Luca Barbato committed
272
    }
273

Luca Barbato's avatar
Luca Barbato committed
274 275
    sum
  }
276 277
}

Romain Vimont's avatar
Romain Vimont committed
278 279 280 281 282 283 284 285 286 287 288 289 290 291 292
#[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,
    }
  }
293 294 295 296 297 298 299 300 301 302

  #[inline(always)]
  pub fn as_tile_motion_vectors(&self) -> TileMotionVectors<'_> {
    TileMotionVectors::new(self, 0, 0, self.cols, self.rows)
  }

  #[inline(always)]
  pub fn as_tile_motion_vectors_mut(&mut self) -> TileMotionVectorsMut<'_> {
    TileMotionVectorsMut::new(self, 0, 0, self.cols, self.rows)
  }
Romain Vimont's avatar
Romain Vimont committed
303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319
}

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

320
fn get_mv_range(
Romain Vimont's avatar
Romain Vimont committed
321
  w_in_b: usize, h_in_b: usize, bo: BlockOffset, blk_w: usize, blk_h: usize
322
) -> (isize, isize, isize, isize) {
323 324 325
  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;
326
  let mvx_max = (w_in_b - bo.x - blk_w / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_w;
327
  let mvy_min = -(bo.y as isize) * (8 * MI_SIZE) as isize - border_h;
328
  let mvy_max = (h_in_b - bo.y - blk_h / MI_SIZE) as isize * (8 * MI_SIZE) as isize + border_h;
329 330 331 332

  (mvx_min, mvx_max, mvy_min, mvy_max)
}

333
pub fn get_subset_predictors<T: Pixel>(
334
  tile_bo: BlockOffset, cmv: MotionVector,
335
  tile_mvs: &TileMotionVectors<'_>, frame_ref_opt: Option<&ReferenceFrame<T>>,
336
  ref_frame_id: usize
337 338
) -> (ArrayVec<[MotionVector; 11]>) {
  let mut predictors = ArrayVec::<[_; 11]>::new();
339

340 341 342 343 344 345
  // Zero motion vector
  predictors.push(MotionVector::default());

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

346 347
  // EPZS subset A and B predictors.

348
  let mut median_preds = ArrayVec::<[_; 3]>::new();
349 350
  if tile_bo.x > 0 {
    let left = tile_mvs[tile_bo.y][tile_bo.x - 1];
351 352
    median_preds.push(left);
    if !left.is_zero() { predictors.push(left); }
353
  }
354 355
  if tile_bo.y > 0 {
    let top = tile_mvs[tile_bo.y - 1][tile_bo.x];
356 357
    median_preds.push(top);
    if !top.is_zero() { predictors.push(top); }
358

359 360
    if tile_bo.x < tile_mvs.cols() - 1 {
      let top_right = tile_mvs[tile_bo.y - 1][tile_bo.x + 1];
361 362
      median_preds.push(top_right);
      if !top_right.is_zero() { predictors.push(top_right); }
363 364 365
    }
  }

366
  if !median_preds.is_empty() {
Vladimir Kazakov's avatar
Vladimir Kazakov committed
367
    let mut median_mv = MotionVector::default();
368
    for mv in median_preds.iter() {
369 370
      median_mv = median_mv + *mv;
    }
371 372 373
    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); }
374 375 376 377 378
  }

  // EPZS subset C predictors.

  if let Some(ref frame_ref) = frame_ref_opt {
379
    let prev_frame_mvs = &frame_ref.frame_mvs[ref_frame_id];
380

381
    let frame_bo = BlockOffset {
382 383
      x: tile_mvs.x() + tile_bo.x,
      y: tile_mvs.y() + tile_bo.y,
384 385 386
    };
    if frame_bo.x > 0 {
      let left = prev_frame_mvs[frame_bo.y][frame_bo.x - 1];
387
      if !left.is_zero() { predictors.push(left); }
388
    }
389 390
    if frame_bo.y > 0 {
      let top = prev_frame_mvs[frame_bo.y - 1][frame_bo.x];
391
      if !top.is_zero() { predictors.push(top); }
392
    }
393 394
    if frame_bo.x < prev_frame_mvs.cols - 1 {
      let right = prev_frame_mvs[frame_bo.y][frame_bo.x + 1];
395
      if !right.is_zero() { predictors.push(right); }
396
    }
397 398
    if frame_bo.y < prev_frame_mvs.rows - 1 {
      let bottom = prev_frame_mvs[frame_bo.y + 1][frame_bo.x];
399
      if !bottom.is_zero() { predictors.push(bottom); }
400 401
    }

402
    let previous = prev_frame_mvs[frame_bo.y][frame_bo.x];
403
    if !previous.is_zero() { predictors.push(previous); }
404 405 406 407 408
  }

  predictors
}

409
pub trait MotionEstimation {
410
  fn full_pixel_me<T: Pixel>(
411
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, rec: &ReferenceFrame<T>,
412
    tile_bo: BlockOffset, lambda: u32,
413
    cmv: MotionVector, pmv: [MotionVector; 2],
414 415
    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
416
    lowest_cost: &mut u64, ref_frame: RefType
417
  );
418

419
  fn sub_pixel_me<T: Pixel>(
420
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, rec: &ReferenceFrame<T>,
421
    tile_bo: BlockOffset, lambda: u32, pmv: [MotionVector; 2],
422 423
    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
424
    lowest_cost: &mut u64, ref_frame: RefType
425 426
  );

427
  fn motion_estimation<T: Pixel> (
428
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, bsize: BlockSize,
429
    tile_bo: BlockOffset, ref_frame: RefType, cmv: MotionVector,
430
    pmv: [MotionVector; 2]
431
  ) -> MotionVector {
Thomas Daede's avatar
Thomas Daede committed
432
    match fi.rec_buffer.frames[fi.ref_frames[ref_frame.to_index()] as usize]
433 434 435 436
    {
      Some(ref rec) => {
        let blk_w = bsize.width();
        let blk_h = bsize.height();
437
        let frame_bo = ts.to_frame_block_offset(tile_bo);
438
        let (mvx_min, mvx_max, mvy_min, mvy_max) =
439
          get_mv_range(fi.w_in_b, fi.h_in_b, frame_bo, blk_w, blk_h);
440 441 442 443 444 445 446 447 448

        // 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();

449
        Self::full_pixel_me(fi, ts, rec, tile_bo, lambda, cmv, pmv,
450 451
                           mvx_min, mvx_max, mvy_min, mvy_max, blk_w, blk_h,
                           &mut best_mv, &mut lowest_cost, ref_frame);
452

453
        Self::sub_pixel_me(fi, ts, rec, tile_bo, lambda, pmv,
454
                           mvx_min, mvx_max, mvy_min, mvy_max, blk_w, blk_h,
455
                           &mut best_mv, &mut lowest_cost, ref_frame);
456 457

        best_mv
Frank Bossen's avatar
Frank Bossen committed
458 459
      }

460
      None => MotionVector::default()
461
    }
462
  }
463 464

  fn estimate_motion_ss2<T: Pixel>(
465
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, bsize: BlockSize, ref_idx: usize,
466
    tile_bo: BlockOffset, pmvs: &[Option<MotionVector>; 3], ref_frame: usize
467 468 469 470
  ) -> Option<MotionVector> {
    if let Some(ref rec) = fi.rec_buffer.frames[ref_idx] {
      let blk_w = bsize.width();
      let blk_h = bsize.height();
471 472
      let tile_bo_adj = adjust_bo(tile_bo, ts.mi_width, ts.mi_height, blk_w, blk_h);
      let frame_bo_adj = ts.to_frame_block_offset(tile_bo_adj);
473
      let (mvx_min, mvx_max, mvy_min, mvy_max) = get_mv_range(fi.w_in_b, fi.h_in_b, frame_bo_adj, blk_w, blk_h);
474 475

      let global_mv = [MotionVector{row: 0, col: 0}; 2];
476
      let tile_mvs = &ts.mvs[ref_frame].as_const();
477
      let frame_ref_opt = fi.rec_buffer.frames[fi.ref_frames[0] as usize].as_ref().map(Arc::as_ref);
478 479 480 481 482 483 484 485

      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(
486
        fi, ts, pmvs, tile_bo_adj,
Romain Vimont's avatar
Romain Vimont committed
487
        tile_mvs, frame_ref_opt, rec, global_mv, lambda,
488 489 490 491 492 493 494 495 496 497 498
        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>(
499
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>,
500
    pmvs: &[Option<MotionVector>; 3], tile_bo_adj: BlockOffset,
501
    tile_mvs: &TileMotionVectors<'_>, frame_ref_opt: Option<&ReferenceFrame<T>>,
502
    rec: &ReferenceFrame<T>, global_mv: [MotionVector; 2], lambda: u32,
503 504 505 506
    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
  );
507
}
508

509 510 511
pub struct DiamondSearch {}
pub struct FullSearch {}

512 513
impl MotionEstimation for DiamondSearch {
  fn full_pixel_me<T: Pixel>(
514
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, rec: &ReferenceFrame<T>,
515
    tile_bo: BlockOffset, lambda: u32,
516 517
    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
518
    best_mv: &mut MotionVector, lowest_cost: &mut u64, ref_frame: RefType
519
  ) {
520
    let tile_mvs = &ts.mvs[ref_frame.to_index()].as_const();
521
    let frame_ref = fi.rec_buffer.frames[fi.ref_frames[0] as usize].as_ref().map(Arc::as_ref);
522
    let predictors =
Romain Vimont's avatar
Romain Vimont committed
523
      get_subset_predictors(tile_bo, cmv, tile_mvs, frame_ref, ref_frame.to_index());
524

525
    let frame_bo = ts.to_frame_block_offset(tile_bo);
526 527
    diamond_me_search(
      fi,
528 529
      frame_bo.to_luma_plane_offset(),
      &ts.input.planes[0],
530 531 532 533 534 535 536 537 538 539 540 541 542
      &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,
543
      false,
544 545 546
      ref_frame
    );
  }
547 548

  fn sub_pixel_me<T: Pixel>(
549
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, rec: &ReferenceFrame<T>,
550
    tile_bo: BlockOffset, lambda: u32,
551 552
    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 556
  )
  {
    let predictors = vec![*best_mv];
557
    let frame_bo = ts.to_frame_block_offset(tile_bo);
558 559
    diamond_me_search(
      fi,
560 561
      frame_bo.to_luma_plane_offset(),
      &ts.input.planes[0],
562 563 564 565 566 567 568 569 570 571 572 573 574
      &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,
575
      true,
576 577 578
      ref_frame
    );
  }
579 580

  fn me_ss2<T: Pixel>(
581
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>,
582
    pmvs: &[Option<MotionVector>; 3], tile_bo_adj: BlockOffset,
583
    tile_mvs: &TileMotionVectors<'_>, frame_ref_opt: Option<&ReferenceFrame<T>>,
584
    rec: &ReferenceFrame<T>, global_mv: [MotionVector; 2], lambda: u32,
585 586 587 588
    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
  ) {
589 590
    let frame_bo_adj = ts.to_frame_block_offset(tile_bo_adj);
    let frame_po = PlaneOffset {
591 592
      x: (frame_bo_adj.x as isize) << BLOCK_TO_PLANE_SHIFT >> 1,
      y: (frame_bo_adj.y as isize) << BLOCK_TO_PLANE_SHIFT >> 1,
593
    };
594 595 596
    for omv in pmvs.iter() {
      if let Some(pmv) = omv {
        let mut predictors = get_subset_predictors::<T>(
597
          tile_bo_adj,
598
          MotionVector{row: pmv.row, col: pmv.col},
599
          &tile_mvs, frame_ref_opt, 0
600 601 602 603 604 605 606 607
        );

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

        diamond_me_search(
608
          fi, frame_po,
609
          &ts.input_hres, &rec.input_hres,
610 611 612 613 614
          &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
615
          false, LAST_FRAME
616 617 618 619
        );
      }
    }
  }
620 621 622 623
}

impl MotionEstimation for FullSearch {
  fn full_pixel_me<T: Pixel>(
624
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, rec: &ReferenceFrame<T>,
625
    tile_bo: BlockOffset, lambda: u32,
626 627
    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
628
    best_mv: &mut MotionVector, lowest_cost: &mut u64, _ref_frame: RefType
629
  ) {
630 631
    let frame_bo = ts.to_frame_block_offset(tile_bo);
    let frame_po = frame_bo.to_luma_plane_offset();
632
    let range = 16;
633
    let x_lo = frame_po.x
634
      + ((-range + (cmv.col / 8) as isize).max(mvx_min / 8).min(mvx_max / 8));
635
    let x_hi = frame_po.x
636
      + ((range + (cmv.col / 8) as isize).max(mvx_min / 8).min(mvx_max / 8));
637
    let y_lo = frame_po.y
638
      + ((-range + (cmv.row / 8) as isize).max(mvy_min / 8).min(mvy_max / 8));
639
    let y_hi = frame_po.y
640 641 642 643 644 645 646 647 648
      + ((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,
649
      &ts.input.planes[0],
650 651 652
      &rec.frame.planes[0],
      best_mv,
      lowest_cost,
653
      frame_po,
654 655 656 657 658 659 660
      2,
      fi.sequence.bit_depth,
      lambda,
      pmv,
      fi.allow_high_precision_mv
    );
  }
661 662

  fn sub_pixel_me<T: Pixel>(
663
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, _rec: &ReferenceFrame<T>,
664
    tile_bo: BlockOffset, lambda: u32,
665
    pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
666
    mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
Thomas Daede's avatar
Thomas Daede committed
667
    best_mv: &mut MotionVector, lowest_cost: &mut u64, ref_frame: RefType,
668 669
  )
  {
670
    let frame_bo = ts.to_frame_block_offset(tile_bo);
671 672
    telescopic_subpel_search(
      fi,
673 674
      ts,
      frame_bo.to_luma_plane_offset(),
675 676 677 678 679 680 681
      lambda,
      ref_frame,
      pmv,
      mvx_min,
      mvx_max,
      mvy_min,
      mvy_max,
682 683
      blk_w,
      blk_h,
684 685 686 687
      best_mv,
      lowest_cost
    );
  }
688 689

  fn me_ss2<T: Pixel>(
690
    fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>,
691
    pmvs: &[Option<MotionVector>; 3], tile_bo_adj: BlockOffset,
692
    _tile_mvs: &TileMotionVectors<'_>, _frame_ref_opt: Option<&ReferenceFrame<T>>,
693
    rec: &ReferenceFrame<T>, _global_mv: [MotionVector; 2], lambda: u32,
694 695 696 697
    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
  ) {
698 699
    let frame_bo_adj = ts.to_frame_block_offset(tile_bo_adj);
    let frame_po = PlaneOffset {
700 701
      x: (frame_bo_adj.x as isize) << BLOCK_TO_PLANE_SHIFT >> 1,
      y: (frame_bo_adj.y as isize) << BLOCK_TO_PLANE_SHIFT >> 1,
702
    };
703 704 705
    let range = 16;
    for omv in pmvs.iter() {
      if let Some(pmv) = omv {
706 707 708 709
        let x_lo = frame_po.x + (((pmv.col as isize / 8 - range).max(mvx_min / 8).min(mvx_max / 8)) >> 1);
        let x_hi = frame_po.x + (((pmv.col as isize / 8 + range).max(mvx_min / 8).min(mvx_max / 8)) >> 1);
        let y_lo = frame_po.y + (((pmv.row as isize / 8 - range).max(mvy_min / 8).min(mvy_max / 8)) >> 1);
        let y_hi = frame_po.y + (((pmv.row as isize / 8 + range).max(mvy_min / 8).min(mvy_max / 8)) >> 1);
710 711 712 713 714 715 716
        full_search(
          x_lo,
          x_hi,
          y_lo,
          y_hi,
          blk_h >> 1,
          blk_w >> 1,
717
          &ts.input_hres,
718 719 720
          &rec.input_hres,
          best_mv,
          lowest_cost,
721
          frame_po,
722 723 724 725 726 727 728 729 730
          1,
          fi.sequence.bit_depth,
          lambda,
          [MotionVector::default(); 2],
          fi.allow_high_precision_mv
        );
      }
    }
  }
731
}
732

733 734
fn get_best_predictor<T: Pixel>(
  fi: &FrameInvariants<T>,
Romain Vimont's avatar
Romain Vimont committed
735
  po: PlaneOffset, p_org: &Plane<T>, p_ref: &Plane<T>,
736
  predictors: &[MotionVector],
737
  bit_depth: usize, pmv: [MotionVector; 2], lambda: u32,
738 739
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
  blk_w: usize, blk_h: usize,
740
  center_mv: &mut MotionVector, center_mv_cost: &mut u64,
Thomas Daede's avatar
Thomas Daede committed
741
  tmp_plane_opt: &mut Option<Plane<T>>, ref_frame: RefType) {
Vladimir Kazakov's avatar
Vladimir Kazakov committed
742
  *center_mv = MotionVector::default();
743 744 745 746 747 748
  *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,
749
      blk_w, blk_h, init_mv, tmp_plane_opt, ref_frame);
750 751 752 753 754 755 756 757

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

758 759
fn diamond_me_search<T: Pixel>(
  fi: &FrameInvariants<T>,
Romain Vimont's avatar
Romain Vimont committed
760
  po: PlaneOffset, p_org: &Plane<T>, p_ref: &Plane<T>,
761
  predictors: &[MotionVector],
762
  bit_depth: usize, pmv: [MotionVector; 2], lambda: u32,
763 764
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
  blk_w: usize, blk_h: usize,
765
  center_mv: &mut MotionVector, center_mv_cost: &mut u64,
Thomas Daede's avatar
Thomas Daede committed
766
  subpixel: bool, ref_frame: RefType)
767 768
{
  let diamond_pattern = [(1i16, 0i16), (0, 1), (-1, 0), (0, -1)];
769 770
  let (mut diamond_radius, diamond_radius_end, mut tmp_plane_opt) = {
    if subpixel {
771
      // Sub-pixel motion estimation
772 773 774 775 776
      (
        4i16,
        if fi.allow_high_precision_mv {1i16} else {2i16},
        Some(Plane::new(blk_w, blk_h, 0, 0, 0, 0)),
      )
777 778
    } else {
      // Full pixel motion estimation
779
      (16i16, 8i16, None)
780 781
    }
  };
782 783 784 785

  get_best_predictor(
    fi, po, p_org, p_ref, &predictors,
    bit_depth, pmv, lambda, mvx_min, mvx_max, mvy_min, mvy_max,
786
    blk_w, blk_h, center_mv, center_mv_cost,
787
    &mut tmp_plane_opt, ref_frame);
788 789 790

  loop {
    let mut best_diamond_rd_cost = std::u64::MAX;
Vladimir Kazakov's avatar
Vladimir Kazakov committed
791
    let mut best_diamond_mv = MotionVector::default();
792 793 794 795 796 797 798 799 800

    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
801
          fi, po, p_org, p_ref, bit_depth,
802
          pmv, lambda, mvx_min, mvx_max, mvy_min, mvy_max,
803
          blk_w, blk_h, cand_mv, &mut tmp_plane_opt, ref_frame);
804 805 806 807 808 809 810 811

        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 {
812
      if diamond_radius == diamond_radius_end {
813 814 815 816 817 818 819 820 821 822 823 824 825 826
        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);
}

827 828
fn get_mv_rd_cost<T: Pixel>(
  fi: &FrameInvariants<T>,
Romain Vimont's avatar
Romain Vimont committed
829
  po: PlaneOffset, p_org: &Plane<T>, p_ref: &Plane<T>, bit_depth: usize,
830
  pmv: [MotionVector; 2], lambda: u32,
831 832
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
  blk_w: usize, blk_h: usize,
833
  cand_mv: MotionVector, tmp_plane_opt: &mut Option<Plane<T>>,
Thomas Daede's avatar
Thomas Daede committed
834
  ref_frame: RefType) -> u64
835 836 837 838 839 840 841 842
{
  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;
  }

Romain Vimont's avatar
Romain Vimont committed
843
  let plane_org = p_org.region(Area::StartingAt { x: po.x, y: po.y });
844

845
  if let Some(ref mut tmp_plane) = tmp_plane_opt {
846 847 848 849 850 851 852
    let tile_rect = TileRect {
      x: 0,
      y: 0,
      width: tmp_plane.cfg.width,
      height: tmp_plane.cfg.height
    };

853 854
    PredictionMode::NEWMV.predict_inter(
      fi,
855
      tile_rect,
856
      0,
Romain Vimont's avatar
Romain Vimont committed
857
      po,
Romain Vimont's avatar
Romain Vimont committed
858
      &mut tmp_plane.as_region_mut(),
859 860 861 862 863
      blk_w,
      blk_h,
      [ref_frame, NONE_FRAME],
      [cand_mv, MotionVector { row: 0, col: 0 }]
    );
Romain Vimont's avatar
Romain Vimont committed
864
    let plane_ref = tmp_plane.as_region();
865 866 867 868 869 870
    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
871
    let plane_ref = p_ref.region(Area::StartingAt {
872 873 874 875 876 877 878 879 880 881 882 883 884 885
      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,
Romain Vimont's avatar
Romain Vimont committed
886
  plane_org: &PlaneRegion<'_, T>, plane_ref: &PlaneRegion<'_, T>
887 888
) -> u64
{
889
  let sad = get_sad(&plane_org, &plane_ref, blk_w, blk_h, bit_depth);
890 891 892 893 894 895 896 897

  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
}

898
fn telescopic_subpel_search<T: Pixel>(
899
  fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, po: PlaneOffset,
Thomas Daede's avatar
Thomas Daede committed
900
  lambda: u32, ref_frame: RefType, pmv: [MotionVector; 2],
901
  mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
902
  blk_w: usize, blk_h: usize,
Romain Vimont's avatar
Romain Vimont committed
903
  best_mv: &mut MotionVector, lowest_cost: &mut u64
904 905 906 907 908 909 910 911
) {
  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
912
  let mut tmp_plane = Plane::new(blk_w, blk_h, 0, 0, 0, 0);
913 914 915 916 917 918
  let tile_rect = TileRect {
    x: 0,
    y: 0,
    width: tmp_plane.cfg.width,
    height: tmp_plane.cfg.height
  };
Romain Vimont's avatar
Romain Vimont committed
919

920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943
  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;
        }

        {
          mode.predict_inter(
            fi,
944
            tile_rect,
945
            0,
Romain Vimont's avatar
Romain Vimont committed
946
            po,
Romain Vimont's avatar
Romain Vimont committed
947
            &mut tmp_plane.as_region_mut(),
948 949 950 951 952 953 954
            blk_w,
            blk_h,
            [ref_frame, NONE_FRAME],
            [cand_mv, MotionVector { row: 0, col: 0 }]
          );
        }

Romain Vimont's avatar
Romain Vimont committed
955 956
        let plane_org = ts.input.planes[0].region(Area::StartingAt { x: po.x, y: po.y });
        let plane_ref = tmp_plane.as_region();
957

958
        let sad = get_sad(&plane_org, &plane_ref, blk_w, blk_h, fi.sequence.bit_depth);
959 960 961 962 963 964 965 966 967 968 969 970 971 972 973

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

974
fn full_search<T: Pixel>(
Kyle Siefring's avatar
Kyle Siefring committed
975
  x_lo: isize, x_hi: isize, y_lo: isize, y_hi: isize, blk_h: usize,
976
  blk_w: usize, p_org: &Plane<T>, p_ref: &Plane<T>, best_mv: &mut MotionVector,
Romain Vimont's avatar
Romain Vimont committed
977
  lowest_cost: &mut u64, po: PlaneOffset, step: usize, bit_depth: usize,
978
  lambda: u32, pmv: [MotionVector; 2], allow_high_precision_mv: bool
Kyle Siefring's avatar
Kyle Siefring committed
979
) {
Luca Barbato's avatar
Luca Barbato committed
980 981 982 983 984
    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)| {
Romain Vimont's avatar
Romain Vimont committed
985 986
      let plane_org = p_org.region(Area::StartingAt { x: po.x, y: po.y });
      let plane_ref = p_ref.region(Area::StartingAt { x, y });
987
      let sad = get_sad(&plane_org, &plane_ref, blk_w, blk_h, bit_depth);
988

Frank Bossen's avatar
Frank Bossen committed
989 990 991 992 993 994 995 996
      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);
997
      let cost = 256 * sad as u64 + rate as u64 * lambda as u64;
Frank Bossen's avatar
Frank Bossen committed
998

Luca Barbato's avatar
Luca Barbato committed
999 1000 1001 1002 1003
      (cost, mv)
  }).min_by_key(|(c, _)| *c).unwrap();

    *lowest_cost = cost;
    *best_mv = mv;
1004 1005
}

1006 1007
// Adjust block offset such that entire block lies within boundaries
fn adjust_bo(bo: BlockOffset, mi_width: usize, mi_height: usize, blk_w: usize, blk_h: usize) -> BlockOffset {
1008
  BlockOffset {
1009 1010
    x: (bo.x as isize).min(mi_width as isize - blk_w as isize / 4).max(0) as usize,
    y: (bo.y as isize).min(mi_height as isize - blk_h as isize / 4).max(0) as usize