Commit 70dd1176e5 for aom

commit 70dd1176e54184b05fd0b16b3d3941a61a46b3fc
Author: Diksha Singh <diksha.singh@ittiam.com>
Date:   Thu Oct 22 09:18:40 2026 +0530

    Speed up av1_optimize_txb()

    - get_lower_levels_ctx(): Return early when all the neighboring
      levels read by get_nz_mag() are zero.
    - get_lower_levels_ctx(): Use a precomputed lookup table for padded
      levels[] indices.
    - av1_get_ext_tx_set_type(): Replace conditional branches with a
      lookup table indexed by use_reduced_set, is_inter, and tx_size.
    - av1_get_adjusted_tx_size(): Replace the switch statement with
      a lookup table.
    - av1_get_tx_scale(): Replace the runtime computation with a lookup
      table.
    - av1_txb_init_levels_avx2(): Zero the bottom padding using up
      to five unrolled AVX2 stores instead of a loop, which GCC
      was converting into a memset() call.

    Encoder performance results averaged over all resolutions are as
    follows:
      Instruction Count
    cpu   Reduction(%)
     1       1.61
     2       1.33
     3       1.04

    This change is bit-exact across all presets.

    Change-Id: If9ee38ddec24c1163f253a6d7223439c300428c0

diff --git a/av1/common/blockd.h b/av1/common/blockd.h
index db912704d2..629a1de81d 100644
--- a/av1/common/blockd.h
+++ b/av1/common/blockd.h
@@ -1094,21 +1094,48 @@ static const uint16_t av1_ext_tx_used_flag[EXT_TX_SET_TYPES] = {
   0xFFFF,  // 1111 1111 1111 1111
 };

-static const TxSetType av1_ext_tx_set_lookup[2][2] = {
-  { EXT_TX_SET_DTT4_IDTX_1DDCT, EXT_TX_SET_DTT4_IDTX },
-  { EXT_TX_SET_ALL16, EXT_TX_SET_DTT9_IDTX_1DDCT },
+// av1_ext_tx_set_lookup[use_reduced_set][is_inter][tx_size] (EXT_TX_SET_
+// prefix omitted):
+//
+// use_reduced_set = 0
+// TX_SIZE                         | is_inter = 0    | is_inter = 1
+// 4X4, 8X8, 4X8, 8X4, 8X16, 16X8, | DTT4_IDTX_1DDCT | ALL16
+// 4X16, 16X4                      |                 |
+// 16X16                           | DTT4_IDTX       | DTT9_IDTX_1DDCT
+// 32X32, 16X32, 32X16, 8X32, 32X8 | DCTONLY         | DCT_IDTX
+// 64X64, 32X64, 64X32, 16X64,     | DCTONLY         | DCTONLY
+// 64X16                           |                 |
+//
+// use_reduced_set = 1
+// TX_SIZE                         | is_inter = 0    | is_inter = 1
+// 4X4, 8X8, 4X8, 8X4, 8X16, 16X8, | DTT4_IDTX       | DCT_IDTX
+// 4X16, 16X4, 16X16               |                 |
+// 32X32, 16X32, 32X16, 8X32, 32X8 | DCTONLY         | DCT_IDTX
+// 64X64, 32X64, 64X32, 16X64,     | DCTONLY         | DCTONLY
+// 64X16                           |                 |
+//
+// EXT_TX_SET_LIST() expands one row of the table above to all TX_SIZES_ALL
+// entries. Its arguments are the sets for the four groups of sizes, in the
+// order of the rows above: s4, s16 (16X16), s32 and s64.
+#define EXT_TX_SET_LIST(s4, s16, s32, s64)           \
+  { s4,  s4,  s16, s32, s64, s4,  s4,  s4,  s4, s32, \
+    s32, s64, s64, s4,  s4,  s32, s32, s64, s64 }
+static const TxSetType av1_ext_tx_set_lookup[2][2][TX_SIZES_ALL] = {
+  { EXT_TX_SET_LIST(EXT_TX_SET_DTT4_IDTX_1DDCT, EXT_TX_SET_DTT4_IDTX,
+                    EXT_TX_SET_DCTONLY, EXT_TX_SET_DCTONLY),
+    EXT_TX_SET_LIST(EXT_TX_SET_ALL16, EXT_TX_SET_DTT9_IDTX_1DDCT,
+                    EXT_TX_SET_DCT_IDTX, EXT_TX_SET_DCTONLY) },
+  { EXT_TX_SET_LIST(EXT_TX_SET_DTT4_IDTX, EXT_TX_SET_DTT4_IDTX,
+                    EXT_TX_SET_DCTONLY, EXT_TX_SET_DCTONLY),
+    EXT_TX_SET_LIST(EXT_TX_SET_DCT_IDTX, EXT_TX_SET_DCT_IDTX,
+                    EXT_TX_SET_DCT_IDTX, EXT_TX_SET_DCTONLY) },
 };
+#undef EXT_TX_SET_LIST

 static inline TxSetType av1_get_ext_tx_set_type(TX_SIZE tx_size, int is_inter,
                                                 int use_reduced_set) {
-  const TX_SIZE tx_size_sqr_up = txsize_sqr_up_map[tx_size];
-  if (tx_size_sqr_up > TX_32X32) return EXT_TX_SET_DCTONLY;
-  if (tx_size_sqr_up == TX_32X32)
-    return is_inter ? EXT_TX_SET_DCT_IDTX : EXT_TX_SET_DCTONLY;
-  if (use_reduced_set)
-    return is_inter ? EXT_TX_SET_DCT_IDTX : EXT_TX_SET_DTT4_IDTX;
-  const TX_SIZE tx_size_sqr = txsize_sqr_map[tx_size];
-  return av1_ext_tx_set_lookup[is_inter][tx_size_sqr == TX_16X16];
+  assert(tx_size < TX_SIZES_ALL);
+  return av1_ext_tx_set_lookup[use_reduced_set != 0][is_inter != 0][tx_size];
 }

 // Maps tx set types to the indices.
@@ -1364,14 +1391,13 @@ static inline TX_SIZE depth_to_tx_size(int depth, BLOCK_SIZE bsize) {
 }

 static inline TX_SIZE av1_get_adjusted_tx_size(TX_SIZE tx_size) {
-  switch (tx_size) {
-    case TX_64X64:
-    case TX_64X32:
-    case TX_32X64: return TX_32X32;
-    case TX_64X16: return TX_32X16;
-    case TX_16X64: return TX_16X32;
-    default: return tx_size;
-  }
+  static const TX_SIZE adjusted_tx_size[TX_SIZES_ALL] = {
+    TX_4X4,  TX_8X8,  TX_16X16, TX_32X32, TX_32X32, TX_4X8,   TX_8X4,
+    TX_8X16, TX_16X8, TX_16X32, TX_32X16, TX_32X32, TX_32X32, TX_4X16,
+    TX_16X4, TX_8X32, TX_32X8,  TX_16X32, TX_32X16,
+  };
+  assert(tx_size < TX_SIZES_ALL);
+  return adjusted_tx_size[tx_size];
 }

 static inline TX_SIZE av1_get_max_uv_txsize(BLOCK_SIZE bsize, int subsampling_x,
diff --git a/av1/common/idct.c b/av1/common/idct.c
index bfd7d5c1c7..7c8e3d6e5a 100644
--- a/av1/common/idct.c
+++ b/av1/common/idct.c
@@ -21,12 +21,6 @@
 #include "av1/common/enums.h"
 #include "av1/common/idct.h"

-int av1_get_tx_scale(const TX_SIZE tx_size) {
-  const int pels = tx_size_2d[tx_size];
-  // Largest possible pels is 4096 (64x64).
-  return (pels > 256) + (pels > 1024);
-}
-
 // NOTE: The implementation of all inverses need to be aware of the fact
 // that input and output could be the same buffer.

diff --git a/av1/common/idct.h b/av1/common/idct.h
index 2dda093a68..3f4c98af6e 100644
--- a/av1/common/idct.h
+++ b/av1/common/idct.h
@@ -32,7 +32,13 @@ typedef struct {
 } transform_2d;

 #define MAX_TX_SCALE 1
-int av1_get_tx_scale(const TX_SIZE tx_size);
+static inline int av1_get_tx_scale(const TX_SIZE tx_size) {
+  // (pels > 256) + (pels > 1024), where pels = tx_size_2d[tx_size].
+  static const uint8_t tx_scale[TX_SIZES_ALL] = { 0, 0, 0, 1, 2, 0, 0, 0, 0, 1,
+                                                  1, 2, 2, 0, 0, 0, 0, 1, 1 };
+  assert(tx_size < TX_SIZES_ALL);
+  return tx_scale[tx_size];
+}

 void av1_inverse_transform_block(const MACROBLOCKD *xd,
                                  const tran_low_t *dqcoeff, int plane,
diff --git a/av1/common/txb_common.c b/av1/common/txb_common.c
index 75dc36e223..219e89332f 100644
--- a/av1/common/txb_common.c
+++ b/av1/common/txb_common.c
@@ -12,6 +12,165 @@
 #include "av1/common/av1_common_int.h"
 #include "av1/common/txb_common.h"

+// Maps a coefficient index to its index in the padded levels[] buffer,
+// which has TX_PAD_HOR bytes of padding after each column.
+static const int16_t padded_idx_h4[128] = {
+  0,   1,   2,   3,   8,   9,   10,  11,  16,  17,  18,  19,  24,  25,  26,
+  27,  32,  33,  34,  35,  40,  41,  42,  43,  48,  49,  50,  51,  56,  57,
+  58,  59,  64,  65,  66,  67,  72,  73,  74,  75,  80,  81,  82,  83,  88,
+  89,  90,  91,  96,  97,  98,  99,  104, 105, 106, 107, 112, 113, 114, 115,
+  120, 121, 122, 123, 128, 129, 130, 131, 136, 137, 138, 139, 144, 145, 146,
+  147, 152, 153, 154, 155, 160, 161, 162, 163, 168, 169, 170, 171, 176, 177,
+  178, 179, 184, 185, 186, 187, 192, 193, 194, 195, 200, 201, 202, 203, 208,
+  209, 210, 211, 216, 217, 218, 219, 224, 225, 226, 227, 232, 233, 234, 235,
+  240, 241, 242, 243, 248, 249, 250, 251,
+};
+
+static const int16_t padded_idx_h8[256] = {
+  0,   1,   2,   3,   4,   5,   6,   7,   12,  13,  14,  15,  16,  17,  18,
+  19,  24,  25,  26,  27,  28,  29,  30,  31,  36,  37,  38,  39,  40,  41,
+  42,  43,  48,  49,  50,  51,  52,  53,  54,  55,  60,  61,  62,  63,  64,
+  65,  66,  67,  72,  73,  74,  75,  76,  77,  78,  79,  84,  85,  86,  87,
+  88,  89,  90,  91,  96,  97,  98,  99,  100, 101, 102, 103, 108, 109, 110,
+  111, 112, 113, 114, 115, 120, 121, 122, 123, 124, 125, 126, 127, 132, 133,
+  134, 135, 136, 137, 138, 139, 144, 145, 146, 147, 148, 149, 150, 151, 156,
+  157, 158, 159, 160, 161, 162, 163, 168, 169, 170, 171, 172, 173, 174, 175,
+  180, 181, 182, 183, 184, 185, 186, 187, 192, 193, 194, 195, 196, 197, 198,
+  199, 204, 205, 206, 207, 208, 209, 210, 211, 216, 217, 218, 219, 220, 221,
+  222, 223, 228, 229, 230, 231, 232, 233, 234, 235, 240, 241, 242, 243, 244,
+  245, 246, 247, 252, 253, 254, 255, 256, 257, 258, 259, 264, 265, 266, 267,
+  268, 269, 270, 271, 276, 277, 278, 279, 280, 281, 282, 283, 288, 289, 290,
+  291, 292, 293, 294, 295, 300, 301, 302, 303, 304, 305, 306, 307, 312, 313,
+  314, 315, 316, 317, 318, 319, 324, 325, 326, 327, 328, 329, 330, 331, 336,
+  337, 338, 339, 340, 341, 342, 343, 348, 349, 350, 351, 352, 353, 354, 355,
+  360, 361, 362, 363, 364, 365, 366, 367, 372, 373, 374, 375, 376, 377, 378,
+  379,
+};
+
+static const int16_t padded_idx_h16[512] = {
+  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,   10,  11,  12,  13,  14,
+  15,  20,  21,  22,  23,  24,  25,  26,  27,  28,  29,  30,  31,  32,  33,
+  34,  35,  40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,
+  53,  54,  55,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,  70,  71,
+  72,  73,  74,  75,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,
+  91,  92,  93,  94,  95,  100, 101, 102, 103, 104, 105, 106, 107, 108, 109,
+  110, 111, 112, 113, 114, 115, 120, 121, 122, 123, 124, 125, 126, 127, 128,
+  129, 130, 131, 132, 133, 134, 135, 140, 141, 142, 143, 144, 145, 146, 147,
+  148, 149, 150, 151, 152, 153, 154, 155, 160, 161, 162, 163, 164, 165, 166,
+  167, 168, 169, 170, 171, 172, 173, 174, 175, 180, 181, 182, 183, 184, 185,
+  186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 200, 201, 202, 203, 204,
+  205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 220, 221, 222, 223,
+  224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 240, 241, 242,
+  243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 260, 261,
+  262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 280,
+  281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295,
+  300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314,
+  315, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333,
+  334, 335, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352,
+  353, 354, 355, 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371,
+  372, 373, 374, 375, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390,
+  391, 392, 393, 394, 395, 400, 401, 402, 403, 404, 405, 406, 407, 408, 409,
+  410, 411, 412, 413, 414, 415, 420, 421, 422, 423, 424, 425, 426, 427, 428,
+  429, 430, 431, 432, 433, 434, 435, 440, 441, 442, 443, 444, 445, 446, 447,
+  448, 449, 450, 451, 452, 453, 454, 455, 460, 461, 462, 463, 464, 465, 466,
+  467, 468, 469, 470, 471, 472, 473, 474, 475, 480, 481, 482, 483, 484, 485,
+  486, 487, 488, 489, 490, 491, 492, 493, 494, 495, 500, 501, 502, 503, 504,
+  505, 506, 507, 508, 509, 510, 511, 512, 513, 514, 515, 520, 521, 522, 523,
+  524, 525, 526, 527, 528, 529, 530, 531, 532, 533, 534, 535, 540, 541, 542,
+  543, 544, 545, 546, 547, 548, 549, 550, 551, 552, 553, 554, 555, 560, 561,
+  562, 563, 564, 565, 566, 567, 568, 569, 570, 571, 572, 573, 574, 575, 580,
+  581, 582, 583, 584, 585, 586, 587, 588, 589, 590, 591, 592, 593, 594, 595,
+  600, 601, 602, 603, 604, 605, 606, 607, 608, 609, 610, 611, 612, 613, 614,
+  615, 620, 621, 622, 623, 624, 625, 626, 627, 628, 629, 630, 631, 632, 633,
+  634, 635,
+};
+
+static const int16_t padded_idx_h32[1024] = {
+  0,    1,    2,    3,    4,    5,    6,    7,    8,    9,    10,   11,   12,
+  13,   14,   15,   16,   17,   18,   19,   20,   21,   22,   23,   24,   25,
+  26,   27,   28,   29,   30,   31,   36,   37,   38,   39,   40,   41,   42,
+  43,   44,   45,   46,   47,   48,   49,   50,   51,   52,   53,   54,   55,
+  56,   57,   58,   59,   60,   61,   62,   63,   64,   65,   66,   67,   72,
+  73,   74,   75,   76,   77,   78,   79,   80,   81,   82,   83,   84,   85,
+  86,   87,   88,   89,   90,   91,   92,   93,   94,   95,   96,   97,   98,
+  99,   100,  101,  102,  103,  108,  109,  110,  111,  112,  113,  114,  115,
+  116,  117,  118,  119,  120,  121,  122,  123,  124,  125,  126,  127,  128,
+  129,  130,  131,  132,  133,  134,  135,  136,  137,  138,  139,  144,  145,
+  146,  147,  148,  149,  150,  151,  152,  153,  154,  155,  156,  157,  158,
+  159,  160,  161,  162,  163,  164,  165,  166,  167,  168,  169,  170,  171,
+  172,  173,  174,  175,  180,  181,  182,  183,  184,  185,  186,  187,  188,
+  189,  190,  191,  192,  193,  194,  195,  196,  197,  198,  199,  200,  201,
+  202,  203,  204,  205,  206,  207,  208,  209,  210,  211,  216,  217,  218,
+  219,  220,  221,  222,  223,  224,  225,  226,  227,  228,  229,  230,  231,
+  232,  233,  234,  235,  236,  237,  238,  239,  240,  241,  242,  243,  244,
+  245,  246,  247,  252,  253,  254,  255,  256,  257,  258,  259,  260,  261,
+  262,  263,  264,  265,  266,  267,  268,  269,  270,  271,  272,  273,  274,
+  275,  276,  277,  278,  279,  280,  281,  282,  283,  288,  289,  290,  291,
+  292,  293,  294,  295,  296,  297,  298,  299,  300,  301,  302,  303,  304,
+  305,  306,  307,  308,  309,  310,  311,  312,  313,  314,  315,  316,  317,
+  318,  319,  324,  325,  326,  327,  328,  329,  330,  331,  332,  333,  334,
+  335,  336,  337,  338,  339,  340,  341,  342,  343,  344,  345,  346,  347,
+  348,  349,  350,  351,  352,  353,  354,  355,  360,  361,  362,  363,  364,
+  365,  366,  367,  368,  369,  370,  371,  372,  373,  374,  375,  376,  377,
+  378,  379,  380,  381,  382,  383,  384,  385,  386,  387,  388,  389,  390,
+  391,  396,  397,  398,  399,  400,  401,  402,  403,  404,  405,  406,  407,
+  408,  409,  410,  411,  412,  413,  414,  415,  416,  417,  418,  419,  420,
+  421,  422,  423,  424,  425,  426,  427,  432,  433,  434,  435,  436,  437,
+  438,  439,  440,  441,  442,  443,  444,  445,  446,  447,  448,  449,  450,
+  451,  452,  453,  454,  455,  456,  457,  458,  459,  460,  461,  462,  463,
+  468,  469,  470,  471,  472,  473,  474,  475,  476,  477,  478,  479,  480,
+  481,  482,  483,  484,  485,  486,  487,  488,  489,  490,  491,  492,  493,
+  494,  495,  496,  497,  498,  499,  504,  505,  506,  507,  508,  509,  510,
+  511,  512,  513,  514,  515,  516,  517,  518,  519,  520,  521,  522,  523,
+  524,  525,  526,  527,  528,  529,  530,  531,  532,  533,  534,  535,  540,
+  541,  542,  543,  544,  545,  546,  547,  548,  549,  550,  551,  552,  553,
+  554,  555,  556,  557,  558,  559,  560,  561,  562,  563,  564,  565,  566,
+  567,  568,  569,  570,  571,  576,  577,  578,  579,  580,  581,  582,  583,
+  584,  585,  586,  587,  588,  589,  590,  591,  592,  593,  594,  595,  596,
+  597,  598,  599,  600,  601,  602,  603,  604,  605,  606,  607,  612,  613,
+  614,  615,  616,  617,  618,  619,  620,  621,  622,  623,  624,  625,  626,
+  627,  628,  629,  630,  631,  632,  633,  634,  635,  636,  637,  638,  639,
+  640,  641,  642,  643,  648,  649,  650,  651,  652,  653,  654,  655,  656,
+  657,  658,  659,  660,  661,  662,  663,  664,  665,  666,  667,  668,  669,
+  670,  671,  672,  673,  674,  675,  676,  677,  678,  679,  684,  685,  686,
+  687,  688,  689,  690,  691,  692,  693,  694,  695,  696,  697,  698,  699,
+  700,  701,  702,  703,  704,  705,  706,  707,  708,  709,  710,  711,  712,
+  713,  714,  715,  720,  721,  722,  723,  724,  725,  726,  727,  728,  729,
+  730,  731,  732,  733,  734,  735,  736,  737,  738,  739,  740,  741,  742,
+  743,  744,  745,  746,  747,  748,  749,  750,  751,  756,  757,  758,  759,
+  760,  761,  762,  763,  764,  765,  766,  767,  768,  769,  770,  771,  772,
+  773,  774,  775,  776,  777,  778,  779,  780,  781,  782,  783,  784,  785,
+  786,  787,  792,  793,  794,  795,  796,  797,  798,  799,  800,  801,  802,
+  803,  804,  805,  806,  807,  808,  809,  810,  811,  812,  813,  814,  815,
+  816,  817,  818,  819,  820,  821,  822,  823,  828,  829,  830,  831,  832,
+  833,  834,  835,  836,  837,  838,  839,  840,  841,  842,  843,  844,  845,
+  846,  847,  848,  849,  850,  851,  852,  853,  854,  855,  856,  857,  858,
+  859,  864,  865,  866,  867,  868,  869,  870,  871,  872,  873,  874,  875,
+  876,  877,  878,  879,  880,  881,  882,  883,  884,  885,  886,  887,  888,
+  889,  890,  891,  892,  893,  894,  895,  900,  901,  902,  903,  904,  905,
+  906,  907,  908,  909,  910,  911,  912,  913,  914,  915,  916,  917,  918,
+  919,  920,  921,  922,  923,  924,  925,  926,  927,  928,  929,  930,  931,
+  936,  937,  938,  939,  940,  941,  942,  943,  944,  945,  946,  947,  948,
+  949,  950,  951,  952,  953,  954,  955,  956,  957,  958,  959,  960,  961,
+  962,  963,  964,  965,  966,  967,  972,  973,  974,  975,  976,  977,  978,
+  979,  980,  981,  982,  983,  984,  985,  986,  987,  988,  989,  990,  991,
+  992,  993,  994,  995,  996,  997,  998,  999,  1000, 1001, 1002, 1003, 1008,
+  1009, 1010, 1011, 1012, 1013, 1014, 1015, 1016, 1017, 1018, 1019, 1020, 1021,
+  1022, 1023, 1024, 1025, 1026, 1027, 1028, 1029, 1030, 1031, 1032, 1033, 1034,
+  1035, 1036, 1037, 1038, 1039, 1044, 1045, 1046, 1047, 1048, 1049, 1050, 1051,
+  1052, 1053, 1054, 1055, 1056, 1057, 1058, 1059, 1060, 1061, 1062, 1063, 1064,
+  1065, 1066, 1067, 1068, 1069, 1070, 1071, 1072, 1073, 1074, 1075, 1080, 1081,
+  1082, 1083, 1084, 1085, 1086, 1087, 1088, 1089, 1090, 1091, 1092, 1093, 1094,
+  1095, 1096, 1097, 1098, 1099, 1100, 1101, 1102, 1103, 1104, 1105, 1106, 1107,
+  1108, 1109, 1110, 1111, 1116, 1117, 1118, 1119, 1120, 1121, 1122, 1123, 1124,
+  1125, 1126, 1127, 1128, 1129, 1130, 1131, 1132, 1133, 1134, 1135, 1136, 1137,
+  1138, 1139, 1140, 1141, 1142, 1143, 1144, 1145, 1146, 1147,
+};
+
+const int16_t *const av1_txb_padded_idx[MAX_TX_SIZE_LOG2 - MIN_TX_SIZE_LOG2] = {
+  padded_idx_h4, padded_idx_h8, padded_idx_h16, padded_idx_h32
+};
+
 // The ctx offset table when TX is TX_CLASS_2D.
 // TX col and row indices are clamped to 4

diff --git a/av1/common/txb_common.h b/av1/common/txb_common.h
index 3325805058..a8a34b00b2 100644
--- a/av1/common/txb_common.h
+++ b/av1/common/txb_common.h
@@ -19,6 +19,10 @@ extern const int16_t av1_eob_offset_bits[12];

 extern const int8_t *const av1_nz_map_ctx_offset[TX_SIZES_ALL];

+// Indexed by [bhl - MIN_TX_SIZE_LOG2][coeff_idx].
+extern const int16_t *const
+    av1_txb_padded_idx[MAX_TX_SIZE_LOG2 - MIN_TX_SIZE_LOG2];
+
 typedef struct txb_ctx {
   int txb_skip_ctx;
   int dc_sign_ctx;
@@ -67,7 +71,7 @@ static inline uint8_t *set_levels(uint8_t *const levels_buf, const int height) {
 }

 static inline int get_padded_idx(const int idx, const int bhl) {
-  return idx + ((idx >> bhl) << TX_PAD_HOR_LOG2);
+  return av1_txb_padded_idx[bhl - MIN_TX_SIZE_LOG2][idx];
 }

 static inline int get_br_ctx_2d(const uint8_t *const levels,
@@ -253,8 +257,24 @@ static AOM_FORCE_INLINE int get_lower_levels_ctx(const uint8_t *levels,
                                                  int coeff_idx, int bhl,
                                                  TX_SIZE tx_size,
                                                  TX_CLASS tx_class) {
-  const int stats =
-      get_nz_mag(levels + get_padded_idx(coeff_idx, bhl), bhl, tx_class);
+  assert(bhl >= MIN_TX_SIZE_LOG2 && bhl <= MAX_TX_SIZE_LOG2 - 1);
+  const uint8_t *const nb = levels + get_padded_idx(coeff_idx, bhl);
+  const int stride = (1 << bhl) + TX_PAD_HOR;
+  // When all the neighbors read by get_nz_mag() are zero, stats is 0 and the
+  // context is the offset alone.
+  if (tx_class == TX_CLASS_2D) {
+    if (coeff_idx == 0) return 0;
+    if ((nb[1] | nb[2] | nb[stride] | nb[stride + 1] | nb[2 * stride]) == 0)
+      return av1_nz_map_ctx_offset[tx_size][coeff_idx];
+  } else if (tx_class == TX_CLASS_VERT) {
+    if ((nb[stride] | nb[1] | nb[2] | nb[3] | nb[4]) == 0)
+      return nz_map_ctx_offset_1d[coeff_idx & ((1 << bhl) - 1)];
+  } else {
+    if ((nb[stride] | nb[1] | nb[2 * stride] | nb[3 * stride] |
+         nb[4 * stride]) == 0)
+      return nz_map_ctx_offset_1d[coeff_idx >> bhl];
+  }
+  const int stats = get_nz_mag(nb, bhl, tx_class);
   return get_nz_map_ctx_from_stats(stats, coeff_idx, bhl, tx_size, tx_class);
 }

diff --git a/av1/encoder/x86/encodetxb_avx2.c b/av1/encoder/x86/encodetxb_avx2.c
index cee9f121af..ce61df7ffa 100644
--- a/av1/encoder/x86/encodetxb_avx2.c
+++ b/av1/encoder/x86/encodetxb_avx2.c
@@ -26,14 +26,16 @@ void av1_txb_init_levels_avx2(const tran_low_t *const coeff, const int width,
   const int stride = height + TX_PAD_HOR;
   const __m256i y_zeros = _mm256_setzero_si256();

+  // Zero bottom padding using up to five unrolled stores.
   const int32_t bottom_len = sizeof(*levels) * (TX_PAD_BOTTOM * stride);
-  uint8_t *bottom_buf_end = levels + (width + TX_PAD_BOTTOM) * stride;
-  uint8_t *bottom_buf = bottom_buf_end - ((bottom_len + 31) & (~31));
-
-  do {
-    yy_storeu_256(bottom_buf, y_zeros);
-    bottom_buf += 32;
-  } while (bottom_buf < bottom_buf_end);
+  uint8_t *const bottom_buf_end = levels + (width + TX_PAD_BOTTOM) * stride;
+  const int num_stores = (bottom_len + 31) >> 5;
+  assert(num_stores >= 1 && num_stores <= 5);
+  yy_storeu_256(bottom_buf_end - 32, y_zeros);
+  if (num_stores > 1) yy_storeu_256(bottom_buf_end - 64, y_zeros);
+  if (num_stores > 2) yy_storeu_256(bottom_buf_end - 96, y_zeros);
+  if (num_stores > 3) yy_storeu_256(bottom_buf_end - 128, y_zeros);
+  if (num_stores > 4) yy_storeu_256(bottom_buf_end - 160, y_zeros);

   int i = 0;
   uint8_t *ls = levels;