| /* |
| * Copyright (c) 2021, Alliance for Open Media. All rights reserved |
| * |
| * This source code is subject to the terms of the BSD 3-Clause Clear License |
| * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear |
| * License was not distributed with this source code in the LICENSE file, you |
| * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/. If the |
| * Alliance for Open Media Patent License 1.0 was not distributed with this |
| * source code in the PATENTS file, you can obtain it at |
| * aomedia.org/license/patent-license/. |
| */ |
| |
| #ifndef AOM_AV1_ENCODER_TRANSFORM_SEARCH_H_ |
| #define AOM_AV1_ENCODER_TRANSFORM_SEARCH_H_ |
| |
| #include "av1/common/pred_common.h" |
| #include "av1/encoder/encoder.h" |
| |
| #ifdef __cplusplus |
| extern "C" { |
| #endif |
| |
| // Set this macro as 1 to collect data about tx size selection. |
| #define COLLECT_TX_SIZE_DATA 0 |
| |
| #if COLLECT_TX_SIZE_DATA |
| static const char av1_tx_size_data_output_file[] = "tx_size_data.txt"; |
| #endif |
| |
| enum { |
| FTXS_NONE = 0, |
| FTXS_DCT_AND_1D_DCT_ONLY = 1 << 0, |
| FTXS_DISABLE_TRELLIS_OPT = 1 << 1, |
| FTXS_USE_TRANSFORM_DOMAIN = 1 << 2 |
| } UENUM1BYTE(FAST_TX_SEARCH_MODE); |
| |
| #if CONFIG_NEW_TX_PARTITION |
| #if CONFIG_TX_PARTITION_CTX |
| static AOM_INLINE int inter_tx_partition_cost(const MACROBLOCK *const x, |
| TX_PARTITION_TYPE partition, |
| BLOCK_SIZE bsize, |
| TX_SIZE max_tx_size) { |
| #else |
| static AOM_INLINE int inter_tx_partition_cost( |
| const MACROBLOCK *const x, int is_rect, TX_PARTITION_TYPE partition, |
| const TXFM_CONTEXT *const above_ctx, const TXFM_CONTEXT *const left_ctx, |
| BLOCK_SIZE bsize, TX_SIZE max_tx_size) { |
| #endif // CONFIG_TX_PARTITION_CTX |
| int cost = 0; |
| const int allow_horz = allow_tx_horz_split(max_tx_size); |
| const int allow_vert = allow_tx_vert_split(max_tx_size); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| const MACROBLOCKD *const xd = &x->e_mbd; |
| const MB_MODE_INFO *const mbmi = xd->mi[0]; |
| const int is_fsc = (mbmi->fsc_mode[xd->tree_type == CHROMA_PART]); |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| #if CONFIG_TX_PARTITION_CTX |
| #if CONFIG_TX_PARTITION_TYPE_EXT |
| const int bsize_group = size_to_tx_part_group_lookup[bsize]; |
| const int txsize_group = size_to_tx_type_group_lookup[bsize]; |
| int do_partition = 0; |
| if (allow_horz || allow_vert) { |
| do_partition = (partition != TX_PARTITION_NONE); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_do_partition_cost[is_fsc][1][bsize_group][do_partition]; |
| #else |
| cost += x->mode_costs.txfm_do_partition_cost[1][bsize_group][do_partition]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } |
| |
| if (do_partition) { |
| if (allow_horz && allow_vert) { |
| assert(txsize_group > 0); |
| const TX_PARTITION_TYPE split4_partition = |
| get_split4_partition(partition); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_4way_partition_type_cost[is_fsc][1][txsize_group - 1] |
| [split4_partition - 1]; |
| #else |
| cost += x->mode_costs.txfm_4way_partition_type_cost[1][txsize_group - 1] |
| [split4_partition - 1]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } else if (allow_horz || allow_vert) { |
| int has_first_split = 0; |
| if (partition == TX_PARTITION_VERT_M || partition == TX_PARTITION_HORZ_M) |
| has_first_split = 1; |
| |
| if (txsize_group) { |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_4way_partition_type_cost[is_fsc][1][txsize_group - 1] |
| [has_first_split]; |
| #else |
| cost += x->mode_costs.txfm_4way_partition_type_cost[1][txsize_group - 1] |
| [has_first_split]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } |
| } |
| } |
| #else |
| const int bsize_group = size_to_tx_part_group_lookup[bsize]; |
| int do_partition = 0; |
| if (allow_horz || allow_vert) { |
| do_partition = (partition != TX_PARTITION_NONE); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_do_partition_cost[is_fsc][1][bsize_group][do_partition]; |
| #else |
| cost += x->mode_costs.txfm_do_partition_cost[1][bsize_group][do_partition]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } |
| |
| if (do_partition) { |
| if (allow_horz && allow_vert) { |
| assert(bsize_group > 0); |
| const TX_PARTITION_TYPE split4_partition = |
| get_split4_partition(partition); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_4way_partition_type_cost[is_fsc][1][bsize_group - 1] |
| [split4_partition - 1]; |
| #else |
| cost += x->mode_costs.txfm_4way_partition_type_cost[1][bsize_group - 1] |
| [split4_partition - 1]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } |
| } |
| #endif // CONFIG_TX_PARTITION_TYPE_EXT |
| #else |
| if (allow_horz && allow_vert) { |
| const int split4_ctx_0 = txfm_partition_split4_inter_context( |
| above_ctx, left_ctx, bsize, max_tx_size); |
| const TX_PARTITION_TYPE split4_partition = get_split4_partition(partition); |
| cost += x->mode_costs.inter_4way_txfm_partition_cost[is_rect][split4_ctx_0] |
| [split4_partition]; |
| } else if (allow_horz || allow_vert) { |
| const int has_first_split = partition != TX_PARTITION_NONE; |
| cost += x->mode_costs.inter_2way_txfm_partition_cost[has_first_split]; |
| } else { |
| assert(!allow_horz && !allow_vert); |
| assert(partition == PARTITION_NONE); |
| } |
| #endif // CONFIG_TX_PARTITION_CTX |
| return cost; |
| } |
| |
| static AOM_INLINE int intra_tx_partition_cost(const MACROBLOCK *const x, |
| #if !CONFIG_TX_PARTITION_CTX |
| int is_rect, |
| #endif // !CONFIG_TX_PARTITION_CTX |
| TX_PARTITION_TYPE partition, |
| TX_SIZE max_tx_size) { |
| int cost = 0; |
| const MACROBLOCKD *const xd = &x->e_mbd; |
| const int allow_horz = allow_tx_horz_split(max_tx_size); |
| const int allow_vert = allow_tx_vert_split(max_tx_size); |
| #if CONFIG_TX_PARTITION_CTX |
| #if CONFIG_TX_PARTITION_TYPE_EXT |
| const MB_MODE_INFO *const mbmi = xd->mi[0]; |
| const BLOCK_SIZE bsize = mbmi->sb_type[PLANE_TYPE_Y]; |
| const int bsize_group = size_to_tx_part_group_lookup[bsize]; |
| const int txsize_group = size_to_tx_type_group_lookup[bsize]; |
| #if CONFIG_IMPROVEIDTX_CTXS |
| const int is_fsc = (mbmi->fsc_mode[xd->tree_type == CHROMA_PART]); |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| int do_partition = 0; |
| if (allow_horz || allow_vert) { |
| do_partition = (partition != TX_PARTITION_NONE); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_do_partition_cost[is_fsc][0][bsize_group][do_partition]; |
| #else |
| cost += x->mode_costs.txfm_do_partition_cost[0][bsize_group][do_partition]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } |
| |
| if (do_partition) { |
| if (allow_horz && allow_vert) { |
| assert(txsize_group > 0); |
| const TX_PARTITION_TYPE split4_partition = |
| get_split4_partition(partition); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_4way_partition_type_cost[is_fsc][0][txsize_group - 1] |
| [split4_partition - 1]; |
| #else |
| cost += x->mode_costs.txfm_4way_partition_type_cost[0][txsize_group - 1] |
| [split4_partition - 1]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } else if (allow_horz || allow_vert) { |
| int has_first_split = 0; |
| if (partition == TX_PARTITION_VERT_M || partition == TX_PARTITION_HORZ_M) |
| has_first_split = 1; |
| |
| if (txsize_group) { |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_4way_partition_type_cost[is_fsc][0][txsize_group - 1] |
| [has_first_split]; |
| #else |
| cost += x->mode_costs.txfm_4way_partition_type_cost[0][txsize_group - 1] |
| [has_first_split]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } |
| } |
| } |
| #else |
| const MB_MODE_INFO *const mbmi = xd->mi[0]; |
| const BLOCK_SIZE bsize = mbmi->sb_type[PLANE_TYPE_Y]; |
| const int bsize_group = size_to_tx_part_group_lookup[bsize]; |
| #if CONFIG_IMPROVEIDTX_CTXS |
| const int is_fsc = (mbmi->fsc_mode[xd->tree_type == CHROMA_PART]); |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| int do_partition = 0; |
| if (allow_horz || allow_vert) { |
| do_partition = (partition != TX_PARTITION_NONE); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_do_partition_cost[is_fsc][0][bsize_group][do_partition]; |
| #else |
| cost += x->mode_costs.txfm_do_partition_cost[0][bsize_group][do_partition]; |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| } |
| |
| if (do_partition) { |
| if (allow_horz && allow_vert) { |
| assert(bsize_group > 0); |
| const TX_PARTITION_TYPE split4_partition = |
| get_split4_partition(partition); |
| #if CONFIG_IMPROVEIDTX_CTXS |
| cost += x->mode_costs |
| .txfm_4way_partition_type_cost[is_fsc][0][bsize_group - 1] |
| #else |
| cost += x->mode_costs.txfm_4way_partition_type_cost[0][bsize_group - 1] |
| #endif // CONFIG_IMPROVEIDTX_CTXS |
| [split4_partition - 1]; |
| } |
| } |
| #endif // CONFIG_TX_PARTITION_TYPE_EXT |
| #else |
| if (allow_horz && allow_vert) { |
| const int split4_ctx_0 = get_tx_size_context(xd); |
| const TX_PARTITION_TYPE split4_partition = get_split4_partition(partition); |
| cost += x->mode_costs.intra_4way_txfm_partition_cost[is_rect][split4_ctx_0] |
| [split4_partition]; |
| } else if (allow_horz || allow_vert) { |
| const int has_first_split = partition != TX_PARTITION_NONE; |
| cost += x->mode_costs.intra_2way_txfm_partition_cost[has_first_split]; |
| } else { |
| assert(!allow_horz && !allow_vert); |
| assert(partition == PARTITION_NONE); |
| } |
| #endif // CONFIG_TX_PARTITION_CTX |
| return cost; |
| } |
| #endif // CONFIG_NEW_TX_PARTITION |
| |
| // Calculate the energy of 32-bit signed data in 1d buffer |
| // TODO(any): Implement x86 SIMD for the following function |
| static AOM_INLINE uint64_t sum_squares_i32(const int32_t *src, int32_t n) { |
| uint64_t ss = 0; |
| do { |
| const int32_t v = *src++; |
| ss += ((int64_t)v) * v; |
| } while (--n); |
| |
| return ss; |
| } |
| |
| static AOM_INLINE int tx_size_cost(const MACROBLOCK *const x, BLOCK_SIZE bsize, |
| TX_SIZE tx_size) { |
| assert(bsize == x->e_mbd.mi[0]->sb_type[PLANE_TYPE_Y]); |
| if (x->txfm_search_params.tx_mode_search_type != TX_MODE_SELECT || |
| !block_signals_txsize(bsize)) |
| return 0; |
| |
| const MACROBLOCKD *const xd = &x->e_mbd; |
| #if CONFIG_NEW_TX_PARTITION |
| if (bsize >= BLOCK_SIZES_ALL) return INT_MAX; |
| |
| (void)tx_size; |
| MB_MODE_INFO *const mbmi = xd->mi[0]; |
| const TX_SIZE max_tx_size = max_txsize_rect_lookup[bsize]; |
| #if CONFIG_TX_PARTITION_CTX |
| return intra_tx_partition_cost(x, mbmi->tx_partition_type[0], max_tx_size); |
| #else |
| const int is_rect = is_rect_tx(max_tx_size); |
| return intra_tx_partition_cost(x, is_rect, mbmi->tx_partition_type[0], |
| max_tx_size); |
| #endif // CONFIG_TX_PARTITION_CTX |
| #else |
| const int32_t tx_size_cat = bsize_to_tx_size_cat(bsize); |
| const int depth = tx_size_to_depth(tx_size, bsize); |
| const int tx_size_ctx = get_tx_size_context(xd); |
| return x->mode_costs.tx_size_cost[tx_size_cat][tx_size_ctx][depth]; |
| #endif // CONFIG_NEW_TX_PARTITION |
| } |
| |
| static AOM_INLINE int skip_cctx_eval_based_on_eob(int plane, int is_inter, |
| uint16_t eob_c1, |
| CctxType cctx_type) { |
| if (plane != AOM_PLANE_U) return 0; |
| if (eob_c1 == 0) return 1; |
| if (eob_c1 == 1 && !is_inter && cctx_type != CCTX_NONE) return 1; |
| return 0; |
| } |
| |
| /*!\brief Transform type search for luma macroblock with fixed transform size. |
| * |
| * \ingroup transform_search |
| * Search for the best transform type and return the transform coefficients RD |
| * cost of current luma macroblock with the given uniform transform size. |
| * |
| * \param[in] x Pointer to structure holding the data for the |
| current encoding macroblock |
| * \param[in] cpi Top-level encoder structure |
| * \param[in] rd_stats Pointer to struct to keep track of the RD stats |
| * \param[in] ref_best_rd Best RD cost seen for this block so far |
| * \param[in] bs Size of the current macroblock |
| * \param[in] tx_size The given transform size |
| * \param[in] ftxs_mode Transform search mode specifying desired speed |
| and quality tradeoff |
| * \param[in] skip_trellis Binary flag indicating if trellis optimization |
| should be skipped |
| * \return An int64_t value that is the best RD cost found. |
| */ |
| int64_t av1_uniform_txfm_yrd(const AV1_COMP *const cpi, MACROBLOCK *x, |
| RD_STATS *rd_stats, int64_t ref_best_rd, |
| BLOCK_SIZE bs, TX_SIZE tx_size, |
| FAST_TX_SEARCH_MODE ftxs_mode, int skip_trellis); |
| |
| /*!\brief Recursive transform size and type search. |
| * |
| * \ingroup transform_search |
| * Search for best transform size and type for luma inter blocks. The transform |
| * block partitioning can be recursive resulting in non-uniform transform sizes. |
| * The best transform size and type, if found, will be saved in the MB_MODE_INFO |
| * structure, and the corresponding RD stats will be saved in rd_stats. |
| * |
| * \param[in] cpi Top-level encoder structure |
| * \param[in] x Pointer to structure holding the data for the |
| current encoding macroblock |
| * \param[in] rd_stats Pointer to struct to keep track of the RD stats |
| * \param[in] bsize Current macroblock size |
| * \param[in] ref_best_rd Best RD cost seen for this block so far |
| * Nothing is returned. The selected transform size and type will be saved |
| * in the MB_MODE_INFO structure. |
| */ |
| void av1_pick_recursive_tx_size_type_yrd(const AV1_COMP *cpi, MACROBLOCK *x, |
| RD_STATS *rd_stats, BLOCK_SIZE bsize, |
| int64_t ref_best_rd); |
| |
| /*!\brief Uniform transform size and type search. |
| * |
| * \ingroup transform_search |
| * Search for the best transform size and type for current macroblock block, |
| * with the assumption that all the transform blocks have a uniform size |
| * (VP9 style). The selected transform size and type will be saved in the |
| * MB_MODE_INFO structure; the corresponding RD stats will be saved in rd_stats. |
| * This function may be used for both intra and inter predicted blocks. |
| * |
| * \param[in] cpi Top-level encoder structure |
| * \param[in] x Pointer to structure holding the data for the |
| current encoding macroblock |
| * \param[in] rd_stats Pointer to struct to keep track of the RD stats |
| * \param[in] bs Current macroblock size |
| * \param[in] ref_best_rd Best RD cost seen for this block so far |
| * Nothing is returned. The selected transform size and type will be saved |
| * in the MB_MODE_INFO structure. |
| */ |
| void av1_pick_uniform_tx_size_type_yrd(const AV1_COMP *const cpi, MACROBLOCK *x, |
| RD_STATS *rd_stats, BLOCK_SIZE bs, |
| int64_t ref_best_rd); |
| |
| /*!\brief Chroma block transform search. |
| * |
| * \ingroup transform_search |
| * Calculate the transform coefficient RD cost for the given chroma macroblock |
| * If the current mode is intra, then this function will compute the predictor. |
| * |
| * \param[in] cpi Top-level encoder structure |
| * \param[in] x Pointer to structure holding the data for the |
| current encoding macroblock |
| * \param[in] rd_stats Pointer to struct to keep track of the RD stats |
| * \param[in] ref_best_rd Best RD cost seen for this block so far |
| * \return An integer value is returned. 0: early termination triggered, |
| no valid rd cost available; 1: rd cost values are valid. |
| */ |
| int av1_txfm_uvrd(const AV1_COMP *const cpi, MACROBLOCK *x, RD_STATS *rd_stats, |
| int64_t ref_best_rd); |
| |
| /*!\brief Transform type search with fixed transform size. |
| * |
| * \ingroup transform_search |
| * Search for the best transform type and calculate the transform coefficients |
| * RD cost of the current transform block with the specified (uniform) transform |
| * size and plane. The RD results will be saved in rd_stats. |
| * |
| * \param[in] x Pointer to structure holding the data for the |
| current encoding macroblock |
| * \param[in] cpi Top-level encoder structure |
| * \param[in] rd_stats Pointer to struct to keep track of the RD stats |
| * \param[in] ref_best_rd Best RD cost seen for this block so far |
| * \param[in] current_rd Current RD cost for this block so far |
| * \param[in] plane Plane index |
| * \param[in] plane_bsize Size of the current macroblock considering |
| sup-sampling |
| * \param[in] tx_size The given transform size |
| * \param[in] ftxs_mode Transform search mode specifying desired speed |
| and quality tradeoff |
| * \param[in] skip_trellis Binary flag indicating if trellis optimization |
| should be skipped |
| * |
| * Nothing is returned. The RD results will be saved in rd_stats. |
| */ |
| void av1_txfm_rd_in_plane(MACROBLOCK *x, const AV1_COMP *cpi, |
| RD_STATS *rd_stats, int64_t ref_best_rd, |
| int64_t current_rd, int plane, BLOCK_SIZE plane_bsize, |
| TX_SIZE tx_size, FAST_TX_SEARCH_MODE ftxs_mode, |
| int skip_trellis); |
| |
| /*!\brief Recursive transform size and type search. |
| * |
| * \ingroup transform_search |
| * This function combines y and uv planes' transform search processes together |
| * for inter-predicted blocks (including IntraBC), when the prediction is |
| * already generated. It first does subtraction to obtain the prediction error. |
| * Then it calls |
| * av1_pick_recursive_tx_size_type_yrd/av1_pick_uniform_tx_size_type_yrd and |
| * av1_txfm_uvrd sequentially and handles possible early terminations. |
| * The RD metrics are calculated and stored in rd_stats/_y/_uv. |
| * |
| * \param[in] cpi Top-level encoder structure |
| * \param[in] x Pointer to structure holding the data for the |
| current encoding macroblock |
| * \param[in] bsize Current macroblock size |
| * \param[in] rd_stats Pointer to struct to keep track of the overal RD |
| stats |
| * \param[in] rd_stats_y Pointer to struct to keep track of the RD |
| stats for the luma plane |
| * \param[in] rd_stats_uv Pointer to struct to keep track of the RD |
| stats for the chroma planes |
| * \param[in] mode_rate Rate cost to encode the prediction mode info. of |
| the current macroblock |
| * \param[in] ref_best_rd Best RD cost seen for this block so far |
| * |
| * \return An integer value is returned indicating if a valid transform |
| candidate is found (1) or not (0). |
| */ |
| int av1_txfm_search(const AV1_COMP *cpi, MACROBLOCK *x, BLOCK_SIZE bsize, |
| RD_STATS *rd_stats, RD_STATS *rd_stats_y, |
| RD_STATS *rd_stats_uv, int mode_rate, int64_t ref_best_rd); |
| |
| #ifdef __cplusplus |
| } // extern "C" |
| #endif |
| |
| #endif // AOM_AV1_ENCODER_TRANSFORM_SEARCH_H_ |
