Coverage Report

Created: 2022-08-24 06:15

/src/aom/av1/encoder/tx_search.c
Line
Count
Source (jump to first uncovered line)
1
/*
2
 * Copyright (c) 2020, Alliance for Open Media. All rights reserved
3
 *
4
 * This source code is subject to the terms of the BSD 2 Clause License and
5
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6
 * was not distributed with this source code in the LICENSE file, you can
7
 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
8
 * Media Patent License 1.0 was not distributed with this source code in the
9
 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10
 */
11
12
#include "av1/common/cfl.h"
13
#include "av1/common/reconintra.h"
14
#include "av1/encoder/block.h"
15
#include "av1/encoder/hybrid_fwd_txfm.h"
16
#include "av1/common/idct.h"
17
#include "av1/encoder/model_rd.h"
18
#include "av1/encoder/random.h"
19
#include "av1/encoder/rdopt_utils.h"
20
#include "av1/encoder/sorting_network.h"
21
#include "av1/encoder/tx_prune_model_weights.h"
22
#include "av1/encoder/tx_search.h"
23
#include "av1/encoder/txb_rdopt.h"
24
25
0
#define PROB_THRESH_OFFSET_TX_TYPE 100
26
27
struct rdcost_block_args {
28
  const AV1_COMP *cpi;
29
  MACROBLOCK *x;
30
  ENTROPY_CONTEXT t_above[MAX_MIB_SIZE];
31
  ENTROPY_CONTEXT t_left[MAX_MIB_SIZE];
32
  RD_STATS rd_stats;
33
  int64_t current_rd;
34
  int64_t best_rd;
35
  int exit_early;
36
  int incomplete_exit;
37
  FAST_TX_SEARCH_MODE ftxs_mode;
38
  int skip_trellis;
39
};
40
41
typedef struct {
42
  int64_t rd;
43
  int txb_entropy_ctx;
44
  TX_TYPE tx_type;
45
} TxCandidateInfo;
46
47
// origin_threshold * 128 / 100
48
static const uint32_t skip_pred_threshold[3][BLOCK_SIZES_ALL] = {
49
  {
50
      64, 64, 64, 70, 60, 60, 68, 68, 68, 68, 68,
51
      68, 68, 68, 68, 68, 64, 64, 70, 70, 68, 68,
52
  },
53
  {
54
      88, 88, 88, 86, 87, 87, 68, 68, 68, 68, 68,
55
      68, 68, 68, 68, 68, 88, 88, 86, 86, 68, 68,
56
  },
57
  {
58
      90, 93, 93, 90, 93, 93, 74, 74, 74, 74, 74,
59
      74, 74, 74, 74, 74, 90, 90, 90, 90, 74, 74,
60
  },
61
};
62
63
// lookup table for predict_skip_txfm
64
// int max_tx_size = max_txsize_rect_lookup[bsize];
65
// if (tx_size_high[max_tx_size] > 16 || tx_size_wide[max_tx_size] > 16)
66
//   max_tx_size = AOMMIN(max_txsize_lookup[bsize], TX_16X16);
67
static const TX_SIZE max_predict_sf_tx_size[BLOCK_SIZES_ALL] = {
68
  TX_4X4,   TX_4X8,   TX_8X4,   TX_8X8,   TX_8X16,  TX_16X8,
69
  TX_16X16, TX_16X16, TX_16X16, TX_16X16, TX_16X16, TX_16X16,
70
  TX_16X16, TX_16X16, TX_16X16, TX_16X16, TX_4X16,  TX_16X4,
71
  TX_8X8,   TX_8X8,   TX_16X16, TX_16X16,
72
};
73
74
// look-up table for sqrt of number of pixels in a transform block
75
// rounded up to the nearest integer.
76
static const int sqrt_tx_pixels_2d[TX_SIZES_ALL] = { 4,  8,  16, 32, 32, 6,  6,
77
                                                     12, 12, 23, 23, 32, 32, 8,
78
                                                     8,  16, 16, 23, 23 };
79
80
0
static INLINE uint32_t get_block_residue_hash(MACROBLOCK *x, BLOCK_SIZE bsize) {
81
0
  const int rows = block_size_high[bsize];
82
0
  const int cols = block_size_wide[bsize];
83
0
  const int16_t *diff = x->plane[0].src_diff;
84
0
  const uint32_t hash =
85
0
      av1_get_crc32c_value(&x->txfm_search_info.mb_rd_record->crc_calculator,
86
0
                           (uint8_t *)diff, 2 * rows * cols);
87
0
  return (hash << 5) + bsize;
88
0
}
89
90
static INLINE int32_t find_mb_rd_info(const MB_RD_RECORD *const mb_rd_record,
91
                                      const int64_t ref_best_rd,
92
0
                                      const uint32_t hash) {
93
0
  int32_t match_index = -1;
94
0
  if (ref_best_rd != INT64_MAX) {
95
0
    for (int i = 0; i < mb_rd_record->num; ++i) {
96
0
      const int index = (mb_rd_record->index_start + i) % RD_RECORD_BUFFER_LEN;
97
      // If there is a match in the mb_rd_record, fetch the RD decision and
98
      // terminate early.
99
0
      if (mb_rd_record->mb_rd_info[index].hash_value == hash) {
100
0
        match_index = index;
101
0
        break;
102
0
      }
103
0
    }
104
0
  }
105
0
  return match_index;
106
0
}
107
108
static AOM_INLINE void fetch_mb_rd_info(int n4,
109
                                        const MB_RD_INFO *const mb_rd_info,
110
                                        RD_STATS *const rd_stats,
111
0
                                        MACROBLOCK *const x) {
112
0
  MACROBLOCKD *const xd = &x->e_mbd;
113
0
  MB_MODE_INFO *const mbmi = xd->mi[0];
114
0
  mbmi->tx_size = mb_rd_info->tx_size;
115
0
  memcpy(x->txfm_search_info.blk_skip, mb_rd_info->blk_skip,
116
0
         sizeof(mb_rd_info->blk_skip[0]) * n4);
117
0
  av1_copy(mbmi->inter_tx_size, mb_rd_info->inter_tx_size);
118
0
  av1_copy_array(xd->tx_type_map, mb_rd_info->tx_type_map, n4);
119
0
  *rd_stats = mb_rd_info->rd_stats;
120
0
}
121
122
// Compute the pixel domain distortion from diff on all visible 4x4s in the
123
// transform block.
124
static INLINE int64_t pixel_diff_dist(const MACROBLOCK *x, int plane,
125
                                      int blk_row, int blk_col,
126
                                      const BLOCK_SIZE plane_bsize,
127
                                      const BLOCK_SIZE tx_bsize,
128
44.3M
                                      unsigned int *block_mse_q8) {
129
44.3M
  int visible_rows, visible_cols;
130
44.3M
  const MACROBLOCKD *xd = &x->e_mbd;
131
44.3M
  get_txb_dimensions(xd, plane, plane_bsize, blk_row, blk_col, tx_bsize, NULL,
132
44.3M
                     NULL, &visible_cols, &visible_rows);
133
44.3M
  const int diff_stride = block_size_wide[plane_bsize];
134
44.3M
  const int16_t *diff = x->plane[plane].src_diff;
135
136
44.3M
  diff += ((blk_row * diff_stride + blk_col) << MI_SIZE_LOG2);
137
44.3M
  uint64_t sse =
138
44.3M
      aom_sum_squares_2d_i16(diff, diff_stride, visible_cols, visible_rows);
139
44.3M
  if (block_mse_q8 != NULL) {
140
44.3M
    if (visible_cols > 0 && visible_rows > 0)
141
44.3M
      *block_mse_q8 =
142
44.3M
          (unsigned int)((256 * sse) / (visible_cols * visible_rows));
143
18.4E
    else
144
18.4E
      *block_mse_q8 = UINT_MAX;
145
44.3M
  }
146
44.3M
  return sse;
147
44.3M
}
148
149
// Computes the residual block's SSE and mean on all visible 4x4s in the
150
// transform block
151
static INLINE int64_t pixel_diff_stats(
152
    MACROBLOCK *x, int plane, int blk_row, int blk_col,
153
    const BLOCK_SIZE plane_bsize, const BLOCK_SIZE tx_bsize,
154
0
    unsigned int *block_mse_q8, int64_t *per_px_mean, uint64_t *block_var) {
155
0
  int visible_rows, visible_cols;
156
0
  const MACROBLOCKD *xd = &x->e_mbd;
157
0
  get_txb_dimensions(xd, plane, plane_bsize, blk_row, blk_col, tx_bsize, NULL,
158
0
                     NULL, &visible_cols, &visible_rows);
159
0
  const int diff_stride = block_size_wide[plane_bsize];
160
0
  const int16_t *diff = x->plane[plane].src_diff;
161
162
0
  diff += ((blk_row * diff_stride + blk_col) << MI_SIZE_LOG2);
163
0
  uint64_t sse = 0;
164
0
  int sum = 0;
165
0
  sse = aom_sum_sse_2d_i16(diff, diff_stride, visible_cols, visible_rows, &sum);
166
0
  if (visible_cols > 0 && visible_rows > 0) {
167
0
    double norm_factor = 1.0 / (visible_cols * visible_rows);
168
0
    int sign_sum = sum > 0 ? 1 : -1;
169
    // Conversion to transform domain
170
0
    *per_px_mean = (int64_t)(norm_factor * abs(sum)) << 7;
171
0
    *per_px_mean = sign_sum * (*per_px_mean);
172
0
    *block_mse_q8 = (unsigned int)(norm_factor * (256 * sse));
173
0
    *block_var = (uint64_t)(sse - (uint64_t)(norm_factor * sum * sum));
174
0
  } else {
175
0
    *block_mse_q8 = UINT_MAX;
176
0
  }
177
0
  return sse;
178
0
}
179
180
// Uses simple features on top of DCT coefficients to quickly predict
181
// whether optimal RD decision is to skip encoding the residual.
182
// The sse value is stored in dist.
183
static int predict_skip_txfm(MACROBLOCK *x, BLOCK_SIZE bsize, int64_t *dist,
184
0
                             int reduced_tx_set) {
185
0
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
186
0
  const int bw = block_size_wide[bsize];
187
0
  const int bh = block_size_high[bsize];
188
0
  const MACROBLOCKD *xd = &x->e_mbd;
189
0
  const int16_t dc_q = av1_dc_quant_QTX(x->qindex, 0, xd->bd);
190
191
0
  *dist = pixel_diff_dist(x, 0, 0, 0, bsize, bsize, NULL);
192
193
0
  const int64_t mse = *dist / bw / bh;
194
  // Normalized quantizer takes the transform upscaling factor (8 for tx size
195
  // smaller than 32) into account.
196
0
  const int16_t normalized_dc_q = dc_q >> 3;
197
0
  const int64_t mse_thresh = (int64_t)normalized_dc_q * normalized_dc_q / 8;
198
  // For faster early skip decision, use dist to compare against threshold so
199
  // that quality risk is less for the skip=1 decision. Otherwise, use mse
200
  // since the fwd_txfm coeff checks will take care of quality
201
  // TODO(any): Use dist to return 0 when skip_txfm_level is 1
202
0
  int64_t pred_err = (txfm_params->skip_txfm_level >= 2) ? *dist : mse;
203
  // Predict not to skip when error is larger than threshold.
204
0
  if (pred_err > mse_thresh) return 0;
205
  // Return as skip otherwise for aggressive early skip
206
0
  else if (txfm_params->skip_txfm_level >= 2)
207
0
    return 1;
208
209
0
  const int max_tx_size = max_predict_sf_tx_size[bsize];
210
0
  const int tx_h = tx_size_high[max_tx_size];
211
0
  const int tx_w = tx_size_wide[max_tx_size];
212
0
  DECLARE_ALIGNED(32, tran_low_t, coefs[32 * 32]);
213
0
  TxfmParam param;
214
0
  param.tx_type = DCT_DCT;
215
0
  param.tx_size = max_tx_size;
216
0
  param.bd = xd->bd;
217
0
  param.is_hbd = is_cur_buf_hbd(xd);
218
0
  param.lossless = 0;
219
0
  param.tx_set_type = av1_get_ext_tx_set_type(
220
0
      param.tx_size, is_inter_block(xd->mi[0]), reduced_tx_set);
221
0
  const int bd_idx = (xd->bd == 8) ? 0 : ((xd->bd == 10) ? 1 : 2);
222
0
  const uint32_t max_qcoef_thresh = skip_pred_threshold[bd_idx][bsize];
223
0
  const int16_t *src_diff = x->plane[0].src_diff;
224
0
  const int n_coeff = tx_w * tx_h;
225
0
  const int16_t ac_q = av1_ac_quant_QTX(x->qindex, 0, xd->bd);
226
0
  const uint32_t dc_thresh = max_qcoef_thresh * dc_q;
227
0
  const uint32_t ac_thresh = max_qcoef_thresh * ac_q;
228
0
  for (int row = 0; row < bh; row += tx_h) {
229
0
    for (int col = 0; col < bw; col += tx_w) {
230
0
      av1_fwd_txfm(src_diff + col, coefs, bw, &param);
231
      // Operating on TX domain, not pixels; we want the QTX quantizers
232
0
      const uint32_t dc_coef = (((uint32_t)abs(coefs[0])) << 7);
233
0
      if (dc_coef >= dc_thresh) return 0;
234
0
      for (int i = 1; i < n_coeff; ++i) {
235
0
        const uint32_t ac_coef = (((uint32_t)abs(coefs[i])) << 7);
236
0
        if (ac_coef >= ac_thresh) return 0;
237
0
      }
238
0
    }
239
0
    src_diff += tx_h * bw;
240
0
  }
241
0
  return 1;
242
0
}
243
244
// Used to set proper context for early termination with skip = 1.
245
static AOM_INLINE void set_skip_txfm(MACROBLOCK *x, RD_STATS *rd_stats,
246
0
                                     int bsize, int64_t dist) {
247
0
  MACROBLOCKD *const xd = &x->e_mbd;
248
0
  MB_MODE_INFO *const mbmi = xd->mi[0];
249
0
  const int n4 = bsize_to_num_blk(bsize);
250
0
  const TX_SIZE tx_size = max_txsize_rect_lookup[bsize];
251
0
  memset(xd->tx_type_map, DCT_DCT, sizeof(xd->tx_type_map[0]) * n4);
252
0
  memset(mbmi->inter_tx_size, tx_size, sizeof(mbmi->inter_tx_size));
253
0
  mbmi->tx_size = tx_size;
254
0
  for (int i = 0; i < n4; ++i)
255
0
    set_blk_skip(x->txfm_search_info.blk_skip, 0, i, 1);
256
0
  rd_stats->skip_txfm = 1;
257
0
  if (is_cur_buf_hbd(xd)) dist = ROUND_POWER_OF_TWO(dist, (xd->bd - 8) * 2);
258
0
  rd_stats->dist = rd_stats->sse = (dist << 4);
259
  // Though decision is to make the block as skip based on luma stats,
260
  // it is possible that block becomes non skip after chroma rd. In addition
261
  // intermediate non skip costs calculated by caller function will be
262
  // incorrect, if rate is set as  zero (i.e., if zero_blk_rate is not
263
  // accounted). Hence intermediate rate is populated to code the luma tx blks
264
  // as skip, the caller function based on final rd decision (i.e., skip vs
265
  // non-skip) sets the final rate accordingly. Here the rate populated
266
  // corresponds to coding all the tx blocks with zero_blk_rate (based on max tx
267
  // size possible) in the current block. Eg: For 128*128 block, rate would be
268
  // 4 * zero_blk_rate where zero_blk_rate corresponds to coding of one 64x64 tx
269
  // block as 'all zeros'
270
0
  ENTROPY_CONTEXT ctxa[MAX_MIB_SIZE];
271
0
  ENTROPY_CONTEXT ctxl[MAX_MIB_SIZE];
272
0
  av1_get_entropy_contexts(bsize, &xd->plane[0], ctxa, ctxl);
273
0
  ENTROPY_CONTEXT *ta = ctxa;
274
0
  ENTROPY_CONTEXT *tl = ctxl;
275
0
  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
276
0
  TXB_CTX txb_ctx;
277
0
  get_txb_ctx(bsize, tx_size, 0, ta, tl, &txb_ctx);
278
0
  const int zero_blk_rate = x->coeff_costs.coeff_costs[txs_ctx][PLANE_TYPE_Y]
279
0
                                .txb_skip_cost[txb_ctx.txb_skip_ctx][1];
280
0
  rd_stats->rate = zero_blk_rate *
281
0
                   (block_size_wide[bsize] >> tx_size_wide_log2[tx_size]) *
282
0
                   (block_size_high[bsize] >> tx_size_high_log2[tx_size]);
283
0
}
284
285
static AOM_INLINE void save_mb_rd_info(int n4, uint32_t hash,
286
                                       const MACROBLOCK *const x,
287
                                       const RD_STATS *const rd_stats,
288
0
                                       MB_RD_RECORD *mb_rd_record) {
289
0
  int index;
290
0
  if (mb_rd_record->num < RD_RECORD_BUFFER_LEN) {
291
0
    index =
292
0
        (mb_rd_record->index_start + mb_rd_record->num) % RD_RECORD_BUFFER_LEN;
293
0
    ++mb_rd_record->num;
294
0
  } else {
295
0
    index = mb_rd_record->index_start;
296
0
    mb_rd_record->index_start =
297
0
        (mb_rd_record->index_start + 1) % RD_RECORD_BUFFER_LEN;
298
0
  }
299
0
  MB_RD_INFO *const mb_rd_info = &mb_rd_record->mb_rd_info[index];
300
0
  const MACROBLOCKD *const xd = &x->e_mbd;
301
0
  const MB_MODE_INFO *const mbmi = xd->mi[0];
302
0
  mb_rd_info->hash_value = hash;
303
0
  mb_rd_info->tx_size = mbmi->tx_size;
304
0
  memcpy(mb_rd_info->blk_skip, x->txfm_search_info.blk_skip,
305
0
         sizeof(mb_rd_info->blk_skip[0]) * n4);
306
0
  av1_copy(mb_rd_info->inter_tx_size, mbmi->inter_tx_size);
307
0
  av1_copy_array(mb_rd_info->tx_type_map, xd->tx_type_map, n4);
308
0
  mb_rd_info->rd_stats = *rd_stats;
309
0
}
310
311
static int get_search_init_depth(int mi_width, int mi_height, int is_inter,
312
                                 const SPEED_FEATURES *sf,
313
94.9k
                                 int tx_size_search_method) {
314
94.9k
  if (tx_size_search_method == USE_LARGESTALL) return MAX_VARTX_DEPTH;
315
316
94.9k
  if (sf->tx_sf.tx_size_search_lgr_block) {
317
0
    if (mi_width > mi_size_wide[BLOCK_64X64] ||
318
0
        mi_height > mi_size_high[BLOCK_64X64])
319
0
      return MAX_VARTX_DEPTH;
320
0
  }
321
322
94.9k
  if (is_inter) {
323
0
    return (mi_height != mi_width)
324
0
               ? sf->tx_sf.inter_tx_size_search_init_depth_rect
325
0
               : sf->tx_sf.inter_tx_size_search_init_depth_sqr;
326
94.9k
  } else {
327
94.9k
    return (mi_height != mi_width)
328
94.9k
               ? sf->tx_sf.intra_tx_size_search_init_depth_rect
329
94.9k
               : sf->tx_sf.intra_tx_size_search_init_depth_sqr;
330
94.9k
  }
331
94.9k
}
332
333
static AOM_INLINE void select_tx_block(
334
    const AV1_COMP *cpi, MACROBLOCK *x, int blk_row, int blk_col, int block,
335
    TX_SIZE tx_size, int depth, BLOCK_SIZE plane_bsize, ENTROPY_CONTEXT *ta,
336
    ENTROPY_CONTEXT *tl, TXFM_CONTEXT *tx_above, TXFM_CONTEXT *tx_left,
337
    RD_STATS *rd_stats, int64_t prev_level_rd, int64_t ref_best_rd,
338
    int *is_cost_valid, FAST_TX_SEARCH_MODE ftxs_mode);
339
340
// NOTE: CONFIG_COLLECT_RD_STATS has 3 possible values
341
// 0: Do not collect any RD stats
342
// 1: Collect RD stats for transform units
343
// 2: Collect RD stats for partition units
344
#if CONFIG_COLLECT_RD_STATS
345
346
static AOM_INLINE void get_energy_distribution_fine(
347
    const AV1_COMP *cpi, BLOCK_SIZE bsize, const uint8_t *src, int src_stride,
348
    const uint8_t *dst, int dst_stride, int need_4th, double *hordist,
349
    double *verdist) {
350
  const int bw = block_size_wide[bsize];
351
  const int bh = block_size_high[bsize];
352
  unsigned int esq[16] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
353
354
  if (bsize < BLOCK_16X16 || (bsize >= BLOCK_4X16 && bsize <= BLOCK_32X8)) {
355
    // Special cases: calculate 'esq' values manually, as we don't have 'vf'
356
    // functions for the 16 (very small) sub-blocks of this block.
357
    const int w_shift = (bw == 4) ? 0 : (bw == 8) ? 1 : (bw == 16) ? 2 : 3;
358
    const int h_shift = (bh == 4) ? 0 : (bh == 8) ? 1 : (bh == 16) ? 2 : 3;
359
    assert(bw <= 32);
360
    assert(bh <= 32);
361
    assert(((bw - 1) >> w_shift) + (((bh - 1) >> h_shift) << 2) == 15);
362
    if (cpi->common.seq_params->use_highbitdepth) {
363
      const uint16_t *src16 = CONVERT_TO_SHORTPTR(src);
364
      const uint16_t *dst16 = CONVERT_TO_SHORTPTR(dst);
365
      for (int i = 0; i < bh; ++i)
366
        for (int j = 0; j < bw; ++j) {
367
          const int index = (j >> w_shift) + ((i >> h_shift) << 2);
368
          esq[index] +=
369
              (src16[j + i * src_stride] - dst16[j + i * dst_stride]) *
370
              (src16[j + i * src_stride] - dst16[j + i * dst_stride]);
371
        }
372
    } else {
373
      for (int i = 0; i < bh; ++i)
374
        for (int j = 0; j < bw; ++j) {
375
          const int index = (j >> w_shift) + ((i >> h_shift) << 2);
376
          esq[index] += (src[j + i * src_stride] - dst[j + i * dst_stride]) *
377
                        (src[j + i * src_stride] - dst[j + i * dst_stride]);
378
        }
379
    }
380
  } else {  // Calculate 'esq' values using 'vf' functions on the 16 sub-blocks.
381
    const int f_index =
382
        (bsize < BLOCK_SIZES) ? bsize - BLOCK_16X16 : bsize - BLOCK_8X16;
383
    assert(f_index >= 0 && f_index < BLOCK_SIZES_ALL);
384
    const BLOCK_SIZE subsize = (BLOCK_SIZE)f_index;
385
    assert(block_size_wide[bsize] == 4 * block_size_wide[subsize]);
386
    assert(block_size_high[bsize] == 4 * block_size_high[subsize]);
387
    cpi->ppi->fn_ptr[subsize].vf(src, src_stride, dst, dst_stride, &esq[0]);
388
    cpi->ppi->fn_ptr[subsize].vf(src + bw / 4, src_stride, dst + bw / 4,
389
                                 dst_stride, &esq[1]);
390
    cpi->ppi->fn_ptr[subsize].vf(src + bw / 2, src_stride, dst + bw / 2,
391
                                 dst_stride, &esq[2]);
392
    cpi->ppi->fn_ptr[subsize].vf(src + 3 * bw / 4, src_stride, dst + 3 * bw / 4,
393
                                 dst_stride, &esq[3]);
394
    src += bh / 4 * src_stride;
395
    dst += bh / 4 * dst_stride;
396
397
    cpi->ppi->fn_ptr[subsize].vf(src, src_stride, dst, dst_stride, &esq[4]);
398
    cpi->ppi->fn_ptr[subsize].vf(src + bw / 4, src_stride, dst + bw / 4,
399
                                 dst_stride, &esq[5]);
400
    cpi->ppi->fn_ptr[subsize].vf(src + bw / 2, src_stride, dst + bw / 2,
401
                                 dst_stride, &esq[6]);
402
    cpi->ppi->fn_ptr[subsize].vf(src + 3 * bw / 4, src_stride, dst + 3 * bw / 4,
403
                                 dst_stride, &esq[7]);
404
    src += bh / 4 * src_stride;
405
    dst += bh / 4 * dst_stride;
406
407
    cpi->ppi->fn_ptr[subsize].vf(src, src_stride, dst, dst_stride, &esq[8]);
408
    cpi->ppi->fn_ptr[subsize].vf(src + bw / 4, src_stride, dst + bw / 4,
409
                                 dst_stride, &esq[9]);
410
    cpi->ppi->fn_ptr[subsize].vf(src + bw / 2, src_stride, dst + bw / 2,
411
                                 dst_stride, &esq[10]);
412
    cpi->ppi->fn_ptr[subsize].vf(src + 3 * bw / 4, src_stride, dst + 3 * bw / 4,
413
                                 dst_stride, &esq[11]);
414
    src += bh / 4 * src_stride;
415
    dst += bh / 4 * dst_stride;
416
417
    cpi->ppi->fn_ptr[subsize].vf(src, src_stride, dst, dst_stride, &esq[12]);
418
    cpi->ppi->fn_ptr[subsize].vf(src + bw / 4, src_stride, dst + bw / 4,
419
                                 dst_stride, &esq[13]);
420
    cpi->ppi->fn_ptr[subsize].vf(src + bw / 2, src_stride, dst + bw / 2,
421
                                 dst_stride, &esq[14]);
422
    cpi->ppi->fn_ptr[subsize].vf(src + 3 * bw / 4, src_stride, dst + 3 * bw / 4,
423
                                 dst_stride, &esq[15]);
424
  }
425
426
  double total = (double)esq[0] + esq[1] + esq[2] + esq[3] + esq[4] + esq[5] +
427
                 esq[6] + esq[7] + esq[8] + esq[9] + esq[10] + esq[11] +
428
                 esq[12] + esq[13] + esq[14] + esq[15];
429
  if (total > 0) {
430
    const double e_recip = 1.0 / total;
431
    hordist[0] = ((double)esq[0] + esq[4] + esq[8] + esq[12]) * e_recip;
432
    hordist[1] = ((double)esq[1] + esq[5] + esq[9] + esq[13]) * e_recip;
433
    hordist[2] = ((double)esq[2] + esq[6] + esq[10] + esq[14]) * e_recip;
434
    if (need_4th) {
435
      hordist[3] = ((double)esq[3] + esq[7] + esq[11] + esq[15]) * e_recip;
436
    }
437
    verdist[0] = ((double)esq[0] + esq[1] + esq[2] + esq[3]) * e_recip;
438
    verdist[1] = ((double)esq[4] + esq[5] + esq[6] + esq[7]) * e_recip;
439
    verdist[2] = ((double)esq[8] + esq[9] + esq[10] + esq[11]) * e_recip;
440
    if (need_4th) {
441
      verdist[3] = ((double)esq[12] + esq[13] + esq[14] + esq[15]) * e_recip;
442
    }
443
  } else {
444
    hordist[0] = verdist[0] = 0.25;
445
    hordist[1] = verdist[1] = 0.25;
446
    hordist[2] = verdist[2] = 0.25;
447
    if (need_4th) {
448
      hordist[3] = verdist[3] = 0.25;
449
    }
450
  }
451
}
452
453
static double get_sse_norm(const int16_t *diff, int stride, int w, int h) {
454
  double sum = 0.0;
455
  for (int j = 0; j < h; ++j) {
456
    for (int i = 0; i < w; ++i) {
457
      const int err = diff[j * stride + i];
458
      sum += err * err;
459
    }
460
  }
461
  assert(w > 0 && h > 0);
462
  return sum / (w * h);
463
}
464
465
static double get_sad_norm(const int16_t *diff, int stride, int w, int h) {
466
  double sum = 0.0;
467
  for (int j = 0; j < h; ++j) {
468
    for (int i = 0; i < w; ++i) {
469
      sum += abs(diff[j * stride + i]);
470
    }
471
  }
472
  assert(w > 0 && h > 0);
473
  return sum / (w * h);
474
}
475
476
static AOM_INLINE void get_2x2_normalized_sses_and_sads(
477
    const AV1_COMP *const cpi, BLOCK_SIZE tx_bsize, const uint8_t *const src,
478
    int src_stride, const uint8_t *const dst, int dst_stride,
479
    const int16_t *const src_diff, int diff_stride, double *const sse_norm_arr,
480
    double *const sad_norm_arr) {
481
  const BLOCK_SIZE tx_bsize_half =
482
      get_partition_subsize(tx_bsize, PARTITION_SPLIT);
483
  if (tx_bsize_half == BLOCK_INVALID) {  // manually calculate stats
484
    const int half_width = block_size_wide[tx_bsize] / 2;
485
    const int half_height = block_size_high[tx_bsize] / 2;
486
    for (int row = 0; row < 2; ++row) {
487
      for (int col = 0; col < 2; ++col) {
488
        const int16_t *const this_src_diff =
489
            src_diff + row * half_height * diff_stride + col * half_width;
490
        if (sse_norm_arr) {
491
          sse_norm_arr[row * 2 + col] =
492
              get_sse_norm(this_src_diff, diff_stride, half_width, half_height);
493
        }
494
        if (sad_norm_arr) {
495
          sad_norm_arr[row * 2 + col] =
496
              get_sad_norm(this_src_diff, diff_stride, half_width, half_height);
497
        }
498
      }
499
    }
500
  } else {  // use function pointers to calculate stats
501
    const int half_width = block_size_wide[tx_bsize_half];
502
    const int half_height = block_size_high[tx_bsize_half];
503
    const int num_samples_half = half_width * half_height;
504
    for (int row = 0; row < 2; ++row) {
505
      for (int col = 0; col < 2; ++col) {
506
        const uint8_t *const this_src =
507
            src + row * half_height * src_stride + col * half_width;
508
        const uint8_t *const this_dst =
509
            dst + row * half_height * dst_stride + col * half_width;
510
511
        if (sse_norm_arr) {
512
          unsigned int this_sse;
513
          cpi->ppi->fn_ptr[tx_bsize_half].vf(this_src, src_stride, this_dst,
514
                                             dst_stride, &this_sse);
515
          sse_norm_arr[row * 2 + col] = (double)this_sse / num_samples_half;
516
        }
517
518
        if (sad_norm_arr) {
519
          const unsigned int this_sad = cpi->ppi->fn_ptr[tx_bsize_half].sdf(
520
              this_src, src_stride, this_dst, dst_stride);
521
          sad_norm_arr[row * 2 + col] = (double)this_sad / num_samples_half;
522
        }
523
      }
524
    }
525
  }
526
}
527
528
#if CONFIG_COLLECT_RD_STATS == 1
529
static double get_mean(const int16_t *diff, int stride, int w, int h) {
530
  double sum = 0.0;
531
  for (int j = 0; j < h; ++j) {
532
    for (int i = 0; i < w; ++i) {
533
      sum += diff[j * stride + i];
534
    }
535
  }
536
  assert(w > 0 && h > 0);
537
  return sum / (w * h);
538
}
539
static AOM_INLINE void PrintTransformUnitStats(
540
    const AV1_COMP *const cpi, MACROBLOCK *x, const RD_STATS *const rd_stats,
541
    int blk_row, int blk_col, BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
542
    TX_TYPE tx_type, int64_t rd) {
543
  if (rd_stats->rate == INT_MAX || rd_stats->dist == INT64_MAX) return;
544
545
  // Generate small sample to restrict output size.
546
  static unsigned int seed = 21743;
547
  if (lcg_rand16(&seed) % 256 > 0) return;
548
549
  const char output_file[] = "tu_stats.txt";
550
  FILE *fout = fopen(output_file, "a");
551
  if (!fout) return;
552
553
  const BLOCK_SIZE tx_bsize = txsize_to_bsize[tx_size];
554
  const MACROBLOCKD *const xd = &x->e_mbd;
555
  const int plane = 0;
556
  struct macroblock_plane *const p = &x->plane[plane];
557
  const struct macroblockd_plane *const pd = &xd->plane[plane];
558
  const int txw = tx_size_wide[tx_size];
559
  const int txh = tx_size_high[tx_size];
560
  const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3;
561
  const int q_step = p->dequant_QTX[1] >> dequant_shift;
562
  const int num_samples = txw * txh;
563
564
  const double rate_norm = (double)rd_stats->rate / num_samples;
565
  const double dist_norm = (double)rd_stats->dist / num_samples;
566
567
  fprintf(fout, "%g %g", rate_norm, dist_norm);
568
569
  const int src_stride = p->src.stride;
570
  const uint8_t *const src =
571
      &p->src.buf[(blk_row * src_stride + blk_col) << MI_SIZE_LOG2];
572
  const int dst_stride = pd->dst.stride;
573
  const uint8_t *const dst =
574
      &pd->dst.buf[(blk_row * dst_stride + blk_col) << MI_SIZE_LOG2];
575
  unsigned int sse;
576
  cpi->ppi->fn_ptr[tx_bsize].vf(src, src_stride, dst, dst_stride, &sse);
577
  const double sse_norm = (double)sse / num_samples;
578
579
  const unsigned int sad =
580
      cpi->ppi->fn_ptr[tx_bsize].sdf(src, src_stride, dst, dst_stride);
581
  const double sad_norm = (double)sad / num_samples;
582
583
  fprintf(fout, " %g %g", sse_norm, sad_norm);
584
585
  const int diff_stride = block_size_wide[plane_bsize];
586
  const int16_t *const src_diff =
587
      &p->src_diff[(blk_row * diff_stride + blk_col) << MI_SIZE_LOG2];
588
589
  double sse_norm_arr[4], sad_norm_arr[4];
590
  get_2x2_normalized_sses_and_sads(cpi, tx_bsize, src, src_stride, dst,
591
                                   dst_stride, src_diff, diff_stride,
592
                                   sse_norm_arr, sad_norm_arr);
593
  for (int i = 0; i < 4; ++i) {
594
    fprintf(fout, " %g", sse_norm_arr[i]);
595
  }
596
  for (int i = 0; i < 4; ++i) {
597
    fprintf(fout, " %g", sad_norm_arr[i]);
598
  }
599
600
  const TX_TYPE_1D tx_type_1d_row = htx_tab[tx_type];
601
  const TX_TYPE_1D tx_type_1d_col = vtx_tab[tx_type];
602
603
  fprintf(fout, " %d %d %d %d %d", q_step, tx_size_wide[tx_size],
604
          tx_size_high[tx_size], tx_type_1d_row, tx_type_1d_col);
605
606
  int model_rate;
607
  int64_t model_dist;
608
  model_rd_sse_fn[MODELRD_CURVFIT](cpi, x, tx_bsize, plane, sse, num_samples,
609
                                   &model_rate, &model_dist);
610
  const double model_rate_norm = (double)model_rate / num_samples;
611
  const double model_dist_norm = (double)model_dist / num_samples;
612
  fprintf(fout, " %g %g", model_rate_norm, model_dist_norm);
613
614
  const double mean = get_mean(src_diff, diff_stride, txw, txh);
615
  float hor_corr, vert_corr;
616
  av1_get_horver_correlation_full(src_diff, diff_stride, txw, txh, &hor_corr,
617
                                  &vert_corr);
618
  fprintf(fout, " %g %g %g", mean, hor_corr, vert_corr);
619
620
  double hdist[4] = { 0 }, vdist[4] = { 0 };
621
  get_energy_distribution_fine(cpi, tx_bsize, src, src_stride, dst, dst_stride,
622
                               1, hdist, vdist);
623
  fprintf(fout, " %g %g %g %g %g %g %g %g", hdist[0], hdist[1], hdist[2],
624
          hdist[3], vdist[0], vdist[1], vdist[2], vdist[3]);
625
626
  fprintf(fout, " %d %" PRId64, x->rdmult, rd);
627
628
  fprintf(fout, "\n");
629
  fclose(fout);
630
}
631
#endif  // CONFIG_COLLECT_RD_STATS == 1
632
633
#if CONFIG_COLLECT_RD_STATS >= 2
634
static int64_t get_sse(const AV1_COMP *cpi, const MACROBLOCK *x) {
635
  const AV1_COMMON *cm = &cpi->common;
636
  const int num_planes = av1_num_planes(cm);
637
  const MACROBLOCKD *xd = &x->e_mbd;
638
  const MB_MODE_INFO *mbmi = xd->mi[0];
639
  int64_t total_sse = 0;
640
  for (int plane = 0; plane < num_planes; ++plane) {
641
    const struct macroblock_plane *const p = &x->plane[plane];
642
    const struct macroblockd_plane *const pd = &xd->plane[plane];
643
    const BLOCK_SIZE bs =
644
        get_plane_block_size(mbmi->bsize, pd->subsampling_x, pd->subsampling_y);
645
    unsigned int sse;
646
647
    if (x->skip_chroma_rd && plane) continue;
648
649
    cpi->ppi->fn_ptr[bs].vf(p->src.buf, p->src.stride, pd->dst.buf,
650
                            pd->dst.stride, &sse);
651
    total_sse += sse;
652
  }
653
  total_sse <<= 4;
654
  return total_sse;
655
}
656
657
static int get_est_rate_dist(const TileDataEnc *tile_data, BLOCK_SIZE bsize,
658
                             int64_t sse, int *est_residue_cost,
659
                             int64_t *est_dist) {
660
  const InterModeRdModel *md = &tile_data->inter_mode_rd_models[bsize];
661
  if (md->ready) {
662
    if (sse < md->dist_mean) {
663
      *est_residue_cost = 0;
664
      *est_dist = sse;
665
    } else {
666
      *est_dist = (int64_t)round(md->dist_mean);
667
      const double est_ld = md->a * sse + md->b;
668
      // Clamp estimated rate cost by INT_MAX / 2.
669
      // TODO(angiebird@google.com): find better solution than clamping.
670
      if (fabs(est_ld) < 1e-2) {
671
        *est_residue_cost = INT_MAX / 2;
672
      } else {
673
        double est_residue_cost_dbl = ((sse - md->dist_mean) / est_ld);
674
        if (est_residue_cost_dbl < 0) {
675
          *est_residue_cost = 0;
676
        } else {
677
          *est_residue_cost =
678
              (int)AOMMIN((int64_t)round(est_residue_cost_dbl), INT_MAX / 2);
679
        }
680
      }
681
      if (*est_residue_cost <= 0) {
682
        *est_residue_cost = 0;
683
        *est_dist = sse;
684
      }
685
    }
686
    return 1;
687
  }
688
  return 0;
689
}
690
691
static double get_highbd_diff_mean(const uint8_t *src8, int src_stride,
692
                                   const uint8_t *dst8, int dst_stride, int w,
693
                                   int h) {
694
  const uint16_t *src = CONVERT_TO_SHORTPTR(src8);
695
  const uint16_t *dst = CONVERT_TO_SHORTPTR(dst8);
696
  double sum = 0.0;
697
  for (int j = 0; j < h; ++j) {
698
    for (int i = 0; i < w; ++i) {
699
      const int diff = src[j * src_stride + i] - dst[j * dst_stride + i];
700
      sum += diff;
701
    }
702
  }
703
  assert(w > 0 && h > 0);
704
  return sum / (w * h);
705
}
706
707
static double get_diff_mean(const uint8_t *src, int src_stride,
708
                            const uint8_t *dst, int dst_stride, int w, int h) {
709
  double sum = 0.0;
710
  for (int j = 0; j < h; ++j) {
711
    for (int i = 0; i < w; ++i) {
712
      const int diff = src[j * src_stride + i] - dst[j * dst_stride + i];
713
      sum += diff;
714
    }
715
  }
716
  assert(w > 0 && h > 0);
717
  return sum / (w * h);
718
}
719
720
static AOM_INLINE void PrintPredictionUnitStats(const AV1_COMP *const cpi,
721
                                                const TileDataEnc *tile_data,
722
                                                MACROBLOCK *x,
723
                                                const RD_STATS *const rd_stats,
724
                                                BLOCK_SIZE plane_bsize) {
725
  if (rd_stats->rate == INT_MAX || rd_stats->dist == INT64_MAX) return;
726
727
  if (cpi->sf.inter_sf.inter_mode_rd_model_estimation == 1 &&
728
      (tile_data == NULL ||
729
       !tile_data->inter_mode_rd_models[plane_bsize].ready))
730
    return;
731
  (void)tile_data;
732
  // Generate small sample to restrict output size.
733
  static unsigned int seed = 95014;
734
735
  if ((lcg_rand16(&seed) % (1 << (14 - num_pels_log2_lookup[plane_bsize]))) !=
736
      1)
737
    return;
738
739
  const char output_file[] = "pu_stats.txt";
740
  FILE *fout = fopen(output_file, "a");
741
  if (!fout) return;
742
743
  MACROBLOCKD *const xd = &x->e_mbd;
744
  const int plane = 0;
745
  struct macroblock_plane *const p = &x->plane[plane];
746
  struct macroblockd_plane *pd = &xd->plane[plane];
747
  const int diff_stride = block_size_wide[plane_bsize];
748
  int bw, bh;
749
  get_txb_dimensions(xd, plane, plane_bsize, 0, 0, plane_bsize, NULL, NULL, &bw,
750
                     &bh);
751
  const int num_samples = bw * bh;
752
  const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3;
753
  const int q_step = p->dequant_QTX[1] >> dequant_shift;
754
  const int shift = (xd->bd - 8);
755
756
  const double rate_norm = (double)rd_stats->rate / num_samples;
757
  const double dist_norm = (double)rd_stats->dist / num_samples;
758
  const double rdcost_norm =
759
      (double)RDCOST(x->rdmult, rd_stats->rate, rd_stats->dist) / num_samples;
760
761
  fprintf(fout, "%g %g %g", rate_norm, dist_norm, rdcost_norm);
762
763
  const int src_stride = p->src.stride;
764
  const uint8_t *const src = p->src.buf;
765
  const int dst_stride = pd->dst.stride;
766
  const uint8_t *const dst = pd->dst.buf;
767
  const int16_t *const src_diff = p->src_diff;
768
769
  int64_t sse = calculate_sse(xd, p, pd, bw, bh);
770
  const double sse_norm = (double)sse / num_samples;
771
772
  const unsigned int sad =
773
      cpi->ppi->fn_ptr[plane_bsize].sdf(src, src_stride, dst, dst_stride);
774
  const double sad_norm =
775
      (double)sad / (1 << num_pels_log2_lookup[plane_bsize]);
776
777
  fprintf(fout, " %g %g", sse_norm, sad_norm);
778
779
  double sse_norm_arr[4], sad_norm_arr[4];
780
  get_2x2_normalized_sses_and_sads(cpi, plane_bsize, src, src_stride, dst,
781
                                   dst_stride, src_diff, diff_stride,
782
                                   sse_norm_arr, sad_norm_arr);
783
  if (shift) {
784
    for (int k = 0; k < 4; ++k) sse_norm_arr[k] /= (1 << (2 * shift));
785
    for (int k = 0; k < 4; ++k) sad_norm_arr[k] /= (1 << shift);
786
  }
787
  for (int i = 0; i < 4; ++i) {
788
    fprintf(fout, " %g", sse_norm_arr[i]);
789
  }
790
  for (int i = 0; i < 4; ++i) {
791
    fprintf(fout, " %g", sad_norm_arr[i]);
792
  }
793
794
  fprintf(fout, " %d %d %d %d", q_step, x->rdmult, bw, bh);
795
796
  int model_rate;
797
  int64_t model_dist;
798
  model_rd_sse_fn[MODELRD_CURVFIT](cpi, x, plane_bsize, plane, sse, num_samples,
799
                                   &model_rate, &model_dist);
800
  const double model_rdcost_norm =
801
      (double)RDCOST(x->rdmult, model_rate, model_dist) / num_samples;
802
  const double model_rate_norm = (double)model_rate / num_samples;
803
  const double model_dist_norm = (double)model_dist / num_samples;
804
  fprintf(fout, " %g %g %g", model_rate_norm, model_dist_norm,
805
          model_rdcost_norm);
806
807
  double mean;
808
  if (is_cur_buf_hbd(xd)) {
809
    mean = get_highbd_diff_mean(p->src.buf, p->src.stride, pd->dst.buf,
810
                                pd->dst.stride, bw, bh);
811
  } else {
812
    mean = get_diff_mean(p->src.buf, p->src.stride, pd->dst.buf, pd->dst.stride,
813
                         bw, bh);
814
  }
815
  mean /= (1 << shift);
816
  float hor_corr, vert_corr;
817
  av1_get_horver_correlation_full(src_diff, diff_stride, bw, bh, &hor_corr,
818
                                  &vert_corr);
819
  fprintf(fout, " %g %g %g", mean, hor_corr, vert_corr);
820
821
  double hdist[4] = { 0 }, vdist[4] = { 0 };
822
  get_energy_distribution_fine(cpi, plane_bsize, src, src_stride, dst,
823
                               dst_stride, 1, hdist, vdist);
824
  fprintf(fout, " %g %g %g %g %g %g %g %g", hdist[0], hdist[1], hdist[2],
825
          hdist[3], vdist[0], vdist[1], vdist[2], vdist[3]);
826
827
  if (cpi->sf.inter_sf.inter_mode_rd_model_estimation == 1) {
828
    assert(tile_data->inter_mode_rd_models[plane_bsize].ready);
829
    const int64_t overall_sse = get_sse(cpi, x);
830
    int est_residue_cost = 0;
831
    int64_t est_dist = 0;
832
    get_est_rate_dist(tile_data, plane_bsize, overall_sse, &est_residue_cost,
833
                      &est_dist);
834
    const double est_residue_cost_norm = (double)est_residue_cost / num_samples;
835
    const double est_dist_norm = (double)est_dist / num_samples;
836
    const double est_rdcost_norm =
837
        (double)RDCOST(x->rdmult, est_residue_cost, est_dist) / num_samples;
838
    fprintf(fout, " %g %g %g", est_residue_cost_norm, est_dist_norm,
839
            est_rdcost_norm);
840
  }
841
842
  fprintf(fout, "\n");
843
  fclose(fout);
844
}
845
#endif  // CONFIG_COLLECT_RD_STATS >= 2
846
#endif  // CONFIG_COLLECT_RD_STATS
847
848
static AOM_INLINE void inverse_transform_block_facade(MACROBLOCK *const x,
849
                                                      int plane, int block,
850
                                                      int blk_row, int blk_col,
851
                                                      int eob,
852
47.8k
                                                      int reduced_tx_set) {
853
47.8k
  if (!eob) return;
854
47.8k
  struct macroblock_plane *const p = &x->plane[plane];
855
47.8k
  MACROBLOCKD *const xd = &x->e_mbd;
856
47.8k
  tran_low_t *dqcoeff = p->dqcoeff + BLOCK_OFFSET(block);
857
47.8k
  const PLANE_TYPE plane_type = get_plane_type(plane);
858
47.8k
  const TX_SIZE tx_size = av1_get_tx_size(plane, xd);
859
47.8k
  const TX_TYPE tx_type = av1_get_tx_type(xd, plane_type, blk_row, blk_col,
860
47.8k
                                          tx_size, reduced_tx_set);
861
862
47.8k
  struct macroblockd_plane *const pd = &xd->plane[plane];
863
47.8k
  const int dst_stride = pd->dst.stride;
864
47.8k
  uint8_t *dst = &pd->dst.buf[(blk_row * dst_stride + blk_col) << MI_SIZE_LOG2];
865
47.8k
  av1_inverse_transform_block(xd, dqcoeff, plane, tx_type, tx_size, dst,
866
47.8k
                              dst_stride, eob, reduced_tx_set);
867
47.8k
}
868
869
static INLINE void recon_intra(const AV1_COMP *cpi, MACROBLOCK *x, int plane,
870
                               int block, int blk_row, int blk_col,
871
                               BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
872
                               const TXB_CTX *const txb_ctx, int skip_trellis,
873
                               TX_TYPE best_tx_type, int do_quant,
874
44.6M
                               int *rate_cost, uint16_t best_eob) {
875
44.6M
  const AV1_COMMON *cm = &cpi->common;
876
44.6M
  MACROBLOCKD *xd = &x->e_mbd;
877
44.6M
  MB_MODE_INFO *mbmi = xd->mi[0];
878
44.6M
  const int is_inter = is_inter_block(mbmi);
879
44.6M
  if (!is_inter && best_eob &&
880
44.6M
      (blk_row + tx_size_high_unit[tx_size] < mi_size_high[plane_bsize] ||
881
313k
       blk_col + tx_size_wide_unit[tx_size] < mi_size_wide[plane_bsize])) {
882
    // if the quantized coefficients are stored in the dqcoeff buffer, we don't
883
    // need to do transform and quantization again.
884
47.8k
    if (do_quant) {
885
0
      TxfmParam txfm_param_intra;
886
0
      QUANT_PARAM quant_param_intra;
887
0
      av1_setup_xform(cm, x, tx_size, best_tx_type, &txfm_param_intra);
888
0
      av1_setup_quant(tx_size, !skip_trellis,
889
0
                      skip_trellis
890
0
                          ? (USE_B_QUANT_NO_TRELLIS ? AV1_XFORM_QUANT_B
891
0
                                                    : AV1_XFORM_QUANT_FP)
892
0
                          : AV1_XFORM_QUANT_FP,
893
0
                      cpi->oxcf.q_cfg.quant_b_adapt, &quant_param_intra);
894
0
      av1_setup_qmatrix(&cm->quant_params, xd, plane, tx_size, best_tx_type,
895
0
                        &quant_param_intra);
896
0
      av1_xform_quant(x, plane, block, blk_row, blk_col, plane_bsize,
897
0
                      &txfm_param_intra, &quant_param_intra);
898
0
      if (quant_param_intra.use_optimize_b) {
899
0
        av1_optimize_b(cpi, x, plane, block, tx_size, best_tx_type, txb_ctx,
900
0
                       rate_cost);
901
0
      }
902
0
    }
903
904
47.8k
    inverse_transform_block_facade(x, plane, block, blk_row, blk_col,
905
47.8k
                                   x->plane[plane].eobs[block],
906
47.8k
                                   cm->features.reduced_tx_set_used);
907
908
    // This may happen because of hash collision. The eob stored in the hash
909
    // table is non-zero, but the real eob is zero. We need to make sure tx_type
910
    // is DCT_DCT in this case.
911
47.8k
    if (plane == 0 && x->plane[plane].eobs[block] == 0 &&
912
47.8k
        best_tx_type != DCT_DCT) {
913
0
      update_txk_array(xd, blk_row, blk_col, tx_size, DCT_DCT);
914
0
    }
915
47.8k
  }
916
44.6M
}
917
918
static unsigned pixel_dist_visible_only(
919
    const AV1_COMP *const cpi, const MACROBLOCK *x, const uint8_t *src,
920
    const int src_stride, const uint8_t *dst, const int dst_stride,
921
    const BLOCK_SIZE tx_bsize, int txb_rows, int txb_cols, int visible_rows,
922
166k
    int visible_cols) {
923
166k
  unsigned sse;
924
925
166k
  if (txb_rows == visible_rows && txb_cols == visible_cols) {
926
147k
    cpi->ppi->fn_ptr[tx_bsize].vf(src, src_stride, dst, dst_stride, &sse);
927
147k
    return sse;
928
147k
  }
929
930
18.6k
#if CONFIG_AV1_HIGHBITDEPTH
931
18.6k
  const MACROBLOCKD *xd = &x->e_mbd;
932
18.6k
  if (is_cur_buf_hbd(xd)) {
933
0
    uint64_t sse64 = aom_highbd_sse_odd_size(src, src_stride, dst, dst_stride,
934
0
                                             visible_cols, visible_rows);
935
0
    return (unsigned int)ROUND_POWER_OF_TWO(sse64, (xd->bd - 8) * 2);
936
0
  }
937
#else
938
  (void)x;
939
#endif
940
18.6k
  sse = aom_sse_odd_size(src, src_stride, dst, dst_stride, visible_cols,
941
18.6k
                         visible_rows);
942
18.6k
  return sse;
943
18.6k
}
944
945
// Compute the pixel domain distortion from src and dst on all visible 4x4s in
946
// the
947
// transform block.
948
static unsigned pixel_dist(const AV1_COMP *const cpi, const MACROBLOCK *x,
949
                           int plane, const uint8_t *src, const int src_stride,
950
                           const uint8_t *dst, const int dst_stride,
951
                           int blk_row, int blk_col,
952
                           const BLOCK_SIZE plane_bsize,
953
166k
                           const BLOCK_SIZE tx_bsize) {
954
166k
  int txb_rows, txb_cols, visible_rows, visible_cols;
955
166k
  const MACROBLOCKD *xd = &x->e_mbd;
956
957
166k
  get_txb_dimensions(xd, plane, plane_bsize, blk_row, blk_col, tx_bsize,
958
166k
                     &txb_cols, &txb_rows, &visible_cols, &visible_rows);
959
166k
  assert(visible_rows > 0);
960
166k
  assert(visible_cols > 0);
961
962
166k
  unsigned sse = pixel_dist_visible_only(cpi, x, src, src_stride, dst,
963
166k
                                         dst_stride, tx_bsize, txb_rows,
964
166k
                                         txb_cols, visible_rows, visible_cols);
965
966
166k
  return sse;
967
166k
}
968
969
static INLINE int64_t dist_block_px_domain(const AV1_COMP *cpi, MACROBLOCK *x,
970
                                           int plane, BLOCK_SIZE plane_bsize,
971
                                           int block, int blk_row, int blk_col,
972
166k
                                           TX_SIZE tx_size) {
973
166k
  MACROBLOCKD *const xd = &x->e_mbd;
974
166k
  const struct macroblock_plane *const p = &x->plane[plane];
975
166k
  const uint16_t eob = p->eobs[block];
976
166k
  const BLOCK_SIZE tx_bsize = txsize_to_bsize[tx_size];
977
166k
  const int bsw = block_size_wide[tx_bsize];
978
166k
  const int bsh = block_size_high[tx_bsize];
979
166k
  const int src_stride = x->plane[plane].src.stride;
980
166k
  const int dst_stride = xd->plane[plane].dst.stride;
981
  // Scale the transform block index to pixel unit.
982
166k
  const int src_idx = (blk_row * src_stride + blk_col) << MI_SIZE_LOG2;
983
166k
  const int dst_idx = (blk_row * dst_stride + blk_col) << MI_SIZE_LOG2;
984
166k
  const uint8_t *src = &x->plane[plane].src.buf[src_idx];
985
166k
  const uint8_t *dst = &xd->plane[plane].dst.buf[dst_idx];
986
166k
  const tran_low_t *dqcoeff = p->dqcoeff + BLOCK_OFFSET(block);
987
988
166k
  assert(cpi != NULL);
989
166k
  assert(tx_size_wide_log2[0] == tx_size_high_log2[0]);
990
991
166k
  uint8_t *recon;
992
166k
  DECLARE_ALIGNED(16, uint16_t, recon16[MAX_TX_SQUARE]);
993
994
166k
#if CONFIG_AV1_HIGHBITDEPTH
995
166k
  if (is_cur_buf_hbd(xd)) {
996
0
    recon = CONVERT_TO_BYTEPTR(recon16);
997
0
    aom_highbd_convolve_copy(CONVERT_TO_SHORTPTR(dst), dst_stride,
998
0
                             CONVERT_TO_SHORTPTR(recon), MAX_TX_SIZE, bsw, bsh);
999
166k
  } else {
1000
166k
    recon = (uint8_t *)recon16;
1001
166k
    aom_convolve_copy(dst, dst_stride, recon, MAX_TX_SIZE, bsw, bsh);
1002
166k
  }
1003
#else
1004
  recon = (uint8_t *)recon16;
1005
  aom_convolve_copy(dst, dst_stride, recon, MAX_TX_SIZE, bsw, bsh);
1006
#endif
1007
1008
166k
  const PLANE_TYPE plane_type = get_plane_type(plane);
1009
166k
  TX_TYPE tx_type = av1_get_tx_type(xd, plane_type, blk_row, blk_col, tx_size,
1010
166k
                                    cpi->common.features.reduced_tx_set_used);
1011
166k
  av1_inverse_transform_block(xd, dqcoeff, plane, tx_type, tx_size, recon,
1012
166k
                              MAX_TX_SIZE, eob,
1013
166k
                              cpi->common.features.reduced_tx_set_used);
1014
1015
166k
  return 16 * pixel_dist(cpi, x, plane, src, src_stride, recon, MAX_TX_SIZE,
1016
166k
                         blk_row, blk_col, plane_bsize, tx_bsize);
1017
166k
}
1018
1019
// pruning thresholds for prune_txk_type and prune_txk_type_separ
1020
static const int prune_factors[5] = { 200, 200, 120, 80, 40 };  // scale 1000
1021
static const int mul_factors[5] = { 80, 80, 70, 50, 30 };       // scale 100
1022
1023
// R-D costs are sorted in ascending order.
1024
280k
static INLINE void sort_rd(int64_t rds[], int txk[], int len) {
1025
280k
  int i, j, k;
1026
1027
1.40M
  for (i = 1; i <= len - 1; ++i) {
1028
3.70M
    for (j = 0; j < i; ++j) {
1029
2.71M
      if (rds[j] > rds[i]) {
1030
127k
        int64_t temprd;
1031
127k
        int tempi;
1032
1033
127k
        temprd = rds[i];
1034
127k
        tempi = txk[i];
1035
1036
343k
        for (k = i; k > j; k--) {
1037
216k
          rds[k] = rds[k - 1];
1038
216k
          txk[k] = txk[k - 1];
1039
216k
        }
1040
1041
127k
        rds[j] = temprd;
1042
127k
        txk[j] = tempi;
1043
127k
        break;
1044
127k
      }
1045
2.71M
    }
1046
1.12M
  }
1047
280k
}
1048
1049
static INLINE void dist_block_tx_domain(MACROBLOCK *x, int plane, int block,
1050
                                        TX_SIZE tx_size, int64_t *out_dist,
1051
1.70M
                                        int64_t *out_sse) {
1052
1.70M
  const struct macroblock_plane *const p = &x->plane[plane];
1053
  // Transform domain distortion computation is more efficient as it does
1054
  // not involve an inverse transform, but it is less accurate.
1055
1.70M
  const int buffer_length = av1_get_max_eob(tx_size);
1056
1.70M
  int64_t this_sse;
1057
  // TX-domain results need to shift down to Q2/D10 to match pixel
1058
  // domain distortion values which are in Q2^2
1059
1.70M
  int shift = (MAX_TX_SCALE - av1_get_tx_scale(tx_size)) * 2;
1060
1.70M
  const int block_offset = BLOCK_OFFSET(block);
1061
1.70M
  tran_low_t *const coeff = p->coeff + block_offset;
1062
1.70M
  tran_low_t *const dqcoeff = p->dqcoeff + block_offset;
1063
1.70M
#if CONFIG_AV1_HIGHBITDEPTH
1064
1.70M
  MACROBLOCKD *const xd = &x->e_mbd;
1065
1.70M
  if (is_cur_buf_hbd(xd))
1066
0
    *out_dist = av1_highbd_block_error(coeff, dqcoeff, buffer_length, &this_sse,
1067
0
                                       xd->bd);
1068
1.70M
  else
1069
1.70M
#endif
1070
1.70M
    *out_dist = av1_block_error(coeff, dqcoeff, buffer_length, &this_sse);
1071
1072
1.70M
  *out_dist = RIGHT_SIGNED_SHIFT(*out_dist, shift);
1073
1.70M
  *out_sse = RIGHT_SIGNED_SHIFT(this_sse, shift);
1074
1.70M
}
1075
1076
uint16_t prune_txk_type_separ(const AV1_COMP *cpi, MACROBLOCK *x, int plane,
1077
                              int block, TX_SIZE tx_size, int blk_row,
1078
                              int blk_col, BLOCK_SIZE plane_bsize, int *txk_map,
1079
                              int16_t allowed_tx_mask, int prune_factor,
1080
                              const TXB_CTX *const txb_ctx,
1081
                              int reduced_tx_set_used, int64_t ref_best_rd,
1082
0
                              int num_sel) {
1083
0
  const AV1_COMMON *cm = &cpi->common;
1084
0
  MACROBLOCKD *xd = &x->e_mbd;
1085
1086
0
  int idx;
1087
1088
0
  int64_t rds_v[4];
1089
0
  int64_t rds_h[4];
1090
0
  int idx_v[4] = { 0, 1, 2, 3 };
1091
0
  int idx_h[4] = { 0, 1, 2, 3 };
1092
0
  int skip_v[4] = { 0 };
1093
0
  int skip_h[4] = { 0 };
1094
0
  const int idx_map[16] = {
1095
0
    DCT_DCT,      DCT_ADST,      DCT_FLIPADST,      V_DCT,
1096
0
    ADST_DCT,     ADST_ADST,     ADST_FLIPADST,     V_ADST,
1097
0
    FLIPADST_DCT, FLIPADST_ADST, FLIPADST_FLIPADST, V_FLIPADST,
1098
0
    H_DCT,        H_ADST,        H_FLIPADST,        IDTX
1099
0
  };
1100
1101
0
  const int sel_pattern_v[16] = {
1102
0
    0, 0, 1, 1, 0, 2, 1, 2, 2, 0, 3, 1, 3, 2, 3, 3
1103
0
  };
1104
0
  const int sel_pattern_h[16] = {
1105
0
    0, 1, 0, 1, 2, 0, 2, 1, 2, 3, 0, 3, 1, 3, 2, 3
1106
0
  };
1107
1108
0
  QUANT_PARAM quant_param;
1109
0
  TxfmParam txfm_param;
1110
0
  av1_setup_xform(cm, x, tx_size, DCT_DCT, &txfm_param);
1111
0
  av1_setup_quant(tx_size, 1, AV1_XFORM_QUANT_B, cpi->oxcf.q_cfg.quant_b_adapt,
1112
0
                  &quant_param);
1113
0
  int tx_type;
1114
  // to ensure we can try ones even outside of ext_tx_set of current block
1115
  // this function should only be called for size < 16
1116
0
  assert(txsize_sqr_up_map[tx_size] <= TX_16X16);
1117
0
  txfm_param.tx_set_type = EXT_TX_SET_ALL16;
1118
1119
0
  int rate_cost = 0;
1120
0
  int64_t dist = 0, sse = 0;
1121
  // evaluate horizontal with vertical DCT
1122
0
  for (idx = 0; idx < 4; ++idx) {
1123
0
    tx_type = idx_map[idx];
1124
0
    txfm_param.tx_type = tx_type;
1125
1126
0
    av1_setup_qmatrix(&cm->quant_params, xd, plane, tx_size, tx_type,
1127
0
                      &quant_param);
1128
1129
0
    av1_xform_quant(x, plane, block, blk_row, blk_col, plane_bsize, &txfm_param,
1130
0
                    &quant_param);
1131
1132
0
    dist_block_tx_domain(x, plane, block, tx_size, &dist, &sse);
1133
1134
0
    rate_cost = av1_cost_coeffs_txb_laplacian(x, plane, block, tx_size, tx_type,
1135
0
                                              txb_ctx, reduced_tx_set_used, 0);
1136
1137
0
    rds_h[idx] = RDCOST(x->rdmult, rate_cost, dist);
1138
1139
0
    if ((rds_h[idx] - (rds_h[idx] >> 2)) > ref_best_rd) {
1140
0
      skip_h[idx] = 1;
1141
0
    }
1142
0
  }
1143
0
  sort_rd(rds_h, idx_h, 4);
1144
0
  for (idx = 1; idx < 4; idx++) {
1145
0
    if (rds_h[idx] > rds_h[0] * 1.2) skip_h[idx_h[idx]] = 1;
1146
0
  }
1147
1148
0
  if (skip_h[idx_h[0]]) return (uint16_t)0xFFFF;
1149
1150
  // evaluate vertical with the best horizontal chosen
1151
0
  rds_v[0] = rds_h[0];
1152
0
  int start_v = 1, end_v = 4;
1153
0
  const int *idx_map_v = idx_map + idx_h[0];
1154
1155
0
  for (idx = start_v; idx < end_v; ++idx) {
1156
0
    tx_type = idx_map_v[idx_v[idx] * 4];
1157
0
    txfm_param.tx_type = tx_type;
1158
1159
0
    av1_setup_qmatrix(&cm->quant_params, xd, plane, tx_size, tx_type,
1160
0
                      &quant_param);
1161
1162
0
    av1_xform_quant(x, plane, block, blk_row, blk_col, plane_bsize, &txfm_param,
1163
0
                    &quant_param);
1164
1165
0
    dist_block_tx_domain(x, plane, block, tx_size, &dist, &sse);
1166
1167
0
    rate_cost = av1_cost_coeffs_txb_laplacian(x, plane, block, tx_size, tx_type,
1168
0
                                              txb_ctx, reduced_tx_set_used, 0);
1169
1170
0
    rds_v[idx] = RDCOST(x->rdmult, rate_cost, dist);
1171
1172
0
    if ((rds_v[idx] - (rds_v[idx] >> 2)) > ref_best_rd) {
1173
0
      skip_v[idx] = 1;
1174
0
    }
1175
0
  }
1176
0
  sort_rd(rds_v, idx_v, 4);
1177
0
  for (idx = 1; idx < 4; idx++) {
1178
0
    if (rds_v[idx] > rds_v[0] * 1.2) skip_v[idx_v[idx]] = 1;
1179
0
  }
1180
1181
  // combine rd_h and rd_v to prune tx candidates
1182
0
  int i_v, i_h;
1183
0
  int64_t rds[16];
1184
0
  int num_cand = 0, last = TX_TYPES - 1;
1185
1186
0
  for (int i = 0; i < 16; i++) {
1187
0
    i_v = sel_pattern_v[i];
1188
0
    i_h = sel_pattern_h[i];
1189
0
    tx_type = idx_map[idx_v[i_v] * 4 + idx_h[i_h]];
1190
0
    if (!(allowed_tx_mask & (1 << tx_type)) || skip_h[idx_h[i_h]] ||
1191
0
        skip_v[idx_v[i_v]]) {
1192
0
      txk_map[last] = tx_type;
1193
0
      last--;
1194
0
    } else {
1195
0
      txk_map[num_cand] = tx_type;
1196
0
      rds[num_cand] = rds_v[i_v] + rds_h[i_h];
1197
0
      if (rds[num_cand] == 0) rds[num_cand] = 1;
1198
0
      num_cand++;
1199
0
    }
1200
0
  }
1201
0
  sort_rd(rds, txk_map, num_cand);
1202
1203
0
  uint16_t prune = (uint16_t)(~(1 << txk_map[0]));
1204
0
  num_sel = AOMMIN(num_sel, num_cand);
1205
1206
0
  for (int i = 1; i < num_sel; i++) {
1207
0
    int64_t factor = 1800 * (rds[i] - rds[0]) / (rds[0]);
1208
0
    if (factor < (int64_t)prune_factor)
1209
0
      prune &= ~(1 << txk_map[i]);
1210
0
    else
1211
0
      break;
1212
0
  }
1213
0
  return prune;
1214
0
}
1215
1216
uint16_t prune_txk_type(const AV1_COMP *cpi, MACROBLOCK *x, int plane,
1217
                        int block, TX_SIZE tx_size, int blk_row, int blk_col,
1218
                        BLOCK_SIZE plane_bsize, int *txk_map,
1219
                        uint16_t allowed_tx_mask, int prune_factor,
1220
280k
                        const TXB_CTX *const txb_ctx, int reduced_tx_set_used) {
1221
280k
  const AV1_COMMON *cm = &cpi->common;
1222
280k
  MACROBLOCKD *xd = &x->e_mbd;
1223
280k
  int tx_type;
1224
1225
280k
  int64_t rds[TX_TYPES];
1226
1227
280k
  int num_cand = 0;
1228
280k
  int last = TX_TYPES - 1;
1229
1230
280k
  TxfmParam txfm_param;
1231
280k
  QUANT_PARAM quant_param;
1232
280k
  av1_setup_xform(cm, x, tx_size, DCT_DCT, &txfm_param);
1233
280k
  av1_setup_quant(tx_size, 1, AV1_XFORM_QUANT_B, cpi->oxcf.q_cfg.quant_b_adapt,
1234
280k
                  &quant_param);
1235
1236
4.75M
  for (int idx = 0; idx < TX_TYPES; idx++) {
1237
4.47M
    tx_type = idx;
1238
4.47M
    int rate_cost = 0;
1239
4.47M
    int64_t dist = 0, sse = 0;
1240
4.47M
    if (!(allowed_tx_mask & (1 << tx_type))) {
1241
3.08M
      txk_map[last] = tx_type;
1242
3.08M
      last--;
1243
3.08M
      continue;
1244
3.08M
    }
1245
1.39M
    txfm_param.tx_type = tx_type;
1246
1247
1.39M
    av1_setup_qmatrix(&cm->quant_params, xd, plane, tx_size, tx_type,
1248
1.39M
                      &quant_param);
1249
1250
    // do txfm and quantization
1251
1.39M
    av1_xform_quant(x, plane, block, blk_row, blk_col, plane_bsize, &txfm_param,
1252
1.39M
                    &quant_param);
1253
    // estimate rate cost
1254
1.39M
    rate_cost = av1_cost_coeffs_txb_laplacian(x, plane, block, tx_size, tx_type,
1255
1.39M
                                              txb_ctx, reduced_tx_set_used, 0);
1256
    // tx domain dist
1257
1.39M
    dist_block_tx_domain(x, plane, block, tx_size, &dist, &sse);
1258
1259
1.39M
    txk_map[num_cand] = tx_type;
1260
1.39M
    rds[num_cand] = RDCOST(x->rdmult, rate_cost, dist);
1261
1.39M
    if (rds[num_cand] == 0) rds[num_cand] = 1;
1262
1.39M
    num_cand++;
1263
1.39M
  }
1264
1265
280k
  if (num_cand == 0) return (uint16_t)0xFFFF;
1266
1267
280k
  sort_rd(rds, txk_map, num_cand);
1268
280k
  uint16_t prune = (uint16_t)(~(1 << txk_map[0]));
1269
1270
  // 0 < prune_factor <= 1000 controls aggressiveness
1271
280k
  int64_t factor = 0;
1272
1.34M
  for (int idx = 1; idx < num_cand; idx++) {
1273
1.08M
    factor = 1000 * (rds[idx] - rds[0]) / rds[0];
1274
1.08M
    if (factor < (int64_t)prune_factor)
1275
1.06M
      prune &= ~(1 << txk_map[idx]);
1276
18.8k
    else
1277
18.8k
      break;
1278
1.08M
  }
1279
280k
  return prune;
1280
280k
}
1281
1282
// These thresholds were calibrated to provide a certain number of TX types
1283
// pruned by the model on average, i.e. selecting a threshold with index i
1284
// will lead to pruning i+1 TX types on average
1285
static const float *prune_2D_adaptive_thresholds[] = {
1286
  // TX_4X4
1287
  (float[]){ 0.00549f, 0.01306f, 0.02039f, 0.02747f, 0.03406f, 0.04065f,
1288
             0.04724f, 0.05383f, 0.06067f, 0.06799f, 0.07605f, 0.08533f,
1289
             0.09778f, 0.11780f },
1290
  // TX_8X8
1291
  (float[]){ 0.00037f, 0.00183f, 0.00525f, 0.01038f, 0.01697f, 0.02502f,
1292
             0.03381f, 0.04333f, 0.05286f, 0.06287f, 0.07434f, 0.08850f,
1293
             0.10803f, 0.14124f },
1294
  // TX_16X16
1295
  (float[]){ 0.01404f, 0.02000f, 0.04211f, 0.05164f, 0.05798f, 0.06335f,
1296
             0.06897f, 0.07629f, 0.08875f, 0.11169f },
1297
  // TX_32X32
1298
  NULL,
1299
  // TX_64X64
1300
  NULL,
1301
  // TX_4X8
1302
  (float[]){ 0.00183f, 0.00745f, 0.01428f, 0.02185f, 0.02966f, 0.03723f,
1303
             0.04456f, 0.05188f, 0.05920f, 0.06702f, 0.07605f, 0.08704f,
1304
             0.10168f, 0.12585f },
1305
  // TX_8X4
1306
  (float[]){ 0.00085f, 0.00476f, 0.01135f, 0.01892f, 0.02698f, 0.03528f,
1307
             0.04358f, 0.05164f, 0.05994f, 0.06848f, 0.07849f, 0.09021f,
1308
             0.10583f, 0.13123f },
1309
  // TX_8X16
1310
  (float[]){ 0.00037f, 0.00232f, 0.00671f, 0.01257f, 0.01965f, 0.02722f,
1311
             0.03552f, 0.04382f, 0.05237f, 0.06189f, 0.07336f, 0.08728f,
1312
             0.10730f, 0.14221f },
1313
  // TX_16X8
1314
  (float[]){ 0.00061f, 0.00330f, 0.00818f, 0.01453f, 0.02185f, 0.02966f,
1315
             0.03772f, 0.04578f, 0.05383f, 0.06262f, 0.07288f, 0.08582f,
1316
             0.10339f, 0.13464f },
1317
  // TX_16X32
1318
  NULL,
1319
  // TX_32X16
1320
  NULL,
1321
  // TX_32X64
1322
  NULL,
1323
  // TX_64X32
1324
  NULL,
1325
  // TX_4X16
1326
  (float[]){ 0.00232f, 0.00671f, 0.01257f, 0.01941f, 0.02673f, 0.03430f,
1327
             0.04211f, 0.04968f, 0.05750f, 0.06580f, 0.07507f, 0.08655f,
1328
             0.10242f, 0.12878f },
1329
  // TX_16X4
1330
  (float[]){ 0.00110f, 0.00525f, 0.01208f, 0.01990f, 0.02795f, 0.03601f,
1331
             0.04358f, 0.05115f, 0.05896f, 0.06702f, 0.07629f, 0.08752f,
1332
             0.10217f, 0.12610f },
1333
  // TX_8X32
1334
  NULL,
1335
  // TX_32X8
1336
  NULL,
1337
  // TX_16X64
1338
  NULL,
1339
  // TX_64X16
1340
  NULL,
1341
};
1342
1343
static INLINE float get_adaptive_thresholds(
1344
    TX_SIZE tx_size, TxSetType tx_set_type,
1345
0
    TX_TYPE_PRUNE_MODE prune_2d_txfm_mode) {
1346
0
  const int prune_aggr_table[5][2] = {
1347
0
    { 4, 1 }, { 6, 3 }, { 9, 6 }, { 9, 6 }, { 12, 9 }
1348
0
  };
1349
0
  int pruning_aggressiveness = 0;
1350
0
  if (tx_set_type == EXT_TX_SET_ALL16)
1351
0
    pruning_aggressiveness =
1352
0
        prune_aggr_table[prune_2d_txfm_mode - TX_TYPE_PRUNE_1][0];
1353
0
  else if (tx_set_type == EXT_TX_SET_DTT9_IDTX_1DDCT)
1354
0
    pruning_aggressiveness =
1355
0
        prune_aggr_table[prune_2d_txfm_mode - TX_TYPE_PRUNE_1][1];
1356
1357
0
  return prune_2D_adaptive_thresholds[tx_size][pruning_aggressiveness];
1358
0
}
1359
1360
static AOM_INLINE void get_energy_distribution_finer(const int16_t *diff,
1361
                                                     int stride, int bw, int bh,
1362
                                                     float *hordist,
1363
0
                                                     float *verdist) {
1364
  // First compute downscaled block energy values (esq); downscale factors
1365
  // are defined by w_shift and h_shift.
1366
0
  unsigned int esq[256];
1367
0
  const int w_shift = bw <= 8 ? 0 : 1;
1368
0
  const int h_shift = bh <= 8 ? 0 : 1;
1369
0
  const int esq_w = bw >> w_shift;
1370
0
  const int esq_h = bh >> h_shift;
1371
0
  const int esq_sz = esq_w * esq_h;
1372
0
  int i, j;
1373
0
  memset(esq, 0, esq_sz * sizeof(esq[0]));
1374
0
  if (w_shift) {
1375
0
    for (i = 0; i < bh; i++) {
1376
0
      unsigned int *cur_esq_row = esq + (i >> h_shift) * esq_w;
1377
0
      const int16_t *cur_diff_row = diff + i * stride;
1378
0
      for (j = 0; j < bw; j += 2) {
1379
0
        cur_esq_row[j >> 1] += (cur_diff_row[j] * cur_diff_row[j] +
1380
0
                                cur_diff_row[j + 1] * cur_diff_row[j + 1]);
1381
0
      }
1382
0
    }
1383
0
  } else {
1384
0
    for (i = 0; i < bh; i++) {
1385
0
      unsigned int *cur_esq_row = esq + (i >> h_shift) * esq_w;
1386
0
      const int16_t *cur_diff_row = diff + i * stride;
1387
0
      for (j = 0; j < bw; j++) {
1388
0
        cur_esq_row[j] += cur_diff_row[j] * cur_diff_row[j];
1389
0
      }
1390
0
    }
1391
0
  }
1392
1393
0
  uint64_t total = 0;
1394
0
  for (i = 0; i < esq_sz; i++) total += esq[i];
1395
1396
  // Output hordist and verdist arrays are normalized 1D projections of esq
1397
0
  if (total == 0) {
1398
0
    float hor_val = 1.0f / esq_w;
1399
0
    for (j = 0; j < esq_w - 1; j++) hordist[j] = hor_val;
1400
0
    float ver_val = 1.0f / esq_h;
1401
0
    for (i = 0; i < esq_h - 1; i++) verdist[i] = ver_val;
1402
0
    return;
1403
0
  }
1404
1405
0
  const float e_recip = 1.0f / (float)total;
1406
0
  memset(hordist, 0, (esq_w - 1) * sizeof(hordist[0]));
1407
0
  memset(verdist, 0, (esq_h - 1) * sizeof(verdist[0]));
1408
0
  const unsigned int *cur_esq_row;
1409
0
  for (i = 0; i < esq_h - 1; i++) {
1410
0
    cur_esq_row = esq + i * esq_w;
1411
0
    for (j = 0; j < esq_w - 1; j++) {
1412
0
      hordist[j] += (float)cur_esq_row[j];
1413
0
      verdist[i] += (float)cur_esq_row[j];
1414
0
    }
1415
0
    verdist[i] += (float)cur_esq_row[j];
1416
0
  }
1417
0
  cur_esq_row = esq + i * esq_w;
1418
0
  for (j = 0; j < esq_w - 1; j++) hordist[j] += (float)cur_esq_row[j];
1419
1420
0
  for (j = 0; j < esq_w - 1; j++) hordist[j] *= e_recip;
1421
0
  for (i = 0; i < esq_h - 1; i++) verdist[i] *= e_recip;
1422
0
}
1423
1424
50.3M
static AOM_INLINE bool check_bit_mask(uint16_t mask, int val) {
1425
50.3M
  return mask & (1 << val);
1426
50.3M
}
1427
1428
0
static AOM_INLINE void set_bit_mask(uint16_t *mask, int val) {
1429
0
  *mask |= (1 << val);
1430
0
}
1431
1432
0
static AOM_INLINE void unset_bit_mask(uint16_t *mask, int val) {
1433
0
  *mask &= ~(1 << val);
1434
0
}
1435
1436
static void prune_tx_2D(MACROBLOCK *x, BLOCK_SIZE bsize, TX_SIZE tx_size,
1437
                        int blk_row, int blk_col, TxSetType tx_set_type,
1438
                        TX_TYPE_PRUNE_MODE prune_2d_txfm_mode, int *txk_map,
1439
0
                        uint16_t *allowed_tx_mask) {
1440
  // This table is used because the search order is different from the enum
1441
  // order.
1442
0
  static const int tx_type_table_2D[16] = {
1443
0
    DCT_DCT,      DCT_ADST,      DCT_FLIPADST,      V_DCT,
1444
0
    ADST_DCT,     ADST_ADST,     ADST_FLIPADST,     V_ADST,
1445
0
    FLIPADST_DCT, FLIPADST_ADST, FLIPADST_FLIPADST, V_FLIPADST,
1446
0
    H_DCT,        H_ADST,        H_FLIPADST,        IDTX
1447
0
  };
1448
0
  if (tx_set_type != EXT_TX_SET_ALL16 &&
1449
0
      tx_set_type != EXT_TX_SET_DTT9_IDTX_1DDCT)
1450
0
    return;
1451
#if CONFIG_NN_V2
1452
  NN_CONFIG_V2 *nn_config_hor = av1_tx_type_nnconfig_map_hor[tx_size];
1453
  NN_CONFIG_V2 *nn_config_ver = av1_tx_type_nnconfig_map_ver[tx_size];
1454
#else
1455
0
  const NN_CONFIG *nn_config_hor = av1_tx_type_nnconfig_map_hor[tx_size];
1456
0
  const NN_CONFIG *nn_config_ver = av1_tx_type_nnconfig_map_ver[tx_size];
1457
0
#endif
1458
0
  if (!nn_config_hor || !nn_config_ver) return;  // Model not established yet.
1459
1460
0
  float hfeatures[16], vfeatures[16];
1461
0
  float hscores[4], vscores[4];
1462
0
  float scores_2D_raw[16];
1463
0
  const int bw = tx_size_wide[tx_size];
1464
0
  const int bh = tx_size_high[tx_size];
1465
0
  const int hfeatures_num = bw <= 8 ? bw : bw / 2;
1466
0
  const int vfeatures_num = bh <= 8 ? bh : bh / 2;
1467
0
  assert(hfeatures_num <= 16);
1468
0
  assert(vfeatures_num <= 16);
1469
1470
0
  const struct macroblock_plane *const p = &x->plane[0];
1471
0
  const int diff_stride = block_size_wide[bsize];
1472
0
  const int16_t *diff = p->src_diff + 4 * blk_row * diff_stride + 4 * blk_col;
1473
0
  get_energy_distribution_finer(diff, diff_stride, bw, bh, hfeatures,
1474
0
                                vfeatures);
1475
1476
0
  av1_get_horver_correlation_full(diff, diff_stride, bw, bh,
1477
0
                                  &hfeatures[hfeatures_num - 1],
1478
0
                                  &vfeatures[vfeatures_num - 1]);
1479
1480
#if CONFIG_NN_V2
1481
  av1_nn_predict_v2(hfeatures, nn_config_hor, 0, hscores);
1482
  av1_nn_predict_v2(vfeatures, nn_config_ver, 0, vscores);
1483
#else
1484
0
  av1_nn_predict(hfeatures, nn_config_hor, 1, hscores);
1485
0
  av1_nn_predict(vfeatures, nn_config_ver, 1, vscores);
1486
0
#endif
1487
1488
0
  for (int i = 0; i < 4; i++) {
1489
0
    float *cur_scores_2D = scores_2D_raw + i * 4;
1490
0
    cur_scores_2D[0] = vscores[i] * hscores[0];
1491
0
    cur_scores_2D[1] = vscores[i] * hscores[1];
1492
0
    cur_scores_2D[2] = vscores[i] * hscores[2];
1493
0
    cur_scores_2D[3] = vscores[i] * hscores[3];
1494
0
  }
1495
1496
0
  assert(TX_TYPES == 16);
1497
  // This version of the function only works when there are at most 16 classes.
1498
  // So we will need to change the optimization or use av1_nn_softmax instead if
1499
  // this ever gets changed.
1500
0
  av1_nn_fast_softmax_16(scores_2D_raw, scores_2D_raw);
1501
1502
0
  const float score_thresh =
1503
0
      get_adaptive_thresholds(tx_size, tx_set_type, prune_2d_txfm_mode);
1504
1505
  // Always keep the TX type with the highest score, prune all others with
1506
  // score below score_thresh.
1507
0
  int max_score_i = 0;
1508
0
  float max_score = 0.0f;
1509
0
  uint16_t allow_bitmask = 0;
1510
0
  float sum_score = 0.0;
1511
  // Calculate sum of allowed tx type score and Populate allow bit mask based
1512
  // on score_thresh and allowed_tx_mask
1513
0
  int allow_count = 0;
1514
0
  int tx_type_allowed[16] = { TX_TYPE_INVALID, TX_TYPE_INVALID, TX_TYPE_INVALID,
1515
0
                              TX_TYPE_INVALID, TX_TYPE_INVALID, TX_TYPE_INVALID,
1516
0
                              TX_TYPE_INVALID, TX_TYPE_INVALID, TX_TYPE_INVALID,
1517
0
                              TX_TYPE_INVALID, TX_TYPE_INVALID, TX_TYPE_INVALID,
1518
0
                              TX_TYPE_INVALID, TX_TYPE_INVALID, TX_TYPE_INVALID,
1519
0
                              TX_TYPE_INVALID };
1520
0
  float scores_2D[16] = {
1521
0
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
1522
0
  };
1523
0
  for (int tx_idx = 0; tx_idx < TX_TYPES; tx_idx++) {
1524
0
    const int allow_tx_type =
1525
0
        check_bit_mask(*allowed_tx_mask, tx_type_table_2D[tx_idx]);
1526
0
    if (!allow_tx_type) {
1527
0
      continue;
1528
0
    }
1529
0
    if (scores_2D_raw[tx_idx] > max_score) {
1530
0
      max_score = scores_2D_raw[tx_idx];
1531
0
      max_score_i = tx_idx;
1532
0
    }
1533
0
    if (scores_2D_raw[tx_idx] >= score_thresh) {
1534
      // Set allow mask based on score_thresh
1535
0
      set_bit_mask(&allow_bitmask, tx_type_table_2D[tx_idx]);
1536
1537
      // Accumulate score of allowed tx type
1538
0
      sum_score += scores_2D_raw[tx_idx];
1539
1540
0
      scores_2D[allow_count] = scores_2D_raw[tx_idx];
1541
0
      tx_type_allowed[allow_count] = tx_type_table_2D[tx_idx];
1542
0
      allow_count += 1;
1543
0
    }
1544
0
  }
1545
0
  if (!check_bit_mask(allow_bitmask, tx_type_table_2D[max_score_i])) {
1546
    // If even the tx_type with max score is pruned, this means that no other
1547
    // tx_type is feasible. When this happens, we force enable max_score_i and
1548
    // end the search.
1549
0
    set_bit_mask(&allow_bitmask, tx_type_table_2D[max_score_i]);
1550
0
    memcpy(txk_map, tx_type_table_2D, sizeof(tx_type_table_2D));
1551
0
    *allowed_tx_mask = allow_bitmask;
1552
0
    return;
1553
0
  }
1554
1555
  // Sort tx type probability of all types
1556
0
  if (allow_count <= 8) {
1557
0
    av1_sort_fi32_8(scores_2D, tx_type_allowed);
1558
0
  } else {
1559
0
    av1_sort_fi32_16(scores_2D, tx_type_allowed);
1560
0
  }
1561
1562
  // Enable more pruning based on tx type probability and number of allowed tx
1563
  // types
1564
0
  if (prune_2d_txfm_mode >= TX_TYPE_PRUNE_4) {
1565
0
    float temp_score = 0.0;
1566
0
    float score_ratio = 0.0;
1567
0
    int tx_idx, tx_count = 0;
1568
0
    const float inv_sum_score = 100 / sum_score;
1569
    // Get allowed tx types based on sorted probability score and tx count
1570
0
    for (tx_idx = 0; tx_idx < allow_count; tx_idx++) {
1571
      // Skip the tx type which has more than 30% of cumulative
1572
      // probability and allowed tx type count is more than 2
1573
0
      if (score_ratio > 30.0 && tx_count >= 2) break;
1574
1575
0
      assert(check_bit_mask(allow_bitmask, tx_type_allowed[tx_idx]));
1576
      // Calculate cumulative probability
1577
0
      temp_score += scores_2D[tx_idx];
1578
1579
      // Calculate percentage of cumulative probability of allowed tx type
1580
0
      score_ratio = temp_score * inv_sum_score;
1581
0
      tx_count++;
1582
0
    }
1583
    // Set remaining tx types as pruned
1584
0
    for (; tx_idx < allow_count; tx_idx++)
1585
0
      unset_bit_mask(&allow_bitmask, tx_type_allowed[tx_idx]);
1586
0
  }
1587
1588
0
  memcpy(txk_map, tx_type_allowed, sizeof(tx_type_table_2D));
1589
0
  *allowed_tx_mask = allow_bitmask;
1590
0
}
1591
1592
0
static float get_dev(float mean, double x2_sum, int num) {
1593
0
  const float e_x2 = (float)(x2_sum / num);
1594
0
  const float diff = e_x2 - mean * mean;
1595
0
  const float dev = (diff > 0) ? sqrtf(diff) : 0;
1596
0
  return dev;
1597
0
}
1598
1599
// Feature used by the model to predict tx split: the mean and standard
1600
// deviation values of the block and sub-blocks.
1601
static AOM_INLINE void get_mean_dev_features(const int16_t *data, int stride,
1602
0
                                             int bw, int bh, float *feature) {
1603
0
  const int16_t *const data_ptr = &data[0];
1604
0
  const int subh = (bh >= bw) ? (bh >> 1) : bh;
1605
0
  const int subw = (bw >= bh) ? (bw >> 1) : bw;
1606
0
  const int num = bw * bh;
1607
0
  const int sub_num = subw * subh;
1608
0
  int feature_idx = 2;
1609
0
  int total_x_sum = 0;
1610
0
  int64_t total_x2_sum = 0;
1611
0
  int blk_idx = 0;
1612
0
  double mean2_sum = 0.0f;
1613
0
  float dev_sum = 0.0f;
1614
1615
0
  for (int row = 0; row < bh; row += subh) {
1616
0
    for (int col = 0; col < bw; col += subw) {
1617
0
      int x_sum;
1618
0
      int64_t x2_sum;
1619
      // TODO(any): Write a SIMD version. Clear registers.
1620
0
      aom_get_blk_sse_sum(data_ptr + row * stride + col, stride, subw, subh,
1621
0
                          &x_sum, &x2_sum);
1622
0
      total_x_sum += x_sum;
1623
0
      total_x2_sum += x2_sum;
1624
1625
0
      const float mean = (float)x_sum / sub_num;
1626
0
      const float dev = get_dev(mean, (double)x2_sum, sub_num);
1627
0
      feature[feature_idx++] = mean;
1628
0
      feature[feature_idx++] = dev;
1629
0
      mean2_sum += (double)(mean * mean);
1630
0
      dev_sum += dev;
1631
0
      blk_idx++;
1632
0
    }
1633
0
  }
1634
1635
0
  const float lvl0_mean = (float)total_x_sum / num;
1636
0
  feature[0] = lvl0_mean;
1637
0
  feature[1] = get_dev(lvl0_mean, (double)total_x2_sum, num);
1638
1639
0
  if (blk_idx > 1) {
1640
    // Deviation of means.
1641
0
    feature[feature_idx++] = get_dev(lvl0_mean, mean2_sum, blk_idx);
1642
    // Mean of deviations.
1643
0
    feature[feature_idx++] = dev_sum / blk_idx;
1644
0
  }
1645
0
}
1646
1647
static int ml_predict_tx_split(MACROBLOCK *x, BLOCK_SIZE bsize, int blk_row,
1648
0
                               int blk_col, TX_SIZE tx_size) {
1649
0
  const NN_CONFIG *nn_config = av1_tx_split_nnconfig_map[tx_size];
1650
0
  if (!nn_config) return -1;
1651
1652
0
  const int diff_stride = block_size_wide[bsize];
1653
0
  const int16_t *diff =
1654
0
      x->plane[0].src_diff + 4 * blk_row * diff_stride + 4 * blk_col;
1655
0
  const int bw = tx_size_wide[tx_size];
1656
0
  const int bh = tx_size_high[tx_size];
1657
1658
0
  float features[64] = { 0.0f };
1659
0
  get_mean_dev_features(diff, diff_stride, bw, bh, features);
1660
1661
0
  float score = 0.0f;
1662
0
  av1_nn_predict(features, nn_config, 1, &score);
1663
1664
0
  int int_score = (int)(score * 10000);
1665
0
  return clamp(int_score, -80000, 80000);
1666
0
}
1667
1668
static INLINE uint16_t
1669
get_tx_mask(const AV1_COMP *cpi, MACROBLOCK *x, int plane, int block,
1670
            int blk_row, int blk_col, BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
1671
            const TXB_CTX *const txb_ctx, FAST_TX_SEARCH_MODE ftxs_mode,
1672
44.3M
            int64_t ref_best_rd, TX_TYPE *allowed_txk_types, int *txk_map) {
1673
44.3M
  const AV1_COMMON *cm = &cpi->common;
1674
44.3M
  MACROBLOCKD *xd = &x->e_mbd;
1675
44.3M
  MB_MODE_INFO *mbmi = xd->mi[0];
1676
44.3M
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
1677
44.3M
  const int is_inter = is_inter_block(mbmi);
1678
44.3M
  const int fast_tx_search = ftxs_mode & FTXS_DCT_AND_1D_DCT_ONLY;
1679
  // if txk_allowed = TX_TYPES, >1 tx types are allowed, else, if txk_allowed <
1680
  // TX_TYPES, only that specific tx type is allowed.
1681
44.3M
  TX_TYPE txk_allowed = TX_TYPES;
1682
1683
44.3M
  const FRAME_UPDATE_TYPE update_type =
1684
44.3M
      get_frame_update_type(&cpi->ppi->gf_group, cpi->gf_frame_index);
1685
44.3M
  int use_actual_frame_probs = 1;
1686
44.3M
  const int *tx_type_probs;
1687
#if CONFIG_FRAME_PARALLEL_ENCODE && CONFIG_FPMT_TEST
1688
  use_actual_frame_probs =
1689
      (cpi->ppi->fpmt_unit_test_cfg == PARALLEL_SIMULATION_ENCODE) ? 0 : 1;
1690
  if (!use_actual_frame_probs) {
1691
    tx_type_probs =
1692
        (int *)cpi->ppi->temp_frame_probs.tx_type_probs[update_type][tx_size];
1693
  }
1694
#endif
1695
44.3M
  if (use_actual_frame_probs) {
1696
44.2M
    tx_type_probs = cpi->ppi->frame_probs.tx_type_probs[update_type][tx_size];
1697
44.2M
  }
1698
1699
44.3M
  if ((!is_inter && txfm_params->use_default_intra_tx_type) ||
1700
44.3M
      (is_inter && txfm_params->default_inter_tx_type_prob_thresh == 0)) {
1701
41.0M
    txk_allowed =
1702
41.0M
        get_default_tx_type(0, xd, tx_size, cpi->use_screen_content_tools);
1703
41.0M
  } else if (is_inter &&
1704
3.33M
             txfm_params->default_inter_tx_type_prob_thresh != INT_MAX) {
1705
0
    if (tx_type_probs[DEFAULT_INTER_TX_TYPE] >
1706
0
        txfm_params->default_inter_tx_type_prob_thresh) {
1707
0
      txk_allowed = DEFAULT_INTER_TX_TYPE;
1708
0
    } else {
1709
0
      int force_tx_type = 0;
1710
0
      int max_prob = 0;
1711
0
      const int tx_type_prob_threshold =
1712
0
          txfm_params->default_inter_tx_type_prob_thresh +
1713
0
          PROB_THRESH_OFFSET_TX_TYPE;
1714
0
      for (int i = 1; i < TX_TYPES; i++) {  // find maximum probability.
1715
0
        if (tx_type_probs[i] > max_prob) {
1716
0
          max_prob = tx_type_probs[i];
1717
0
          force_tx_type = i;
1718
0
        }
1719
0
      }
1720
0
      if (max_prob > tx_type_prob_threshold)  // force tx type with max prob.
1721
0
        txk_allowed = force_tx_type;
1722
0
      else if (x->rd_model == LOW_TXFM_RD) {
1723
0
        if (plane == 0) txk_allowed = DCT_DCT;
1724
0
      }
1725
0
    }
1726
3.33M
  } else if (x->rd_model == LOW_TXFM_RD) {
1727
0
    if (plane == 0) txk_allowed = DCT_DCT;
1728
0
  }
1729
1730
44.3M
  const TxSetType tx_set_type = av1_get_ext_tx_set_type(
1731
44.3M
      tx_size, is_inter, cm->features.reduced_tx_set_used);
1732
1733
44.3M
  TX_TYPE uv_tx_type = DCT_DCT;
1734
44.3M
  if (plane) {
1735
    // tx_type of PLANE_TYPE_UV should be the same as PLANE_TYPE_Y
1736
1.34M
    uv_tx_type = txk_allowed =
1737
1.34M
        av1_get_tx_type(xd, get_plane_type(plane), blk_row, blk_col, tx_size,
1738
1.34M
                        cm->features.reduced_tx_set_used);
1739
1.34M
  }
1740
44.3M
  PREDICTION_MODE intra_dir =
1741
44.3M
      mbmi->filter_intra_mode_info.use_filter_intra
1742
44.3M
          ? fimode_to_intradir[mbmi->filter_intra_mode_info.filter_intra_mode]
1743
44.3M
          : mbmi->mode;
1744
44.3M
  uint16_t ext_tx_used_flag =
1745
44.3M
      cpi->sf.tx_sf.tx_type_search.use_reduced_intra_txset != 0 &&
1746
44.3M
              tx_set_type == EXT_TX_SET_DTT4_IDTX_1DDCT
1747
44.3M
          ? av1_reduced_intra_tx_used_flag[intra_dir]
1748
44.3M
          : av1_ext_tx_used_flag[tx_set_type];
1749
1750
44.3M
  if (cpi->sf.tx_sf.tx_type_search.use_reduced_intra_txset == 2)
1751
0
    ext_tx_used_flag &= av1_derived_intra_tx_used_flag[intra_dir];
1752
1753
44.3M
  if (xd->lossless[mbmi->segment_id] || txsize_sqr_up_map[tx_size] > TX_32X32 ||
1754
44.3M
      ext_tx_used_flag == 0x0001 ||
1755
44.3M
      (is_inter && cpi->oxcf.txfm_cfg.use_inter_dct_only) ||
1756
44.3M
      (!is_inter && cpi->oxcf.txfm_cfg.use_intra_dct_only)) {
1757
41.5M
    txk_allowed = DCT_DCT;
1758
41.5M
  }
1759
1760
44.3M
  if (cpi->oxcf.txfm_cfg.enable_flip_idtx == 0)
1761
0
    ext_tx_used_flag &= DCT_ADST_TX_MASK;
1762
1763
44.3M
  uint16_t allowed_tx_mask = 0;  // 1: allow; 0: skip.
1764
44.3M
  if (txk_allowed < TX_TYPES) {
1765
43.7M
    allowed_tx_mask = 1 << txk_allowed;
1766
43.7M
    allowed_tx_mask &= ext_tx_used_flag;
1767
43.7M
  } else if (fast_tx_search) {
1768
0
    allowed_tx_mask = 0x0c01;  // V_DCT, H_DCT, DCT_DCT
1769
0
    allowed_tx_mask &= ext_tx_used_flag;
1770
587k
  } else {
1771
587k
    assert(plane == 0);
1772
587k
    allowed_tx_mask = ext_tx_used_flag;
1773
587k
    int num_allowed = 0;
1774
587k
    int i;
1775
1776
587k
    if (cpi->sf.tx_sf.tx_type_search.prune_tx_type_using_stats) {
1777
0
      static const int thresh_arr[2][7] = { { 10, 15, 15, 10, 15, 15, 15 },
1778
0
                                            { 10, 17, 17, 10, 17, 17, 17 } };
1779
0
      const int thresh =
1780
0
          thresh_arr[cpi->sf.tx_sf.tx_type_search.prune_tx_type_using_stats - 1]
1781
0
                    [update_type];
1782
0
      uint16_t prune = 0;
1783
0
      int max_prob = -1;
1784
0
      int max_idx = 0;
1785
0
      for (i = 0; i < TX_TYPES; i++) {
1786
0
        if (tx_type_probs[i] > max_prob && (allowed_tx_mask & (1 << i))) {
1787
0
          max_prob = tx_type_probs[i];
1788
0
          max_idx = i;
1789
0
        }
1790
0
        if (tx_type_probs[i] < thresh) prune |= (1 << i);
1791
0
      }
1792
0
      if ((prune >> max_idx) & 0x01) prune &= ~(1 << max_idx);
1793
0
      allowed_tx_mask &= (~prune);
1794
0
    }
1795
5.06M
    for (i = 0; i < TX_TYPES; i++) {
1796
4.48M
      if (allowed_tx_mask & (1 << i)) num_allowed++;
1797
4.48M
    }
1798
587k
    assert(num_allowed > 0);
1799
1800
587k
    if (num_allowed > 2 && cpi->sf.tx_sf.tx_type_search.prune_tx_type_est_rd) {
1801
280k
      int pf = prune_factors[txfm_params->prune_2d_txfm_mode];
1802
280k
      int mf = mul_factors[txfm_params->prune_2d_txfm_mode];
1803
280k
      if (num_allowed <= 7) {
1804
280k
        const uint16_t prune =
1805
280k
            prune_txk_type(cpi, x, plane, block, tx_size, blk_row, blk_col,
1806
280k
                           plane_bsize, txk_map, allowed_tx_mask, pf, txb_ctx,
1807
280k
                           cm->features.reduced_tx_set_used);
1808
280k
        allowed_tx_mask &= (~prune);
1809
280k
      } else {
1810
0
        const int num_sel = (num_allowed * mf + 50) / 100;
1811
0
        const uint16_t prune = prune_txk_type_separ(
1812
0
            cpi, x, plane, block, tx_size, blk_row, blk_col, plane_bsize,
1813
0
            txk_map, allowed_tx_mask, pf, txb_ctx,
1814
0
            cm->features.reduced_tx_set_used, ref_best_rd, num_sel);
1815
1816
0
        allowed_tx_mask &= (~prune);
1817
0
      }
1818
307k
    } else {
1819
307k
      assert(num_allowed > 0);
1820
307k
      int allowed_tx_count =
1821
307k
          (txfm_params->prune_2d_txfm_mode >= TX_TYPE_PRUNE_4) ? 1 : 5;
1822
      // !fast_tx_search && txk_end != txk_start && plane == 0
1823
307k
      if (txfm_params->prune_2d_txfm_mode >= TX_TYPE_PRUNE_1 && is_inter &&
1824
307k
          num_allowed > allowed_tx_count) {
1825
0
        prune_tx_2D(x, plane_bsize, tx_size, blk_row, blk_col, tx_set_type,
1826
0
                    txfm_params->prune_2d_txfm_mode, txk_map, &allowed_tx_mask);
1827
0
      }
1828
307k
    }
1829
587k
  }
1830
1831
  // Need to have at least one transform type allowed.
1832
44.3M
  if (allowed_tx_mask == 0) {
1833
1.02M
    txk_allowed = (plane ? uv_tx_type : DCT_DCT);
1834
1.02M
    allowed_tx_mask = (1 << txk_allowed);
1835
1.02M
  }
1836
1837
44.3M
  assert(IMPLIES(txk_allowed < TX_TYPES, allowed_tx_mask == 1 << txk_allowed));
1838
44.3M
  *allowed_txk_types = txk_allowed;
1839
44.3M
  return allowed_tx_mask;
1840
44.3M
}
1841
1842
#if CONFIG_RD_DEBUG
1843
static INLINE void update_txb_coeff_cost(RD_STATS *rd_stats, int plane,
1844
                                         int txb_coeff_cost) {
1845
  rd_stats->txb_coeff_cost[plane] += txb_coeff_cost;
1846
}
1847
#endif
1848
1849
static INLINE int cost_coeffs(MACROBLOCK *x, int plane, int block,
1850
                              TX_SIZE tx_size, const TX_TYPE tx_type,
1851
                              const TXB_CTX *const txb_ctx,
1852
41.2M
                              int reduced_tx_set_used) {
1853
#if TXCOEFF_COST_TIMER
1854
  struct aom_usec_timer timer;
1855
  aom_usec_timer_start(&timer);
1856
#endif
1857
41.2M
  const int cost = av1_cost_coeffs_txb(x, plane, block, tx_size, tx_type,
1858
41.2M
                                       txb_ctx, reduced_tx_set_used);
1859
#if TXCOEFF_COST_TIMER
1860
  AV1_COMMON *tmp_cm = (AV1_COMMON *)&cpi->common;
1861
  aom_usec_timer_mark(&timer);
1862
  const int64_t elapsed_time = aom_usec_timer_elapsed(&timer);
1863
  tmp_cm->txcoeff_cost_timer += elapsed_time;
1864
  ++tmp_cm->txcoeff_cost_count;
1865
#endif
1866
41.2M
  return cost;
1867
41.2M
}
1868
1869
static int skip_trellis_opt_based_on_satd(MACROBLOCK *x,
1870
                                          QUANT_PARAM *quant_param, int plane,
1871
                                          int block, TX_SIZE tx_size,
1872
                                          int quant_b_adapt, int qstep,
1873
                                          unsigned int coeff_opt_satd_threshold,
1874
44.2M
                                          int skip_trellis, int dc_only_blk) {
1875
44.2M
  if (skip_trellis || (coeff_opt_satd_threshold == UINT_MAX))
1876
41.1M
    return skip_trellis;
1877
1878
3.03M
  const struct macroblock_plane *const p = &x->plane[plane];
1879
3.03M
  const int block_offset = BLOCK_OFFSET(block);
1880
3.03M
  tran_low_t *const coeff_ptr = p->coeff + block_offset;
1881
3.03M
  const int n_coeffs = av1_get_max_eob(tx_size);
1882
3.03M
  const int shift = (MAX_TX_SCALE - av1_get_tx_scale(tx_size));
1883
3.03M
  int satd = (dc_only_blk) ? abs(coeff_ptr[0]) : aom_satd(coeff_ptr, n_coeffs);
1884
3.03M
  satd = RIGHT_SIGNED_SHIFT(satd, shift);
1885
3.03M
  satd >>= (x->e_mbd.bd - 8);
1886
1887
3.03M
  const int skip_block_trellis =
1888
3.03M
      ((uint64_t)satd >
1889
3.03M
       (uint64_t)coeff_opt_satd_threshold * qstep * sqrt_tx_pixels_2d[tx_size]);
1890
1891
3.03M
  av1_setup_quant(
1892
3.03M
      tx_size, !skip_block_trellis,
1893
3.03M
      skip_block_trellis
1894
18.4E
          ? (USE_B_QUANT_NO_TRELLIS ? AV1_XFORM_QUANT_B : AV1_XFORM_QUANT_FP)
1895
3.03M
          : AV1_XFORM_QUANT_FP,
1896
3.03M
      quant_b_adapt, quant_param);
1897
1898
3.03M
  return skip_block_trellis;
1899
44.2M
}
1900
1901
// Predict DC only blocks if the residual variance is below a qstep based
1902
// threshold.For such blocks, transform type search is bypassed.
1903
static INLINE void predict_dc_only_block(
1904
    MACROBLOCK *x, int plane, BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
1905
    int block, int blk_row, int blk_col, RD_STATS *best_rd_stats,
1906
    int64_t *block_sse, unsigned int *block_mse_q8, int64_t *per_px_mean,
1907
0
    int *dc_only_blk) {
1908
0
  MACROBLOCKD *xd = &x->e_mbd;
1909
0
  MB_MODE_INFO *mbmi = xd->mi[0];
1910
0
  const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3;
1911
0
  const int qstep = x->plane[plane].dequant_QTX[1] >> dequant_shift;
1912
0
  uint64_t block_var = UINT64_MAX;
1913
0
  const int dc_qstep = x->plane[plane].dequant_QTX[0] >> 3;
1914
0
  *block_sse = pixel_diff_stats(x, plane, blk_row, blk_col, plane_bsize,
1915
0
                                txsize_to_bsize[tx_size], block_mse_q8,
1916
0
                                per_px_mean, &block_var);
1917
0
  assert((*block_mse_q8) != UINT_MAX);
1918
0
  uint64_t var_threshold = (uint64_t)(1.8 * qstep * qstep);
1919
0
  if (is_cur_buf_hbd(xd))
1920
0
    block_var = ROUND_POWER_OF_TWO(block_var, (xd->bd - 8) * 2);
1921
  // Early prediction of skip block if residual mean and variance are less
1922
  // than qstep based threshold
1923
0
  if (((llabs(*per_px_mean) * dc_coeff_scale[tx_size]) < (dc_qstep << 12)) &&
1924
0
      (block_var < var_threshold)) {
1925
    // If the normalized mean of residual block is less than the dc qstep and
1926
    // the  normalized block variance is less than ac qstep, then the block is
1927
    // assumed to be a skip block and its rdcost is updated accordingly.
1928
0
    best_rd_stats->skip_txfm = 1;
1929
1930
0
    x->plane[plane].eobs[block] = 0;
1931
1932
0
    if (is_cur_buf_hbd(xd))
1933
0
      *block_sse = ROUND_POWER_OF_TWO((*block_sse), (xd->bd - 8) * 2);
1934
1935
0
    best_rd_stats->dist = (*block_sse) << 4;
1936
0
    best_rd_stats->sse = best_rd_stats->dist;
1937
1938
0
    ENTROPY_CONTEXT ctxa[MAX_MIB_SIZE];
1939
0
    ENTROPY_CONTEXT ctxl[MAX_MIB_SIZE];
1940
0
    av1_get_entropy_contexts(plane_bsize, &xd->plane[plane], ctxa, ctxl);
1941
0
    ENTROPY_CONTEXT *ta = ctxa;
1942
0
    ENTROPY_CONTEXT *tl = ctxl;
1943
0
    const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
1944
0
    TXB_CTX txb_ctx_tmp;
1945
0
    const PLANE_TYPE plane_type = get_plane_type(plane);
1946
0
    get_txb_ctx(plane_bsize, tx_size, plane, ta, tl, &txb_ctx_tmp);
1947
0
    const int zero_blk_rate = x->coeff_costs.coeff_costs[txs_ctx][plane_type]
1948
0
                                  .txb_skip_cost[txb_ctx_tmp.txb_skip_ctx][1];
1949
0
    best_rd_stats->rate = zero_blk_rate;
1950
1951
0
    best_rd_stats->rdcost =
1952
0
        RDCOST(x->rdmult, best_rd_stats->rate, best_rd_stats->sse);
1953
1954
0
    x->plane[plane].txb_entropy_ctx[block] = 0;
1955
0
  } else if (block_var < var_threshold) {
1956
    // Predict DC only blocks based on residual variance.
1957
    // For chroma plane, this early prediction is disabled for intra blocks.
1958
0
    if ((plane == 0) || (plane > 0 && is_inter_block(mbmi))) *dc_only_blk = 1;
1959
0
  }
1960
0
}
1961
1962
// Search for the best transform type for a given transform block.
1963
// This function can be used for both inter and intra, both luma and chroma.
1964
static void search_tx_type(const AV1_COMP *cpi, MACROBLOCK *x, int plane,
1965
                           int block, int blk_row, int blk_col,
1966
                           BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
1967
                           const TXB_CTX *const txb_ctx,
1968
                           FAST_TX_SEARCH_MODE ftxs_mode, int skip_trellis,
1969
44.5M
                           int64_t ref_best_rd, RD_STATS *best_rd_stats) {
1970
44.5M
  const AV1_COMMON *cm = &cpi->common;
1971
44.5M
  MACROBLOCKD *xd = &x->e_mbd;
1972
44.5M
  MB_MODE_INFO *mbmi = xd->mi[0];
1973
44.5M
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
1974
44.5M
  int64_t best_rd = INT64_MAX;
1975
44.5M
  uint16_t best_eob = 0;
1976
44.5M
  TX_TYPE best_tx_type = DCT_DCT;
1977
44.5M
  int rate_cost = 0;
1978
  // The buffer used to swap dqcoeff in macroblockd_plane so we can keep dqcoeff
1979
  // of the best tx_type
1980
44.5M
  DECLARE_ALIGNED(32, tran_low_t, this_dqcoeff[MAX_SB_SQUARE]);
1981
44.5M
  struct macroblock_plane *const p = &x->plane[plane];
1982
44.5M
  tran_low_t *orig_dqcoeff = p->dqcoeff;
1983
44.5M
  tran_low_t *best_dqcoeff = this_dqcoeff;
1984
44.5M
  const int tx_type_map_idx =
1985
44.5M
      plane ? 0 : blk_row * xd->tx_type_map_stride + blk_col;
1986
44.5M
  av1_invalid_rd_stats(best_rd_stats);
1987
1988
44.5M
  skip_trellis |= !is_trellis_used(cpi->optimize_seg_arr[xd->mi[0]->segment_id],
1989
44.5M
                                   DRY_RUN_NORMAL);
1990
1991
44.5M
  uint8_t best_txb_ctx = 0;
1992
  // txk_allowed = TX_TYPES: >1 tx types are allowed
1993
  // txk_allowed < TX_TYPES: only that specific tx type is allowed.
1994
44.5M
  TX_TYPE txk_allowed = TX_TYPES;
1995
44.5M
  int txk_map[TX_TYPES] = {
1996
44.5M
    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15
1997
44.5M
  };
1998
44.5M
  const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3;
1999
44.5M
  const int qstep = x->plane[plane].dequant_QTX[1] >> dequant_shift;
2000
2001
44.5M
  const uint8_t txw = tx_size_wide[tx_size];
2002
44.5M
  const uint8_t txh = tx_size_high[tx_size];
2003
44.5M
  int64_t block_sse;
2004
44.5M
  unsigned int block_mse_q8;
2005
44.5M
  int dc_only_blk = 0;
2006
44.5M
  const bool predict_dc_block =
2007
44.5M
      txfm_params->predict_dc_level && txw != 64 && txh != 64;
2008
44.5M
  int64_t per_px_mean = INT64_MAX;
2009
44.5M
  if (predict_dc_block) {
2010
0
    predict_dc_only_block(x, plane, plane_bsize, tx_size, block, blk_row,
2011
0
                          blk_col, best_rd_stats, &block_sse, &block_mse_q8,
2012
0
                          &per_px_mean, &dc_only_blk);
2013
0
    if (best_rd_stats->skip_txfm == 1) {
2014
0
      const TX_TYPE tx_type = DCT_DCT;
2015
0
      if (plane == 0) xd->tx_type_map[tx_type_map_idx] = tx_type;
2016
0
      return;
2017
0
    }
2018
44.5M
  } else {
2019
44.5M
    block_sse = pixel_diff_dist(x, plane, blk_row, blk_col, plane_bsize,
2020
44.5M
                                txsize_to_bsize[tx_size], &block_mse_q8);
2021
44.5M
    assert(block_mse_q8 != UINT_MAX);
2022
44.5M
  }
2023
2024
  // Bit mask to indicate which transform types are allowed in the RD search.
2025
44.5M
  uint16_t tx_mask;
2026
2027
  // Use DCT_DCT transform for DC only block.
2028
44.5M
  if (dc_only_blk)
2029
0
    tx_mask = 1 << DCT_DCT;
2030
44.5M
  else
2031
44.5M
    tx_mask = get_tx_mask(cpi, x, plane, block, blk_row, blk_col, plane_bsize,
2032
44.5M
                          tx_size, txb_ctx, ftxs_mode, ref_best_rd,
2033
44.5M
                          &txk_allowed, txk_map);
2034
44.5M
  const uint16_t allowed_tx_mask = tx_mask;
2035
2036
44.5M
  if (is_cur_buf_hbd(xd)) {
2037
0
    block_sse = ROUND_POWER_OF_TWO(block_sse, (xd->bd - 8) * 2);
2038
0
    block_mse_q8 = ROUND_POWER_OF_TWO(block_mse_q8, (xd->bd - 8) * 2);
2039
0
  }
2040
44.5M
  block_sse *= 16;
2041
  // Use mse / qstep^2 based threshold logic to take decision of R-D
2042
  // optimization of coeffs. For smaller residuals, coeff optimization
2043
  // would be helpful. For larger residuals, R-D optimization may not be
2044
  // effective.
2045
  // TODO(any): Experiment with variance and mean based thresholds
2046
44.5M
  const int perform_block_coeff_opt =
2047
44.5M
      ((uint64_t)block_mse_q8 <=
2048
44.5M
       (uint64_t)txfm_params->coeff_opt_thresholds[0] * qstep * qstep);
2049
44.5M
  skip_trellis |= !perform_block_coeff_opt;
2050
2051
  // Flag to indicate if distortion should be calculated in transform domain or
2052
  // not during iterating through transform type candidates.
2053
  // Transform domain distortion is accurate for higher residuals.
2054
  // TODO(any): Experiment with variance and mean based thresholds
2055
44.5M
  int use_transform_domain_distortion =
2056
44.5M
      (txfm_params->use_transform_domain_distortion > 0) &&
2057
44.5M
      (block_mse_q8 >= txfm_params->tx_domain_dist_threshold) &&
2058
      // Any 64-pt transforms only preserves half the coefficients.
2059
      // Therefore transform domain distortion is not valid for these
2060
      // transform sizes.
2061
44.5M
      (txsize_sqr_up_map[tx_size] != TX_64X64) &&
2062
      // Use pixel domain distortion for DC only blocks
2063
44.5M
      !dc_only_blk;
2064
  // Flag to indicate if an extra calculation of distortion in the pixel domain
2065
  // should be performed at the end, after the best transform type has been
2066
  // decided.
2067
44.5M
  int calc_pixel_domain_distortion_final =
2068
44.5M
      txfm_params->use_transform_domain_distortion == 1 &&
2069
44.5M
      use_transform_domain_distortion && x->rd_model != LOW_TXFM_RD;
2070
44.5M
  if (calc_pixel_domain_distortion_final &&
2071
44.5M
      (txk_allowed < TX_TYPES || allowed_tx_mask == 0x0001))
2072
1.34M
    calc_pixel_domain_distortion_final = use_transform_domain_distortion = 0;
2073
2074
44.5M
  const uint16_t *eobs_ptr = x->plane[plane].eobs;
2075
2076
44.5M
  TxfmParam txfm_param;
2077
44.5M
  QUANT_PARAM quant_param;
2078
44.5M
  int skip_trellis_based_on_satd[TX_TYPES] = { 0 };
2079
44.5M
  av1_setup_xform(cm, x, tx_size, DCT_DCT, &txfm_param);
2080
44.5M
  av1_setup_quant(tx_size, !skip_trellis,
2081
44.5M
                  skip_trellis ? (USE_B_QUANT_NO_TRELLIS ? AV1_XFORM_QUANT_B
2082
18.4E
                                                         : AV1_XFORM_QUANT_FP)
2083
44.5M
                               : AV1_XFORM_QUANT_FP,
2084
44.5M
                  cpi->oxcf.q_cfg.quant_b_adapt, &quant_param);
2085
2086
  // Iterate through all transform type candidates.
2087
50.4M
  for (int idx = 0; idx < TX_TYPES; ++idx) {
2088
50.4M
    const TX_TYPE tx_type = (TX_TYPE)txk_map[idx];
2089
50.4M
    if (tx_type == TX_TYPE_INVALID || !check_bit_mask(allowed_tx_mask, tx_type))
2090
6.40M
      continue;
2091
43.9M
    txfm_param.tx_type = tx_type;
2092
43.9M
    if (av1_use_qmatrix(&cm->quant_params, xd, mbmi->segment_id)) {
2093
0
      av1_setup_qmatrix(&cm->quant_params, xd, plane, tx_size, tx_type,
2094
0
                        &quant_param);
2095
0
    }
2096
43.9M
    if (plane == 0) xd->tx_type_map[tx_type_map_idx] = tx_type;
2097
43.9M
    RD_STATS this_rd_stats;
2098
43.9M
    av1_invalid_rd_stats(&this_rd_stats);
2099
2100
43.9M
    if (!dc_only_blk)
2101
43.9M
      av1_xform(x, plane, block, blk_row, blk_col, plane_bsize, &txfm_param);
2102
90.9k
    else
2103
90.9k
      av1_xform_dc_only(x, plane, block, &txfm_param, per_px_mean);
2104
2105
43.9M
    skip_trellis_based_on_satd[tx_type] = skip_trellis_opt_based_on_satd(
2106
43.9M
        x, &quant_param, plane, block, tx_size, cpi->oxcf.q_cfg.quant_b_adapt,
2107
43.9M
        qstep, txfm_params->coeff_opt_thresholds[1], skip_trellis, dc_only_blk);
2108
2109
43.9M
    av1_quant(x, plane, block, &txfm_param, &quant_param);
2110
2111
    // Calculate rate cost of quantized coefficients.
2112
43.9M
    if (quant_param.use_optimize_b) {
2113
3.41M
      av1_optimize_b(cpi, x, plane, block, tx_size, tx_type, txb_ctx,
2114
3.41M
                     &rate_cost);
2115
40.5M
    } else {
2116
40.5M
      rate_cost = cost_coeffs(x, plane, block, tx_size, tx_type, txb_ctx,
2117
40.5M
                              cm->features.reduced_tx_set_used);
2118
40.5M
    }
2119
2120
    // If rd cost based on coeff rate alone is already more than best_rd,
2121
    // terminate early.
2122
43.9M
    if (RDCOST(x->rdmult, rate_cost, 0) > best_rd) continue;
2123
2124
    // Calculate distortion.
2125
44.2M
    if (eobs_ptr[block] == 0) {
2126
      // When eob is 0, pixel domain distortion is more efficient and accurate.
2127
44.2M
      this_rd_stats.dist = this_rd_stats.sse = block_sse;
2128
18.4E
    } else if (dc_only_blk) {
2129
0
      this_rd_stats.sse = block_sse;
2130
0
      this_rd_stats.dist = dist_block_px_domain(
2131
0
          cpi, x, plane, plane_bsize, block, blk_row, blk_col, tx_size);
2132
18.4E
    } else if (use_transform_domain_distortion) {
2133
160k
      dist_block_tx_domain(x, plane, block, tx_size, &this_rd_stats.dist,
2134
160k
                           &this_rd_stats.sse);
2135
18.4E
    } else {
2136
18.4E
      int64_t sse_diff = INT64_MAX;
2137
      // high_energy threshold assumes that every pixel within a txfm block
2138
      // has a residue energy of at least 25% of the maximum, i.e. 128 * 128
2139
      // for 8 bit.
2140
18.4E
      const int64_t high_energy_thresh =
2141
18.4E
          ((int64_t)128 * 128 * tx_size_2d[tx_size]);
2142
18.4E
      const int is_high_energy = (block_sse >= high_energy_thresh);
2143
18.4E
      if (tx_size == TX_64X64 || is_high_energy) {
2144
        // Because 3 out 4 quadrants of transform coefficients are forced to
2145
        // zero, the inverse transform has a tendency to overflow. sse_diff
2146
        // is effectively the energy of those 3 quadrants, here we use it
2147
        // to decide if we should do pixel domain distortion. If the energy
2148
        // is mostly in first quadrant, then it is unlikely that we have
2149
        // overflow issue in inverse transform.
2150
148k
        dist_block_tx_domain(x, plane, block, tx_size, &this_rd_stats.dist,
2151
148k
                             &this_rd_stats.sse);
2152
148k
        sse_diff = block_sse - this_rd_stats.sse;
2153
148k
      }
2154
18.4E
      if (tx_size != TX_64X64 || !is_high_energy ||
2155
18.4E
          (sse_diff * 2) < this_rd_stats.sse) {
2156
166k
        const int64_t tx_domain_dist = this_rd_stats.dist;
2157
166k
        this_rd_stats.dist = dist_block_px_domain(
2158
166k
            cpi, x, plane, plane_bsize, block, blk_row, blk_col, tx_size);
2159
        // For high energy blocks, occasionally, the pixel domain distortion
2160
        // can be artificially low due to clamping at reconstruction stage
2161
        // even when inverse transform output is hugely different from the
2162
        // actual residue.
2163
166k
        if (is_high_energy && this_rd_stats.dist < tx_domain_dist)
2164
86.2k
          this_rd_stats.dist = tx_domain_dist;
2165
18.4E
      } else {
2166
18.4E
        assert(sse_diff < INT64_MAX);
2167
18.4E
        this_rd_stats.dist += sse_diff;
2168
18.4E
      }
2169
18.4E
      this_rd_stats.sse = block_sse;
2170
18.4E
    }
2171
2172
43.9M
    this_rd_stats.rate = rate_cost;
2173
2174
43.9M
    const int64_t rd =
2175
43.9M
        RDCOST(x->rdmult, this_rd_stats.rate, this_rd_stats.dist);
2176
2177
44.5M
    if (rd < best_rd) {
2178
44.5M
      best_rd = rd;
2179
44.5M
      *best_rd_stats = this_rd_stats;
2180
44.5M
      best_tx_type = tx_type;
2181
44.5M
      best_txb_ctx = x->plane[plane].txb_entropy_ctx[block];
2182
44.5M
      best_eob = x->plane[plane].eobs[block];
2183
      // Swap dqcoeff buffers
2184
44.5M
      tran_low_t *const tmp_dqcoeff = best_dqcoeff;
2185
44.5M
      best_dqcoeff = p->dqcoeff;
2186
44.5M
      p->dqcoeff = tmp_dqcoeff;
2187
44.5M
    }
2188
2189
#if CONFIG_COLLECT_RD_STATS == 1
2190
    if (plane == 0) {
2191
      PrintTransformUnitStats(cpi, x, &this_rd_stats, blk_row, blk_col,
2192
                              plane_bsize, tx_size, tx_type, rd);
2193
    }
2194
#endif  // CONFIG_COLLECT_RD_STATS == 1
2195
2196
#if COLLECT_TX_SIZE_DATA
2197
    // Generate small sample to restrict output size.
2198
    static unsigned int seed = 21743;
2199
    if (lcg_rand16(&seed) % 200 == 0) {
2200
      FILE *fp = NULL;
2201
2202
      if (within_border) {
2203
        fp = fopen(av1_tx_size_data_output_file, "a");
2204
      }
2205
2206
      if (fp) {
2207
        // Transform info and RD
2208
        const int txb_w = tx_size_wide[tx_size];
2209
        const int txb_h = tx_size_high[tx_size];
2210
2211
        // Residue signal.
2212
        const int diff_stride = block_size_wide[plane_bsize];
2213
        struct macroblock_plane *const p = &x->plane[plane];
2214
        const int16_t *src_diff =
2215
            &p->src_diff[(blk_row * diff_stride + blk_col) * 4];
2216
2217
        for (int r = 0; r < txb_h; ++r) {
2218
          for (int c = 0; c < txb_w; ++c) {
2219
            fprintf(fp, "%d,", src_diff[c]);
2220
          }
2221
          src_diff += diff_stride;
2222
        }
2223
2224
        fprintf(fp, "%d,%d,%d,%" PRId64, txb_w, txb_h, tx_type, rd);
2225
        fprintf(fp, "\n");
2226
        fclose(fp);
2227
      }
2228
    }
2229
#endif  // COLLECT_TX_SIZE_DATA
2230
2231
    // If the current best RD cost is much worse than the reference RD cost,
2232
    // terminate early.
2233
44.9M
    if (cpi->sf.tx_sf.adaptive_txb_search_level) {
2234
44.9M
      if ((best_rd - (best_rd >> cpi->sf.tx_sf.adaptive_txb_search_level)) >
2235
44.9M
          ref_best_rd) {
2236
580k
        break;
2237
580k
      }
2238
44.9M
    }
2239
2240
    // Terminate transform type search if the block has been quantized to
2241
    // all zero.
2242
44.3M
    if (cpi->sf.tx_sf.tx_type_search.skip_tx_search && !best_eob) break;
2243
43.4M
  }
2244
2245
44.5M
  assert(best_rd != INT64_MAX);
2246
2247
44.5M
  best_rd_stats->skip_txfm = best_eob == 0;
2248
44.5M
  if (plane == 0) update_txk_array(xd, blk_row, blk_col, tx_size, best_tx_type);
2249
44.5M
  x->plane[plane].txb_entropy_ctx[block] = best_txb_ctx;
2250
44.5M
  x->plane[plane].eobs[block] = best_eob;
2251
44.5M
  skip_trellis = skip_trellis_based_on_satd[best_tx_type];
2252
2253
  // Point dqcoeff to the quantized coefficients corresponding to the best
2254
  // transform type, then we can skip transform and quantization, e.g. in the
2255
  // final pixel domain distortion calculation and recon_intra().
2256
44.5M
  p->dqcoeff = best_dqcoeff;
2257
2258
44.5M
  if (calc_pixel_domain_distortion_final && best_eob) {
2259
0
    best_rd_stats->dist = dist_block_px_domain(
2260
0
        cpi, x, plane, plane_bsize, block, blk_row, blk_col, tx_size);
2261
0
    best_rd_stats->sse = block_sse;
2262
0
  }
2263
2264
  // Intra mode needs decoded pixels such that the next transform block
2265
  // can use them for prediction.
2266
44.5M
  recon_intra(cpi, x, plane, block, blk_row, blk_col, plane_bsize, tx_size,
2267
44.5M
              txb_ctx, skip_trellis, best_tx_type, 0, &rate_cost, best_eob);
2268
44.5M
  p->dqcoeff = orig_dqcoeff;
2269
44.5M
}
2270
2271
// Pick transform type for a luma transform block of tx_size. Note this function
2272
// is used only for inter-predicted blocks.
2273
static AOM_INLINE void tx_type_rd(const AV1_COMP *cpi, MACROBLOCK *x,
2274
                                  TX_SIZE tx_size, int blk_row, int blk_col,
2275
                                  int block, int plane_bsize, TXB_CTX *txb_ctx,
2276
                                  RD_STATS *rd_stats,
2277
                                  FAST_TX_SEARCH_MODE ftxs_mode,
2278
0
                                  int64_t ref_rdcost) {
2279
0
  assert(is_inter_block(x->e_mbd.mi[0]));
2280
0
  RD_STATS this_rd_stats;
2281
0
  const int skip_trellis = 0;
2282
0
  search_tx_type(cpi, x, 0, block, blk_row, blk_col, plane_bsize, tx_size,
2283
0
                 txb_ctx, ftxs_mode, skip_trellis, ref_rdcost, &this_rd_stats);
2284
2285
0
  av1_merge_rd_stats(rd_stats, &this_rd_stats);
2286
0
}
2287
2288
static AOM_INLINE void try_tx_block_no_split(
2289
    const AV1_COMP *cpi, MACROBLOCK *x, int blk_row, int blk_col, int block,
2290
    TX_SIZE tx_size, int depth, BLOCK_SIZE plane_bsize,
2291
    const ENTROPY_CONTEXT *ta, const ENTROPY_CONTEXT *tl,
2292
    int txfm_partition_ctx, RD_STATS *rd_stats, int64_t ref_best_rd,
2293
0
    FAST_TX_SEARCH_MODE ftxs_mode, TxCandidateInfo *no_split) {
2294
0
  MACROBLOCKD *const xd = &x->e_mbd;
2295
0
  MB_MODE_INFO *const mbmi = xd->mi[0];
2296
0
  struct macroblock_plane *const p = &x->plane[0];
2297
0
  const int bw = mi_size_wide[plane_bsize];
2298
0
  const ENTROPY_CONTEXT *const pta = ta + blk_col;
2299
0
  const ENTROPY_CONTEXT *const ptl = tl + blk_row;
2300
0
  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
2301
0
  TXB_CTX txb_ctx;
2302
0
  get_txb_ctx(plane_bsize, tx_size, 0, pta, ptl, &txb_ctx);
2303
0
  const int zero_blk_rate = x->coeff_costs.coeff_costs[txs_ctx][PLANE_TYPE_Y]
2304
0
                                .txb_skip_cost[txb_ctx.txb_skip_ctx][1];
2305
0
  rd_stats->zero_rate = zero_blk_rate;
2306
0
  const int index = av1_get_txb_size_index(plane_bsize, blk_row, blk_col);
2307
0
  mbmi->inter_tx_size[index] = tx_size;
2308
0
  tx_type_rd(cpi, x, tx_size, blk_row, blk_col, block, plane_bsize, &txb_ctx,
2309
0
             rd_stats, ftxs_mode, ref_best_rd);
2310
0
  assert(rd_stats->rate < INT_MAX);
2311
2312
0
  const int pick_skip_txfm =
2313
0
      !xd->lossless[mbmi->segment_id] &&
2314
0
      (rd_stats->skip_txfm == 1 ||
2315
0
       RDCOST(x->rdmult, rd_stats->rate, rd_stats->dist) >=
2316
0
           RDCOST(x->rdmult, zero_blk_rate, rd_stats->sse));
2317
0
  if (pick_skip_txfm) {
2318
#if CONFIG_RD_DEBUG
2319
    update_txb_coeff_cost(rd_stats, 0, zero_blk_rate - rd_stats->rate);
2320
#endif  // CONFIG_RD_DEBUG
2321
0
    rd_stats->rate = zero_blk_rate;
2322
0
    rd_stats->dist = rd_stats->sse;
2323
0
    p->eobs[block] = 0;
2324
0
    update_txk_array(xd, blk_row, blk_col, tx_size, DCT_DCT);
2325
0
  }
2326
0
  rd_stats->skip_txfm = pick_skip_txfm;
2327
0
  set_blk_skip(x->txfm_search_info.blk_skip, 0, blk_row * bw + blk_col,
2328
0
               pick_skip_txfm);
2329
2330
0
  if (tx_size > TX_4X4 && depth < MAX_VARTX_DEPTH)
2331
0
    rd_stats->rate += x->mode_costs.txfm_partition_cost[txfm_partition_ctx][0];
2332
2333
0
  no_split->rd = RDCOST(x->rdmult, rd_stats->rate, rd_stats->dist);
2334
0
  no_split->txb_entropy_ctx = p->txb_entropy_ctx[block];
2335
0
  no_split->tx_type =
2336
0
      xd->tx_type_map[blk_row * xd->tx_type_map_stride + blk_col];
2337
0
}
2338
2339
static AOM_INLINE void try_tx_block_split(
2340
    const AV1_COMP *cpi, MACROBLOCK *x, int blk_row, int blk_col, int block,
2341
    TX_SIZE tx_size, int depth, BLOCK_SIZE plane_bsize, ENTROPY_CONTEXT *ta,
2342
    ENTROPY_CONTEXT *tl, TXFM_CONTEXT *tx_above, TXFM_CONTEXT *tx_left,
2343
    int txfm_partition_ctx, int64_t no_split_rd, int64_t ref_best_rd,
2344
0
    FAST_TX_SEARCH_MODE ftxs_mode, RD_STATS *split_rd_stats) {
2345
0
  assert(tx_size < TX_SIZES_ALL);
2346
0
  MACROBLOCKD *const xd = &x->e_mbd;
2347
0
  const int max_blocks_high = max_block_high(xd, plane_bsize, 0);
2348
0
  const int max_blocks_wide = max_block_wide(xd, plane_bsize, 0);
2349
0
  const int txb_width = tx_size_wide_unit[tx_size];
2350
0
  const int txb_height = tx_size_high_unit[tx_size];
2351
  // Transform size after splitting current block.
2352
0
  const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
2353
0
  const int sub_txb_width = tx_size_wide_unit[sub_txs];
2354
0
  const int sub_txb_height = tx_size_high_unit[sub_txs];
2355
0
  const int sub_step = sub_txb_width * sub_txb_height;
2356
0
  const int nblks = (txb_height / sub_txb_height) * (txb_width / sub_txb_width);
2357
0
  assert(nblks > 0);
2358
0
  av1_init_rd_stats(split_rd_stats);
2359
0
  split_rd_stats->rate =
2360
0
      x->mode_costs.txfm_partition_cost[txfm_partition_ctx][1];
2361
2362
0
  for (int r = 0, blk_idx = 0; r < txb_height; r += sub_txb_height) {
2363
0
    const int offsetr = blk_row + r;
2364
0
    if (offsetr >= max_blocks_high) break;
2365
0
    for (int c = 0; c < txb_width; c += sub_txb_width, ++blk_idx) {
2366
0
      assert(blk_idx < 4);
2367
0
      const int offsetc = blk_col + c;
2368
0
      if (offsetc >= max_blocks_wide) continue;
2369
2370
0
      RD_STATS this_rd_stats;
2371
0
      int this_cost_valid = 1;
2372
0
      select_tx_block(cpi, x, offsetr, offsetc, block, sub_txs, depth + 1,
2373
0
                      plane_bsize, ta, tl, tx_above, tx_left, &this_rd_stats,
2374
0
                      no_split_rd / nblks, ref_best_rd - split_rd_stats->rdcost,
2375
0
                      &this_cost_valid, ftxs_mode);
2376
0
      if (!this_cost_valid) {
2377
0
        split_rd_stats->rdcost = INT64_MAX;
2378
0
        return;
2379
0
      }
2380
0
      av1_merge_rd_stats(split_rd_stats, &this_rd_stats);
2381
0
      split_rd_stats->rdcost =
2382
0
          RDCOST(x->rdmult, split_rd_stats->rate, split_rd_stats->dist);
2383
0
      if (split_rd_stats->rdcost > ref_best_rd) {
2384
0
        split_rd_stats->rdcost = INT64_MAX;
2385
0
        return;
2386
0
      }
2387
0
      block += sub_step;
2388
0
    }
2389
0
  }
2390
0
}
2391
2392
0
static float get_var(float mean, double x2_sum, int num) {
2393
0
  const float e_x2 = (float)(x2_sum / num);
2394
0
  const float diff = e_x2 - mean * mean;
2395
0
  return diff;
2396
0
}
2397
2398
static AOM_INLINE void get_blk_var_dev(const int16_t *data, int stride, int bw,
2399
                                       int bh, float *dev_of_mean,
2400
0
                                       float *var_of_vars) {
2401
0
  const int16_t *const data_ptr = &data[0];
2402
0
  const int subh = (bh >= bw) ? (bh >> 1) : bh;
2403
0
  const int subw = (bw >= bh) ? (bw >> 1) : bw;
2404
0
  const int num = bw * bh;
2405
0
  const int sub_num = subw * subh;
2406
0
  int total_x_sum = 0;
2407
0
  int64_t total_x2_sum = 0;
2408
0
  int blk_idx = 0;
2409
0
  float var_sum = 0.0f;
2410
0
  float mean_sum = 0.0f;
2411
0
  double var2_sum = 0.0f;
2412
0
  double mean2_sum = 0.0f;
2413
2414
0
  for (int row = 0; row < bh; row += subh) {
2415
0
    for (int col = 0; col < bw; col += subw) {
2416
0
      int x_sum;
2417
0
      int64_t x2_sum;
2418
0
      aom_get_blk_sse_sum(data_ptr + row * stride + col, stride, subw, subh,
2419
0
                          &x_sum, &x2_sum);
2420
0
      total_x_sum += x_sum;
2421
0
      total_x2_sum += x2_sum;
2422
2423
0
      const float mean = (float)x_sum / sub_num;
2424
0
      const float var = get_var(mean, (double)x2_sum, sub_num);
2425
0
      mean_sum += mean;
2426
0
      mean2_sum += (double)(mean * mean);
2427
0
      var_sum += var;
2428
0
      var2_sum += var * var;
2429
0
      blk_idx++;
2430
0
    }
2431
0
  }
2432
2433
0
  const float lvl0_mean = (float)total_x_sum / num;
2434
0
  const float block_var = get_var(lvl0_mean, (double)total_x2_sum, num);
2435
0
  mean_sum += lvl0_mean;
2436
0
  mean2_sum += (double)(lvl0_mean * lvl0_mean);
2437
0
  var_sum += block_var;
2438
0
  var2_sum += block_var * block_var;
2439
0
  const float av_mean = mean_sum / 5;
2440
2441
0
  if (blk_idx > 1) {
2442
    // Deviation of means.
2443
0
    *dev_of_mean = get_dev(av_mean, mean2_sum, (blk_idx + 1));
2444
    // Variance of variances.
2445
0
    const float mean_var = var_sum / (blk_idx + 1);
2446
0
    *var_of_vars = get_var(mean_var, var2_sum, (blk_idx + 1));
2447
0
  }
2448
0
}
2449
2450
static void prune_tx_split_no_split(MACROBLOCK *x, BLOCK_SIZE bsize,
2451
                                    int blk_row, int blk_col, TX_SIZE tx_size,
2452
                                    int *try_no_split, int *try_split,
2453
0
                                    int pruning_level) {
2454
0
  const int diff_stride = block_size_wide[bsize];
2455
0
  const int16_t *diff =
2456
0
      x->plane[0].src_diff + 4 * blk_row * diff_stride + 4 * blk_col;
2457
0
  const int bw = tx_size_wide[tx_size];
2458
0
  const int bh = tx_size_high[tx_size];
2459
0
  float dev_of_means = 0.0f;
2460
0
  float var_of_vars = 0.0f;
2461
2462
  // This function calculates the deviation of means, and the variance of pixel
2463
  // variances of the block as well as it's sub-blocks.
2464
0
  get_blk_var_dev(diff, diff_stride, bw, bh, &dev_of_means, &var_of_vars);
2465
0
  const int dc_q = x->plane[0].dequant_QTX[0] >> 3;
2466
0
  const int ac_q = x->plane[0].dequant_QTX[1] >> 3;
2467
0
  const int no_split_thresh_scales[4] = { 0, 24, 8, 8 };
2468
0
  const int no_split_thresh_scale = no_split_thresh_scales[pruning_level];
2469
0
  const int split_thresh_scales[4] = { 0, 24, 10, 8 };
2470
0
  const int split_thresh_scale = split_thresh_scales[pruning_level];
2471
2472
0
  if ((dev_of_means <= dc_q) &&
2473
0
      (split_thresh_scale * var_of_vars <= ac_q * ac_q)) {
2474
0
    *try_split = 0;
2475
0
  }
2476
0
  if ((dev_of_means > no_split_thresh_scale * dc_q) &&
2477
0
      (var_of_vars > no_split_thresh_scale * ac_q * ac_q)) {
2478
0
    *try_no_split = 0;
2479
0
  }
2480
0
}
2481
2482
// Search for the best transform partition(recursive)/type for a given
2483
// inter-predicted luma block. The obtained transform selection will be saved
2484
// in xd->mi[0], the corresponding RD stats will be saved in rd_stats.
2485
static AOM_INLINE void select_tx_block(
2486
    const AV1_COMP *cpi, MACROBLOCK *x, int blk_row, int blk_col, int block,
2487
    TX_SIZE tx_size, int depth, BLOCK_SIZE plane_bsize, ENTROPY_CONTEXT *ta,
2488
    ENTROPY_CONTEXT *tl, TXFM_CONTEXT *tx_above, TXFM_CONTEXT *tx_left,
2489
    RD_STATS *rd_stats, int64_t prev_level_rd, int64_t ref_best_rd,
2490
0
    int *is_cost_valid, FAST_TX_SEARCH_MODE ftxs_mode) {
2491
0
  assert(tx_size < TX_SIZES_ALL);
2492
0
  av1_init_rd_stats(rd_stats);
2493
0
  if (ref_best_rd < 0) {
2494
0
    *is_cost_valid = 0;
2495
0
    return;
2496
0
  }
2497
2498
0
  MACROBLOCKD *const xd = &x->e_mbd;
2499
0
  assert(blk_row < max_block_high(xd, plane_bsize, 0) &&
2500
0
         blk_col < max_block_wide(xd, plane_bsize, 0));
2501
0
  MB_MODE_INFO *const mbmi = xd->mi[0];
2502
0
  const int ctx = txfm_partition_context(tx_above + blk_col, tx_left + blk_row,
2503
0
                                         mbmi->bsize, tx_size);
2504
0
  struct macroblock_plane *const p = &x->plane[0];
2505
2506
0
  int try_no_split = (cpi->oxcf.txfm_cfg.enable_tx64 ||
2507
0
                      txsize_sqr_up_map[tx_size] != TX_64X64) &&
2508
0
                     (cpi->oxcf.txfm_cfg.enable_rect_tx ||
2509
0
                      tx_size_wide[tx_size] == tx_size_high[tx_size]);
2510
0
  int try_split = tx_size > TX_4X4 && depth < MAX_VARTX_DEPTH;
2511
0
  TxCandidateInfo no_split = { INT64_MAX, 0, TX_TYPES };
2512
2513
  // Prune tx_split and no-split based on sub-block properties.
2514
0
  if (tx_size != TX_4X4 && try_split == 1 && try_no_split == 1 &&
2515
0
      cpi->sf.tx_sf.prune_tx_size_level > 0) {
2516
0
    prune_tx_split_no_split(x, plane_bsize, blk_row, blk_col, tx_size,
2517
0
                            &try_no_split, &try_split,
2518
0
                            cpi->sf.tx_sf.prune_tx_size_level);
2519
0
  }
2520
2521
0
  if (cpi->sf.rt_sf.skip_tx_no_split_var_based_partition) {
2522
0
    if (x->try_merge_partition && try_split && p->eobs[block]) try_no_split = 0;
2523
0
  }
2524
2525
  // Try using current block as a single transform block without split.
2526
0
  if (try_no_split) {
2527
0
    try_tx_block_no_split(cpi, x, blk_row, blk_col, block, tx_size, depth,
2528
0
                          plane_bsize, ta, tl, ctx, rd_stats, ref_best_rd,
2529
0
                          ftxs_mode, &no_split);
2530
2531
    // Speed features for early termination.
2532
0
    const int search_level = cpi->sf.tx_sf.adaptive_txb_search_level;
2533
0
    if (search_level) {
2534
0
      if ((no_split.rd - (no_split.rd >> (1 + search_level))) > ref_best_rd) {
2535
0
        *is_cost_valid = 0;
2536
0
        return;
2537
0
      }
2538
0
      if (no_split.rd - (no_split.rd >> (2 + search_level)) > prev_level_rd) {
2539
0
        try_split = 0;
2540
0
      }
2541
0
    }
2542
0
    if (cpi->sf.tx_sf.txb_split_cap) {
2543
0
      if (p->eobs[block] == 0) try_split = 0;
2544
0
    }
2545
0
  }
2546
2547
  // ML based speed feature to skip searching for split transform blocks.
2548
0
  if (x->e_mbd.bd == 8 && try_split &&
2549
0
      !(ref_best_rd == INT64_MAX && no_split.rd == INT64_MAX)) {
2550
0
    const int threshold = cpi->sf.tx_sf.tx_type_search.ml_tx_split_thresh;
2551
0
    if (threshold >= 0) {
2552
0
      const int split_score =
2553
0
          ml_predict_tx_split(x, plane_bsize, blk_row, blk_col, tx_size);
2554
0
      if (split_score < -threshold) try_split = 0;
2555
0
    }
2556
0
  }
2557
2558
0
  RD_STATS split_rd_stats;
2559
0
  split_rd_stats.rdcost = INT64_MAX;
2560
  // Try splitting current block into smaller transform blocks.
2561
0
  if (try_split) {
2562
0
    try_tx_block_split(cpi, x, blk_row, blk_col, block, tx_size, depth,
2563
0
                       plane_bsize, ta, tl, tx_above, tx_left, ctx, no_split.rd,
2564
0
                       AOMMIN(no_split.rd, ref_best_rd), ftxs_mode,
2565
0
                       &split_rd_stats);
2566
0
  }
2567
2568
0
  if (no_split.rd < split_rd_stats.rdcost) {
2569
0
    ENTROPY_CONTEXT *pta = ta + blk_col;
2570
0
    ENTROPY_CONTEXT *ptl = tl + blk_row;
2571
0
    p->txb_entropy_ctx[block] = no_split.txb_entropy_ctx;
2572
0
    av1_set_txb_context(x, 0, block, tx_size, pta, ptl);
2573
0
    txfm_partition_update(tx_above + blk_col, tx_left + blk_row, tx_size,
2574
0
                          tx_size);
2575
0
    for (int idy = 0; idy < tx_size_high_unit[tx_size]; ++idy) {
2576
0
      for (int idx = 0; idx < tx_size_wide_unit[tx_size]; ++idx) {
2577
0
        const int index =
2578
0
            av1_get_txb_size_index(plane_bsize, blk_row + idy, blk_col + idx);
2579
0
        mbmi->inter_tx_size[index] = tx_size;
2580
0
      }
2581
0
    }
2582
0
    mbmi->tx_size = tx_size;
2583
0
    update_txk_array(xd, blk_row, blk_col, tx_size, no_split.tx_type);
2584
0
    const int bw = mi_size_wide[plane_bsize];
2585
0
    set_blk_skip(x->txfm_search_info.blk_skip, 0, blk_row * bw + blk_col,
2586
0
                 rd_stats->skip_txfm);
2587
0
  } else {
2588
0
    *rd_stats = split_rd_stats;
2589
0
    if (split_rd_stats.rdcost == INT64_MAX) *is_cost_valid = 0;
2590
0
  }
2591
0
}
2592
2593
static AOM_INLINE void choose_largest_tx_size(const AV1_COMP *const cpi,
2594
                                              MACROBLOCK *x, RD_STATS *rd_stats,
2595
                                              int64_t ref_best_rd,
2596
2.62M
                                              BLOCK_SIZE bs) {
2597
2.62M
  MACROBLOCKD *const xd = &x->e_mbd;
2598
2.62M
  MB_MODE_INFO *const mbmi = xd->mi[0];
2599
2.62M
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
2600
2.62M
  mbmi->tx_size = tx_size_from_tx_mode(bs, txfm_params->tx_mode_search_type);
2601
2602
  // If tx64 is not enabled, we need to go down to the next available size
2603
2.62M
  if (!cpi->oxcf.txfm_cfg.enable_tx64 && cpi->oxcf.txfm_cfg.enable_rect_tx) {
2604
0
    static const TX_SIZE tx_size_max_32[TX_SIZES_ALL] = {
2605
0
      TX_4X4,    // 4x4 transform
2606
0
      TX_8X8,    // 8x8 transform
2607
0
      TX_16X16,  // 16x16 transform
2608
0
      TX_32X32,  // 32x32 transform
2609
0
      TX_32X32,  // 64x64 transform
2610
0
      TX_4X8,    // 4x8 transform
2611
0
      TX_8X4,    // 8x4 transform
2612
0
      TX_8X16,   // 8x16 transform
2613
0
      TX_16X8,   // 16x8 transform
2614
0
      TX_16X32,  // 16x32 transform
2615
0
      TX_32X16,  // 32x16 transform
2616
0
      TX_32X32,  // 32x64 transform
2617
0
      TX_32X32,  // 64x32 transform
2618
0
      TX_4X16,   // 4x16 transform
2619
0
      TX_16X4,   // 16x4 transform
2620
0
      TX_8X32,   // 8x32 transform
2621
0
      TX_32X8,   // 32x8 transform
2622
0
      TX_16X32,  // 16x64 transform
2623
0
      TX_32X16,  // 64x16 transform
2624
0
    };
2625
0
    mbmi->tx_size = tx_size_max_32[mbmi->tx_size];
2626
2.62M
  } else if (cpi->oxcf.txfm_cfg.enable_tx64 &&
2627
2.62M
             !cpi->oxcf.txfm_cfg.enable_rect_tx) {
2628
0
    static const TX_SIZE tx_size_max_square[TX_SIZES_ALL] = {
2629
0
      TX_4X4,    // 4x4 transform
2630
0
      TX_8X8,    // 8x8 transform
2631
0
      TX_16X16,  // 16x16 transform
2632
0
      TX_32X32,  // 32x32 transform
2633
0
      TX_64X64,  // 64x64 transform
2634
0
      TX_4X4,    // 4x8 transform
2635
0
      TX_4X4,    // 8x4 transform
2636
0
      TX_8X8,    // 8x16 transform
2637
0
      TX_8X8,    // 16x8 transform
2638
0
      TX_16X16,  // 16x32 transform
2639
0
      TX_16X16,  // 32x16 transform
2640
0
      TX_32X32,  // 32x64 transform
2641
0
      TX_32X32,  // 64x32 transform
2642
0
      TX_4X4,    // 4x16 transform
2643
0
      TX_4X4,    // 16x4 transform
2644
0
      TX_8X8,    // 8x32 transform
2645
0
      TX_8X8,    // 32x8 transform
2646
0
      TX_16X16,  // 16x64 transform
2647
0
      TX_16X16,  // 64x16 transform
2648
0
    };
2649
0
    mbmi->tx_size = tx_size_max_square[mbmi->tx_size];
2650
2.62M
  } else if (!cpi->oxcf.txfm_cfg.enable_tx64 &&
2651
2.62M
             !cpi->oxcf.txfm_cfg.enable_rect_tx) {
2652
0
    static const TX_SIZE tx_size_max_32_square[TX_SIZES_ALL] = {
2653
0
      TX_4X4,    // 4x4 transform
2654
0
      TX_8X8,    // 8x8 transform
2655
0
      TX_16X16,  // 16x16 transform
2656
0
      TX_32X32,  // 32x32 transform
2657
0
      TX_32X32,  // 64x64 transform
2658
0
      TX_4X4,    // 4x8 transform
2659
0
      TX_4X4,    // 8x4 transform
2660
0
      TX_8X8,    // 8x16 transform
2661
0
      TX_8X8,    // 16x8 transform
2662
0
      TX_16X16,  // 16x32 transform
2663
0
      TX_16X16,  // 32x16 transform
2664
0
      TX_32X32,  // 32x64 transform
2665
0
      TX_32X32,  // 64x32 transform
2666
0
      TX_4X4,    // 4x16 transform
2667
0
      TX_4X4,    // 16x4 transform
2668
0
      TX_8X8,    // 8x32 transform
2669
0
      TX_8X8,    // 32x8 transform
2670
0
      TX_16X16,  // 16x64 transform
2671
0
      TX_16X16,  // 64x16 transform
2672
0
    };
2673
2674
0
    mbmi->tx_size = tx_size_max_32_square[mbmi->tx_size];
2675
0
  }
2676
2677
2.62M
  const int skip_ctx = av1_get_skip_txfm_context(xd);
2678
2.62M
  const int no_skip_txfm_rate = x->mode_costs.skip_txfm_cost[skip_ctx][0];
2679
2.62M
  const int skip_txfm_rate = x->mode_costs.skip_txfm_cost[skip_ctx][1];
2680
  // Skip RDcost is used only for Inter blocks
2681
2.62M
  const int64_t skip_txfm_rd =
2682
2.62M
      is_inter_block(mbmi) ? RDCOST(x->rdmult, skip_txfm_rate, 0) : INT64_MAX;
2683
2.62M
  const int64_t no_skip_txfm_rd = RDCOST(x->rdmult, no_skip_txfm_rate, 0);
2684
2.62M
  const int skip_trellis = 0;
2685
2.62M
  av1_txfm_rd_in_plane(x, cpi, rd_stats, ref_best_rd,
2686
2.62M
                       AOMMIN(no_skip_txfm_rd, skip_txfm_rd), AOM_PLANE_Y, bs,
2687
2.62M
                       mbmi->tx_size, FTXS_NONE, skip_trellis);
2688
2.62M
}
2689
2690
static AOM_INLINE void choose_smallest_tx_size(const AV1_COMP *const cpi,
2691
                                               MACROBLOCK *x,
2692
                                               RD_STATS *rd_stats,
2693
                                               int64_t ref_best_rd,
2694
1.16M
                                               BLOCK_SIZE bs) {
2695
1.16M
  MACROBLOCKD *const xd = &x->e_mbd;
2696
1.16M
  MB_MODE_INFO *const mbmi = xd->mi[0];
2697
2698
1.16M
  mbmi->tx_size = TX_4X4;
2699
  // TODO(any) : Pass this_rd based on skip/non-skip cost
2700
1.16M
  const int skip_trellis = 0;
2701
1.16M
  av1_txfm_rd_in_plane(x, cpi, rd_stats, ref_best_rd, 0, 0, bs, mbmi->tx_size,
2702
1.16M
                       FTXS_NONE, skip_trellis);
2703
1.16M
}
2704
2705
// Search for the best uniform transform size and type for current coding block.
2706
static AOM_INLINE void choose_tx_size_type_from_rd(const AV1_COMP *const cpi,
2707
                                                   MACROBLOCK *x,
2708
                                                   RD_STATS *rd_stats,
2709
                                                   int64_t ref_best_rd,
2710
94.9k
                                                   BLOCK_SIZE bs) {
2711
94.9k
  av1_invalid_rd_stats(rd_stats);
2712
2713
94.9k
  MACROBLOCKD *const xd = &x->e_mbd;
2714
94.9k
  MB_MODE_INFO *const mbmi = xd->mi[0];
2715
94.9k
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
2716
94.9k
  const TX_SIZE max_rect_tx_size = max_txsize_rect_lookup[bs];
2717
94.9k
  const int tx_select = txfm_params->tx_mode_search_type == TX_MODE_SELECT;
2718
94.9k
  int start_tx;
2719
  // The split depth can be at most MAX_TX_DEPTH, so the init_depth controls
2720
  // how many times of splitting is allowed during the RD search.
2721
94.9k
  int init_depth;
2722
2723
94.9k
  if (tx_select) {
2724
94.9k
    start_tx = max_rect_tx_size;
2725
94.9k
    init_depth = get_search_init_depth(mi_size_wide[bs], mi_size_high[bs],
2726
94.9k
                                       is_inter_block(mbmi), &cpi->sf,
2727
94.9k
                                       txfm_params->tx_size_search_method);
2728
94.9k
    if (init_depth == MAX_TX_DEPTH && !cpi->oxcf.txfm_cfg.enable_tx64 &&
2729
94.9k
        txsize_sqr_up_map[start_tx] == TX_64X64) {
2730
0
      start_tx = sub_tx_size_map[start_tx];
2731
0
    }
2732
94.9k
  } else {
2733
3
    const TX_SIZE chosen_tx_size =
2734
3
        tx_size_from_tx_mode(bs, txfm_params->tx_mode_search_type);
2735
3
    start_tx = chosen_tx_size;
2736
3
    init_depth = MAX_TX_DEPTH;
2737
3
  }
2738
2739
94.9k
  const int skip_trellis = 0;
2740
94.9k
  uint8_t best_txk_type_map[MAX_MIB_SIZE * MAX_MIB_SIZE];
2741
94.9k
  uint8_t best_blk_skip[MAX_MIB_SIZE * MAX_MIB_SIZE];
2742
94.9k
  TX_SIZE best_tx_size = max_rect_tx_size;
2743
94.9k
  int64_t best_rd = INT64_MAX;
2744
94.9k
  const int num_blks = bsize_to_num_blk(bs);
2745
94.9k
  x->rd_model = FULL_TXFM_RD;
2746
94.9k
  int64_t rd[MAX_TX_DEPTH + 1] = { INT64_MAX, INT64_MAX, INT64_MAX };
2747
94.9k
  TxfmSearchInfo *txfm_info = &x->txfm_search_info;
2748
242k
  for (int tx_size = start_tx, depth = init_depth; depth <= MAX_TX_DEPTH;
2749
182k
       depth++, tx_size = sub_tx_size_map[tx_size]) {
2750
182k
    if ((!cpi->oxcf.txfm_cfg.enable_tx64 &&
2751
182k
         txsize_sqr_up_map[tx_size] == TX_64X64) ||
2752
182k
        (!cpi->oxcf.txfm_cfg.enable_rect_tx &&
2753
182k
         tx_size_wide[tx_size] != tx_size_high[tx_size])) {
2754
0
      continue;
2755
0
    }
2756
2757
182k
    RD_STATS this_rd_stats;
2758
182k
    rd[depth] = av1_uniform_txfm_yrd(cpi, x, &this_rd_stats, ref_best_rd, bs,
2759
182k
                                     tx_size, FTXS_NONE, skip_trellis);
2760
182k
    if (rd[depth] < best_rd) {
2761
158k
      av1_copy_array(best_blk_skip, txfm_info->blk_skip, num_blks);
2762
158k
      av1_copy_array(best_txk_type_map, xd->tx_type_map, num_blks);
2763
158k
      best_tx_size = tx_size;
2764
158k
      best_rd = rd[depth];
2765
158k
      *rd_stats = this_rd_stats;
2766
158k
    }
2767
182k
    if (tx_size == TX_4X4) break;
2768
    // If we are searching three depths, prune the smallest size depending
2769
    // on rd results for the first two depths for low contrast blocks.
2770
147k
    if (depth > init_depth && depth != MAX_TX_DEPTH &&
2771
147k
        x->source_variance < 256) {
2772
0
      if (rd[depth - 1] != INT64_MAX && rd[depth] > rd[depth - 1]) break;
2773
0
    }
2774
147k
  }
2775
2776
94.9k
  if (rd_stats->rate != INT_MAX) {
2777
94.9k
    mbmi->tx_size = best_tx_size;
2778
94.9k
    av1_copy_array(xd->tx_type_map, best_txk_type_map, num_blks);
2779
94.9k
    av1_copy_array(txfm_info->blk_skip, best_blk_skip, num_blks);
2780
94.9k
  }
2781
94.9k
}
2782
2783
// Search for the best transform type for the given transform block in the
2784
// given plane/channel, and calculate the corresponding RD cost.
2785
static AOM_INLINE void block_rd_txfm(int plane, int block, int blk_row,
2786
                                     int blk_col, BLOCK_SIZE plane_bsize,
2787
44.3M
                                     TX_SIZE tx_size, void *arg) {
2788
44.3M
  struct rdcost_block_args *args = arg;
2789
44.3M
  if (args->exit_early) {
2790
47.4k
    args->incomplete_exit = 1;
2791
47.4k
    return;
2792
47.4k
  }
2793
2794
44.2M
  MACROBLOCK *const x = args->x;
2795
44.2M
  MACROBLOCKD *const xd = &x->e_mbd;
2796
44.2M
  const int is_inter = is_inter_block(xd->mi[0]);
2797
44.2M
  const AV1_COMP *cpi = args->cpi;
2798
44.2M
  ENTROPY_CONTEXT *a = args->t_above + blk_col;
2799
44.2M
  ENTROPY_CONTEXT *l = args->t_left + blk_row;
2800
44.2M
  const AV1_COMMON *cm = &cpi->common;
2801
44.2M
  RD_STATS this_rd_stats;
2802
44.2M
  av1_init_rd_stats(&this_rd_stats);
2803
2804
44.2M
  if (!is_inter) {
2805
44.1M
    av1_predict_intra_block_facade(cm, xd, plane, blk_col, blk_row, tx_size);
2806
44.1M
    av1_subtract_txb(x, plane, plane_bsize, blk_col, blk_row, tx_size);
2807
44.1M
  }
2808
2809
44.2M
  TXB_CTX txb_ctx;
2810
44.2M
  get_txb_ctx(plane_bsize, tx_size, plane, a, l, &txb_ctx);
2811
44.2M
  search_tx_type(cpi, x, plane, block, blk_row, blk_col, plane_bsize, tx_size,
2812
44.2M
                 &txb_ctx, args->ftxs_mode, args->skip_trellis,
2813
44.2M
                 args->best_rd - args->current_rd, &this_rd_stats);
2814
2815
44.2M
  if (plane == AOM_PLANE_Y && xd->cfl.store_y) {
2816
0
    assert(!is_inter || plane_bsize < BLOCK_8X8);
2817
0
    cfl_store_tx(xd, blk_row, blk_col, tx_size, plane_bsize);
2818
0
  }
2819
2820
#if CONFIG_RD_DEBUG
2821
  update_txb_coeff_cost(&this_rd_stats, plane, this_rd_stats.rate);
2822
#endif  // CONFIG_RD_DEBUG
2823
44.2M
  av1_set_txb_context(x, plane, block, tx_size, a, l);
2824
2825
44.2M
  const int blk_idx =
2826
44.2M
      blk_row * (block_size_wide[plane_bsize] >> MI_SIZE_LOG2) + blk_col;
2827
2828
44.2M
  TxfmSearchInfo *txfm_info = &x->txfm_search_info;
2829
44.2M
  if (plane == 0)
2830
43.1M
    set_blk_skip(txfm_info->blk_skip, plane, blk_idx,
2831
43.1M
                 x->plane[plane].eobs[block] == 0);
2832
1.18M
  else
2833
1.18M
    set_blk_skip(txfm_info->blk_skip, plane, blk_idx, 0);
2834
2835
44.2M
  int64_t rd;
2836
44.2M
  if (is_inter) {
2837
0
    const int64_t no_skip_txfm_rd =
2838
0
        RDCOST(x->rdmult, this_rd_stats.rate, this_rd_stats.dist);
2839
0
    const int64_t skip_txfm_rd = RDCOST(x->rdmult, 0, this_rd_stats.sse);
2840
0
    rd = AOMMIN(no_skip_txfm_rd, skip_txfm_rd);
2841
0
    this_rd_stats.skip_txfm &= !x->plane[plane].eobs[block];
2842
44.2M
  } else {
2843
    // Signal non-skip_txfm for Intra blocks
2844
44.2M
    rd = RDCOST(x->rdmult, this_rd_stats.rate, this_rd_stats.dist);
2845
44.2M
    this_rd_stats.skip_txfm = 0;
2846
44.2M
  }
2847
2848
44.2M
  av1_merge_rd_stats(&args->rd_stats, &this_rd_stats);
2849
2850
44.2M
  args->current_rd += rd;
2851
44.2M
  if (args->current_rd > args->best_rd) args->exit_early = 1;
2852
44.2M
}
2853
2854
int64_t av1_estimate_txfm_yrd(const AV1_COMP *const cpi, MACROBLOCK *x,
2855
                              RD_STATS *rd_stats, int64_t ref_best_rd,
2856
0
                              BLOCK_SIZE bs, TX_SIZE tx_size) {
2857
0
  MACROBLOCKD *const xd = &x->e_mbd;
2858
0
  MB_MODE_INFO *const mbmi = xd->mi[0];
2859
0
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
2860
0
  const ModeCosts *mode_costs = &x->mode_costs;
2861
0
  const int is_inter = is_inter_block(mbmi);
2862
0
  const int tx_select = txfm_params->tx_mode_search_type == TX_MODE_SELECT &&
2863
0
                        block_signals_txsize(mbmi->bsize);
2864
0
  int tx_size_rate = 0;
2865
0
  if (tx_select) {
2866
0
    const int ctx = txfm_partition_context(
2867
0
        xd->above_txfm_context, xd->left_txfm_context, mbmi->bsize, tx_size);
2868
0
    tx_size_rate = mode_costs->txfm_partition_cost[ctx][0];
2869
0
  }
2870
0
  const int skip_ctx = av1_get_skip_txfm_context(xd);
2871
0
  const int no_skip_txfm_rate = mode_costs->skip_txfm_cost[skip_ctx][0];
2872
0
  const int skip_txfm_rate = mode_costs->skip_txfm_cost[skip_ctx][1];
2873
0
  const int64_t skip_txfm_rd = RDCOST(x->rdmult, skip_txfm_rate, 0);
2874
0
  const int64_t no_this_rd =
2875
0
      RDCOST(x->rdmult, no_skip_txfm_rate + tx_size_rate, 0);
2876
0
  mbmi->tx_size = tx_size;
2877
2878
0
  const uint8_t txw_unit = tx_size_wide_unit[tx_size];
2879
0
  const uint8_t txh_unit = tx_size_high_unit[tx_size];
2880
0
  const int step = txw_unit * txh_unit;
2881
0
  const int max_blocks_wide = max_block_wide(xd, bs, 0);
2882
0
  const int max_blocks_high = max_block_high(xd, bs, 0);
2883
2884
0
  struct rdcost_block_args args;
2885
0
  av1_zero(args);
2886
0
  args.x = x;
2887
0
  args.cpi = cpi;
2888
0
  args.best_rd = ref_best_rd;
2889
0
  args.current_rd = AOMMIN(no_this_rd, skip_txfm_rd);
2890
0
  av1_init_rd_stats(&args.rd_stats);
2891
0
  av1_get_entropy_contexts(bs, &xd->plane[0], args.t_above, args.t_left);
2892
0
  int i = 0;
2893
0
  for (int blk_row = 0; blk_row < max_blocks_high && !args.incomplete_exit;
2894
0
       blk_row += txh_unit) {
2895
0
    for (int blk_col = 0; blk_col < max_blocks_wide; blk_col += txw_unit) {
2896
0
      RD_STATS this_rd_stats;
2897
0
      av1_init_rd_stats(&this_rd_stats);
2898
2899
0
      if (args.exit_early) {
2900
0
        args.incomplete_exit = 1;
2901
0
        break;
2902
0
      }
2903
2904
0
      ENTROPY_CONTEXT *a = args.t_above + blk_col;
2905
0
      ENTROPY_CONTEXT *l = args.t_left + blk_row;
2906
0
      TXB_CTX txb_ctx;
2907
0
      get_txb_ctx(bs, tx_size, 0, a, l, &txb_ctx);
2908
2909
0
      TxfmParam txfm_param;
2910
0
      QUANT_PARAM quant_param;
2911
0
      av1_setup_xform(&cpi->common, x, tx_size, DCT_DCT, &txfm_param);
2912
0
      av1_setup_quant(tx_size, 0, AV1_XFORM_QUANT_B, 0, &quant_param);
2913
2914
0
      av1_xform(x, 0, i, blk_row, blk_col, bs, &txfm_param);
2915
0
      av1_quant(x, 0, i, &txfm_param, &quant_param);
2916
2917
0
      this_rd_stats.rate =
2918
0
          cost_coeffs(x, 0, i, tx_size, txfm_param.tx_type, &txb_ctx, 0);
2919
2920
0
      dist_block_tx_domain(x, 0, i, tx_size, &this_rd_stats.dist,
2921
0
                           &this_rd_stats.sse);
2922
2923
0
      const int64_t no_skip_txfm_rd =
2924
0
          RDCOST(x->rdmult, this_rd_stats.rate, this_rd_stats.dist);
2925
0
      const int64_t skip_rd = RDCOST(x->rdmult, 0, this_rd_stats.sse);
2926
2927
0
      this_rd_stats.skip_txfm &= !x->plane[0].eobs[i];
2928
2929
0
      av1_merge_rd_stats(&args.rd_stats, &this_rd_stats);
2930
0
      args.current_rd += AOMMIN(no_skip_txfm_rd, skip_rd);
2931
2932
0
      if (args.current_rd > ref_best_rd) {
2933
0
        args.exit_early = 1;
2934
0
        break;
2935
0
      }
2936
2937
0
      av1_set_txb_context(x, 0, i, tx_size, a, l);
2938
0
      i += step;
2939
0
    }
2940
0
  }
2941
2942
0
  if (args.incomplete_exit) av1_invalid_rd_stats(&args.rd_stats);
2943
2944
0
  *rd_stats = args.rd_stats;
2945
0
  if (rd_stats->rate == INT_MAX) return INT64_MAX;
2946
2947
0
  int64_t rd;
2948
  // rdstats->rate should include all the rate except skip/non-skip cost as the
2949
  // same is accounted in the caller functions after rd evaluation of all
2950
  // planes. However the decisions should be done after considering the
2951
  // skip/non-skip header cost
2952
0
  if (rd_stats->skip_txfm && is_inter) {
2953
0
    rd = RDCOST(x->rdmult, skip_txfm_rate, rd_stats->sse);
2954
0
  } else {
2955
    // Intra blocks are always signalled as non-skip
2956
0
    rd = RDCOST(x->rdmult, rd_stats->rate + no_skip_txfm_rate + tx_size_rate,
2957
0
                rd_stats->dist);
2958
0
    rd_stats->rate += tx_size_rate;
2959
0
  }
2960
  // Check if forcing the block to skip transform leads to smaller RD cost.
2961
0
  if (is_inter && !rd_stats->skip_txfm && !xd->lossless[mbmi->segment_id]) {
2962
0
    int64_t temp_skip_txfm_rd =
2963
0
        RDCOST(x->rdmult, skip_txfm_rate, rd_stats->sse);
2964
0
    if (temp_skip_txfm_rd <= rd) {
2965
0
      rd = temp_skip_txfm_rd;
2966
0
      rd_stats->rate = 0;
2967
0
      rd_stats->dist = rd_stats->sse;
2968
0
      rd_stats->skip_txfm = 1;
2969
0
    }
2970
0
  }
2971
2972
0
  return rd;
2973
0
}
2974
2975
int64_t av1_uniform_txfm_yrd(const AV1_COMP *const cpi, MACROBLOCK *x,
2976
                             RD_STATS *rd_stats, int64_t ref_best_rd,
2977
                             BLOCK_SIZE bs, TX_SIZE tx_size,
2978
182k
                             FAST_TX_SEARCH_MODE ftxs_mode, int skip_trellis) {
2979
182k
  assert(IMPLIES(is_rect_tx(tx_size), is_rect_tx_allowed_bsize(bs)));
2980
182k
  MACROBLOCKD *const xd = &x->e_mbd;
2981
182k
  MB_MODE_INFO *const mbmi = xd->mi[0];
2982
182k
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
2983
182k
  const ModeCosts *mode_costs = &x->mode_costs;
2984
182k
  const int is_inter = is_inter_block(mbmi);
2985
182k
  const int tx_select = txfm_params->tx_mode_search_type == TX_MODE_SELECT &&
2986
182k
                        block_signals_txsize(mbmi->bsize);
2987
182k
  int tx_size_rate = 0;
2988
182k
  if (tx_select) {
2989
175k
    const int ctx = txfm_partition_context(
2990
175k
        xd->above_txfm_context, xd->left_txfm_context, mbmi->bsize, tx_size);
2991
175k
    tx_size_rate = is_inter ? mode_costs->txfm_partition_cost[ctx][0]
2992
175k
                            : tx_size_cost(x, bs, tx_size);
2993
175k
  }
2994
182k
  const int skip_ctx = av1_get_skip_txfm_context(xd);
2995
182k
  const int no_skip_txfm_rate = mode_costs->skip_txfm_cost[skip_ctx][0];
2996
182k
  const int skip_txfm_rate = mode_costs->skip_txfm_cost[skip_ctx][1];
2997
182k
  const int64_t skip_txfm_rd =
2998
182k
      is_inter ? RDCOST(x->rdmult, skip_txfm_rate, 0) : INT64_MAX;
2999
182k
  const int64_t no_this_rd =
3000
182k
      RDCOST(x->rdmult, no_skip_txfm_rate + tx_size_rate, 0);
3001
3002
182k
  mbmi->tx_size = tx_size;
3003
182k
  av1_txfm_rd_in_plane(x, cpi, rd_stats, ref_best_rd,
3004
182k
                       AOMMIN(no_this_rd, skip_txfm_rd), AOM_PLANE_Y, bs,
3005
182k
                       tx_size, ftxs_mode, skip_trellis);
3006
182k
  if (rd_stats->rate == INT_MAX) return INT64_MAX;
3007
3008
182k
  int64_t rd;
3009
  // rdstats->rate should include all the rate except skip/non-skip cost as the
3010
  // same is accounted in the caller functions after rd evaluation of all
3011
  // planes. However the decisions should be done after considering the
3012
  // skip/non-skip header cost
3013
182k
  if (rd_stats->skip_txfm && is_inter) {
3014
0
    rd = RDCOST(x->rdmult, skip_txfm_rate, rd_stats->sse);
3015
182k
  } else {
3016
    // Intra blocks are always signalled as non-skip
3017
182k
    rd = RDCOST(x->rdmult, rd_stats->rate + no_skip_txfm_rate + tx_size_rate,
3018
182k
                rd_stats->dist);
3019
182k
    rd_stats->rate += tx_size_rate;
3020
182k
  }
3021
  // Check if forcing the block to skip transform leads to smaller RD cost.
3022
182k
  if (is_inter && !rd_stats->skip_txfm && !xd->lossless[mbmi->segment_id]) {
3023
0
    int64_t temp_skip_txfm_rd =
3024
0
        RDCOST(x->rdmult, skip_txfm_rate, rd_stats->sse);
3025
0
    if (temp_skip_txfm_rd <= rd) {
3026
0
      rd = temp_skip_txfm_rd;
3027
0
      rd_stats->rate = 0;
3028
0
      rd_stats->dist = rd_stats->sse;
3029
0
      rd_stats->skip_txfm = 1;
3030
0
    }
3031
0
  }
3032
3033
182k
  return rd;
3034
182k
}
3035
3036
// Search for the best transform type for a luma inter-predicted block, given
3037
// the transform block partitions.
3038
// This function is used only when some speed features are enabled.
3039
static AOM_INLINE void tx_block_yrd(
3040
    const AV1_COMP *cpi, MACROBLOCK *x, int blk_row, int blk_col, int block,
3041
    TX_SIZE tx_size, BLOCK_SIZE plane_bsize, int depth,
3042
    ENTROPY_CONTEXT *above_ctx, ENTROPY_CONTEXT *left_ctx,
3043
    TXFM_CONTEXT *tx_above, TXFM_CONTEXT *tx_left, int64_t ref_best_rd,
3044
0
    RD_STATS *rd_stats, FAST_TX_SEARCH_MODE ftxs_mode) {
3045
0
  assert(tx_size < TX_SIZES_ALL);
3046
0
  MACROBLOCKD *const xd = &x->e_mbd;
3047
0
  MB_MODE_INFO *const mbmi = xd->mi[0];
3048
0
  assert(is_inter_block(mbmi));
3049
0
  const int max_blocks_high = max_block_high(xd, plane_bsize, 0);
3050
0
  const int max_blocks_wide = max_block_wide(xd, plane_bsize, 0);
3051
3052
0
  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
3053
3054
0
  const TX_SIZE plane_tx_size = mbmi->inter_tx_size[av1_get_txb_size_index(
3055
0
      plane_bsize, blk_row, blk_col)];
3056
0
  const int ctx = txfm_partition_context(tx_above + blk_col, tx_left + blk_row,
3057
0
                                         mbmi->bsize, tx_size);
3058
3059
0
  av1_init_rd_stats(rd_stats);
3060
0
  if (tx_size == plane_tx_size) {
3061
0
    ENTROPY_CONTEXT *ta = above_ctx + blk_col;
3062
0
    ENTROPY_CONTEXT *tl = left_ctx + blk_row;
3063
0
    const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
3064
0
    TXB_CTX txb_ctx;
3065
0
    get_txb_ctx(plane_bsize, tx_size, 0, ta, tl, &txb_ctx);
3066
3067
0
    const int zero_blk_rate =
3068
0
        x->coeff_costs.coeff_costs[txs_ctx][get_plane_type(0)]
3069
0
            .txb_skip_cost[txb_ctx.txb_skip_ctx][1];
3070
0
    rd_stats->zero_rate = zero_blk_rate;
3071
0
    tx_type_rd(cpi, x, tx_size, blk_row, blk_col, block, plane_bsize, &txb_ctx,
3072
0
               rd_stats, ftxs_mode, ref_best_rd);
3073
0
    const int mi_width = mi_size_wide[plane_bsize];
3074
0
    TxfmSearchInfo *txfm_info = &x->txfm_search_info;
3075
0
    if (RDCOST(x->rdmult, rd_stats->rate, rd_stats->dist) >=
3076
0
            RDCOST(x->rdmult, zero_blk_rate, rd_stats->sse) ||
3077
0
        rd_stats->skip_txfm == 1) {
3078
0
      rd_stats->rate = zero_blk_rate;
3079
0
      rd_stats->dist = rd_stats->sse;
3080
0
      rd_stats->skip_txfm = 1;
3081
0
      set_blk_skip(txfm_info->blk_skip, 0, blk_row * mi_width + blk_col, 1);
3082
0
      x->plane[0].eobs[block] = 0;
3083
0
      x->plane[0].txb_entropy_ctx[block] = 0;
3084
0
      update_txk_array(xd, blk_row, blk_col, tx_size, DCT_DCT);
3085
0
    } else {
3086
0
      rd_stats->skip_txfm = 0;
3087
0
      set_blk_skip(txfm_info->blk_skip, 0, blk_row * mi_width + blk_col, 0);
3088
0
    }
3089
0
    if (tx_size > TX_4X4 && depth < MAX_VARTX_DEPTH)
3090
0
      rd_stats->rate += x->mode_costs.txfm_partition_cost[ctx][0];
3091
0
    av1_set_txb_context(x, 0, block, tx_size, ta, tl);
3092
0
    txfm_partition_update(tx_above + blk_col, tx_left + blk_row, tx_size,
3093
0
                          tx_size);
3094
0
  } else {
3095
0
    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
3096
0
    const int txb_width = tx_size_wide_unit[sub_txs];
3097
0
    const int txb_height = tx_size_high_unit[sub_txs];
3098
0
    const int step = txb_height * txb_width;
3099
0
    const int row_end =
3100
0
        AOMMIN(tx_size_high_unit[tx_size], max_blocks_high - blk_row);
3101
0
    const int col_end =
3102
0
        AOMMIN(tx_size_wide_unit[tx_size], max_blocks_wide - blk_col);
3103
0
    RD_STATS pn_rd_stats;
3104
0
    int64_t this_rd = 0;
3105
0
    assert(txb_width > 0 && txb_height > 0);
3106
3107
0
    for (int row = 0; row < row_end; row += txb_height) {
3108
0
      const int offsetr = blk_row + row;
3109
0
      for (int col = 0; col < col_end; col += txb_width) {
3110
0
        const int offsetc = blk_col + col;
3111
3112
0
        av1_init_rd_stats(&pn_rd_stats);
3113
0
        tx_block_yrd(cpi, x, offsetr, offsetc, block, sub_txs, plane_bsize,
3114
0
                     depth + 1, above_ctx, left_ctx, tx_above, tx_left,
3115
0
                     ref_best_rd - this_rd, &pn_rd_stats, ftxs_mode);
3116
0
        if (pn_rd_stats.rate == INT_MAX) {
3117
0
          av1_invalid_rd_stats(rd_stats);
3118
0
          return;
3119
0
        }
3120
0
        av1_merge_rd_stats(rd_stats, &pn_rd_stats);
3121
0
        this_rd += RDCOST(x->rdmult, pn_rd_stats.rate, pn_rd_stats.dist);
3122
0
        block += step;
3123
0
      }
3124
0
    }
3125
3126
0
    if (tx_size > TX_4X4 && depth < MAX_VARTX_DEPTH)
3127
0
      rd_stats->rate += x->mode_costs.txfm_partition_cost[ctx][1];
3128
0
  }
3129
0
}
3130
3131
// search for tx type with tx sizes already decided for a inter-predicted luma
3132
// partition block. It's used only when some speed features are enabled.
3133
// Return value 0: early termination triggered, no valid rd cost available;
3134
//              1: rd cost values are valid.
3135
static int inter_block_yrd(const AV1_COMP *cpi, MACROBLOCK *x,
3136
                           RD_STATS *rd_stats, BLOCK_SIZE bsize,
3137
0
                           int64_t ref_best_rd, FAST_TX_SEARCH_MODE ftxs_mode) {
3138
0
  if (ref_best_rd < 0) {
3139
0
    av1_invalid_rd_stats(rd_stats);
3140
0
    return 0;
3141
0
  }
3142
3143
0
  av1_init_rd_stats(rd_stats);
3144
3145
0
  MACROBLOCKD *const xd = &x->e_mbd;
3146
0
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
3147
0
  const struct macroblockd_plane *const pd = &xd->plane[0];
3148
0
  const int mi_width = mi_size_wide[bsize];
3149
0
  const int mi_height = mi_size_high[bsize];
3150
0
  const TX_SIZE max_tx_size = get_vartx_max_txsize(xd, bsize, 0);
3151
0
  const int bh = tx_size_high_unit[max_tx_size];
3152
0
  const int bw = tx_size_wide_unit[max_tx_size];
3153
0
  const int step = bw * bh;
3154
0
  const int init_depth = get_search_init_depth(
3155
0
      mi_width, mi_height, 1, &cpi->sf, txfm_params->tx_size_search_method);
3156
0
  ENTROPY_CONTEXT ctxa[MAX_MIB_SIZE];
3157
0
  ENTROPY_CONTEXT ctxl[MAX_MIB_SIZE];
3158
0
  TXFM_CONTEXT tx_above[MAX_MIB_SIZE];
3159
0
  TXFM_CONTEXT tx_left[MAX_MIB_SIZE];
3160
0
  av1_get_entropy_contexts(bsize, pd, ctxa, ctxl);
3161
0
  memcpy(tx_above, xd->above_txfm_context, sizeof(TXFM_CONTEXT) * mi_width);
3162
0
  memcpy(tx_left, xd->left_txfm_context, sizeof(TXFM_CONTEXT) * mi_height);
3163
3164
0
  int64_t this_rd = 0;
3165
0
  for (int idy = 0, block = 0; idy < mi_height; idy += bh) {
3166
0
    for (int idx = 0; idx < mi_width; idx += bw) {
3167
0
      RD_STATS pn_rd_stats;
3168
0
      av1_init_rd_stats(&pn_rd_stats);
3169
0
      tx_block_yrd(cpi, x, idy, idx, block, max_tx_size, bsize, init_depth,
3170
0
                   ctxa, ctxl, tx_above, tx_left, ref_best_rd - this_rd,
3171
0
                   &pn_rd_stats, ftxs_mode);
3172
0
      if (pn_rd_stats.rate == INT_MAX) {
3173
0
        av1_invalid_rd_stats(rd_stats);
3174
0
        return 0;
3175
0
      }
3176
0
      av1_merge_rd_stats(rd_stats, &pn_rd_stats);
3177
0
      this_rd +=
3178
0
          AOMMIN(RDCOST(x->rdmult, pn_rd_stats.rate, pn_rd_stats.dist),
3179
0
                 RDCOST(x->rdmult, pn_rd_stats.zero_rate, pn_rd_stats.sse));
3180
0
      block += step;
3181
0
    }
3182
0
  }
3183
3184
0
  const int skip_ctx = av1_get_skip_txfm_context(xd);
3185
0
  const int no_skip_txfm_rate = x->mode_costs.skip_txfm_cost[skip_ctx][0];
3186
0
  const int skip_txfm_rate = x->mode_costs.skip_txfm_cost[skip_ctx][1];
3187
0
  const int64_t skip_txfm_rd = RDCOST(x->rdmult, skip_txfm_rate, rd_stats->sse);
3188
0
  this_rd =
3189
0
      RDCOST(x->rdmult, rd_stats->rate + no_skip_txfm_rate, rd_stats->dist);
3190
0
  if (skip_txfm_rd < this_rd) {
3191
0
    this_rd = skip_txfm_rd;
3192
0
    rd_stats->rate = 0;
3193
0
    rd_stats->dist = rd_stats->sse;
3194
0
    rd_stats->skip_txfm = 1;
3195
0
  }
3196
3197
0
  const int is_cost_valid = this_rd > ref_best_rd;
3198
0
  if (!is_cost_valid) {
3199
    // reset cost value
3200
0
    av1_invalid_rd_stats(rd_stats);
3201
0
  }
3202
0
  return is_cost_valid;
3203
0
}
3204
3205
// Search for the best transform size and type for current inter-predicted
3206
// luma block with recursive transform block partitioning. The obtained
3207
// transform selection will be saved in xd->mi[0], the corresponding RD stats
3208
// will be saved in rd_stats. The returned value is the corresponding RD cost.
3209
static int64_t select_tx_size_and_type(const AV1_COMP *cpi, MACROBLOCK *x,
3210
                                       RD_STATS *rd_stats, BLOCK_SIZE bsize,
3211
0
                                       int64_t ref_best_rd) {
3212
0
  MACROBLOCKD *const xd = &x->e_mbd;
3213
0
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
3214
0
  assert(is_inter_block(xd->mi[0]));
3215
0
  assert(bsize < BLOCK_SIZES_ALL);
3216
0
  const int fast_tx_search = txfm_params->tx_size_search_method > USE_FULL_RD;
3217
0
  int64_t rd_thresh = ref_best_rd;
3218
0
  if (rd_thresh == 0) {
3219
0
    av1_invalid_rd_stats(rd_stats);
3220
0
    return INT64_MAX;
3221
0
  }
3222
0
  if (fast_tx_search && rd_thresh < INT64_MAX) {
3223
0
    if (INT64_MAX - rd_thresh > (rd_thresh >> 3)) rd_thresh += (rd_thresh >> 3);
3224
0
  }
3225
0
  assert(rd_thresh > 0);
3226
0
  const FAST_TX_SEARCH_MODE ftxs_mode =
3227
0
      fast_tx_search ? FTXS_DCT_AND_1D_DCT_ONLY : FTXS_NONE;
3228
0
  const struct macroblockd_plane *const pd = &xd->plane[0];
3229
0
  assert(bsize < BLOCK_SIZES_ALL);
3230
0
  const int mi_width = mi_size_wide[bsize];
3231
0
  const int mi_height = mi_size_high[bsize];
3232
0
  ENTROPY_CONTEXT ctxa[MAX_MIB_SIZE];
3233
0
  ENTROPY_CONTEXT ctxl[MAX_MIB_SIZE];
3234
0
  TXFM_CONTEXT tx_above[MAX_MIB_SIZE];
3235
0
  TXFM_CONTEXT tx_left[MAX_MIB_SIZE];
3236
0
  av1_get_entropy_contexts(bsize, pd, ctxa, ctxl);
3237
0
  memcpy(tx_above, xd->above_txfm_context, sizeof(TXFM_CONTEXT) * mi_width);
3238
0
  memcpy(tx_left, xd->left_txfm_context, sizeof(TXFM_CONTEXT) * mi_height);
3239
0
  const int init_depth = get_search_init_depth(
3240
0
      mi_width, mi_height, 1, &cpi->sf, txfm_params->tx_size_search_method);
3241
0
  const TX_SIZE max_tx_size = max_txsize_rect_lookup[bsize];
3242
0
  const int bh = tx_size_high_unit[max_tx_size];
3243
0
  const int bw = tx_size_wide_unit[max_tx_size];
3244
0
  const int step = bw * bh;
3245
0
  const int skip_ctx = av1_get_skip_txfm_context(xd);
3246
0
  const int no_skip_txfm_cost = x->mode_costs.skip_txfm_cost[skip_ctx][0];
3247
0
  const int skip_txfm_cost = x->mode_costs.skip_txfm_cost[skip_ctx][1];
3248
0
  int64_t skip_txfm_rd = RDCOST(x->rdmult, skip_txfm_cost, 0);
3249
0
  int64_t no_skip_txfm_rd = RDCOST(x->rdmult, no_skip_txfm_cost, 0);
3250
0
  int block = 0;
3251
3252
0
  av1_init_rd_stats(rd_stats);
3253
0
  for (int idy = 0; idy < max_block_high(xd, bsize, 0); idy += bh) {
3254
0
    for (int idx = 0; idx < max_block_wide(xd, bsize, 0); idx += bw) {
3255
0
      const int64_t best_rd_sofar =
3256
0
          (rd_thresh == INT64_MAX)
3257
0
              ? INT64_MAX
3258
0
              : (rd_thresh - (AOMMIN(skip_txfm_rd, no_skip_txfm_rd)));
3259
0
      int is_cost_valid = 1;
3260
0
      RD_STATS pn_rd_stats;
3261
      // Search for the best transform block size and type for the sub-block.
3262
0
      select_tx_block(cpi, x, idy, idx, block, max_tx_size, init_depth, bsize,
3263
0
                      ctxa, ctxl, tx_above, tx_left, &pn_rd_stats, INT64_MAX,
3264
0
                      best_rd_sofar, &is_cost_valid, ftxs_mode);
3265
0
      if (!is_cost_valid || pn_rd_stats.rate == INT_MAX) {
3266
0
        av1_invalid_rd_stats(rd_stats);
3267
0
        return INT64_MAX;
3268
0
      }
3269
0
      av1_merge_rd_stats(rd_stats, &pn_rd_stats);
3270
0
      skip_txfm_rd = RDCOST(x->rdmult, skip_txfm_cost, rd_stats->sse);
3271
0
      no_skip_txfm_rd =
3272
0
          RDCOST(x->rdmult, rd_stats->rate + no_skip_txfm_cost, rd_stats->dist);
3273
0
      block += step;
3274
0
    }
3275
0
  }
3276
3277
0
  if (rd_stats->rate == INT_MAX) return INT64_MAX;
3278
3279
0
  rd_stats->skip_txfm = (skip_txfm_rd <= no_skip_txfm_rd);
3280
3281
  // If fast_tx_search is true, only DCT and 1D DCT were tested in
3282
  // select_inter_block_yrd() above. Do a better search for tx type with
3283
  // tx sizes already decided.
3284
0
  if (fast_tx_search && cpi->sf.tx_sf.refine_fast_tx_search_results) {
3285
0
    if (!inter_block_yrd(cpi, x, rd_stats, bsize, ref_best_rd, FTXS_NONE))
3286
0
      return INT64_MAX;
3287
0
  }
3288
3289
0
  int64_t final_rd;
3290
0
  if (rd_stats->skip_txfm) {
3291
0
    final_rd = RDCOST(x->rdmult, skip_txfm_cost, rd_stats->sse);
3292
0
  } else {
3293
0
    final_rd =
3294
0
        RDCOST(x->rdmult, rd_stats->rate + no_skip_txfm_cost, rd_stats->dist);
3295
0
    if (!xd->lossless[xd->mi[0]->segment_id]) {
3296
0
      final_rd =
3297
0
          AOMMIN(final_rd, RDCOST(x->rdmult, skip_txfm_cost, rd_stats->sse));
3298
0
    }
3299
0
  }
3300
3301
0
  return final_rd;
3302
0
}
3303
3304
// Return 1 to terminate transform search early. The decision is made based on
3305
// the comparison with the reference RD cost and the model-estimated RD cost.
3306
static AOM_INLINE int model_based_tx_search_prune(const AV1_COMP *cpi,
3307
                                                  MACROBLOCK *x,
3308
                                                  BLOCK_SIZE bsize,
3309
0
                                                  int64_t ref_best_rd) {
3310
0
  const int level = cpi->sf.tx_sf.model_based_prune_tx_search_level;
3311
0
  assert(level >= 0 && level <= 2);
3312
0
  int model_rate;
3313
0
  int64_t model_dist;
3314
0
  int model_skip;
3315
0
  MACROBLOCKD *const xd = &x->e_mbd;
3316
0
  model_rd_sb_fn[MODELRD_TYPE_TX_SEARCH_PRUNE](
3317
0
      cpi, bsize, x, xd, 0, 0, &model_rate, &model_dist, &model_skip, NULL,
3318
0
      NULL, NULL, NULL);
3319
0
  if (model_skip) return 0;
3320
0
  const int64_t model_rd = RDCOST(x->rdmult, model_rate, model_dist);
3321
  // TODO(debargha, urvang): Improve the model and make the check below
3322
  // tighter.
3323
0
  static const int prune_factor_by8[] = { 3, 5 };
3324
0
  const int factor = prune_factor_by8[level - 1];
3325
0
  return ((model_rd * factor) >> 3) > ref_best_rd;
3326
0
}
3327
3328
void av1_pick_recursive_tx_size_type_yrd(const AV1_COMP *cpi, MACROBLOCK *x,
3329
                                         RD_STATS *rd_stats, BLOCK_SIZE bsize,
3330
0
                                         int64_t ref_best_rd) {
3331
0
  MACROBLOCKD *const xd = &x->e_mbd;
3332
0
  const TxfmSearchParams *txfm_params = &x->txfm_search_params;
3333
0
  assert(is_inter_block(xd->mi[0]));
3334
3335
0
  av1_invalid_rd_stats(rd_stats);
3336
3337
  // If modeled RD cost is a lot worse than the best so far, terminate early.
3338
0
  if (cpi->sf.tx_sf.model_based_prune_tx_search_level &&
3339
0
      ref_best_rd != INT64_MAX) {
3340
0
    if (model_based_tx_search_prune(cpi, x, bsize, ref_best_rd)) return;
3341
0
  }
3342
3343
  // Hashing based speed feature. If the hash of the prediction residue block is
3344
  // found in the hash table, use previous search results and terminate early.
3345
0
  uint32_t hash = 0;
3346
0
  MB_RD_RECORD *mb_rd_record = NULL;
3347
0
  const int mi_row = x->e_mbd.mi_row;
3348
0
  const int mi_col = x->e_mbd.mi_col;
3349
0
  const int within_border =
3350
0
      mi_row >= xd->tile.mi_row_start &&
3351
0
      (mi_row + mi_size_high[bsize] < xd->tile.mi_row_end) &&
3352
0
      mi_col >= xd->tile.mi_col_start &&
3353
0
      (mi_col + mi_size_wide[bsize] < xd->tile.mi_col_end);
3354
0
  const int is_mb_rd_hash_enabled =
3355
0
      (within_border && cpi->sf.rd_sf.use_mb_rd_hash);
3356
0
  const int n4 = bsize_to_num_blk(bsize);
3357
0
  if (is_mb_rd_hash_enabled) {
3358
0
    hash = get_block_residue_hash(x, bsize);
3359
0
    mb_rd_record = x->txfm_search_info.mb_rd_record;
3360
0
    const int match_index = find_mb_rd_info(mb_rd_record, ref_best_rd, hash);
3361
0
    if (match_index != -1) {
3362
0
      MB_RD_INFO *mb_rd_info = &mb_rd_record->mb_rd_info[match_index];
3363
0
      fetch_mb_rd_info(n4, mb_rd_info, rd_stats, x);
3364
0
      return;
3365
0
    }
3366
0
  }
3367
3368
  // If we predict that skip is the optimal RD decision - set the respective
3369
  // context and terminate early.
3370
0
  int64_t dist;
3371
0
  if (txfm_params->skip_txfm_level &&
3372
0
      predict_skip_txfm(x, bsize, &dist,
3373
0
                        cpi->common.features.reduced_tx_set_used)) {
3374
0
    set_skip_txfm(x, rd_stats, bsize, dist);
3375
    // Save the RD search results into mb_rd_record.
3376
0
    if (is_mb_rd_hash_enabled)
3377
0
      save_mb_rd_info(n4, hash, x, rd_stats, mb_rd_record);
3378
0
    return;
3379
0
  }
3380
#if CONFIG_SPEED_STATS
3381
  ++x->txfm_search_info.tx_search_count;
3382
#endif  // CONFIG_SPEED_STATS
3383
3384
0
  const int64_t rd =
3385
0
      select_tx_size_and_type(cpi, x, rd_stats, bsize, ref_best_rd);
3386
3387
0
  if (rd == INT64_MAX) {
3388
    // We should always find at least one candidate unless ref_best_rd is less
3389
    // than INT64_MAX (in which case, all the calls to select_tx_size_fix_type
3390
    // might have failed to find something better)
3391
0
    assert(ref_best_rd != INT64_MAX);
3392
0
    av1_invalid_rd_stats(rd_stats);
3393
0
    return;
3394
0
  }
3395
3396
  // Save the RD search results into mb_rd_record.
3397
0
  if (is_mb_rd_hash_enabled) {
3398
0
    assert(mb_rd_record != NULL);
3399
0
    save_mb_rd_info(n4, hash, x, rd_stats, mb_rd_record);
3400
0
  }
3401
0
}
3402
3403
void av1_pick_uniform_tx_size_type_yrd(const AV1_COMP *const cpi, MACROBLOCK *x,
3404
                                       RD_STATS *rd_stats, BLOCK_SIZE bs,
3405
3.88M
                                       int64_t ref_best_rd) {
3406
3.88M
  MACROBLOCKD *const xd = &x->e_mbd;
3407
3.88M
  MB_MODE_INFO *const mbmi = xd->mi[0];
3408
3.88M
  const TxfmSearchParams *tx_params = &x->txfm_search_params;
3409
3.88M
  assert(bs == mbmi->bsize);
3410
3.88M
  const int is_inter = is_inter_block(mbmi);
3411
3.88M
  const int mi_row = xd->mi_row;
3412
3.88M
  const int mi_col = xd->mi_col;
3413
3414
3.88M
  av1_init_rd_stats(rd_stats);
3415
3416
  // Hashing based speed feature for inter blocks. If the hash of the residue
3417
  // block is found in the table, use previously saved search results and
3418
  // terminate early.
3419
3.88M
  uint32_t hash = 0;
3420
3.88M
  MB_RD_RECORD *mb_rd_record = NULL;
3421
3.88M
  const int num_blks = bsize_to_num_blk(bs);
3422
3.88M
  if (is_inter && cpi->sf.rd_sf.use_mb_rd_hash) {
3423
0
    const int within_border =
3424
0
        mi_row >= xd->tile.mi_row_start &&
3425
0
        (mi_row + mi_size_high[bs] < xd->tile.mi_row_end) &&
3426
0
        mi_col >= xd->tile.mi_col_start &&
3427
0
        (mi_col + mi_size_wide[bs] < xd->tile.mi_col_end);
3428
0
    if (within_border) {
3429
0
      hash = get_block_residue_hash(x, bs);
3430
0
      mb_rd_record = x->txfm_search_info.mb_rd_record;
3431
0
      const int match_index = find_mb_rd_info(mb_rd_record, ref_best_rd, hash);
3432
0
      if (match_index != -1) {
3433
0
        MB_RD_INFO *mb_rd_info = &mb_rd_record->mb_rd_info[match_index];
3434
0
        fetch_mb_rd_info(num_blks, mb_rd_info, rd_stats, x);
3435
0
        return;
3436
0
      }
3437
0
    }
3438
0
  }
3439
3440
  // If we predict that skip is the optimal RD decision - set the respective
3441
  // context and terminate early.
3442
3.88M
  int64_t dist;
3443
3.88M
  if (tx_params->skip_txfm_level && is_inter &&
3444
3.88M
      !xd->lossless[mbmi->segment_id] &&
3445
3.88M
      predict_skip_txfm(x, bs, &dist,
3446
0
                        cpi->common.features.reduced_tx_set_used)) {
3447
    // Populate rdstats as per skip decision
3448
0
    set_skip_txfm(x, rd_stats, bs, dist);
3449
    // Save the RD search results into mb_rd_record.
3450
0
    if (mb_rd_record) {
3451
0
      save_mb_rd_info(num_blks, hash, x, rd_stats, mb_rd_record);
3452
0
    }
3453
0
    return;
3454
0
  }
3455
3456
3.88M
  if (xd->lossless[mbmi->segment_id]) {
3457
    // Lossless mode can only pick the smallest (4x4) transform size.
3458
1.16M
    choose_smallest_tx_size(cpi, x, rd_stats, ref_best_rd, bs);
3459
2.72M
  } else if (tx_params->tx_size_search_method == USE_LARGESTALL) {
3460
2.62M
    choose_largest_tx_size(cpi, x, rd_stats, ref_best_rd, bs);
3461
2.62M
  } else {
3462
95.2k
    choose_tx_size_type_from_rd(cpi, x, rd_stats, ref_best_rd, bs);
3463
95.2k
  }
3464
3465
  // Save the RD search results into mb_rd_record for possible reuse in future.
3466
3.88M
  if (mb_rd_record) {
3467
0
    save_mb_rd_info(num_blks, hash, x, rd_stats, mb_rd_record);
3468
0
  }
3469
3.88M
}
3470
3471
int av1_txfm_uvrd(const AV1_COMP *const cpi, MACROBLOCK *x, RD_STATS *rd_stats,
3472
130k
                  BLOCK_SIZE bsize, int64_t ref_best_rd) {
3473
130k
  av1_init_rd_stats(rd_stats);
3474
130k
  if (ref_best_rd < 0) return 0;
3475
130k
  if (!x->e_mbd.is_chroma_ref) return 1;
3476
3477
130k
  MACROBLOCKD *const xd = &x->e_mbd;
3478
130k
  MB_MODE_INFO *const mbmi = xd->mi[0];
3479
130k
  struct macroblockd_plane *const pd = &xd->plane[AOM_PLANE_U];
3480
130k
  const int is_inter = is_inter_block(mbmi);
3481
130k
  int64_t this_rd = 0, skip_txfm_rd = 0;
3482
130k
  const BLOCK_SIZE plane_bsize =
3483
130k
      get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);
3484
3485
130k
  if (is_inter) {
3486
0
    for (int plane = 1; plane < MAX_MB_PLANE; ++plane)
3487
0
      av1_subtract_plane(x, plane_bsize, plane);
3488
0
  }
3489
3490
130k
  const int skip_trellis = 0;
3491
130k
  const TX_SIZE uv_tx_size = av1_get_tx_size(AOM_PLANE_U, xd);
3492
130k
  int is_cost_valid = 1;
3493
377k
  for (int plane = 1; plane < MAX_MB_PLANE; ++plane) {
3494
255k
    RD_STATS this_rd_stats;
3495
255k
    int64_t chroma_ref_best_rd = ref_best_rd;
3496
    // For inter blocks, refined ref_best_rd is used for early exit
3497
    // For intra blocks, even though current rd crosses ref_best_rd, early
3498
    // exit is not recommended as current rd is used for gating subsequent
3499
    // modes as well (say, for angular modes)
3500
    // TODO(any): Extend the early exit mechanism for intra modes as well
3501
255k
    if (cpi->sf.inter_sf.perform_best_rd_based_gating_for_chroma && is_inter &&
3502
255k
        chroma_ref_best_rd != INT64_MAX)
3503
0
      chroma_ref_best_rd = ref_best_rd - AOMMIN(this_rd, skip_txfm_rd);
3504
255k
    av1_txfm_rd_in_plane(x, cpi, &this_rd_stats, chroma_ref_best_rd, 0, plane,
3505
255k
                         plane_bsize, uv_tx_size, FTXS_NONE, skip_trellis);
3506
255k
    if (this_rd_stats.rate == INT_MAX) {
3507
5.49k
      is_cost_valid = 0;
3508
5.49k
      break;
3509
5.49k
    }
3510
249k
    av1_merge_rd_stats(rd_stats, &this_rd_stats);
3511
249k
    this_rd = RDCOST(x->rdmult, rd_stats->rate, rd_stats->dist);
3512
249k
    skip_txfm_rd = RDCOST(x->rdmult, 0, rd_stats->sse);
3513
249k
    if (AOMMIN(this_rd, skip_txfm_rd) > ref_best_rd) {
3514
2.57k
      is_cost_valid = 0;
3515
2.57k
      break;
3516
2.57k
    }
3517
249k
  }
3518
3519
130k
  if (!is_cost_valid) {
3520
    // reset cost value
3521
8.06k
    av1_invalid_rd_stats(rd_stats);
3522
8.06k
  }
3523
3524
130k
  return is_cost_valid;
3525
130k
}
3526
3527
void av1_txfm_rd_in_plane(MACROBLOCK *x, const AV1_COMP *cpi,
3528
                          RD_STATS *rd_stats, int64_t ref_best_rd,
3529
                          int64_t current_rd, int plane, BLOCK_SIZE plane_bsize,
3530
                          TX_SIZE tx_size, FAST_TX_SEARCH_MODE ftxs_mode,
3531
4.61M
                          int skip_trellis) {
3532
4.61M
  assert(IMPLIES(plane == 0, x->e_mbd.mi[0]->tx_size == tx_size));
3533
3534
4.61M
  if (!cpi->oxcf.txfm_cfg.enable_tx64 &&
3535
4.61M
      txsize_sqr_up_map[tx_size] == TX_64X64) {
3536
0
    av1_invalid_rd_stats(rd_stats);
3537
0
    return;
3538
0
  }
3539
3540
4.61M
  if (current_rd > ref_best_rd) {
3541
2.77k
    av1_invalid_rd_stats(rd_stats);
3542
2.77k
    return;
3543
2.77k
  }
3544
3545
4.60M
  MACROBLOCKD *const xd = &x->e_mbd;
3546
4.60M
  const struct macroblockd_plane *const pd = &xd->plane[plane];
3547
4.60M
  struct rdcost_block_args args;
3548
4.60M
  av1_zero(args);
3549
4.60M
  args.x = x;
3550
4.60M
  args.cpi = cpi;
3551
4.60M
  args.best_rd = ref_best_rd;
3552
4.60M
  args.current_rd = current_rd;
3553
4.60M
  args.ftxs_mode = ftxs_mode;
3554
4.60M
  args.skip_trellis = skip_trellis;
3555
4.60M
  av1_init_rd_stats(&args.rd_stats);
3556
3557
4.60M
  av1_get_entropy_contexts(plane_bsize, pd, args.t_above, args.t_left);
3558
4.60M
  av1_foreach_transformed_block_in_plane(xd, plane_bsize, plane, block_rd_txfm,
3559
4.60M
                                         &args);
3560
3561
4.60M
  MB_MODE_INFO *const mbmi = xd->mi[0];
3562
4.60M
  const int is_inter = is_inter_block(mbmi);
3563
4.60M
  const int invalid_rd = is_inter ? args.incomplete_exit : args.exit_early;
3564
3565
4.60M
  if (invalid_rd) {
3566
730k
    av1_invalid_rd_stats(rd_stats);
3567
3.87M
  } else {
3568
3.87M
    *rd_stats = args.rd_stats;
3569
3.87M
  }
3570
4.60M
}
3571
3572
int av1_txfm_search(const AV1_COMP *cpi, MACROBLOCK *x, BLOCK_SIZE bsize,
3573
                    RD_STATS *rd_stats, RD_STATS *rd_stats_y,
3574
0
                    RD_STATS *rd_stats_uv, int mode_rate, int64_t ref_best_rd) {
3575
0
  MACROBLOCKD *const xd = &x->e_mbd;
3576
0
  TxfmSearchParams *txfm_params = &x->txfm_search_params;
3577
0
  const int skip_ctx = av1_get_skip_txfm_context(xd);
3578
0
  const int skip_txfm_cost[2] = { x->mode_costs.skip_txfm_cost[skip_ctx][0],
3579
0
                                  x->mode_costs.skip_txfm_cost[skip_ctx][1] };
3580
0
  const int64_t min_header_rate =
3581
0
      mode_rate + AOMMIN(skip_txfm_cost[0], skip_txfm_cost[1]);
3582
  // Account for minimum skip and non_skip rd.
3583
  // Eventually either one of them will be added to mode_rate
3584
0
  const int64_t min_header_rd_possible = RDCOST(x->rdmult, min_header_rate, 0);
3585
0
  if (min_header_rd_possible > ref_best_rd) {
3586
0
    av1_invalid_rd_stats(rd_stats_y);
3587
0
    return 0;
3588
0
  }
3589
3590
0
  const AV1_COMMON *cm = &cpi->common;
3591
0
  MB_MODE_INFO *const mbmi = xd->mi[0];
3592
0
  const int64_t mode_rd = RDCOST(x->rdmult, mode_rate, 0);
3593
0
  const int64_t rd_thresh =
3594
0
      ref_best_rd == INT64_MAX ? INT64_MAX : ref_best_rd - mode_rd;
3595
0
  av1_init_rd_stats(rd_stats);
3596
0
  av1_init_rd_stats(rd_stats_y);
3597
0
  rd_stats->rate = mode_rate;
3598
3599
  // cost and distortion
3600
0
  av1_subtract_plane(x, bsize, 0);
3601
0
  if (txfm_params->tx_mode_search_type == TX_MODE_SELECT &&
3602
0
      !xd->lossless[mbmi->segment_id]) {
3603
0
    av1_pick_recursive_tx_size_type_yrd(cpi, x, rd_stats_y, bsize, rd_thresh);
3604
#if CONFIG_COLLECT_RD_STATS == 2
3605
    PrintPredictionUnitStats(cpi, tile_data, x, rd_stats_y, bsize);
3606
#endif  // CONFIG_COLLECT_RD_STATS == 2
3607
0
  } else {
3608
0
    av1_pick_uniform_tx_size_type_yrd(cpi, x, rd_stats_y, bsize, rd_thresh);
3609
0
    memset(mbmi->inter_tx_size, mbmi->tx_size, sizeof(mbmi->inter_tx_size));
3610
0
    for (int i = 0; i < xd->height * xd->width; ++i)
3611
0
      set_blk_skip(x->txfm_search_info.blk_skip, 0, i, rd_stats_y->skip_txfm);
3612
0
  }
3613
3614
0
  if (rd_stats_y->rate == INT_MAX) return 0;
3615
3616
0
  av1_merge_rd_stats(rd_stats, rd_stats_y);
3617
3618
0
  const int64_t non_skip_txfm_rdcosty =
3619
0
      RDCOST(x->rdmult, rd_stats->rate + skip_txfm_cost[0], rd_stats->dist);
3620
0
  const int64_t skip_txfm_rdcosty =
3621
0
      RDCOST(x->rdmult, mode_rate + skip_txfm_cost[1], rd_stats->sse);
3622
0
  const int64_t min_rdcosty = AOMMIN(non_skip_txfm_rdcosty, skip_txfm_rdcosty);
3623
0
  if (min_rdcosty > ref_best_rd) return 0;
3624
3625
0
  av1_init_rd_stats(rd_stats_uv);
3626
0
  const int num_planes = av1_num_planes(cm);
3627
0
  if (num_planes > 1) {
3628
0
    int64_t ref_best_chroma_rd = ref_best_rd;
3629
    // Calculate best rd cost possible for chroma
3630
0
    if (cpi->sf.inter_sf.perform_best_rd_based_gating_for_chroma &&
3631
0
        (ref_best_chroma_rd != INT64_MAX)) {
3632
0
      ref_best_chroma_rd = (ref_best_chroma_rd -
3633
0
                            AOMMIN(non_skip_txfm_rdcosty, skip_txfm_rdcosty));
3634
0
    }
3635
0
    const int is_cost_valid_uv =
3636
0
        av1_txfm_uvrd(cpi, x, rd_stats_uv, bsize, ref_best_chroma_rd);
3637
0
    if (!is_cost_valid_uv) return 0;
3638
0
    av1_merge_rd_stats(rd_stats, rd_stats_uv);
3639
0
  }
3640
3641
0
  int choose_skip_txfm = rd_stats->skip_txfm;
3642
0
  if (!choose_skip_txfm && !xd->lossless[mbmi->segment_id]) {
3643
0
    const int64_t rdcost_no_skip_txfm = RDCOST(
3644
0
        x->rdmult, rd_stats_y->rate + rd_stats_uv->rate + skip_txfm_cost[0],
3645
0
        rd_stats->dist);
3646
0
    const int64_t rdcost_skip_txfm =
3647
0
        RDCOST(x->rdmult, skip_txfm_cost[1], rd_stats->sse);
3648
0
    if (rdcost_no_skip_txfm >= rdcost_skip_txfm) choose_skip_txfm = 1;
3649
0
  }
3650
0
  if (choose_skip_txfm) {
3651
0
    rd_stats_y->rate = 0;
3652
0
    rd_stats_uv->rate = 0;
3653
0
    rd_stats->rate = mode_rate + skip_txfm_cost[1];
3654
0
    rd_stats->dist = rd_stats->sse;
3655
0
    rd_stats_y->dist = rd_stats_y->sse;
3656
0
    rd_stats_uv->dist = rd_stats_uv->sse;
3657
0
    mbmi->skip_txfm = 1;
3658
0
    if (rd_stats->skip_txfm) {
3659
0
      const int64_t tmprd = RDCOST(x->rdmult, rd_stats->rate, rd_stats->dist);
3660
0
      if (tmprd > ref_best_rd) return 0;
3661
0
    }
3662
0
  } else {
3663
0
    rd_stats->rate += skip_txfm_cost[0];
3664
0
    mbmi->skip_txfm = 0;
3665
0
  }
3666
3667
0
  return 1;
3668
0
}