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 AOM_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 AOM_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,
89 int *rate, 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 AOM_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 AOM_INLINE void model_rd_for_sb(
159 const AV1_COMP *const cpi, BLOCK_SIZE bsize, MACROBLOCK *x, MACROBLOCKD *xd,
160 int plane_from, int plane_to, int *out_rate_sum, int64_t *out_dist_sum,
161 uint8_t *skip_txfm_sb, int64_t *skip_sse_sb, int *plane_rate,
162 int64_t *plane_sse, int64_t *plane_dist) {
163 // Note our transform coeffs are 8 times an orthogonal transform.
164 // Hence quantizer step is also 8 times. To get effective quantizer
165 // we need to divide by 8 before sending to modeling function.
166 int plane;
167 const int ref = xd->mi[0]->ref_frame[0];
168
169 int64_t rate_sum = 0;
170 int64_t dist_sum = 0;
171 int64_t total_sse = 0;
172
173 assert(bsize < BLOCK_SIZES_ALL);
174
175 for (plane = plane_from; plane <= plane_to; ++plane) {
176 if (plane && !xd->is_chroma_ref) break;
177 struct macroblock_plane *const p = &x->plane[plane];
178 struct macroblockd_plane *const pd = &xd->plane[plane];
179 const BLOCK_SIZE plane_bsize =
180 get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);
181 assert(plane_bsize < BLOCK_SIZES_ALL);
182 const int bw = block_size_wide[plane_bsize];
183 const int bh = block_size_high[plane_bsize];
184 int64_t sse;
185 int rate;
186 int64_t dist;
187
188 sse = calculate_sse(xd, p, pd, bw, bh);
189
190 model_rd_from_sse(cpi, x, plane_bsize, plane, sse, bw * bh, &rate, &dist);
191
192 if (plane == 0) x->pred_sse[ref] = (unsigned int)AOMMIN(sse, UINT_MAX);
193
194 total_sse += sse;
195 rate_sum += rate;
196 dist_sum += dist;
197 if (plane_rate) plane_rate[plane] = rate;
198 if (plane_sse) plane_sse[plane] = sse;
199 if (plane_dist) plane_dist[plane] = dist;
200 assert(rate_sum >= 0);
201 }
202
203 if (skip_txfm_sb) *skip_txfm_sb = total_sse == 0;
204 if (skip_sse_sb) *skip_sse_sb = total_sse << 4;
205 rate_sum = AOMMIN(rate_sum, INT_MAX);
206 *out_rate_sum = (int)rate_sum;
207 *out_dist_sum = dist_sum;
208 }
209
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)210 static AOM_INLINE void model_rd_for_sb_with_curvfit(
211 const AV1_COMP *const cpi, BLOCK_SIZE bsize, MACROBLOCK *x, MACROBLOCKD *xd,
212 int plane_from, int plane_to, int *out_rate_sum, int64_t *out_dist_sum,
213 uint8_t *skip_txfm_sb, int64_t *skip_sse_sb, int *plane_rate,
214 int64_t *plane_sse, int64_t *plane_dist) {
215 // Note our transform coeffs are 8 times an orthogonal transform.
216 // Hence quantizer step is also 8 times. To get effective quantizer
217 // we need to divide by 8 before sending to modeling function.
218 const int ref = xd->mi[0]->ref_frame[0];
219
220 int64_t rate_sum = 0;
221 int64_t dist_sum = 0;
222 int64_t total_sse = 0;
223
224 for (int plane = plane_from; plane <= plane_to; ++plane) {
225 if (plane && !xd->is_chroma_ref) break;
226 struct macroblockd_plane *const pd = &xd->plane[plane];
227 const BLOCK_SIZE plane_bsize =
228 get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);
229 int64_t dist, sse;
230 int rate;
231 int bw, bh;
232 const struct macroblock_plane *const p = &x->plane[plane];
233 get_txb_dimensions(xd, plane, plane_bsize, 0, 0, plane_bsize, NULL, NULL,
234 &bw, &bh);
235
236 sse = calculate_sse(xd, p, pd, bw, bh);
237 model_rd_with_curvfit(cpi, x, plane_bsize, plane, sse, bw * bh, &rate,
238 &dist);
239
240 if (plane == 0) x->pred_sse[ref] = (unsigned int)AOMMIN(sse, UINT_MAX);
241
242 total_sse += sse;
243 rate_sum += rate;
244 dist_sum += dist;
245
246 if (plane_rate) plane_rate[plane] = rate;
247 if (plane_sse) plane_sse[plane] = sse;
248 if (plane_dist) plane_dist[plane] = dist;
249 }
250
251 if (skip_txfm_sb) *skip_txfm_sb = rate_sum == 0;
252 if (skip_sse_sb) *skip_sse_sb = total_sse << 4;
253 *out_rate_sum = (int)rate_sum;
254 *out_dist_sum = dist_sum;
255 }
256
257 enum { MODELRD_LEGACY, MODELRD_CURVFIT, MODELRD_TYPES } UENUM1BYTE(ModelRdType);
258
259 static const model_rd_for_sb_type model_rd_sb_fn[MODELRD_TYPES] = {
260 model_rd_for_sb, model_rd_for_sb_with_curvfit
261 };
262
263 static const model_rd_from_sse_type model_rd_sse_fn[MODELRD_TYPES] = {
264 model_rd_from_sse, model_rd_with_curvfit
265 };
266
267 #ifdef __cplusplus
268 } // extern "C"
269 #endif
270 #endif // AOM_AV1_ENCODER_MODEL_RD_H_
271