Add nonzero prob/count tables for adaptive scan order

The nonzero prob is updated by update_scan_prob()

Change-Id: I0c96bab2dc29938103b4b668653585198ff7a5c8

av1/common/entropy.c

@@ -14,6 +14,7 @@
 #include "av1/common/blockd.h"
 #include "av1/common/onyxc_int.h"
 #include "av1/common/entropymode.h"
+#include "av1/common/scan.h"
 #if CONFIG_RANS
 #include "aom_dsp/ans.h"
 #endif  // CONFIG_RANS
@@ -1633,6 +1634,10 @@ void av1_default_coef_probs(AV1_COMMON *cm) {
 #define COEF_COUNT_SAT_AFTER_KEY 24
 #define COEF_MAX_UPDATE_FACTOR_AFTER_KEY 128
 
+#if CONFIG_ADAPT_SCAN
+#define ADAPT_SCAN_UPDATE_RATE_16 (1 << 13)
+#endif
+
 static void adapt_coef_probs(AV1_COMMON *cm, TX_SIZE tx_size,
                              unsigned int count_sat,
                              unsigned int update_factor) {
@@ -1663,6 +1668,10 @@ static void adapt_coef_probs(AV1_COMMON *cm, TX_SIZE tx_size,
 }
 
 void av1_adapt_coef_probs(AV1_COMMON *cm) {
+#if CONFIG_ADAPT_SCAN
+  TX_TYPE tx_type;
+  TX_SIZE tx_size;
+#endif
   TX_SIZE t;
   unsigned int count_sat, update_factor;
 
@@ -1676,6 +1685,13 @@ void av1_adapt_coef_probs(AV1_COMMON *cm) {
     update_factor = COEF_MAX_UPDATE_FACTOR;
     count_sat = COEF_COUNT_SAT;
   }
+
+#if CONFIG_ADAPT_SCAN
+  for (tx_size = 0; tx_size < TX_SIZES; ++tx_size)
+    for (tx_type = 0; tx_type < TX_TYPES; ++tx_type)
+      update_scan_prob(cm, tx_size, tx_type, ADAPT_SCAN_UPDATE_RATE_16);
+#endif
+
   for (t = 0; t <= TX_32X32; t++)
     adapt_coef_probs(cm, t, count_sat, update_factor);
 #if CONFIG_RANS || CONFIG_DAALA_EC

av1/common/entropymode.h

@@ -68,6 +68,14 @@ typedef struct frame_contexts {
   aom_prob switchable_interp_prob[SWITCHABLE_FILTER_CONTEXTS]
                                  [SWITCHABLE_FILTERS - 1];
 
+#if CONFIG_ADAPT_SCAN
+  // TODO(angiebird): try aom_prob
+  uint32_t non_zero_prob_4X4[TX_TYPES][16];
+  uint32_t non_zero_prob_8X8[TX_TYPES][64];
+  uint32_t non_zero_prob_16X16[TX_TYPES][256];
+  uint32_t non_zero_prob_32X32[TX_TYPES][1024];
+#endif
+
 #if CONFIG_REF_MV
   aom_prob newmv_prob[NEWMV_MODE_CONTEXTS];
   aom_prob zeromv_prob[ZEROMV_MODE_CONTEXTS];
@@ -127,6 +135,14 @@ typedef struct FRAME_COUNTS {
   unsigned int switchable_interp[SWITCHABLE_FILTER_CONTEXTS]
                                 [SWITCHABLE_FILTERS];
 
+#if CONFIG_ADAPT_SCAN
+  unsigned int non_zero_count_4X4[TX_TYPES][16];
+  unsigned int non_zero_count_8X8[TX_TYPES][64];
+  unsigned int non_zero_count_16X16[TX_TYPES][256];
+  unsigned int non_zero_count_32X32[TX_TYPES][1024];
+  unsigned int txb_count[TX_SIZES][TX_TYPES];
+#endif
+
 #if CONFIG_REF_MV
   unsigned int newmv_mode[NEWMV_MODE_CONTEXTS][2];
   unsigned int zeromv_mode[ZEROMV_MODE_CONTEXTS][2];

av1/common/scan.c

@@ -736,3 +736,62 @@ const SCAN_ORDER av1_scan_orders[TX_SIZES][TX_TYPES] = {
     { default_scan_32x32, av1_default_iscan_32x32,
       default_scan_32x32_neighbors } },
 };
+
+#if CONFIG_ADAPT_SCAN
+int get_tx1d_size(TX_SIZE tx_size) { return 1 << (tx_size + 2); }
+
+int get_tx2d_size(TX_SIZE tx_size) { return 1 << ((tx_size + 2) * 2); }
+
+uint32_t *get_non_zero_prob(FRAME_CONTEXT *fc, TX_SIZE tx_size,
+                            TX_TYPE tx_type) {
+  switch (tx_size) {
+    case TX_4X4:
+      return fc->non_zero_prob_4X4[tx_type];
+    case TX_8X8:
+      return fc->non_zero_prob_8X8[tx_type];
+    case TX_16X16:
+      return fc->non_zero_prob_16X16[tx_type];
+    case TX_32X32:
+      return fc->non_zero_prob_32X32[tx_type];
+    default:
+      assert(0);
+      return NULL;
+  }
+}
+
+uint32_t *get_non_zero_counts(FRAME_COUNTS *counts, TX_SIZE tx_size,
+                              TX_TYPE tx_type) {
+  switch (tx_size) {
+    case TX_4X4:
+      return counts->non_zero_count_4X4[tx_type];
+    case TX_8X8:
+      return counts->non_zero_count_8X8[tx_type];
+    case TX_16X16:
+      return counts->non_zero_count_16X16[tx_type];
+    case TX_32X32:
+      return counts->non_zero_count_32X32[tx_type];
+    default:
+      assert(0);
+      return NULL;
+  }
+}
+
+void update_scan_prob(AV1_COMMON *cm, TX_SIZE tx_size, TX_TYPE tx_type,
+                      int rate_16) {
+  FRAME_CONTEXT *pre_fc = &cm->frame_contexts[cm->frame_context_idx];
+  uint32_t *prev_non_zero_prob = get_non_zero_prob(pre_fc, tx_size, tx_type);
+  uint32_t *non_zero_prob = get_non_zero_prob(cm->fc, tx_size, tx_type);
+  uint32_t *non_zero_count = get_non_zero_counts(&cm->counts, tx_size, tx_type);
+  int tx2d_size = get_tx2d_size(tx_size);
+  unsigned int block_num = cm->counts.txb_count[tx_size][tx_type];
+  int i;
+  for (i = 0; i < tx2d_size; i++) {
+    int64_t curr_prob =
+        block_num == 0 ? 0 : (non_zero_count[i] << 16) / block_num;
+    int64_t prev_prob = prev_non_zero_prob[i];
+    int64_t pred_prob =
+        (curr_prob * rate_16 + prev_prob * ((1 << 16) - rate_16)) >> 16;
+    non_zero_prob[i] = clamp(pred_prob, 0, UINT16_MAX);
+  }
+}
+#endif

av1/common/scan.h

@@ -15,8 +15,9 @@
 #include "aom/aom_integer.h"
 #include "aom_ports/mem.h"
 
-#include "av1/common/enums.h"
 #include "av1/common/blockd.h"
+#include "av1/common/enums.h"
+#include "av1/common/onyxc_int.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -33,6 +34,11 @@ typedef struct {
 extern const SCAN_ORDER av1_default_scan_orders[TX_SIZES];
 extern const SCAN_ORDER av1_scan_orders[TX_SIZES][TX_TYPES];
 
+#if CONFIG_ADAPT_SCAN
+void update_scan_prob(AV1_COMMON *cm, TX_SIZE tx_size, TX_TYPE tx_type,
+                      int rate_16);
+#endif
+
 static INLINE int get_coef_context(const int16_t *neighbors,
                                    const uint8_t *token_cache, int c) {
   return (1 + token_cache[neighbors[MAX_NEIGHBORS * c + 0]] +