1 /* 2 * Copyright (c) 2021, 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 #ifndef AOM_AOM_AOM_EXTERNAL_PARTITION_H_ 12 #define AOM_AOM_AOM_EXTERNAL_PARTITION_H_ 13 14 /*!\defgroup aom_encoder AOMedia AOM/AV1 Encoder 15 * \ingroup aom 16 * 17 * @{ 18 */ 19 #include <stdint.h> 20 21 /*!\file 22 * \brief Provides function pointer definitions for the external partition. 23 * 24 * \note The external partition API should be considered experimental. Until the 25 * external partition API is declared stable, breaking changes may be made to 26 * this API in a future libaom release. 27 */ 28 29 /*!\brief Current ABI version number 30 * 31 * \internal 32 * If this file is altered in any way that changes the ABI, this value 33 * must be bumped. Examples include, but are not limited to, changing 34 * types, removing or reassigning enums, adding/removing/rearranging 35 * fields to structures. 36 */ 37 #define AOM_EXT_PART_ABI_VERSION 8 38 39 #ifdef __cplusplus 40 extern "C" { 41 #endif 42 43 /*!\brief Abstract external partition model handler 44 */ 45 typedef void *aom_ext_part_model_t; 46 47 /*!\brief Number of features to determine whether to skip partition none and 48 * do partition split directly. The same as "FEATURE_SIZE_SMS_SPLIT". 49 */ 50 #define AOM_EXT_PART_SIZE_DIRECT_SPLIT 17 51 52 /*!\brief Number of features to use simple motion search to prune out 53 * rectangular partition in some direction. The same as 54 * "FEATURE_SIZE_SMS_PRUNE_PART". 55 */ 56 #define AOM_EXT_PART_SIZE_PRUNE_PART 25 57 58 /*!\brief Number of features to prune split and rectangular partition 59 * after PARTITION_NONE. 60 */ 61 #define AOM_EXT_PART_SIZE_PRUNE_NONE 4 62 63 /*!\brief Number of features to terminates partition after partition none using 64 * simple_motion_search features and the rate, distortion, and rdcost of 65 * PARTITION_NONE. The same as "FEATURE_SIZE_SMS_TERM_NONE". 66 */ 67 #define AOM_EXT_PART_SIZE_TERM_NONE 28 68 69 /*!\brief Number of features to terminates partition after partition split. 70 */ 71 #define AOM_EXT_PART_SIZE_TERM_SPLIT 31 72 73 /*!\brief Number of features to prune rectangular partition using stats 74 * collected after partition split. 75 */ 76 #define AOM_EXT_PART_SIZE_PRUNE_RECT 9 77 78 /*!\brief Number of features to prune AB partition using stats 79 * collected after rectangular partition.. 80 */ 81 #define AOM_EXT_PART_SIZE_PRUNE_AB 10 82 83 /*!\brief Number of features to prune 4-way partition using stats 84 * collected after AB partition. 85 */ 86 #define AOM_EXT_PART_SIZE_PRUNE_4_WAY 18 87 88 /*!\brief Decision mode of the external partition model. 89 * AOM_EXT_PART_WHOLE_TREE: the external partition model should provide the 90 * whole partition tree for the superblock. 91 * 92 * AOM_EXT_PART_RECURSIVE: the external partition model provides the partition 93 * decision of the current block only. The decision process starts from 94 * the superblock size, down to the smallest block size (4x4) recursively. 95 */ 96 typedef enum aom_ext_part_decision_mode { 97 AOM_EXT_PART_WHOLE_TREE = 0, 98 AOM_EXT_PART_RECURSIVE = 1, 99 } aom_ext_part_decision_mode_t; 100 101 /*!\brief Config information sent to the external partition model. 102 * 103 * For example, the maximum superblock size determined by the sequence header. 104 */ 105 typedef struct aom_ext_part_config { 106 int superblock_size; ///< super block size (either 64x64 or 128x128) 107 } aom_ext_part_config_t; 108 109 /*!\brief Features pass to the external model to make partition decisions. 110 * Specifically, features collected before NONE partition. 111 * Features "f" are used to determine: 112 * partition_none_allowed, partition_horz_allowed, partition_vert_allowed, 113 * do_rectangular_split, do_square_split 114 * Features "f_part2" are used to determine: 115 * prune_horz, prune_vert. 116 */ 117 typedef struct aom_partition_features_before_none { 118 /*! features to determine whether skip partition none and do split directly */ 119 float f[AOM_EXT_PART_SIZE_DIRECT_SPLIT]; 120 /*! features to determine whether to prune rectangular partition */ 121 float f_part2[AOM_EXT_PART_SIZE_PRUNE_PART]; 122 } aom_partition_features_before_none_t; 123 124 /*!\brief Features pass to the external model to make partition decisions. 125 * Specifically, features collected after NONE partition. 126 */ 127 typedef struct aom_partition_features_none { 128 /*! features to prune split and rectangular partition */ 129 float f[AOM_EXT_PART_SIZE_PRUNE_NONE]; 130 /*! features to determine termination of partition */ 131 float f_terminate[AOM_EXT_PART_SIZE_TERM_NONE]; 132 } aom_partition_features_none_t; 133 134 /*!\brief Features pass to the external model to make partition decisions. 135 * Specifically, features collected after SPLIT partition. 136 */ 137 typedef struct aom_partition_features_split { 138 /*! features to determine termination of partition */ 139 float f_terminate[AOM_EXT_PART_SIZE_TERM_SPLIT]; 140 /*! features to determine pruning rect partition */ 141 float f_prune_rect[AOM_EXT_PART_SIZE_PRUNE_RECT]; 142 } aom_partition_features_split_t; 143 144 /*!\brief Features pass to the external model to make partition decisions. 145 * Specifically, features collected after RECTANGULAR partition. 146 */ 147 typedef struct aom_partition_features_rect { 148 /*! features to determine pruning AB partition */ 149 float f[AOM_EXT_PART_SIZE_PRUNE_AB]; 150 } aom_partition_features_rect_t; 151 152 /*!\brief Features pass to the external model to make partition decisions. 153 * Specifically, features collected after AB partition: HORZ_A, HORZ_B, VERT_A, 154 * VERT_B. 155 */ 156 typedef struct aom_partition_features_ab { 157 /*! features to determine pruning 4-way partition */ 158 float f[AOM_EXT_PART_SIZE_PRUNE_4_WAY]; 159 } aom_partition_features_ab_t; 160 161 /*!\brief Feature id to tell the external model the current stage in partition 162 * pruning and what features to use to make decisions accordingly. 163 */ 164 typedef enum { 165 AOM_EXT_PART_FEATURE_BEFORE_NONE, 166 AOM_EXT_PART_FEATURE_BEFORE_NONE_PART2, 167 AOM_EXT_PART_FEATURE_AFTER_NONE, 168 AOM_EXT_PART_FEATURE_AFTER_NONE_PART2, 169 AOM_EXT_PART_FEATURE_AFTER_SPLIT, 170 AOM_EXT_PART_FEATURE_AFTER_SPLIT_PART2, 171 AOM_EXT_PART_FEATURE_AFTER_RECT, 172 AOM_EXT_PART_FEATURE_AFTER_AB 173 } AOM_EXT_PART_FEATURE_ID; 174 175 /*!\brief Features collected from the tpl process. 176 * 177 * The tpl process collects information that help measure the inter-frame 178 * dependency. 179 * The tpl process is computed in the unit of tpl_bsize_1d (16x16). 180 * Therefore, the max number of units inside a superblock is 181 * 128x128 / (16x16) = 64. Change it if the tpl process changes. 182 */ 183 typedef struct aom_sb_tpl_features { 184 int available; ///< If tpl stats are available 185 int tpl_unit_length; ///< The block length of tpl process 186 int num_units; ///< The number of units inside the current superblock 187 int64_t intra_cost[64]; ///< The intra cost of each unit 188 int64_t inter_cost[64]; ///< The inter cost of each unit 189 int64_t mc_dep_cost[64]; ///< The motion compensated dependency cost 190 } aom_sb_tpl_features_t; 191 192 /*!\brief Features collected from the simple motion process. 193 * 194 * The simple motion process collects information by applying motion compensated 195 * prediction on each block. 196 * The block size is 16x16, which could be changed. If it is changed, update 197 * comments and the array size here. 198 */ 199 typedef struct aom_sb_simple_motion_features { 200 int unit_length; ///< The block length of the simple motion process 201 int num_units; ///< The number of units inside the current superblock 202 int block_sse[64]; ///< Sum of squared error of each unit 203 int block_var[64]; ///< Variance of each unit 204 } aom_sb_simple_motion_features_t; 205 206 /*!\brief Features of each super block. 207 * 208 * Features collected for each super block before partition search. 209 */ 210 typedef struct aom_sb_features { 211 /*! Features from motion search */ 212 aom_sb_simple_motion_features_t motion_features; 213 /*! Features from tpl process */ 214 aom_sb_tpl_features_t tpl_features; 215 } aom_sb_features_t; 216 217 /*!\brief Features pass to the external model to make partition decisions. 218 * 219 * The encoder sends these features to the external model through 220 * "func()" defined in ..... 221 * 222 * NOTE: new member variables may be added to this structure in the future. 223 * Once new features are finalized, bump the major version of libaom. 224 */ 225 typedef struct aom_partition_features { 226 // Features for the current supervised multi-stage ML model. 227 /*! Feature ID to indicate active features */ 228 AOM_EXT_PART_FEATURE_ID id; 229 /*! Features collected before NONE partition */ 230 aom_partition_features_before_none_t before_part_none; 231 /*! Features collected after NONE partition */ 232 aom_partition_features_none_t after_part_none; 233 /*! Features collected after SPLIT partition */ 234 aom_partition_features_split_t after_part_split; 235 /*! Features collected after RECTANGULAR partition */ 236 aom_partition_features_rect_t after_part_rect; 237 /*! Features collected after AB partition */ 238 aom_partition_features_ab_t after_part_ab; 239 240 // Features for a new ML model. 241 aom_sb_features_t sb_features; ///< Features collected for the super block 242 int mi_row; ///< Mi_row position of the block 243 int mi_col; ///< Mi_col position of the block 244 int frame_width; ///< Frame width 245 int frame_height; ///< Frame height 246 int block_size; ///< As "BLOCK_SIZE" in av1/common/enums.h 247 /*! 248 * Valid partition types. A bitmask is used. "1" represents the 249 * corresponding type is valid. The bitmask follows the enum order for 250 * PARTITION_TYPE in "enums.h" to represent one partition type at a bit. 251 * For example, 0x01 stands for only PARTITION_NONE is valid, 252 * 0x09 (00...001001) stands for PARTITION_NONE and PARTITION_SPLIT are valid. 253 */ 254 int valid_partition_types; 255 int update_type; ///< Frame update type, defined in ratectrl.h 256 int qindex; ///< Quantization index, range: [0, 255] 257 int rdmult; ///< Rate-distortion multiplier 258 int pyramid_level; ///< The level of this frame in the hierarchical structure 259 int has_above_block; ///< Has above neighbor block 260 int above_block_width; ///< Width of the above block, -1 if not exist 261 int above_block_height; ///< Height of the above block, -1 if not exist 262 int has_left_block; ///< Has left neighbor block 263 int left_block_width; ///< Width of the left block, -1 if not exist 264 int left_block_height; ///< Height of the left block, -1 if not exist 265 /*! 266 * The following parameters are collected from applying simple motion search. 267 * Sum of squared error (SSE) and variance of motion compensated residual 268 * are good indicators of block partitioning. 269 * If a block is a square, we also apply motion search for its 4 sub blocks. 270 * If not a square, their values are -1. 271 * If a block is able to split horizontally, we apply motion search and get 272 * stats for horizontal blocks. If not, their values are -1. 273 * If a block is able to split vertically, we apply motion search and get 274 * stats for vertical blocks. If not, their values are -1. 275 */ 276 unsigned int block_sse; ///< SSE of motion compensated residual 277 unsigned int block_var; ///< Variance of motion compensated residual 278 unsigned int sub_block_sse[4]; ///< SSE of sub blocks. 279 unsigned int sub_block_var[4]; ///< Variance of sub blocks. 280 unsigned int horz_block_sse[2]; ///< SSE of horz sub blocks 281 unsigned int horz_block_var[2]; ///< Variance of horz sub blocks 282 unsigned int vert_block_sse[2]; ///< SSE of vert sub blocks 283 unsigned int vert_block_var[2]; ///< Variance of vert sub blocks 284 /*! 285 * The following parameters are calculated from tpl model. 286 * If tpl model is not available, their values are -1. 287 */ 288 int64_t tpl_intra_cost; ///< Intra cost, ref to "TplDepStats" in tpl_model.h 289 int64_t tpl_inter_cost; ///< Inter cost in tpl model 290 int64_t tpl_mc_dep_cost; ///< Motion compensated dependency cost in tpl model 291 } aom_partition_features_t; 292 293 /*!\brief Partition decisions received from the external model. 294 * 295 * The encoder receives partition decisions and encodes the superblock 296 * with the given partition type. 297 * The encoder receives it from "func()" define in .... 298 * 299 * NOTE: new member variables may be added to this structure in the future. 300 * Once new features are finalized, bump the major version of libaom. 301 */ 302 typedef struct aom_partition_decision { 303 // Decisions for directly set partition types 304 int is_final_decision; ///< The flag whether it's the final decision 305 int num_nodes; ///< The number of leaf nodes 306 int partition_decision[2048]; ///< Partition decisions 307 int current_decision; ///< Partition decision for the current block 308 309 // Decisions for partition type pruning 310 int terminate_partition_search; ///< Terminate further partition search 311 int partition_none_allowed; ///< Allow partition none type 312 int partition_rect_allowed[2]; ///< Allow rectangular partitions 313 int do_rectangular_split; ///< Try rectangular split partition 314 int do_square_split; ///< Try square split partition 315 int prune_rect_part[2]; ///< Prune rectangular partition 316 int horza_partition_allowed; ///< Allow HORZ_A partition 317 int horzb_partition_allowed; ///< Allow HORZ_B partition 318 int verta_partition_allowed; ///< Allow VERT_A partition 319 int vertb_partition_allowed; ///< Allow VERT_B partition 320 int partition_horz4_allowed; ///< Allow HORZ4 partition 321 int partition_vert4_allowed; ///< Allow VERT4 partition 322 } aom_partition_decision_t; 323 324 /*!\brief Encoding stats for the given partition decision. 325 * 326 * The encoding stats collected by encoding the superblock with the 327 * given partition types. 328 * The encoder sends the stats to the external model for training 329 * or inference through "func()" defined in .... 330 */ 331 typedef struct aom_partition_stats { 332 int rate; ///< Rate cost of the block 333 int64_t dist; ///< Distortion of the block 334 int64_t rdcost; ///< Rate-distortion cost of the block 335 } aom_partition_stats_t; 336 337 /*!\brief Enum for return status. 338 */ 339 typedef enum aom_ext_part_status { 340 AOM_EXT_PART_OK = 0, ///< Status of success 341 AOM_EXT_PART_ERROR = 1, ///< Status of failure 342 AOM_EXT_PART_TEST = 2, ///< Status used for tests 343 } aom_ext_part_status_t; 344 345 /*!\brief Callback of creating an external partition model. 346 * 347 * The callback is invoked by the encoder to create an external partition 348 * model. 349 * 350 * \param[in] priv Callback's private data 351 * \param[in] part_config Config information pointer for model creation 352 * \param[out] ext_part_model Pointer to the model 353 */ 354 typedef aom_ext_part_status_t (*aom_ext_part_create_model_fn_t)( 355 void *priv, const aom_ext_part_config_t *part_config, 356 aom_ext_part_model_t *ext_part_model); 357 358 /*!\brief Callback of sending features to the external partition model. 359 * 360 * The callback is invoked by the encoder to send features to the external 361 * partition model. 362 * 363 * \param[in] ext_part_model The external model 364 * \param[in] part_features Pointer to the features 365 */ 366 typedef aom_ext_part_status_t (*aom_ext_part_send_features_fn_t)( 367 aom_ext_part_model_t ext_part_model, 368 const aom_partition_features_t *part_features); 369 370 /*!\brief Callback of receiving partition decisions from the external 371 * partition model. 372 * 373 * The callback is invoked by the encoder to receive partition decisions from 374 * the external partition model. 375 * 376 * \param[in] ext_part_model The external model 377 * \param[in] ext_part_decision Pointer to the partition decisions 378 */ 379 typedef aom_ext_part_status_t (*aom_ext_part_get_decision_fn_t)( 380 aom_ext_part_model_t ext_part_model, 381 aom_partition_decision_t *ext_part_decision); 382 383 /*!\brief Callback of sending stats to the external partition model. 384 * 385 * The callback is invoked by the encoder to send encoding stats to 386 * the external partition model. 387 * 388 * \param[in] ext_part_model The external model 389 * \param[in] ext_part_stats Pointer to the encoding stats 390 */ 391 typedef aom_ext_part_status_t (*aom_ext_part_send_partition_stats_fn_t)( 392 aom_ext_part_model_t ext_part_model, 393 const aom_partition_stats_t *ext_part_stats); 394 395 /*!\brief Callback of deleting the external partition model. 396 * 397 * The callback is invoked by the encoder to delete the external partition 398 * model. 399 * 400 * \param[in] ext_part_model The external model 401 */ 402 typedef aom_ext_part_status_t (*aom_ext_part_delete_model_fn_t)( 403 aom_ext_part_model_t ext_part_model); 404 405 /*!\brief Callback function set for external partition model. 406 * 407 * Uses can enable external partition model by registering a set of 408 * callback functions with the flag: AV1E_SET_EXTERNAL_PARTITION_MODEL 409 */ 410 typedef struct aom_ext_part_funcs { 411 /*! 412 * Create an external partition model. 413 */ 414 aom_ext_part_create_model_fn_t create_model; 415 416 /*! 417 * Send features to the external partition model to make partition decisions. 418 */ 419 aom_ext_part_send_features_fn_t send_features; 420 421 /*! 422 * Get partition decisions from the external partition model. 423 */ 424 aom_ext_part_get_decision_fn_t get_partition_decision; 425 426 /*! 427 * Send stats of the current partition to the external model. 428 */ 429 aom_ext_part_send_partition_stats_fn_t send_partition_stats; 430 431 /*! 432 * Delete the external partition model. 433 */ 434 aom_ext_part_delete_model_fn_t delete_model; 435 436 /*! 437 * The decision mode of the model. 438 */ 439 aom_ext_part_decision_mode_t decision_mode; 440 441 /*! 442 * Private data for the external partition model. 443 */ 444 void *priv; 445 } aom_ext_part_funcs_t; 446 447 /*!@} - end defgroup aom_encoder*/ 448 #ifdef __cplusplus 449 } // extern "C" 450 #endif 451 452 #endif // AOM_AOM_AOM_EXTERNAL_PARTITION_H_ 453