bd481cf4b16db7f2da0eb8d4ab13dac2a2c54f87
[x262.git] / encoder / rdo.c
1 /*****************************************************************************
2  * rdo.c: h264 encoder library (rate-distortion optimization)
3  *****************************************************************************
4  * Copyright (C) 2005 x264 project
5  *
6  * Authors: Loren Merritt <lorenm@u.washington.edu>
7  *
8  * This program is free software; you can redistribute it and/or modify
9  * it under the terms of the GNU General Public License as published by
10  * the Free Software Foundation; either version 2 of the License, or
11  * (at your option) any later version.
12  *
13  * This program is distributed in the hope that it will be useful,
14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
16  * GNU General Public License for more details.
17  *
18  * You should have received a copy of the GNU General Public License
19  * along with this program; if not, write to the Free Software
20  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111, USA.
21  *****************************************************************************/
22
23 /* duplicate all the writer functions, just calculating bit cost
24  * instead of writing the bitstream.
25  * TODO: use these for fast 1st pass too. */
26
27 #define RDO_SKIP_BS
28
29 static uint8_t cabac_prefix_transition[15][128];
30 static uint16_t cabac_prefix_size[15][128];
31
32 /* CAVLC: produces exactly the same bit count as a normal encode */
33 /* this probably still leaves some unnecessary computations */
34 #define bs_write1(s,v)     ((s)->i_bits_encoded += 1)
35 #define bs_write(s,n,v)    ((s)->i_bits_encoded += (n))
36 #define bs_write_ue(s,v)   ((s)->i_bits_encoded += bs_size_ue(v))
37 #define bs_write_se(s,v)   ((s)->i_bits_encoded += bs_size_se(v))
38 #define bs_write_te(s,v,l) ((s)->i_bits_encoded += bs_size_te(v,l))
39 #define x264_macroblock_write_cavlc  x264_macroblock_size_cavlc
40 #include "cavlc.c"
41
42 /* CABAC: not exactly the same. x264_cabac_size_decision() keeps track of
43  * fractional bits, but only finite precision. */
44 #define x264_cabac_encode_decision(c,x,v) x264_cabac_size_decision(c,x,v)
45 #define x264_cabac_encode_terminal(c)     x264_cabac_size_decision(c,276,0)
46 #define x264_cabac_encode_bypass(c,v)     ((c)->f8_bits_encoded += 256)
47 #define x264_cabac_encode_ue_bypass(c,e,v) ((c)->f8_bits_encoded += (bs_size_ue(v+(1<<e)-1)-e)<<8)
48 #define x264_cabac_encode_flush(h,c)
49 #define x264_macroblock_write_cabac  x264_macroblock_size_cabac
50 #include "cabac.c"
51
52 #define COPY_CABAC h->mc.memcpy_aligned( &cabac_tmp.f8_bits_encoded, &h->cabac.f8_bits_encoded, \
53         sizeof(x264_cabac_t) - offsetof(x264_cabac_t,f8_bits_encoded) )
54     
55 static int ssd_mb( x264_t *h )
56 {
57     return h->pixf.ssd[PIXEL_16x16]( h->mb.pic.p_fenc[0], FENC_STRIDE,
58                                      h->mb.pic.p_fdec[0], FDEC_STRIDE )
59          + h->pixf.ssd[PIXEL_8x8](   h->mb.pic.p_fenc[1], FENC_STRIDE,
60                                      h->mb.pic.p_fdec[1], FDEC_STRIDE )
61          + h->pixf.ssd[PIXEL_8x8](   h->mb.pic.p_fenc[2], FENC_STRIDE,
62                                      h->mb.pic.p_fdec[2], FDEC_STRIDE );
63 }
64
65 static int ssd_plane( x264_t *h, int size, int p, int x, int y )
66 {
67     return h->pixf.ssd[size]( h->mb.pic.p_fenc[p] + x+y*FENC_STRIDE, FENC_STRIDE,
68                               h->mb.pic.p_fdec[p] + x+y*FDEC_STRIDE, FDEC_STRIDE );
69 }
70
71 static int x264_rd_cost_mb( x264_t *h, int i_lambda2 )
72 {
73     int b_transform_bak = h->mb.b_transform_8x8;
74     int i_ssd;
75     int i_bits;
76
77     x264_macroblock_encode( h );
78
79     i_ssd = ssd_mb( h );
80
81     if( IS_SKIP( h->mb.i_type ) )
82     {
83         i_bits = (1 * i_lambda2 + 128) >> 8;
84     }
85     else if( h->param.b_cabac )
86     {
87         x264_cabac_t cabac_tmp;
88         COPY_CABAC;
89         x264_macroblock_size_cabac( h, &cabac_tmp );
90         i_bits = ( (uint64_t)cabac_tmp.f8_bits_encoded * i_lambda2 + 32768 ) >> 16;
91     }
92     else
93     {
94         bs_t bs_tmp = h->out.bs;
95         bs_tmp.i_bits_encoded = 0;
96         x264_macroblock_size_cavlc( h, &bs_tmp );
97         i_bits = ( bs_tmp.i_bits_encoded * i_lambda2 + 128 ) >> 8;
98     }
99
100     h->mb.b_transform_8x8 = b_transform_bak;
101
102     return i_ssd + i_bits;
103 }
104
105 int x264_rd_cost_part( x264_t *h, int i_lambda2, int i8, int i_pixel )
106 {
107     int i_ssd, i_bits;
108
109     if( i_pixel == PIXEL_16x16 )
110     {
111         int type_bak = h->mb.i_type;
112         int i_cost = x264_rd_cost_mb( h, i_lambda2 );
113         h->mb.i_type = type_bak;
114         return i_cost;
115     }
116
117     x264_macroblock_encode_p8x8( h, i8 );
118     if( i_pixel == PIXEL_16x8 )
119         x264_macroblock_encode_p8x8( h, i8+1 );
120     if( i_pixel == PIXEL_8x16 )
121         x264_macroblock_encode_p8x8( h, i8+2 );
122
123     i_ssd = ssd_plane( h, i_pixel,   0, (i8&1)*8, (i8>>1)*8 )
124           + ssd_plane( h, i_pixel+3, 1, (i8&1)*4, (i8>>1)*4 )
125           + ssd_plane( h, i_pixel+3, 2, (i8&1)*4, (i8>>1)*4 );
126
127     if( h->param.b_cabac )
128     {
129         x264_cabac_t cabac_tmp;
130         COPY_CABAC;
131         x264_partition_size_cabac( h, &cabac_tmp, i8, i_pixel );
132         i_bits = ( (uint64_t)cabac_tmp.f8_bits_encoded * i_lambda2 + 32768 ) >> 16;
133     }
134     else
135     {
136         i_bits = ( x264_partition_size_cavlc( h, i8, i_pixel ) * i_lambda2 + 128 ) >> 8;
137     }
138
139     return i_ssd + i_bits;
140 }
141
142 int x264_rd_cost_i8x8( x264_t *h, int i_lambda2, int i8, int i_mode )
143 {
144     int i_ssd, i_bits;
145
146     x264_mb_encode_i8x8( h, i8, h->mb.i_qp );
147     i_ssd = ssd_plane( h, PIXEL_8x8, 0, (i8&1)*8, (i8>>1)*8 );
148
149     if( h->param.b_cabac )
150     {
151         x264_cabac_t cabac_tmp;
152         COPY_CABAC;
153         x264_partition_i8x8_size_cabac( h, &cabac_tmp, i8, i_mode );
154         i_bits = ( (uint64_t)cabac_tmp.f8_bits_encoded * i_lambda2 + 32768 ) >> 16;
155     }
156     else
157     {
158         i_bits = ( x264_partition_i8x8_size_cavlc( h, i8, i_mode ) * i_lambda2 + 128 ) >> 8;
159     }
160
161     return i_ssd + i_bits;
162 }
163
164 int x264_rd_cost_i4x4( x264_t *h, int i_lambda2, int i4, int i_mode )
165 {
166     int i_ssd, i_bits;
167
168     x264_mb_encode_i4x4( h, i4, h->mb.i_qp );
169     i_ssd = ssd_plane( h, PIXEL_4x4, 0, block_idx_x[i4]*4, block_idx_y[i4]*4 );
170
171     if( h->param.b_cabac )
172     {
173         x264_cabac_t cabac_tmp;
174         COPY_CABAC;
175         x264_partition_i4x4_size_cabac( h, &cabac_tmp, i4, i_mode );
176         i_bits = ( (uint64_t)cabac_tmp.f8_bits_encoded * i_lambda2 + 32768 ) >> 16;
177     }
178     else
179     {
180         i_bits = ( x264_partition_i4x4_size_cavlc( h, i4, i_mode ) * i_lambda2 + 128 ) >> 8;
181     }
182
183     return i_ssd + i_bits;
184 }
185
186 int x264_rd_cost_i8x8_chroma( x264_t *h, int i_lambda2, int i_mode, int b_dct )
187 {
188     int i_ssd, i_bits;
189
190     if( b_dct )
191         x264_mb_encode_8x8_chroma( h, 0, h->mb.i_chroma_qp );
192     i_ssd = ssd_plane( h, PIXEL_8x8, 1, 0, 0 ) +
193             ssd_plane( h, PIXEL_8x8, 2, 0, 0 );
194
195     h->mb.i_chroma_pred_mode = i_mode;
196
197     if( h->param.b_cabac )
198     {
199         x264_cabac_t cabac_tmp;
200         COPY_CABAC;
201         x264_i8x8_chroma_size_cabac( h, &cabac_tmp );
202         i_bits = ( (uint64_t)cabac_tmp.f8_bits_encoded * i_lambda2 + 32768 ) >> 16;
203     }
204     else
205     {
206         i_bits = ( x264_i8x8_chroma_size_cavlc( h ) * i_lambda2 + 128 ) >> 8;
207     }
208
209     return i_ssd + i_bits;
210 }
211 /****************************************************************************
212  * Trellis RD quantization
213  ****************************************************************************/
214
215 #define TRELLIS_SCORE_MAX ((uint64_t)1<<50)
216 #define CABAC_SIZE_BITS 8
217 #define SSD_WEIGHT_BITS 5
218 #define LAMBDA_BITS 4
219
220 /* precalculate the cost of coding abs_level_m1 */
221 void x264_rdo_init( )
222 {
223     int i_prefix;
224     int i_ctx;
225     for( i_prefix = 0; i_prefix < 15; i_prefix++ )
226     {
227         for( i_ctx = 0; i_ctx < 128; i_ctx++ )
228         {
229             int f8_bits = 0;
230             uint8_t ctx = i_ctx;
231             int i;
232
233             for( i = 1; i < i_prefix; i++ )
234                 f8_bits += x264_cabac_size_decision2( &ctx, 1 );
235             if( i_prefix > 0 && i_prefix < 14 )
236                 f8_bits += x264_cabac_size_decision2( &ctx, 0 );
237             f8_bits += 1 << CABAC_SIZE_BITS; //sign
238
239             cabac_prefix_size[i_prefix][i_ctx] = f8_bits;
240             cabac_prefix_transition[i_prefix][i_ctx] = ctx;
241         }
242     }
243 }
244
245 // should the intra and inter lambdas be different?
246 // I'm just matching the behaviour of deadzone quant.
247 static const int lambda2_tab[2][52] = {
248     // inter lambda = .85 * .85 * 2**(qp/3. + 10 - LAMBDA_BITS)
249     {    46,      58,      73,      92,     117,     147, 
250         185,     233,     294,     370,     466,     587, 
251         740,     932,    1174,    1480,    1864,    2349, 
252        2959,    3728,    4697,    5918,    7457,    9395, 
253       11837,   14914,   18790,   23674,   29828,   37581, 
254       47349,   59656,   75163,   94699,  119313,  150326, 
255      189399,  238627,  300652,  378798,  477255,  601304, 
256      757596,  954511, 1202608, 1515192, 1909022, 2405217, 
257     3030384, 3818045, 4810435, 6060769 },
258     // intra lambda = .65 * .65 * 2**(qp/3. + 10 - LAMBDA_BITS)
259     {    27,      34,      43,      54,      68,      86, 
260         108,     136,     172,     216,     273,     343, 
261         433,     545,     687,     865,    1090,    1374, 
262        1731,    2180,    2747,    3461,    4361,    5494, 
263        6922,    8721,   10988,   13844,   17442,   21976, 
264       27688,   34885,   43953,   55377,   69771,   87906, 
265      110755,  139543,  175813,  221511,  279087,  351627, 
266      443023,  558174,  703255,  886046, 1116348, 1406511, 
267     1772093, 2232697, 2813022, 3544186 }
268 };
269
270 typedef struct {
271     uint64_t score;
272     int level_idx; // index into level_tree[]
273     uint8_t cabac_state[10]; //just the contexts relevant to coding abs_level_m1
274 } trellis_node_t;
275
276 // TODO:
277 // support chroma and i16x16 DC
278 // save cabac state between blocks?
279 // use trellis' RD score instead of x264_mb_decimate_score?
280 // code 8x8 sig/last flags forwards with deadzone and save the contexts at
281 //   each position?
282 // change weights when using CQMs?
283
284 // possible optimizations:
285 // make scores fit in 32bit
286 // save quantized coefs during rd, to avoid a duplicate trellis in the final encode
287 // if trellissing all MBRD modes, finish SSD calculation so we can skip all of
288 //   the normal dequant/idct/ssd/cabac
289
290 // the unquant_mf here is not the same as dequant_mf:
291 // in normal operation (dct->quant->dequant->idct) the dct and idct are not
292 // normalized. quant/dequant absorb those scaling factors.
293 // in this function, we just do (quant->unquant) and want the output to be
294 // comparable to the input. so unquant is the direct inverse of quant,
295 // and uses the dct scaling factors, not the idct ones.
296
297 static void quant_trellis_cabac( x264_t *h, int16_t *dct,
298                                  const uint16_t *quant_mf, const int *unquant_mf,
299                                  const int *coef_weight, const uint8_t *zigzag,
300                                  int i_ctxBlockCat, int i_lambda2, int b_ac, int i_coefs )
301 {
302     int abs_coefs[64], signs[64];
303     trellis_node_t nodes[2][8];
304     trellis_node_t *nodes_cur = nodes[0];
305     trellis_node_t *nodes_prev = nodes[1];
306     trellis_node_t *bnode;
307     uint8_t cabac_state_sig[64];
308     uint8_t cabac_state_last[64];
309     const int b_interlaced = h->mb.b_interlaced;
310     const int f = 1 << 15; // no deadzone
311     int i_last_nnz;
312     int i, j;
313
314     // (# of coefs) * (# of ctx) * (# of levels tried) = 1024
315     // we don't need to keep all of those: (# of coefs) * (# of ctx) would be enough,
316     // but it takes more time to remove dead states than you gain in reduced memory.
317     struct {
318         uint16_t abs_level;
319         uint16_t next;
320     } level_tree[64*8*2];
321     int i_levels_used = 1;
322
323     /* init coefs */
324     for( i = i_coefs-1; i >= b_ac; i-- )
325         if( (unsigned)(dct[zigzag[i]] * quant_mf[zigzag[i]] + f-1) >= 2*f )
326             break;
327
328     if( i < b_ac )
329     {
330         memset( dct, 0, i_coefs * sizeof(*dct) );
331         return;
332     }
333
334     i_last_nnz = i;
335
336     for( ; i >= b_ac; i-- )
337     {
338         int coef = dct[zigzag[i]];
339         abs_coefs[i] = abs(coef);
340         signs[i] = coef < 0 ? -1 : 1;
341     }
342
343     /* init trellis */
344     for( i = 1; i < 8; i++ )
345         nodes_cur[i].score = TRELLIS_SCORE_MAX;
346     nodes_cur[0].score = 0;
347     nodes_cur[0].level_idx = 0;
348     level_tree[0].abs_level = 0;
349     level_tree[0].next = 0;
350
351     // coefs are processed in reverse order, because that's how the abs value is coded.
352     // last_coef and significant_coef flags are normally coded in forward order, but
353     // we have to reverse them to match the levels.
354     // in 4x4 blocks, last_coef and significant_coef use a separate context for each
355     // position, so the order doesn't matter, and we don't even have to update their contexts.
356     // in 8x8 blocks, some positions share contexts, so we'll just have to hope that
357     // cabac isn't too sensitive.
358
359     if( i_coefs == 64 )
360     {
361         const uint8_t *ctx_sig  = &h->cabac.state[ significant_coeff_flag_offset[b_interlaced][i_ctxBlockCat] ];
362         const uint8_t *ctx_last = &h->cabac.state[ last_coeff_flag_offset[b_interlaced][i_ctxBlockCat] ];
363         for( i = 0; i < 63; i++ )
364         {
365             cabac_state_sig[i]  = ctx_sig[ significant_coeff_flag_offset_8x8[b_interlaced][i] ];
366             cabac_state_last[i] = ctx_last[ last_coeff_flag_offset_8x8[i] ];
367         }
368     }
369     else
370     {
371         memcpy( cabac_state_sig,  &h->cabac.state[ significant_coeff_flag_offset[b_interlaced][i_ctxBlockCat] ], 15 );
372         memcpy( cabac_state_last, &h->cabac.state[ last_coeff_flag_offset[b_interlaced][i_ctxBlockCat] ], 15 );
373     }
374     memcpy( nodes_cur[0].cabac_state, &h->cabac.state[ coeff_abs_level_m1_offset[i_ctxBlockCat] ], 10 );
375
376     for( i = i_last_nnz; i >= b_ac; i-- )
377     {
378         int i_coef = abs_coefs[i];
379         int q = ( f + i_coef * quant_mf[zigzag[i]] ) >> 16;
380         int abs_level;
381         int cost_sig[2], cost_last[2];
382         trellis_node_t n;
383
384         // skip 0s: this doesn't affect the output, but saves some unnecessary computation.
385         if( q == 0 )
386         {
387             // no need to calculate ssd of 0s: it's the same in all nodes.
388             // no need to modify level_tree for ctx=0: it starts with an infinite loop of 0s.
389             const uint32_t cost_sig0 = x264_cabac_size_decision_noup( &cabac_state_sig[i], 0 )
390                                      * (uint64_t)i_lambda2 >> ( CABAC_SIZE_BITS - LAMBDA_BITS );
391             for( j = 1; j < 8; j++ )
392             {
393                 if( nodes_cur[j].score != TRELLIS_SCORE_MAX )
394                 {
395 #define SET_LEVEL(n,l) \
396                     level_tree[i_levels_used].abs_level = l; \
397                     level_tree[i_levels_used].next = n.level_idx; \
398                     n.level_idx = i_levels_used; \
399                     i_levels_used++;
400
401                     SET_LEVEL( nodes_cur[j], 0 );
402                     nodes_cur[j].score += cost_sig0;
403                 }
404             }
405             continue;
406         }
407
408         XCHG( trellis_node_t*, nodes_cur, nodes_prev );
409
410         for( j = 0; j < 8; j++ )
411             nodes_cur[j].score = TRELLIS_SCORE_MAX;
412
413         if( i < i_coefs-1 )
414         {
415             cost_sig[0] = x264_cabac_size_decision_noup( &cabac_state_sig[i], 0 );
416             cost_sig[1] = x264_cabac_size_decision_noup( &cabac_state_sig[i], 1 );
417             cost_last[0] = x264_cabac_size_decision_noup( &cabac_state_last[i], 0 );
418             cost_last[1] = x264_cabac_size_decision_noup( &cabac_state_last[i], 1 );
419         }
420         else
421         {
422             cost_sig[0] = cost_sig[1] = 0;
423             cost_last[0] = cost_last[1] = 0;
424         }
425
426         // there are a few cases where increasing the coeff magnitude helps,
427         // but it's only around .003 dB, and skipping them ~doubles the speed of trellis.
428         // could also try q-2: that sometimes helps, but also sometimes decimates blocks
429         // that are better left coded, especially at QP > 40.
430         for( abs_level = q; abs_level >= q-1; abs_level-- )
431         {
432             int d = i_coef - ((unquant_mf[zigzag[i]] * abs_level + 128) >> 8);
433             uint64_t ssd = (int64_t)d*d * coef_weight[i];
434
435             for( j = 0; j < 8; j++ )
436             {
437                 int node_ctx = j;
438                 if( nodes_prev[j].score == TRELLIS_SCORE_MAX )
439                     continue;
440                 n = nodes_prev[j];
441
442                 /* code the proposed level, and count how much entropy it would take */
443                 if( abs_level || node_ctx )
444                 {
445                     unsigned f8_bits = cost_sig[ abs_level != 0 ];
446                     if( abs_level )
447                     {
448                         const int i_prefix = X264_MIN( abs_level - 1, 14 );
449                         f8_bits += cost_last[ node_ctx == 0 ];
450                         f8_bits += x264_cabac_size_decision2( &n.cabac_state[coeff_abs_level1_ctx[node_ctx]], i_prefix > 0 );
451                         if( i_prefix > 0 )
452                         {
453                             uint8_t *ctx = &n.cabac_state[coeff_abs_levelgt1_ctx[node_ctx]];
454                             f8_bits += cabac_prefix_size[i_prefix][*ctx];
455                             *ctx = cabac_prefix_transition[i_prefix][*ctx];
456                             if( abs_level >= 15 )
457                                 f8_bits += bs_size_ue( abs_level - 15 ) << CABAC_SIZE_BITS;
458                             node_ctx = coeff_abs_level_transition[1][node_ctx];
459                         }
460                         else
461                         {
462                             f8_bits += 1 << CABAC_SIZE_BITS;
463                             node_ctx = coeff_abs_level_transition[0][node_ctx];
464                         }
465                     }
466                     n.score += (uint64_t)f8_bits * i_lambda2 >> ( CABAC_SIZE_BITS - LAMBDA_BITS );
467                 }
468
469                 n.score += ssd;
470
471                 /* save the node if it's better than any existing node with the same cabac ctx */
472                 if( n.score < nodes_cur[node_ctx].score )
473                 {
474                     SET_LEVEL( n, abs_level );
475                     nodes_cur[node_ctx] = n;
476                 }
477             }
478         }
479     }
480
481     /* output levels from the best path through the trellis */
482     bnode = &nodes_cur[0];
483     for( j = 1; j < 8; j++ )
484         if( nodes_cur[j].score < bnode->score )
485             bnode = &nodes_cur[j];
486
487     j = bnode->level_idx;
488     for( i = b_ac; i < i_coefs; i++ )
489     {
490         dct[zigzag[i]] = level_tree[j].abs_level * signs[i];
491         j = level_tree[j].next;
492     }
493 }
494
495
496 void x264_quant_4x4_trellis( x264_t *h, int16_t dct[4][4], int i_quant_cat,
497                              int i_qp, int i_ctxBlockCat, int b_intra )
498 {
499     int b_ac = (i_ctxBlockCat == DCT_LUMA_AC);
500     quant_trellis_cabac( h, (int16_t*)dct,
501         h->quant4_mf[i_quant_cat][i_qp], h->unquant4_mf[i_quant_cat][i_qp],
502         x264_dct4_weight2_zigzag[h->mb.b_interlaced],
503         x264_zigzag_scan4[h->mb.b_interlaced],
504         i_ctxBlockCat, lambda2_tab[b_intra][i_qp], b_ac, 16 );
505 }
506
507
508 void x264_quant_8x8_trellis( x264_t *h, int16_t dct[8][8], int i_quant_cat,
509                              int i_qp, int b_intra )
510 {
511     quant_trellis_cabac( h, (int16_t*)dct,
512         h->quant8_mf[i_quant_cat][i_qp], h->unquant8_mf[i_quant_cat][i_qp],
513         x264_dct8_weight2_zigzag[h->mb.b_interlaced],
514         x264_zigzag_scan8[h->mb.b_interlaced],
515         DCT_LUMA_8x8, lambda2_tab[b_intra][i_qp], 0, 64 );
516 }
517