rav1e/api/
lookahead.rs

1use crate::api::internal::InterConfig;
2use crate::config::EncoderConfig;
3use crate::context::{BlockOffset, FrameBlocks, TileBlockOffset};
4use crate::cpu_features::CpuFeatureLevel;
5use crate::dist::get_satd;
6use crate::encoder::{
7  FrameInvariants, FrameState, Sequence, IMPORTANCE_BLOCK_SIZE,
8};
9use crate::frame::{AsRegion, PlaneOffset};
10use crate::me::{estimate_tile_motion, RefMEStats};
11use crate::partition::{get_intra_edges, BlockSize};
12use crate::predict::{IntraParam, PredictionMode};
13use crate::tiling::{Area, PlaneRegion, TileRect};
14use crate::transform::TxSize;
15use crate::util::Aligned;
16use crate::Pixel;
17use rayon::iter::*;
18use std::sync::Arc;
19use v_frame::frame::Frame;
20use v_frame::pixel::CastFromPrimitive;
21use v_frame::plane::Plane;
22
23pub(crate) const IMP_BLOCK_MV_UNITS_PER_PIXEL: i64 = 8;
24pub(crate) const IMP_BLOCK_SIZE_IN_MV_UNITS: i64 =
25  IMPORTANCE_BLOCK_SIZE as i64 * IMP_BLOCK_MV_UNITS_PER_PIXEL;
26pub(crate) const IMP_BLOCK_AREA_IN_MV_UNITS: i64 =
27  IMP_BLOCK_SIZE_IN_MV_UNITS * IMP_BLOCK_SIZE_IN_MV_UNITS;
28
29#[profiling::function]
30pub(crate) fn estimate_intra_costs<T: Pixel>(
31  temp_plane: &mut Plane<T>, frame: &Frame<T>, bit_depth: usize,
32  cpu_feature_level: CpuFeatureLevel,
33) -> Box<[u32]> {
34  let plane = &frame.planes[0];
35  let plane_after_prediction = temp_plane;
36
37  let bsize = BlockSize::from_width_and_height(
38    IMPORTANCE_BLOCK_SIZE,
39    IMPORTANCE_BLOCK_SIZE,
40  );
41  let tx_size = bsize.tx_size();
42
43  let h_in_imp_b = plane.cfg.height / IMPORTANCE_BLOCK_SIZE;
44  let w_in_imp_b = plane.cfg.width / IMPORTANCE_BLOCK_SIZE;
45  let mut intra_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);
46
47  for y in 0..h_in_imp_b {
48    for x in 0..w_in_imp_b {
49      let plane_org = plane.region(Area::Rect {
50        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
51        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
52        width: IMPORTANCE_BLOCK_SIZE,
53        height: IMPORTANCE_BLOCK_SIZE,
54      });
55
56      // TODO: other intra prediction modes.
57      let mut edge_buf = Aligned::uninit_array();
58      let edge_buf = get_intra_edges(
59        &mut edge_buf,
60        &plane.as_region(),
61        TileBlockOffset(BlockOffset { x, y }),
62        0,
63        0,
64        bsize,
65        PlaneOffset {
66          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
67          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
68        },
69        TxSize::TX_8X8,
70        bit_depth,
71        Some(PredictionMode::DC_PRED),
72        false,
73        IntraParam::None,
74      );
75
76      let mut plane_after_prediction_region = plane_after_prediction
77        .region_mut(Area::Rect {
78          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
79          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
80          width: IMPORTANCE_BLOCK_SIZE,
81          height: IMPORTANCE_BLOCK_SIZE,
82        });
83
84      PredictionMode::DC_PRED.predict_intra(
85        TileRect {
86          x: x * IMPORTANCE_BLOCK_SIZE,
87          y: y * IMPORTANCE_BLOCK_SIZE,
88          width: IMPORTANCE_BLOCK_SIZE,
89          height: IMPORTANCE_BLOCK_SIZE,
90        },
91        &mut plane_after_prediction_region,
92        tx_size,
93        bit_depth,
94        &[], // Not used by DC_PRED
95        IntraParam::None,
96        None, // Not used by DC_PRED
97        &edge_buf,
98        cpu_feature_level,
99      );
100
101      let plane_after_prediction_region =
102        plane_after_prediction.region(Area::Rect {
103          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
104          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
105          width: IMPORTANCE_BLOCK_SIZE,
106          height: IMPORTANCE_BLOCK_SIZE,
107        });
108
109      let intra_cost = get_satd(
110        &plane_org,
111        &plane_after_prediction_region,
112        bsize.width(),
113        bsize.height(),
114        bit_depth,
115        cpu_feature_level,
116      );
117
118      intra_costs.push(intra_cost);
119    }
120  }
121
122  intra_costs.into_boxed_slice()
123}
124
125#[profiling::function]
126pub(crate) fn estimate_importance_block_difference<T: Pixel>(
127  frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>,
128) -> f64 {
129  let plane_org = &frame.planes[0];
130  let plane_ref = &ref_frame.planes[0];
131  let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
132  let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
133
134  let mut imp_block_costs = 0;
135
136  (0..h_in_imp_b).for_each(|y| {
137    (0..w_in_imp_b).for_each(|x| {
138      // Coordinates of the top-left corner of the reference block, in MV
139      // units.
140      let region_org = plane_org.region(Area::Rect {
141        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
142        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
143        width: IMPORTANCE_BLOCK_SIZE,
144        height: IMPORTANCE_BLOCK_SIZE,
145      });
146
147      let region_ref = plane_ref.region(Area::Rect {
148        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
149        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
150        width: IMPORTANCE_BLOCK_SIZE,
151        height: IMPORTANCE_BLOCK_SIZE,
152      });
153
154      let sum_8x8_block = |region: &PlaneRegion<T>| {
155        region
156          .rows_iter()
157          .map(|row| {
158            // 16-bit precision is sufficient for an 8px row, as IMPORTANCE_BLOCK_SIZE * (2^12 - 1) < 2^16 - 1,
159            // so overflow is not possible
160            row.iter().map(|pixel| u16::cast_from(*pixel)).sum::<u16>() as i64
161          })
162          .sum::<i64>()
163      };
164
165      let histogram_org_sum = sum_8x8_block(&region_org);
166      let histogram_ref_sum = sum_8x8_block(&region_ref);
167
168      let count = (IMPORTANCE_BLOCK_SIZE * IMPORTANCE_BLOCK_SIZE) as i64;
169
170      let mean = (((histogram_org_sum + count / 2) / count)
171        - ((histogram_ref_sum + count / 2) / count))
172        .abs();
173
174      imp_block_costs += mean as u64;
175    });
176  });
177
178  imp_block_costs as f64 / (w_in_imp_b * h_in_imp_b) as f64
179}
180
181#[profiling::function]
182pub(crate) fn estimate_inter_costs<T: Pixel>(
183  frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>, bit_depth: usize,
184  mut config: EncoderConfig, sequence: Arc<Sequence>, buffer: RefMEStats,
185) -> f64 {
186  config.low_latency = true;
187  config.speed_settings.multiref = false;
188  let inter_cfg = InterConfig::new(&config);
189  let last_fi = FrameInvariants::new_key_frame(
190    Arc::new(config),
191    sequence,
192    0,
193    Box::new([]),
194  );
195  let mut fi = FrameInvariants::new_inter_frame(
196    &last_fi,
197    &inter_cfg,
198    0,
199    1,
200    2,
201    false,
202    Box::new([]),
203  )
204  .unwrap();
205
206  // Compute the motion vectors.
207  let mut fs = FrameState::new_with_frame_and_me_stats_and_rec(
208    &fi,
209    Arc::clone(&frame),
210    buffer,
211    // We do not use this field, so we can avoid the expensive allocation
212    Arc::new(Frame {
213      planes: [
214        Plane::new(0, 0, 0, 0, 0, 0),
215        Plane::new(0, 0, 0, 0, 0, 0),
216        Plane::new(0, 0, 0, 0, 0, 0),
217      ],
218    }),
219  );
220  compute_motion_vectors(&mut fi, &mut fs, &inter_cfg);
221
222  // Estimate inter costs
223  let plane_org = &frame.planes[0];
224  let plane_ref = &ref_frame.planes[0];
225  let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
226  let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
227  let stats = &fs.frame_me_stats.read().expect("poisoned lock")[0];
228  let bsize = BlockSize::from_width_and_height(
229    IMPORTANCE_BLOCK_SIZE,
230    IMPORTANCE_BLOCK_SIZE,
231  );
232
233  let mut inter_costs = 0;
234  (0..h_in_imp_b).for_each(|y| {
235    (0..w_in_imp_b).for_each(|x| {
236      let mv = stats[y * 2][x * 2].mv;
237
238      // Coordinates of the top-left corner of the reference block, in MV
239      // units.
240      let reference_x = x as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.col as i64;
241      let reference_y = y as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.row as i64;
242
243      let region_org = plane_org.region(Area::Rect {
244        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
245        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
246        width: IMPORTANCE_BLOCK_SIZE,
247        height: IMPORTANCE_BLOCK_SIZE,
248      });
249
250      let region_ref = plane_ref.region(Area::Rect {
251        x: reference_x as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
252        y: reference_y as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
253        width: IMPORTANCE_BLOCK_SIZE,
254        height: IMPORTANCE_BLOCK_SIZE,
255      });
256
257      inter_costs += get_satd(
258        &region_org,
259        &region_ref,
260        bsize.width(),
261        bsize.height(),
262        bit_depth,
263        fi.cpu_feature_level,
264      ) as u64;
265    });
266  });
267  inter_costs as f64 / (w_in_imp_b * h_in_imp_b) as f64
268}
269
270#[profiling::function]
271pub(crate) fn compute_motion_vectors<T: Pixel>(
272  fi: &mut FrameInvariants<T>, fs: &mut FrameState<T>, inter_cfg: &InterConfig,
273) {
274  let mut blocks = FrameBlocks::new(fi.w_in_b, fi.h_in_b);
275  fi.sequence
276    .tiling
277    .tile_iter_mut(fs, &mut blocks)
278    .collect::<Vec<_>>()
279    .into_par_iter()
280    .for_each(|mut ctx| {
281      let ts = &mut ctx.ts;
282      estimate_tile_motion(fi, ts, inter_cfg);
283    });
284}
rav1e/api/lookahead.rs

rav1e/api/
lookahead.rs
