From 81c389e79035988436fa29121cee54e4a2d7186d Mon Sep 17 00:00:00 2001
From: Jingning Han <jingning@google.com>
Date: Thu, 18 Jun 2015 12:39:09 -0700
Subject: [PATCH] Make tx partition entropy coder account for block size

This commit allows the entropy coder for transform block partition
to account for its relative position with respect to the block size.

Change-Id: I2b5019c378bfb58c11b926fa50c0db1933f35852
---
 vp9/common/vp9_blockd.h       | 2 +-
 vp9/common/vp9_entropymode.c  | 3 +--
 vp9/common/vp9_onyxc_int.h    | 6 +++++-
 vp9/decoder/vp9_decodemv.c    | 1 +
 vp9/encoder/vp9_bitstream.c   | 1 +
 vp9/encoder/vp9_encodeframe.c | 1 +
 vp9/encoder/vp9_rdopt.c       | 1 +
 7 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/vp9/common/vp9_blockd.h b/vp9/common/vp9_blockd.h
index d559d4364..142b2145e 100644
--- a/vp9/common/vp9_blockd.h
+++ b/vp9/common/vp9_blockd.h
@@ -29,7 +29,7 @@ extern "C" {
 #define BLOCK_SIZE_GROUPS 4
 #define SKIP_CONTEXTS 3
 #define INTER_MODE_CONTEXTS 7
-#define TXFM_PARTITION_CONTEXTS 9
+#define TXFM_PARTITION_CONTEXTS 12
 
 /* Segment Feature Masks */
 #define MAX_MV_REF_CANDIDATES 2
diff --git a/vp9/common/vp9_entropymode.c b/vp9/common/vp9_entropymode.c
index abd489871..15185e8a2 100644
--- a/vp9/common/vp9_entropymode.c
+++ b/vp9/common/vp9_entropymode.c
@@ -303,8 +303,7 @@ void tx_counts_to_branch_counts_8x8(const unsigned int *tx_count_8x8p,
 }
 
 static const vp9_prob default_txfm_partition_probs[TXFM_PARTITION_CONTEXTS] = {
-    // 64, 128, 192, 64, 128, 192, 64, 128, 192,
-    192, 128, 64, 192, 128, 64, 192, 128, 64,
+    141, 139, 175, 87, 196, 165, 177, 75, 220, 179, 205, 197
 };
 
 static const vp9_prob default_skip_probs[SKIP_CONTEXTS] = {
diff --git a/vp9/common/vp9_onyxc_int.h b/vp9/common/vp9_onyxc_int.h
index a748723f2..d1ad5c080 100644
--- a/vp9/common/vp9_onyxc_int.h
+++ b/vp9/common/vp9_onyxc_int.h
@@ -441,12 +441,16 @@ static void txfm_partition_update(TXFM_CONTEXT *above_ctx,
   }
 }
 
+static int max_tx_size_offset[TX_SIZES] = {0, 0, 2, 6};
+
 static int txfm_partition_context(const TXFM_CONTEXT *above_ctx,
                                   const TXFM_CONTEXT *left_ctx,
+                                  TX_SIZE max_tx_size,
                                   TX_SIZE tx_size) {
   int above = *above_ctx < tx_size;
   int left = *left_ctx < tx_size;
-  return (tx_size - 1) * 3 + above + left;
+  return max_tx_size_offset[max_tx_size] +
+      2 * (max_tx_size - tx_size) + (above || left);
 }
 
 #ifdef __cplusplus
diff --git a/vp9/decoder/vp9_decodemv.c b/vp9/decoder/vp9_decodemv.c
index 4c13ac519..0ea7337e3 100644
--- a/vp9/decoder/vp9_decodemv.c
+++ b/vp9/decoder/vp9_decodemv.c
@@ -73,6 +73,7 @@ static void read_tx_size_inter(VP9_COMMON *cm, MACROBLOCKD *xd,
   int i, j;
   int ctx = txfm_partition_context(xd->above_txfm_context + (blk_col / 2),
                                    xd->left_txfm_context + (blk_row / 2),
+                                   mbmi->max_tx_size,
                                    tx_size);
   if (xd->mb_to_bottom_edge < 0)
     max_blocks_high += xd->mb_to_bottom_edge >> 5;
diff --git a/vp9/encoder/vp9_bitstream.c b/vp9/encoder/vp9_bitstream.c
index 71d61b773..e74774768 100644
--- a/vp9/encoder/vp9_bitstream.c
+++ b/vp9/encoder/vp9_bitstream.c
@@ -85,6 +85,7 @@ static void write_tx_size_inter(const VP9_COMMON *cm, MACROBLOCKD *xd,
   int max_blocks_wide = num_4x4_blocks_wide_lookup[mbmi->sb_type];
   int ctx = txfm_partition_context(xd->above_txfm_context + (blk_col / 2),
                                    xd->left_txfm_context + (blk_row / 2),
+                                   mbmi->max_tx_size,
                                    tx_size);
   if (xd->mb_to_bottom_edge < 0)
     max_blocks_high += xd->mb_to_bottom_edge >> 5;
diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index ea3f47739..4b13aa374 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -4080,6 +4080,7 @@ static void update_txfm_count(MACROBLOCKD *xd, FRAME_COUNTS *counts,
   int max_blocks_wide = num_4x4_blocks_wide_lookup[mbmi->sb_type];
   int ctx = txfm_partition_context(xd->above_txfm_context + (blk_col / 2),
                                    xd->left_txfm_context + (blk_row / 2),
+                                   mbmi->max_tx_size,
                                    tx_size);
   TX_SIZE plane_tx_size = mbmi->inter_tx_size[tx_idx];
 
diff --git a/vp9/encoder/vp9_rdopt.c b/vp9/encoder/vp9_rdopt.c
index 49550672a..912a92a6b 100644
--- a/vp9/encoder/vp9_rdopt.c
+++ b/vp9/encoder/vp9_rdopt.c
@@ -1326,6 +1326,7 @@ static void select_tx_block(const VP9_COMP *cpi, MACROBLOCK *x,
   TXFM_CONTEXT stxa[8], stxl[8];
   int ctx = txfm_partition_context(txa + (blk_col / 2),
                                    txl + (blk_row / 2),
+                                   mbmi->max_tx_size,
                                    tx_size);
   int zero_blk_rate;
 
-- 
2.49.0