vp9_rdopt.c 166 KB
Newer Older
John Koleszar's avatar
John Koleszar committed
1
/*
2
 *  Copyright (c) 2010 The WebM project authors. All Rights Reserved.
John Koleszar's avatar
John Koleszar committed
3
 *
4
 *  Use of this source code is governed by a BSD-style license
5
6
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
7
 *  in the file PATENTS.  All contributing project authors may
8
 *  be found in the AUTHORS file in the root of the source tree.
John Koleszar's avatar
John Koleszar committed
9
10
11
12
13
14
 */

#include <stdio.h>
#include <math.h>
#include <limits.h>
#include <assert.h>
15

16
#include "vp9/common/vp9_pragmas.h"
17
18
19
20
21
#include "vp9/encoder/vp9_tokenize.h"
#include "vp9/encoder/vp9_treewriter.h"
#include "vp9/encoder/vp9_onyx_int.h"
#include "vp9/encoder/vp9_modecosts.h"
#include "vp9/encoder/vp9_encodeintra.h"
22
23
24
25
26
#include "vp9/common/vp9_entropymode.h"
#include "vp9/common/vp9_reconinter.h"
#include "vp9/common/vp9_reconintra.h"
#include "vp9/common/vp9_findnearmv.h"
#include "vp9/common/vp9_quant_common.h"
27
28
29
30
31
32
#include "vp9/encoder/vp9_encodemb.h"
#include "vp9/encoder/vp9_quantize.h"
#include "vp9/encoder/vp9_variance.h"
#include "vp9/encoder/vp9_mcomp.h"
#include "vp9/encoder/vp9_rdopt.h"
#include "vp9/encoder/vp9_ratectrl.h"
John Koleszar's avatar
John Koleszar committed
33
#include "vpx_mem/vpx_mem.h"
34
35
36
37
38
#include "vp9/common/vp9_systemdependent.h"
#include "vp9/encoder/vp9_encodemv.h"
#include "vp9/common/vp9_seg_common.h"
#include "vp9/common/vp9_pred_common.h"
#include "vp9/common/vp9_entropy.h"
39
#include "./vp9_rtcd.h"
40
#include "vp9/common/vp9_mvref_common.h"
Ronald S. Bultje's avatar
Ronald S. Bultje committed
41
#include "vp9/common/vp9_common.h"
Paul Wilkins's avatar
Paul Wilkins committed
42

43
44
#define INVALID_MV 0x80008000

45
46
47
/* Factor to weigh the rate for switchable interp filters */
#define SWITCHABLE_INTERP_RATE_FACTOR 1

48
49
50
#define LAST_FRAME_MODE_MASK    0xFFEDCD60
#define GOLDEN_FRAME_MODE_MASK  0xFFDA3BB0
#define ALT_REF_MODE_MASK       0xFFC648D0
51

Paul Wilkins's avatar
Paul Wilkins committed
52
53
#define MIN_EARLY_TERM_INDEX    3

54
const MODE_DEFINITION vp9_mode_order[MAX_MODES] = {
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
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
  {NEARESTMV, LAST_FRAME,   NONE},
  {NEARESTMV, ALTREF_FRAME, NONE},
  {NEARESTMV, GOLDEN_FRAME, NONE},

  {DC_PRED,   INTRA_FRAME,  NONE},

  {NEWMV,     LAST_FRAME,   NONE},
  {NEWMV,     ALTREF_FRAME, NONE},
  {NEWMV,     GOLDEN_FRAME, NONE},

  {NEARMV,    LAST_FRAME,   NONE},
  {NEARMV,    ALTREF_FRAME, NONE},
  {NEARESTMV, LAST_FRAME,   ALTREF_FRAME},
  {NEARESTMV, GOLDEN_FRAME, ALTREF_FRAME},

  {TM_PRED,   INTRA_FRAME,  NONE},

  {NEARMV,    LAST_FRAME,   ALTREF_FRAME},
  {NEWMV,     LAST_FRAME,   ALTREF_FRAME},
  {NEARMV,    GOLDEN_FRAME, NONE},
  {NEARMV,    GOLDEN_FRAME, ALTREF_FRAME},
  {NEWMV,     GOLDEN_FRAME, ALTREF_FRAME},

  {ZEROMV,    LAST_FRAME,   NONE},
  {ZEROMV,    GOLDEN_FRAME, NONE},
  {ZEROMV,    ALTREF_FRAME, NONE},
  {ZEROMV,    LAST_FRAME,   ALTREF_FRAME},
  {ZEROMV,    GOLDEN_FRAME, ALTREF_FRAME},

  {H_PRED,    INTRA_FRAME,  NONE},
  {V_PRED,    INTRA_FRAME,  NONE},
  {D135_PRED, INTRA_FRAME,  NONE},
  {D207_PRED, INTRA_FRAME,  NONE},
  {D153_PRED, INTRA_FRAME,  NONE},
  {D63_PRED,  INTRA_FRAME,  NONE},
  {D117_PRED, INTRA_FRAME,  NONE},
  {D45_PRED,  INTRA_FRAME,  NONE},
};

const REF_DEFINITION vp9_ref_order[MAX_REFS] = {
  {LAST_FRAME,   NONE},
  {GOLDEN_FRAME, NONE},
  {ALTREF_FRAME, NONE},
  {LAST_FRAME,   ALTREF_FRAME},
  {GOLDEN_FRAME, ALTREF_FRAME},
  {INTRA_FRAME,  NONE},
John Koleszar's avatar
John Koleszar committed
101
102
};

103
104
105
106
// The baseline rd thresholds for breaking out of the rd loop for
// certain modes are assumed to be based on 8x8 blocks.
// This table is used to correct for blocks size.
// The factors here are << 2 (2 = x0.5, 32 = x8 etc).
107
static int rd_thresh_block_size_factor[BLOCK_SIZES] =
108
109
  {2, 3, 3, 4, 6, 6, 8, 12, 12, 16, 24, 24, 32};

Deb Mukherjee's avatar
Deb Mukherjee committed
110
111
112
#define RD_THRESH_MAX_FACT 64
#define RD_THRESH_INC      1
#define RD_THRESH_POW      1.25
113
#define RD_MULT_EPB_RATIO  64
Deb Mukherjee's avatar
Deb Mukherjee committed
114

115
116
#define MV_COST_WEIGHT      108
#define MV_COST_WEIGHT_SUB  120
117

118
static void fill_token_costs(vp9_coeff_cost *c,
119
                             vp9_coeff_probs_model (*p)[BLOCK_TYPES]) {
120
  int i, j, k, l;
121
122
123
124
125
126
127
128
  TX_SIZE t;
  for (t = TX_4X4; t <= TX_32X32; t++)
    for (i = 0; i < BLOCK_TYPES; i++)
      for (j = 0; j < REF_TYPES; j++)
        for (k = 0; k < COEF_BANDS; k++)
          for (l = 0; l < PREV_COEF_CONTEXTS; l++) {
            vp9_prob probs[ENTROPY_NODES];
            vp9_model_to_full_probs(p[t][i][j][k][l], probs);
129
            vp9_cost_tokens((int *)c[t][i][j][k][0][l], probs,
130
                            vp9_coef_tree);
131
            vp9_cost_tokens_skip((int *)c[t][i][j][k][1][l], probs,
132
                                 vp9_coef_tree);
133
134
            assert(c[t][i][j][k][0][l][DCT_EOB_TOKEN] ==
                   c[t][i][j][k][1][l][DCT_EOB_TOKEN]);
135
          }
136
137
}

138
139
140
141
142
143
static const int rd_iifactor[32] = {
  4, 4, 3, 2, 1, 0, 0, 0,
  0, 0, 0, 0, 0, 0, 0, 0,
  0, 0, 0, 0, 0, 0, 0, 0,
  0, 0, 0, 0, 0, 0, 0, 0,
};
John Koleszar's avatar
John Koleszar committed
144

145
// 3* dc_qlookup[Q]*dc_qlookup[Q];
146

147
/* values are now correlated to quantizer */
Paul Wilkins's avatar
Paul Wilkins committed
148
149
150
static int sad_per_bit16lut[QINDEX_RANGE];
static int sad_per_bit4lut[QINDEX_RANGE];

151
void vp9_init_me_luts() {
John Koleszar's avatar
John Koleszar committed
152
153
154
155
156
157
158
  int i;

  // Initialize the sad lut tables using a formulaic calculation for now
  // This is to make it easier to resolve the impact of experimental changes
  // to the quantizer tables.
  for (i = 0; i < QINDEX_RANGE; i++) {
    sad_per_bit16lut[i] =
159
      (int)((0.0418 * vp9_convert_qindex_to_q(i)) + 2.4107);
160
    sad_per_bit4lut[i] = (int)(0.063 * vp9_convert_qindex_to_q(i) + 2.742);
John Koleszar's avatar
John Koleszar committed
161
  }
Paul Wilkins's avatar
Paul Wilkins committed
162
}
John Koleszar's avatar
John Koleszar committed
163

164
static int compute_rd_mult(int qindex) {
165
  const int q = vp9_dc_quant(qindex, 0);
166
167
168
169
170
171
172
173
174
175
176
  // TODO(debargha): Adjust the function below
  return (88 * q * q / 25);
}

static int compute_rd_thresh_factor(int qindex) {
  int q;
  // TODO(debargha): Adjust the function below
  q = (int)(pow(vp9_dc_quant(qindex, 0) / 4.0, RD_THRESH_POW) * 5.12);
  if (q < 8)
    q = 8;
  return q;
177
178
}

Dmitry Kovalev's avatar
Dmitry Kovalev committed
179
180
181
void vp9_initialize_me_consts(VP9_COMP *cpi, int qindex) {
  cpi->mb.sadperbit16 = sad_per_bit16lut[qindex];
  cpi->mb.sadperbit4 = sad_per_bit4lut[qindex];
John Koleszar's avatar
John Koleszar committed
182
183
}

Deb Mukherjee's avatar
Deb Mukherjee committed
184
static void set_block_thresholds(VP9_COMP *cpi, int qindex) {
185
  int q, i, bsize;
186
  q = compute_rd_thresh_factor(qindex);
187

188
  for (bsize = 0; bsize < BLOCK_SIZES; ++bsize) {
189
    for (i = 0; i < MAX_MODES; ++i) {
190
191
192
      // Threshold here seem unecessarily harsh but fine given actual
      // range of values used for cpi->sf.thresh_mult[]
      int thresh_max = INT_MAX / (q * rd_thresh_block_size_factor[bsize]);
John Koleszar's avatar
John Koleszar committed
193

194
195
      if (cpi->sf.thresh_mult[i] < thresh_max) {
        cpi->rd_threshes[bsize][i] =
196
197
            cpi->sf.thresh_mult[i] * q *
            rd_thresh_block_size_factor[bsize] / 4;
198
199
      } else {
        cpi->rd_threshes[bsize][i] = INT_MAX;
John Koleszar's avatar
John Koleszar committed
200
      }
John Koleszar's avatar
John Koleszar committed
201
    }
202
203
204
205
206
207
208
209
210
211
212
213

    for (i = 0; i < MAX_REFS; ++i) {
      int thresh_max = INT_MAX / (q * rd_thresh_block_size_factor[bsize]);

      if (cpi->sf.thresh_mult_sub8x8[i] < thresh_max) {
        cpi->rd_thresh_sub8x8[bsize][i] =
            cpi->sf.thresh_mult_sub8x8[i] * q *
            rd_thresh_block_size_factor[bsize] / 4;
      } else {
        cpi->rd_thresh_sub8x8[bsize][i] = INT_MAX;
      }
    }
John Koleszar's avatar
John Koleszar committed
214
  }
Deb Mukherjee's avatar
Deb Mukherjee committed
215
216
217
218
219
220
221
222
223
224
225
226
227
}

void vp9_initialize_rd_consts(VP9_COMP *cpi, int qindex) {
  int i;

  vp9_clear_system_state();  // __asm emms;

  // Further tests required to see if optimum is different
  // for key frames, golden frames and arf frames.
  // if (cpi->common.refresh_golden_frame ||
  //     cpi->common.refresh_alt_ref_frame)
  qindex = clamp(qindex, 0, MAXQ);

228
  cpi->RDDIV = RDDIV_BITS;  // in bits (to multiply D by 128)
Deb Mukherjee's avatar
Deb Mukherjee committed
229
230
231
232
233
234
235
236
  cpi->RDMULT = compute_rd_mult(qindex);
  if (cpi->pass == 2 && (cpi->common.frame_type != KEY_FRAME)) {
    if (cpi->twopass.next_iiratio > 31)
      cpi->RDMULT += (cpi->RDMULT * rd_iifactor[31]) >> 4;
    else
      cpi->RDMULT +=
          (cpi->RDMULT * rd_iifactor[cpi->twopass.next_iiratio]) >> 4;
  }
237
  cpi->mb.errorperbit = cpi->RDMULT / RD_MULT_EPB_RATIO;
Deb Mukherjee's avatar
Deb Mukherjee committed
238
239
240
241
242
  cpi->mb.errorperbit += (cpi->mb.errorperbit == 0);

  vp9_set_speed_features(cpi);

  set_block_thresholds(cpi, qindex);
John Koleszar's avatar
John Koleszar committed
243

244
  fill_token_costs(cpi->mb.token_costs, cpi->common.fc.coef_probs);
245

246
  for (i = 0; i < NUM_PARTITION_CONTEXTS; i++)
247
    vp9_cost_tokens(cpi->mb.partition_cost[i],
248
                    cpi->common.fc.partition_prob[cpi->common.frame_type][i],
249
250
                    vp9_partition_tree);

John Koleszar's avatar
John Koleszar committed
251
  /*rough estimate for costing*/
252
  vp9_init_mode_costs(cpi);
John Koleszar's avatar
John Koleszar committed
253

254
  if (!frame_is_intra_only(&cpi->common)) {
255
    vp9_build_nmv_cost_table(
256
257
258
259
260
        cpi->mb.nmvjointcost,
        cpi->mb.e_mbd.allow_high_precision_mv ?
        cpi->mb.nmvcost_hp : cpi->mb.nmvcost,
        &cpi->common.fc.nmvc,
        cpi->mb.e_mbd.allow_high_precision_mv, 1, 1);
261
262
263
264
265

    for (i = 0; i < INTER_MODE_CONTEXTS; i++) {
      MB_PREDICTION_MODE m;

      for (m = NEARESTMV; m < MB_MODE_COUNT; m++)
266
        cpi->mb.inter_mode_cost[i][inter_mode_offset(m)] =
267
            cost_token(vp9_inter_mode_tree,
268
                       cpi->common.fc.inter_mode_probs[i],
269
                       vp9_inter_mode_encodings + inter_mode_offset(m));
270
    }
271
  }
John Koleszar's avatar
John Koleszar committed
272
273
}

Yaowu Xu's avatar
Yaowu Xu committed
274
static INLINE void linear_interpolate2(double x, int ntab, int inv_step,
275
276
                                       const double *tab1, const double *tab2,
                                       double *v1, double *v2) {
Deb Mukherjee's avatar
Deb Mukherjee committed
277
278
279
  double y = x * inv_step;
  int d = (int) y;
  if (d >= ntab - 1) {
280
281
    *v1 = tab1[ntab - 1];
    *v2 = tab2[ntab - 1];
Deb Mukherjee's avatar
Deb Mukherjee committed
282
283
  } else {
    double a = y - d;
284
285
    *v1 = tab1[d] * (1 - a) + tab1[d + 1] * a;
    *v2 = tab2[d] * (1 - a) + tab2[d + 1] * a;
Deb Mukherjee's avatar
Deb Mukherjee committed
286
287
288
  }
}

289
290
291
292
293
static void model_rd_norm(double x, double *R, double *D) {
  static const int inv_tab_step = 8;
  static const int tab_size = 120;
  // NOTE: The tables below must be of the same size
  //
Deb Mukherjee's avatar
Deb Mukherjee committed
294
  // Normalized rate
295
  // This table models the rate for a Laplacian source
Deb Mukherjee's avatar
Deb Mukherjee committed
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
  // source with given variance when quantized with a uniform quantizer
  // with given stepsize. The closed form expression is:
  // Rn(x) = H(sqrt(r)) + sqrt(r)*[1 + H(r)/(1 - r)],
  // where r = exp(-sqrt(2) * x) and x = qpstep / sqrt(variance),
  // and H(x) is the binary entropy function.
  static const double rate_tab[] = {
    64.00, 4.944, 3.949, 3.372, 2.966, 2.655, 2.403, 2.194,
    2.014, 1.858, 1.720, 1.596, 1.485, 1.384, 1.291, 1.206,
    1.127, 1.054, 0.986, 0.923, 0.863, 0.808, 0.756, 0.708,
    0.662, 0.619, 0.579, 0.541, 0.506, 0.473, 0.442, 0.412,
    0.385, 0.359, 0.335, 0.313, 0.291, 0.272, 0.253, 0.236,
    0.220, 0.204, 0.190, 0.177, 0.165, 0.153, 0.142, 0.132,
    0.123, 0.114, 0.106, 0.099, 0.091, 0.085, 0.079, 0.073,
    0.068, 0.063, 0.058, 0.054, 0.050, 0.047, 0.043, 0.040,
    0.037, 0.034, 0.032, 0.029, 0.027, 0.025, 0.023, 0.022,
    0.020, 0.019, 0.017, 0.016, 0.015, 0.014, 0.013, 0.012,
    0.011, 0.010, 0.009, 0.008, 0.008, 0.007, 0.007, 0.006,
    0.006, 0.005, 0.005, 0.005, 0.004, 0.004, 0.004, 0.003,
    0.003, 0.003, 0.003, 0.002, 0.002, 0.002, 0.002, 0.002,
    0.002, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001,
    0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.000,
  };
  // Normalized distortion
319
  // This table models the normalized distortion for a Laplacian source
Deb Mukherjee's avatar
Deb Mukherjee committed
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
  // source with given variance when quantized with a uniform quantizer
  // with given stepsize. The closed form expression is:
  // Dn(x) = 1 - 1/sqrt(2) * x / sinh(x/sqrt(2))
  // where x = qpstep / sqrt(variance)
  // Note the actual distortion is Dn * variance.
  static const double dist_tab[] = {
    0.000, 0.001, 0.005, 0.012, 0.021, 0.032, 0.045, 0.061,
    0.079, 0.098, 0.119, 0.142, 0.166, 0.190, 0.216, 0.242,
    0.269, 0.296, 0.324, 0.351, 0.378, 0.405, 0.432, 0.458,
    0.484, 0.509, 0.534, 0.557, 0.580, 0.603, 0.624, 0.645,
    0.664, 0.683, 0.702, 0.719, 0.735, 0.751, 0.766, 0.780,
    0.794, 0.807, 0.819, 0.830, 0.841, 0.851, 0.861, 0.870,
    0.878, 0.886, 0.894, 0.901, 0.907, 0.913, 0.919, 0.925,
    0.930, 0.935, 0.939, 0.943, 0.947, 0.951, 0.954, 0.957,
    0.960, 0.963, 0.966, 0.968, 0.971, 0.973, 0.975, 0.976,
    0.978, 0.980, 0.981, 0.982, 0.984, 0.985, 0.986, 0.987,
    0.988, 0.989, 0.990, 0.990, 0.991, 0.992, 0.992, 0.993,
    0.993, 0.994, 0.994, 0.995, 0.995, 0.996, 0.996, 0.996,
    0.996, 0.997, 0.997, 0.997, 0.997, 0.998, 0.998, 0.998,
    0.998, 0.998, 0.998, 0.999, 0.999, 0.999, 0.999, 0.999,
    0.999, 0.999, 0.999, 0.999, 0.999, 0.999, 0.999, 1.000,
  };
342
343
344
345
346
  /*
  assert(sizeof(rate_tab) == tab_size * sizeof(rate_tab[0]);
  assert(sizeof(dist_tab) == tab_size * sizeof(dist_tab[0]);
  assert(sizeof(rate_tab) == sizeof(dist_tab));
  */
Deb Mukherjee's avatar
Deb Mukherjee committed
347
  assert(x >= 0.0);
348
349
  linear_interpolate2(x, tab_size, inv_tab_step,
                      rate_tab, dist_tab, R, D);
Deb Mukherjee's avatar
Deb Mukherjee committed
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
}

static void model_rd_from_var_lapndz(int var, int n, int qstep,
                                     int *rate, int64_t *dist) {
  // This function models the rate and distortion for a Laplacian
  // source with given variance when quantized with a uniform quantizer
  // with given stepsize. The closed form expressions are in:
  // Hang and Chen, "Source Model for transform video coder and its
  // application - Part I: Fundamental Theory", IEEE Trans. Circ.
  // Sys. for Video Tech., April 1997.
  vp9_clear_system_state();
  if (var == 0 || n == 0) {
    *rate = 0;
    *dist = 0;
  } else {
    double D, R;
    double s2 = (double) var / n;
    double x = qstep / sqrt(s2);
368
    model_rd_norm(x, &R, &D);
Yaowu Xu's avatar
Yaowu Xu committed
369
370
    *rate = (int)((n << 8) * R + 0.5);
    *dist = (int)(var * D + 0.5);
Deb Mukherjee's avatar
Deb Mukherjee committed
371
372
373
374
  }
  vp9_clear_system_state();
}

375
static void model_rd_for_sb(VP9_COMP *cpi, BLOCK_SIZE bsize,
Deb Mukherjee's avatar
Deb Mukherjee committed
376
377
378
379
380
381
382
383
384
385
                            MACROBLOCK *x, MACROBLOCKD *xd,
                            int *out_rate_sum, int64_t *out_dist_sum) {
  // Note our transform coeffs are 8 times an orthogonal transform.
  // Hence quantizer step is also 8 times. To get effective quantizer
  // we need to divide by 8 before sending to modeling function.
  int i, rate_sum = 0, dist_sum = 0;

  for (i = 0; i < MAX_MB_PLANE; ++i) {
    struct macroblock_plane *const p = &x->plane[i];
    struct macroblockd_plane *const pd = &xd->plane[i];
386
    const BLOCK_SIZE bs = get_plane_block_size(bsize, pd);
Deb Mukherjee's avatar
Deb Mukherjee committed
387
388
389
390
391
392
    unsigned int sse;
    int rate;
    int64_t dist;
    (void) cpi->fn_ptr[bs].vf(p->src.buf, p->src.stride,
                              pd->dst.buf, pd->dst.stride, &sse);
    // sse works better than var, since there is no dc prediction used
Dmitry Kovalev's avatar
Dmitry Kovalev committed
393
    model_rd_from_var_lapndz(sse, 1 << num_pels_log2_lookup[bs],
394
                             pd->dequant[1] >> 3, &rate, &dist);
Deb Mukherjee's avatar
Deb Mukherjee committed
395
396

    rate_sum += rate;
Yaowu Xu's avatar
Yaowu Xu committed
397
    dist_sum += (int)dist;
Deb Mukherjee's avatar
Deb Mukherjee committed
398
399
400
401
402
403
  }

  *out_rate_sum = rate_sum;
  *out_dist_sum = dist_sum << 4;
}

404
static void model_rd_for_sb_y_tx(VP9_COMP *cpi, BLOCK_SIZE bsize,
Deb Mukherjee's avatar
Deb Mukherjee committed
405
406
407
408
                                 TX_SIZE tx_size,
                                 MACROBLOCK *x, MACROBLOCKD *xd,
                                 int *out_rate_sum, int64_t *out_dist_sum,
                                 int *out_skip) {
409
  int j, k;
410
  BLOCK_SIZE bs;
Deb Mukherjee's avatar
Deb Mukherjee committed
411
412
  struct macroblock_plane *const p = &x->plane[0];
  struct macroblockd_plane *const pd = &xd->plane[0];
413
414
  const int width = 4 << num_4x4_blocks_wide_lookup[bsize];
  const int height = 4 << num_4x4_blocks_high_lookup[bsize];
Deb Mukherjee's avatar
Deb Mukherjee committed
415
416
  int rate_sum = 0;
  int64_t dist_sum = 0;
417
  const int t = 4 << tx_size;
Deb Mukherjee's avatar
Deb Mukherjee committed
418
419
420
421
422
423
424
425
426
427
428
429

  if (tx_size == TX_4X4) {
    bs = BLOCK_4X4;
  } else if (tx_size == TX_8X8) {
    bs = BLOCK_8X8;
  } else if (tx_size == TX_16X16) {
    bs = BLOCK_16X16;
  } else if (tx_size == TX_32X32) {
    bs = BLOCK_32X32;
  } else {
    assert(0);
  }
430

Deb Mukherjee's avatar
Deb Mukherjee committed
431
  *out_skip = 1;
Jim Bankoski's avatar
Jim Bankoski committed
432
433
  for (j = 0; j < height; j += t) {
    for (k = 0; k < width; k += t) {
Deb Mukherjee's avatar
Deb Mukherjee committed
434
435
436
      int rate;
      int64_t dist;
      unsigned int sse;
437
438
439
      cpi->fn_ptr[bs].vf(&p->src.buf[j * p->src.stride + k], p->src.stride,
                         &pd->dst.buf[j * pd->dst.stride + k], pd->dst.stride,
                         &sse);
Deb Mukherjee's avatar
Deb Mukherjee committed
440
      // sse works better than var, since there is no dc prediction used
441
      model_rd_from_var_lapndz(sse, t * t, pd->dequant[1] >> 3, &rate, &dist);
Deb Mukherjee's avatar
Deb Mukherjee committed
442
443
444
445
446
      rate_sum += rate;
      dist_sum += dist;
      *out_skip &= (rate < 1024);
    }
  }
447

Deb Mukherjee's avatar
Deb Mukherjee committed
448
  *out_rate_sum = rate_sum;
449
  *out_dist_sum = dist_sum << 4;
Deb Mukherjee's avatar
Deb Mukherjee committed
450
451
}

Ronald S. Bultje's avatar
Ronald S. Bultje committed
452
int64_t vp9_block_error_c(int16_t *coeff, int16_t *dqcoeff,
453
                          intptr_t block_size, int64_t *ssz) {
Ronald S. Bultje's avatar
Ronald S. Bultje committed
454
  int i;
455
  int64_t error = 0, sqcoeff = 0;
John Koleszar's avatar
John Koleszar committed
456

457
  for (i = 0; i < block_size; i++) {
John Koleszar's avatar
John Koleszar committed
458
    int this_diff = coeff[i] - dqcoeff[i];
Ronald S. Bultje's avatar
Ronald S. Bultje committed
459
    error += (unsigned)this_diff * this_diff;
460
    sqcoeff += (unsigned) coeff[i] * coeff[i];
John Koleszar's avatar
John Koleszar committed
461
  }
John Koleszar's avatar
John Koleszar committed
462

463
  *ssz = sqcoeff;
John Koleszar's avatar
John Koleszar committed
464
  return error;
John Koleszar's avatar
John Koleszar committed
465
466
}

467
468
469
470
471
/* The trailing '0' is a terminator which is used inside cost_coeffs() to
 * decide whether to include cost of a trailing EOB node or not (i.e. we
 * can skip this if the last coefficient in this transform block, e.g. the
 * 16th coefficient in a 4x4 block or the 64th coefficient in a 8x8 block,
 * were non-zero). */
472
static const int16_t band_counts[TX_SIZES][8] = {
473
474
475
476
  { 1, 2, 3, 4,  3,   16 - 13, 0 },
  { 1, 2, 3, 4, 11,   64 - 21, 0 },
  { 1, 2, 3, 4, 11,  256 - 21, 0 },
  { 1, 2, 3, 4, 11, 1024 - 21, 0 },
477
478
};

479
static INLINE int cost_coeffs(MACROBLOCK *x,
480
                              int plane, int block,
481
                              ENTROPY_CONTEXT *A, ENTROPY_CONTEXT *L,
John Koleszar's avatar
John Koleszar committed
482
                              TX_SIZE tx_size,
483
                              const int16_t *scan, const int16_t *nb) {
484
  MACROBLOCKD *const xd = &x->e_mbd;
485
  MB_MODE_INFO *mbmi = &xd->this_mi->mbmi;
486
487
  struct macroblockd_plane *pd = &xd->plane[plane];
  const PLANE_TYPE type = pd->plane_type;
488
  const int16_t *band_count = &band_counts[tx_size][1];
489
490
  const int eob = pd->eobs[block];
  const int16_t *const qcoeff_ptr = BLOCK_OFFSET(pd->qcoeff, block);
Ronald S. Bultje's avatar
Ronald S. Bultje committed
491
  const int ref = mbmi->ref_frame[0] != INTRA_FRAME;
492
  unsigned int (*token_costs)[2][PREV_COEF_CONTEXTS][MAX_ENTROPY_TOKENS] =
493
                   x->token_costs[tx_size][type][ref];
494
  const ENTROPY_CONTEXT above_ec = !!*A, left_ec = !!*L;
495
  uint8_t *p_tok = x->token_cache;
496
497
  int pt = combine_entropy_contexts(above_ec, left_ec);
  int c, cost;
498
499

  // Check for consistency of tx_size with mode info
500
  assert(type == PLANE_TYPE_Y_WITH_DC ? mbmi->tx_size == tx_size
501
                                      : get_uv_tx_size(mbmi) == tx_size);
502

503
504
  if (eob == 0) {
    // single eob token
505
506
    cost = token_costs[0][0][pt][DCT_EOB_TOKEN];
    c = 0;
507
  } else {
508
    int band_left = *band_count++;
509
510

    // dc token
511
512
    int v = qcoeff_ptr[0];
    int prev_t = vp9_dct_value_tokens_ptr[v].token;
513
    cost = (*token_costs)[0][pt][prev_t] + vp9_dct_value_cost_ptr[v];
514
    p_tok[0] = vp9_pt_energy_class[prev_t];
515
    ++token_costs;
516
517
518
519

    // ac tokens
    for (c = 1; c < eob; c++) {
      const int rc = scan[c];
520
      int t;
521
522
523

      v = qcoeff_ptr[rc];
      t = vp9_dct_value_tokens_ptr[v].token;
524
      pt = get_coef_context(nb, p_tok, c);
525
      cost += (*token_costs)[!prev_t][pt][t] + vp9_dct_value_cost_ptr[v];
526
      p_tok[rc] = vp9_pt_energy_class[t];
527
      prev_t = t;
528
      if (!--band_left) {
529
530
        band_left = *band_count++;
        ++token_costs;
531
      }
532
    }
533
534

    // eob token
535
    if (band_left) {
536
      pt = get_coef_context(nb, p_tok, c);
537
      cost += (*token_costs)[0][pt][DCT_EOB_TOKEN];
538
    }
539
540
  }

541
  // is eob first coefficient;
542
  *A = *L = (c > 0);
543

544
545
546
  return cost;
}

547
static void dist_block(int plane, int block, TX_SIZE tx_size, void *arg) {
548
  const int ss_txfrm_size = tx_size << 1;
Deb Mukherjee's avatar
Deb Mukherjee committed
549
550
551
  struct rdcost_block_args* args = arg;
  MACROBLOCK* const x = args->x;
  MACROBLOCKD* const xd = &x->e_mbd;
552
553
  struct macroblock_plane *const p = &x->plane[plane];
  struct macroblockd_plane *const pd = &xd->plane[plane];
Deb Mukherjee's avatar
Deb Mukherjee committed
554
555
  int64_t this_sse;
  int shift = args->tx_size == TX_32X32 ? 0 : 2;
556
557
  int16_t *const coeff = BLOCK_OFFSET(p->coeff, block);
  int16_t *const dqcoeff = BLOCK_OFFSET(pd->dqcoeff, block);
558
559
560
  args->dist = vp9_block_error(coeff, dqcoeff, 16 << ss_txfrm_size,
                               &this_sse) >> shift;
  args->sse  = this_sse >> shift;
561
562

  if (x->skip_encode &&
563
      xd->this_mi->mbmi.ref_frame[0] == INTRA_FRAME) {
564
565
    // TODO(jingning): tune the model to better capture the distortion.
    int64_t p = (pd->dequant[1] * pd->dequant[1] *
566
                    (1 << ss_txfrm_size)) >> (shift + 2);
567
568
    args->dist += (p >> 4);
    args->sse  += p;
569
  }
Deb Mukherjee's avatar
Deb Mukherjee committed
570
571
}

572
static void rate_block(int plane, int block, BLOCK_SIZE plane_bsize,
573
                       TX_SIZE tx_size, void *arg) {
Deb Mukherjee's avatar
Deb Mukherjee committed
574
  struct rdcost_block_args* args = arg;
575

Deb Mukherjee's avatar
Deb Mukherjee committed
576
  int x_idx, y_idx;
577
  txfrm_block_to_raster_xy(plane_bsize, args->tx_size, block, &x_idx, &y_idx);
Deb Mukherjee's avatar
Deb Mukherjee committed
578

579
580
581
  args->rate = cost_coeffs(args->x, plane, block, args->t_above + x_idx,
                           args->t_left + y_idx, args->tx_size,
                           args->scan, args->nb);
Deb Mukherjee's avatar
Deb Mukherjee committed
582
583
}

584
static void block_yrd_txfm(int plane, int block, BLOCK_SIZE plane_bsize,
585
                           TX_SIZE tx_size, void *arg) {
Deb Mukherjee's avatar
Deb Mukherjee committed
586
587
588
  struct rdcost_block_args *args = arg;
  MACROBLOCK *const x = args->x;
  MACROBLOCKD *const xd = &x->e_mbd;
589
  struct encode_b_args encode_args = {x, NULL};
590
  int64_t rd1, rd2, rd;
Deb Mukherjee's avatar
Deb Mukherjee committed
591

592
593
594
  if (args->skip)
    return;

595
  if (!is_inter_block(&xd->this_mi->mbmi))
596
    vp9_encode_block_intra(plane, block, plane_bsize, tx_size, &encode_args);
Deb Mukherjee's avatar
Deb Mukherjee committed
597
  else
598
    vp9_xform_quant(plane, block, plane_bsize, tx_size, &encode_args);
Deb Mukherjee's avatar
Deb Mukherjee committed
599

600
601
  dist_block(plane, block, tx_size, args);
  rate_block(plane, block, plane_bsize, tx_size, args);
602
603
  rd1 = RDCOST(x->rdmult, x->rddiv, args->rate, args->dist);
  rd2 = RDCOST(x->rdmult, x->rddiv, 0, args->sse);
604
605

  // TODO(jingning): temporarily enabled only for luma component
606
  rd = MIN(rd1, rd2);
607
608
609
  if (plane == 0)
    x->zcoeff_blk[tx_size][block] = rd1 > rd2;

610
611
612
  args->this_rate += args->rate;
  args->this_dist += args->dist;
  args->this_sse  += args->sse;
613
614
615
616
617
618
  args->this_rd += rd;

  if (args->this_rd > args->best_rd) {
    args->skip = 1;
    return;
  }
Deb Mukherjee's avatar
Deb Mukherjee committed
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
644
645
646
647
648
649
650
651
652
653
void vp9_get_entropy_contexts(TX_SIZE tx_size,
    ENTROPY_CONTEXT t_above[16], ENTROPY_CONTEXT t_left[16],
    const ENTROPY_CONTEXT *above, const ENTROPY_CONTEXT *left,
    int num_4x4_w, int num_4x4_h) {
  int i;
  switch (tx_size) {
    case TX_4X4:
      vpx_memcpy(t_above, above, sizeof(ENTROPY_CONTEXT) * num_4x4_w);
      vpx_memcpy(t_left, left, sizeof(ENTROPY_CONTEXT) * num_4x4_h);
      break;
    case TX_8X8:
      for (i = 0; i < num_4x4_w; i += 2)
        t_above[i] = !!*(const uint16_t *)&above[i];
      for (i = 0; i < num_4x4_h; i += 2)
        t_left[i] = !!*(const uint16_t *)&left[i];
      break;
    case TX_16X16:
      for (i = 0; i < num_4x4_w; i += 4)
        t_above[i] = !!*(const uint32_t *)&above[i];
      for (i = 0; i < num_4x4_h; i += 4)
        t_left[i] = !!*(const uint32_t *)&left[i];
      break;
    case TX_32X32:
      for (i = 0; i < num_4x4_w; i += 8)
        t_above[i] = !!*(const uint64_t *)&above[i];
      for (i = 0; i < num_4x4_h; i += 8)
        t_left[i] = !!*(const uint64_t *)&left[i];
      break;
    default:
      assert(!"Invalid transform size.");
  }
}

654
655
656
657
658
659
660
661
662
663
664
665
static void init_rdcost_stack(MACROBLOCK *x, TX_SIZE tx_size,
                              const int num_4x4_w, const int num_4x4_h,
                              const int64_t ref_rdcost,
                              struct rdcost_block_args *arg) {
  vpx_memset(arg, 0, sizeof(struct rdcost_block_args));
  arg->x = x;
  arg->tx_size = tx_size;
  arg->bw = num_4x4_w;
  arg->bh = num_4x4_h;
  arg->best_rd = ref_rdcost;
}

666
static void txfm_rd_in_plane(MACROBLOCK *x,
667
                             struct rdcost_block_args *rd_stack,
668
669
670
                             int *rate, int64_t *distortion,
                             int *skippable, int64_t *sse,
                             int64_t ref_best_rd, int plane,
671
                             BLOCK_SIZE bsize, TX_SIZE tx_size) {
Deb Mukherjee's avatar
Deb Mukherjee committed
672
  MACROBLOCKD *const xd = &x->e_mbd;
673
  struct macroblockd_plane *const pd = &xd->plane[plane];
674
  const BLOCK_SIZE bs = get_plane_block_size(bsize, pd);
675
676
  const int num_4x4_w = num_4x4_blocks_wide_lookup[bs];
  const int num_4x4_h = num_4x4_blocks_high_lookup[bs];
677
  const uint8_t *band_translate;  // just for the get_scan_and_band call
678

679
680
  init_rdcost_stack(x, tx_size, num_4x4_w, num_4x4_h,
                    ref_best_rd, rd_stack);
681
  if (plane == 0)
682
    xd->this_mi->mbmi.tx_size = tx_size;
683

684
  vp9_get_entropy_contexts(tx_size, rd_stack->t_above, rd_stack->t_left,
685
686
                           pd->above_context, pd->left_context,
                           num_4x4_w, num_4x4_h);
687

688
689
  get_scan_and_band(xd, tx_size, pd->plane_type, 0, &rd_stack->scan,
                    &rd_stack->nb, &band_translate);
Deb Mukherjee's avatar
Deb Mukherjee committed
690

691
692
693
  foreach_transformed_block_in_plane(xd, bsize, plane,
                                     block_yrd_txfm, rd_stack);
  if (rd_stack->skip) {
694
695
696
697
698
    *rate       = INT_MAX;
    *distortion = INT64_MAX;
    *sse        = INT64_MAX;
    *skippable  = 0;
  } else {
699
700
701
    *distortion = rd_stack->this_dist;
    *rate       = rd_stack->this_rate;
    *sse        = rd_stack->this_sse;
702
703
    *skippable  = vp9_is_skippable_in_plane(xd, bsize, plane);
  }
Deb Mukherjee's avatar
Deb Mukherjee committed
704
705
706
707
708
}

static void choose_largest_txfm_size(VP9_COMP *cpi, MACROBLOCK *x,
                                     int *rate, int64_t *distortion,
                                     int *skip, int64_t *sse,
709
                                     int64_t ref_best_rd,
710
                                     BLOCK_SIZE bs) {
Dmitry Kovalev's avatar
Dmitry Kovalev committed
711
  const TX_SIZE max_tx_size = max_txsize_lookup[bs];
Deb Mukherjee's avatar
Deb Mukherjee committed
712
713
  VP9_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;
714
  MB_MODE_INFO *const mbmi = &xd->this_mi->mbmi;
Dmitry Kovalev's avatar
Dmitry Kovalev committed
715
  if (max_tx_size == TX_32X32 &&
716
717
      (cm->tx_mode == ALLOW_32X32 ||
       cm->tx_mode == TX_MODE_SELECT)) {
718
    mbmi->tx_size = TX_32X32;
Dmitry Kovalev's avatar
Dmitry Kovalev committed
719
  } else if (max_tx_size >= TX_16X16 &&
720
721
722
             (cm->tx_mode == ALLOW_16X16 ||
              cm->tx_mode == ALLOW_32X32 ||
              cm->tx_mode == TX_MODE_SELECT)) {
723
    mbmi->tx_size = TX_16X16;
724
  } else if (cm->tx_mode != ONLY_4X4) {
725
    mbmi->tx_size = TX_8X8;
Deb Mukherjee's avatar
Deb Mukherjee committed
726
  } else {
727
    mbmi->tx_size = TX_4X4;
Deb Mukherjee's avatar
Deb Mukherjee committed
728
  }
729
  txfm_rd_in_plane(x, &cpi->rdcost_stack, rate, distortion, skip,
730
731
                   &sse[mbmi->tx_size], ref_best_rd, 0, bs,
                   mbmi->tx_size);
Dmitry Kovalev's avatar
Dmitry Kovalev committed
732
  cpi->tx_stepdown_count[0]++;
Deb Mukherjee's avatar
Deb Mukherjee committed
733
734
}

735
static void choose_txfm_size_from_rd(VP9_COMP *cpi, MACROBLOCK *x,
736
                                     int (*r)[2], int *rate,
Ronald S. Bultje's avatar
Ronald S. Bultje committed
737
                                     int64_t *d, int64_t *distortion,
738
                                     int *s, int *skip,
739
                                     int64_t tx_cache[TX_MODES],
740
                                     BLOCK_SIZE bs) {
741
  const TX_SIZE max_tx_size = max_txsize_lookup[bs];
742
743
  VP9_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;
744
  MB_MODE_INFO *const mbmi = &xd->this_mi->mbmi;
745
  vp9_prob skip_prob = vp9_get_pred_prob_mbskip(cm, xd);
746
  int64_t rd[TX_SIZES][2];
747
  int n, m;
748
  int s0, s1;
749

750
  const vp9_prob *tx_probs = get_tx_probs2(xd, &cm->fc.tx_probs, xd->this_mi);
751

752
  for (n = TX_4X4; n <= max_tx_size; n++) {
753
    r[n][1] = r[n][0];
754
755
    if (r[n][0] == INT_MAX)
      continue;
756
    for (m = 0; m <= n - (n == max_tx_size); m++) {
757
      if (m == n)
758
        r[n][1] += vp9_cost_zero(tx_probs[m]);
759
      else
760
        r[n][1] += vp9_cost_one(tx_probs[m]);
761
762
    }
  }
763

764
765
766
  assert(skip_prob > 0);
  s0 = vp9_cost_bit(skip_prob, 0);
  s1 = vp9_cost_bit(skip_prob, 1);
767

768
  for (n = TX_4X4; n <= max_tx_size; n++) {
769
770
771
772
    if (d[n] == INT64_MAX) {
      rd[n][0] = rd[n][1] = INT64_MAX;
      continue;
    }
773
774
775
776
777
    if (s[n]) {
      rd[n][0] = rd[n][1] = RDCOST(x->rdmult, x->rddiv, s1, d[n]);
    } else {
      rd[n][0] = RDCOST(x->rdmult, x->rddiv, r[n][0] + s0, d[n]);
      rd[n][1] = RDCOST(x->rdmult, x->rddiv, r[n][1] + s0, d[n]);
778
779
780
    }
  }

781
  if (max_tx_size == TX_32X32 &&
782
783
      (cm->tx_mode == ALLOW_32X32 ||
       (cm->tx_mode == TX_MODE_SELECT &&
784
785
        rd[TX_32X32][1] < rd[TX_16X16][1] && rd[TX_32X32][1] < rd[TX_8X8][1] &&
        rd[TX_32X32][1] < rd[TX_4X4][1]))) {
786
    mbmi->tx_size = TX_32X32;
787
  } else if (max_tx_size >= TX_16X16 &&
788
789
790
             (cm->tx_mode == ALLOW_16X16 ||
              cm->tx_mode == ALLOW_32X32 ||
              (cm->tx_mode == TX_MODE_SELECT &&
791
792
               rd[TX_16X16][1] < rd[TX_8X8][1] &&
               rd[TX_16X16][1] < rd[TX_4X4][1]))) {
793
    mbmi->tx_size = TX_16X16;
794
795
796
797
  } else if (cm->tx_mode == ALLOW_8X8 ||
             cm->tx_mode == ALLOW_16X16 ||
             cm->tx_mode == ALLOW_32X32 ||
           (cm->tx_mode == TX_MODE_SELECT && rd[TX_8X8][1] < rd[TX_4X4][1])) {
798
    mbmi->tx_size = TX_8X8;
799
  } else {
800
    mbmi->tx_size = TX_4X4;
801
802
  }

803
804
805
  *distortion = d[mbmi->tx_size];
  *rate       = r[mbmi->tx_size][cm->tx_mode == TX_MODE_SELECT];
  *skip       = s[mbmi->tx_size];
806

807
808
809
810
811
  tx_cache[ONLY_4X4] = rd[TX_4X4][0];
  tx_cache[ALLOW_8X8] = rd[TX_8X8][0];
  tx_cache[ALLOW_16X16] = rd[MIN(max_tx_size, TX_16X16)][0];
  tx_cache[ALLOW_32X32] = rd[MIN(max_tx_size, TX_32X32)][0];
  if (max_tx_size == TX_32X32 &&
812
813
      rd[TX_32X32][1] < rd[TX_16X16][1] && rd[TX_32X32][1] < rd[TX_8X8][1] &&
      rd[TX_32X32][1] < rd[TX_4X4][1])
814
815
    tx_cache[TX_MODE_SELECT] = rd[TX_32X32][1];
  else if (max_tx_size >= TX_16X16 &&
816
           rd[TX_16X16][1] < rd[TX_8X8][1] && rd[TX_16X16][1] < rd[TX_4X4][1])
817
    tx_cache[TX_MODE_SELECT] = rd[TX_16X16][1];
818
  else
819
    tx_cache[TX_MODE_SELECT] = rd[TX_4X4][1] < rd[TX_8X8][1] ?
820
                                 rd[TX_4X4][1] : rd[TX_8X8][1];
821

822
  if (max_tx_size == TX_32X32 &&
Deb Mukherjee's avatar
Deb Mukherjee committed
823
824
825
      rd[TX_32X32][1] < rd[TX_16X16][1] &&
      rd[TX_32X32][1] < rd[TX_8X8][1] &&
      rd[TX_32X32][1] < rd[TX_4X4][1]) {
Dmitry Kovalev's avatar
Dmitry Kovalev committed
826
    cpi->tx_stepdown_count[0]++;
827
  } else if (max_tx_size >= TX_16X16 &&
Deb Mukherjee's avatar
Deb Mukherjee committed
828
829
             rd[TX_16X16][1] < rd[TX_8X8][1] &&
             rd[TX_16X16][1] < rd[TX_4X4][1]) {
Dmitry Kovalev's avatar
Dmitry Kovalev committed
830
    cpi->tx_stepdown_count[max_tx_size - TX_16X16]++;
Deb Mukherjee's avatar
Deb Mukherjee committed
831
  } else if (rd[TX_8X8][1] < rd[TX_4X4][1]) {
Dmitry Kovalev's avatar
Dmitry Kovalev committed
832
    cpi->tx_stepdown_count[max_tx_size - TX_8X8]++;
Deb Mukherjee's avatar
Deb Mukherjee committed
833
  } else {
Dmitry Kovalev's avatar
Dmitry Kovalev committed
834
    cpi->tx_stepdown_count[max_tx_size - TX_4X4]++;
835
  }
Deb Mukherjee's avatar
Deb Mukherjee committed
836
}
837

Deb Mukherjee's avatar
Deb Mukherjee committed
838
839
840
841
static void choose_txfm_size_from_modelrd(VP9_COMP *cpi, MACROBLOCK *x,
                                          int (*r)[2], int *rate,
                                          int64_t *d, int64_t *distortion,
                                          int *s, int *skip, int64_t *sse,
842
                                          int64_t ref_best_rd,
843
                                          BLOCK_SIZE bs) {
Dmitry Kovalev's avatar
Dmitry Kovalev committed
844
  const TX_SIZE max_tx_size = max_txsize_lookup[bs];
Deb Mukherjee's avatar
Deb Mukherjee committed
845
846
  VP9_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;
847
  MB_MODE_INFO *const mbmi = &xd->this_mi->mbmi;
848
  vp9_prob skip_prob = vp9_get_pred_prob_mbskip(cm, xd);
849
  int64_t rd[TX_SIZES][2];
Deb Mukherjee's avatar
Deb Mukherjee committed
850
851
  int n, m;
  int s0, s1;
852
853
  double scale_rd[TX_SIZES] = {1.73, 1.44, 1.20, 1.00};
  // double scale_r[TX_SIZES] = {2.82, 2.00, 1.41, 1.00};
854

855
  const vp9_prob *tx_probs = get_tx_probs2(xd, &cm->fc.tx_probs,  xd->this_mi);
856

Deb Mukherjee's avatar
Deb Mukherjee committed
857
858
  // for (n = TX_4X4; n <= max_txfm_size; n++)
  //   r[n][0] = (r[n][0] * scale_r[n]);
859

Dmitry Kovalev's avatar
Dmitry Kovalev committed
860
  for (n = TX_4X4; n <= max_tx_size; n++) {
Deb Mukherjee's avatar
Deb Mukherjee committed
861
    r[n][1] = r[n][0];
Dmitry Kovalev's avatar
Dmitry Kovalev committed
862
    for (m = 0; m <= n - (n == max_tx_size); m++) {
Deb Mukherjee's avatar
Deb Mukherjee committed
863
864
865
866
867
      if (m == n)
        r[n][1] += vp9_cost_zero(tx_probs[m]);
      else
        r[n][1] += vp9_cost_one(tx_probs[m]);
    }
868
  }
869

Deb Mukherjee's avatar
Deb Mukherjee committed
870
871
872
  assert(skip_prob > 0);
  s0 = vp9_cost_bit(skip_prob, 0);
  s1 = vp9_cost_bit(skip_prob, 1);
873

Dmitry Kovalev's avatar
Dmitry Kovalev committed
874
  for (n = TX_4X4; n <= max_tx_size; n++) {
Deb Mukherjee's avatar
Deb Mukherjee committed
875
876
877
878
879
880
881
    if (s[n]) {
      rd[n][0] = rd[n][1] = RDCOST(x->rdmult, x->rddiv, s1, d[n]);
    } else {
      rd[n][0] = RDCOST(x->rdmult, x->rddiv, r[n][0] + s0, d[n]);
      rd[n][1] = RDCOST(x->rdmult, x->rddiv, r[n][1] + s0, d[n]);
    }
  }
Dmitry Kovalev's avatar
Dmitry Kovalev committed
882
  for (n = TX_4X4; n <= max_tx_size; n++) {
Yaowu Xu's avatar
Yaowu Xu committed
883
884
    rd[n][0] = (int64_t)(scale_rd[n] * rd[n][0]);
    rd[n][1] = (int64_t)(scale_rd[n] * rd[n][1]);
Deb Mukherjee's avatar
Deb Mukherjee committed
885
  }
886

Dmitry Kovalev's avatar
Dmitry Kovalev committed
887
  if (max_tx_size == TX_32X32 &&
888
889
      (cm->tx_mode == ALLOW_32X32 ||
       (cm->tx_mode == TX_MODE_SELECT &&
Deb Mukherjee's avatar
Deb Mukherjee committed
890
891
892
        rd[TX_32X32][1] <= rd[TX_16X16][1] &&
        rd[TX_32X32][1] <= rd[TX_8X8][1] &&
        rd[TX_32X32][1] <= rd[TX_4X4][1]))) {
893
    mbmi->tx_size = TX_32X32;
Dmitry Kovalev's avatar
Dmitry Kovalev committed
894
  } else if (max_tx_size >= TX_16X16 &&
895
896
897
             (cm->tx_mode == ALLOW_16X16 ||
              cm->tx_mode == ALLOW_32X32 ||
              (cm->tx_mode == TX_MODE_SELECT &&
Deb Mukherjee's avatar
Deb Mukherjee committed
898
899
               rd[TX_16X16][1] <= rd[TX_8X8][1] &&
               rd[TX_16X16][1] <= rd[TX_4X4][1]))) {
900
    mbmi->tx_size = TX_16X16;
901
902
903
904
  } else if (cm->tx_mode == ALLOW_8X8 ||
             cm->tx_mode == ALLOW_16X16 ||
             cm->tx_mode == ALLOW_32X32 ||
           (cm->tx_mode == TX_MODE_SELECT &&
Deb Mukherjee's avatar
Deb Mukherjee committed
905
            rd[TX_8X8][1] <= rd[TX_4X4][1])) {
906
    mbmi->tx_size = TX_8X8;
Deb Mukherjee's avatar
Deb Mukherjee committed
907
  } else {
908
    mbmi->tx_size = TX_4X4;
Deb Mukherjee's avatar
Deb Mukherjee committed
909
  }
910

911
912
  // Actually encode using the chosen mode if a model was used, but do not
  // update the r, d costs