cb4x4: Don't assume TX units are half MI_SIZE in skip context indexing.

Currently the "transform units" used to measure transform sizes
in some parts of the code are based on the smallest defined
transform size.

cb4x4 currently defines a 2x2 transform size, even when chroma_2x2
is not enabled, which means that the scale of the transform units
was always double that of MODEINFO units.

Several areas of the code were hard-coding this assumption instead
of converting from one to the other using appropriate constants.

Change-Id: Ibc55671aa5bc3ad272cb8a036f9c4f9621df85ab

av1/common/alloccommon.c

@@ -155,7 +155,8 @@ int av1_alloc_context_buffers(AV1_COMMON *cm, int width, int height) {
     for (i = 0; i < MAX_MB_PLANE; i++) {
       aom_free(cm->above_context[i]);
       cm->above_context[i] = (ENTROPY_CONTEXT *)aom_calloc(
-          2 * aligned_mi_cols, sizeof(*cm->above_context[0]));
+          aligned_mi_cols << (MI_SIZE_LOG2 - tx_size_wide_log2[0]),
+          sizeof(*cm->above_context[0]));
       if (!cm->above_context[i]) goto fail;
     }
 

av1/common/enums.h

@@ -51,7 +51,6 @@ extern "C" {
 
 // Mask to extract MI offset within max MIB
 #define MAX_MIB_MASK (MAX_MIB_SIZE - 1)
-#define MAX_MIB_MASK_2 (MAX_MIB_SIZE * 2 - 1)
 
 // Maximum number of tile rows and tile columns
 #if CONFIG_EXT_TILE

av1/common/onyxc_int.h

@@ -570,8 +570,9 @@ static INLINE void set_skip_context(MACROBLOCKD *xd, int mi_row, int mi_col) {
       if (pd->subsampling_x && (mi_col & 0x01)) mi_col -= 1;
     }
 #endif
-    int above_idx = mi_col * 2;
-    int left_idx = (mi_row * 2) & MAX_MIB_MASK_2;
+    int above_idx = mi_col << (MI_SIZE_LOG2 - tx_size_wide_log2[0]);
+    int left_idx = (mi_row & MAX_MIB_MASK)
+                   << (MI_SIZE_LOG2 - tx_size_high_log2[0]);
     pd->above_context = &xd->above_context[i][above_idx >> pd->subsampling_x];
     pd->left_context = &xd->left_context[i][left_idx >> pd->subsampling_y];
   }
@@ -856,8 +857,8 @@ static INLINE void av1_zero_above_context(AV1_COMMON *const cm,
   const int width = mi_col_end - mi_col_start;
   const int aligned_width = ALIGN_POWER_OF_TWO(width, cm->mib_size_log2);
 
-  const int offset_y = 2 * mi_col_start;
-  const int width_y = 2 * aligned_width;
+  const int offset_y = mi_col_start << (MI_SIZE_LOG2 - tx_size_wide_log2[0]);
+  const int width_y = width << (MI_SIZE_LOG2 - tx_size_wide_log2[0]);
   const int offset_uv = offset_y >> cm->subsampling_x;
   const int width_uv = width_y >> cm->subsampling_x;
 

av1/encoder/encodeframe.c

@@ -2332,12 +2332,15 @@ static void restore_context(MACROBLOCK *x,
   int mi_width = mi_size_wide[bsize];
   int mi_height = mi_size_high[bsize];
   for (p = 0; p < MAX_MB_PLANE; p++) {
-    memcpy(xd->above_context[p] + ((mi_col * 2) >> xd->plane[p].subsampling_x),
+    int tx_col;
+    int tx_row;
+    tx_col = mi_col << (MI_SIZE_LOG2 - tx_size_wide_log2[0]);
+    tx_row = (mi_row & MAX_MIB_MASK) << (MI_SIZE_LOG2 - tx_size_high_log2[0]);
+    memcpy(xd->above_context[p] + (tx_col >> xd->plane[p].subsampling_x),
            ctx->a + num_4x4_blocks_wide * p,
            (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_wide) >>
                xd->plane[p].subsampling_x);
-    memcpy(xd->left_context[p] +
-               ((mi_row & MAX_MIB_MASK) * 2 >> xd->plane[p].subsampling_y),
+    memcpy(xd->left_context[p] + (tx_row >> xd->plane[p].subsampling_y),
            ctx->l + num_4x4_blocks_high * p,
            (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_high) >>
                xd->plane[p].subsampling_y);
@@ -2376,13 +2379,16 @@ static void save_context(const MACROBLOCK *x, RD_SEARCH_MACROBLOCK_CONTEXT *ctx,
 
   // buffer the above/left context information of the block in search.
   for (p = 0; p < MAX_MB_PLANE; ++p) {
+    int tx_col;
+    int tx_row;
+    tx_col = mi_col << (MI_SIZE_LOG2 - tx_size_wide_log2[0]);
+    tx_row = (mi_row & MAX_MIB_MASK) << (MI_SIZE_LOG2 - tx_size_high_log2[0]);
     memcpy(ctx->a + num_4x4_blocks_wide * p,
-           xd->above_context[p] + (mi_col * 2 >> xd->plane[p].subsampling_x),
+           xd->above_context[p] + (tx_col >> xd->plane[p].subsampling_x),
            (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_wide) >>
                xd->plane[p].subsampling_x);
     memcpy(ctx->l + num_4x4_blocks_high * p,
-           xd->left_context[p] +
-               ((mi_row & MAX_MIB_MASK) * 2 >> xd->plane[p].subsampling_y),
+           xd->left_context[p] + (tx_row >> xd->plane[p].subsampling_y),
            (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_high) >>
                xd->plane[p].subsampling_y);
   }