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