xref: /aosp_15_r20/external/libaom/av1/encoder/context_tree.c (revision 77c1e3ccc04c968bd2bc212e87364f250e820521)
1 /*
2  * Copyright (c) 2016, 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/encoder/context_tree.h"
13 #include "av1/encoder/encoder.h"
14 #include "av1/encoder/rd.h"
15 #include <assert.h>
16 
av1_copy_tree_context(PICK_MODE_CONTEXT * dst_ctx,PICK_MODE_CONTEXT * src_ctx)17 void av1_copy_tree_context(PICK_MODE_CONTEXT *dst_ctx,
18                            PICK_MODE_CONTEXT *src_ctx) {
19   dst_ctx->mic = src_ctx->mic;
20   dst_ctx->mbmi_ext_best = src_ctx->mbmi_ext_best;
21 
22   dst_ctx->num_4x4_blk = src_ctx->num_4x4_blk;
23   dst_ctx->skippable = src_ctx->skippable;
24 #if CONFIG_INTERNAL_STATS
25   dst_ctx->best_mode_index = src_ctx->best_mode_index;
26 #endif  // CONFIG_INTERNAL_STATS
27 
28   memcpy(dst_ctx->blk_skip, src_ctx->blk_skip,
29          sizeof(uint8_t) * src_ctx->num_4x4_blk);
30   av1_copy_array(dst_ctx->tx_type_map, src_ctx->tx_type_map,
31                  src_ctx->num_4x4_blk);
32 
33   dst_ctx->rd_stats = src_ctx->rd_stats;
34   dst_ctx->rd_mode_is_ready = src_ctx->rd_mode_is_ready;
35 }
36 
av1_setup_shared_coeff_buffer(const SequenceHeader * const seq_params,PC_TREE_SHARED_BUFFERS * shared_bufs,struct aom_internal_error_info * error)37 void av1_setup_shared_coeff_buffer(const SequenceHeader *const seq_params,
38                                    PC_TREE_SHARED_BUFFERS *shared_bufs,
39                                    struct aom_internal_error_info *error) {
40   const int num_planes = seq_params->monochrome ? 1 : MAX_MB_PLANE;
41   const int max_sb_square_y = 1 << num_pels_log2_lookup[seq_params->sb_size];
42   const int max_sb_square_uv = max_sb_square_y >> (seq_params->subsampling_x +
43                                                    seq_params->subsampling_y);
44   for (int i = 0; i < num_planes; i++) {
45     const int max_num_pix =
46         (i == AOM_PLANE_Y) ? max_sb_square_y : max_sb_square_uv;
47     AOM_CHECK_MEM_ERROR(error, shared_bufs->coeff_buf[i],
48                         aom_memalign(32, max_num_pix * sizeof(tran_low_t)));
49     AOM_CHECK_MEM_ERROR(error, shared_bufs->qcoeff_buf[i],
50                         aom_memalign(32, max_num_pix * sizeof(tran_low_t)));
51     AOM_CHECK_MEM_ERROR(error, shared_bufs->dqcoeff_buf[i],
52                         aom_memalign(32, max_num_pix * sizeof(tran_low_t)));
53   }
54 }
55 
av1_free_shared_coeff_buffer(PC_TREE_SHARED_BUFFERS * shared_bufs)56 void av1_free_shared_coeff_buffer(PC_TREE_SHARED_BUFFERS *shared_bufs) {
57   for (int i = 0; i < 3; i++) {
58     aom_free(shared_bufs->coeff_buf[i]);
59     aom_free(shared_bufs->qcoeff_buf[i]);
60     aom_free(shared_bufs->dqcoeff_buf[i]);
61     shared_bufs->coeff_buf[i] = NULL;
62     shared_bufs->qcoeff_buf[i] = NULL;
63     shared_bufs->dqcoeff_buf[i] = NULL;
64   }
65 }
66 
av1_alloc_pmc(const struct AV1_COMP * const cpi,BLOCK_SIZE bsize,PC_TREE_SHARED_BUFFERS * shared_bufs)67 PICK_MODE_CONTEXT *av1_alloc_pmc(const struct AV1_COMP *const cpi,
68                                  BLOCK_SIZE bsize,
69                                  PC_TREE_SHARED_BUFFERS *shared_bufs) {
70   PICK_MODE_CONTEXT *volatile ctx = NULL;
71   const AV1_COMMON *const cm = &cpi->common;
72   struct aom_internal_error_info error;
73 
74   if (setjmp(error.jmp)) {
75     av1_free_pmc(ctx, av1_num_planes(cm));
76     return NULL;
77   }
78   error.setjmp = 1;
79 
80   AOM_CHECK_MEM_ERROR(&error, ctx, aom_calloc(1, sizeof(*ctx)));
81   ctx->rd_mode_is_ready = 0;
82 
83   const int num_planes = av1_num_planes(cm);
84   const int num_pix = block_size_wide[bsize] * block_size_high[bsize];
85   const int num_blk = num_pix / 16;
86 
87   AOM_CHECK_MEM_ERROR(&error, ctx->blk_skip,
88                       aom_calloc(num_blk, sizeof(*ctx->blk_skip)));
89   AOM_CHECK_MEM_ERROR(&error, ctx->tx_type_map,
90                       aom_calloc(num_blk, sizeof(*ctx->tx_type_map)));
91   ctx->num_4x4_blk = num_blk;
92 
93   for (int i = 0; i < num_planes; ++i) {
94     ctx->coeff[i] = shared_bufs->coeff_buf[i];
95     ctx->qcoeff[i] = shared_bufs->qcoeff_buf[i];
96     ctx->dqcoeff[i] = shared_bufs->dqcoeff_buf[i];
97     AOM_CHECK_MEM_ERROR(&error, ctx->eobs[i],
98                         aom_memalign(32, num_blk * sizeof(*ctx->eobs[i])));
99     AOM_CHECK_MEM_ERROR(
100         &error, ctx->txb_entropy_ctx[i],
101         aom_memalign(32, num_blk * sizeof(*ctx->txb_entropy_ctx[i])));
102   }
103 
104   if (num_pix <= MAX_PALETTE_SQUARE) {
105     for (int i = 0; i < 2; ++i) {
106       if (cm->features.allow_screen_content_tools) {
107         AOM_CHECK_MEM_ERROR(
108             &error, ctx->color_index_map[i],
109             aom_memalign(32, num_pix * sizeof(*ctx->color_index_map[i])));
110       } else {
111         ctx->color_index_map[i] = NULL;
112       }
113     }
114   }
115 
116   av1_invalid_rd_stats(&ctx->rd_stats);
117 
118   return ctx;
119 }
120 
av1_reset_pmc(PICK_MODE_CONTEXT * ctx)121 void av1_reset_pmc(PICK_MODE_CONTEXT *ctx) {
122   av1_zero_array(ctx->blk_skip, ctx->num_4x4_blk);
123   av1_zero_array(ctx->tx_type_map, ctx->num_4x4_blk);
124   av1_invalid_rd_stats(&ctx->rd_stats);
125 }
126 
av1_free_pmc(PICK_MODE_CONTEXT * ctx,int num_planes)127 void av1_free_pmc(PICK_MODE_CONTEXT *ctx, int num_planes) {
128   if (ctx == NULL) return;
129 
130   aom_free(ctx->blk_skip);
131   ctx->blk_skip = NULL;
132   aom_free(ctx->tx_type_map);
133   for (int i = 0; i < num_planes; ++i) {
134     ctx->coeff[i] = NULL;
135     ctx->qcoeff[i] = NULL;
136     ctx->dqcoeff[i] = NULL;
137     aom_free(ctx->eobs[i]);
138     ctx->eobs[i] = NULL;
139     aom_free(ctx->txb_entropy_ctx[i]);
140     ctx->txb_entropy_ctx[i] = NULL;
141   }
142 
143   for (int i = 0; i < 2; ++i) {
144     if (ctx->color_index_map[i]) {
145       aom_free(ctx->color_index_map[i]);
146       ctx->color_index_map[i] = NULL;
147     }
148   }
149 
150   aom_free(ctx);
151 }
152 
av1_alloc_pc_tree_node(BLOCK_SIZE bsize)153 PC_TREE *av1_alloc_pc_tree_node(BLOCK_SIZE bsize) {
154   PC_TREE *pc_tree = aom_calloc(1, sizeof(*pc_tree));
155   if (pc_tree == NULL) return NULL;
156 
157   pc_tree->partitioning = PARTITION_NONE;
158   pc_tree->block_size = bsize;
159 
160   return pc_tree;
161 }
162 
163 #define FREE_PMC_NODE(CTX)         \
164   do {                             \
165     av1_free_pmc(CTX, num_planes); \
166     CTX = NULL;                    \
167   } while (0)
168 
av1_free_pc_tree_recursive(PC_TREE * pc_tree,int num_planes,int keep_best,int keep_none,PARTITION_SEARCH_TYPE partition_search_type)169 void av1_free_pc_tree_recursive(PC_TREE *pc_tree, int num_planes, int keep_best,
170                                 int keep_none,
171                                 PARTITION_SEARCH_TYPE partition_search_type) {
172   if (pc_tree == NULL) return;
173 
174   // Avoid freeing of extended partitions as they are not supported when
175   // partition_search_type is VAR_BASED_PARTITION.
176   if (partition_search_type == VAR_BASED_PARTITION && !keep_best &&
177       !keep_none) {
178     FREE_PMC_NODE(pc_tree->none);
179 
180     for (int i = 0; i < 2; ++i) {
181       FREE_PMC_NODE(pc_tree->horizontal[i]);
182       FREE_PMC_NODE(pc_tree->vertical[i]);
183     }
184 
185 #if !defined(NDEBUG) && !CONFIG_REALTIME_ONLY
186     for (int i = 0; i < 3; ++i) {
187       assert(pc_tree->horizontala[i] == NULL);
188       assert(pc_tree->horizontalb[i] == NULL);
189       assert(pc_tree->verticala[i] == NULL);
190       assert(pc_tree->verticalb[i] == NULL);
191     }
192     for (int i = 0; i < 4; ++i) {
193       assert(pc_tree->horizontal4[i] == NULL);
194       assert(pc_tree->vertical4[i] == NULL);
195     }
196 #endif
197 
198     for (int i = 0; i < 4; ++i) {
199       if (pc_tree->split[i] != NULL) {
200         av1_free_pc_tree_recursive(pc_tree->split[i], num_planes, 0, 0,
201                                    partition_search_type);
202         pc_tree->split[i] = NULL;
203       }
204     }
205     aom_free(pc_tree);
206     return;
207   }
208 
209   const PARTITION_TYPE partition = pc_tree->partitioning;
210 
211   if (!keep_none && (!keep_best || (partition != PARTITION_NONE)))
212     FREE_PMC_NODE(pc_tree->none);
213 
214   for (int i = 0; i < 2; ++i) {
215     if (!keep_best || (partition != PARTITION_HORZ))
216       FREE_PMC_NODE(pc_tree->horizontal[i]);
217     if (!keep_best || (partition != PARTITION_VERT))
218       FREE_PMC_NODE(pc_tree->vertical[i]);
219   }
220 #if !CONFIG_REALTIME_ONLY
221   for (int i = 0; i < 3; ++i) {
222     if (!keep_best || (partition != PARTITION_HORZ_A))
223       FREE_PMC_NODE(pc_tree->horizontala[i]);
224     if (!keep_best || (partition != PARTITION_HORZ_B))
225       FREE_PMC_NODE(pc_tree->horizontalb[i]);
226     if (!keep_best || (partition != PARTITION_VERT_A))
227       FREE_PMC_NODE(pc_tree->verticala[i]);
228     if (!keep_best || (partition != PARTITION_VERT_B))
229       FREE_PMC_NODE(pc_tree->verticalb[i]);
230   }
231   for (int i = 0; i < 4; ++i) {
232     if (!keep_best || (partition != PARTITION_HORZ_4))
233       FREE_PMC_NODE(pc_tree->horizontal4[i]);
234     if (!keep_best || (partition != PARTITION_VERT_4))
235       FREE_PMC_NODE(pc_tree->vertical4[i]);
236   }
237 #endif
238   if (!keep_best || (partition != PARTITION_SPLIT)) {
239     for (int i = 0; i < 4; ++i) {
240       if (pc_tree->split[i] != NULL) {
241         av1_free_pc_tree_recursive(pc_tree->split[i], num_planes, 0, 0,
242                                    partition_search_type);
243         pc_tree->split[i] = NULL;
244       }
245     }
246   }
247 
248   if (!keep_best && !keep_none) aom_free(pc_tree);
249 }
250 
av1_setup_sms_tree(AV1_COMP * const cpi,ThreadData * td)251 int av1_setup_sms_tree(AV1_COMP *const cpi, ThreadData *td) {
252   // The structure 'sms_tree' is used to store the simple motion search data for
253   // partition pruning in inter frames. Hence, the memory allocations and
254   // initializations related to it are avoided for allintra encoding mode.
255   if (cpi->oxcf.kf_cfg.key_freq_max == 0) return 0;
256 
257   AV1_COMMON *const cm = &cpi->common;
258   const int stat_generation_stage = is_stat_generation_stage(cpi);
259   const int is_sb_size_128 = cm->seq_params->sb_size == BLOCK_128X128;
260   const int tree_nodes =
261       av1_get_pc_tree_nodes(is_sb_size_128, stat_generation_stage);
262   int sms_tree_index = 0;
263   SIMPLE_MOTION_DATA_TREE *this_sms;
264   int square_index = 1;
265   int nodes;
266 
267   aom_free(td->sms_tree);
268   td->sms_tree =
269       (SIMPLE_MOTION_DATA_TREE *)aom_calloc(tree_nodes, sizeof(*td->sms_tree));
270   if (!td->sms_tree) return -1;
271   this_sms = &td->sms_tree[0];
272 
273   if (!stat_generation_stage) {
274     const int leaf_factor = is_sb_size_128 ? 4 : 1;
275     const int leaf_nodes = 256 * leaf_factor;
276 
277     // Sets up all the leaf nodes in the tree.
278     for (sms_tree_index = 0; sms_tree_index < leaf_nodes; ++sms_tree_index) {
279       SIMPLE_MOTION_DATA_TREE *const tree = &td->sms_tree[sms_tree_index];
280       tree->block_size = square[0];
281     }
282 
283     // Each node has 4 leaf nodes, fill each block_size level of the tree
284     // from leafs to the root.
285     for (nodes = leaf_nodes >> 2; nodes > 0; nodes >>= 2) {
286       for (int i = 0; i < nodes; ++i) {
287         SIMPLE_MOTION_DATA_TREE *const tree = &td->sms_tree[sms_tree_index];
288         tree->block_size = square[square_index];
289         for (int j = 0; j < 4; j++) tree->split[j] = this_sms++;
290         ++sms_tree_index;
291       }
292       ++square_index;
293     }
294   } else {
295     // Allocation for firstpass/LAP stage
296     // TODO(Mufaddal): refactor square_index to use a common block_size macro
297     // from firstpass.c
298     SIMPLE_MOTION_DATA_TREE *const tree = &td->sms_tree[sms_tree_index];
299     square_index = 2;
300     tree->block_size = square[square_index];
301   }
302 
303   // Set up the root node for the largest superblock size
304   td->sms_root = &td->sms_tree[tree_nodes - 1];
305   return 0;
306 }
307 
av1_free_sms_tree(ThreadData * td)308 void av1_free_sms_tree(ThreadData *td) {
309   aom_free(td->sms_tree);
310   td->sms_tree = NULL;
311 }
312