JNT_COMP: 4. add context and entropy read/write

Change-Id: I0e6f7ab981e31f7120105515f6204568b6dc82d3

av1/common/entropymode.c

@@ -2354,6 +2354,12 @@ static const aom_cdf_prob default_skip_cdfs[SKIP_CONTEXTS][CDF_SIZE(2)] = {
 static const aom_prob default_skip_probs[SKIP_CONTEXTS] = { 192, 128, 64 };
 #endif  // CONFIG_NEW_MULTISYMBOL
 
+#if CONFIG_JNT_COMP
+static const aom_prob default_compound_idx_probs[COMP_INDEX_CONTEXTS] = {
+  192, 128, 64, 192, 128, 64, 192, 128, 64,
+};
+#endif  // CONFIG_JNT_COMP
+
 #if CONFIG_LGT_FROM_PRED
 static const aom_prob default_intra_lgt_prob[LGT_SIZES][INTRA_MODES] = {
   { 255, 208, 208, 180, 230, 208, 194, 214, 220, 255,
@@ -6224,6 +6230,9 @@ static void init_mode_probs(FRAME_CONTEXT *fc) {
 #if CONFIG_NEW_MULTISYMBOL
   av1_copy(fc->txfm_partition_cdf, default_txfm_partition_cdf);
 #endif
+#if CONFIG_JNT_COMP
+  av1_copy(fc->compound_index_probs, default_compound_idx_probs);
+#endif  // CONFIG_JNT_COMP
   av1_copy(fc->newmv_prob, default_newmv_prob);
   av1_copy(fc->zeromv_prob, default_zeromv_prob);
   av1_copy(fc->refmv_prob, default_refmv_prob);
@@ -6455,6 +6464,12 @@ void av1_adapt_inter_frame_probs(AV1_COMMON *cm) {
     }
   }
 #endif  // CONFIG_COMPOUND_SEGMENT || CONFIG_WEDGE
+
+#if CONFIG_JNT_COMP
+  for (i = 0; i < COMP_INDEX_CONTEXTS; ++i)
+    fc->compound_index_probs[i] = av1_mode_mv_merge_probs(
+        pre_fc->compound_index_probs[i], counts->compound_index[i]);
+#endif  // CONFIG_JNT_COMP
 }
 
 void av1_adapt_intra_frame_probs(AV1_COMMON *cm) {

av1/common/entropymode.h

@@ -292,6 +292,9 @@ typedef struct frame_contexts {
 #if CONFIG_NEW_MULTISYMBOL
   aom_cdf_prob txfm_partition_cdf[TXFM_PARTITION_CONTEXTS][CDF_SIZE(2)];
 #endif
+#if CONFIG_JNT_COMP
+  aom_prob compound_index_probs[COMP_INDEX_CONTEXTS];
+#endif  // CONFIG_JNT_COMP
 #if CONFIG_NEW_MULTISYMBOL
   aom_cdf_prob skip_cdfs[SKIP_CONTEXTS][CDF_SIZE(2)];
   aom_cdf_prob intra_inter_cdf[INTRA_INTER_CONTEXTS][CDF_SIZE(2)];
@@ -477,6 +480,9 @@ typedef struct FRAME_COUNTS {
   unsigned int txfm_partition[TXFM_PARTITION_CONTEXTS][2];
   unsigned int skip[SKIP_CONTEXTS][2];
   nmv_context_counts mv[NMV_CONTEXTS];
+#if CONFIG_JNT_COMP
+  unsigned int compound_index[COMP_INDEX_CONTEXTS][2];
+#endif  // CONFIG_JNT_COMP
 #if CONFIG_INTRABC
   unsigned int intrabc[2];
   nmv_context_counts dv;

av1/common/enums.h

@@ -640,6 +640,10 @@ typedef enum ATTRIBUTE_PACKED {
 
 #define SKIP_CONTEXTS 3
 
+#if CONFIG_JNT_COMP
+#define COMP_INDEX_CONTEXTS 9
+#endif  // CONFIG_JNT_COMP
+
 #define NMV_CONTEXTS 3
 
 #define NEWMV_MODE_CONTEXTS 7

av1/common/pred_common.h

@@ -54,6 +54,53 @@ static INLINE aom_prob av1_get_pred_prob_seg_id(
   return segp->pred_probs[av1_get_pred_context_seg_id(xd)];
 }
 
+#if CONFIG_JNT_COMP
+static INLINE int get_comp_index_context(const AV1_COMMON *cm,
+                                         const MACROBLOCKD *xd) {
+  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
+  int bck_idx = cm->frame_refs[mbmi->ref_frame[0] - LAST_FRAME].idx;
+  int fwd_idx = cm->frame_refs[mbmi->ref_frame[1] - LAST_FRAME].idx;
+  int bck_frame_index = 0, fwd_frame_index = 0;
+  int cur_frame_index = cm->cur_frame->cur_frame_offset;
+
+  if (bck_idx >= 0)
+    bck_frame_index = cm->buffer_pool->frame_bufs[bck_idx].cur_frame_offset;
+
+  if (fwd_idx >= 0)
+    fwd_frame_index = cm->buffer_pool->frame_bufs[fwd_idx].cur_frame_offset;
+  int fwd = abs(fwd_frame_index - cur_frame_index);
+  int bck = abs(cur_frame_index - bck_frame_index);
+
+  const MODE_INFO *const above_mi = xd->above_mi;
+  const MODE_INFO *const left_mi = xd->left_mi;
+
+  int above_ctx = 0, left_ctx = 0;
+  int offset = (fwd > bck) ? 0 : 1;
+
+  if (fwd < (bck >> 1) + bck && bck < (fwd >> 1) + fwd) {
+    offset = 2;
+  }
+
+  if (above_mi) {
+    const MB_MODE_INFO *above_mbmi = &above_mi->mbmi;
+    if (has_second_ref(above_mbmi))
+      above_ctx = above_mbmi->compound_idx;
+    else if (above_mbmi->ref_frame[0] == ALTREF_FRAME)
+      above_ctx = 1;
+  }
+
+  if (left_mi) {
+    const MB_MODE_INFO *left_mbmi = &left_mi->mbmi;
+    if (has_second_ref(left_mbmi))
+      left_ctx = left_mbmi->compound_idx;
+    else if (left_mbmi->ref_frame[0] == ALTREF_FRAME)
+      left_ctx = 1;
+  }
+
+  return above_ctx + left_ctx + 3 * offset;
+}
+#endif  // CONFIG_JNT_COMP
+
 #if CONFIG_NEW_MULTISYMBOL
 static INLINE aom_cdf_prob *av1_get_pred_cdf_seg_id(
     struct segmentation_probs *segp, const MACROBLOCKD *xd) {

av1/decoder/decodeframe.c

@@ -3531,6 +3531,11 @@ static int read_compressed_header(AV1Decoder *pbi, const uint8_t *data,
       av1_diff_update_prob(&r, &fc->txfm_partition_prob[i], ACCT_STR);
   for (int i = 0; i < SKIP_CONTEXTS; ++i)
     av1_diff_update_prob(&r, &fc->skip_probs[i], ACCT_STR);
+
+#if CONFIG_JNT_COMP
+  for (int i = 0; i < COMP_INDEX_CONTEXTS; ++i)
+    av1_diff_update_prob(&r, &fc->compound_index_probs[i], ACCT_STR);
+#endif  // CONFIG_JNT_COMP
 #endif
 
   if (!frame_is_intra_only(cm)) {

av1/decoder/decodemv.c

@@ -2221,6 +2221,16 @@ static void read_inter_block_mode_info(AV1Decoder *const pbi,
   read_ref_frames(cm, xd, r, mbmi->segment_id, mbmi->ref_frame);
   is_compound = has_second_ref(mbmi);
 
+#if CONFIG_JNT_COMP
+  if (is_compound) {
+    const int comp_index_ctx = get_comp_index_context(cm, xd);
+    mbmi->compound_idx =
+        aom_read(r, ec_ctx->compound_index_probs[comp_index_ctx], ACCT_STR);
+    if (xd->counts)
+      ++xd->counts->compound_index[comp_index_ctx][mbmi->compound_idx];
+  }
+#endif  // CONFIG_JNT_COMP
+
 #if CONFIG_EXT_COMP_REFS
 #if !USE_UNI_COMP_REFS
   // NOTE: uni-directional comp refs disabled

av1/encoder/bitstream.c

@@ -1605,6 +1605,14 @@ static void pack_inter_mode_mvs(AV1_COMP *cpi, const int mi_row,
     int16_t mode_ctx;
     write_ref_frames(cm, xd, w);
 
+#if CONFIG_JNT_COMP
+    if (has_second_ref(mbmi)) {
+      const int comp_index_ctx = get_comp_index_context(cm, xd);
+      aom_write(w, mbmi->compound_idx,
+                ec_ctx->compound_index_probs[comp_index_ctx]);
+    }
+#endif  // CONFIG_JNT_COMP
+
 #if CONFIG_COMPOUND_SINGLEREF
     if (!segfeature_active(seg, segment_id, SEG_LVL_REF_FRAME)) {
       // NOTE: Handle single ref comp mode
@@ -4570,6 +4578,11 @@ static uint32_t write_compressed_header(AV1_COMP *cpi, uint8_t *data) {
 
 #if !CONFIG_NEW_MULTISYMBOL
   update_skip_probs(cm, header_bc, counts);
+#if CONFIG_JNT_COMP
+  for (int k = 0; k < COMP_INDEX_CONTEXTS; ++k)
+    av1_cond_prob_diff_update(header_bc, &cm->fc->compound_index_probs[k],
+                              counts->compound_index[k], probwt);
+#endif  // CONFIG_JNT_COMP
 #endif
 
   if (!frame_is_intra_only(cm)) {

av1/encoder/encodeframe.c

@@ -590,6 +590,13 @@ static void update_state(const AV1_COMP *const cpi, ThreadData *td,
   const int x_mis = AOMMIN(bw, cm->mi_cols - mi_col);
   const int y_mis = AOMMIN(bh, cm->mi_rows - mi_row);
   av1_copy_frame_mvs(cm, mi, mi_row, mi_col, x_mis, y_mis);
+
+#if CONFIG_JNT_COMP
+  if (has_second_ref(mbmi)) {
+    const int comp_index_ctx = get_comp_index_context(cm, xd);
+    ++td->counts->compound_index[comp_index_ctx][mbmi->compound_idx];
+  }
+#endif  // CONFIG_JNT_COMP
 }
 
 #if CONFIG_MOTION_VAR && NC_MODE_INFO
@@ -3443,7 +3450,7 @@ static void init_encode_frame_mb_context(AV1_COMP *cpi) {
   av1_setup_block_planes(xd, cm->subsampling_x, cm->subsampling_y);
 }
 
-#if !CONFIG_REF_ADAPT
+#if !CONFIG_REF_ADAPT && !CONFIG_JNT_COMP
 static int check_dual_ref_flags(AV1_COMP *cpi) {
   const int ref_flags = cpi->ref_frame_flags;
 
@@ -4162,12 +4169,14 @@ void av1_encode_frame(AV1_COMP *cpi) {
     if (is_alt_ref || !cpi->allow_comp_inter_inter)
 #endif  // CONFIG_BGSPRITE
       cm->reference_mode = SINGLE_REFERENCE;
+#if !CONFIG_JNT_COMP
     else if (mode_thrs[COMPOUND_REFERENCE] > mode_thrs[SINGLE_REFERENCE] &&
              mode_thrs[COMPOUND_REFERENCE] > mode_thrs[REFERENCE_MODE_SELECT] &&
              check_dual_ref_flags(cpi) && cpi->static_mb_pct == 100)
       cm->reference_mode = COMPOUND_REFERENCE;
     else if (mode_thrs[SINGLE_REFERENCE] > mode_thrs[REFERENCE_MODE_SELECT])
       cm->reference_mode = SINGLE_REFERENCE;
+#endif  // CONFIG_JNT_COMP
     else
       cm->reference_mode = REFERENCE_MODE_SELECT;
 #endif  // CONFIG_REF_ADAPT

av1/encoder/rdopt.c

@@ -8462,6 +8462,14 @@ static int64_t handle_inter_mode(const AV1_COMP *const cpi, MACROBLOCK *x,
   }
 #endif  // CONFIG_COMPOUND_SINGLEREF
 
+#if CONFIG_JNT_COMP
+  if (is_comp_pred) {
+    const int comp_index_ctx = get_comp_index_context(cm, xd);
+    rd_stats->rate += av1_cost_bit(cm->fc->compound_index_probs[comp_index_ctx],
+                                   mbmi->compound_idx);
+  }
+#endif  // CONFIG_JNT_COMP
+
   if (this_mode == NEAREST_NEARESTMV) {
     if (mbmi_ext->ref_mv_count[ref_frame_type] > 0) {
       cur_mv[0] = mbmi_ext->ref_mv_stack[ref_frame_type][0].this_mv;