1 files changed, 27 insertions, 20 deletions

diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index f3a03f3c8..b6f611c87 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -503,21 +503,19 @@ static unsigned find_seg_id(uint8_t *buf, BLOCK_SIZE_TYPE bsize,
 void vp9_setup_src_planes(MACROBLOCK *x,
                           const YV12_BUFFER_CONFIG *src,
                           int mb_row, int mb_col) {
-  setup_pred_plane(&x->plane[0].src,
-                   src->y_buffer, src->y_stride,
-                   mb_row, mb_col, NULL,
-                   x->e_mbd.plane[0].subsampling_x,
-                   x->e_mbd.plane[0].subsampling_y);
-  setup_pred_plane(&x->plane[1].src,
-                   src->u_buffer, src->uv_stride,
-                   mb_row, mb_col, NULL,
-                   x->e_mbd.plane[1].subsampling_x,
-                   x->e_mbd.plane[1].subsampling_y);
-  setup_pred_plane(&x->plane[2].src,
-                   src->v_buffer, src->uv_stride,
-                   mb_row, mb_col, NULL,
-                   x->e_mbd.plane[2].subsampling_x,
-                   x->e_mbd.plane[2].subsampling_y);
+  uint8_t *buffers[4] = {src->y_buffer, src->u_buffer, src->v_buffer,
+                         src->alpha_buffer};
+  int strides[4] = {src->y_stride, src->uv_stride, src->uv_stride,
+                    src->alpha_stride};
+  int i;
+
+  for (i = 0; i < MAX_MB_PLANE; i++) {
+    setup_pred_plane(&x->plane[i].src,
+                     buffers[i], strides[i],
+                     mb_row, mb_col, NULL,
+                     x->e_mbd.plane[i].subsampling_x,
+                     x->e_mbd.plane[i].subsampling_y);
+  }
 }
 
 static void set_offsets(VP9_COMP *cpi,
@@ -1614,11 +1612,20 @@ static void sum_intra_stats(VP9_COMP *cpi, MACROBLOCK *x) {
   }
     ++cpi->y_uv_mode_count[m][uvm];
   if (m == I4X4_PRED) {
-    int b = 0;
-    do {
-      int m = xd->mode_info_context->bmi[b].as_mode.first;
-      ++cpi->bmode_count[m];
-    } while (++b < 4);
+    int idx, idy;
+    int bw = 1 << b_width_log2(xd->mode_info_context->mbmi.sb_type);
+    int bh = 1 << b_height_log2(xd->mode_info_context->mbmi.sb_type);
+    // FIXME(jingning): fix intra4x4 rate-distortion optimization, then
+    // use bw and bh as the increment values.
+#if !CONFIG_AB4X4 || CONFIG_AB4X4
+    bw = 1, bh = 1;
+#endif
+    for (idy = 0; idy < 2; idy += bh) {
+      for (idx = 0; idx < 2; idx += bw) {
+        int m = xd->mode_info_context->bmi[idy * 2 + idx].as_mode.first;
+        ++cpi->sb_ymode_count[m];
+      }
+    }
   }
 }