Commit 02f1eafda5 for aom

commit 02f1eafda52841082a4d04e00d0e4f338b37c73b
Author: Frank Galligan <fgalligan@google.com>
Date:   Fri Oct 30 12:07:26 2026 -0700

    av1: Optimize txb_rdopt and entropy contexts

    Force-inline update_coeff_general, check abs_dqc < abs_tqc before
    computing lower-level costs in update_coeff_simple, add an eob == 1 fast
    path in warehouse_efficients_txb, inline av1_get_entropy_contexts, and
    use a sentinel bin for dx == 0 in generate_hog_using_gradient_cache.

    Benchmarking results (incremental vs prev CL / cumulative vs baseline):
    - Stock avifenc still-image (Default 1T): -2.01% / -21.01% encode time
    - Stock avifenc still-image (Default 8T): -0.88% / -16.25% encode time
    - 48-image multi-rate corpus (Full Suite 1T): -2.02% / -21.43% time
    - Bit-exact output (1T and 8T)

    TAG=agy
    CONV=b929a630-a6cd-45a0-979c-1767c30178ee
    Change-Id: I82c3f5e01c1dc0ef8517df6b0a8335d93627f186

diff --git a/av1/common/txb_common.h b/av1/common/txb_common.h
index 243979169e..84abf4faac 100644
--- a/av1/common/txb_common.h
+++ b/av1/common/txb_common.h
@@ -84,11 +84,9 @@ static inline int get_br_ctx_2d(const uint8_t *const levels,
   assert(c > 0);
   const int col = c >> bhl;
   const int row = c - (col << bhl);
-  const int stride = (1 << bhl) + TX_PAD_HOR;
-  const int pos = c + (col << TX_PAD_HOR_LOG2);
-  int mag = AOMMIN(levels[pos + 1], MAX_BASE_BR_RANGE) +
-            AOMMIN(levels[pos + stride], MAX_BASE_BR_RANGE) +
-            AOMMIN(levels[pos + 1 + stride], MAX_BASE_BR_RANGE);
+  const intptr_t stride = (1 << bhl) + TX_PAD_HOR;
+  const uint8_t *const lv = levels + c + (col << TX_PAD_HOR_LOG2);
+  int mag = lv[1] + lv[stride] + lv[stride + 1];
   mag = AOMMIN((mag + 1) >> 1, 6);
   //((row | col) < 2) is equivalent to ((row < 2) && (col < 2))
   if ((row | col) < 2) return mag + 7;
@@ -98,10 +96,10 @@ static inline int get_br_ctx_2d(const uint8_t *const levels,
 static AOM_FORCE_INLINE int get_br_ctx_eob(const int c,  // raster order
                                            const int bhl,
                                            const TX_CLASS tx_class) {
+  if (c == 0) return 0;
   const int col = c >> bhl;
   const int row = c - (col << bhl);
-  if (c == 0) return 0;
-  if ((tx_class == TX_CLASS_2D && row < 2 && col < 2) ||
+  if ((tx_class == TX_CLASS_2D && ((row | col) < 2)) ||
       (tx_class == TX_CLASS_HORIZ && col == 0) ||
       (tx_class == TX_CLASS_VERT && row == 0))
     return 7;
@@ -113,25 +111,24 @@ static AOM_FORCE_INLINE int get_br_ctx(const uint8_t *const levels,
                                        const int bhl, const TX_CLASS tx_class) {
   const int col = c >> bhl;
   const int row = c - (col << bhl);
-  const int stride = (1 << bhl) + TX_PAD_HOR;
-  const int pos = c + (col << TX_PAD_HOR_LOG2);
-  int mag = levels[pos + 1];
-  mag += levels[pos + stride];
+  const intptr_t stride = (1 << bhl) + TX_PAD_HOR;
+  const uint8_t *const lv = levels + c + (col << TX_PAD_HOR_LOG2);
+  int mag = lv[1] + lv[stride];
   switch (tx_class) {
     case TX_CLASS_2D:
-      mag += levels[pos + stride + 1];
+      mag += lv[stride + 1];
       mag = AOMMIN((mag + 1) >> 1, 6);
       if (c == 0) return mag;
-      if ((row < 2) && (col < 2)) return mag + 7;
+      if ((row | col) < 2) return mag + 7;
       break;
     case TX_CLASS_HORIZ:
-      mag += levels[pos + (stride << 1)];
+      mag += lv[stride << 1];
       mag = AOMMIN((mag + 1) >> 1, 6);
       if (c == 0) return mag;
       if (col == 0) return mag + 7;
       break;
     case TX_CLASS_VERT:
-      mag += levels[pos + 2];
+      mag += lv[2];
       mag = AOMMIN((mag + 1) >> 1, 6);
       if (c == 0) return mag;
       if (row == 0) return mag + 7;
@@ -142,6 +139,33 @@ static AOM_FORCE_INLINE int get_br_ctx(const uint8_t *const levels,
   return mag + 14;
 }

+static AOM_FORCE_INLINE int get_br_ctx_simple(const uint8_t *const levels,
+                                              const int c,  // raster order > 0
+                                              const int bhl,
+                                              const TX_CLASS tx_class) {
+  assert(c > 0);
+  const int col = c >> bhl;
+  const int row = c - (col << bhl);
+  const intptr_t stride = (1 << bhl) + TX_PAD_HOR;
+  const uint8_t *const lv = levels + c + (col << TX_PAD_HOR_LOG2);
+  int mag = lv[1] + lv[stride];
+  switch (tx_class) {
+    case TX_CLASS_2D:
+      mag += lv[stride + 1];
+      mag = AOMMIN((mag + 1) >> 1, 6);
+      return mag + (((row | col) < 2) ? 7 : 14);
+    case TX_CLASS_HORIZ:
+      mag += lv[stride << 1];
+      mag = AOMMIN((mag + 1) >> 1, 6);
+      return mag + ((col == 0) ? 7 : 14);
+    case TX_CLASS_VERT:
+      mag += lv[2];
+      mag = AOMMIN((mag + 1) >> 1, 6);
+      return mag + ((row == 0) ? 7 : 14);
+    default: return mag + 14;
+  }
+}
+
 static const uint8_t clip_max3[256] = {
   0, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
   3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
diff --git a/av1/encoder/cost.c b/av1/encoder/cost.c
index 4bf8fce07f..6552c7ea9e 100644
--- a/av1/encoder/cost.c
+++ b/av1/encoder/cost.c
@@ -14,8 +14,9 @@
 #include "av1/encoder/cost.h"
 #include "av1/common/entropy.h"

-// round(-log2(i/256.) * (1 << AV1_PROB_COST_SHIFT)); i = 128~255.
-const uint16_t av1_prob_cost[128] = {
+// round(-log2(i/256.) * (1 << AV1_PROB_COST_SHIFT)); i = 128~255 (plus 128th
+// duplicate for 256).
+const uint16_t av1_prob_cost[129] = {
   512, 506, 501, 495, 489, 484, 478, 473, 467, 462, 456, 451, 446, 441, 435,
   430, 425, 420, 415, 410, 405, 400, 395, 390, 385, 380, 375, 371, 366, 361,
   356, 352, 347, 343, 338, 333, 329, 324, 320, 316, 311, 307, 302, 298, 294,
@@ -24,24 +25,29 @@ const uint16_t av1_prob_cost[128] = {
   171, 168, 164, 161, 157, 153, 150, 146, 143, 139, 136, 132, 129, 125, 122,
   119, 115, 112, 109, 105, 102, 99,  95,  92,  89,  86,  82,  79,  76,  73,
   70,  66,  63,  60,  57,  54,  51,  48,  45,  42,  38,  35,  32,  29,  26,
-  23,  20,  18,  15,  12,  9,   6,   3,
+  23,  20,  18,  15,  12,  9,   6,   3,   3,
 };

 void av1_cost_tokens_from_cdf(int *costs, const aom_cdf_prob *cdf,
                               const int *inv_map) {
-  int i;
-  aom_cdf_prob prev_cdf = 0;
-  for (i = 0;; ++i) {
-    aom_cdf_prob p15 = AOM_ICDF(cdf[i]) - prev_cdf;
-    p15 = (p15 < EC_MIN_PROB) ? EC_MIN_PROB : p15;
-    prev_cdf = AOM_ICDF(cdf[i]);
-
-    if (inv_map)
-      costs[inv_map[i]] = av1_cost_symbol(p15);
-    else
-      costs[i] = av1_cost_symbol(p15);
-
-    // Stop once we reach the end of the CDF
-    if (cdf[i] == AOM_ICDF(CDF_PROB_TOP)) break;
+  int prev_icdf = CDF_PROB_TOP;
+  if (inv_map) {
+    for (int i = 0;; ++i) {
+      const int cur_icdf = cdf[i];
+      const int p15 =
+          clamp(prev_icdf - cur_icdf, EC_MIN_PROB, CDF_PROB_TOP - 1);
+      prev_icdf = cur_icdf;
+      costs[inv_map[i]] = av1_cost_symbol_clamped(p15);
+      if (cur_icdf == 0) break;
+    }
+  } else {
+    for (int i = 0;; ++i) {
+      const int cur_icdf = cdf[i];
+      const int p15 =
+          clamp(prev_icdf - cur_icdf, EC_MIN_PROB, CDF_PROB_TOP - 1);
+      prev_icdf = cur_icdf;
+      costs[i] = av1_cost_symbol_clamped(p15);
+      if (cur_icdf == 0) break;
+    }
   }
 }
diff --git a/av1/encoder/cost.h b/av1/encoder/cost.h
index 9ce4b1083d..cdadb893cc 100644
--- a/av1/encoder/cost.h
+++ b/av1/encoder/cost.h
@@ -19,7 +19,7 @@
 extern "C" {
 #endif

-extern const uint16_t av1_prob_cost[128];
+extern const uint16_t av1_prob_cost[129];

 // The factor to scale from cost in bits to cost in av1_prob_cost units.
 #define AV1_PROB_COST_SHIFT 9
@@ -28,18 +28,23 @@ extern const uint16_t av1_prob_cost[128];
 // for each bit.
 #define av1_cost_literal(n) ((n) * (1 << AV1_PROB_COST_SHIFT))

+static inline int av1_cost_symbol_clamped(int p15) {
+  assert(0 < p15 && p15 < CDF_PROB_TOP);
+  const int shift = CDF_PROB_BITS - 1 - get_msb(p15);
+  // Since CDF_PROB_TOP == 32768 and av1_prob_cost has 129 entries (including a
+  // duplicate 128th entry for prob_idx == 128), get_prob(p15 << shift, 32768) -
+  // 128 simplifies to ((p15 << shift) - 16320) >> 7.
+  const int prob_idx = ((p15 << shift) - 16320) >> 7;
+  return av1_prob_cost[prob_idx] + av1_cost_literal(shift);
+}
+
 // Calculate the cost of a symbol with probability p15 / 2^15
 static inline int av1_cost_symbol(aom_cdf_prob p15) {
   // p15 can be out of range [1, CDF_PROB_TOP - 1]. Clamping it, so that the
   // following cost calculation works correctly. Otherwise, if p15 =
   // CDF_PROB_TOP, shift would be -1, and "p15 << shift" would be wrong.
-  p15 = (aom_cdf_prob)clamp(p15, 1, CDF_PROB_TOP - 1);
-  assert(0 < p15 && p15 < CDF_PROB_TOP);
-  const int shift = CDF_PROB_BITS - 1 - get_msb(p15);
-  // Since CDF_PROB_TOP == 32768, get_prob(p15 << shift, 32768) - 128 simplifies
-  // to min(((p15 << shift) - 16320) >> 7, 127).
-  const int prob_idx = AOMMIN(((p15 << shift) - 16320) >> 7, 127);
-  return av1_prob_cost[prob_idx] + av1_cost_literal(shift);
+  const int clamped = clamp((int)p15, 1, CDF_PROB_TOP - 1);
+  return av1_cost_symbol_clamped(clamped);
 }

 void av1_cost_tokens_from_cdf(int *costs, const aom_cdf_prob *cdf,
diff --git a/av1/encoder/intra_mode_search_utils.h b/av1/encoder/intra_mode_search_utils.h
index d64cc8c831..df139a0264 100644
--- a/av1/encoder/intra_mode_search_utils.h
+++ b/av1/encoder/intra_mode_search_utils.h
@@ -228,13 +228,8 @@ static AOM_FORCE_INLINE void lowbd_compute_gradient_info_sb(
       d1 = d2;

       row_info[c].abs_dx_abs_dy_sum = (uint16_t)(abs(dx) + abs(dy));
-      if (dx == 0) {
-        row_info[c].is_dx_zero = 1;
-        row_info[c].hist_bin_idx = -1;
-      } else {
-        row_info[c].is_dx_zero = 0;
-        row_info[c].hist_bin_idx = (int8_t)get_bin_idx(dx, dy);
-      }
+      row_info[c].hist_bin_idx =
+          (dx == 0) ? HOG_BINS : (int8_t)get_bin_idx(dx, dy);
     }
     src += stride;
   }
@@ -334,13 +329,8 @@ static AOM_FORCE_INLINE void highbd_compute_gradient_info_sb(
       d1 = d2;

       row_info[c].abs_dx_abs_dy_sum = (uint16_t)(abs(dx) + abs(dy));
-      if (dx == 0) {
-        row_info[c].is_dx_zero = 1;
-        row_info[c].hist_bin_idx = -1;
-      } else {
-        row_info[c].is_dx_zero = 0;
-        row_info[c].hist_bin_idx = (int8_t)get_bin_idx(dx, dy);
-      }
+      row_info[c].hist_bin_idx =
+          (dx == 0) ? HOG_BINS : (int8_t)get_bin_idx(dx, dy);
     }
     src += stride;
   }
@@ -442,26 +432,24 @@ static inline void generate_hog_using_gradient_cache(const MACROBLOCK *x,
                                                 plane * MAX_SB_SQUARE +
                                                 block_offset_in_grad_cache;

-  // Accumulate cached gradients in integers first to avoid per-pixel
-  // int-to-float conversions inside the inner loop.
-  uint32_t int_hist[HOG_BINS] = { 0 };
+  // Accumulate cached gradients in integers first, using bin HOG_BINS as a
+  // sentinel bucket for dx == 0 so the inner loop is branchless.
+  uint32_t int_hist[HOG_BINS + 1] = { 0 };
   for (int r = 1; r < rows - 1; ++r) {
     const PixelLevelGradientInfo *const row_info = grad_info_blk + r * sb_width;
     for (int c = 1; c < cols - 1; ++c) {
       const uint16_t abs_dx_abs_dy_sum = row_info[c].abs_dx_abs_dy_sum;
       if (!abs_dx_abs_dy_sum) continue;
       total += abs_dx_abs_dy_sum;
-      if (row_info[c].is_dx_zero) {
-        const uint32_t half = abs_dx_abs_dy_sum >> 1;
-        int_hist[0] += half;
-        int_hist[HOG_BINS - 1] += half;
-      } else {
-        const int8_t idx = row_info[c].hist_bin_idx;
-        assert(idx >= 0 && idx < HOG_BINS);
-        int_hist[idx] += abs_dx_abs_dy_sum;
-      }
+      const uint8_t idx = (uint8_t)row_info[c].hist_bin_idx;
+      assert(idx <= HOG_BINS);
+      int_hist[idx] += abs_dx_abs_dy_sum;
     }
   }
+
+  const uint32_t dx_zero_half = int_hist[HOG_BINS] >> 1;
+  int_hist[0] += dx_zero_half;
+  int_hist[HOG_BINS - 1] += dx_zero_half;
   for (int b = 0; b < HOG_BINS; ++b) {
     hist[b] = (float)int_hist[b];
   }
diff --git a/av1/encoder/rd.c b/av1/encoder/rd.c
index 1263d35dcd..9fafdafae6 100644
--- a/av1/encoder/rd.c
+++ b/av1/encoder/rd.c
@@ -1097,27 +1097,6 @@ void av1_model_rd_curvfit(BLOCK_SIZE bsize, double sse_norm, double xqr,
   av1_interp_cubic_rate_dist(prate, pdist, xo, rate_dist_f);
 }

-static void get_entropy_contexts_plane(BLOCK_SIZE plane_bsize,
-                                       const struct macroblockd_plane *pd,
-                                       ENTROPY_CONTEXT t_above[MAX_MIB_SIZE],
-                                       ENTROPY_CONTEXT t_left[MAX_MIB_SIZE]) {
-  const int num_4x4_w = mi_size_wide[plane_bsize];
-  const int num_4x4_h = mi_size_high[plane_bsize];
-  const ENTROPY_CONTEXT *const above = pd->above_entropy_context;
-  const ENTROPY_CONTEXT *const left = pd->left_entropy_context;
-
-  memcpy(t_above, above, sizeof(ENTROPY_CONTEXT) * num_4x4_w);
-  memcpy(t_left, left, sizeof(ENTROPY_CONTEXT) * num_4x4_h);
-}
-
-void av1_get_entropy_contexts(BLOCK_SIZE plane_bsize,
-                              const struct macroblockd_plane *pd,
-                              ENTROPY_CONTEXT t_above[MAX_MIB_SIZE],
-                              ENTROPY_CONTEXT t_left[MAX_MIB_SIZE]) {
-  assert(plane_bsize < BLOCK_SIZES_ALL);
-  get_entropy_contexts_plane(plane_bsize, pd, t_above, t_left);
-}
-
 // Special clamping used in the encoder when calculating a prediction
 //
 // Logically, all pixel fetches used for prediction are clamped against the
diff --git a/av1/encoder/rd.h b/av1/encoder/rd.h
index 55b17fce52..f0b667f523 100644
--- a/av1/encoder/rd.h
+++ b/av1/encoder/rd.h
@@ -312,10 +312,28 @@ void av1_init_me_luts(void);

 void av1_set_mvcost(MACROBLOCK *x, int ref, int ref_mv_idx);

-void av1_get_entropy_contexts(BLOCK_SIZE plane_bsize,
-                              const struct macroblockd_plane *pd,
-                              ENTROPY_CONTEXT t_above[MAX_MIB_SIZE],
-                              ENTROPY_CONTEXT t_left[MAX_MIB_SIZE]);
+static inline void copy_entropy_ctx(ENTROPY_CONTEXT *dst,
+                                    const ENTROPY_CONTEXT *src, int n) {
+  switch (n) {
+    case 1: *dst = *src; break;
+    case 2: memcpy(dst, src, 2); break;
+    case 4: memcpy(dst, src, 4); break;
+    case 8: memcpy(dst, src, 8); break;
+    case 16: memcpy(dst, src, 16); break;
+    case 32: memcpy(dst, src, 32); break;
+    default: memcpy(dst, src, (size_t)n); break;
+  }
+}
+
+static inline void av1_get_entropy_contexts(
+    BLOCK_SIZE plane_bsize, const struct macroblockd_plane *pd,
+    ENTROPY_CONTEXT t_above[MAX_MIB_SIZE],
+    ENTROPY_CONTEXT t_left[MAX_MIB_SIZE]) {
+  assert(plane_bsize < BLOCK_SIZES_ALL);
+  copy_entropy_ctx(t_above, pd->above_entropy_context,
+                   mi_size_wide[plane_bsize]);
+  copy_entropy_ctx(t_left, pd->left_entropy_context, mi_size_high[plane_bsize]);
+}

 void av1_set_rd_speed_thresholds(struct AV1_COMP *cpi);

diff --git a/av1/encoder/tx_search.c b/av1/encoder/tx_search.c
index 68505e483d..c41bcbf5ff 100644
--- a/av1/encoder/tx_search.c
+++ b/av1/encoder/tx_search.c
@@ -285,14 +285,10 @@ static inline void set_skip_txfm(MACROBLOCK *x, RD_STATS *rd_stats,
   // size possible) in the current block. Eg: For 128*128 block, rate would be
   // 4 * zero_blk_rate where zero_blk_rate corresponds to coding of one 64x64 tx
   // block as 'all zeros'
-  ENTROPY_CONTEXT ctxa[MAX_MIB_SIZE];
-  ENTROPY_CONTEXT ctxl[MAX_MIB_SIZE];
-  av1_get_entropy_contexts(bsize, &xd->plane[0], ctxa, ctxl);
-  ENTROPY_CONTEXT *ta = ctxa;
-  ENTROPY_CONTEXT *tl = ctxl;
   const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
   TXB_CTX txb_ctx;
-  get_txb_ctx(bsize, tx_size, 0, ta, tl, &txb_ctx);
+  get_txb_ctx(bsize, tx_size, 0, xd->plane[0].above_entropy_context,
+              xd->plane[0].left_entropy_context, &txb_ctx);
   const int zero_blk_rate = x->coeff_costs.coeff_costs[txs_ctx][PLANE_TYPE_Y]
                                 .txb_skip_cost[txb_ctx.txb_skip_ctx][1];
   rd_stats->rate = zero_blk_rate *
@@ -2030,15 +2026,12 @@ static inline void predict_dc_only_block(
     best_rd_stats->dist = (*block_sse) << 4;
     best_rd_stats->sse = best_rd_stats->dist;

-    ENTROPY_CONTEXT ctxa[MAX_MIB_SIZE];
-    ENTROPY_CONTEXT ctxl[MAX_MIB_SIZE];
-    av1_get_entropy_contexts(plane_bsize, &xd->plane[plane], ctxa, ctxl);
-    ENTROPY_CONTEXT *ta = ctxa;
-    ENTROPY_CONTEXT *tl = ctxl;
     const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
     TXB_CTX txb_ctx_tmp;
     const PLANE_TYPE plane_type = get_plane_type(plane);
-    get_txb_ctx(plane_bsize, tx_size, plane, ta, tl, &txb_ctx_tmp);
+    get_txb_ctx(plane_bsize, tx_size, plane,
+                xd->plane[plane].above_entropy_context,
+                xd->plane[plane].left_entropy_context, &txb_ctx_tmp);
     const int zero_blk_rate = x->coeff_costs.coeff_costs[txs_ctx][plane_type]
                                   .txb_skip_cost[txb_ctx_tmp.txb_skip_ctx][1];
     best_rd_stats->rate = zero_blk_rate;
diff --git a/av1/encoder/txb_rdopt.c b/av1/encoder/txb_rdopt.c
index 3996f86fff..cd94156f87 100644
--- a/av1/encoder/txb_rdopt.c
+++ b/av1/encoder/txb_rdopt.c
@@ -16,7 +16,7 @@
 #include "aom_ports/mem.h"
 #include "av1/common/idct.h"

-static inline void update_coeff_general(
+static AOM_FORCE_INLINE void update_coeff_general(
     int *accu_rate, int64_t *accu_dist, int si, int eob, TX_SIZE tx_size,
     TX_CLASS tx_class, int bhl, int width, int64_t rdmult, int shift,
     int dc_sign_ctx, const int16_t *dequant, const int16_t *scan,
@@ -45,7 +45,7 @@ static inline void update_coeff_general(
     if (abs_qc == 1) {
       rate_low = txb_costs->base_cost[coeff_ctx][0];
     } else {
-      const int dqv = get_dqv(dequant, scan[si], iqmatrix);
+      const int dqv = get_dqv(dequant, ci, iqmatrix);
       get_qc_dqc_low(abs_qc, sign, dqv, shift, &qc_low, &dqc_low);
       abs_qc_low = abs_qc - 1;
       rate_low =
@@ -62,12 +62,13 @@ static inline void update_coeff_general(
               ? 0
               : (((int64_t)dqc_low * (dqc_low - tqc2)) * (1LL << (2 * shift)));
     } else {
-      const int64_t dist0 = get_coeff_dist(tqc, 0, shift, qmatrix, ci);
-      dist_diff_0 = get_coeff_dist(tqc, dqc, shift, qmatrix, ci) - dist0;
+      const int64_t qm_val = qmatrix[ci];
+      const int64_t dist0 = get_coeff_dist_qm(tqc, 0, shift, qm_val);
+      dist_diff_0 = get_coeff_dist_qm(tqc, dqc, shift, qm_val) - dist0;
       dist_diff_low_0 =
           (abs_qc == 1)
               ? 0
-              : (get_coeff_dist(tqc, dqc_low, shift, qmatrix, ci) - dist0);
+              : (get_coeff_dist_qm(tqc, dqc_low, shift, qm_val) - dist0);
     }

     const int64_t rd = RDCOST(rdmult, rate, dist_diff_0);
@@ -123,9 +124,10 @@ static AOM_FORCE_INLINE int update_coeff_simple(
               ((int64_t)abs_dqc * (abs_dqc - ((int64_t)abs_tqc * 2))) *
               (1LL << (2 * shift));
         } else {
-          const int64_t dist0 = get_coeff_dist(abs_tqc, 0, shift, qmatrix, ci);
+          const int64_t qm_val = qmatrix[ci];
+          const int64_t dist0 = get_coeff_dist_qm(abs_tqc, 0, shift, qm_val);
           dist_diff_0 =
-              get_coeff_dist(abs_tqc, abs_dqc, shift, qmatrix, ci) - dist0;
+              get_coeff_dist_qm(abs_tqc, abs_dqc, shift, qm_val) - dist0;
         }
         const int rate_low = rate - base_cost[5];
         const int64_t rd = RDCOST(rdmult, rate, dist_diff_0);
@@ -143,7 +145,7 @@ static AOM_FORCE_INLINE int update_coeff_simple(
         int rate = txb_costs->base_cost[coeff_ctx][AOMMIN(abs_qc, 3)] +
                    av1_cost_literal(1);
         if (abs_qc > NUM_BASE_LEVELS) {
-          const int br_ctx = get_br_ctx(levels, ci, bhl, tx_class);
+          const int br_ctx = get_br_ctx_simple(levels, ci, bhl, tx_class);
           rate += get_br_cost(abs_qc, txb_costs->lps_cost[br_ctx]);
         }
         return rate;
@@ -168,11 +170,12 @@ static AOM_FORCE_INLINE int update_coeff_simple(
         dist_diff_low_0 = ((int64_t)abs_dqc_low * (abs_dqc_low - abs_tqc2)) *
                           (1LL << (2 * shift));
       } else {
-        const int64_t dist0 = get_coeff_dist(abs_tqc, 0, shift, qmatrix, ci);
+        const int64_t qm_val = qmatrix[ci];
+        const int64_t dist0 = get_coeff_dist_qm(abs_tqc, 0, shift, qm_val);
         dist_diff_0 =
-            get_coeff_dist(abs_tqc, abs_dqc, shift, qmatrix, ci) - dist0;
+            get_coeff_dist_qm(abs_tqc, abs_dqc, shift, qm_val) - dist0;
         dist_diff_low_0 =
-            get_coeff_dist(abs_tqc, abs_dqc_low, shift, qmatrix, ci) - dist0;
+            get_coeff_dist_qm(abs_tqc, abs_dqc_low, shift, qm_val) - dist0;
       }
       const int64_t rd = RDCOST(rdmult, rate, dist_diff_0);
       const int64_t rd_low = RDCOST(rdmult, rate_low, dist_diff_low_0);
@@ -296,12 +299,13 @@ static AOM_FORCE_INLINE void update_coeff_eob_facade(
                                  : (((int64_t)dqc_low * (dqc_low - tqc2)) *
                                     (1LL << (2 * shift)));
       } else {
-        const int64_t dist0 = get_coeff_dist(tqc, 0, shift, qmatrix, ci);
-        dist = get_coeff_dist(tqc, dqc, shift, qmatrix, ci) - dist0;
+        const int64_t qm_val = qmatrix[ci];
+        const int64_t dist0 = get_coeff_dist_qm(tqc, 0, shift, qm_val);
+        dist = get_coeff_dist_qm(tqc, dqc, shift, qm_val) - dist0;
         dist_low =
             (abs_qc == 1)
                 ? 0
-                : (get_coeff_dist(tqc, dqc_low, shift, qmatrix, ci) - dist0);
+                : (get_coeff_dist_qm(tqc, dqc_low, shift, qm_val) - dist0);
       }

       int rate =
@@ -610,7 +614,6 @@ static AOM_FORCE_INLINE int warehouse_efficients_txb(
   const LV_MAP_EOB_COST *const eob_costs =
       &x->coeff_costs.eob_costs[eob_multi_size][plane_type];
   int cost = coeff_costs->txb_skip_cost[txb_skip_ctx][0];
-
   cost += get_tx_type_cost(x, xd, plane, tx_size, tx_type, reduced_tx_set_used);
   cost += get_eob_cost(eob, eob_costs, coeff_costs, tx_class);

@@ -626,9 +629,8 @@ static AOM_FORCE_INLINE int warehouse_efficients_txb(
       if (level > NUM_BASE_LEVELS) {
         cost += get_br_cost(level, lps_cost[0]);
       }
-      const int sign01 = (sign ^ sign) - sign;
-      const int dc_sign_ctx = txb_ctx->dc_sign_ctx;
-      cost += coeff_costs->dc_sign_cost[dc_sign_ctx][sign01];
+      const int sign01 = -sign;
+      cost += coeff_costs->dc_sign_cost[txb_ctx->dc_sign_ctx][sign01];
     }
     return cost;
   }
diff --git a/av1/encoder/txb_rdopt_utils.h b/av1/encoder/txb_rdopt_utils.h
index 2b6874e7f6..97c87b8af2 100644
--- a/av1/encoder/txb_rdopt_utils.h
+++ b/av1/encoder/txb_rdopt_utils.h
@@ -45,22 +45,21 @@ static inline int get_dqv(const int16_t *dequant, int coeff_idx,
   return dqv;
 }

+static inline int64_t get_coeff_dist_qm(tran_low_t tcoeff, tran_low_t dqcoeff,
+                                        int shift, int64_t qm_val) {
+  const int64_t diff = (int64_t)(tcoeff - dqcoeff) * (1 << shift) * qm_val;
+  return (diff * diff + (1 << (2 * AOM_QM_BITS - 1))) >> (2 * AOM_QM_BITS);
+}
+
 static inline int64_t get_coeff_dist(tran_low_t tcoeff, tran_low_t dqcoeff,
                                      int shift, const qm_val_t *qmatrix,
                                      int coeff_idx) {
-  int64_t diff = (tcoeff - dqcoeff) * (1 << shift);
+  const int64_t diff = (int64_t)(tcoeff - dqcoeff) * (1 << shift);
   if (qmatrix == NULL) {
     return diff * diff;
   }
-  // When AOM_DIST_METRIC_QM_PSNR is enabled, this mirrors the rate-distortion
-  // computation done in av1_block_error_qm, improving visual quality.
-  // The maximum value of `shift` is 2, `tcoeff` and `dqcoeff` are at most 22
-  // bits, and AOM_QM_BITS is 5, so `diff` should fit in 29-bits. The
-  // multiplication `diff * diff` then does not risk overflowing.
-  diff *= qmatrix[coeff_idx];
-  const int64_t error =
-      (diff * diff + (1 << (2 * AOM_QM_BITS - 1))) >> (2 * AOM_QM_BITS);
-  return error;
+  const int64_t wdiff = diff * qmatrix[coeff_idx];
+  return (wdiff * wdiff + (1 << (2 * AOM_QM_BITS - 1))) >> (2 * AOM_QM_BITS);
 }

 static inline int get_eob_cost(int eob, const LV_MAP_EOB_COST *txb_eob_costs,
@@ -91,29 +90,34 @@ static inline int get_golomb_cost(int abs_qc) {
 }

 static inline int get_br_cost(tran_low_t level, const int *coeff_lps) {
-  const int base_range = AOMMIN(level - 1 - NUM_BASE_LEVELS, COEFF_BASE_RANGE);
-  return coeff_lps[base_range] + get_golomb_cost(level);
+  if (level < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+    return coeff_lps[level - 1 - NUM_BASE_LEVELS];
+  }
+  const int r = level - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
+  return coeff_lps[COEFF_BASE_RANGE] + av1_cost_literal((get_msb(r) << 1) + 1);
 }

 static inline int get_br_cost_with_diff(tran_low_t level, const int *coeff_lps,
                                         int *diff) {
-  const int base_range = AOMMIN(level - 1 - NUM_BASE_LEVELS, COEFF_BASE_RANGE);
-  int golomb_bits = 0;
-  if (level <= COEFF_BASE_RANGE + 1 + NUM_BASE_LEVELS)
+  if (level < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+    const int base_range = level - 1 - NUM_BASE_LEVELS;
     *diff += coeff_lps[base_range + COEFF_BASE_RANGE + 1];
+    return coeff_lps[base_range];
+  }
+  int golomb_bits = 0;
+  if (level == COEFF_BASE_RANGE + 1 + NUM_BASE_LEVELS)
+    *diff += coeff_lps[COEFF_BASE_RANGE + COEFF_BASE_RANGE + 1];

-  if (level >= COEFF_BASE_RANGE + 1 + NUM_BASE_LEVELS) {
-    int r = level - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
-    if (r < 32) {
-      golomb_bits = golomb_bits_cost[r];
-      *diff += golomb_cost_diff[r];
-    } else {
-      golomb_bits = get_golomb_cost(level);
-      *diff += (r & (r - 1)) == 0 ? 1024 : 0;
-    }
+  const int r = level - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
+  if (r < 32) {
+    golomb_bits = golomb_bits_cost[r];
+    *diff += golomb_cost_diff[r];
+  } else {
+    golomb_bits = av1_cost_literal((get_msb(r) << 1) + 1);
+    *diff += (r & (r - 1)) == 0 ? 1024 : 0;
   }

-  return coeff_lps[base_range] + golomb_bits;
+  return coeff_lps[COEFF_BASE_RANGE] + golomb_bits;
 }

 static AOM_FORCE_INLINE int get_two_coeff_cost_simple(
@@ -129,7 +133,7 @@ static AOM_FORCE_INLINE int get_two_coeff_cost_simple(
   if (abs_qc) {
     cost += av1_cost_literal(1);
     if (abs_qc > NUM_BASE_LEVELS) {
-      const int br_ctx = get_br_ctx(levels, ci, bhl, tx_class);
+      const int br_ctx = get_br_ctx_simple(levels, ci, bhl, tx_class);
       int brcost_diff = 0;
       cost += get_br_cost_with_diff(abs_qc, txb_costs->lps_cost[br_ctx],
                                     &brcost_diff);