1 /*
2 * Copyright (c) 2019, Alliance for Open Media. All rights reserved
3 *
4 * This source code is subject to the terms of the BSD 2 Clause License and
5 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6 * was not distributed with this source code in the LICENSE file, you can
7 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
8 * Media Patent License 1.0 was not distributed with this source code in the
9 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10 */
11
12 #include <float.h>
13
14 #include "config/aom_dsp_rtcd.h"
15
16 #include "aom_ports/system_state.h"
17
18 #include "av1/common/enums.h"
19 #include "av1/common/reconinter.h"
20
21 #if !CONFIG_REALTIME_ONLY
22 #include "av1/encoder/cnn.h"
23 #include "av1/encoder/partition_model_weights.h"
24 #include "av1/encoder/partition_cnn_weights.h"
25 #endif
26 #include "av1/encoder/encoder.h"
27
28 #include "av1/encoder/motion_search_facade.h"
29 #include "av1/encoder/partition_strategy.h"
30 #include "av1/encoder/rdopt.h"
31
32 #if !CONFIG_REALTIME_ONLY
33 static AOM_INLINE void simple_motion_search_prune_part_features(
34 AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
35 int mi_col, BLOCK_SIZE bsize, float *features, int features_to_get);
36 #endif
37
convert_bsize_to_idx(BLOCK_SIZE bsize)38 static INLINE int convert_bsize_to_idx(BLOCK_SIZE bsize) {
39 switch (bsize) {
40 case BLOCK_128X128: return 0;
41 case BLOCK_64X64: return 1;
42 case BLOCK_32X32: return 2;
43 case BLOCK_16X16: return 3;
44 case BLOCK_8X8: return 4;
45 default: assert(0 && "Invalid bsize"); return -1;
46 }
47 }
48
49 #if !CONFIG_REALTIME_ONLY
50 // TODO(chiyotsai@google.com): This is very much a work in progress. We still
51 // need to the following:
52 // -- add support for hdres
53 // -- add support for pruning rectangular partitions
54 // -- use reconstructed pixels instead of source pixels for padding
55 // -- use chroma pixels in addition to luma pixels
av1_intra_mode_cnn_partition(const AV1_COMMON * const cm,MACROBLOCK * x,int bsize,int quad_tree_idx,int * partition_none_allowed,int * partition_horz_allowed,int * partition_vert_allowed,int * do_rectangular_split,int * do_square_split)56 void av1_intra_mode_cnn_partition(const AV1_COMMON *const cm, MACROBLOCK *x,
57 int bsize, int quad_tree_idx,
58 int *partition_none_allowed,
59 int *partition_horz_allowed,
60 int *partition_vert_allowed,
61 int *do_rectangular_split,
62 int *do_square_split) {
63 assert(cm->seq_params.sb_size >= BLOCK_64X64 &&
64 "Invalid sb_size for intra_cnn!");
65 const int bsize_idx = convert_bsize_to_idx(bsize);
66
67 if (bsize == BLOCK_128X128) {
68 return;
69 }
70
71 // Precompute the CNN part and cache the result in MACROBLOCK
72 if (bsize == BLOCK_64X64 && !x->cnn_output_valid) {
73 aom_clear_system_state();
74 const CNN_CONFIG *cnn_config = &av1_intra_mode_cnn_partition_cnn_config;
75
76 // Prepare the output
77 const CNN_THREAD_DATA thread_data = { .num_workers = 1, .workers = NULL };
78 const int num_outputs = 4;
79 const int output_dims[4] = { 1, 2, 4, 8 };
80 const int out_chs[4] = { CNN_BRANCH_0_OUT_CH, CNN_BRANCH_1_OUT_CH,
81 CNN_BRANCH_2_OUT_CH, CNN_BRANCH_3_OUT_CH };
82 float *output_buffer[CNN_TOT_OUT_CH];
83
84 float **cur_output_buf = output_buffer;
85 float *curr_buf_ptr = x->cnn_buffer;
86 for (int output_idx = 0; output_idx < num_outputs; output_idx++) {
87 const int num_chs = out_chs[output_idx];
88 const int ch_size = output_dims[output_idx] * output_dims[output_idx];
89 for (int ch = 0; ch < num_chs; ch++) {
90 cur_output_buf[ch] = curr_buf_ptr;
91 curr_buf_ptr += ch_size;
92 }
93 cur_output_buf += num_chs;
94 }
95
96 CNN_MULTI_OUT output = {
97 .num_outputs = 4,
98 .output_channels = out_chs,
99 .output_strides = output_dims,
100 .output_buffer = output_buffer,
101 };
102
103 // Prepare the input
104 const MACROBLOCKD *xd = &x->e_mbd;
105 const int bit_depth = xd->bd;
106 const int dc_q =
107 av1_dc_quant_QTX(x->qindex, 0, bit_depth) >> (bit_depth - 8);
108 x->log_q = logf(1.0f + (float)(dc_q * dc_q) / 256.0f);
109 x->log_q = (x->log_q - av1_intra_mode_cnn_partition_mean[0]) /
110 av1_intra_mode_cnn_partition_std[0];
111
112 const int width = 65, height = 65,
113 stride = x->plane[AOM_PLANE_Y].src.stride;
114
115 if (xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH) {
116 uint16_t *image[1] = {
117 CONVERT_TO_SHORTPTR(x->plane[AOM_PLANE_Y].src.buf) - stride - 1
118 };
119
120 av1_cnn_predict_img_multi_out_highbd(image, width, height, stride,
121 cnn_config, &thread_data, bit_depth,
122 &output);
123 } else {
124 uint8_t *image[1] = { x->plane[AOM_PLANE_Y].src.buf - stride - 1 };
125
126 av1_cnn_predict_img_multi_out(image, width, height, stride, cnn_config,
127 &thread_data, &output);
128 }
129
130 x->cnn_output_valid = 1;
131 }
132
133 if (!x->cnn_output_valid) {
134 return;
135 }
136
137 const NN_CONFIG *dnn_configs[5] = {
138 NULL,
139 &av1_intra_mode_cnn_partition_branch_0_dnn_config,
140 &av1_intra_mode_cnn_partition_branch_1_dnn_config,
141 &av1_intra_mode_cnn_partition_branch_2_dnn_config,
142 &av1_intra_mode_cnn_partition_branch_3_dnn_config,
143 };
144
145 const NN_CONFIG *dnn_config = dnn_configs[bsize_idx];
146
147 aom_clear_system_state();
148 float dnn_features[100];
149 float logits[4] = { 0.0f };
150
151 const float *branch_0 = x->cnn_buffer;
152 const float *branch_1 = branch_0 + CNN_BRANCH_0_OUT_SIZE;
153 const float *branch_2 = branch_1 + CNN_BRANCH_1_OUT_SIZE;
154 const float *branch_3 = branch_2 + CNN_BRANCH_2_OUT_SIZE;
155
156 if (bsize == BLOCK_64X64) {
157 int f_idx = 0;
158 for (int ch_idx = 0; ch_idx < CNN_BRANCH_0_OUT_CH; ch_idx++) {
159 dnn_features[f_idx++] = branch_0[ch_idx];
160 }
161
162 const int spa_stride = 2 * 2;
163 for (int lin_idx = 0; lin_idx < spa_stride; lin_idx++) {
164 for (int ch_idx = 0; ch_idx < CNN_BRANCH_1_OUT_CH; ch_idx++) {
165 dnn_features[f_idx++] = branch_1[lin_idx + ch_idx * spa_stride];
166 }
167 }
168 dnn_features[f_idx++] = x->log_q;
169 } else if (bsize == BLOCK_32X32) {
170 int f_idx = 0;
171 for (int idx = 0; idx < CNN_BRANCH_0_OUT_CH; idx++) {
172 dnn_features[f_idx++] = branch_0[idx];
173 }
174
175 const int curr_lin_idx = quad_to_linear_1[quad_tree_idx - 1];
176 const int spa_stride = 2 * 2;
177 for (int ch_idx = 0; ch_idx < CNN_BRANCH_1_OUT_CH; ch_idx++) {
178 dnn_features[f_idx++] = branch_1[curr_lin_idx + ch_idx * spa_stride];
179 }
180 dnn_features[f_idx++] = x->log_q;
181 } else if (bsize == BLOCK_16X16) {
182 int f_idx = 0;
183 const int prev_quad_idx = (quad_tree_idx - 1) / 4;
184 const int prev_lin_idx = quad_to_linear_1[prev_quad_idx - 1];
185 const int prev_spa_stride = 2 * 2;
186 for (int ch_idx = 0; ch_idx < CNN_BRANCH_1_OUT_CH; ch_idx++) {
187 dnn_features[f_idx++] = branch_1[prev_lin_idx + ch_idx * prev_spa_stride];
188 }
189
190 const int curr_lin_idx = quad_to_linear_2[quad_tree_idx - 5];
191 const int spa_stride = 4 * 4;
192 for (int ch_idx = 0; ch_idx < CNN_BRANCH_2_OUT_CH; ch_idx++) {
193 dnn_features[f_idx++] = branch_2[curr_lin_idx + ch_idx * spa_stride];
194 }
195 dnn_features[f_idx++] = x->log_q;
196 } else if (bsize == BLOCK_8X8) {
197 int f_idx = 0;
198 const int prev_quad_idx = (quad_tree_idx - 1) / 4;
199 const int prev_lin_idx = quad_to_linear_2[prev_quad_idx - 5];
200 const int prev_spa_stride = 4 * 4;
201 for (int ch_idx = 0; ch_idx < CNN_BRANCH_2_OUT_CH; ch_idx++) {
202 dnn_features[f_idx++] = branch_2[prev_lin_idx + ch_idx * prev_spa_stride];
203 }
204
205 const int curr_lin_idx = quad_to_linear_3[quad_tree_idx - 21];
206 const int spa_stride = 8 * 8;
207 for (int ch_idx = 0; ch_idx < CNN_BRANCH_3_OUT_CH; ch_idx++) {
208 dnn_features[f_idx++] = branch_3[curr_lin_idx + ch_idx * spa_stride];
209 }
210 dnn_features[f_idx++] = x->log_q;
211 } else {
212 assert(0 && "Invalid bsize in intra_cnn partition");
213 }
214
215 // Make decision
216 av1_nn_predict(dnn_features, dnn_config, 1, logits);
217 aom_clear_system_state();
218
219 const int is_720p_or_larger = AOMMIN(cm->width, cm->height) >= 720;
220 const int is_480p_or_larger = AOMMIN(cm->width, cm->height) >= 480;
221 float split_only_thresh = 100.0f, no_split_thresh = -100.0f;
222 if (is_720p_or_larger) {
223 split_only_thresh =
224 av1_intra_mode_cnn_partition_split_thresh_hdres[bsize_idx];
225 no_split_thresh =
226 av1_intra_mode_cnn_partition_no_split_thresh_hdres[bsize_idx];
227 } else if (is_480p_or_larger) {
228 split_only_thresh =
229 av1_intra_mode_cnn_partition_split_thresh_midres[bsize_idx];
230 no_split_thresh =
231 av1_intra_mode_cnn_partition_no_split_thresh_midres[bsize_idx];
232 } else {
233 split_only_thresh =
234 av1_intra_mode_cnn_partition_split_thresh_lowres[bsize_idx];
235 no_split_thresh =
236 av1_intra_mode_cnn_partition_no_split_thresh_lowres[bsize_idx];
237 }
238
239 if (logits[0] > split_only_thresh) {
240 *partition_none_allowed = 0;
241 *partition_horz_allowed = 0;
242 *partition_vert_allowed = 0;
243 *do_rectangular_split = 0;
244 }
245
246 if (logits[0] < no_split_thresh) {
247 *do_square_split = 0;
248 }
249 }
250
av1_simple_motion_search_based_split(AV1_COMP * const cpi,MACROBLOCK * x,PC_TREE * pc_tree,int mi_row,int mi_col,BLOCK_SIZE bsize,int * partition_none_allowed,int * partition_horz_allowed,int * partition_vert_allowed,int * do_rectangular_split,int * do_square_split)251 void av1_simple_motion_search_based_split(
252 AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
253 int mi_col, BLOCK_SIZE bsize, int *partition_none_allowed,
254 int *partition_horz_allowed, int *partition_vert_allowed,
255 int *do_rectangular_split, int *do_square_split) {
256 aom_clear_system_state();
257
258 const AV1_COMMON *const cm = &cpi->common;
259 const int bsize_idx = convert_bsize_to_idx(bsize);
260 const int is_720p_or_larger = AOMMIN(cm->width, cm->height) >= 720;
261 const int is_480p_or_larger = AOMMIN(cm->width, cm->height) >= 480;
262 // res_idx is 0 for res < 480p, 1 for 480p, 2 for 720p+
263 const int res_idx = is_480p_or_larger + is_720p_or_larger;
264
265 assert(bsize_idx >= 0 && bsize_idx <= 4 &&
266 "Invalid bsize in simple_motion_search_based_split");
267
268 const float *ml_mean = av1_simple_motion_search_split_mean[bsize_idx];
269 const float *ml_std = av1_simple_motion_search_split_std[bsize_idx];
270 const NN_CONFIG *nn_config =
271 av1_simple_motion_search_split_nn_config[bsize_idx];
272 const int agg = cpi->sf.part_sf.simple_motion_search_prune_agg;
273
274 const float split_only_thresh =
275 av1_simple_motion_search_split_thresh[agg][res_idx][bsize_idx];
276 const float no_split_thresh =
277 av1_simple_motion_search_no_split_thresh[agg][res_idx][bsize_idx];
278
279 float features[FEATURE_SIZE_SMS_SPLIT] = { 0.0f };
280 simple_motion_search_prune_part_features(cpi, x, pc_tree, mi_row, mi_col,
281 bsize, features,
282 FEATURE_SMS_SPLIT_MODEL_FLAG);
283 for (int idx = 0; idx < FEATURE_SIZE_SMS_SPLIT; idx++) {
284 features[idx] = (features[idx] - ml_mean[idx]) / ml_std[idx];
285 }
286
287 float score = 0.0f;
288
289 av1_nn_predict(features, nn_config, 1, &score);
290 aom_clear_system_state();
291
292 if (score > split_only_thresh) {
293 *partition_none_allowed = 0;
294 *partition_horz_allowed = 0;
295 *partition_vert_allowed = 0;
296 *do_rectangular_split = 0;
297 }
298
299 if (cpi->sf.part_sf.simple_motion_search_split >= 2 &&
300 score < no_split_thresh) {
301 *do_square_split = 0;
302 }
303 }
304
305 // Given a list of ref frames in refs, performs simple_motion_search on each of
306 // the refs and returns the ref with the smallest sse. Returns -1 if none of the
307 // ref in the list is available. Also stores the best sse and var in best_sse,
308 // best_var, respectively. If save_mv is 0, don't update mv_ref_fulls in
309 // pc_tree. If save_mv is 1, update mv_ref_fulls under pc_tree and the
310 // subtrees.
simple_motion_search_get_best_ref(AV1_COMP * const cpi,MACROBLOCK * x,PC_TREE * pc_tree,int mi_row,int mi_col,BLOCK_SIZE bsize,const int * const refs,int num_refs,int use_subpixel,int save_mv,unsigned int * best_sse,unsigned int * best_var)311 static int simple_motion_search_get_best_ref(
312 AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
313 int mi_col, BLOCK_SIZE bsize, const int *const refs, int num_refs,
314 int use_subpixel, int save_mv, unsigned int *best_sse,
315 unsigned int *best_var) {
316 const AV1_COMMON *const cm = &cpi->common;
317 int best_ref = -1;
318
319 if (mi_col >= cm->mi_params.mi_cols || mi_row >= cm->mi_params.mi_rows) {
320 // If the whole block is outside of the image, set the var and sse to 0.
321 *best_var = 0;
322 *best_sse = 0;
323
324 return best_ref;
325 }
326
327 // Otherwise do loop through the reference frames and find the one with the
328 // minimum SSE
329 const MACROBLOCKD *xd = &x->e_mbd;
330
331 const int num_planes = 1;
332
333 *best_sse = INT_MAX;
334
335 for (int ref_idx = 0; ref_idx < num_refs; ref_idx++) {
336 const int ref = refs[ref_idx];
337
338 if (cpi->ref_frame_flags & av1_ref_frame_flag_list[ref]) {
339 const FULLPEL_MV *start_mvs = pc_tree->start_mvs;
340 unsigned int curr_sse = 0, curr_var = 0;
341 int_mv best_mv =
342 av1_simple_motion_search(cpi, x, mi_row, mi_col, bsize, ref,
343 start_mvs[ref], num_planes, use_subpixel);
344 curr_var = cpi->fn_ptr[bsize].vf(
345 x->plane[0].src.buf, x->plane[0].src.stride, xd->plane[0].dst.buf,
346 xd->plane[0].dst.stride, &curr_sse);
347 if (curr_sse < *best_sse) {
348 *best_sse = curr_sse;
349 *best_var = curr_var;
350 best_ref = ref;
351 }
352
353 if (save_mv) {
354 pc_tree->start_mvs[ref].row = best_mv.as_mv.row / 8;
355 pc_tree->start_mvs[ref].col = best_mv.as_mv.col / 8;
356
357 if (bsize >= BLOCK_8X8) {
358 for (int r_idx = 0; r_idx < 4; r_idx++) {
359 // Propagate the new motion vectors to a lower level
360 PC_TREE *sub_tree = pc_tree->split[r_idx];
361 sub_tree->start_mvs[ref] = pc_tree->start_mvs[ref];
362 }
363 }
364 }
365 }
366 }
367
368 return best_ref;
369 }
370
371 // Collects features using simple_motion_search and store them in features. The
372 // features are also cached in PC_TREE. By default, the features collected are
373 // the sse and var from the subblocks flagged by features_to_get. Furthermore,
374 // if features is not NULL, then 7 more features are appended to the end of
375 // features:
376 // - log(1.0 + dc_q ** 2)
377 // - whether an above macroblock exists
378 // - width of above macroblock
379 // - height of above macroblock
380 // - whether a left marcoblock exists
381 // - width of left macroblock
382 // - height of left macroblock
simple_motion_search_prune_part_features(AV1_COMP * const cpi,MACROBLOCK * x,PC_TREE * pc_tree,int mi_row,int mi_col,BLOCK_SIZE bsize,float * features,int features_to_get)383 static AOM_INLINE void simple_motion_search_prune_part_features(
384 AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
385 int mi_col, BLOCK_SIZE bsize, float *features, int features_to_get) {
386 const int w_mi = mi_size_wide[bsize];
387 const int h_mi = mi_size_high[bsize];
388 assert(mi_size_wide[bsize] == mi_size_high[bsize]);
389 assert(cpi->ref_frame_flags & av1_ref_frame_flag_list[LAST_FRAME] ||
390 cpi->ref_frame_flags & av1_ref_frame_flag_list[ALTREF_FRAME]);
391
392 // Setting up motion search
393 const int ref_list[] = { cpi->rc.is_src_frame_alt_ref ? ALTREF_FRAME
394 : LAST_FRAME };
395 const int num_refs = 1;
396 const int use_subpixel = 1;
397
398 // Doing whole block first to update the mv
399 if (!pc_tree->sms_none_valid && features_to_get & FEATURE_SMS_NONE_FLAG) {
400 simple_motion_search_get_best_ref(cpi, x, pc_tree, mi_row, mi_col, bsize,
401 ref_list, num_refs, use_subpixel, 1,
402 &pc_tree->sms_none_feat[0],
403 &pc_tree->sms_none_feat[1]);
404 pc_tree->sms_none_valid = 1;
405 }
406
407 // Split subblocks
408 if (features_to_get & FEATURE_SMS_SPLIT_FLAG) {
409 const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
410 for (int r_idx = 0; r_idx < 4; r_idx++) {
411 const int sub_mi_col = mi_col + (r_idx & 1) * w_mi / 2;
412 const int sub_mi_row = mi_row + (r_idx >> 1) * h_mi / 2;
413 PC_TREE *sub_tree = pc_tree->split[r_idx];
414
415 if (!sub_tree->sms_none_valid) {
416 simple_motion_search_get_best_ref(
417 cpi, x, sub_tree, sub_mi_row, sub_mi_col, subsize, ref_list,
418 num_refs, use_subpixel, 1, &sub_tree->sms_none_feat[0],
419 &sub_tree->sms_none_feat[1]);
420 sub_tree->sms_none_valid = 1;
421 }
422 }
423 }
424
425 // Rectangular subblocks
426 if (!pc_tree->sms_rect_valid && features_to_get & FEATURE_SMS_RECT_FLAG) {
427 // Horz subblock
428 BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_HORZ);
429 for (int r_idx = 0; r_idx < 2; r_idx++) {
430 const int sub_mi_col = mi_col + 0;
431 const int sub_mi_row = mi_row + r_idx * h_mi / 2;
432
433 simple_motion_search_get_best_ref(
434 cpi, x, pc_tree, sub_mi_row, sub_mi_col, subsize, ref_list, num_refs,
435 use_subpixel, 0, &pc_tree->sms_rect_feat[2 * r_idx],
436 &pc_tree->sms_rect_feat[2 * r_idx + 1]);
437 }
438
439 // Vert subblock
440 subsize = get_partition_subsize(bsize, PARTITION_VERT);
441 for (int r_idx = 0; r_idx < 2; r_idx++) {
442 const int sub_mi_col = mi_col + r_idx * w_mi / 2;
443 const int sub_mi_row = mi_row + 0;
444
445 simple_motion_search_get_best_ref(
446 cpi, x, pc_tree, sub_mi_row, sub_mi_col, subsize, ref_list, num_refs,
447 use_subpixel, 0, &pc_tree->sms_rect_feat[4 + 2 * r_idx],
448 &pc_tree->sms_rect_feat[4 + 2 * r_idx + 1]);
449 }
450 pc_tree->sms_rect_valid = 1;
451 }
452
453 if (!features) return;
454
455 aom_clear_system_state();
456 int f_idx = 0;
457 if (features_to_get & FEATURE_SMS_NONE_FLAG) {
458 for (int sub_idx = 0; sub_idx < 2; sub_idx++) {
459 features[f_idx++] = logf(1.0f + pc_tree->sms_none_feat[sub_idx]);
460 }
461 }
462
463 if (features_to_get & FEATURE_SMS_SPLIT_FLAG) {
464 for (int sub_idx = 0; sub_idx < 4; sub_idx++) {
465 PC_TREE *sub_tree = pc_tree->split[sub_idx];
466 features[f_idx++] = logf(1.0f + sub_tree->sms_none_feat[0]);
467 features[f_idx++] = logf(1.0f + sub_tree->sms_none_feat[1]);
468 }
469 }
470
471 if (features_to_get & FEATURE_SMS_RECT_FLAG) {
472 for (int sub_idx = 0; sub_idx < 8; sub_idx++) {
473 features[f_idx++] = logf(1.0f + pc_tree->sms_rect_feat[sub_idx]);
474 }
475 }
476
477 const MACROBLOCKD *xd = &x->e_mbd;
478 set_offsets_for_motion_search(cpi, x, mi_row, mi_col, bsize);
479
480 // Q_INDEX
481 const int dc_q = av1_dc_quant_QTX(x->qindex, 0, xd->bd) >> (xd->bd - 8);
482 features[f_idx++] = logf(1.0f + (float)(dc_q * dc_q) / 256.0f);
483
484 // Neighbor stuff
485 const int has_above = !!xd->above_mbmi;
486 const int has_left = !!xd->left_mbmi;
487 const BLOCK_SIZE above_bsize = has_above ? xd->above_mbmi->sb_type : bsize;
488 const BLOCK_SIZE left_bsize = has_left ? xd->left_mbmi->sb_type : bsize;
489 features[f_idx++] = (float)has_above;
490 features[f_idx++] = (float)mi_size_wide_log2[above_bsize];
491 features[f_idx++] = (float)mi_size_high_log2[above_bsize];
492 features[f_idx++] = (float)has_left;
493 features[f_idx++] = (float)mi_size_wide_log2[left_bsize];
494 features[f_idx++] = (float)mi_size_high_log2[left_bsize];
495 }
496
av1_simple_motion_search_prune_rect(AV1_COMP * const cpi,MACROBLOCK * x,PC_TREE * pc_tree,int mi_row,int mi_col,BLOCK_SIZE bsize,int * partition_horz_allowed,int * partition_vert_allowed,int * prune_horz,int * prune_vert)497 void av1_simple_motion_search_prune_rect(AV1_COMP *const cpi, MACROBLOCK *x,
498 PC_TREE *pc_tree, int mi_row,
499 int mi_col, BLOCK_SIZE bsize,
500 int *partition_horz_allowed,
501 int *partition_vert_allowed,
502 int *prune_horz, int *prune_vert) {
503 aom_clear_system_state();
504 const AV1_COMMON *const cm = &cpi->common;
505 const int bsize_idx = convert_bsize_to_idx(bsize);
506 const int is_720p_or_larger = AOMMIN(cm->width, cm->height) >= 720;
507 const int is_480p_or_larger = AOMMIN(cm->width, cm->height) >= 480;
508 // res_idx is 0 for lowres, 1 for 48p, 2 for 720p+
509 const int res_idx = is_480p_or_larger + is_720p_or_larger;
510
511 // Get model parameters
512 const NN_CONFIG *nn_config =
513 av1_simple_motion_search_prune_rect_nn_config[bsize_idx];
514 const float *ml_mean = av1_simple_motion_search_prune_rect_mean[bsize_idx],
515 *ml_std = av1_simple_motion_search_prune_rect_std[bsize_idx];
516
517 const int agg = cpi->sf.part_sf.simple_motion_search_prune_agg;
518 const float prune_thresh =
519 av1_simple_motion_search_prune_rect_thresh[agg][res_idx][bsize_idx];
520
521 // If there is no valid threshold, return immediately.
522 if (!nn_config || prune_thresh == 0.0f) {
523 return;
524 }
525
526 // Get features
527 float features[FEATURE_SIZE_SMS_PRUNE_PART] = { 0.0f };
528 simple_motion_search_prune_part_features(cpi, x, pc_tree, mi_row, mi_col,
529 bsize, features,
530 FEATURE_SMS_PRUNE_PART_FLAG);
531 for (int f_idx = 0; f_idx < FEATURE_SIZE_SMS_PRUNE_PART; f_idx++) {
532 features[f_idx] = (features[f_idx] - ml_mean[f_idx]) / ml_std[f_idx];
533 }
534
535 // Get probabilities
536 float scores[EXT_PARTITION_TYPES] = { 0.0f },
537 probs[EXT_PARTITION_TYPES] = { 0.0f };
538 const int num_classes = (bsize == BLOCK_128X128 || bsize == BLOCK_8X8)
539 ? PARTITION_TYPES
540 : EXT_PARTITION_TYPES;
541
542 av1_nn_predict(features, nn_config, 1, scores);
543 aom_clear_system_state();
544
545 av1_nn_softmax(scores, probs, num_classes);
546
547 // Determine if we should prune rectangular partitions.
548 if (cpi->sf.part_sf.simple_motion_search_prune_rect &&
549 !frame_is_intra_only(cm) &&
550 (*partition_horz_allowed || *partition_vert_allowed) &&
551 bsize >= BLOCK_8X8 && !av1_superres_scaled(cm)) {
552 *prune_horz = probs[PARTITION_HORZ] <= prune_thresh;
553 *prune_vert = probs[PARTITION_VERT] <= prune_thresh;
554 }
555 }
556
557 // Early terminates PARTITION_NONE using simple_motion_search features and the
558 // rate, distortion, and rdcost of PARTITION_NONE. This is only called when:
559 // - The frame is a show frame
560 // - The frame is not intra only
561 // - The current bsize is > BLOCK_8X8
562 // - blk_row + blk_height/2 < total_rows and blk_col + blk_width/2 < total_cols
av1_simple_motion_search_early_term_none(AV1_COMP * const cpi,MACROBLOCK * x,PC_TREE * pc_tree,int mi_row,int mi_col,BLOCK_SIZE bsize,const RD_STATS * none_rdc,int * early_terminate)563 void av1_simple_motion_search_early_term_none(AV1_COMP *const cpi,
564 MACROBLOCK *x, PC_TREE *pc_tree,
565 int mi_row, int mi_col,
566 BLOCK_SIZE bsize,
567 const RD_STATS *none_rdc,
568 int *early_terminate) {
569 // TODO(chiyotsai@google.com): There are other features we can extract from
570 // PARTITION_NONE. Play with this later.
571 float features[FEATURE_SIZE_SMS_TERM_NONE] = { 0.0f };
572 simple_motion_search_prune_part_features(cpi, x, pc_tree, mi_row, mi_col,
573 bsize, features,
574 FEATURE_SMS_PRUNE_PART_FLAG);
575 int f_idx = FEATURE_SIZE_SMS_PRUNE_PART;
576
577 features[f_idx++] = logf(1.0f + (float)none_rdc->rate);
578 features[f_idx++] = logf(1.0f + (float)none_rdc->dist);
579 features[f_idx++] = logf(1.0f + (float)none_rdc->rdcost);
580
581 assert(f_idx == FEATURE_SIZE_SMS_TERM_NONE);
582
583 const float *ml_mean = NULL;
584 const float *ml_std = NULL;
585 const float *ml_model = NULL;
586
587 if (bsize == BLOCK_128X128) {
588 ml_mean = av1_simple_motion_search_term_none_mean_128;
589 ml_std = av1_simple_motion_search_term_none_std_128;
590 ml_model = av1_simple_motion_search_term_none_model_128;
591 } else if (bsize == BLOCK_64X64) {
592 ml_mean = av1_simple_motion_search_term_none_mean_64;
593 ml_std = av1_simple_motion_search_term_none_std_64;
594 ml_model = av1_simple_motion_search_term_none_model_64;
595 } else if (bsize == BLOCK_32X32) {
596 ml_mean = av1_simple_motion_search_term_none_mean_32;
597 ml_std = av1_simple_motion_search_term_none_std_32;
598 ml_model = av1_simple_motion_search_term_none_model_32;
599 } else if (bsize == BLOCK_16X16) {
600 ml_mean = av1_simple_motion_search_term_none_mean_16;
601 ml_std = av1_simple_motion_search_term_none_std_16;
602 ml_model = av1_simple_motion_search_term_none_model_16;
603 } else {
604 assert(0 && "Unexpected block size in simple_motion_term_none");
605 }
606
607 if (ml_model) {
608 float score = 0.0f;
609 for (f_idx = 0; f_idx < FEATURE_SIZE_SMS_TERM_NONE; f_idx++) {
610 score +=
611 ml_model[f_idx] * (features[f_idx] - ml_mean[f_idx]) / ml_std[f_idx];
612 }
613 score += ml_model[FEATURE_SIZE_SMS_TERM_NONE];
614
615 if (score >= 0.0f) {
616 *early_terminate = 1;
617 }
618 }
619 }
620
av1_get_max_min_partition_features(AV1_COMP * const cpi,MACROBLOCK * x,int mi_row,int mi_col,float * features)621 void av1_get_max_min_partition_features(AV1_COMP *const cpi, MACROBLOCK *x,
622 int mi_row, int mi_col,
623 float *features) {
624 AV1_COMMON *const cm = &cpi->common;
625 MACROBLOCKD *xd = &x->e_mbd;
626 const BLOCK_SIZE sb_size = cm->seq_params.sb_size;
627
628 assert(sb_size == BLOCK_128X128);
629
630 int f_idx = 0;
631
632 const int dc_q = av1_dc_quant_QTX(x->qindex, 0, xd->bd) >> (xd->bd - 8);
633 aom_clear_system_state();
634 const float log_q_sq = logf(1.0f + (float)(dc_q * dc_q) / 256.0f);
635
636 // Perform full-pixel single motion search in Y plane of 16x16 mbs in the sb
637 float sum_mv_row_sq = 0;
638 float sum_mv_row = 0;
639 float min_abs_mv_row = FLT_MAX;
640 float max_abs_mv_row = 0;
641
642 float sum_mv_col_sq = 0;
643 float sum_mv_col = 0;
644 float min_abs_mv_col = FLT_MAX;
645 float max_abs_mv_col = 0;
646
647 float sum_log_sse_sq = 0;
648 float sum_log_sse = 0;
649 float min_log_sse = FLT_MAX;
650 float max_log_sse = 0;
651
652 const BLOCK_SIZE mb_size = BLOCK_16X16;
653 const int mb_rows = block_size_high[sb_size] / block_size_high[mb_size];
654 const int mb_cols = block_size_wide[sb_size] / block_size_wide[mb_size];
655 const int mb_in_mi_size_high_log2 = mi_size_high_log2[mb_size];
656 const int mb_in_mi_size_wide_log2 = mi_size_wide_log2[mb_size];
657
658 for (int mb_row = 0; mb_row < mb_rows; mb_row++)
659 for (int mb_col = 0; mb_col < mb_cols; mb_col++) {
660 const int this_mi_row = mi_row + (mb_row << mb_in_mi_size_high_log2);
661 const int this_mi_col = mi_col + (mb_col << mb_in_mi_size_wide_log2);
662 unsigned int sse = 0;
663 unsigned int var = 0;
664 const FULLPEL_MV start_mv = kZeroFullMv;
665 int_mv best_mv = av1_simple_motion_sse_var(
666 cpi, x, this_mi_row, this_mi_col, mb_size, start_mv, 0, &sse, &var);
667
668 aom_clear_system_state();
669 const float mv_row = (float)(best_mv.as_mv.row / 8);
670 const float mv_col = (float)(best_mv.as_mv.col / 8);
671 const float log_sse = logf(1.0f + (float)sse);
672 const float abs_mv_row = fabsf(mv_row);
673 const float abs_mv_col = fabsf(mv_col);
674
675 sum_mv_row_sq += mv_row * mv_row;
676 sum_mv_row += mv_row;
677 sum_mv_col_sq += mv_col * mv_col;
678 sum_mv_col += mv_col;
679
680 if (abs_mv_row < min_abs_mv_row) min_abs_mv_row = abs_mv_row;
681 if (abs_mv_row > max_abs_mv_row) max_abs_mv_row = abs_mv_row;
682 if (abs_mv_col < min_abs_mv_col) min_abs_mv_col = abs_mv_col;
683 if (abs_mv_col > max_abs_mv_col) max_abs_mv_col = abs_mv_col;
684
685 sum_log_sse_sq += log_sse * log_sse;
686 sum_log_sse += log_sse;
687 if (log_sse < min_log_sse) min_log_sse = log_sse;
688 if (log_sse > max_log_sse) max_log_sse = log_sse;
689 }
690 aom_clear_system_state();
691 const float avg_mv_row = sum_mv_row / 64.0f;
692 const float var_mv_row = sum_mv_row_sq / 64.0f - avg_mv_row * avg_mv_row;
693
694 const float avg_mv_col = sum_mv_col / 64.0f;
695 const float var_mv_col = sum_mv_col_sq / 64.0f - avg_mv_col * avg_mv_col;
696
697 const float avg_log_sse = sum_log_sse / 64.0f;
698 const float var_log_sse = sum_log_sse_sq / 64.0f - avg_log_sse * avg_log_sse;
699
700 features[f_idx++] = avg_log_sse;
701 features[f_idx++] = avg_mv_col;
702 features[f_idx++] = avg_mv_row;
703 features[f_idx++] = log_q_sq;
704 features[f_idx++] = max_abs_mv_col;
705 features[f_idx++] = max_abs_mv_row;
706 features[f_idx++] = max_log_sse;
707 features[f_idx++] = min_abs_mv_col;
708 features[f_idx++] = min_abs_mv_row;
709 features[f_idx++] = min_log_sse;
710 features[f_idx++] = var_log_sse;
711 features[f_idx++] = var_mv_col;
712 features[f_idx++] = var_mv_row;
713
714 assert(f_idx == FEATURE_SIZE_MAX_MIN_PART_PRED);
715 }
716
av1_predict_max_partition(AV1_COMP * const cpi,MACROBLOCK * const x,const float * features)717 BLOCK_SIZE av1_predict_max_partition(AV1_COMP *const cpi, MACROBLOCK *const x,
718 const float *features) {
719 float scores[MAX_NUM_CLASSES_MAX_MIN_PART_PRED] = { 0.0f },
720 probs[MAX_NUM_CLASSES_MAX_MIN_PART_PRED] = { 0.0f };
721 const NN_CONFIG *nn_config = &av1_max_part_pred_nn_config;
722
723 assert(cpi->sf.part_sf.auto_max_partition_based_on_simple_motion !=
724 NOT_IN_USE);
725
726 aom_clear_system_state();
727 av1_nn_predict(features, nn_config, 1, scores);
728 av1_nn_softmax(scores, probs, MAX_NUM_CLASSES_MAX_MIN_PART_PRED);
729
730 int result = MAX_NUM_CLASSES_MAX_MIN_PART_PRED - 1;
731 if (cpi->sf.part_sf.auto_max_partition_based_on_simple_motion ==
732 DIRECT_PRED) {
733 result = 0;
734 float max_prob = probs[0];
735 for (int i = 1; i < MAX_NUM_CLASSES_MAX_MIN_PART_PRED; ++i) {
736 if (probs[i] > max_prob) {
737 max_prob = probs[i];
738 result = i;
739 }
740 }
741 } else if (cpi->sf.part_sf.auto_max_partition_based_on_simple_motion ==
742 RELAXED_PRED) {
743 for (result = MAX_NUM_CLASSES_MAX_MIN_PART_PRED - 1; result >= 0;
744 --result) {
745 if (result < MAX_NUM_CLASSES_MAX_MIN_PART_PRED - 1) {
746 probs[result] += probs[result + 1];
747 }
748 if (probs[result] > 0.2) break;
749 }
750 } else if (cpi->sf.part_sf.auto_max_partition_based_on_simple_motion ==
751 ADAPT_PRED) {
752 const BLOCK_SIZE sb_size = cpi->common.seq_params.sb_size;
753 MACROBLOCKD *const xd = &x->e_mbd;
754 // TODO(debargha): x->source_variance is unavailable at this point,
755 // so compute. The redundant recomputation later can be removed.
756 const unsigned int source_variance =
757 is_cur_buf_hbd(xd)
758 ? av1_high_get_sby_perpixel_variance(cpi, &x->plane[0].src, sb_size,
759 xd->bd)
760 : av1_get_sby_perpixel_variance(cpi, &x->plane[0].src, sb_size);
761 if (source_variance > 16) {
762 const double thresh = source_variance < 128 ? 0.05 : 0.1;
763 for (result = MAX_NUM_CLASSES_MAX_MIN_PART_PRED - 1; result >= 0;
764 --result) {
765 if (result < MAX_NUM_CLASSES_MAX_MIN_PART_PRED - 1) {
766 probs[result] += probs[result + 1];
767 }
768 if (probs[result] > thresh) break;
769 }
770 }
771 }
772
773 return (BLOCK_SIZE)((result + 2) * 3);
774 }
775
776 // Get the minimum partition block width and height(in log scale) under a
777 // PC_TREE.
get_min_bsize(const PC_TREE * pc_tree,int * min_bw,int * min_bh)778 static AOM_INLINE void get_min_bsize(const PC_TREE *pc_tree, int *min_bw,
779 int *min_bh) {
780 if (!pc_tree) return;
781
782 const BLOCK_SIZE bsize = pc_tree->block_size;
783 if (bsize == BLOCK_4X4) {
784 *min_bw = 0;
785 *min_bh = 0;
786 return;
787 }
788
789 PARTITION_TYPE part_type = pc_tree->partitioning;
790 if (part_type == PARTITION_INVALID) return;
791
792 if (part_type == PARTITION_SPLIT) {
793 for (int i = 0; i < 4; ++i) {
794 get_min_bsize(pc_tree->split[i], min_bw, min_bh);
795 }
796 } else {
797 if (part_type == PARTITION_HORZ_A || part_type == PARTITION_HORZ_B ||
798 part_type == PARTITION_VERT_A || part_type == PARTITION_VERT_B)
799 part_type = PARTITION_SPLIT;
800 const BLOCK_SIZE subsize = get_partition_subsize(bsize, part_type);
801 if (subsize != BLOCK_INVALID) {
802 *min_bw = AOMMIN(*min_bw, mi_size_wide_log2[subsize]);
803 *min_bh = AOMMIN(*min_bh, mi_size_high_log2[subsize]);
804 }
805 }
806 }
807
add_rd_feature(int64_t rd,int64_t best_rd,float * features,int * feature_idx)808 static INLINE void add_rd_feature(int64_t rd, int64_t best_rd, float *features,
809 int *feature_idx) {
810 const int rd_valid = rd > 0 && rd < INT64_MAX;
811 const float rd_ratio = rd_valid ? (float)rd / best_rd : 1.0f;
812 features[(*feature_idx)++] = (float)rd_valid;
813 features[(*feature_idx)++] = rd_ratio;
814 }
815
816 #define FEATURES 31
av1_ml_early_term_after_split(AV1_COMP * const cpi,MACROBLOCK * const x,PC_TREE * const pc_tree,BLOCK_SIZE bsize,int64_t best_rd,int64_t part_none_rd,int64_t part_split_rd,int64_t * split_block_rd,int mi_row,int mi_col,int * const terminate_partition_search)817 void av1_ml_early_term_after_split(AV1_COMP *const cpi, MACROBLOCK *const x,
818 PC_TREE *const pc_tree, BLOCK_SIZE bsize,
819 int64_t best_rd, int64_t part_none_rd,
820 int64_t part_split_rd,
821 int64_t *split_block_rd, int mi_row,
822 int mi_col,
823 int *const terminate_partition_search) {
824 if (best_rd <= 0 || best_rd == INT64_MAX || *terminate_partition_search)
825 return;
826
827 const AV1_COMMON *const cm = &cpi->common;
828 const int is_480p_or_larger = AOMMIN(cm->width, cm->height) >= 480;
829 const NN_CONFIG *nn_config = NULL;
830 float thresh = -1e6;
831 switch (bsize) {
832 case BLOCK_128X128: break;
833 case BLOCK_64X64:
834 nn_config = &av1_early_term_after_split_nnconfig_64;
835 thresh = is_480p_or_larger ? -2.0f : -1.2f;
836 break;
837 case BLOCK_32X32:
838 nn_config = &av1_early_term_after_split_nnconfig_32;
839 thresh = is_480p_or_larger ? -2.6f : -2.3f;
840 break;
841 case BLOCK_16X16:
842 nn_config = &av1_early_term_after_split_nnconfig_16;
843 thresh = is_480p_or_larger ? -2.0f : -2.4f;
844 break;
845 case BLOCK_8X8:
846 nn_config = &av1_early_term_after_split_nnconfig_8;
847 thresh = is_480p_or_larger ? -1.0f : -1.4f;
848 break;
849 case BLOCK_4X4: break;
850 default:
851 assert(0 && "Invalid block size in av1_ml_early_term_after_split().");
852 break;
853 }
854 if (!nn_config) return;
855
856 // Use more conservative threshold for level 1.
857 if (cpi->sf.part_sf.ml_early_term_after_part_split_level < 2) thresh -= 0.3f;
858
859 const MACROBLOCKD *const xd = &x->e_mbd;
860 const int dc_q = av1_dc_quant_QTX(x->qindex, 0, xd->bd) >> (xd->bd - 8);
861 const int bs = block_size_wide[bsize];
862 int f_idx = 0;
863 float features[FEATURES] = { 0.0f };
864
865 aom_clear_system_state();
866
867 features[f_idx++] = logf(1.0f + (float)dc_q / 4.0f);
868 features[f_idx++] = logf(1.0f + (float)best_rd / bs / bs / 1024.0f);
869
870 add_rd_feature(part_none_rd, best_rd, features, &f_idx);
871 add_rd_feature(part_split_rd, best_rd, features, &f_idx);
872
873 for (int i = 0; i < 4; ++i) {
874 add_rd_feature(split_block_rd[i], best_rd, features, &f_idx);
875 int min_bw = MAX_SB_SIZE_LOG2;
876 int min_bh = MAX_SB_SIZE_LOG2;
877 get_min_bsize(pc_tree->split[i], &min_bw, &min_bh);
878 features[f_idx++] = (float)min_bw;
879 features[f_idx++] = (float)min_bh;
880 }
881
882 simple_motion_search_prune_part_features(cpi, x, pc_tree, mi_row, mi_col,
883 bsize, NULL,
884 FEATURE_SMS_PRUNE_PART_FLAG);
885
886 features[f_idx++] = logf(1.0f + (float)pc_tree->sms_none_feat[1]);
887
888 features[f_idx++] = logf(1.0f + (float)pc_tree->split[0]->sms_none_feat[1]);
889 features[f_idx++] = logf(1.0f + (float)pc_tree->split[1]->sms_none_feat[1]);
890 features[f_idx++] = logf(1.0f + (float)pc_tree->split[2]->sms_none_feat[1]);
891 features[f_idx++] = logf(1.0f + (float)pc_tree->split[3]->sms_none_feat[1]);
892
893 features[f_idx++] = logf(1.0f + (float)pc_tree->sms_rect_feat[1]);
894 features[f_idx++] = logf(1.0f + (float)pc_tree->sms_rect_feat[3]);
895 features[f_idx++] = logf(1.0f + (float)pc_tree->sms_rect_feat[5]);
896 features[f_idx++] = logf(1.0f + (float)pc_tree->sms_rect_feat[7]);
897
898 assert(f_idx == FEATURES);
899
900 float score = 0.0f;
901 av1_nn_predict(features, nn_config, 1, &score);
902 // Score is indicator of confidence that we should NOT terminate.
903 if (score < thresh) *terminate_partition_search = 1;
904 }
905 #undef FEATURES
906
av1_ml_prune_rect_partition(const AV1_COMP * const cpi,const MACROBLOCK * const x,BLOCK_SIZE bsize,int64_t best_rd,int64_t none_rd,int64_t * split_rd,int * const dst_prune_horz,int * const dst_prune_vert)907 void av1_ml_prune_rect_partition(const AV1_COMP *const cpi,
908 const MACROBLOCK *const x, BLOCK_SIZE bsize,
909 int64_t best_rd, int64_t none_rd,
910 int64_t *split_rd, int *const dst_prune_horz,
911 int *const dst_prune_vert) {
912 if (bsize < BLOCK_8X8 || best_rd >= 1000000000) return;
913 best_rd = AOMMAX(best_rd, 1);
914 const NN_CONFIG *nn_config = NULL;
915 const float prob_thresholds[5] = { 0.01f, 0.01f, 0.004f, 0.002f, 0.002f };
916 float cur_thresh = 0.0f;
917 switch (bsize) {
918 case BLOCK_8X8:
919 nn_config = &av1_rect_partition_nnconfig_8;
920 cur_thresh = prob_thresholds[0];
921 break;
922 case BLOCK_16X16:
923 nn_config = &av1_rect_partition_nnconfig_16;
924 cur_thresh = prob_thresholds[1];
925 break;
926 case BLOCK_32X32:
927 nn_config = &av1_rect_partition_nnconfig_32;
928 cur_thresh = prob_thresholds[2];
929 break;
930 case BLOCK_64X64:
931 nn_config = &av1_rect_partition_nnconfig_64;
932 cur_thresh = prob_thresholds[3];
933 break;
934 case BLOCK_128X128:
935 nn_config = &av1_rect_partition_nnconfig_128;
936 cur_thresh = prob_thresholds[4];
937 break;
938 default: assert(0 && "Unexpected bsize.");
939 }
940 if (!nn_config) return;
941 aom_clear_system_state();
942
943 // 1. Compute input features
944 float features[9];
945
946 // RD cost ratios
947 for (int i = 0; i < 5; i++) features[i] = 1.0f;
948 if (none_rd > 0 && none_rd < 1000000000)
949 features[0] = (float)none_rd / (float)best_rd;
950 for (int i = 0; i < 4; i++) {
951 if (split_rd[i] > 0 && split_rd[i] < 1000000000)
952 features[1 + i] = (float)split_rd[i] / (float)best_rd;
953 }
954
955 // Variance ratios
956 const MACROBLOCKD *const xd = &x->e_mbd;
957 int whole_block_variance;
958 if (is_cur_buf_hbd(xd)) {
959 whole_block_variance = av1_high_get_sby_perpixel_variance(
960 cpi, &x->plane[0].src, bsize, xd->bd);
961 } else {
962 whole_block_variance =
963 av1_get_sby_perpixel_variance(cpi, &x->plane[0].src, bsize);
964 }
965 whole_block_variance = AOMMAX(whole_block_variance, 1);
966
967 int split_variance[4];
968 const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
969 struct buf_2d buf;
970 buf.stride = x->plane[0].src.stride;
971 const int bw = block_size_wide[bsize];
972 for (int i = 0; i < 4; ++i) {
973 const int x_idx = (i & 1) * bw / 2;
974 const int y_idx = (i >> 1) * bw / 2;
975 buf.buf = x->plane[0].src.buf + x_idx + y_idx * buf.stride;
976 if (is_cur_buf_hbd(xd)) {
977 split_variance[i] =
978 av1_high_get_sby_perpixel_variance(cpi, &buf, subsize, xd->bd);
979 } else {
980 split_variance[i] = av1_get_sby_perpixel_variance(cpi, &buf, subsize);
981 }
982 }
983
984 for (int i = 0; i < 4; i++)
985 features[5 + i] = (float)split_variance[i] / (float)whole_block_variance;
986
987 // 2. Do the prediction and prune 0-2 partitions based on their probabilities
988 float raw_scores[3] = { 0.0f };
989 av1_nn_predict(features, nn_config, 1, raw_scores);
990 aom_clear_system_state();
991 float probs[3] = { 0.0f };
992 av1_nn_softmax(raw_scores, probs, 3);
993
994 // probs[0] is the probability of the fact that both rectangular partitions
995 // are worse than current best_rd
996 if (probs[1] <= cur_thresh) (*dst_prune_horz) = 1;
997 if (probs[2] <= cur_thresh) (*dst_prune_vert) = 1;
998 }
999
1000 // Use a ML model to predict if horz_a, horz_b, vert_a, and vert_b should be
1001 // considered.
av1_ml_prune_ab_partition(BLOCK_SIZE bsize,int part_ctx,int var_ctx,int64_t best_rd,int64_t horz_rd[2],int64_t vert_rd[2],int64_t split_rd[4],int * const horza_partition_allowed,int * const horzb_partition_allowed,int * const verta_partition_allowed,int * const vertb_partition_allowed)1002 void av1_ml_prune_ab_partition(BLOCK_SIZE bsize, int part_ctx, int var_ctx,
1003 int64_t best_rd, int64_t horz_rd[2],
1004 int64_t vert_rd[2], int64_t split_rd[4],
1005 int *const horza_partition_allowed,
1006 int *const horzb_partition_allowed,
1007 int *const verta_partition_allowed,
1008 int *const vertb_partition_allowed) {
1009 if (bsize < BLOCK_8X8 || best_rd >= 1000000000) return;
1010 const NN_CONFIG *nn_config = NULL;
1011 switch (bsize) {
1012 case BLOCK_8X8: nn_config = NULL; break;
1013 case BLOCK_16X16: nn_config = &av1_ab_partition_nnconfig_16; break;
1014 case BLOCK_32X32: nn_config = &av1_ab_partition_nnconfig_32; break;
1015 case BLOCK_64X64: nn_config = &av1_ab_partition_nnconfig_64; break;
1016 case BLOCK_128X128: nn_config = &av1_ab_partition_nnconfig_128; break;
1017 default: assert(0 && "Unexpected bsize.");
1018 }
1019 if (!nn_config) return;
1020
1021 aom_clear_system_state();
1022
1023 // Generate features.
1024 float features[10];
1025 int feature_index = 0;
1026 features[feature_index++] = (float)part_ctx;
1027 features[feature_index++] = (float)var_ctx;
1028 const int rdcost = (int)AOMMIN(INT_MAX, best_rd);
1029 int sub_block_rdcost[8] = { 0 };
1030 int rd_index = 0;
1031 for (int i = 0; i < 2; ++i) {
1032 if (horz_rd[i] > 0 && horz_rd[i] < 1000000000)
1033 sub_block_rdcost[rd_index] = (int)horz_rd[i];
1034 ++rd_index;
1035 }
1036 for (int i = 0; i < 2; ++i) {
1037 if (vert_rd[i] > 0 && vert_rd[i] < 1000000000)
1038 sub_block_rdcost[rd_index] = (int)vert_rd[i];
1039 ++rd_index;
1040 }
1041 for (int i = 0; i < 4; ++i) {
1042 if (split_rd[i] > 0 && split_rd[i] < 1000000000)
1043 sub_block_rdcost[rd_index] = (int)split_rd[i];
1044 ++rd_index;
1045 }
1046 for (int i = 0; i < 8; ++i) {
1047 // Ratio between the sub-block RD and the whole-block RD.
1048 float rd_ratio = 1.0f;
1049 if (sub_block_rdcost[i] > 0 && sub_block_rdcost[i] < rdcost)
1050 rd_ratio = (float)sub_block_rdcost[i] / (float)rdcost;
1051 features[feature_index++] = rd_ratio;
1052 }
1053 assert(feature_index == 10);
1054
1055 // Calculate scores using the NN model.
1056 float score[16] = { 0.0f };
1057 av1_nn_predict(features, nn_config, 1, score);
1058 aom_clear_system_state();
1059 int int_score[16];
1060 int max_score = -1000;
1061 for (int i = 0; i < 16; ++i) {
1062 int_score[i] = (int)(100 * score[i]);
1063 max_score = AOMMAX(int_score[i], max_score);
1064 }
1065
1066 // Make decisions based on the model scores.
1067 int thresh = max_score;
1068 switch (bsize) {
1069 case BLOCK_16X16: thresh -= 150; break;
1070 case BLOCK_32X32: thresh -= 100; break;
1071 default: break;
1072 }
1073 *horza_partition_allowed = 0;
1074 *horzb_partition_allowed = 0;
1075 *verta_partition_allowed = 0;
1076 *vertb_partition_allowed = 0;
1077 for (int i = 0; i < 16; ++i) {
1078 if (int_score[i] >= thresh) {
1079 if ((i >> 0) & 1) *horza_partition_allowed = 1;
1080 if ((i >> 1) & 1) *horzb_partition_allowed = 1;
1081 if ((i >> 2) & 1) *verta_partition_allowed = 1;
1082 if ((i >> 3) & 1) *vertb_partition_allowed = 1;
1083 }
1084 }
1085 }
1086
1087 #define FEATURES 18
1088 #define LABELS 4
1089 // Use a ML model to predict if horz4 and vert4 should be considered.
av1_ml_prune_4_partition(const AV1_COMP * const cpi,MACROBLOCK * const x,BLOCK_SIZE bsize,int part_ctx,int64_t best_rd,int64_t horz_rd[2],int64_t vert_rd[2],int64_t split_rd[4],int * const partition_horz4_allowed,int * const partition_vert4_allowed,unsigned int pb_source_variance,int mi_row,int mi_col)1090 void av1_ml_prune_4_partition(const AV1_COMP *const cpi, MACROBLOCK *const x,
1091 BLOCK_SIZE bsize, int part_ctx, int64_t best_rd,
1092 int64_t horz_rd[2], int64_t vert_rd[2],
1093 int64_t split_rd[4],
1094 int *const partition_horz4_allowed,
1095 int *const partition_vert4_allowed,
1096 unsigned int pb_source_variance, int mi_row,
1097 int mi_col) {
1098 if (best_rd >= 1000000000) return;
1099 const NN_CONFIG *nn_config = NULL;
1100 switch (bsize) {
1101 case BLOCK_16X16: nn_config = &av1_4_partition_nnconfig_16; break;
1102 case BLOCK_32X32: nn_config = &av1_4_partition_nnconfig_32; break;
1103 case BLOCK_64X64: nn_config = &av1_4_partition_nnconfig_64; break;
1104 default: assert(0 && "Unexpected bsize.");
1105 }
1106 if (!nn_config) return;
1107
1108 aom_clear_system_state();
1109
1110 // Generate features.
1111 float features[FEATURES];
1112 int feature_index = 0;
1113 features[feature_index++] = (float)part_ctx;
1114 features[feature_index++] = (float)get_unsigned_bits(pb_source_variance);
1115
1116 const int rdcost = (int)AOMMIN(INT_MAX, best_rd);
1117 int sub_block_rdcost[8] = { 0 };
1118 int rd_index = 0;
1119 for (int i = 0; i < 2; ++i) {
1120 if (horz_rd[i] > 0 && horz_rd[i] < 1000000000)
1121 sub_block_rdcost[rd_index] = (int)horz_rd[i];
1122 ++rd_index;
1123 }
1124 for (int i = 0; i < 2; ++i) {
1125 if (vert_rd[i] > 0 && vert_rd[i] < 1000000000)
1126 sub_block_rdcost[rd_index] = (int)vert_rd[i];
1127 ++rd_index;
1128 }
1129 for (int i = 0; i < 4; ++i) {
1130 if (split_rd[i] > 0 && split_rd[i] < 1000000000)
1131 sub_block_rdcost[rd_index] = (int)split_rd[i];
1132 ++rd_index;
1133 }
1134 for (int i = 0; i < 8; ++i) {
1135 // Ratio between the sub-block RD and the whole-block RD.
1136 float rd_ratio = 1.0f;
1137 if (sub_block_rdcost[i] > 0 && sub_block_rdcost[i] < rdcost)
1138 rd_ratio = (float)sub_block_rdcost[i] / (float)rdcost;
1139 features[feature_index++] = rd_ratio;
1140 }
1141
1142 // Get variance of the 1:4 and 4:1 sub-blocks.
1143 unsigned int horz_4_source_var[4] = { 0 };
1144 unsigned int vert_4_source_var[4] = { 0 };
1145 {
1146 BLOCK_SIZE horz_4_bs = get_partition_subsize(bsize, PARTITION_HORZ_4);
1147 BLOCK_SIZE vert_4_bs = get_partition_subsize(bsize, PARTITION_VERT_4);
1148 av1_setup_src_planes(x, cpi->source, mi_row, mi_col,
1149 av1_num_planes(&cpi->common), bsize);
1150 const int src_stride = x->plane[0].src.stride;
1151 uint8_t *src = x->plane[0].src.buf;
1152 const MACROBLOCKD *const xd = &x->e_mbd;
1153
1154 struct buf_2d horz_4_src, vert_4_src;
1155 horz_4_src.stride = src_stride;
1156 vert_4_src.stride = src_stride;
1157
1158 for (int i = 0; i < 4; ++i) {
1159 horz_4_src.buf = src + i * block_size_high[horz_4_bs] * src_stride;
1160 vert_4_src.buf = src + i * block_size_wide[vert_4_bs];
1161
1162 if (is_cur_buf_hbd(xd)) {
1163 horz_4_source_var[i] = av1_high_get_sby_perpixel_variance(
1164 cpi, &horz_4_src, horz_4_bs, xd->bd);
1165 vert_4_source_var[i] = av1_high_get_sby_perpixel_variance(
1166 cpi, &vert_4_src, vert_4_bs, xd->bd);
1167 } else {
1168 horz_4_source_var[i] =
1169 av1_get_sby_perpixel_variance(cpi, &horz_4_src, horz_4_bs);
1170 vert_4_source_var[i] =
1171 av1_get_sby_perpixel_variance(cpi, &vert_4_src, vert_4_bs);
1172 }
1173 }
1174 }
1175
1176 const float denom = (float)(pb_source_variance + 1);
1177 const float low_b = 0.1f;
1178 const float high_b = 10.0f;
1179 for (int i = 0; i < 4; ++i) {
1180 // Ratio between the 4:1 sub-block variance and the whole-block variance.
1181 float var_ratio = (float)(horz_4_source_var[i] + 1) / denom;
1182 if (var_ratio < low_b) var_ratio = low_b;
1183 if (var_ratio > high_b) var_ratio = high_b;
1184 features[feature_index++] = var_ratio;
1185 }
1186 for (int i = 0; i < 4; ++i) {
1187 // Ratio between the 1:4 sub-block RD and the whole-block RD.
1188 float var_ratio = (float)(vert_4_source_var[i] + 1) / denom;
1189 if (var_ratio < low_b) var_ratio = low_b;
1190 if (var_ratio > high_b) var_ratio = high_b;
1191 features[feature_index++] = var_ratio;
1192 }
1193 assert(feature_index == FEATURES);
1194
1195 // Calculate scores using the NN model.
1196 float score[LABELS] = { 0.0f };
1197 av1_nn_predict(features, nn_config, 1, score);
1198 aom_clear_system_state();
1199 int int_score[LABELS];
1200 int max_score = -1000;
1201 for (int i = 0; i < LABELS; ++i) {
1202 int_score[i] = (int)(100 * score[i]);
1203 max_score = AOMMAX(int_score[i], max_score);
1204 }
1205
1206 // Make decisions based on the model scores.
1207 int thresh = max_score;
1208 switch (bsize) {
1209 case BLOCK_16X16: thresh -= 500; break;
1210 case BLOCK_32X32: thresh -= 500; break;
1211 case BLOCK_64X64: thresh -= 200; break;
1212 default: break;
1213 }
1214 *partition_horz4_allowed = 0;
1215 *partition_vert4_allowed = 0;
1216 for (int i = 0; i < LABELS; ++i) {
1217 if (int_score[i] >= thresh) {
1218 if ((i >> 0) & 1) *partition_horz4_allowed = 1;
1219 if ((i >> 1) & 1) *partition_vert4_allowed = 1;
1220 }
1221 }
1222 }
1223 #undef FEATURES
1224 #undef LABELS
1225
1226 #define FEATURES 4
av1_ml_predict_breakout(const AV1_COMP * const cpi,BLOCK_SIZE bsize,const MACROBLOCK * const x,const RD_STATS * const rd_stats,unsigned int pb_source_variance)1227 int av1_ml_predict_breakout(const AV1_COMP *const cpi, BLOCK_SIZE bsize,
1228 const MACROBLOCK *const x,
1229 const RD_STATS *const rd_stats,
1230 unsigned int pb_source_variance) {
1231 const NN_CONFIG *nn_config = NULL;
1232 int thresh = 0;
1233 switch (bsize) {
1234 case BLOCK_8X8:
1235 nn_config = &av1_partition_breakout_nnconfig_8;
1236 thresh = cpi->sf.part_sf.ml_partition_search_breakout_thresh[0];
1237 break;
1238 case BLOCK_16X16:
1239 nn_config = &av1_partition_breakout_nnconfig_16;
1240 thresh = cpi->sf.part_sf.ml_partition_search_breakout_thresh[1];
1241 break;
1242 case BLOCK_32X32:
1243 nn_config = &av1_partition_breakout_nnconfig_32;
1244 thresh = cpi->sf.part_sf.ml_partition_search_breakout_thresh[2];
1245 break;
1246 case BLOCK_64X64:
1247 nn_config = &av1_partition_breakout_nnconfig_64;
1248 thresh = cpi->sf.part_sf.ml_partition_search_breakout_thresh[3];
1249 break;
1250 case BLOCK_128X128:
1251 nn_config = &av1_partition_breakout_nnconfig_128;
1252 thresh = cpi->sf.part_sf.ml_partition_search_breakout_thresh[4];
1253 break;
1254 default: assert(0 && "Unexpected bsize.");
1255 }
1256 if (!nn_config || thresh < 0) return 0;
1257
1258 // Generate feature values.
1259 float features[FEATURES];
1260 int feature_index = 0;
1261 aom_clear_system_state();
1262
1263 const int num_pels_log2 = num_pels_log2_lookup[bsize];
1264 float rate_f = (float)AOMMIN(rd_stats->rate, INT_MAX);
1265 rate_f = ((float)x->rdmult / 128.0f / 512.0f / (float)(1 << num_pels_log2)) *
1266 rate_f;
1267 features[feature_index++] = rate_f;
1268
1269 const float dist_f =
1270 (float)(AOMMIN(rd_stats->dist, INT_MAX) >> num_pels_log2);
1271 features[feature_index++] = dist_f;
1272
1273 features[feature_index++] = (float)pb_source_variance;
1274
1275 const int dc_q = (int)x->plane[0].dequant_QTX[0];
1276 features[feature_index++] = (float)(dc_q * dc_q) / 256.0f;
1277 assert(feature_index == FEATURES);
1278
1279 // Calculate score using the NN model.
1280 float score = 0.0f;
1281 av1_nn_predict(features, nn_config, 1, &score);
1282 aom_clear_system_state();
1283
1284 // Make decision.
1285 return (int)(score * 100) >= thresh;
1286 }
1287 #undef FEATURES
1288 #endif // !CONFIG_REALTIME_ONLY
1289