diff --git a/vp9/common/vp9_onyxc_int.h b/vp9/common/vp9_onyxc_int.h
index 57f9978f4d644f897ca8f52f66bb3ca8478189eb..2d4cd30ccd446f0b70be24815522b5a36c5d2204 100644
--- a/vp9/common/vp9_onyxc_int.h
+++ b/vp9/common/vp9_onyxc_int.h
@@ -306,6 +306,13 @@ static int mi_cols_aligned_to_sb(VP9_COMMON *cm) {
   return 2 * ((cm->mb_cols + 3) & ~3);
 }
 
+static INLINE void set_partition_seg_context(VP9_COMMON *cm,
+                                             MACROBLOCKD *xd,
+                                             int mi_row, int mi_col) {
+  xd->above_seg_context = cm->above_seg_context + mi_col;
+  xd->left_seg_context  = cm->left_seg_context + (mi_row & MI_MASK);
+}
+
 static void set_mi_row_col(VP9_COMMON *cm, MACROBLOCKD *xd,
                        int mi_row, int bh,
                        int mi_col, int bw) {
diff --git a/vp9/decoder/vp9_decodframe.c b/vp9/decoder/vp9_decodframe.c
index dc4ac632b479a639497d330e133bc7e1052f5256..ea5905bd8b4be04ff76eb3d96b9f87d9277e430b 100644
--- a/vp9/decoder/vp9_decodframe.c
+++ b/vp9/decoder/vp9_decodframe.c
@@ -476,12 +476,11 @@ static void decode_modes_sb(VP9D_COMP *pbi, int mi_row, int mi_col,
       assert(0);
   }
   // update partition context
-  if ((partition == PARTITION_SPLIT) && (bsize > BLOCK_SIZE_MB16X16))
-    return;
-
-  xd->left_seg_context = pc->left_seg_context + (mi_row & MI_MASK);
-  xd->above_seg_context = pc->above_seg_context + mi_col;
-  update_partition_context(xd, subsize, bsize);
+  if (bsize > BLOCK_SIZE_SB8X8 &&
+      (bsize == BLOCK_SIZE_MB16X16 || partition != PARTITION_SPLIT)) {
+    set_partition_seg_context(pc, xd, mi_row, mi_col);
+    update_partition_context(xd, subsize, bsize);
+  }
 }
 
 static void setup_token_decoder(VP9D_COMP *pbi,
diff --git a/vp9/encoder/vp9_bitstream.c b/vp9/encoder/vp9_bitstream.c
index 0c3d225c8879e339500fdf6c337e07155df3ce1d..9222133ce19e55d5a1c3c23707d226260934ae1f 100644
--- a/vp9/encoder/vp9_bitstream.c
+++ b/vp9/encoder/vp9_bitstream.c
@@ -915,12 +915,11 @@ static void write_modes_sb(VP9_COMP *cpi, MODE_INFO *m, vp9_writer *bc,
   }
 
   // update partition context
-  if ((partition == PARTITION_SPLIT) && (bsize > BLOCK_SIZE_MB16X16))
-    return;
-
-  xd->left_seg_context = cm->left_seg_context + (mi_row & MI_MASK);
-  xd->above_seg_context = cm->above_seg_context + mi_col;
-  update_partition_context(xd, subsize, bsize);
+  if (bsize > BLOCK_SIZE_SB8X8 &&
+      (bsize == BLOCK_SIZE_MB16X16 || partition != PARTITION_SPLIT)) {
+    set_partition_seg_context(cm, xd, mi_row, mi_col);
+    update_partition_context(xd, subsize, bsize);
+  }
 }
 
 static void write_modes(VP9_COMP *cpi, vp9_writer* const bc,
diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index 1f7b259a0bdc2645c8beea8f36d8bd12464b9bb2..2edeb780719675fdd3e2332d9887d941bef0655e 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -539,15 +539,6 @@ void vp9_setup_src_planes(MACROBLOCK *x,
                    x->e_mbd.plane[2].subsampling_y);
 }
 
-static INLINE void set_partition_seg_context(VP9_COMP *cpi,
-                                             int mi_row, int mi_col) {
-  VP9_COMMON *const cm = &cpi->common;
-  MACROBLOCKD *const xd = &cpi->mb.e_mbd;
-
-  xd->above_seg_context = cm->above_seg_context + mi_col;
-  xd->left_seg_context  = cm->left_seg_context + (mi_row & MI_MASK);
-}
-
 static void set_offsets(VP9_COMP *cpi,
                         int mi_row, int mi_col, BLOCK_SIZE_TYPE bsize) {
   MACROBLOCK *const x = &cpi->mb;
@@ -571,7 +562,7 @@ static void set_offsets(VP9_COMP *cpi,
   }
 
   // partition contexts
-  set_partition_seg_context(cpi, mi_row, mi_col);
+  set_partition_seg_context(cm, xd, mi_row, mi_col);
 
   // Activity map pointer
   x->mb_activity_ptr = &cpi->mb_activity_map[idx_map];
@@ -850,7 +841,7 @@ static void encode_sb(VP9_COMP *cpi, TOKENEXTRA **tp,
     return;
 
   if (bsize > BLOCK_SIZE_SB8X8) {
-    set_partition_seg_context(cpi, mi_row, mi_col);
+    set_partition_seg_context(cm, xd, mi_row, mi_col);
     pl = partition_plane_context(xd, bsize);
     c1 = *(get_sb_partitioning(x, bsize));
   }
@@ -899,7 +890,7 @@ static void encode_sb(VP9_COMP *cpi, TOKENEXTRA **tp,
 
   if (bsize > BLOCK_SIZE_SB8X8 &&
       (bsize == BLOCK_SIZE_MB16X16 || bsl == bwl || bsl == bhl)) {
-    set_partition_seg_context(cpi, mi_row, mi_col);
+    set_partition_seg_context(cm, xd, mi_row, mi_col);
     update_partition_context(xd, c1, bsize);
   }
 }
@@ -960,7 +951,7 @@ static void rd_pick_partition(VP9_COMP *cpi, TOKENEXTRA **tp,
       r4 += r;
       d4 += d;
     }
-    set_partition_seg_context(cpi, mi_row, mi_col);
+    set_partition_seg_context(cm, xd, mi_row, mi_col);
     pl = partition_plane_context(xd, bsize);
     r4 += x->partition_cost[pl][PARTITION_SPLIT];
 
@@ -992,7 +983,7 @@ static void rd_pick_partition(VP9_COMP *cpi, TOKENEXTRA **tp,
       if (mi_row + (ms >> 1) != cm->mi_rows)
         mb_skip = 1;
     }
-    set_partition_seg_context(cpi, mi_row, mi_col);
+    set_partition_seg_context(cm, xd, mi_row, mi_col);
     pl = partition_plane_context(xd, bsize);
     r2 += x->partition_cost[pl][PARTITION_HORZ];
 
@@ -1027,7 +1018,7 @@ static void rd_pick_partition(VP9_COMP *cpi, TOKENEXTRA **tp,
       if (mi_col + (ms >> 1) != cm->mi_cols)
         mb_skip = 1;
     }
-    set_partition_seg_context(cpi, mi_row, mi_col);
+    set_partition_seg_context(cm, xd, mi_row, mi_col);
     pl = partition_plane_context(xd, bsize);
     r2 += x->partition_cost[pl][PARTITION_VERT];
 
@@ -1046,7 +1037,7 @@ static void rd_pick_partition(VP9_COMP *cpi, TOKENEXTRA **tp,
     pick_sb_modes(cpi, mi_row, mi_col, tp, &r, &d, bsize,
                   get_block_context(x, bsize));
     if (bsize >= BLOCK_SIZE_MB16X16) {
-      set_partition_seg_context(cpi, mi_row, mi_col);
+      set_partition_seg_context(cm, xd, mi_row, mi_col);
       pl = partition_plane_context(xd, bsize);
       r += x->partition_cost[pl][PARTITION_NONE];
     }