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)(const AV1_COMP *const cpi,
39 BLOCK_SIZE bsize, MACROBLOCK *x,
40 MACROBLOCKD *xd, int plane_from,
41 int plane_to, int *out_rate_sum,
42 int64_t *out_dist_sum, int *skip_txfm_sb,
43 int64_t *skip_sse_sb, int *plane_rate,
44 int64_t *plane_sse, int64_t *plane_dist);
45 typedef void (*model_rd_from_sse_type)(const AV1_COMP *const cpi,
46 const MACROBLOCK *const x,
47 BLOCK_SIZE plane_bsize, int plane,
48 int64_t sse, int num_samples, int *rate,
49 int64_t *dist);
50
calculate_sse(MACROBLOCKD * const xd,const struct macroblock_plane * p,struct macroblockd_plane * pd,const int bw,const int bh)51 static int64_t calculate_sse(MACROBLOCKD *const xd,
52 const struct macroblock_plane *p,
53 struct macroblockd_plane *pd, const int bw,
54 const int bh) {
55 int64_t sse = 0;
56 const int shift = xd->bd - 8;
57 #if CONFIG_AV1_HIGHBITDEPTH
58 if (is_cur_buf_hbd(xd)) {
59 sse = aom_highbd_sse(p->src.buf, p->src.stride, pd->dst.buf, pd->dst.stride,
60 bw, bh);
61 } else {
62 sse =
63 aom_sse(p->src.buf, p->src.stride, pd->dst.buf, pd->dst.stride, bw, bh);
64 }
65 #else
66 sse = aom_sse(p->src.buf, p->src.stride, pd->dst.buf, pd->dst.stride, bw, bh);
67 #endif
68 sse = ROUND_POWER_OF_TWO(sse, shift * 2);
69 return sse;
70 }
71
compute_sse_plane(MACROBLOCK * x,MACROBLOCKD * xd,int plane,const BLOCK_SIZE bsize)72 static AOM_INLINE int64_t compute_sse_plane(MACROBLOCK *x, MACROBLOCKD *xd,
73 int plane, const BLOCK_SIZE bsize) {
74 struct macroblockd_plane *const pd = &xd->plane[plane];
75 const BLOCK_SIZE plane_bsize =
76 get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);
77 int bw, bh;
78 const struct macroblock_plane *const p = &x->plane[plane];
79 get_txb_dimensions(xd, plane, plane_bsize, 0, 0, plane_bsize, NULL, NULL, &bw,
80 &bh);
81
82 int64_t sse = calculate_sse(xd, p, pd, bw, bh);
83
84 return sse;
85 }
86
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)87 static AOM_INLINE void model_rd_from_sse(const AV1_COMP *const cpi,
88 const MACROBLOCK *const x,
89 BLOCK_SIZE plane_bsize, int plane,
90 int64_t sse, int num_samples,
91 int *rate, int64_t *dist) {
92 (void)num_samples;
93 const MACROBLOCKD *const xd = &x->e_mbd;
94 const struct macroblock_plane *const p = &x->plane[plane];
95 const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3;
96
97 // Fast approximate the modelling function.
98 if (cpi->sf.rd_sf.simple_model_rd_from_var) {
99 const int64_t square_error = sse;
100 int quantizer = p->dequant_QTX[1] >> dequant_shift;
101 if (quantizer < 120)
102 *rate = (int)AOMMIN(
103 (square_error * (280 - quantizer)) >> (16 - AV1_PROB_COST_SHIFT),
104 INT_MAX);
105 else
106 *rate = 0;
107 assert(*rate >= 0);
108 *dist = (square_error * quantizer) >> 8;
109 } else {
110 av1_model_rd_from_var_lapndz(sse, num_pels_log2_lookup[plane_bsize],
111 p->dequant_QTX[1] >> dequant_shift, rate,
112 dist);
113 }
114 *dist <<= 4;
115 }
116
117 // Fits a curve for rate and distortion using as feature:
118 // 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)119 static AOM_INLINE void model_rd_with_curvfit(const AV1_COMP *const cpi,
120 const MACROBLOCK *const x,
121 BLOCK_SIZE plane_bsize, int plane,
122 int64_t sse, int num_samples,
123 int *rate, int64_t *dist) {
124 (void)cpi;
125 (void)plane_bsize;
126 const MACROBLOCKD *const xd = &x->e_mbd;
127 const struct macroblock_plane *const p = &x->plane[plane];
128 const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3;
129 const int qstep = AOMMAX(p->dequant_QTX[1] >> dequant_shift, 1);
130
131 if (sse == 0) {
132 if (rate) *rate = 0;
133 if (dist) *dist = 0;
134 return;
135 }
136 const double sse_norm = (double)sse / num_samples;
137 const double qstepsqr = (double)qstep * qstep;
138 const double xqr = log2(sse_norm / qstepsqr);
139 double rate_f, dist_by_sse_norm_f;
140 av1_model_rd_curvfit(plane_bsize, sse_norm, xqr, &rate_f,
141 &dist_by_sse_norm_f);
142
143 const double dist_f = dist_by_sse_norm_f * sse_norm;
144 int rate_i = (int)(AOMMAX(0.0, rate_f * num_samples) + 0.5);
145 int64_t dist_i = (int64_t)(AOMMAX(0.0, dist_f * num_samples) + 0.5);
146
147 // Check if skip is better
148 if (rate_i == 0) {
149 dist_i = sse << 4;
150 } else if (RDCOST(x->rdmult, rate_i, dist_i) >=
151 RDCOST(x->rdmult, 0, sse << 4)) {
152 rate_i = 0;
153 dist_i = sse << 4;
154 }
155
156 if (rate) *rate = rate_i;
157 if (dist) *dist = dist_i;
158 }
159
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)160 static AOM_INLINE void model_rd_for_sb(
161 const AV1_COMP *const cpi, BLOCK_SIZE bsize, MACROBLOCK *x, MACROBLOCKD *xd,
162 int plane_from, int plane_to, int *out_rate_sum, int64_t *out_dist_sum,
163 int *skip_txfm_sb, int64_t *skip_sse_sb, int *plane_rate,
164 int64_t *plane_sse, 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
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)212 static AOM_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 int *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_
273