aom/aom_external_partition.h

   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 8
  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   /*!
 244    * Valid partition types. A bitmask is used.  "1" represents the
 245    * corresponding type is vaild. The bitmask follows the enum order for
 246    * PARTITION_TYPE in "enums.h" to represent one partition type at a bit.
 247    * For example, 0x01 stands for only PARTITION_NONE is valid,
 248    * 0x09 (00...001001) stands for PARTITION_NONE and PARTITION_SPLIT are valid.
 249    */
 250   int valid_partition_types;
 251   int update_type;    ///< Frame update type, defined in ratectrl.h
 252   int qindex;         ///< Quantization index, range: [0, 255]
 253   int rdmult;         ///< Rate-distortion multiplier
 254   int pyramid_level;  ///< The level of this frame in the hierarchical structure
 255   int has_above_block;     ///< Has above neighbor block
 256   int above_block_width;   ///< Width of the above block, -1 if not exist
 257   int above_block_height;  ///< Height of the above block, -1 if not exist
 258   int has_left_block;      ///< Has left neighbor block
 259   int left_block_width;    ///< Width of the left block, -1 if not exist
 260   int left_block_height;   ///< Height of the left block, -1 if not exist
 261   /*!
 262    * The following parameters are collected from applying simple motion search.
 263    * Sum of squared error (SSE) and variance of motion compensated residual
 264    * are good indicators of block partitioning.
 265    * If a block is a square, we also apply motion search for its 4 sub blocks.
 266    * If not a square, their values are -1.
 267    * If a block is able to split horizontally, we apply motion search and get
 268    * stats for horizontal blocks. If not, their values are -1.
 269    * If a block is able to split vertically, we apply motion search and get
 270    * stats for vertical blocks. If not, their values are -1.
 271    */
 272   unsigned int block_sse;          ///< SSE of motion compensated residual
 273   unsigned int block_var;          ///< Variance of motion compensated residual
 274   unsigned int sub_block_sse[4];   ///< SSE of sub blocks.
 275   unsigned int sub_block_var[4];   ///< Variance of sub blocks.
 276   unsigned int horz_block_sse[2];  ///< SSE of horz sub blocks
 277   unsigned int horz_block_var[2];  ///< Variance of horz sub blocks
 278   unsigned int vert_block_sse[2];  ///< SSE of vert sub blocks
 279   unsigned int vert_block_var[2];  ///< Variance of vert sub blocks
 280   /*!
 281    * The following parameters are calculated from tpl model.
 282    * If tpl model is not available, their values are -1.
 283    */
 284   int64_t tpl_intra_cost;   ///< Intra cost, ref to "TplDepStats" in tpl_model.h
 285   int64_t tpl_inter_cost;   ///< Inter cost in tpl model
 286   int64_t tpl_mc_dep_cost;  ///< Motion compensated dependency cost in tpl model
 287 } aom_partition_features_t;
 288
 289 /*!\brief Partition decisions received from the external model.
 290  *
 291  * The encoder receives partition decisions and encodes the superblock
 292  * with the given partition type.
 293  * The encoder receives it from "func()" define in ....
 294  *
 295  * NOTE: new member variables may be added to this structure in the future.
 296  * Once new features are finalized, bump the major version of libaom.
 297  */
 298 typedef struct aom_partition_decision {
 299   // Decisions for directly set partition types
 300   int is_final_decision;         ///< The flag whether it's the final decision
 301   int num_nodes;                 ///< The number of leaf nodes
 302   int partition_decision[2048];  ///< Partition decisions
 303   int current_decision;          ///< Partition decision for the current block
 304
 305   // Decisions for partition type pruning
 306   int terminate_partition_search;  ///< Terminate further partition search
 307   int partition_none_allowed;      ///< Allow partition none type
 308   int partition_rect_allowed[2];   ///< Allow rectangular partitions
 309   int do_rectangular_split;        ///< Try rectangular split partition
 310   int do_square_split;             ///< Try square split partition
 311   int prune_rect_part[2];          ///< Prune rectangular partition
 312   int horza_partition_allowed;     ///< Allow HORZ_A partitioin
 313   int horzb_partition_allowed;     ///< Allow HORZ_B partitioin
 314   int verta_partition_allowed;     ///< Allow VERT_A partitioin
 315   int vertb_partition_allowed;     ///< Allow VERT_B partitioin
 316   int partition_horz4_allowed;     ///< Allow HORZ4 partition
 317   int partition_vert4_allowed;     ///< Allow VERT4 partition
 318 } aom_partition_decision_t;
 319
 320 /*!\brief Encoding stats for the given partition decision.
 321  *
 322  * The encoding stats collected by encoding the superblock with the
 323  * given partition types.
 324  * The encoder sends the stats to the external model for training
 325  * or inference though "func()" defined in ....
 326  */
 327 typedef struct aom_partition_stats {
 328   int rate;        ///< Rate cost of the block
 329   int64_t dist;    ///< Distortion of the block
 330   int64_t rdcost;  ///< Rate-distortion cost of the block
 331 } aom_partition_stats_t;
 332
 333 /*!\brief Enum for return status.
 334  */
 335 typedef enum aom_ext_part_status {
 336   AOM_EXT_PART_OK = 0,     ///< Status of success
 337   AOM_EXT_PART_ERROR = 1,  ///< Status of failure
 338   AOM_EXT_PART_TEST = 2,   ///< Status used for tests
 339 } aom_ext_part_status_t;
 340
 341 /*!\brief Callback of creating an external partition model.
 342  *
 343  * The callback is invoked by the encoder to create an external partition
 344  * model.
 345  *
 346  * \param[in] priv Callback's private data
 347  * \param[in] part_config Config information pointer for model creation
 348  * \param[out] ext_part_model Pointer to the model
 349  */
 350 typedef aom_ext_part_status_t (*aom_ext_part_create_model_fn_t)(
 351     void *priv, const aom_ext_part_config_t *part_config,
 352     aom_ext_part_model_t *ext_part_model);
 353
 354 /*!\brief Callback of sending features to the external partition model.
 355  *
 356  * The callback is invoked by the encoder to send features to the external
 357  * partition model.
 358  *
 359  * \param[in] ext_part_model The external model
 360  * \param[in] part_features Pointer to the features
 361  */
 362 typedef aom_ext_part_status_t (*aom_ext_part_send_features_fn_t)(
 363     aom_ext_part_model_t ext_part_model,
 364     const aom_partition_features_t *part_features);
 365
 366 /*!\brief Callback of receiving partition decisions from the external
 367  * partition model.
 368  *
 369  * The callback is invoked by the encoder to receive partition decisions from
 370  * the external partition model.
 371  *
 372  * \param[in] ext_part_model The external model
 373  * \param[in] ext_part_decision Pointer to the partition decisions
 374  */
 375 typedef aom_ext_part_status_t (*aom_ext_part_get_decision_fn_t)(
 376     aom_ext_part_model_t ext_part_model,
 377     aom_partition_decision_t *ext_part_decision);
 378
 379 /*!\brief Callback of sending stats to the external partition model.
 380  *
 381  * The callback is invoked by the encoder to send encoding stats to
 382  * the external partition model.
 383  *
 384  * \param[in] ext_part_model The external model
 385  * \param[in] ext_part_stats Pointer to the encoding stats
 386  */
 387 typedef aom_ext_part_status_t (*aom_ext_part_send_partition_stats_fn_t)(
 388     aom_ext_part_model_t ext_part_model,
 389     const aom_partition_stats_t *ext_part_stats);
 390
 391 /*!\brief Callback of deleting the external partition model.
 392  *
 393  * The callback is invoked by the encoder to delete the external partition
 394  * model.
 395  *
 396  * \param[in] ext_part_model The external model
 397  */
 398 typedef aom_ext_part_status_t (*aom_ext_part_delete_model_fn_t)(
 399     aom_ext_part_model_t ext_part_model);
 400
 401 /*!\brief Callback function set for external partition model.
 402  *
 403  * Uses can enable external partition model by registering a set of
 404  * callback functions with the flag: AV1E_SET_EXTERNAL_PARTITION_MODEL
 405  */
 406 typedef struct aom_ext_part_funcs {
 407   /*!
 408    * Create an external partition model.
 409    */
 410   aom_ext_part_create_model_fn_t create_model;
 411
 412   /*!
 413    * Send features to the external partition model to make partition decisions.
 414    */
 415   aom_ext_part_send_features_fn_t send_features;
 416
 417   /*!
 418    * Get partition decisions from the external partition model.
 419    */
 420   aom_ext_part_get_decision_fn_t get_partition_decision;
 421
 422   /*!
 423    * Send stats of the current partition to the external model.
 424    */
 425   aom_ext_part_send_partition_stats_fn_t send_partition_stats;
 426
 427   /*!
 428    * Delete the external partition model.
 429    */
 430   aom_ext_part_delete_model_fn_t delete_model;
 431
 432   /*!
 433    * The decision mode of the model.
 434    */
 435   aom_ext_part_decision_mode_t decision_mode;
 436
 437   /*!
 438    * Private data for the external partition model.
 439    */
 440   void *priv;
 441 } aom_ext_part_funcs_t;
 442
 443 /*!@} - end defgroup aom_encoder*/
 444 #ifdef __cplusplus
 445 }  // extern "C"
 446 #endif
 447
 448 #endif  // AOM_AOM_AOM_EXTERNAL_PARTITION_H_