Decoder's code cleanup.

Using vp9_set_pred_flag function instead of custom code, adding
decode_tokens function which is now called from decode_atom,
decode_sb_intra, and decode_sb.

Change-Id: Ie163a7106c0241099da9c5fe03069bd71f9d9ff8

diff --git a/vp9/common/vp9_entropymode.c b/vp9/common/vp9_entropymode.c
index 33028146ac47615f9bb9aeba9bb601fb8299ce2a..9c8d3552de6095a3d1ab69d35e0597c2c38342e4 100644 (file)
--- a/vp9/common/vp9_entropymode.c
+++ b/vp9/common/vp9_entropymode.c
@@ -181,52 +181,35 @@ void tx_counts_to_branch_counts_32x32(unsigned int *tx_count_32x32p,
 void tx_counts_to_branch_counts_16x16(unsigned int *tx_count_16x16p,
                                       unsigned int (*ct_16x16p)[2]) {
   ct_16x16p[0][0] = tx_count_16x16p[TX_4X4];
-  ct_16x16p[0][1] = tx_count_16x16p[TX_8X8] +
-                    tx_count_16x16p[TX_16X16];
+  ct_16x16p[0][1] = tx_count_16x16p[TX_8X8] + tx_count_16x16p[TX_16X16];
   ct_16x16p[1][0] = tx_count_16x16p[TX_8X8];
   ct_16x16p[1][1] = tx_count_16x16p[TX_16X16];
 }
 
 void tx_counts_to_branch_counts_8x8(unsigned int *tx_count_8x8p,
                                     unsigned int (*ct_8x8p)[2]) {
-  ct_8x8p[0][0] =   tx_count_8x8p[TX_4X4];
-  ct_8x8p[0][1] =   tx_count_8x8p[TX_8X8];
+  ct_8x8p[0][0] = tx_count_8x8p[TX_4X4];
+  ct_8x8p[0][1] = tx_count_8x8p[TX_8X8];
 }
 
 const vp9_prob vp9_default_mbskip_probs[MBSKIP_CONTEXTS] = {
   192, 128, 64
 };
 
-void vp9_init_mbmode_probs(VP9_COMMON *x) {
-  vpx_memcpy(x->fc.uv_mode_prob, default_if_uv_probs,
-             sizeof(default_if_uv_probs));
-  vpx_memcpy(x->kf_uv_mode_prob, default_kf_uv_probs,
-             sizeof(default_kf_uv_probs));
-  vpx_memcpy(x->fc.y_mode_prob, default_if_y_probs,
-             sizeof(default_if_y_probs));
-
-  vpx_memcpy(x->fc.switchable_interp_prob, vp9_switchable_interp_prob,
-             sizeof(vp9_switchable_interp_prob));
-
-  vpx_memcpy(x->fc.partition_prob, vp9_partition_probs,
-             sizeof(vp9_partition_probs));
-
-  vpx_memcpy(x->fc.intra_inter_prob, default_intra_inter_p,
-             sizeof(default_intra_inter_p));
-  vpx_memcpy(x->fc.comp_inter_prob, default_comp_inter_p,
-             sizeof(default_comp_inter_p));
-  vpx_memcpy(x->fc.comp_ref_prob, default_comp_ref_p,
-             sizeof(default_comp_ref_p));
-  vpx_memcpy(x->fc.single_ref_prob, default_single_ref_p,
-             sizeof(default_single_ref_p));
-  vpx_memcpy(x->fc.tx_probs_32x32p, vp9_default_tx_probs_32x32p,
-             sizeof(vp9_default_tx_probs_32x32p));
-  vpx_memcpy(x->fc.tx_probs_16x16p, vp9_default_tx_probs_16x16p,
-             sizeof(vp9_default_tx_probs_16x16p));
-  vpx_memcpy(x->fc.tx_probs_8x8p, vp9_default_tx_probs_8x8p,
-             sizeof(vp9_default_tx_probs_8x8p));
-  vpx_memcpy(x->fc.mbskip_probs, vp9_default_mbskip_probs,
-             sizeof(vp9_default_mbskip_probs));
+void vp9_init_mbmode_probs(VP9_COMMON *cm) {
+  vp9_copy(cm->fc.uv_mode_prob, default_if_uv_probs);
+  vp9_copy(cm->kf_uv_mode_prob, default_kf_uv_probs);
+  vp9_copy(cm->fc.y_mode_prob, default_if_y_probs);
+  vp9_copy(cm->fc.switchable_interp_prob, vp9_switchable_interp_prob);
+  vp9_copy(cm->fc.partition_prob, vp9_partition_probs);
+  vp9_copy(cm->fc.intra_inter_prob, default_intra_inter_p);
+  vp9_copy(cm->fc.comp_inter_prob, default_comp_inter_p);
+  vp9_copy(cm->fc.comp_ref_prob, default_comp_ref_p);
+  vp9_copy(cm->fc.single_ref_prob, default_single_ref_p);
+  vp9_copy(cm->fc.tx_probs_32x32p, vp9_default_tx_probs_32x32p);
+  vp9_copy(cm->fc.tx_probs_16x16p, vp9_default_tx_probs_16x16p);
+  vp9_copy(cm->fc.tx_probs_8x8p, vp9_default_tx_probs_8x8p);
+  vp9_copy(cm->fc.mbskip_probs, vp9_default_mbskip_probs);
 }
 
 const vp9_tree_index vp9_switchable_interp_tree[VP9_SWITCHABLE_FILTERS*2-2] = {
@@ -253,15 +236,13 @@ void vp9_entropy_mode_init() {
   vp9_tokens_from_tree(vp9_switchable_interp_encodings,
                        vp9_switchable_interp_tree);
   vp9_tokens_from_tree(vp9_partition_encodings, vp9_partition_tree);
-
   vp9_tokens_from_tree_offset(vp9_sb_mv_ref_encoding_array,
                               vp9_sb_mv_ref_tree, NEARESTMV);
 }
 
 void vp9_init_mode_contexts(VP9_COMMON *pc) {
   vpx_memset(pc->fc.inter_mode_counts, 0, sizeof(pc->fc.inter_mode_counts));
-  vpx_memcpy(pc->fc.inter_mode_probs,
-             vp9_default_inter_mode_probs,
+  vpx_memcpy(pc->fc.inter_mode_probs, vp9_default_inter_mode_probs,
              sizeof(vp9_default_inter_mode_probs));
 }
 

diff --git a/vp9/common/vp9_pred_common.c b/vp9/common/vp9_pred_common.c
index 3adfaebdbf9456d5e05285f0ac67748089b5e777..e7c91200ded04c0d5598895dbcf5ad354ec67386 100644 (file)
--- a/vp9/common/vp9_pred_common.c
+++ b/vp9/common/vp9_pred_common.c
@@ -455,10 +455,9 @@ unsigned char vp9_get_pred_flag(const MACROBLOCKD *const xd,
 
 // This function sets the status of the given prediction signal.
 // I.e. is the predicted value for the given signal correct.
-void vp9_set_pred_flag(MACROBLOCKD *xd, PRED_ID pred_id,
+void vp9_set_pred_flag(MACROBLOCKD *xd, BLOCK_SIZE_TYPE bsize, PRED_ID pred_id,
                        unsigned char pred_flag) {
   const int mis = xd->mode_info_stride;
-  BLOCK_SIZE_TYPE bsize = xd->mode_info_context->mbmi.sb_type;
   const int bh = 1 << mi_height_log2(bsize);
   const int bw = 1 << mi_width_log2(bsize);
 #define sub(a, b) (b) < 0 ? (a) + (b) : (a)

diff --git a/vp9/common/vp9_pred_common.h b/vp9/common/vp9_pred_common.h
index 2717ff18858b9cb20e77f48d25fd20478b11f0dc..0954ff215ce2b52756b8834679328f8403a0c929 100644 (file)
--- a/vp9/common/vp9_pred_common.h
+++ b/vp9/common/vp9_pred_common.h
@@ -38,7 +38,7 @@ const vp9_prob *vp9_get_pred_probs(const VP9_COMMON *cm, const MACROBLOCKD *xd,
 
 unsigned char vp9_get_pred_flag(const MACROBLOCKD *xd, PRED_ID pred_id);
 
-void vp9_set_pred_flag(MACROBLOCKD *xd, PRED_ID pred_id,
+void vp9_set_pred_flag(MACROBLOCKD *xd, BLOCK_SIZE_TYPE bsize, PRED_ID pred_id,
                        unsigned char pred_flag);
 
 int vp9_get_segment_id(VP9_COMMON *cm, const uint8_t *segment_ids,

diff --git a/vp9/common/vp9_reconinter.c b/vp9/common/vp9_reconinter.c
index f1d155819bd524585e53db898a5f7b0d6f8448fd..094d827e56dae586e7cc8b80845e3c9f6488fc6f 100644 (file)
--- a/vp9/common/vp9_reconinter.c
+++ b/vp9/common/vp9_reconinter.c
@@ -172,9 +172,7 @@ void vp9_setup_interp_filters(MACROBLOCKD *xd,
   if (xd->mode_info_context) {
     MB_MODE_INFO *mbmi = &xd->mode_info_context->mbmi;
 
-    set_scale_factors(xd,
-                      mbmi->ref_frame[0] - 1,
-                      mbmi->ref_frame[1] - 1,
+    set_scale_factors(xd, mbmi->ref_frame[0] - 1, mbmi->ref_frame[1] - 1,
                       cm->active_ref_scale);
   }
 

diff --git a/vp9/common/vp9_reconinter.h b/vp9/common/vp9_reconinter.h
index 82dda599a661d56855077c68ea915d341372f031..c7219467015374c46b82b7887926a0abf0f6de09 100644 (file)
--- a/vp9/common/vp9_reconinter.h
+++ b/vp9/common/vp9_reconinter.h
@@ -109,14 +109,10 @@ static void setup_pre_planes(MACROBLOCKD *xd,
   }
 }
 
-static void set_scale_factors(MACROBLOCKD *xd,
-    int ref0, int ref1,
-    struct scale_factors scale_factor[MAX_REF_FRAMES]) {
-
-  xd->scale_factor[0] = scale_factor[ref0 >= 0 ? ref0 : 0];
-  xd->scale_factor[1] = scale_factor[ref1 >= 0 ? ref1 : 0];
-  xd->scale_factor_uv[0] = xd->scale_factor[0];
-  xd->scale_factor_uv[1] = xd->scale_factor[1];
+static void set_scale_factors(MACROBLOCKD *xd, int ref0, int ref1,
+                              struct scale_factors sf[MAX_REF_FRAMES]) {
+  xd->scale_factor[0] = xd->scale_factor_uv[0] = sf[ref0 >= 0 ? ref0 : 0];
+  xd->scale_factor[1] = xd->scale_factor_uv[1] = sf[ref1 >= 0 ? ref1 : 0];
 }
 
 void vp9_setup_scale_factors(VP9_COMMON *cm, int i);

diff --git a/vp9/decoder/vp9_decodemv.c b/vp9/decoder/vp9_decodemv.c
index 020e9c6cc73c530f2ca1114baccdb5bd1dc99793..3c5abd43b4417468b5cddbe6ca512de8f45541ca 100644 (file)
--- a/vp9/decoder/vp9_decodemv.c
+++ b/vp9/decoder/vp9_decodemv.c
@@ -413,7 +413,7 @@ static int read_inter_segment_id(VP9D_COMP *pbi, int mi_row, int mi_col,
   if (cm->temporal_update) {
     const vp9_prob pred_prob = vp9_get_pred_prob(cm, xd, PRED_SEG_ID);
     const int pred_flag = vp9_read(r, pred_prob);
-    vp9_set_pred_flag(xd, PRED_SEG_ID, pred_flag);
+    vp9_set_pred_flag(xd, bsize, PRED_SEG_ID, pred_flag);
     segment_id = pred_flag ? pred_segment_id
                            : read_segment_id(r, xd);
   } else {

diff --git a/vp9/decoder/vp9_decodframe.c b/vp9/decoder/vp9_decodframe.c
index 08acbf31ddb40954f83ba96927e77a6b1e527112..e578052d84525425a4dd8ff94e7456e3f0143329 100644 (file)
--- a/vp9/decoder/vp9_decodframe.c
+++ b/vp9/decoder/vp9_decodframe.c
@@ -14,15 +14,16 @@
 #include "vpx_mem/vpx_mem.h"
 #include "vpx_scale/vpx_scale.h"
 
-#include "vp9/common/vp9_extend.h"
-#include "vp9/common/vp9_modecont.h"
+#include "vp9/common/vp9_alloccommon.h"
 #include "vp9/common/vp9_common.h"
-#include "vp9/common/vp9_reconintra.h"
-#include "vp9/common/vp9_reconinter.h"
 #include "vp9/common/vp9_entropy.h"
-#include "vp9/common/vp9_alloccommon.h"
 #include "vp9/common/vp9_entropymode.h"
+#include "vp9/common/vp9_extend.h"
+#include "vp9/common/vp9_modecont.h"
+#include "vp9/common/vp9_pred_common.h"
 #include "vp9/common/vp9_quant_common.h"
+#include "vp9/common/vp9_reconintra.h"
+#include "vp9/common/vp9_reconinter.h"
 #include "vp9/common/vp9_seg_common.h"
 #include "vp9/common/vp9_tile_common.h"
 
@@ -133,6 +134,7 @@ static void decode_block_intra(int plane, int block, BLOCK_SIZE_TYPE bsize,
                                int ss_txfrm_size, void *arg) {
   MACROBLOCKD* const xd = arg;
   struct macroblockd_plane *pd = &xd->plane[plane];
+  MODE_INFO *const mi = xd->mode_info_context;
 
   const int raster_block = txfrm_block_to_raster_block(xd, bsize, plane,
                                                        block, ss_txfrm_size);
@@ -143,13 +145,12 @@ static void decode_block_intra(int plane, int block, BLOCK_SIZE_TYPE bsize,
   int b_mode;
   int plane_b_size;
   const int tx_ib = raster_block >> tx_size;
-  const int mode = plane == 0 ? xd->mode_info_context->mbmi.mode
-                              : xd->mode_info_context->mbmi.uv_mode;
-
+  const int mode = plane == 0 ? mi->mbmi.mode
+                              : mi->mbmi.uv_mode;
 
-  if (plane == 0 && xd->mode_info_context->mbmi.sb_type < BLOCK_SIZE_SB8X8) {
+  if (plane == 0 && mi->mbmi.sb_type < BLOCK_SIZE_SB8X8) {
     assert(bsize == BLOCK_SIZE_SB8X8);
-    b_mode = xd->mode_info_context->bmi[raster_block].as_mode.first;
+    b_mode = mi->bmi[raster_block].as_mode.first;
   } else {
     b_mode = mode;
   }
@@ -163,91 +164,66 @@ static void decode_block_intra(int plane, int block, BLOCK_SIZE_TYPE bsize,
                           dst, pd->dst.stride);
 
   // Early exit if there are no coefficients
-  if (xd->mode_info_context->mbmi.mb_skip_coeff)
+  if (mi->mbmi.mb_skip_coeff)
     return;
 
   decode_block(plane, block, bsize, ss_txfrm_size, arg);
 }
 
-static void decode_atom(VP9D_COMP *pbi, MACROBLOCKD *xd,
-                        int mi_row, int mi_col,
-                        vp9_reader *r, BLOCK_SIZE_TYPE bsize) {
-  MB_MODE_INFO *const mbmi = &xd->mode_info_context->mbmi;
-
-  assert(mbmi->ref_frame[0] != INTRA_FRAME);
-  vp9_setup_interp_filters(xd, mbmi->interp_filter, &pbi->common);
-
-  // prediction
-  vp9_build_inter_predictors_sb(xd, mi_row, mi_col, bsize);
+static int decode_tokens(VP9D_COMP *pbi, BLOCK_SIZE_TYPE bsize, vp9_reader *r) {
+  MACROBLOCKD *const xd = &pbi->mb;
 
-  if (mbmi->mb_skip_coeff) {
+  if (xd->mode_info_context->mbmi.mb_skip_coeff) {
     vp9_reset_sb_tokens_context(xd, bsize);
+    return -1;
   } else {
     if (xd->segmentation_enabled)
       mb_init_dequantizer(&pbi->common, xd);
 
-    if (!vp9_reader_has_error(r))
-      vp9_decode_tokens(pbi, r, bsize);
-
-    foreach_transformed_block(xd, bsize, decode_block, xd);
+    // TODO(dkovalev) if (!vp9_reader_has_error(r))
+    return vp9_decode_tokens(pbi, r, bsize);
   }
 }
 
-static void decode_sb_intra(VP9D_COMP *pbi, MACROBLOCKD *xd,
-                          int mi_row, int mi_col,
-                          vp9_reader *r, BLOCK_SIZE_TYPE bsize) {
+static void decode_atom(VP9D_COMP *pbi, int mi_row, int mi_col,
+                        vp9_reader *r) {
+  BLOCK_SIZE_TYPE bsize = BLOCK_SIZE_SB8X8;
+  MACROBLOCKD *const xd = &pbi->mb;
   MB_MODE_INFO *const mbmi = &xd->mode_info_context->mbmi;
-  if (mbmi->mb_skip_coeff) {
-    vp9_reset_sb_tokens_context(xd, bsize);
-  } else {
-    if (xd->segmentation_enabled)
-      mb_init_dequantizer(&pbi->common, xd);
 
-    if (!vp9_reader_has_error(r))
-      vp9_decode_tokens(pbi, r, bsize);
-  }
+  assert(mbmi->ref_frame[0] != INTRA_FRAME);
+  vp9_setup_interp_filters(xd, mbmi->interp_filter, &pbi->common);
+  vp9_build_inter_predictors_sb(xd, mi_row, mi_col, bsize);
+  if (decode_tokens(pbi, bsize, r) >= 0)
+    foreach_transformed_block(xd, bsize, decode_block, xd);
+}
 
+static void decode_sb_intra(VP9D_COMP *pbi, int mi_row, int mi_col,
+                            vp9_reader *r, BLOCK_SIZE_TYPE bsize) {
+  MACROBLOCKD *const xd = &pbi->mb;
+  decode_tokens(pbi, bsize, r);
   foreach_transformed_block(xd, bsize, decode_block_intra, xd);
 }
 
-
-static void decode_sb(VP9D_COMP *pbi, MACROBLOCKD *xd, int mi_row, int mi_col,
-                      vp9_reader *r, BLOCK_SIZE_TYPE bsize) {
-  const int bwl = mi_width_log2(bsize), bhl = mi_height_log2(bsize);
-  const int bw = 1 << bwl, bh = 1 << bhl;
-  int n, eobtotal;
-  VP9_COMMON *const pc = &pbi->common;
+static void decode_sb(VP9D_COMP *pbi, int mi_row, int mi_col, vp9_reader *r,
+                      BLOCK_SIZE_TYPE bsize) {
+  VP9_COMMON *const cm = &pbi->common;
+  MACROBLOCKD *const xd = &pbi->mb;
   MODE_INFO *const mi = xd->mode_info_context;
   MB_MODE_INFO *const mbmi = &mi->mbmi;
-  const int mis = pc->mode_info_stride;
+  int  eobtotal;
 
   assert(mbmi->sb_type == bsize);
   assert(mbmi->ref_frame[0] != INTRA_FRAME);
 
-  vp9_setup_interp_filters(xd, mbmi->interp_filter, pc);
-
-  // generate prediction
+  vp9_setup_interp_filters(xd, mbmi->interp_filter, cm);
   vp9_build_inter_predictors_sb(xd, mi_row, mi_col, bsize);
 
-  if (mbmi->mb_skip_coeff) {
-    vp9_reset_sb_tokens_context(xd, bsize);
-  } else {
-    // re-initialize macroblock dequantizer before detokenization
-    if (xd->segmentation_enabled)
-      mb_init_dequantizer(pc, xd);
-
-    // dequantization and idct
-    eobtotal = vp9_decode_tokens(pbi, r, bsize);
-    if (eobtotal == 0) {  // skip loopfilter
-      for (n = 0; n < bw * bh; n++) {
-        const int x_idx = n & (bw - 1), y_idx = n >> bwl;
-
-        if (mi_col + x_idx < pc->mi_cols && mi_row + y_idx < pc->mi_rows)
-          mi[y_idx * mis + x_idx].mbmi.mb_skip_coeff = 1;
-      }
-    } else {
-      foreach_transformed_block(xd, bsize, decode_block, xd);
-    }
+  eobtotal = decode_tokens(pbi, bsize, r);
+  if (eobtotal == 0) {
+    vp9_set_pred_flag(xd, bsize, PRED_MBSKIP, 1);  //  skip loopfilter
+  } else if (eobtotal > 0) {
+    foreach_transformed_block(xd, bsize, decode_block, xd);
   }
 }
 
@@ -316,18 +292,19 @@ static void decode_modes_b(VP9D_COMP *pbi, int mi_row, int mi_col,
   if (bsize < BLOCK_SIZE_SB8X8)
     if (xd->ab_index > 0)
       return;
+
   set_offsets(pbi, bsize, mi_row, mi_col);
   vp9_decode_mb_mode_mv(pbi, xd, mi_row, mi_col, r);
 
   if (xd->mode_info_context->mbmi.ref_frame[0] == INTRA_FRAME) {
-    decode_sb_intra(pbi, xd, mi_row, mi_col, r, (bsize < BLOCK_SIZE_SB8X8) ?
-                                     BLOCK_SIZE_SB8X8 : bsize);
+    decode_sb_intra(pbi, mi_row, mi_col, r,
+                    (bsize < BLOCK_SIZE_SB8X8) ? BLOCK_SIZE_SB8X8 : bsize);
   } else {
     set_refs(pbi, mi_row, mi_col);
     if (bsize < BLOCK_SIZE_SB8X8)
-      decode_atom(pbi, xd, mi_row, mi_col, r, BLOCK_SIZE_SB8X8);
+      decode_atom(pbi, mi_row, mi_col, r);
     else
-      decode_sb(pbi, xd, mi_row, mi_col, r, bsize);
+      decode_sb(pbi, mi_row, mi_col, r, bsize);
   }
   xd->corrupted |= vp9_reader_has_error(r);
 }

diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index 168b460127261526fdcb2bdbdf5f4dbd86d433ee..f46d2e3ce42d250a24b5ff741239260b1f5f20ca 100644 (file)
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -2053,7 +2053,6 @@ static void encode_superblock(VP9_COMP *cpi, TOKENEXTRA **t, int output_enabled,
   VP9_COMMON * const cm = &cpi->common;
   MACROBLOCK * const x = &cpi->mb;
   MACROBLOCKD * const xd = &x->e_mbd;
-  int n;
   MODE_INFO *mi = xd->mode_info_context;
   MB_MODE_INFO *mbmi = &mi->mbmi;
   unsigned int segment_id = mbmi->segment_id;
@@ -2145,11 +2144,7 @@ static void encode_superblock(VP9_COMP *cpi, TOKENEXTRA **t, int output_enabled,
 
   // copy skip flag on all mb_mode_info contexts in this SB
   // if this was a skip at this txfm size
-  for (n = 1; n < bw * bh; n++) {
-    const int x_idx = n & (bw - 1), y_idx = n >> bwl;
-    if (mi_col + x_idx < cm->mi_cols && mi_row + y_idx < cm->mi_rows)
-      mi[x_idx + y_idx * mis].mbmi.mb_skip_coeff = mi->mbmi.mb_skip_coeff;
-  }
+  vp9_set_pred_flag(xd, bsize, PRED_MBSKIP, mi->mbmi.mb_skip_coeff);
 
   if (output_enabled) {
     if (cm->txfm_mode == TX_MODE_SELECT && mbmi->sb_type >= BLOCK_SIZE_SB8X8

diff --git a/vp9/encoder/vp9_segmentation.c b/vp9/encoder/vp9_segmentation.c
index eaa1d534ab58238d2028e4d15b1ac22b13c04af3..51bd362b06cabe5c9e6d1aace0e0ad11e3a4058d 100644 (file)
--- a/vp9/encoder/vp9_segmentation.c
+++ b/vp9/encoder/vp9_segmentation.c
@@ -136,16 +136,16 @@ static void count_segs(VP9_COMP *cpi, MODE_INFO *mi,
 
   // Temporal prediction not allowed on key frames
   if (cm->frame_type != KEY_FRAME) {
+    const BLOCK_SIZE_TYPE bsize = mi->mbmi.sb_type;
     // Test to see if the segment id matches the predicted value.
     const int pred_segment_id = vp9_get_segment_id(cm, cm->last_frame_seg_map,
-                                                   mi->mbmi.sb_type,
-                                                   mi_row, mi_col);
+                                                   bsize, mi_row, mi_col);
     const int pred_flag = pred_segment_id == segment_id;
     const int pred_context = vp9_get_pred_context(cm, xd, PRED_SEG_ID);
 
     // Store the prediction status for this mb and update counts
     // as appropriate
-    vp9_set_pred_flag(xd, PRED_SEG_ID, pred_flag);
+    vp9_set_pred_flag(xd, bsize, PRED_SEG_ID, pred_flag);
     temporal_predictor_count[pred_context][pred_flag]++;
 
     if (!pred_flag)