vp9_bitstream.c 43.4 KB
Newer Older
John Koleszar's avatar
John Koleszar committed
1
/*
2
 *  Copyright (c) 2010 The WebM project authors. All Rights Reserved.
John Koleszar's avatar
John Koleszar committed
3
 *
4
 *  Use of this source code is governed by a BSD-style license
5 6
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
7
 *  in the file PATENTS.  All contributing project authors may
8
 *  be found in the AUTHORS file in the root of the source tree.
John Koleszar's avatar
John Koleszar committed
9 10
 */

11 12 13
#include <assert.h>
#include <stdio.h>
#include <limits.h>
John Koleszar's avatar
John Koleszar committed
14

15
#include "vpx/vpx_encoder.h"
16
#include "vpx_dsp/bitwriter_buffer.h"
17
#include "vpx_dsp/vpx_dsp_common.h"
18
#include "vpx_mem/vpx_mem.h"
19
#include "vpx_ports/mem_ops.h"
20
#include "vpx_ports/system_state.h"
21

22
#include "vp9/common/vp9_entropy.h"
23
#include "vp9/common/vp9_entropymode.h"
24
#include "vp9/common/vp9_entropymv.h"
25
#include "vp9/common/vp9_mvref_common.h"
26 27 28
#include "vp9/common/vp9_pred_common.h"
#include "vp9/common/vp9_seg_common.h"
#include "vp9/common/vp9_tile_common.h"
29

Dmitry Kovalev's avatar
Dmitry Kovalev committed
30
#include "vp9/encoder/vp9_cost.h"
31
#include "vp9/encoder/vp9_bitstream.h"
32 33
#include "vp9/encoder/vp9_encodemv.h"
#include "vp9/encoder/vp9_mcomp.h"
34
#include "vp9/encoder/vp9_segmentation.h"
35
#include "vp9/encoder/vp9_subexp.h"
36
#include "vp9/encoder/vp9_tokenize.h"
37

38 39 40 41 42 43 44 45 46
static const struct vp9_token intra_mode_encodings[INTRA_MODES] = {
  {0, 1}, {6, 3}, {28, 5}, {30, 5}, {58, 6}, {59, 6}, {126, 7}, {127, 7},
  {62, 6}, {2, 2}};
static const struct vp9_token switchable_interp_encodings[SWITCHABLE_FILTERS] =
  {{0, 1}, {2, 2}, {3, 2}};
static const struct vp9_token partition_encodings[PARTITION_TYPES] =
  {{0, 1}, {2, 2}, {6, 3}, {7, 3}};
static const struct vp9_token inter_mode_encodings[INTER_MODES] =
  {{2, 2}, {6, 3}, {0, 1}, {7, 3}};
47

Yaowu Xu's avatar
Yaowu Xu committed
48
static void write_intra_mode(vpx_writer *w, PREDICTION_MODE mode,
Yaowu Xu's avatar
Yaowu Xu committed
49
                             const vpx_prob *probs) {
50
  vp9_write_token(w, vp9_intra_mode_tree, probs, &intra_mode_encodings[mode]);
51 52
}

Yaowu Xu's avatar
Yaowu Xu committed
53
static void write_inter_mode(vpx_writer *w, PREDICTION_MODE mode,
Yaowu Xu's avatar
Yaowu Xu committed
54
                             const vpx_prob *probs) {
55
  assert(is_inter_mode(mode));
56 57
  vp9_write_token(w, vp9_inter_mode_tree, probs,
                  &inter_mode_encodings[INTER_OFFSET(mode)]);
58 59
}

60
static void encode_unsigned_max(struct vpx_write_bit_buffer *wb,
61
                                int data, int max) {
62
  vpx_wb_write_literal(wb, data, get_unsigned_bits(max));
63 64
}

Yaowu Xu's avatar
Yaowu Xu committed
65
static void prob_diff_update(const vpx_tree_index *tree,
Yaowu Xu's avatar
Yaowu Xu committed
66
                             vpx_prob probs[/*n - 1*/],
67
                             const unsigned int counts[/*n - 1*/],
Yaowu Xu's avatar
Yaowu Xu committed
68
                             int n, vpx_writer *w) {
69 70
  int i;
  unsigned int branch_ct[32][2];
71 72

  // Assuming max number of probabilities <= 32
73
  assert(n <= 32);
74

75
  vp9_tree_probs_from_distribution(tree, branch_ct, counts);
76
  for (i = 0; i < n - 1; ++i)
77
    vp9_cond_prob_diff_update(w, &probs[i], branch_ct[i]);
78 79
}

80
static void write_selected_tx_size(const VP9_COMMON *cm,
Yaowu Xu's avatar
Yaowu Xu committed
81
                                   const MACROBLOCKD *xd, vpx_writer *w) {
82 83
  TX_SIZE tx_size = xd->mi[0]->mbmi.tx_size;
  BLOCK_SIZE bsize = xd->mi[0]->mbmi.sb_type;
84
  const TX_SIZE max_tx_size = max_txsize_lookup[bsize];
Yaowu Xu's avatar
Yaowu Xu committed
85
  const vpx_prob *const tx_probs = get_tx_probs2(max_tx_size, xd,
86
                                                 &cm->fc->tx_probs);
Yaowu Xu's avatar
Yaowu Xu committed
87
  vpx_write(w, tx_size != TX_4X4, tx_probs[0]);
88
  if (tx_size != TX_4X4 && max_tx_size >= TX_16X16) {
Yaowu Xu's avatar
Yaowu Xu committed
89
    vpx_write(w, tx_size != TX_8X8, tx_probs[1]);
90
    if (tx_size != TX_8X8 && max_tx_size >= TX_32X32)
Yaowu Xu's avatar
Yaowu Xu committed
91
      vpx_write(w, tx_size != TX_16X16, tx_probs[2]);
92 93 94
  }
}

95
static int write_skip(const VP9_COMMON *cm, const MACROBLOCKD *xd,
Yaowu Xu's avatar
Yaowu Xu committed
96
                      int segment_id, const MODE_INFO *mi, vpx_writer *w) {
97
  if (segfeature_active(&cm->seg, segment_id, SEG_LVL_SKIP)) {
98 99
    return 1;
  } else {
100
    const int skip = mi->mbmi.skip;
Yaowu Xu's avatar
Yaowu Xu committed
101
    vpx_write(w, skip, vp9_get_skip_prob(cm, xd));
102
    return skip;
103 104 105
  }
}

Yaowu Xu's avatar
Yaowu Xu committed
106
static void update_skip_probs(VP9_COMMON *cm, vpx_writer *w,
107
                              FRAME_COUNTS *counts) {
108 109
  int k;

Dmitry Kovalev's avatar
Dmitry Kovalev committed
110
  for (k = 0; k < SKIP_CONTEXTS; ++k)
111
    vp9_cond_prob_diff_update(w, &cm->fc->skip_probs[k], counts->skip[k]);
112 113
}

Yaowu Xu's avatar
Yaowu Xu committed
114
static void update_switchable_interp_probs(VP9_COMMON *cm, vpx_writer *w,
115
                                           FRAME_COUNTS *counts) {
116 117 118
  int j;
  for (j = 0; j < SWITCHABLE_FILTER_CONTEXTS; ++j)
    prob_diff_update(vp9_switchable_interp_tree,
119
                     cm->fc->switchable_interp_prob[j],
120
                     counts->switchable_interp[j], SWITCHABLE_FILTERS, w);
121 122
}

Yaowu Xu's avatar
Yaowu Xu committed
123
static void pack_mb_tokens(vpx_writer *w,
124 125
                           TOKENEXTRA **tp, const TOKENEXTRA *const stop,
                           vpx_bit_depth_t bit_depth) {
126
  TOKENEXTRA *p = *tp;
John Koleszar's avatar
John Koleszar committed
127

128
  while (p < stop && p->token != EOSB_TOKEN) {
129
    const int t = p->token;
130
    const struct vp9_token *const a = &vp9_coef_encodings[t];
John Koleszar's avatar
John Koleszar committed
131 132
    int i = 0;
    int v = a->value;
133
    int n = a->len;
134 135 136 137 138 139 140 141 142 143 144 145
#if CONFIG_VP9_HIGHBITDEPTH
    const vp9_extra_bit *b;
    if (bit_depth == VPX_BITS_12)
      b = &vp9_extra_bits_high12[t];
    else if (bit_depth == VPX_BITS_10)
      b = &vp9_extra_bits_high10[t];
    else
      b = &vp9_extra_bits[t];
#else
    const vp9_extra_bit *const b = &vp9_extra_bits[t];
    (void) bit_depth;
#endif  // CONFIG_VP9_HIGHBITDEPTH
146

John Koleszar's avatar
John Koleszar committed
147 148 149 150 151
    /* skip one or two nodes */
    if (p->skip_eob_node) {
      n -= p->skip_eob_node;
      i = 2 * p->skip_eob_node;
    }
John Koleszar's avatar
John Koleszar committed
152

153 154 155 156 157 158 159 160
    // TODO(jbb): expanding this can lead to big gains.  It allows
    // much better branch prediction and would enable us to avoid numerous
    // lookups and compares.

    // If we have a token that's in the constrained set, the coefficient tree
    // is split into two treed writes.  The first treed write takes care of the
    // unconstrained nodes.  The second treed write takes care of the
    // constrained nodes.
161
    if (t >= TWO_TOKEN && t < EOB_TOKEN) {
162 163
      int len = UNCONSTRAINED_NODES - p->skip_eob_node;
      int bits = v >> (n - len);
164 165 166 167
      vp9_write_tree(w, vp9_coef_tree, p->context_tree, bits, len, i);
      vp9_write_tree(w, vp9_coef_con_tree,
                     vp9_pareto8_full[p->context_tree[PIVOT_NODE] - 1],
                     v, n - len, 0);
168
    } else {
169
      vp9_write_tree(w, vp9_coef_tree, p->context_tree, v, n, i);
170
    }
John Koleszar's avatar
John Koleszar committed
171

John Koleszar's avatar
John Koleszar committed
172
    if (b->base_val) {
173
      const int e = p->extra, l = b->len;
John Koleszar's avatar
John Koleszar committed
174

175
      if (l) {
176
        const unsigned char *pb = b->prob;
John Koleszar's avatar
John Koleszar committed
177
        int v = e >> 1;
178
        int n = l;              /* number of bits in v, assumed nonzero */
John Koleszar's avatar
John Koleszar committed
179
        int i = 0;
John Koleszar's avatar
John Koleszar committed
180

John Koleszar's avatar
John Koleszar committed
181 182
        do {
          const int bb = (v >> --n) & 1;
Yaowu Xu's avatar
Yaowu Xu committed
183
          vpx_write(w, bb, pb[i >> 1]);
John Koleszar's avatar
John Koleszar committed
184 185 186
          i = b->tree[i + bb];
        } while (n);
      }
John Koleszar's avatar
John Koleszar committed
187

Yaowu Xu's avatar
Yaowu Xu committed
188
      vpx_write_bit(w, e & 1);
John Koleszar's avatar
John Koleszar committed
189
    }
John Koleszar's avatar
John Koleszar committed
190 191 192
    ++p;
  }

193
  *tp = p + (p->token == EOSB_TOKEN);
John Koleszar's avatar
John Koleszar committed
194 195
}

Yaowu Xu's avatar
Yaowu Xu committed
196
static void write_segment_id(vpx_writer *w, const struct segmentation *seg,
197
                             int segment_id) {
198
  if (seg->enabled && seg->update_map)
199
    vp9_write_tree(w, vp9_segment_tree, seg->tree_probs, segment_id, 3, 0);
John Koleszar's avatar
John Koleszar committed
200 201
}

Paul Wilkins's avatar
Paul Wilkins committed
202
// This function encodes the reference frame
203
static void write_ref_frames(const VP9_COMMON *cm, const MACROBLOCKD *xd,
Yaowu Xu's avatar
Yaowu Xu committed
204
                             vpx_writer *w) {
205
  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
206 207 208
  const int is_compound = has_second_ref(mbmi);
  const int segment_id = mbmi->segment_id;

John Koleszar's avatar
John Koleszar committed
209 210
  // If segment level coding of this signal is disabled...
  // or the segment allows multiple reference frame options
211
  if (segfeature_active(&cm->seg, segment_id, SEG_LVL_REF_FRAME)) {
212 213
    assert(!is_compound);
    assert(mbmi->ref_frame[0] ==
Scott LaVarnway's avatar
Scott LaVarnway committed
214
               get_segdata(&cm->seg, segment_id, SEG_LVL_REF_FRAME));
215
  } else {
Ronald S. Bultje's avatar
Ronald S. Bultje committed
216 217
    // does the feature use compound prediction or not
    // (if not specified at the frame/segment level)
218
    if (cm->reference_mode == REFERENCE_MODE_SELECT) {
Yaowu Xu's avatar
Yaowu Xu committed
219
      vpx_write(w, is_compound, vp9_get_reference_mode_prob(cm, xd));
Ronald S. Bultje's avatar
Ronald S. Bultje committed
220
    } else {
221
      assert(!is_compound == (cm->reference_mode == SINGLE_REFERENCE));
Ronald S. Bultje's avatar
Ronald S. Bultje committed
222
    }
223

224
    if (is_compound) {
Yaowu Xu's avatar
Yaowu Xu committed
225
      vpx_write(w, mbmi->ref_frame[0] == GOLDEN_FRAME,
226
                vp9_get_pred_prob_comp_ref_p(cm, xd));
Ronald S. Bultje's avatar
Ronald S. Bultje committed
227
    } else {
228
      const int bit0 = mbmi->ref_frame[0] != LAST_FRAME;
Yaowu Xu's avatar
Yaowu Xu committed
229
      vpx_write(w, bit0, vp9_get_pred_prob_single_ref_p1(cm, xd));
230 231
      if (bit0) {
        const int bit1 = mbmi->ref_frame[0] != GOLDEN_FRAME;
Yaowu Xu's avatar
Yaowu Xu committed
232
        vpx_write(w, bit1, vp9_get_pred_prob_single_ref_p2(cm, xd));
233
      }
Paul Wilkins's avatar
Paul Wilkins committed
234
    }
John Koleszar's avatar
John Koleszar committed
235
  }
Paul Wilkins's avatar
Paul Wilkins committed
236
}
John Koleszar's avatar
John Koleszar committed
237

238
static void pack_inter_mode_mvs(VP9_COMP *cpi, const MODE_INFO *mi,
Yaowu Xu's avatar
Yaowu Xu committed
239
                                vpx_writer *w) {
240
  VP9_COMMON *const cm = &cpi->common;
241
  const nmv_context *nmvc = &cm->fc->nmvc;
242
  const MACROBLOCK *const x = &cpi->td.mb;
243
  const MACROBLOCKD *const xd = &x->e_mbd;
244
  const struct segmentation *const seg = &cm->seg;
245
  const MB_MODE_INFO *const mbmi = &mi->mbmi;
246
  const MB_MODE_INFO_EXT *const mbmi_ext = x->mbmi_ext;
247
  const PREDICTION_MODE mode = mbmi->mode;
248 249
  const int segment_id = mbmi->segment_id;
  const BLOCK_SIZE bsize = mbmi->sb_type;
250
  const int allow_hp = cm->allow_high_precision_mv;
251 252 253
  const int is_inter = is_inter_block(mbmi);
  const int is_compound = has_second_ref(mbmi);
  int skip, ref;
Adrian Grange's avatar
Adrian Grange committed
254

255 256
  if (seg->update_map) {
    if (seg->temporal_update) {
257
      const int pred_flag = mbmi->seg_id_predicted;
Yaowu Xu's avatar
Yaowu Xu committed
258
      vpx_prob pred_prob = vp9_get_pred_prob_seg_id(seg, xd);
Yaowu Xu's avatar
Yaowu Xu committed
259
      vpx_write(w, pred_flag, pred_prob);
260
      if (!pred_flag)
261
        write_segment_id(w, seg, segment_id);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
262
    } else {
263
      write_segment_id(w, seg, segment_id);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
264 265
    }
  }
266

267
  skip = write_skip(cm, xd, segment_id, mi, w);
John Koleszar's avatar
John Koleszar committed
268

269
  if (!segfeature_active(seg, segment_id, SEG_LVL_REF_FRAME))
Yaowu Xu's avatar
Yaowu Xu committed
270
    vpx_write(w, is_inter, vp9_get_intra_inter_prob(cm, xd));
Paul Wilkins's avatar
Paul Wilkins committed
271

272
  if (bsize >= BLOCK_8X8 && cm->tx_mode == TX_MODE_SELECT &&
Alex Converse's avatar
Alex Converse committed
273
      !(is_inter && skip)) {
274
    write_selected_tx_size(cm, xd, w);
275 276
  }

277
  if (!is_inter) {
278
    if (bsize >= BLOCK_8X8) {
279
      write_intra_mode(w, mode, cm->fc->y_mode_prob[size_group_lookup[bsize]]);
280
    } else {
281
      int idx, idy;
282 283 284 285
      const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize];
      const int num_4x4_h = num_4x4_blocks_high_lookup[bsize];
      for (idy = 0; idy < 2; idy += num_4x4_h) {
        for (idx = 0; idx < 2; idx += num_4x4_w) {
286
          const PREDICTION_MODE b_mode = mi->bmi[idy * 2 + idx].as_mode;
287
          write_intra_mode(w, b_mode, cm->fc->y_mode_prob[0]);
288
        }
Jim Bankoski's avatar
Jim Bankoski committed
289
      }
Ronald S. Bultje's avatar
Ronald S. Bultje committed
290
    }
291
    write_intra_mode(w, mbmi->uv_mode, cm->fc->uv_mode_prob[mode]);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
292
  } else {
293
    const int mode_ctx = mbmi_ext->mode_context[mbmi->ref_frame[0]];
Yaowu Xu's avatar
Yaowu Xu committed
294
    const vpx_prob *const inter_probs = cm->fc->inter_mode_probs[mode_ctx];
295
    write_ref_frames(cm, xd, w);
Yaowu Xu's avatar
Yaowu Xu committed
296

297
    // If segment skip is not enabled code the mode.
298
    if (!segfeature_active(seg, segment_id, SEG_LVL_SKIP)) {
299
      if (bsize >= BLOCK_8X8) {
300
        write_inter_mode(w, mode, inter_probs);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
301
      }
Ronald S. Bultje's avatar
Ronald S. Bultje committed
302
    }
303

304
    if (cm->interp_filter == SWITCHABLE) {
305
      const int ctx = vp9_get_pred_context_switchable_interp(xd);
306
      vp9_write_token(w, vp9_switchable_interp_tree,
307
                      cm->fc->switchable_interp_prob[ctx],
308
                      &switchable_interp_encodings[mbmi->interp_filter]);
309
      ++cpi->interp_filter_selected[0][mbmi->interp_filter];
Ronald S. Bultje's avatar
Ronald S. Bultje committed
310
    } else {
311
      assert(mbmi->interp_filter == cm->interp_filter);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
312
    }
313

314
    if (bsize < BLOCK_8X8) {
315 316
      const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize];
      const int num_4x4_h = num_4x4_blocks_high_lookup[bsize];
Ronald S. Bultje's avatar
Ronald S. Bultje committed
317
      int idx, idy;
318 319
      for (idy = 0; idy < 2; idy += num_4x4_h) {
        for (idx = 0; idx < 2; idx += num_4x4_w) {
320
          const int j = idy * 2 + idx;
321
          const PREDICTION_MODE b_mode = mi->bmi[j].as_mode;
322
          write_inter_mode(w, b_mode, inter_probs);
323
          if (b_mode == NEWMV) {
324 325
            for (ref = 0; ref < 1 + is_compound; ++ref)
              vp9_encode_mv(cpi, w, &mi->bmi[j].as_mv[ref].as_mv,
326
                            &mbmi_ext->ref_mvs[mbmi->ref_frame[ref]][0].as_mv,
327
                            nmvc, allow_hp);
John Koleszar's avatar
John Koleszar committed
328
          }
329
        }
Ronald S. Bultje's avatar
Ronald S. Bultje committed
330
      }
331 332 333 334
    } else {
      if (mode == NEWMV) {
        for (ref = 0; ref < 1 + is_compound; ++ref)
          vp9_encode_mv(cpi, w, &mbmi->mv[ref].as_mv,
335
                        &mbmi_ext->ref_mvs[mbmi->ref_frame[ref]][0].as_mv, nmvc,
336 337
                        allow_hp);
      }
Ronald S. Bultje's avatar
Ronald S. Bultje committed
338 339
    }
  }
John Koleszar's avatar
John Koleszar committed
340
}
341

342
static void write_mb_modes_kf(const VP9_COMMON *cm, const MACROBLOCKD *xd,
Yaowu Xu's avatar
Yaowu Xu committed
343
                              MODE_INFO **mi_8x8, vpx_writer *w) {
344
  const struct segmentation *const seg = &cm->seg;
345
  const MODE_INFO *const mi = mi_8x8[0];
346 347
  const MODE_INFO *const above_mi = xd->above_mi;
  const MODE_INFO *const left_mi = xd->left_mi;
348 349
  const MB_MODE_INFO *const mbmi = &mi->mbmi;
  const BLOCK_SIZE bsize = mbmi->sb_type;
350

351
  if (seg->update_map)
352
    write_segment_id(w, seg, mbmi->segment_id);
353

354
  write_skip(cm, xd, mbmi->segment_id, mi, w);
355

356
  if (bsize >= BLOCK_8X8 && cm->tx_mode == TX_MODE_SELECT)
357
    write_selected_tx_size(cm, xd, w);
358

359 360
  if (bsize >= BLOCK_8X8) {
    write_intra_mode(w, mbmi->mode, get_y_mode_probs(mi, above_mi, left_mi, 0));
361
  } else {
362 363
    const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize];
    const int num_4x4_h = num_4x4_blocks_high_lookup[bsize];
364
    int idx, idy;
365 366 367 368 369 370

    for (idy = 0; idy < 2; idy += num_4x4_h) {
      for (idx = 0; idx < 2; idx += num_4x4_w) {
        const int block = idy * 2 + idx;
        write_intra_mode(w, mi->bmi[block].as_mode,
                         get_y_mode_probs(mi, above_mi, left_mi, block));
371 372
      }
    }
373 374
  }

375
  write_intra_mode(w, mbmi->uv_mode, vp9_kf_uv_mode_prob[mbmi->mode]);
376 377
}

James Zern's avatar
James Zern committed
378
static void write_modes_b(VP9_COMP *cpi, const TileInfo *const tile,
Yaowu Xu's avatar
Yaowu Xu committed
379
                          vpx_writer *w, TOKENEXTRA **tok,
380
                          const TOKENEXTRA *const tok_end,
381
                          int mi_row, int mi_col) {
382
  const VP9_COMMON *const cm = &cpi->common;
383
  MACROBLOCKD *const xd = &cpi->td.mb.e_mbd;
384
  MODE_INFO *m;
Dmitry Kovalev's avatar
Dmitry Kovalev committed
385

386 387
  xd->mi = cm->mi_grid_visible + (mi_row * cm->mi_stride + mi_col);
  m = xd->mi[0];
388

389 390 391
  cpi->td.mb.mbmi_ext = cpi->td.mb.mbmi_ext_base +
      (mi_row * cm->mi_cols + mi_col);

James Zern's avatar
James Zern committed
392
  set_mi_row_col(xd, tile,
Dmitry Kovalev's avatar
Dmitry Kovalev committed
393
                 mi_row, num_8x8_blocks_high_lookup[m->mbmi.sb_type],
James Zern's avatar
James Zern committed
394 395
                 mi_col, num_8x8_blocks_wide_lookup[m->mbmi.sb_type],
                 cm->mi_rows, cm->mi_cols);
396
  if (frame_is_intra_only(cm)) {
397
    write_mb_modes_kf(cm, xd, xd->mi, w);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
398
  } else {
399
    pack_inter_mode_mvs(cpi, m, w);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
400 401 402
  }

  assert(*tok < tok_end);
403
  pack_mb_tokens(w, tok, tok_end, cm->bit_depth);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
404 405
}

406 407
static void write_partition(const VP9_COMMON *const cm,
                            const MACROBLOCKD *const xd,
408
                            int hbs, int mi_row, int mi_col,
Yaowu Xu's avatar
Yaowu Xu committed
409
                            PARTITION_TYPE p, BLOCK_SIZE bsize, vpx_writer *w) {
410
  const int ctx = partition_plane_context(xd, mi_row, mi_col, bsize);
Yaowu Xu's avatar
Yaowu Xu committed
411
  const vpx_prob *const probs = xd->partition_probs[ctx];
412 413
  const int has_rows = (mi_row + hbs) < cm->mi_rows;
  const int has_cols = (mi_col + hbs) < cm->mi_cols;
414 415

  if (has_rows && has_cols) {
416
    vp9_write_token(w, vp9_partition_tree, probs, &partition_encodings[p]);
417
  } else if (!has_rows && has_cols) {
418
    assert(p == PARTITION_SPLIT || p == PARTITION_HORZ);
Yaowu Xu's avatar
Yaowu Xu committed
419
    vpx_write(w, p == PARTITION_SPLIT, probs[1]);
420
  } else if (has_rows && !has_cols) {
421
    assert(p == PARTITION_SPLIT || p == PARTITION_VERT);
Yaowu Xu's avatar
Yaowu Xu committed
422
    vpx_write(w, p == PARTITION_SPLIT, probs[2]);
423
  } else {
424
    assert(p == PARTITION_SPLIT);
425 426 427
  }
}

428
static void write_modes_sb(VP9_COMP *cpi,
Yaowu Xu's avatar
Yaowu Xu committed
429
                           const TileInfo *const tile, vpx_writer *w,
430
                           TOKENEXTRA **tok, const TOKENEXTRA *const tok_end,
431
                           int mi_row, int mi_col, BLOCK_SIZE bsize) {
432
  const VP9_COMMON *const cm = &cpi->common;
433
  MACROBLOCKD *const xd = &cpi->td.mb.e_mbd;
434

435
  const int bsl = b_width_log2_lookup[bsize];
436 437
  const int bs = (1 << bsl) / 4;
  PARTITION_TYPE partition;
438
  BLOCK_SIZE subsize;
439
  const MODE_INFO *m = NULL;
440

441
  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols)
442 443
    return;

444
  m = cm->mi_grid_visible[mi_row * cm->mi_stride + mi_col];
445

Jim Bankoski's avatar
Jim Bankoski committed
446
  partition = partition_lookup[bsl][m->mbmi.sb_type];
447
  write_partition(cm, xd, bs, mi_row, mi_col, partition, bsize, w);
448
  subsize = get_subsize(bsize, partition);
449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477
  if (subsize < BLOCK_8X8) {
    write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col);
  } else {
    switch (partition) {
      case PARTITION_NONE:
        write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col);
        break;
      case PARTITION_HORZ:
        write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col);
        if (mi_row + bs < cm->mi_rows)
          write_modes_b(cpi, tile, w, tok, tok_end, mi_row + bs, mi_col);
        break;
      case PARTITION_VERT:
        write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col);
        if (mi_col + bs < cm->mi_cols)
          write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col + bs);
        break;
      case PARTITION_SPLIT:
        write_modes_sb(cpi, tile, w, tok, tok_end, mi_row, mi_col, subsize);
        write_modes_sb(cpi, tile, w, tok, tok_end, mi_row, mi_col + bs,
                       subsize);
        write_modes_sb(cpi, tile, w, tok, tok_end, mi_row + bs, mi_col,
                       subsize);
        write_modes_sb(cpi, tile, w, tok, tok_end, mi_row + bs, mi_col + bs,
                       subsize);
        break;
      default:
        assert(0);
    }
478
  }
479 480

  // update partition context
481
  if (bsize >= BLOCK_8X8 &&
482
      (bsize == BLOCK_8X8 || partition != PARTITION_SPLIT))
483
    update_partition_context(xd, mi_row, mi_col, subsize, bsize);
484 485
}

486
static void write_modes(VP9_COMP *cpi,
Yaowu Xu's avatar
Yaowu Xu committed
487
                        const TileInfo *const tile, vpx_writer *w,
488
                        TOKENEXTRA **tok, const TOKENEXTRA *const tok_end) {
489
  const VP9_COMMON *const cm = &cpi->common;
490
  MACROBLOCKD *const xd = &cpi->td.mb.e_mbd;
491
  int mi_row, mi_col;
492

493 494
  set_partition_probs(cm, xd);

James Zern's avatar
James Zern committed
495
  for (mi_row = tile->mi_row_start; mi_row < tile->mi_row_end;
496
       mi_row += MI_BLOCK_SIZE) {
497
    vp9_zero(xd->left_seg_context);
James Zern's avatar
James Zern committed
498
    for (mi_col = tile->mi_col_start; mi_col < tile->mi_col_end;
499
         mi_col += MI_BLOCK_SIZE)
500 501
      write_modes_sb(cpi, tile, w, tok, tok_end, mi_row, mi_col,
                     BLOCK_64X64);
John Koleszar's avatar
John Koleszar committed
502
  }
John Koleszar's avatar
John Koleszar committed
503
}
504

505
static void build_tree_distribution(VP9_COMP *cpi, TX_SIZE tx_size,
506 507
                                    vp9_coeff_stats *coef_branch_ct,
                                    vp9_coeff_probs_model *coef_probs) {
508
  vp9_coeff_count *coef_counts = cpi->td.rd_counts.coef_counts[tx_size];
509
  unsigned int (*eob_branch_ct)[REF_TYPES][COEF_BANDS][COEFF_CONTEXTS] =
510
      cpi->common.counts.eob_branch[tx_size];
511
  int i, j, k, l, m;
512

513
  for (i = 0; i < PLANE_TYPES; ++i) {
514 515
    for (j = 0; j < REF_TYPES; ++j) {
      for (k = 0; k < COEF_BANDS; ++k) {
516
        for (l = 0; l < BAND_COEFF_CONTEXTS(k); ++l) {
517
          vp9_tree_probs_from_distribution(vp9_coef_tree,
518
                                           coef_branch_ct[i][j][k][l],
519
                                           coef_counts[i][j][k][l]);
520 521
          coef_branch_ct[i][j][k][l][0][1] = eob_branch_ct[i][j][k][l] -
                                             coef_branch_ct[i][j][k][l][0][0];
522 523 524 525
          for (m = 0; m < UNCONSTRAINED_NODES; ++m)
            coef_probs[i][j][k][l][m] = get_binary_prob(
                                            coef_branch_ct[i][j][k][l][m][0],
                                            coef_branch_ct[i][j][k][l][m][1]);
526
        }
Daniel Kang's avatar
Daniel Kang committed
527 528 529
      }
    }
  }
530 531
}

Yaowu Xu's avatar
Yaowu Xu committed
532
static void update_coef_probs_common(vpx_writer* const bc, VP9_COMP *cpi,
533
                                     TX_SIZE tx_size,
534 535
                                     vp9_coeff_stats *frame_branch_ct,
                                     vp9_coeff_probs_model *new_coef_probs) {
536
  vp9_coeff_probs_model *old_coef_probs = cpi->common.fc->coef_probs[tx_size];
Yaowu Xu's avatar
Yaowu Xu committed
537
  const vpx_prob upd = DIFF_UPDATE_PROB;
538
  const int entropy_nodes_update = UNCONSTRAINED_NODES;
539
  int i, j, k, l, t;
540 541
  int stepsize = cpi->sf.coeff_prob_appx_step;

542
  switch (cpi->sf.use_fast_coef_updates) {
543
    case TWO_LOOP: {
James Zern's avatar
James Zern committed
544
      /* dry run to see if there is any update at all needed */
545 546
      int savings = 0;
      int update[2] = {0, 0};
547
      for (i = 0; i < PLANE_TYPES; ++i) {
548 549
        for (j = 0; j < REF_TYPES; ++j) {
          for (k = 0; k < COEF_BANDS; ++k) {
550
            for (l = 0; l < BAND_COEFF_CONTEXTS(k); ++l) {
551
              for (t = 0; t < entropy_nodes_update; ++t) {
Yaowu Xu's avatar
Yaowu Xu committed
552 553
                vpx_prob newp = new_coef_probs[i][j][k][l][t];
                const vpx_prob oldp = old_coef_probs[i][j][k][l][t];
554 555 556 557 558
                int s;
                int u = 0;
                if (t == PIVOT_NODE)
                  s = vp9_prob_diff_update_savings_search_model(
                      frame_branch_ct[i][j][k][l][0],
559
                      old_coef_probs[i][j][k][l], &newp, upd, stepsize);
560 561 562 563 564 565 566 567 568 569 570 571 572 573 574
                else
                  s = vp9_prob_diff_update_savings_search(
                      frame_branch_ct[i][j][k][l][t], oldp, &newp, upd);
                if (s > 0 && newp != oldp)
                  u = 1;
                if (u)
                  savings += s - (int)(vp9_cost_zero(upd));
                else
                  savings -= (int)(vp9_cost_zero(upd));
                update[u]++;
              }
            }
          }
        }
      }
575

576 577 578
      // printf("Update %d %d, savings %d\n", update[0], update[1], savings);
      /* Is coef updated at all */
      if (update[1] == 0 || savings < 0) {
Yaowu Xu's avatar
Yaowu Xu committed
579
        vpx_write_bit(bc, 0);
580 581
        return;
      }
Yaowu Xu's avatar
Yaowu Xu committed
582
      vpx_write_bit(bc, 1);
583
      for (i = 0; i < PLANE_TYPES; ++i) {
584 585
        for (j = 0; j < REF_TYPES; ++j) {
          for (k = 0; k < COEF_BANDS; ++k) {
586
            for (l = 0; l < BAND_COEFF_CONTEXTS(k); ++l) {
587 588
              // calc probs and branch cts for this frame only
              for (t = 0; t < entropy_nodes_update; ++t) {
Yaowu Xu's avatar
Yaowu Xu committed
589 590 591
                vpx_prob newp = new_coef_probs[i][j][k][l][t];
                vpx_prob *oldp = old_coef_probs[i][j][k][l] + t;
                const vpx_prob upd = DIFF_UPDATE_PROB;
592 593 594 595 596
                int s;
                int u = 0;
                if (t == PIVOT_NODE)
                  s = vp9_prob_diff_update_savings_search_model(
                      frame_branch_ct[i][j][k][l][0],
597
                      old_coef_probs[i][j][k][l], &newp, upd, stepsize);
598 599 600 601 602 603
                else
                  s = vp9_prob_diff_update_savings_search(
                      frame_branch_ct[i][j][k][l][t],
                      *oldp, &newp, upd);
                if (s > 0 && newp != *oldp)
                  u = 1;
Yaowu Xu's avatar
Yaowu Xu committed
604
                vpx_write(bc, u, upd);
605 606 607 608 609 610 611
                if (u) {
                  /* send/use new probability */
                  vp9_write_prob_diff_update(bc, newp, *oldp);
                  *oldp = newp;
                }
              }
            }
612
          }
Daniel Kang's avatar
Daniel Kang committed
613 614
        }
      }
615
      return;
Daniel Kang's avatar
Daniel Kang committed
616
    }
John Koleszar's avatar
John Koleszar committed
617

618
    case ONE_LOOP_REDUCED: {
619 620
      int updates = 0;
      int noupdates_before_first = 0;
621
      for (i = 0; i < PLANE_TYPES; ++i) {
622 623
        for (j = 0; j < REF_TYPES; ++j) {
          for (k = 0; k < COEF_BANDS; ++k) {
624
            for (l = 0; l < BAND_COEFF_CONTEXTS(k); ++l) {
625 626
              // calc probs and branch cts for this frame only
              for (t = 0; t < entropy_nodes_update; ++t) {
Yaowu Xu's avatar
Yaowu Xu committed
627 628
                vpx_prob newp = new_coef_probs[i][j][k][l][t];
                vpx_prob *oldp = old_coef_probs[i][j][k][l] + t;
629 630
                int s;
                int u = 0;
631 632 633 634 635

                if (t == PIVOT_NODE) {
                  s = vp9_prob_diff_update_savings_search_model(
                      frame_branch_ct[i][j][k][l][0],
                      old_coef_probs[i][j][k][l], &newp, upd, stepsize);
636
                } else {
637 638 639
                  s = vp9_prob_diff_update_savings_search(
                      frame_branch_ct[i][j][k][l][t],
                      *oldp, &newp, upd);
640
                }
641 642 643

                if (s > 0 && newp != *oldp)
                  u = 1;
644 645 646 647 648 649 650 651
                updates += u;
                if (u == 0 && updates == 0) {
                  noupdates_before_first++;
                  continue;
                }
                if (u == 1 && updates == 1) {
                  int v;
                  // first update
Yaowu Xu's avatar
Yaowu Xu committed
652
                  vpx_write_bit(bc, 1);
653
                  for (v = 0; v < noupdates_before_first; ++v)
Yaowu Xu's avatar
Yaowu Xu committed
654
                    vpx_write(bc, 0, upd);