model_rd.h (10004B)
1 /* 2 * Copyright (c) 2020, 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 #ifndef AOM_AV1_ENCODER_MODEL_RD_H_ 13 #define AOM_AV1_ENCODER_MODEL_RD_H_ 14 15 #include "aom/aom_integer.h" 16 #include "av1/encoder/block.h" 17 #include "av1/encoder/encoder.h" 18 #include "av1/encoder/pustats.h" 19 #include "av1/encoder/rdopt_utils.h" 20 #include "config/aom_dsp_rtcd.h" 21 22 #ifdef __cplusplus 23 extern "C" { 24 #endif 25 26 // 0: Legacy model 27 // 1: Curve fit model 28 // 2: Surface fit model 29 // 3: DNN regression model 30 // 4: Full rd model 31 #define MODELRD_TYPE_INTERP_FILTER 1 32 #define MODELRD_TYPE_TX_SEARCH_PRUNE 1 33 #define MODELRD_TYPE_MASKED_COMPOUND 1 34 #define MODELRD_TYPE_INTERINTRA 1 35 #define MODELRD_TYPE_INTRA 1 36 #define MODELRD_TYPE_MOTION_MODE_RD 1 37 38 typedef void (*model_rd_for_sb_type)( 39 const AV1_COMP *const cpi, BLOCK_SIZE bsize, MACROBLOCK *x, MACROBLOCKD *xd, 40 int plane_from, int plane_to, int *out_rate_sum, int64_t *out_dist_sum, 41 uint8_t *skip_txfm_sb, int64_t *skip_sse_sb, int *plane_rate, 42 int64_t *plane_sse, int64_t *plane_dist); 43 typedef void (*model_rd_from_sse_type)(const AV1_COMP *const cpi, 44 const MACROBLOCK *const x, 45 BLOCK_SIZE plane_bsize, int plane, 46 int64_t sse, int num_samples, int *rate, 47 int64_t *dist); 48 49 static int64_t calculate_sse(MACROBLOCKD *const xd, 50 const struct macroblock_plane *p, 51 struct macroblockd_plane *pd, const int bw, 52 const int bh) { 53 int64_t sse = 0; 54 const int shift = xd->bd - 8; 55 #if CONFIG_AV1_HIGHBITDEPTH 56 if (is_cur_buf_hbd(xd)) { 57 sse = aom_highbd_sse(p->src.buf, p->src.stride, pd->dst.buf, pd->dst.stride, 58 bw, bh); 59 } else { 60 sse = 61 aom_sse(p->src.buf, p->src.stride, pd->dst.buf, pd->dst.stride, bw, bh); 62 } 63 #else 64 sse = aom_sse(p->src.buf, p->src.stride, pd->dst.buf, pd->dst.stride, bw, bh); 65 #endif 66 sse = ROUND_POWER_OF_TWO(sse, shift * 2); 67 return sse; 68 } 69 70 static inline int64_t compute_sse_plane(MACROBLOCK *x, MACROBLOCKD *xd, 71 int plane, const BLOCK_SIZE bsize) { 72 struct macroblockd_plane *const pd = &xd->plane[plane]; 73 const BLOCK_SIZE plane_bsize = 74 get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y); 75 int bw, bh; 76 const struct macroblock_plane *const p = &x->plane[plane]; 77 get_txb_dimensions(xd, plane, plane_bsize, 0, 0, plane_bsize, NULL, NULL, &bw, 78 &bh); 79 80 int64_t sse = calculate_sse(xd, p, pd, bw, bh); 81 82 return sse; 83 } 84 85 static inline void model_rd_from_sse(const AV1_COMP *const cpi, 86 const MACROBLOCK *const x, 87 BLOCK_SIZE plane_bsize, int plane, 88 int64_t sse, int num_samples, int *rate, 89 int64_t *dist) { 90 (void)num_samples; 91 const MACROBLOCKD *const xd = &x->e_mbd; 92 const struct macroblock_plane *const p = &x->plane[plane]; 93 const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3; 94 95 // Fast approximate the modelling function. 96 if (cpi->sf.rd_sf.simple_model_rd_from_var) { 97 const int64_t square_error = sse; 98 int quantizer = p->dequant_QTX[1] >> dequant_shift; 99 if (quantizer < 120) 100 *rate = (int)AOMMIN( 101 (square_error * (280 - quantizer)) >> (16 - AV1_PROB_COST_SHIFT), 102 INT_MAX); 103 else 104 *rate = 0; 105 assert(*rate >= 0); 106 *dist = (square_error * quantizer) >> 8; 107 } else { 108 av1_model_rd_from_var_lapndz(sse, num_pels_log2_lookup[plane_bsize], 109 p->dequant_QTX[1] >> dequant_shift, rate, 110 dist); 111 } 112 *dist <<= 4; 113 } 114 115 // Fits a curve for rate and distortion using as feature: 116 // log2(sse_norm/qstep^2) 117 static inline void model_rd_with_curvfit(const AV1_COMP *const cpi, 118 const MACROBLOCK *const x, 119 BLOCK_SIZE plane_bsize, int plane, 120 int64_t sse, int num_samples, 121 int *rate, int64_t *dist) { 122 (void)cpi; 123 (void)plane_bsize; 124 const MACROBLOCKD *const xd = &x->e_mbd; 125 const struct macroblock_plane *const p = &x->plane[plane]; 126 const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3; 127 const int qstep = AOMMAX(p->dequant_QTX[1] >> dequant_shift, 1); 128 129 if (sse == 0) { 130 if (rate) *rate = 0; 131 if (dist) *dist = 0; 132 return; 133 } 134 const double sse_norm = (double)sse / num_samples; 135 const double qstepsqr = (double)qstep * qstep; 136 const double xqr = log2(sse_norm / qstepsqr); 137 double rate_f, dist_by_sse_norm_f; 138 av1_model_rd_curvfit(plane_bsize, sse_norm, xqr, &rate_f, 139 &dist_by_sse_norm_f); 140 141 const double dist_f = dist_by_sse_norm_f * sse_norm; 142 int rate_i = (int)(AOMMAX(0.0, rate_f * num_samples) + 0.5); 143 int64_t dist_i = (int64_t)(AOMMAX(0.0, dist_f * num_samples) + 0.5); 144 145 // Check if skip is better 146 if (rate_i == 0) { 147 dist_i = sse << 4; 148 } else if (RDCOST(x->rdmult, rate_i, dist_i) >= 149 RDCOST(x->rdmult, 0, sse << 4)) { 150 rate_i = 0; 151 dist_i = sse << 4; 152 } 153 154 if (rate) *rate = rate_i; 155 if (dist) *dist = dist_i; 156 } 157 158 static inline void model_rd_for_sb(const AV1_COMP *const cpi, BLOCK_SIZE bsize, 159 MACROBLOCK *x, MACROBLOCKD *xd, 160 int plane_from, int plane_to, 161 int *out_rate_sum, int64_t *out_dist_sum, 162 uint8_t *skip_txfm_sb, int64_t *skip_sse_sb, 163 int *plane_rate, int64_t *plane_sse, 164 int64_t *plane_dist) { 165 // Note our transform coeffs are 8 times an orthogonal transform. 166 // Hence quantizer step is also 8 times. To get effective quantizer 167 // we need to divide by 8 before sending to modeling function. 168 int plane; 169 const int ref = xd->mi[0]->ref_frame[0]; 170 171 int64_t rate_sum = 0; 172 int64_t dist_sum = 0; 173 int64_t total_sse = 0; 174 175 assert(bsize < BLOCK_SIZES_ALL); 176 177 for (plane = plane_from; plane <= plane_to; ++plane) { 178 if (plane && !xd->is_chroma_ref) break; 179 struct macroblock_plane *const p = &x->plane[plane]; 180 struct macroblockd_plane *const pd = &xd->plane[plane]; 181 const BLOCK_SIZE plane_bsize = 182 get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y); 183 assert(plane_bsize < BLOCK_SIZES_ALL); 184 const int bw = block_size_wide[plane_bsize]; 185 const int bh = block_size_high[plane_bsize]; 186 int64_t sse; 187 int rate; 188 int64_t dist; 189 190 sse = calculate_sse(xd, p, pd, bw, bh); 191 192 model_rd_from_sse(cpi, x, plane_bsize, plane, sse, bw * bh, &rate, &dist); 193 194 if (plane == 0) x->pred_sse[ref] = (unsigned int)AOMMIN(sse, UINT_MAX); 195 196 total_sse += sse; 197 rate_sum += rate; 198 dist_sum += dist; 199 if (plane_rate) plane_rate[plane] = rate; 200 if (plane_sse) plane_sse[plane] = sse; 201 if (plane_dist) plane_dist[plane] = dist; 202 assert(rate_sum >= 0); 203 } 204 205 if (skip_txfm_sb) *skip_txfm_sb = total_sse == 0; 206 if (skip_sse_sb) *skip_sse_sb = total_sse << 4; 207 rate_sum = AOMMIN(rate_sum, INT_MAX); 208 *out_rate_sum = (int)rate_sum; 209 *out_dist_sum = dist_sum; 210 } 211 212 static inline void model_rd_for_sb_with_curvfit( 213 const AV1_COMP *const cpi, BLOCK_SIZE bsize, MACROBLOCK *x, MACROBLOCKD *xd, 214 int plane_from, int plane_to, int *out_rate_sum, int64_t *out_dist_sum, 215 uint8_t *skip_txfm_sb, int64_t *skip_sse_sb, int *plane_rate, 216 int64_t *plane_sse, int64_t *plane_dist) { 217 // Note our transform coeffs are 8 times an orthogonal transform. 218 // Hence quantizer step is also 8 times. To get effective quantizer 219 // we need to divide by 8 before sending to modeling function. 220 const int ref = xd->mi[0]->ref_frame[0]; 221 222 int64_t rate_sum = 0; 223 int64_t dist_sum = 0; 224 int64_t total_sse = 0; 225 226 for (int plane = plane_from; plane <= plane_to; ++plane) { 227 if (plane && !xd->is_chroma_ref) break; 228 struct macroblockd_plane *const pd = &xd->plane[plane]; 229 const BLOCK_SIZE plane_bsize = 230 get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y); 231 int64_t dist, sse; 232 int rate; 233 int bw, bh; 234 const struct macroblock_plane *const p = &x->plane[plane]; 235 get_txb_dimensions(xd, plane, plane_bsize, 0, 0, plane_bsize, NULL, NULL, 236 &bw, &bh); 237 238 sse = calculate_sse(xd, p, pd, bw, bh); 239 model_rd_with_curvfit(cpi, x, plane_bsize, plane, sse, bw * bh, &rate, 240 &dist); 241 242 if (plane == 0) x->pred_sse[ref] = (unsigned int)AOMMIN(sse, UINT_MAX); 243 244 total_sse += sse; 245 rate_sum += rate; 246 dist_sum += dist; 247 248 if (plane_rate) plane_rate[plane] = rate; 249 if (plane_sse) plane_sse[plane] = sse; 250 if (plane_dist) plane_dist[plane] = dist; 251 } 252 253 if (skip_txfm_sb) *skip_txfm_sb = rate_sum == 0; 254 if (skip_sse_sb) *skip_sse_sb = total_sse << 4; 255 *out_rate_sum = (int)rate_sum; 256 *out_dist_sum = dist_sum; 257 } 258 259 enum { MODELRD_LEGACY, MODELRD_CURVFIT, MODELRD_TYPES } UENUM1BYTE(ModelRdType); 260 261 static const model_rd_for_sb_type model_rd_sb_fn[MODELRD_TYPES] = { 262 model_rd_for_sb, model_rd_for_sb_with_curvfit 263 }; 264 265 static const model_rd_from_sse_type model_rd_sse_fn[MODELRD_TYPES] = { 266 model_rd_from_sse, model_rd_with_curvfit 267 }; 268 269 #ifdef __cplusplus 270 } // extern "C" 271 #endif 272 #endif // AOM_AV1_ENCODER_MODEL_RD_H_