JNT_COMP: change COMPOUND_AVERAGE in cdf

Remove COMPOUND_AVERAGE from compound_type_cdfs since it is now grouped
to compound_idx. However, COMPOUND_AVERAGE is still used elsewhere.

Change-Id: Ie0d460aabf9252e80eb4130cfef9aaf0efc3969d

av1/common/entropymode.c

@@ -760,6 +760,27 @@ static const aom_cdf_prob
       { AOM_CDF8(3456, 9067, 14069, 16907, 18817, 21214, 23139) }
     };
 
+#if CONFIG_JNT_COMP
+static const aom_cdf_prob
+    default_compound_type_cdf[BLOCK_SIZES_ALL][CDF_SIZE(COMPOUND_TYPES - 1)] = {
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) }, { AOM_CDF2(16384) },
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) }, { AOM_CDF2(16384) },
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) }, { AOM_CDF2(16384) },
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) }, { AOM_CDF2(16384) },
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) }, { AOM_CDF2(16384) },
+      { AOM_CDF2(16384) },
+#if CONFIG_EXT_PARTITION
+      { AOM_CDF2(16384) },  // 255, 1
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) },
+#endif  // CONFIG_EXT_PARTITION
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) }, { AOM_CDF2(16384) },
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) },  // 208, 1
+      { AOM_CDF2(16384) },
+#if CONFIG_EXT_PARTITION
+      { AOM_CDF2(16384) }, { AOM_CDF2(16384) },
+#endif
+    };
+#else
 static const aom_cdf_prob
     default_compound_type_cdf[BLOCK_SIZES_ALL][CDF_SIZE(COMPOUND_TYPES)] = {
       { AOM_CDF3(16384, 24576) }, { AOM_CDF3(16384, 24576) },
@@ -782,6 +803,7 @@ static const aom_cdf_prob
       { AOM_CDF3(26624, 26648) }, { AOM_CDF3(26624, 26648) },
 #endif
     };
+#endif  // CONFIG_JNT_COMP
 
 static const aom_cdf_prob default_interintra_cdf[BLOCK_SIZE_GROUPS][CDF_SIZE(
     2)] = { { AOM_CDF2(128 * 128) },
@@ -1024,9 +1046,15 @@ const aom_tree_index av1_inter_compound_mode_tree
 };
 
 
+#if CONFIG_JNT_COMP
+const aom_tree_index av1_compound_type_tree[TREE_SIZE(COMPOUND_TYPES - 1)] = {
+  -COMPOUND_WEDGE, -COMPOUND_SEG
+};
+#else
 const aom_tree_index av1_compound_type_tree[TREE_SIZE(COMPOUND_TYPES)] = {
   -COMPOUND_AVERAGE, 2, -COMPOUND_WEDGE, -COMPOUND_SEG
 };
+#endif  // CONFIG_JNT_COMP
 /* clang-format on */
 
 static const aom_cdf_prob default_intra_inter_cdf[INTRA_INTER_CONTEXTS]

av1/common/entropymode.h

@@ -201,7 +201,13 @@ typedef struct frame_contexts {
 
   aom_cdf_prob inter_compound_mode_cdf[INTER_MODE_CONTEXTS]
                                       [CDF_SIZE(INTER_COMPOUND_MODES)];
+#if CONFIG_JNT_COMP
+  aom_prob compound_type_prob[BLOCK_SIZES_ALL][COMPOUND_TYPES - 2];
+  aom_cdf_prob compound_type_cdf[BLOCK_SIZES_ALL][CDF_SIZE(COMPOUND_TYPES - 1)];
+#else
+  aom_prob compound_type_prob[BLOCK_SIZES_ALL][COMPOUND_TYPES - 1];
   aom_cdf_prob compound_type_cdf[BLOCK_SIZES_ALL][CDF_SIZE(COMPOUND_TYPES)];
+#endif  // CONFIG_JNT_COMP
   aom_cdf_prob interintra_cdf[BLOCK_SIZE_GROUPS][CDF_SIZE(2)];
   aom_cdf_prob wedge_interintra_cdf[BLOCK_SIZES_ALL][CDF_SIZE(2)];
   aom_cdf_prob interintra_mode_cdf[BLOCK_SIZE_GROUPS]
@@ -373,7 +379,11 @@ typedef struct FRAME_COUNTS {
   unsigned int interintra[BLOCK_SIZE_GROUPS][2];
   unsigned int interintra_mode[BLOCK_SIZE_GROUPS][INTERINTRA_MODES];
   unsigned int wedge_interintra[BLOCK_SIZES_ALL][2];
+#if CONFIG_JNT_COMP
+  unsigned int compound_interinter[BLOCK_SIZES_ALL][COMPOUND_TYPES - 1];
+#else
   unsigned int compound_interinter[BLOCK_SIZES_ALL][COMPOUND_TYPES];
+#endif  // CONFIG_JNT_COMP
 #if CONFIG_EXT_WARPED_MOTION
   unsigned int motion_mode[MOTION_MODE_CTX][BLOCK_SIZES_ALL][MOTION_MODES];
 #else
@@ -487,7 +497,12 @@ extern const aom_tree_index
     av1_interintra_mode_tree[TREE_SIZE(INTERINTRA_MODES)];
 extern const aom_tree_index
     av1_inter_compound_mode_tree[TREE_SIZE(INTER_COMPOUND_MODES)];
+#if CONFIG_JNT_COMP
+extern const aom_tree_index
+    av1_compound_type_tree[TREE_SIZE(COMPOUND_TYPES - 1)];
+#else
 extern const aom_tree_index av1_compound_type_tree[TREE_SIZE(COMPOUND_TYPES)];
+#endif  // CONFIG_JNT_COMP
 extern const aom_tree_index av1_motion_mode_tree[TREE_SIZE(MOTION_MODES)];
 
 void av1_setup_frame_contexts(struct AV1Common *cm);

av1/decoder/decodemv.c

@@ -2214,13 +2214,13 @@ static void read_inter_block_mode_info(AV1Decoder *const pbi,
 #if CONFIG_EXT_SKIP
           && !mbmi->skip_mode
 #endif  // CONFIG_EXT_SKIP
-          && mbmi->comp_group_idx) {
+          ) {
         if (is_any_masked_compound_used(bsize)) {
           if (cm->allow_masked_compound) {
             if (is_interinter_compound_used(COMPOUND_WEDGE, bsize))
               mbmi->interinter_compound_type =
-                  aom_read_symbol(r, ec_ctx->compound_type_cdf[bsize],
-                                  COMPOUND_TYPES, ACCT_STR);
+                  1 + aom_read_symbol(r, ec_ctx->compound_type_cdf[bsize],
+                                      COMPOUND_TYPES - 1, ACCT_STR);
             else
               mbmi->interinter_compound_type = COMPOUND_SEG;
 
@@ -2238,8 +2238,8 @@ static void read_inter_block_mode_info(AV1Decoder *const pbi,
         }
 
         if (xd->counts)
-          xd->counts
-              ->compound_interinter[bsize][mbmi->interinter_compound_type]++;
+          xd->counts->compound_interinter[bsize]
+                                         [mbmi->interinter_compound_type - 1]++;
       }
     }
   }

av1/encoder/bitstream.c

@@ -70,7 +70,11 @@ static INLINE void write_uniform(aom_writer *w, int n, int v) {
 }
 
 static struct av1_token interintra_mode_encodings[INTERINTRA_MODES];
+#if CONFIG_JNT_COMP
+static struct av1_token compound_type_encodings[COMPOUND_TYPES - 1];
+#else
 static struct av1_token compound_type_encodings[COMPOUND_TYPES];
+#endif  // CONFIG_JNT_COMP
 #if CONFIG_LOOP_RESTORATION
 static void loop_restoration_write_sb_coeffs(const AV1_COMMON *const cm,
                                              MACROBLOCKD *xd,
@@ -1469,11 +1473,11 @@ static void pack_inter_mode_mvs(AV1_COMP *cpi, const int mi_row,
         if (cpi->common.reference_mode != SINGLE_REFERENCE &&
             is_inter_compound_mode(mbmi->mode) &&
             mbmi->motion_mode == SIMPLE_TRANSLATION &&
-            is_any_masked_compound_used(bsize) && cm->allow_masked_compound &&
-            mbmi->comp_group_idx) {
+            is_any_masked_compound_used(bsize) && cm->allow_masked_compound) {
           if (is_interinter_compound_used(COMPOUND_WEDGE, bsize))
-            aom_write_symbol(w, mbmi->interinter_compound_type,
-                             ec_ctx->compound_type_cdf[bsize], COMPOUND_TYPES);
+            aom_write_symbol(w, mbmi->interinter_compound_type - 1,
+                             ec_ctx->compound_type_cdf[bsize],
+                             COMPOUND_TYPES - 1);
 
           if (is_interinter_compound_used(COMPOUND_WEDGE, bsize) &&
               mbmi->interinter_compound_type == COMPOUND_WEDGE) {

av1/encoder/block.h

@@ -271,7 +271,11 @@ struct macroblock {
   int drl_mode_cost0[DRL_MODE_CONTEXTS][2];
 
   int inter_compound_mode_cost[INTER_MODE_CONTEXTS][INTER_COMPOUND_MODES];
+#if CONFIG_JNT_COMP
+  int compound_type_cost[BLOCK_SIZES_ALL][COMPOUND_TYPES - 1];
+#else
   int compound_type_cost[BLOCK_SIZES_ALL][COMPOUND_TYPES];
+#endif  // CONFIG_JNT_COMP
   int interintra_cost[BLOCK_SIZE_GROUPS][2];
   int wedge_interintra_cost[BLOCK_SIZES_ALL][2];
   int interintra_mode_cost[BLOCK_SIZE_GROUPS][INTERINTRA_MODES];
@@ -347,6 +351,7 @@ struct macroblock {
 #endif  // CONFIG_DIST_8X8
 #if CONFIG_JNT_COMP
   int comp_idx_cost[COMP_INDEX_CONTEXTS][2];
+  int comp_group_idx_cost[COMP_GROUP_IDX_CONTEXTS][2];
 #endif  // CONFIG_JNT_COMP
 };
 

av1/encoder/encodeframe.c

@@ -1271,14 +1271,24 @@ static void update_stats(const AV1_COMMON *const cm, TileDataEnc *tile_data,
                          mbmi->compound_idx, 2);
           }
         }
-#endif  // CONFIG_JNT_COMP
 
         if (cm->reference_mode != SINGLE_REFERENCE &&
-            is_inter_compound_mode(mbmi->mode)
-#if CONFIG_JNT_COMP
-            && mbmi->comp_group_idx
-#endif  // CONFIG_JNT_COMP
-            && mbmi->motion_mode == SIMPLE_TRANSLATION) {
+            is_inter_compound_mode(mbmi->mode) && mbmi->comp_group_idx &&
+            mbmi->motion_mode == SIMPLE_TRANSLATION) {
+          if (is_interinter_compound_used(COMPOUND_WEDGE, bsize)) {
+            counts->compound_interinter[bsize]
+                                       [mbmi->interinter_compound_type - 1]++;
+            if (allow_update_cdf)
+              update_cdf(fc->compound_type_cdf[bsize],
+                         mbmi->interinter_compound_type - 1,
+                         COMPOUND_TYPES - 1);
+          }
+        }
+#else   // CONFIG_JNT_COMP
+
+        if (cm->reference_mode != SINGLE_REFERENCE &&
+            is_inter_compound_mode(mbmi->mode) &&
+            mbmi->motion_mode == SIMPLE_TRANSLATION) {
           if (is_interinter_compound_used(COMPOUND_WEDGE, bsize)) {
             counts
                 ->compound_interinter[bsize][mbmi->interinter_compound_type]++;
@@ -1287,6 +1297,7 @@ static void update_stats(const AV1_COMMON *const cm, TileDataEnc *tile_data,
                          mbmi->interinter_compound_type, COMPOUND_TYPES);
           }
         }
+#endif  // CONFIG_JNT_COMP
       }
     }
 

av1/encoder/rd.c

@@ -284,6 +284,10 @@ void av1_fill_mode_rates(AV1_COMMON *const cm, MACROBLOCK *x,
       av1_cost_tokens_from_cdf(x->comp_idx_cost[i], fc->compound_index_cdf[i],
                                NULL);
     }
+    for (i = 0; i < COMP_GROUP_IDX_CONTEXTS; ++i) {
+      av1_cost_tokens_from_cdf(x->comp_group_idx_cost[i],
+                               fc->comp_group_idx_cdf[i], NULL);
+    }
 #endif  // CONFIG_JNT_COMP
   }
 }

av1/encoder/rdopt.c

@@ -5661,6 +5661,7 @@ static int cost_mv_ref(const MACROBLOCK *const x, PREDICTION_MODE mode,
   }
 }
 
+#if !CONFIG_JNT_COMP
 static int get_interinter_compound_type_bits(BLOCK_SIZE bsize,
                                              COMPOUND_TYPE comp_type) {
   (void)bsize;
@@ -5671,6 +5672,7 @@ static int get_interinter_compound_type_bits(BLOCK_SIZE bsize,
     default: assert(0); return 0;
   }
 }
+#endif
 
 typedef struct {
   int eobs;
@@ -8139,8 +8141,14 @@ static int64_t handle_inter_mode(
 
 #if CONFIG_JNT_COMP
   if (is_comp_pred) {
-    const int comp_index_ctx = get_comp_index_context(cm, xd);
-    rd_stats->rate += x->comp_idx_cost[comp_index_ctx][mbmi->compound_idx];
+    if (mbmi->compound_idx == 0) {
+      mbmi->comp_group_idx = 0;
+      const int comp_group_idx_ctx = get_comp_group_idx_context(xd);
+      rd_stats->rate += x->comp_group_idx_cost[comp_group_idx_ctx][0];
+
+      const int comp_index_ctx = get_comp_index_context(cm, xd);
+      rd_stats->rate += x->comp_idx_cost[comp_index_ctx][0];
+    }
   }
 #endif  // CONFIG_JNT_COMP
 
@@ -8307,6 +8315,28 @@ static int64_t handle_inter_mode(
       tmp_rate_mv = rate_mv;
       best_rd_cur = INT64_MAX;
       mbmi->interinter_compound_type = cur_type;
+#if CONFIG_JNT_COMP
+      const int ctx_comp_group_idx = get_comp_group_idx_context(xd);
+      if (cur_type == COMPOUND_AVERAGE) {
+        mbmi->comp_group_idx = 0;
+        rs2 = x->comp_group_idx_cost[ctx_comp_group_idx][0];
+
+        const int comp_index_ctx = get_comp_index_context(cm, xd);
+        rs2 += x->comp_idx_cost[comp_index_ctx][1];
+      } else {
+        mbmi->comp_group_idx = 1;
+        rs2 = x->comp_group_idx_cost[ctx_comp_group_idx][1];
+
+        int masked_type_cost = 0;
+        if (masked_compound_used) {
+          if (is_interinter_compound_used(COMPOUND_WEDGE, bsize))
+            masked_type_cost +=
+                x->compound_type_cost[bsize]
+                                     [mbmi->interinter_compound_type - 1];
+        }
+        rs2 += masked_type_cost;
+      }
+#else
       int masked_type_cost = 0;
       if (masked_compound_used) {
         if (!is_interinter_compound_used(COMPOUND_WEDGE, bsize))
@@ -8318,6 +8348,7 @@ static int64_t handle_inter_mode(
       rs2 = av1_cost_literal(get_interinter_compound_type_bits(
                 bsize, mbmi->interinter_compound_type)) +
             masked_type_cost;
+#endif  // CONFIG_JNT_COMP
 
       switch (cur_type) {
         case COMPOUND_AVERAGE: