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