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