Update aom to slightly newer commit ID

1 files changed, 1124 insertions, 25 deletions

diff --git a/third_party/aom/av1/encoder/encodetxb.c b/third_party/aom/av1/encoder/encodetxb.c
index 3f71a4472..731642064 100644
--- a/third_party/aom/av1/encoder/encodetxb.c
+++ b/third_party/aom/av1/encoder/encodetxb.c
@@ -21,6 +21,8 @@
 #include "av1/encoder/subexp.h"
 #include "av1/encoder/tokenize.h"
 
+#define TEST_OPTIMIZE_TXB 0
+
 void av1_alloc_txb_buf(AV1_COMP *cpi) {
 #if 0
   AV1_COMMON *cm = &cpi->common;
@@ -159,7 +161,7 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
     }
 
     // level is above 1.
-    ctx = get_level_ctx(tcoeff, scan[c], bwl);
+    ctx = get_br_ctx(tcoeff, scan[c], bwl);
     for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
       if (level == (idx + 1 + NUM_BASE_LEVELS)) {
         aom_write(w, 1, cm->fc->coeff_lps[tx_size][plane_type][ctx]);
@@ -251,6 +253,32 @@ static INLINE void get_base_ctx_set(const tran_low_t *tcoeffs,
   return;
 }
 
+static INLINE int get_br_cost(tran_low_t abs_qc, int ctx,
+                              const aom_prob *coeff_lps) {
+  const tran_low_t min_level = 1 + NUM_BASE_LEVELS;
+  const tran_low_t max_level = 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE;
+  if (abs_qc >= min_level) {
+    const int cost0 = av1_cost_bit(coeff_lps[ctx], 0);
+    const int cost1 = av1_cost_bit(coeff_lps[ctx], 1);
+    if (abs_qc >= max_level)
+      return COEFF_BASE_RANGE * cost0;
+    else
+      return (abs_qc - min_level) * cost0 + cost1;
+  } else {
+    return 0;
+  }
+}
+
+static INLINE int get_base_cost(tran_low_t abs_qc, int ctx,
+                                aom_prob (*coeff_base)[COEFF_BASE_CONTEXTS],
+                                int base_idx) {
+  const int level = base_idx + 1;
+  if (abs_qc < level)
+    return 0;
+  else
+    return av1_cost_bit(coeff_base[base_idx][ctx], abs_qc == level);
+}
+
 int av1_cost_coeffs_txb(const AV1_COMP *const cpi, MACROBLOCK *x, int plane,
                         int block, TXB_CTX *txb_ctx) {
   const AV1_COMMON *const cm = &cpi->common;
@@ -331,7 +359,7 @@ int av1_cost_coeffs_txb(const AV1_COMP *const cpi, MACROBLOCK *x, int plane,
         int idx;
         int ctx;
 
-        ctx = get_level_ctx(qcoeff, scan[c], bwl);
+        ctx = get_br_ctx(qcoeff, scan[c], bwl);
 
         for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
           if (level == (idx + 1 + NUM_BASE_LEVELS)) {
@@ -373,12 +401,1085 @@ int av1_cost_coeffs_txb(const AV1_COMP *const cpi, MACROBLOCK *x, int plane,
   return cost;
 }
 
-typedef struct TxbParams {
-  const AV1_COMP *cpi;
-  ThreadData *td;
-  int rate;
-} TxbParams;
+static INLINE int has_base(tran_low_t qc, int base_idx) {
+  const int level = base_idx + 1;
+  return abs(qc) >= level;
+}
+
+static void gen_base_count_mag_arr(int (*base_count_arr)[MAX_TX_SQUARE],
+                                   int (*base_mag_arr)[2],
+                                   const tran_low_t *qcoeff, int stride,
+                                   int eob, const int16_t *scan) {
+  for (int c = 0; c < eob; ++c) {
+    const int coeff_idx = scan[c];  // raster order
+    if (!has_base(qcoeff[coeff_idx], 0)) continue;
+    const int row = coeff_idx / stride;
+    const int col = coeff_idx % stride;
+    int *mag = base_mag_arr[coeff_idx];
+    get_mag(mag, qcoeff, stride, row, col, base_ref_offset,
+            BASE_CONTEXT_POSITION_NUM);
+    for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
+      if (!has_base(qcoeff[coeff_idx], i)) continue;
+      int *count = base_count_arr[i] + coeff_idx;
+      *count = get_level_count(qcoeff, stride, row, col, i, base_ref_offset,
+                               BASE_CONTEXT_POSITION_NUM);
+    }
+  }
+}
+
+static void gen_nz_count_arr(int(*nz_count_arr), const tran_low_t *qcoeff,
+                             int stride, int eob,
+                             const SCAN_ORDER *scan_order) {
+  const int16_t *scan = scan_order->scan;
+  const int16_t *iscan = scan_order->iscan;
+  for (int c = 0; c < eob; ++c) {
+    const int coeff_idx = scan[c];  // raster order
+    const int row = coeff_idx / stride;
+    const int col = coeff_idx % stride;
+    nz_count_arr[coeff_idx] = get_nz_count(qcoeff, stride, row, col, iscan);
+  }
+}
+
+static void gen_nz_ctx_arr(int (*nz_ctx_arr)[2], int(*nz_count_arr),
+                           const tran_low_t *qcoeff, int bwl, int eob,
+                           const SCAN_ORDER *scan_order) {
+  const int16_t *scan = scan_order->scan;
+  const int16_t *iscan = scan_order->iscan;
+  for (int c = 0; c < eob; ++c) {
+    const int coeff_idx = scan[c];  // raster order
+    const int count = nz_count_arr[coeff_idx];
+    nz_ctx_arr[coeff_idx][0] =
+        get_nz_map_ctx_from_count(count, qcoeff, coeff_idx, bwl, iscan);
+  }
+}
+
+static void gen_base_ctx_arr(int (*base_ctx_arr)[MAX_TX_SQUARE][2],
+                             int (*base_count_arr)[MAX_TX_SQUARE],
+                             int (*base_mag_arr)[2], const tran_low_t *qcoeff,
+                             int stride, int eob, const int16_t *scan) {
+  (void)qcoeff;
+  for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
+    for (int c = 0; c < eob; ++c) {
+      const int coeff_idx = scan[c];  // raster order
+      if (!has_base(qcoeff[coeff_idx], i)) continue;
+      const int row = coeff_idx / stride;
+      const int col = coeff_idx % stride;
+      const int count = base_count_arr[i][coeff_idx];
+      const int *mag = base_mag_arr[coeff_idx];
+      const int level = i + 1;
+      base_ctx_arr[i][coeff_idx][0] =
+          get_base_ctx_from_count_mag(row, col, count, mag[0], level);
+    }
+  }
+}
+
+static INLINE int has_br(tran_low_t qc) {
+  return abs(qc) >= 1 + NUM_BASE_LEVELS;
+}
+
+static void gen_br_count_mag_arr(int *br_count_arr, int (*br_mag_arr)[2],
+                                 const tran_low_t *qcoeff, int stride, int eob,
+                                 const int16_t *scan) {
+  for (int c = 0; c < eob; ++c) {
+    const int coeff_idx = scan[c];  // raster order
+    if (!has_br(qcoeff[coeff_idx])) continue;
+    const int row = coeff_idx / stride;
+    const int col = coeff_idx % stride;
+    int *count = br_count_arr + coeff_idx;
+    int *mag = br_mag_arr[coeff_idx];
+    *count = get_level_count(qcoeff, stride, row, col, NUM_BASE_LEVELS,
+                             br_ref_offset, BR_CONTEXT_POSITION_NUM);
+    get_mag(mag, qcoeff, stride, row, col, br_ref_offset,
+            BR_CONTEXT_POSITION_NUM);
+  }
+}
+
+static void gen_br_ctx_arr(int (*br_ctx_arr)[2], const int *br_count_arr,
+                           int (*br_mag_arr)[2], const tran_low_t *qcoeff,
+                           int stride, int eob, const int16_t *scan) {
+  (void)qcoeff;
+  for (int c = 0; c < eob; ++c) {
+    const int coeff_idx = scan[c];  // raster order
+    if (!has_br(qcoeff[coeff_idx])) continue;
+    const int row = coeff_idx / stride;
+    const int col = coeff_idx % stride;
+    const int count = br_count_arr[coeff_idx];
+    const int *mag = br_mag_arr[coeff_idx];
+    br_ctx_arr[coeff_idx][0] =
+        get_br_ctx_from_count_mag(row, col, count, mag[0]);
+  }
+}
+
+static INLINE int get_sign_bit_cost(tran_low_t qc, int coeff_idx,
+                                    const aom_prob *dc_sign_prob,
+                                    int dc_sign_ctx) {
+  const int sign = (qc < 0) ? 1 : 0;
+  // sign bit cost
+  if (coeff_idx == 0) {
+    return av1_cost_bit(dc_sign_prob[dc_sign_ctx], sign);
+  } else {
+    return av1_cost_bit(128, sign);
+  }
+}
+static INLINE int get_golomb_cost(int abs_qc) {
+  if (abs_qc >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+    // residual cost
+    int r = abs_qc - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
+    int ri = r;
+    int length = 0;
+
+    while (ri) {
+      ri >>= 1;
+      ++length;
+    }
+
+    return av1_cost_literal(2 * length - 1);
+  } else {
+    return 0;
+  }
+}
+
+// TODO(angiebird): add static once this function is called
+void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
+  const int16_t *scan = txb_info->scan_order->scan;
+  gen_nz_count_arr(txb_cache->nz_count_arr, txb_info->qcoeff, txb_info->stride,
+                   txb_info->eob, txb_info->scan_order);
+  gen_nz_ctx_arr(txb_cache->nz_ctx_arr, txb_cache->nz_count_arr,
+                 txb_info->qcoeff, txb_info->bwl, txb_info->eob,
+                 txb_info->scan_order);
+  gen_base_count_mag_arr(txb_cache->base_count_arr, txb_cache->base_mag_arr,
+                         txb_info->qcoeff, txb_info->stride, txb_info->eob,
+                         scan);
+  gen_base_ctx_arr(txb_cache->base_ctx_arr, txb_cache->base_count_arr,
+                   txb_cache->base_mag_arr, txb_info->qcoeff, txb_info->stride,
+                   txb_info->eob, scan);
+  gen_br_count_mag_arr(txb_cache->br_count_arr, txb_cache->br_mag_arr,
+                       txb_info->qcoeff, txb_info->stride, txb_info->eob, scan);
+  gen_br_ctx_arr(txb_cache->br_ctx_arr, txb_cache->br_count_arr,
+                 txb_cache->br_mag_arr, txb_info->qcoeff, txb_info->stride,
+                 txb_info->eob, scan);
+}
+
+static INLINE aom_prob get_level_prob(int level, int coeff_idx,
+                                      const TxbCache *txb_cache,
+                                      const TxbProbs *txb_probs) {
+  if (level == 0) {
+    const int ctx = txb_cache->nz_ctx_arr[coeff_idx][0];
+    return txb_probs->nz_map[ctx];
+  } else if (level >= 1 && level < 1 + NUM_BASE_LEVELS) {
+    const int idx = level - 1;
+    const int ctx = txb_cache->base_ctx_arr[idx][coeff_idx][0];
+    return txb_probs->coeff_base[idx][ctx];
+  } else if (level >= 1 + NUM_BASE_LEVELS &&
+             level < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+    const int ctx = txb_cache->br_ctx_arr[coeff_idx][0];
+    return txb_probs->coeff_lps[ctx];
+  } else if (level >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+    printf("get_level_prob does not support golomb\n");
+    assert(0);
+    return 0;
+  } else {
+    assert(0);
+    return 0;
+  }
+}
+
+static INLINE tran_low_t get_lower_coeff(tran_low_t qc) {
+  if (qc == 0) {
+    return 0;
+  }
+  return qc > 0 ? qc - 1 : qc + 1;
+}
+
+static INLINE void update_mag_arr(int *mag_arr, int abs_qc) {
+  if (mag_arr[0] == abs_qc) {
+    mag_arr[1] -= 1;
+    assert(mag_arr[1] >= 0);
+  }
+}
+
+static INLINE int get_mag_from_mag_arr(const int *mag_arr) {
+  int mag;
+  if (mag_arr[1] > 0) {
+    mag = mag_arr[0];
+  } else if (mag_arr[0] > 0) {
+    mag = mag_arr[0] - 1;
+  } else {
+    // no neighbor
+    assert(mag_arr[0] == 0 && mag_arr[1] == 0);
+    mag = 0;
+  }
+  return mag;
+}
+
+static int neighbor_level_down_update(int *new_count, int *new_mag, int count,
+                                      const int *mag, int coeff_idx,
+                                      tran_low_t abs_nb_coeff, int nb_coeff_idx,
+                                      int level, const TxbInfo *txb_info) {
+  *new_count = count;
+  *new_mag = get_mag_from_mag_arr(mag);
+
+  int update = 0;
+  // check if br_count changes
+  if (abs_nb_coeff == level) {
+    update = 1;
+    *new_count -= 1;
+    assert(*new_count >= 0);
+  }
+  const int row = coeff_idx >> txb_info->bwl;
+  const int col = coeff_idx - (row << txb_info->bwl);
+  const int nb_row = nb_coeff_idx >> txb_info->bwl;
+  const int nb_col = nb_coeff_idx - (nb_row << txb_info->bwl);
+
+  // check if mag changes
+  if (nb_row >= row && nb_col >= col) {
+    if (abs_nb_coeff == mag[0]) {
+      assert(mag[1] > 0);
+      if (mag[1] == 1) {
+        // the nb is the only qc with max mag
+        *new_mag -= 1;
+        assert(*new_mag >= 0);
+        update = 1;
+      }
+    }
+  }
+  return update;
+}
+
+static int try_neighbor_level_down_br(int coeff_idx, int nb_coeff_idx,
+                                      const TxbCache *txb_cache,
+                                      const TxbProbs *txb_probs,
+                                      const TxbInfo *txb_info) {
+  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  const tran_low_t abs_qc = abs(qc);
+  const int level = NUM_BASE_LEVELS + 1;
+  if (abs_qc < level) return 0;
+
+  const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
+  const tran_low_t abs_nb_coeff = abs(nb_coeff);
+  const int count = txb_cache->br_count_arr[coeff_idx];
+  const int *mag = txb_cache->br_mag_arr[coeff_idx];
+  int new_count;
+  int new_mag;
+  const int update =
+      neighbor_level_down_update(&new_count, &new_mag, count, mag, coeff_idx,
+                                 abs_nb_coeff, nb_coeff_idx, level, txb_info);
+  if (update) {
+    const int row = coeff_idx >> txb_info->bwl;
+    const int col = coeff_idx - (row << txb_info->bwl);
+    const int ctx = txb_cache->br_ctx_arr[coeff_idx][0];
+    const int org_cost = get_br_cost(abs_qc, ctx, txb_probs->coeff_lps);
+
+    const int new_ctx = get_br_ctx_from_count_mag(row, col, new_count, new_mag);
+    const int new_cost = get_br_cost(abs_qc, new_ctx, txb_probs->coeff_lps);
+    const int cost_diff = -org_cost + new_cost;
+    return cost_diff;
+  } else {
+    return 0;
+  }
+}
+
+static int try_neighbor_level_down_base(int coeff_idx, int nb_coeff_idx,
+                                        const TxbCache *txb_cache,
+                                        const TxbProbs *txb_probs,
+                                        const TxbInfo *txb_info) {
+  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  const tran_low_t abs_qc = abs(qc);
+
+  int cost_diff = 0;
+  for (int base_idx = 0; base_idx < NUM_BASE_LEVELS; ++base_idx) {
+    const int level = base_idx + 1;
+    if (abs_qc < level) continue;
+
+    const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
+    const tran_low_t abs_nb_coeff = abs(nb_coeff);
+
+    const int count = txb_cache->base_count_arr[base_idx][coeff_idx];
+    const int *mag = txb_cache->base_mag_arr[coeff_idx];
+    int new_count;
+    int new_mag;
+    const int update =
+        neighbor_level_down_update(&new_count, &new_mag, count, mag, coeff_idx,
+                                   abs_nb_coeff, nb_coeff_idx, level, txb_info);
+    if (update) {
+      const int row = coeff_idx >> txb_info->bwl;
+      const int col = coeff_idx - (row << txb_info->bwl);
+      const int ctx = txb_cache->base_ctx_arr[base_idx][coeff_idx][0];
+      const int org_cost =
+          get_base_cost(abs_qc, ctx, txb_probs->coeff_base, base_idx);
+
+      const int new_ctx =
+          get_base_ctx_from_count_mag(row, col, new_count, new_mag, level);
+      const int new_cost =
+          get_base_cost(abs_qc, new_ctx, txb_probs->coeff_base, base_idx);
+      cost_diff += -org_cost + new_cost;
+    }
+  }
+  return cost_diff;
+}
+
+static int try_neighbor_level_down_nz(int coeff_idx, int nb_coeff_idx,
+                                      const TxbCache *txb_cache,
+                                      const TxbProbs *txb_probs,
+                                      TxbInfo *txb_info) {
+  // assume eob doesn't change
+  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  const tran_low_t abs_qc = abs(qc);
+  const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
+  const tran_low_t abs_nb_coeff = abs(nb_coeff);
+  if (abs_nb_coeff != 1) return 0;
+  const int16_t *iscan = txb_info->scan_order->iscan;
+  const int scan_idx = iscan[coeff_idx];
+  if (scan_idx == txb_info->seg_eob) return 0;
+  const int nb_scan_idx = iscan[nb_coeff_idx];
+  if (nb_scan_idx < scan_idx) {
+    const int count = txb_cache->nz_count_arr[coeff_idx];
+    assert(count > 0);
+    txb_info->qcoeff[nb_coeff_idx] = get_lower_coeff(nb_coeff);
+    const int new_ctx = get_nz_map_ctx_from_count(
+        count - 1, txb_info->qcoeff, coeff_idx, txb_info->bwl, iscan);
+    txb_info->qcoeff[nb_coeff_idx] = nb_coeff;
+    const int ctx = txb_cache->nz_ctx_arr[coeff_idx][0];
+    const int is_nz = abs_qc > 0;
+    const int org_cost = av1_cost_bit(txb_probs->nz_map[ctx], is_nz);
+    const int new_cost = av1_cost_bit(txb_probs->nz_map[new_ctx], is_nz);
+    const int cost_diff = new_cost - org_cost;
+    return cost_diff;
+  } else {
+    return 0;
+  }
+}
+
+static int try_self_level_down(tran_low_t *low_coeff, int coeff_idx,
+                               const TxbCache *txb_cache,
+                               const TxbProbs *txb_probs, TxbInfo *txb_info) {
+  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  if (qc == 0) {
+    *low_coeff = 0;
+    return 0;
+  }
+  const tran_low_t abs_qc = abs(qc);
+  *low_coeff = get_lower_coeff(qc);
+  int cost_diff;
+  if (*low_coeff == 0) {
+    const int scan_idx = txb_info->scan_order->iscan[coeff_idx];
+    const aom_prob level_prob =
+        get_level_prob(abs_qc, coeff_idx, txb_cache, txb_probs);
+    const aom_prob low_level_prob =
+        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_probs);
+    if (scan_idx < txb_info->seg_eob) {
+      // When level-0, we code the binary of abs_qc > level
+      // but when level-k k > 0 we code the binary of abs_qc == level
+      // That's why wee need this special treatment for level-0 map
+      // TODO(angiebird): make leve-0 consistent to other levels
+      cost_diff = -av1_cost_bit(level_prob, 1) +
+                  av1_cost_bit(low_level_prob, 0) -
+                  av1_cost_bit(low_level_prob, 1);
+    } else {
+      cost_diff = -av1_cost_bit(level_prob, 1);
+    }
+
+    if (scan_idx < txb_info->seg_eob) {
+      const int eob_ctx =
+          get_eob_ctx(txb_info->qcoeff, coeff_idx, txb_info->bwl);
+      cost_diff -= av1_cost_bit(txb_probs->eob_flag[eob_ctx],
+                                scan_idx == (txb_info->eob - 1));
+    }
+
+    const int sign_cost = get_sign_bit_cost(
+        qc, coeff_idx, txb_probs->dc_sign_prob, txb_info->txb_ctx->dc_sign_ctx);
+    cost_diff -= sign_cost;
+  } else if (abs_qc < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+    const aom_prob level_prob =
+        get_level_prob(abs_qc, coeff_idx, txb_cache, txb_probs);
+    const aom_prob low_level_prob =
+        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_probs);
+    cost_diff = -av1_cost_bit(level_prob, 1) + av1_cost_bit(low_level_prob, 1) -
+                av1_cost_bit(low_level_prob, 0);
+  } else if (abs_qc == 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+    const aom_prob low_level_prob =
+        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_probs);
+    cost_diff = -get_golomb_cost(abs_qc) + av1_cost_bit(low_level_prob, 1) -
+                av1_cost_bit(low_level_prob, 0);
+  } else {
+    assert(abs_qc > 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE);
+    const tran_low_t abs_low_coeff = abs(*low_coeff);
+    cost_diff = -get_golomb_cost(abs_qc) + get_golomb_cost(abs_low_coeff);
+  }
+  return cost_diff;
+}
+
+#define COST_MAP_SIZE 5
+#define COST_MAP_OFFSET 2
+
+static INLINE int check_nz_neighbor(tran_low_t qc) { return abs(qc) == 1; }
+
+static INLINE int check_base_neighbor(tran_low_t qc) {
+  return abs(qc) <= 1 + NUM_BASE_LEVELS;
+}
+
+static INLINE int check_br_neighbor(tran_low_t qc) {
+  return abs(qc) > BR_MAG_OFFSET;
+}
+
+// TODO(angiebird): add static to this function once it's called
+int try_level_down(int coeff_idx, const TxbCache *txb_cache,
+                   const TxbProbs *txb_probs, TxbInfo *txb_info,
+                   int (*cost_map)[COST_MAP_SIZE]) {
+  if (cost_map) {
+    for (int i = 0; i < COST_MAP_SIZE; ++i) av1_zero(cost_map[i]);
+  }
+
+  tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  tran_low_t low_coeff;
+  if (qc == 0) return 0;
+  int accu_cost_diff = 0;
+
+  const int16_t *iscan = txb_info->scan_order->iscan;
+  const int eob = txb_info->eob;
+  const int scan_idx = iscan[coeff_idx];
+  if (scan_idx < eob) {
+    const int cost_diff = try_self_level_down(&low_coeff, coeff_idx, txb_cache,
+                                              txb_probs, txb_info);
+    if (cost_map)
+      cost_map[0 + COST_MAP_OFFSET][0 + COST_MAP_OFFSET] = cost_diff;
+    accu_cost_diff += cost_diff;
+  }
+
+  const int row = coeff_idx >> txb_info->bwl;
+  const int col = coeff_idx - (row << txb_info->bwl);
+  if (check_nz_neighbor(qc)) {
+    for (int i = 0; i < SIG_REF_OFFSET_NUM; ++i) {
+      const int nb_row = row - sig_ref_offset[i][0];
+      const int nb_col = col - sig_ref_offset[i][1];
+      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
+      const int nb_scan_idx = iscan[nb_coeff_idx];
+      if (nb_scan_idx < eob && nb_row >= 0 && nb_col >= 0 &&
+          nb_row < txb_info->stride && nb_col < txb_info->stride) {
+        const int cost_diff = try_neighbor_level_down_nz(
+            nb_coeff_idx, coeff_idx, txb_cache, txb_probs, txb_info);
+        if (cost_map)
+          cost_map[nb_row - row + COST_MAP_OFFSET]
+                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
+        accu_cost_diff += cost_diff;
+      }
+    }
+  }
+
+  if (check_base_neighbor(qc)) {
+    for (int i = 0; i < BASE_CONTEXT_POSITION_NUM; ++i) {
+      const int nb_row = row - base_ref_offset[i][0];
+      const int nb_col = col - base_ref_offset[i][1];
+      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
+      const int nb_scan_idx = iscan[nb_coeff_idx];
+      if (nb_scan_idx < eob && nb_row >= 0 && nb_col >= 0 &&
+          nb_row < txb_info->stride && nb_col < txb_info->stride) {
+        const int cost_diff = try_neighbor_level_down_base(
+            nb_coeff_idx, coeff_idx, txb_cache, txb_probs, txb_info);
+        if (cost_map)
+          cost_map[nb_row - row + COST_MAP_OFFSET]
+                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
+        accu_cost_diff += cost_diff;
+      }
+    }
+  }
+
+  if (check_br_neighbor(qc)) {
+    for (int i = 0; i < BR_CONTEXT_POSITION_NUM; ++i) {
+      const int nb_row = row - br_ref_offset[i][0];
+      const int nb_col = col - br_ref_offset[i][1];
+      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
+      const int nb_scan_idx = iscan[nb_coeff_idx];
+      if (nb_scan_idx < eob && nb_row >= 0 && nb_col >= 0 &&
+          nb_row < txb_info->stride && nb_col < txb_info->stride) {
+        const int cost_diff = try_neighbor_level_down_br(
+            nb_coeff_idx, coeff_idx, txb_cache, txb_probs, txb_info);
+        if (cost_map)
+          cost_map[nb_row - row + COST_MAP_OFFSET]
+                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
+        accu_cost_diff += cost_diff;
+      }
+    }
+  }
+
+  return accu_cost_diff;
+}
+
+static int get_low_coeff_cost(int coeff_idx, const TxbCache *txb_cache,
+                              const TxbProbs *txb_probs,
+                              const TxbInfo *txb_info) {
+  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  const int abs_qc = abs(qc);
+  assert(abs_qc <= 1);
+  int cost = 0;
+  const int scan_idx = txb_info->scan_order->iscan[coeff_idx];
+  if (scan_idx < txb_info->seg_eob) {
+    const aom_prob level_prob =
+        get_level_prob(0, coeff_idx, txb_cache, txb_probs);
+    cost += av1_cost_bit(level_prob, qc != 0);
+  }
+
+  if (qc != 0) {
+    const int base_idx = 0;
+    const int ctx = txb_cache->base_ctx_arr[base_idx][coeff_idx][0];
+    cost += get_base_cost(abs_qc, ctx, txb_probs->coeff_base, base_idx);
+    if (scan_idx < txb_info->seg_eob) {
+      const int eob_ctx =
+          get_eob_ctx(txb_info->qcoeff, coeff_idx, txb_info->bwl);
+      cost += av1_cost_bit(txb_probs->eob_flag[eob_ctx],
+                           scan_idx == (txb_info->eob - 1));
+    }
+    cost += get_sign_bit_cost(qc, coeff_idx, txb_probs->dc_sign_prob,
+                              txb_info->txb_ctx->dc_sign_ctx);
+  }
+  return cost;
+}
+
+static INLINE void set_eob(TxbInfo *txb_info, int eob) {
+  txb_info->eob = eob;
+  txb_info->seg_eob = AOMMIN(eob, tx_size_2d[txb_info->tx_size] - 1);
+}
+
+// TODO(angiebird): add static to this function once it's called
+int try_change_eob(int *new_eob, int coeff_idx, const TxbCache *txb_cache,
+                   const TxbProbs *txb_probs, TxbInfo *txb_info) {
+  assert(txb_info->eob > 0);
+  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  const int abs_qc = abs(qc);
+  if (abs_qc != 1) {
+    *new_eob = -1;
+    return 0;
+  }
+  const int16_t *iscan = txb_info->scan_order->iscan;
+  const int16_t *scan = txb_info->scan_order->scan;
+  const int scan_idx = iscan[coeff_idx];
+  *new_eob = 0;
+  int cost_diff = 0;
+  cost_diff -= get_low_coeff_cost(coeff_idx, txb_cache, txb_probs, txb_info);
+  // int coeff_cost =
+  //     get_coeff_cost(qc, scan_idx, txb_info, txb_probs);
+  // if (-cost_diff != coeff_cost) {
+  //   printf("-cost_diff %d coeff_cost %d\n", -cost_diff, coeff_cost);
+  //   get_low_coeff_cost(coeff_idx, txb_cache, txb_probs, txb_info);
+  //   get_coeff_cost(qc, scan_idx, txb_info, txb_probs);
+  // }
+  for (int si = scan_idx - 1; si >= 0; --si) {
+    const int ci = scan[si];
+    if (txb_info->qcoeff[ci] != 0) {
+      *new_eob = si + 1;
+      break;
+    } else {
+      cost_diff -= get_low_coeff_cost(ci, txb_cache, txb_probs, txb_info);
+    }
+  }
+
+  const int org_eob = txb_info->eob;
+  set_eob(txb_info, *new_eob);
+  cost_diff += try_level_down(coeff_idx, txb_cache, txb_probs, txb_info, NULL);
+  set_eob(txb_info, org_eob);
+
+  if (*new_eob > 0) {
+    // Note that get_eob_ctx does NOT actually account for qcoeff, so we don't
+    // need to lower down the qcoeff here
+    const int eob_ctx =
+        get_eob_ctx(txb_info->qcoeff, scan[*new_eob - 1], txb_info->bwl);
+    cost_diff -= av1_cost_bit(txb_probs->eob_flag[eob_ctx], 0);
+    cost_diff += av1_cost_bit(txb_probs->eob_flag[eob_ctx], 1);
+  } else {
+    const int txb_skip_ctx = txb_info->txb_ctx->txb_skip_ctx;
+    cost_diff -= av1_cost_bit(txb_probs->txb_skip[txb_skip_ctx], 0);
+    cost_diff += av1_cost_bit(txb_probs->txb_skip[txb_skip_ctx], 1);
+  }
+  return cost_diff;
+}
+
+static INLINE tran_low_t qcoeff_to_dqcoeff(tran_low_t qc, int dqv, int shift) {
+  int sgn = qc < 0 ? -1 : 1;
+  return sgn * ((abs(qc) * dqv) >> shift);
+}
+
+// TODO(angiebird): add static to this function it's called
+void update_level_down(int coeff_idx, TxbCache *txb_cache, TxbInfo *txb_info) {
+  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  const int abs_qc = abs(qc);
+  if (qc == 0) return;
+  const tran_low_t low_coeff = get_lower_coeff(qc);
+  txb_info->qcoeff[coeff_idx] = low_coeff;
+  const int dqv = txb_info->dequant[coeff_idx != 0];
+  txb_info->dqcoeff[coeff_idx] =
+      qcoeff_to_dqcoeff(low_coeff, dqv, txb_info->shift);
+
+  const int row = coeff_idx >> txb_info->bwl;
+  const int col = coeff_idx - (row << txb_info->bwl);
+  const int eob = txb_info->eob;
+  const int16_t *iscan = txb_info->scan_order->iscan;
+  for (int i = 0; i < SIG_REF_OFFSET_NUM; ++i) {
+    const int nb_row = row - sig_ref_offset[i][0];
+    const int nb_col = col - sig_ref_offset[i][1];
+    const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
+    const int nb_scan_idx = iscan[nb_coeff_idx];
+    if (nb_scan_idx < eob && nb_row >= 0 && nb_col >= 0 &&
+        nb_row < txb_info->stride && nb_col < txb_info->stride) {
+      const int scan_idx = iscan[coeff_idx];
+      if (scan_idx < nb_scan_idx) {
+        const int level = 1;
+        if (abs_qc == level) {
+          txb_cache->nz_count_arr[nb_coeff_idx] -= 1;
+          assert(txb_cache->nz_count_arr[nb_coeff_idx] >= 0);
+        }
+        const int count = txb_cache->nz_count_arr[nb_coeff_idx];
+        txb_cache->nz_ctx_arr[nb_coeff_idx][0] = get_nz_map_ctx_from_count(
+            count, txb_info->qcoeff, nb_coeff_idx, txb_info->bwl, iscan);
+        // int ref_ctx = get_nz_map_ctx2(txb_info->qcoeff, nb_coeff_idx,
+        // txb_info->bwl, iscan);
+        // if (ref_ctx != txb_cache->nz_ctx_arr[nb_coeff_idx][0])
+        //   printf("nz ctx %d ref_ctx %d\n",
+        //   txb_cache->nz_ctx_arr[nb_coeff_idx][0], ref_ctx);
+      }
+    }
+  }
+
+  for (int i = 0; i < BASE_CONTEXT_POSITION_NUM; ++i) {
+    const int nb_row = row - base_ref_offset[i][0];
+    const int nb_col = col - base_ref_offset[i][1];
+    const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
+    const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
+    if (!has_base(nb_coeff, 0)) continue;
+    const int nb_scan_idx = iscan[nb_coeff_idx];
+    if (nb_scan_idx < eob && nb_row >= 0 && nb_col >= 0 &&
+        nb_row < txb_info->stride && nb_col < txb_info->stride) {
+      if (row >= nb_row && col >= nb_col)
+        update_mag_arr(txb_cache->base_mag_arr[nb_coeff_idx], abs_qc);
+      const int mag =
+          get_mag_from_mag_arr(txb_cache->base_mag_arr[nb_coeff_idx]);
+      for (int base_idx = 0; base_idx < NUM_BASE_LEVELS; ++base_idx) {
+        if (!has_base(nb_coeff, base_idx)) continue;
+        const int level = base_idx + 1;
+        if (abs_qc == level) {
+          txb_cache->base_count_arr[base_idx][nb_coeff_idx] -= 1;
+          assert(txb_cache->base_count_arr[base_idx][nb_coeff_idx] >= 0);
+        }
+        const int count = txb_cache->base_count_arr[base_idx][nb_coeff_idx];
+        txb_cache->base_ctx_arr[base_idx][nb_coeff_idx][0] =
+            get_base_ctx_from_count_mag(nb_row, nb_col, count, mag, level);
+        // int ref_ctx = get_base_ctx(txb_info->qcoeff, nb_coeff_idx,
+        // txb_info->bwl, level);
+        // if (ref_ctx != txb_cache->base_ctx_arr[base_idx][nb_coeff_idx][0]) {
+        //   printf("base ctx %d ref_ctx %d\n",
+        //   txb_cache->base_ctx_arr[base_idx][nb_coeff_idx][0], ref_ctx);
+        // }
+      }
+    }
+  }
+
+  for (int i = 0; i < BR_CONTEXT_POSITION_NUM; ++i) {
+    const int nb_row = row - br_ref_offset[i][0];
+    const int nb_col = col - br_ref_offset[i][1];
+    const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
+    const int nb_scan_idx = iscan[nb_coeff_idx];
+    const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
+    if (!has_br(nb_coeff)) continue;
+    if (nb_scan_idx < eob && nb_row >= 0 && nb_col >= 0 &&
+        nb_row < txb_info->stride && nb_col < txb_info->stride) {
+      const int level = 1 + NUM_BASE_LEVELS;
+      if (abs_qc == level) {
+        txb_cache->br_count_arr[nb_coeff_idx] -= 1;
+        assert(txb_cache->br_count_arr[nb_coeff_idx] >= 0);
+      }
+      if (row >= nb_row && col >= nb_col)
+        update_mag_arr(txb_cache->br_mag_arr[nb_coeff_idx], abs_qc);
+      const int count = txb_cache->br_count_arr[nb_coeff_idx];
+      const int mag = get_mag_from_mag_arr(txb_cache->br_mag_arr[nb_coeff_idx]);
+      txb_cache->br_ctx_arr[nb_coeff_idx][0] =
+          get_br_ctx_from_count_mag(nb_row, nb_col, count, mag);
+      // int ref_ctx = get_level_ctx(txb_info->qcoeff, nb_coeff_idx,
+      // txb_info->bwl);
+      // if (ref_ctx != txb_cache->br_ctx_arr[nb_coeff_idx][0]) {
+      //   printf("base ctx %d ref_ctx %d\n",
+      //   txb_cache->br_ctx_arr[nb_coeff_idx][0], ref_ctx);
+      // }
+    }
+  }
+}
+
+static int get_coeff_cost(tran_low_t qc, int scan_idx, TxbInfo *txb_info,
+                          const TxbProbs *txb_probs) {
+  const TXB_CTX *txb_ctx = txb_info->txb_ctx;
+  const int is_nz = (qc != 0);
+  const tran_low_t abs_qc = abs(qc);
+  int cost = 0;
+  const int16_t *scan = txb_info->scan_order->scan;
+  const int16_t *iscan = txb_info->scan_order->iscan;
+
+  if (scan_idx < txb_info->seg_eob) {
+    int coeff_ctx =
+        get_nz_map_ctx2(txb_info->qcoeff, scan[scan_idx], txb_info->bwl, iscan);
+    cost += av1_cost_bit(txb_probs->nz_map[coeff_ctx], is_nz);
+  }
+
+  if (is_nz) {
+    cost += get_sign_bit_cost(qc, scan_idx, txb_probs->dc_sign_prob,
+                              txb_ctx->dc_sign_ctx);
+
+    int ctx_ls[NUM_BASE_LEVELS] = { 0 };
+    get_base_ctx_set(txb_info->qcoeff, scan[scan_idx], txb_info->bwl, ctx_ls);
+
+    int i;
+    for (i = 0; i < NUM_BASE_LEVELS; ++i) {
+      cost += get_base_cost(abs_qc, ctx_ls[i], txb_probs->coeff_base, i);
+    }
+
+    if (abs_qc > NUM_BASE_LEVELS) {
+      int ctx = get_br_ctx(txb_info->qcoeff, scan[scan_idx], txb_info->bwl);
+      cost += get_br_cost(abs_qc, ctx, txb_probs->coeff_lps);
+      cost += get_golomb_cost(abs_qc);
+    }
+
+    if (scan_idx < txb_info->seg_eob) {
+      int eob_ctx =
+          get_eob_ctx(txb_info->qcoeff, scan[scan_idx], txb_info->bwl);
+      cost += av1_cost_bit(txb_probs->eob_flag[eob_ctx],
+                           scan_idx == (txb_info->eob - 1));
+    }
+  }
+  return cost;
+}
+
+#if TEST_OPTIMIZE_TXB
+#define ALL_REF_OFFSET_NUM 17
+static int all_ref_offset[ALL_REF_OFFSET_NUM][2] = {
+  { 0, 0 },  { -2, -1 }, { -2, 0 }, { -2, 1 }, { -1, -2 }, { -1, -1 },
+  { -1, 0 }, { -1, 1 },  { 0, -2 }, { 0, -1 }, { 1, -2 },  { 1, -1 },
+  { 1, 0 },  { 2, 0 },   { 0, 1 },  { 0, 2 },  { 1, 1 },
+};
+
+static int try_level_down_ref(int coeff_idx, const TxbProbs *txb_probs,
+                              TxbInfo *txb_info,
+                              int (*cost_map)[COST_MAP_SIZE]) {
+  if (cost_map) {
+    for (int i = 0; i < COST_MAP_SIZE; ++i) av1_zero(cost_map[i]);
+  }
+  tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  if (qc == 0) return 0;
+  int row = coeff_idx >> txb_info->bwl;
+  int col = coeff_idx - (row << txb_info->bwl);
+  int org_cost = 0;
+  for (int i = 0; i < ALL_REF_OFFSET_NUM; ++i) {
+    int nb_row = row - all_ref_offset[i][0];
+    int nb_col = col - all_ref_offset[i][1];
+    int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
+    int nb_scan_idx = txb_info->scan_order->iscan[nb_coeff_idx];
+    if (nb_scan_idx < txb_info->eob && nb_row >= 0 && nb_col >= 0 &&
+        nb_row < txb_info->stride && nb_col < txb_info->stride) {
+      tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
+      int cost = get_coeff_cost(nb_coeff, nb_scan_idx, txb_info, txb_probs);
+      if (cost_map)
+        cost_map[nb_row - row + COST_MAP_OFFSET]
+                [nb_col - col + COST_MAP_OFFSET] -= cost;
+      org_cost += cost;
+    }
+  }
+  txb_info->qcoeff[coeff_idx] = get_lower_coeff(qc);
+  int new_cost = 0;
+  for (int i = 0; i < ALL_REF_OFFSET_NUM; ++i) {
+    int nb_row = row - all_ref_offset[i][0];
+    int nb_col = col - all_ref_offset[i][1];
+    int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
+    int nb_scan_idx = txb_info->scan_order->iscan[nb_coeff_idx];
+    if (nb_scan_idx < txb_info->eob && nb_row >= 0 && nb_col >= 0 &&
+        nb_row < txb_info->stride && nb_col < txb_info->stride) {
+      tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
+      int cost = get_coeff_cost(nb_coeff, nb_scan_idx, txb_info, txb_probs);
+      if (cost_map)
+        cost_map[nb_row - row + COST_MAP_OFFSET]
+                [nb_col - col + COST_MAP_OFFSET] += cost;
+      new_cost += cost;
+    }
+  }
+  txb_info->qcoeff[coeff_idx] = qc;
+  return new_cost - org_cost;
+}
 
+static void test_level_down(int coeff_idx, const TxbCache *txb_cache,
+                            const TxbProbs *txb_probs, TxbInfo *txb_info) {
+  int cost_map[COST_MAP_SIZE][COST_MAP_SIZE];
+  int ref_cost_map[COST_MAP_SIZE][COST_MAP_SIZE];
+  const int cost_diff =
+      try_level_down(coeff_idx, txb_cache, txb_probs, txb_info, cost_map);
+  const int cost_diff_ref =
+      try_level_down_ref(coeff_idx, txb_probs, txb_info, ref_cost_map);
+  if (cost_diff != cost_diff_ref) {
+    printf("qc %d cost_diff %d cost_diff_ref %d\n", txb_info->qcoeff[coeff_idx],
+           cost_diff, cost_diff_ref);
+    for (int r = 0; r < COST_MAP_SIZE; ++r) {
+      for (int c = 0; c < COST_MAP_SIZE; ++c) {
+        printf("%d:%d ", cost_map[r][c], ref_cost_map[r][c]);
+      }
+      printf("\n");
+    }
+  }
+}
+#endif
+
+// TODO(angiebird): make this static once it's called
+int get_txb_cost(TxbInfo *txb_info, const TxbProbs *txb_probs) {
+  int cost = 0;
+  int txb_skip_ctx = txb_info->txb_ctx->txb_skip_ctx;
+  const int16_t *scan = txb_info->scan_order->scan;
+  if (txb_info->eob == 0) {
+    cost = av1_cost_bit(txb_probs->txb_skip[txb_skip_ctx], 1);
+    return cost;
+  }
+  cost = av1_cost_bit(txb_probs->txb_skip[txb_skip_ctx], 0);
+  for (int c = 0; c < txb_info->eob; ++c) {
+    tran_low_t qc = txb_info->qcoeff[scan[c]];
+    int coeff_cost = get_coeff_cost(qc, c, txb_info, txb_probs);
+    cost += coeff_cost;
+  }
+  return cost;
+}
+
+#if TEST_OPTIMIZE_TXB
+void test_try_change_eob(TxbInfo *txb_info, TxbProbs *txb_probs,
+                         TxbCache *txb_cache) {
+  int eob = txb_info->eob;
+  const int16_t *scan = txb_info->scan_order->scan;
+  if (eob > 0) {
+    int last_si = eob - 1;
+    int last_ci = scan[last_si];
+    int last_coeff = txb_info->qcoeff[last_ci];
+    if (abs(last_coeff) == 1) {
+      int new_eob;
+      int cost_diff =
+          try_change_eob(&new_eob, last_ci, txb_cache, txb_probs, txb_info);
+      int org_eob = txb_info->eob;
+      int cost = get_txb_cost(txb_info, txb_probs);
+
+      txb_info->qcoeff[last_ci] = get_lower_coeff(last_coeff);
+      set_eob(txb_info, new_eob);
+      int new_cost = get_txb_cost(txb_info, txb_probs);
+      set_eob(txb_info, org_eob);
+      txb_info->qcoeff[last_ci] = last_coeff;
+
+      int ref_cost_diff = -cost + new_cost;
+      if (cost_diff != ref_cost_diff)
+        printf("org_eob %d new_eob %d cost_diff %d ref_cost_diff %d\n", org_eob,
+               new_eob, cost_diff, ref_cost_diff);
+    }
+  }
+}
+#endif
+
+static INLINE int64_t get_coeff_dist(tran_low_t tcoeff, tran_low_t dqcoeff,
+                                     int shift) {
+  const int64_t diff = (tcoeff - dqcoeff) * (1 << shift);
+  const int64_t error = diff * diff;
+  return error;
+}
+
+typedef struct LevelDownStats {
+  int update;
+  tran_low_t low_qc;
+  tran_low_t low_dqc;
+  int64_t rd_diff;
+  int cost_diff;
+  int64_t dist_diff;
+  int new_eob;
+} LevelDownStats;
+
+void try_level_down_facade(LevelDownStats *stats, int scan_idx,
+                           const TxbCache *txb_cache, const TxbProbs *txb_probs,
+                           TxbInfo *txb_info) {
+  const int16_t *scan = txb_info->scan_order->scan;
+  const int coeff_idx = scan[scan_idx];
+  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
+  stats->new_eob = -1;
+  stats->update = 0;
+  if (qc == 0) {
+    return;
+  }
+
+  const tran_low_t tqc = txb_info->tcoeff[coeff_idx];
+  const int dqv = txb_info->dequant[coeff_idx != 0];
+
+  const tran_low_t dqc = qcoeff_to_dqcoeff(qc, dqv, txb_info->shift);
+  const int64_t dqc_dist = get_coeff_dist(tqc, dqc, txb_info->shift);
+
+  stats->low_qc = get_lower_coeff(qc);
+  stats->low_dqc = qcoeff_to_dqcoeff(stats->low_qc, dqv, txb_info->shift);
+  const int64_t low_dqc_dist =
+      get_coeff_dist(tqc, stats->low_dqc, txb_info->shift);
+
+  stats->dist_diff = -dqc_dist + low_dqc_dist;
+  stats->cost_diff = 0;
+  stats->new_eob = txb_info->eob;
+  if (scan_idx == txb_info->eob - 1 && abs(qc) == 1) {
+    stats->cost_diff = try_change_eob(&stats->new_eob, coeff_idx, txb_cache,
+                                      txb_probs, txb_info);
+  } else {
+    stats->cost_diff =
+        try_level_down(coeff_idx, txb_cache, txb_probs, txb_info, NULL);
+#if TEST_OPTIMIZE_TXB
+    test_level_down(coeff_idx, txb_cache, txb_probs, txb_info);
+#endif
+  }
+  stats->rd_diff = RDCOST(txb_info->rdmult, txb_info->rddiv, stats->cost_diff,
+                          stats->dist_diff);
+  if (stats->rd_diff < 0) stats->update = 1;
+  return;
+}
+
+static int optimize_txb(TxbInfo *txb_info, const TxbProbs *txb_probs,
+                        TxbCache *txb_cache, int dry_run) {
+  int update = 0;
+  if (txb_info->eob == 0) return update;
+  int cost_diff = 0;
+  int64_t dist_diff = 0;
+  int64_t rd_diff = 0;
+  const int max_eob = tx_size_2d[txb_info->tx_size];
+
+#if TEST_OPTIMIZE_TXB
+  int64_t sse;
+  int64_t org_dist =
+      av1_block_error_c(txb_info->tcoeff, txb_info->dqcoeff, max_eob, &sse) *
+      (1 << (2 * txb_info->shift));
+  int org_cost = get_txb_cost(txb_info, txb_probs);
+#endif
+
+  tran_low_t *org_qcoeff = txb_info->qcoeff;
+  tran_low_t *org_dqcoeff = txb_info->dqcoeff;
+
+  tran_low_t tmp_qcoeff[MAX_TX_SQUARE];
+  tran_low_t tmp_dqcoeff[MAX_TX_SQUARE];
+  const int org_eob = txb_info->eob;
+  if (dry_run) {
+    memcpy(tmp_qcoeff, org_qcoeff, sizeof(org_qcoeff[0]) * max_eob);
+    memcpy(tmp_dqcoeff, org_dqcoeff, sizeof(org_dqcoeff[0]) * max_eob);
+    txb_info->qcoeff = tmp_qcoeff;
+    txb_info->dqcoeff = tmp_dqcoeff;
+  }
+
+  const int16_t *scan = txb_info->scan_order->scan;
+
+  // forward optimize the nz_map
+  const int cur_eob = txb_info->eob;
+  for (int si = 0; si < cur_eob; ++si) {
+    const int coeff_idx = scan[si];
+    tran_low_t qc = txb_info->qcoeff[coeff_idx];
+    if (abs(qc) == 1) {
+      LevelDownStats stats;
+      try_level_down_facade(&stats, si, txb_cache, txb_probs, txb_info);
+      if (stats.update) {
+        update = 1;
+        cost_diff += stats.cost_diff;
+        dist_diff += stats.dist_diff;
+        rd_diff += stats.rd_diff;
+        update_level_down(coeff_idx, txb_cache, txb_info);
+        set_eob(txb_info, stats.new_eob);
+      }
+    }
+  }
+
+  // backward optimize the level-k map
+  for (int si = txb_info->eob - 1; si >= 0; --si) {
+    LevelDownStats stats;
+    try_level_down_facade(&stats, si, txb_cache, txb_probs, txb_info);
+    const int coeff_idx = scan[si];
+    if (stats.update) {
+#if TEST_OPTIMIZE_TXB
+// printf("si %d low_qc %d cost_diff %d dist_diff %ld rd_diff %ld eob %d new_eob
+// %d\n", si, stats.low_qc, stats.cost_diff, stats.dist_diff, stats.rd_diff,
+// txb_info->eob, stats.new_eob);
+#endif
+      update = 1;
+      cost_diff += stats.cost_diff;
+      dist_diff += stats.dist_diff;
+      rd_diff += stats.rd_diff;
+      update_level_down(coeff_idx, txb_cache, txb_info);
+      set_eob(txb_info, stats.new_eob);
+    }
+    if (si > txb_info->eob) si = txb_info->eob;
+  }
+#if TEST_OPTIMIZE_TXB
+  int64_t new_dist =
+      av1_block_error_c(txb_info->tcoeff, txb_info->dqcoeff, max_eob, &sse) *
+      (1 << (2 * txb_info->shift));
+  int new_cost = get_txb_cost(txb_info, txb_probs);
+  int64_t ref_dist_diff = new_dist - org_dist;
+  int ref_cost_diff = new_cost - org_cost;
+  if (cost_diff != ref_cost_diff || dist_diff != ref_dist_diff)
+    printf(
+        "overall rd_diff %ld\ncost_diff %d ref_cost_diff%d\ndist_diff %ld "
+        "ref_dist_diff %ld\neob %d new_eob %d\n\n",
+        rd_diff, cost_diff, ref_cost_diff, dist_diff, ref_dist_diff, org_eob,
+        txb_info->eob);
+#endif
+  if (dry_run) {
+    txb_info->qcoeff = org_qcoeff;
+    txb_info->dqcoeff = org_dqcoeff;
+    set_eob(txb_info, org_eob);
+  }
+  return update;
+}
+
+// These numbers are empirically obtained.
+static const int plane_rd_mult[REF_TYPES][PLANE_TYPES] = {
+#if CONFIG_EC_ADAPT
+  { 17, 13 }, { 16, 10 },
+#else
+  { 20, 12 }, { 16, 12 },
+#endif
+};
+
+int av1_optimize_txb(const AV1_COMMON *cm, MACROBLOCK *x, int plane, int block,
+                     TX_SIZE tx_size, TXB_CTX *txb_ctx) {
+  MACROBLOCKD *const xd = &x->e_mbd;
+  const PLANE_TYPE plane_type = get_plane_type(plane);
+  const TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
+  const MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
+  const struct macroblock_plane *p = &x->plane[plane];
+  struct macroblockd_plane *pd = &xd->plane[plane];
+  const int eob = p->eobs[block];
+  tran_low_t *qcoeff = BLOCK_OFFSET(p->qcoeff, block);
+  tran_low_t *dqcoeff = BLOCK_OFFSET(pd->dqcoeff, block);
+  const tran_low_t *tcoeff = BLOCK_OFFSET(p->coeff, block);
+  const int16_t *dequant = pd->dequant;
+  const int seg_eob = AOMMIN(eob, tx_size_2d[tx_size] - 1);
+  const aom_prob *nz_map = xd->fc->nz_map[tx_size][plane_type];
+
+  const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
+  const int stride = 1 << bwl;
+  aom_prob(*coeff_base)[COEFF_BASE_CONTEXTS] =
+      xd->fc->coeff_base[tx_size][plane_type];
+
+  const aom_prob *coeff_lps = xd->fc->coeff_lps[tx_size][plane_type];
+
+  const int is_inter = is_inter_block(mbmi);
+  const SCAN_ORDER *const scan_order =
+      get_scan(cm, tx_size, tx_type, is_inter_block(mbmi));
+
+  const TxbProbs txb_probs = { xd->fc->dc_sign[plane_type],
+                               nz_map,
+                               coeff_base,
+                               coeff_lps,
+                               xd->fc->eob_flag[tx_size][plane_type],
+                               xd->fc->txb_skip[tx_size] };
+
+  const int shift = av1_get_tx_scale(tx_size);
+  const int64_t rdmult =
+      (x->rdmult * plane_rd_mult[is_inter][plane_type] + 2) >> 2;
+  const int64_t rddiv = x->rddiv;
+
+  TxbInfo txb_info = { qcoeff,     dqcoeff, tcoeff, dequant, shift,
+                       tx_size,    bwl,     stride, eob,     seg_eob,
+                       scan_order, txb_ctx, rdmult, rddiv };
+  TxbCache txb_cache;
+  gen_txb_cache(&txb_cache, &txb_info);
+
+  const int update = optimize_txb(&txb_info, &txb_probs, &txb_cache, 0);
+  if (update) p->eobs[block] = txb_info.eob;
+  return txb_info.eob;
+}
 int av1_get_txb_entropy_context(const tran_low_t *qcoeff,
                                 const SCAN_ORDER *scan_order, int eob) {
   const int16_t *scan = scan_order->scan;
@@ -394,10 +1495,10 @@ int av1_get_txb_entropy_context(const tran_low_t *qcoeff,
   return cul_level;
 }
 
-static void update_txb_context(int plane, int block, int blk_row, int blk_col,
-                               BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
-                               void *arg) {
-  TxbParams *const args = arg;
+void av1_update_txb_context_b(int plane, int block, int blk_row, int blk_col,
+                              BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
+                              void *arg) {
+  struct tokenize_b_args *const args = arg;
   const AV1_COMP *cpi = args->cpi;
   const AV1_COMMON *cm = &cpi->common;
   ThreadData *const td = args->td;
@@ -418,10 +1519,10 @@ static void update_txb_context(int plane, int block, int blk_row, int blk_col,
   av1_set_contexts(xd, pd, plane, tx_size, cul_level, blk_col, blk_row);
 }
 
-static void update_and_record_txb_context(int plane, int block, int blk_row,
-                                          int blk_col, BLOCK_SIZE plane_bsize,
-                                          TX_SIZE tx_size, void *arg) {
-  TxbParams *const args = arg;
+void av1_update_and_record_txb_context(int plane, int block, int blk_row,
+                                       int blk_col, BLOCK_SIZE plane_bsize,
+                                       TX_SIZE tx_size, void *arg) {
+  struct tokenize_b_args *const args = arg;
   const AV1_COMP *cpi = args->cpi;
   const AV1_COMMON *cm = &cpi->common;
   ThreadData *const td = args->td;
@@ -529,7 +1630,7 @@ static void update_and_record_txb_context(int plane, int block, int blk_row,
     }
 
     // level is above 1.
-    ctx = get_level_ctx(tcoeff, scan[c], bwl);
+    ctx = get_br_ctx(tcoeff, scan[c], bwl);
     for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
       if (level == (idx + 1 + NUM_BASE_LEVELS)) {
         ++td->counts->coeff_lps[tx_size][plane_type][ctx][1];
@@ -568,23 +1669,23 @@ void av1_update_txb_context(const AV1_COMP *cpi, ThreadData *td,
   const int ctx = av1_get_skip_context(xd);
   const int skip_inc =
       !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);
-  struct TxbParams arg = { cpi, td, 0 };
+  struct tokenize_b_args arg = { cpi, td, NULL, 0 };
   (void)rate;
   (void)mi_row;
   (void)mi_col;
   if (mbmi->skip) {
     if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
-    reset_skip_context(xd, bsize);
+    av1_reset_skip_context(xd, mi_row, mi_col, bsize);
     return;
   }
 
   if (!dry_run) {
     td->counts->skip[ctx][0] += skip_inc;
     av1_foreach_transformed_block(xd, bsize, mi_row, mi_col,
-                                  update_and_record_txb_context, &arg);
+                                  av1_update_and_record_txb_context, &arg);
   } else if (dry_run == DRY_RUN_NORMAL) {
-    av1_foreach_transformed_block(xd, bsize, mi_row, mi_col, update_txb_context,
-                                  &arg);
+    av1_foreach_transformed_block(xd, bsize, mi_row, mi_col,
+                                  av1_update_txb_context_b, &arg);
   } else {
     printf("DRY_RUN_COSTCOEFFS is not supported yet\n");
     assert(0);
@@ -749,8 +1850,7 @@ int64_t av1_search_txk_type(const AV1_COMP *cpi, MACROBLOCK *x, int plane,
     av1_invalid_rd_stats(&this_rd_stats);
     av1_xform_quant(cm, x, plane, block, blk_row, blk_col, plane_bsize, tx_size,
                     coeff_ctx, AV1_XFORM_QUANT_FP);
-    if (x->plane[plane].eobs[block] && !xd->lossless[mbmi->segment_id])
-      av1_optimize_b(cm, x, plane, block, tx_size, coeff_ctx);
+    av1_optimize_b(cm, x, plane, block, plane_bsize, tx_size, a, l);
     av1_dist_block(cpi, x, plane, plane_bsize, block, blk_row, blk_col, tx_size,
                    &this_rd_stats.dist, &this_rd_stats.sse,
                    OUTPUT_HAS_PREDICTED_PIXELS);
@@ -771,8 +1871,7 @@ int64_t av1_search_txk_type(const AV1_COMP *cpi, MACROBLOCK *x, int plane,
   // copy the best result in the above tx_type search for loop
   av1_xform_quant(cm, x, plane, block, blk_row, blk_col, plane_bsize, tx_size,
                   coeff_ctx, AV1_XFORM_QUANT_FP);
-  if (x->plane[plane].eobs[block] && !xd->lossless[mbmi->segment_id])
-    av1_optimize_b(cm, x, plane, block, tx_size, coeff_ctx);
+  av1_optimize_b(cm, x, plane, block, plane_bsize, tx_size, a, l);
   if (!is_inter_block(mbmi)) {
     // intra mode needs decoded result such that the next transform block
     // can use it for prediction.