• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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