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
calculate_sse(MACROBLOCKD * const xd,const struct macroblock_plane * p,struct macroblockd_plane * pd,const int bw,const int bh)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
compute_sse_plane(MACROBLOCK * x,MACROBLOCKD * xd,int plane,const BLOCK_SIZE bsize)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
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)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)
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)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
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,uint8_t * skip_txfm_sb,int64_t * skip_sse_sb,int * plane_rate,int64_t * plane_sse,int64_t * plane_dist)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
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,uint8_t * skip_txfm_sb,int64_t * skip_sse_sb,int * plane_rate,int64_t * plane_sse,int64_t * plane_dist)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_
273