diff --git a/vp9/common/vp9_blockd.h b/vp9/common/vp9_blockd.h
index 99272769a..1ad2979f2 100644
--- a/vp9/common/vp9_blockd.h
+++ b/vp9/common/vp9_blockd.h
@@ -686,46 +686,39 @@ static void extend_for_intra(MACROBLOCKD* const xd, int plane, int block,
   }
 }
 static void set_contexts_on_border(MACROBLOCKD *xd, BLOCK_SIZE_TYPE bsize,
-                                   int plane, int ss_tx_size, int eob, int aoff,
-                                   int loff, ENTROPY_CONTEXT *A,
-                                   ENTROPY_CONTEXT *L) {
-  const int bw = b_width_log2(bsize), bh = b_height_log2(bsize);
-  const int sw = bw - xd->plane[plane].subsampling_x;
-  const int sh = bh - xd->plane[plane].subsampling_y;
-  int mi_blocks_wide = 1 << sw;
-  int mi_blocks_high = 1 << sh;
-  int tx_size_in_blocks = (1 << ss_tx_size);
+                                   int plane, int tx_size_in_blocks,
+                                   int eob, int aoff, int loff,
+                                   ENTROPY_CONTEXT *A, ENTROPY_CONTEXT *L) {
+  struct macroblockd_plane *pd = &xd->plane[plane];
   int above_contexts = tx_size_in_blocks;
   int left_contexts = tx_size_in_blocks;
+  int mi_blocks_wide = 1 << plane_block_width_log2by4(bsize, pd);
+  int mi_blocks_high = 1 << plane_block_height_log2by4(bsize, pd);
   int pt;
 
   // xd->mb_to_right_edge is in units of pixels * 8.  This converts
   // it to 4x4 block sizes.
-  if (xd->mb_to_right_edge < 0) {
-    mi_blocks_wide += (xd->mb_to_right_edge
-        >> (5 + xd->plane[plane].subsampling_x));
-  }
+  if (xd->mb_to_right_edge < 0)
+    mi_blocks_wide += (xd->mb_to_right_edge >> (5 + pd->subsampling_x));
 
   // this code attempts to avoid copying into contexts that are outside
   // our border.  Any blocks that do are set to 0...
   if (above_contexts + aoff > mi_blocks_wide)
     above_contexts = mi_blocks_wide - aoff;
 
-  if (xd->mb_to_bottom_edge < 0) {
-    mi_blocks_high += (xd->mb_to_bottom_edge
-        >> (5 + xd->plane[plane].subsampling_y));
-  }
-  if (left_contexts + loff > mi_blocks_high) {
+  if (xd->mb_to_bottom_edge < 0)
+    mi_blocks_high += (xd->mb_to_bottom_edge >> (5 + pd->subsampling_y));
+
+  if (left_contexts + loff > mi_blocks_high)
     left_contexts = mi_blocks_high - loff;
-  }
 
   for (pt = 0; pt < above_contexts; pt++)
     A[pt] = eob > 0;
-  for (pt = above_contexts; pt < (1 << ss_tx_size); pt++)
+  for (pt = above_contexts; pt < tx_size_in_blocks; pt++)
     A[pt] = 0;
   for (pt = 0; pt < left_contexts; pt++)
     L[pt] = eob > 0;
-  for (pt = left_contexts; pt < (1 << ss_tx_size); pt++)
+  for (pt = left_contexts; pt < tx_size_in_blocks; pt++)
     L[pt] = 0;
 }
 
diff --git a/vp9/decoder/vp9_detokenize.c b/vp9/decoder/vp9_detokenize.c
index e004c083c..56786dc3b 100644
--- a/vp9/decoder/vp9_detokenize.c
+++ b/vp9/decoder/vp9_detokenize.c
@@ -269,7 +269,7 @@ static void decode_block(int plane, int block,
   const int mod = bw - ss_tx_size - pd->subsampling_x;
   const int aoff = (off & ((1 << mod) - 1)) << ss_tx_size;
   const int loff = (off >> mod) << ss_tx_size;
-
+  const int tx_size_in_blocks = 1 << ss_tx_size;
   ENTROPY_CONTEXT *A = pd->above_context + aoff;
   ENTROPY_CONTEXT *L = pd->left_context + loff;
   const int eob = decode_coefs(&arg->pbi->common, xd, arg->r, block,
@@ -278,10 +278,11 @@ static void decode_block(int plane, int block,
                                ss_tx_size, pd->dequant, A, L);
 
   if (xd->mb_to_right_edge < 0 || xd->mb_to_bottom_edge < 0) {
-    set_contexts_on_border(xd, bsize, plane, ss_tx_size, eob, aoff, loff, A, L);
+    set_contexts_on_border(xd, bsize, plane, tx_size_in_blocks, eob, aoff, loff,
+                           A, L);
   } else {
     int pt;
-    for (pt = 0; pt < (1 << ss_tx_size); pt++)
+    for (pt = 0; pt < tx_size_in_blocks; pt++)
       A[pt] = L[pt] = eob > 0;
   }
   pd->eobs[block] = eob;
diff --git a/vp9/encoder/vp9_tokenize.c b/vp9/encoder/vp9_tokenize.c
index 0a42a32ae..6b2be04eb 100644
--- a/vp9/encoder/vp9_tokenize.c
+++ b/vp9/encoder/vp9_tokenize.c
@@ -110,12 +110,14 @@ static void set_entropy_context_b(int plane, int block, BLOCK_SIZE_TYPE bsize,
   ENTROPY_CONTEXT *A = xd->plane[plane].above_context + aoff;
   ENTROPY_CONTEXT *L = xd->plane[plane].left_context + loff;
   const int eob = xd->plane[plane].eobs[block];
+  const int tx_size_in_blocks = 1 << tx_size;
 
   if (xd->mb_to_right_edge < 0 || xd->mb_to_bottom_edge < 0) {
-    set_contexts_on_border(xd, bsize, plane, tx_size, eob, aoff, loff, A, L);
+    set_contexts_on_border(xd, bsize, plane, tx_size_in_blocks, eob, aoff, loff,
+                           A, L);
   } else {
-    vpx_memset(A, eob > 0, sizeof(ENTROPY_CONTEXT) * (1 << tx_size));
-    vpx_memset(L, eob > 0, sizeof(ENTROPY_CONTEXT) * (1 << tx_size));
+    vpx_memset(A, eob > 0, sizeof(ENTROPY_CONTEXT) * tx_size_in_blocks);
+    vpx_memset(L, eob > 0, sizeof(ENTROPY_CONTEXT) * tx_size_in_blocks);
   }
 }
 
@@ -125,7 +127,8 @@ static void tokenize_b(int plane, int block, BLOCK_SIZE_TYPE bsize,
   VP9_COMP *cpi = args->cpi;
   MACROBLOCKD *xd = args->xd;
   TOKENEXTRA **tp = args->tp;
-  TX_SIZE tx_size = ss_txfrm_size >> 1;
+  const TX_SIZE tx_size = ss_txfrm_size >> 1;
+  const int tx_size_in_blocks = 1 << tx_size;
   MB_MODE_INFO *mbmi = &xd->mode_info_context->mbmi;
   int pt; /* near block/prev token context index */
   int c = 0, rc = 0;
@@ -225,10 +228,11 @@ static void tokenize_b(int plane, int block, BLOCK_SIZE_TYPE bsize,
 
   *tp = t;
   if (xd->mb_to_right_edge < 0 || xd->mb_to_bottom_edge < 0) {
-    set_contexts_on_border(xd, bsize, plane, tx_size, c, aoff, loff, A, L);
+    set_contexts_on_border(xd, bsize, plane, tx_size_in_blocks, c, aoff, loff,
+                           A, L);
   } else {
-    vpx_memset(A, c > 0, sizeof(ENTROPY_CONTEXT) * (1 << tx_size));
-    vpx_memset(L, c > 0, sizeof(ENTROPY_CONTEXT) * (1 << tx_size));
+    vpx_memset(A, c > 0, sizeof(ENTROPY_CONTEXT) * tx_size_in_blocks);
+    vpx_memset(L, c > 0, sizeof(ENTROPY_CONTEXT) * tx_size_in_blocks);
   }
 }