aom_entropy_optimizer.c 24.4 KB
Newer Older
Yue Chen's avatar
Yue Chen committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/*
 * Copyright (c) 2017, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 2 Clause License and
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
 * was not distributed with this source code in the LICENSE file, you can
 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
 * Media Patent License 1.0 was not distributed with this source code in the
 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
 */

// This tool is a gadget for offline probability training.
// A binary executable aom_entropy_optimizer will be generated in tools/. It
// parses a binary file consisting of counts written in the format of
15 16 17
// FRAME_COUNTS in entropymode.h, and computes optimized probability tables
// and CDF tables, which will be written to a new c file optimized_probs.c
// according to format in the codebase.
Yue Chen's avatar
Yue Chen committed
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36
//
// Command line: ./aom_entropy_optimizer [directory of the count file]
//
// The input file can either be generated by encoding a single clip by
// turning on entropy_stats experiment, or be collected at a larger scale at
// which a python script which will be provided soon can be used to aggregate
// multiple stats output.

#include <assert.h>
#include <stdio.h>
#include "./aom_config.h"
#include "av1/common/entropymode.h"

#define SPACES_PER_TAB 2

typedef unsigned int aom_count_type;
// A log file recording parsed counts
static FILE *logfile;  // TODO(yuec): make it a command line option

37
static INLINE uint8_t get_binary_prob_new(unsigned int n0, unsigned int n1) {
38 39 40 41 42 43
  // The "+1" will prevent this function from generating extreme probability
  // when both n0 and n1 are small
  const unsigned int den = n0 + 1 + n1 + 1;
  return get_prob(n0 + 1, den);
}

Yue Chen's avatar
Yue Chen committed
44 45
static int parse_stats(aom_count_type **ct_ptr, FILE *const probsfile, int tabs,
                       int dim_of_cts, int *cts_each_dim,
46
                       int flatten_last_dim) {
Yue Chen's avatar
Yue Chen committed
47 48 49 50 51 52 53
  if (dim_of_cts < 1) {
    fprintf(stderr, "The dimension of a counts vector should be at least 1!\n");
    return 1;
  }
  if (dim_of_cts == 1) {
    const int total_modes = cts_each_dim[0];
    aom_count_type *counts1d = *ct_ptr;
54
    uint8_t *probs = aom_malloc(sizeof(*probs) * (total_modes - 1));
Yue Chen's avatar
Yue Chen committed
55 56 57 58 59 60 61

    if (probs == NULL) {
      fprintf(stderr, "Allocating prob array failed!\n");
      return 1;
    }

    (*ct_ptr) += total_modes;
62 63
    assert(total_modes == 2);
    probs[0] = get_binary_prob_new(counts1d[0], counts1d[1]);
Yue Chen's avatar
Yue Chen committed
64 65
    if (tabs > 0) fprintf(probsfile, "%*c", tabs * SPACES_PER_TAB, ' ');
    for (int k = 0; k < total_modes - 1; ++k) {
66 67 68 69
      if (k == total_modes - 2)
        fprintf(probsfile, " %3d ", probs[k]);
      else
        fprintf(probsfile, " %3d,", probs[k]);
Yue Chen's avatar
Yue Chen committed
70 71 72 73 74 75 76
      fprintf(logfile, "%d ", counts1d[k]);
    }
    fprintf(logfile, "%d\n", counts1d[total_modes - 1]);
  } else if (dim_of_cts == 2 && flatten_last_dim) {
    assert(cts_each_dim[1] == 2);

    for (int k = 0; k < cts_each_dim[0]; ++k) {
77 78
      if (k == cts_each_dim[0] - 1) {
        fprintf(probsfile, " %3d ",
79
                get_binary_prob_new((*ct_ptr)[0], (*ct_ptr)[1]));
80 81
      } else {
        fprintf(probsfile, " %3d,",
82
                get_binary_prob_new((*ct_ptr)[0], (*ct_ptr)[1]));
83
      }
Yue Chen's avatar
Yue Chen committed
84 85 86 87 88 89 90 91 92 93 94 95 96 97
      fprintf(logfile, "%d %d\n", (*ct_ptr)[0], (*ct_ptr)[1]);
      (*ct_ptr) += 2;
    }
  } else {
    for (int k = 0; k < cts_each_dim[0]; ++k) {
      int tabs_next_level;
      if (dim_of_cts == 2 || (dim_of_cts == 3 && flatten_last_dim)) {
        fprintf(probsfile, "%*c{", tabs * SPACES_PER_TAB, ' ');
        tabs_next_level = 0;
      } else {
        fprintf(probsfile, "%*c{\n", tabs * SPACES_PER_TAB, ' ');
        tabs_next_level = tabs + 1;
      }
      if (parse_stats(ct_ptr, probsfile, tabs_next_level, dim_of_cts - 1,
98
                      cts_each_dim + 1, flatten_last_dim)) {
Yue Chen's avatar
Yue Chen committed
99 100 101
        return 1;
      }
      if (dim_of_cts == 2 || (dim_of_cts == 3 && flatten_last_dim)) {
102 103 104 105
        if (k == cts_each_dim[0] - 1)
          fprintf(probsfile, "}\n");
        else
          fprintf(probsfile, "},\n");
Yue Chen's avatar
Yue Chen committed
106
      } else {
107 108 109 110
        if (k == cts_each_dim[0] - 1)
          fprintf(probsfile, "%*c}\n", tabs * SPACES_PER_TAB, ' ');
        else
          fprintf(probsfile, "%*c},\n", tabs * SPACES_PER_TAB, ' ');
Yue Chen's avatar
Yue Chen committed
111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
      }
    }
  }
  return 0;
}

// This function parses the stats of a syntax, either binary or multi-symbol,
// in different contexts, and writes the optimized probability table to
// probsfile.
//   counts: pointer of the first count element in counts array
//   probsfile: output file
//   dim_of_cts: number of dimensions of counts array
//   cts_each_dim: an array storing size of each dimension of counts array
//   flatten_last_dim: for a binary syntax, if flatten_last_dim is 0, probs in
//                     different contexts will be written separately, e.g.,
//                     {{p1}, {p2}, ...};
//                     otherwise will be grouped together at the second last
//                     dimension, i.e.,
//                     {p1, p2, ...}.
//   prefix: declaration header for the entropy table
static void optimize_entropy_table(aom_count_type *counts,
                                   FILE *const probsfile, int dim_of_cts,
133 134
                                   int *cts_each_dim, int flatten_last_dim,
                                   char *prefix) {
Yue Chen's avatar
Yue Chen committed
135 136 137 138 139
  aom_count_type *ct_ptr = counts;

  assert(!flatten_last_dim || cts_each_dim[dim_of_cts - 1] == 2);

  fprintf(probsfile, "%s = {\n", prefix);
140
  if (parse_stats(&ct_ptr, probsfile, 1, dim_of_cts, cts_each_dim,
Yue Chen's avatar
Yue Chen committed
141 142 143 144 145 146 147
                  flatten_last_dim)) {
    fprintf(probsfile, "Optimizer failed!\n");
  }
  fprintf(probsfile, "};\n\n");
  fprintf(logfile, "\n");
}

148 149 150 151 152 153 154 155
static int counts_to_cdf(const aom_count_type *counts, aom_cdf_prob *cdf,
                         int modes) {
  int64_t *csum = aom_malloc(sizeof(*csum) * modes);

  if (csum == NULL) {
    fprintf(stderr, "Allocating csum array failed!\n");
    return 1;
  }
Dake He's avatar
Dake He committed
156 157
  csum[0] = counts[0] + 1;
  for (int i = 1; i < modes; ++i) csum[i] = counts[i] + 1 + csum[i - 1];
158 159 160 161

  int64_t sum = csum[modes - 1];
  int64_t round_shift = sum >> 1;
  for (int i = 0; i < modes; ++i) {
Dake He's avatar
Dake He committed
162 163 164
    if (sum <= 0) {
      cdf[i] = CDF_PROB_TOP - modes + i + 1;
    } else {
165
      cdf[i] = (csum[i] * CDF_PROB_TOP + round_shift) / sum;
Dake He's avatar
Dake He committed
166 167 168
      cdf[i] = AOMMIN(cdf[i], CDF_PROB_TOP - (modes - 1 + i) * 4);
      cdf[i] = (i == 0) ? AOMMAX(cdf[i], 4) : AOMMAX(cdf[i], cdf[i - 1] + 4);
    }
169
  }
Dake He's avatar
Dake He committed
170 171
  // if (sum <= 0) cdf[0] = CDF_PROB_TOP - 1;

172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195
  return 0;
}

static int parse_counts_for_cdf_opt(aom_count_type **ct_ptr,
                                    FILE *const probsfile, int tabs,
                                    int dim_of_cts, int *cts_each_dim) {
  if (dim_of_cts < 1) {
    fprintf(stderr, "The dimension of a counts vector should be at least 1!\n");
    return 1;
  }
  if (dim_of_cts == 1) {
    const int total_modes = cts_each_dim[0];
    aom_count_type *counts1d = *ct_ptr;
    aom_cdf_prob *cdfs = aom_malloc(sizeof(*cdfs) * total_modes);

    if (cdfs == NULL) {
      fprintf(stderr, "Allocating cdf array failed!\n");
      return 1;
    }

    counts_to_cdf(counts1d, cdfs, total_modes);
    (*ct_ptr) += total_modes;

    if (tabs > 0) fprintf(probsfile, "%*c", tabs * SPACES_PER_TAB, ' ');
196 197 198 199 200 201
    fprintf(probsfile, "AOM_CDF%d( ", total_modes);
    for (int k = 0; k < total_modes - 1; ++k) {
      fprintf(probsfile, "%d", cdfs[k]);
      if (k < total_modes - 2) fprintf(probsfile, ",");
    }
    fprintf(probsfile, " )");
202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246
  } else {
    for (int k = 0; k < cts_each_dim[0]; ++k) {
      int tabs_next_level;

      if (dim_of_cts == 2)
        fprintf(probsfile, "%*c{", tabs * SPACES_PER_TAB, ' ');
      else
        fprintf(probsfile, "%*c{\n", tabs * SPACES_PER_TAB, ' ');
      tabs_next_level = dim_of_cts == 2 ? 0 : tabs + 1;

      if (parse_counts_for_cdf_opt(ct_ptr, probsfile, tabs_next_level,
                                   dim_of_cts - 1, cts_each_dim + 1)) {
        return 1;
      }

      if (dim_of_cts == 2) {
        if (k == cts_each_dim[0] - 1)
          fprintf(probsfile, "}\n");
        else
          fprintf(probsfile, "},\n");
      } else {
        if (k == cts_each_dim[0] - 1)
          fprintf(probsfile, "%*c}\n", tabs * SPACES_PER_TAB, ' ');
        else
          fprintf(probsfile, "%*c},\n", tabs * SPACES_PER_TAB, ' ');
      }
    }
  }

  return 0;
}

static void optimize_cdf_table(aom_count_type *counts, FILE *const probsfile,
                               int dim_of_cts, int *cts_each_dim,
                               char *prefix) {
  aom_count_type *ct_ptr = counts;

  fprintf(probsfile, "%s = {\n", prefix);
  if (parse_counts_for_cdf_opt(&ct_ptr, probsfile, 1, dim_of_cts,
                               cts_each_dim)) {
    fprintf(probsfile, "Optimizer failed!\n");
  }
  fprintf(probsfile, "};\n\n");
}

Yue Chen's avatar
Yue Chen committed
247 248 249 250 251 252 253 254 255 256 257 258 259
int main(int argc, const char **argv) {
  if (argc < 2) {
    fprintf(stderr, "Please specify the input stats file!\n");
    exit(EXIT_FAILURE);
  }

  FILE *const statsfile = fopen(argv[1], "rb");
  if (statsfile == NULL) {
    fprintf(stderr, "Failed to open input file!\n");
    exit(EXIT_FAILURE);
  }

  FRAME_COUNTS fc;
Hui Su's avatar
Hui Su committed
260 261
  const size_t bytes = fread(&fc, sizeof(FRAME_COUNTS), 1, statsfile);
  if (!bytes) return 1;
Yue Chen's avatar
Yue Chen committed
262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277

  FILE *const probsfile = fopen("optimized_probs.c", "w");
  if (probsfile == NULL) {
    fprintf(stderr,
            "Failed to create output file for optimized entropy tables!\n");
    exit(EXIT_FAILURE);
  }

  logfile = fopen("aom_entropy_optimizer_parsed_counts.log", "w");
  if (logfile == NULL) {
    fprintf(stderr, "Failed to create log file for parsed counts!\n");
    exit(EXIT_FAILURE);
  }

  int cts_each_dim[10];

Hui Su's avatar
Hui Su committed
278
  /* Intra mode (keyframe luma) */
279 280 281 282 283 284 285
  cts_each_dim[0] = KF_MODE_CONTEXTS;
  cts_each_dim[1] = KF_MODE_CONTEXTS;
  cts_each_dim[2] = INTRA_MODES;
  optimize_cdf_table(&fc.kf_y_mode[0][0][0], probsfile, 3, cts_each_dim,
                     "const aom_cdf_prob\n"
                     "default_kf_y_mode_cdf[KF_MODE_CONTEXTS][KF_MODE_CONTEXTS]"
                     "[CDF_SIZE(INTRA_MODES)]");
Yue Chen's avatar
Yue Chen committed
286

Joe Young's avatar
Joe Young committed
287 288 289 290 291 292 293
  cts_each_dim[0] = DIRECTIONAL_MODES;
  cts_each_dim[1] = 2 * MAX_ANGLE_DELTA + 1;
  optimize_cdf_table(&fc.angle_delta[0][0], probsfile, 2, cts_each_dim,
                     "const aom_cdf_prob\n"
                     "default_angle_delta_cdf"
                     "[DIRECTIONAL_MODES][CDF_SIZE(2 * MAX_ANGLE_DELTA + 1)]");

Yue Chen's avatar
Yue Chen committed
294 295 296
  /* Intra mode (non-keyframe luma) */
  cts_each_dim[0] = BLOCK_SIZE_GROUPS;
  cts_each_dim[1] = INTRA_MODES;
297 298 299 300
  optimize_cdf_table(
      &fc.y_mode[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_if_y_mode_cdf[BLOCK_SIZE_GROUPS][CDF_SIZE(INTRA_MODES)]");
Yue Chen's avatar
Yue Chen committed
301

302 303 304 305 306 307 308 309 310 311
/* Intra mode (chroma) */
#if CONFIG_CFL
  cts_each_dim[0] = CFL_ALLOWED_TYPES;
  cts_each_dim[1] = INTRA_MODES;
  cts_each_dim[2] = UV_INTRA_MODES;
  optimize_cdf_table(&fc.uv_mode[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob\n"
                     "default_uv_mode_cdf[CFL_ALLOWED_TYPES][INTRA_MODES]"
                     "[CDF_SIZE(UV_INTRA_MODES)]");
#else
Yue Chen's avatar
Yue Chen committed
312
  cts_each_dim[0] = INTRA_MODES;
Luc Trudeau's avatar
Luc Trudeau committed
313 314 315 316 317
  cts_each_dim[1] = UV_INTRA_MODES;
  optimize_cdf_table(
      &fc.uv_mode[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_uv_mode_cdf[INTRA_MODES][CDF_SIZE(UV_INTRA_MODES)]");
318
#endif
Yue Chen's avatar
Yue Chen committed
319 320 321 322 323

  /* Partition */
  cts_each_dim[0] = PARTITION_CONTEXTS;
#if CONFIG_EXT_PARTITION_TYPES
  cts_each_dim[1] = EXT_PARTITION_TYPES;
324 325 326 327
  optimize_cdf_table(&fc.partition[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob\n"
                     "default_partition_cdf[PARTITION_CONTEXTS][CDF_SIZE(EXT_"
                     "PARTITION_TYPES)]");
Yue Chen's avatar
Yue Chen committed
328 329
#else
  cts_each_dim[1] = PARTITION_TYPES;
330 331 332 333
  optimize_cdf_table(
      &fc.partition[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_partition_cdf[PARTITION_CONTEXTS][CDF_SIZE(PARTITION_TYPES)]");
Yue Chen's avatar
Yue Chen committed
334 335 336 337 338
#endif

  /* Interpolation filter */
  cts_each_dim[0] = SWITCHABLE_FILTER_CONTEXTS;
  cts_each_dim[1] = SWITCHABLE_FILTERS;
339 340 341 342
  optimize_cdf_table(&fc.switchable_interp[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob\n"
                     "default_switchable_interp_cdf[SWITCHABLE_FILTER_CONTEXTS]"
                     "[CDF_SIZE(SWITCHABLE_FILTERS)]");
Yue Chen's avatar
Yue Chen committed
343 344 345 346 347

  /* Motion vector referencing */
  cts_each_dim[0] = NEWMV_MODE_CONTEXTS;
  cts_each_dim[1] = 2;
  optimize_entropy_table(
348
      &fc.newmv_mode[0][0], probsfile, 2, cts_each_dim, 1,
Yue Chen's avatar
Yue Chen committed
349
      "static const aom_prob default_newmv_prob[NEWMV_MODE_CONTEXTS]");
350 351 352
  optimize_cdf_table(&fc.newmv_mode[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_newmv_cdf[NEWMV_MODE_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
353

Sarah Parker's avatar
Sarah Parker committed
354
  cts_each_dim[0] = GLOBALMV_MODE_CONTEXTS;
Yue Chen's avatar
Yue Chen committed
355 356
  cts_each_dim[1] = 2;
  optimize_entropy_table(
357
      &fc.zeromv_mode[0][0], probsfile, 2, cts_each_dim, 1,
Sarah Parker's avatar
Sarah Parker committed
358
      "static const aom_prob default_zeromv_prob[GLOBALMV_MODE_CONTEXTS]");
359 360
  optimize_cdf_table(&fc.zeromv_mode[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
Sarah Parker's avatar
Sarah Parker committed
361
                     "default_zeromv_cdf[GLOBALMV_MODE_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
362 363 364 365

  cts_each_dim[0] = REFMV_MODE_CONTEXTS;
  cts_each_dim[1] = 2;
  optimize_entropy_table(
366
      &fc.refmv_mode[0][0], probsfile, 2, cts_each_dim, 1,
Yue Chen's avatar
Yue Chen committed
367
      "static const aom_prob default_refmv_prob[REFMV_MODE_CONTEXTS]");
368 369 370
  optimize_cdf_table(&fc.refmv_mode[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_refmv_cdf[REFMV_MODE_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
371 372 373

  cts_each_dim[0] = DRL_MODE_CONTEXTS;
  cts_each_dim[1] = 2;
374 375 376
  optimize_cdf_table(&fc.drl_mode[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_drl_cdf[DRL_MODE_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
377

378
  /* ext_inter experiment */
Yue Chen's avatar
Yue Chen committed
379 380 381
  /* New compound mode */
  cts_each_dim[0] = INTER_MODE_CONTEXTS;
  cts_each_dim[1] = INTER_COMPOUND_MODES;
382 383 384 385
  optimize_cdf_table(&fc.inter_compound_mode[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob\n"
                     "default_inter_compound_mode_cdf[INTER_MODE_CONTEXTS][CDF_"
                     "SIZE(INTER_COMPOUND_MODES)]");
Yue Chen's avatar
Yue Chen committed
386

Yue Chen's avatar
Yue Chen committed
387 388 389
  /* Interintra */
  cts_each_dim[0] = BLOCK_SIZE_GROUPS;
  cts_each_dim[1] = 2;
390 391 392
  optimize_cdf_table(&fc.interintra[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_interintra_cdf[BLOCK_SIZE_GROUPS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
393 394 395

  cts_each_dim[0] = BLOCK_SIZE_GROUPS;
  cts_each_dim[1] = INTERINTRA_MODES;
396 397 398 399
  optimize_cdf_table(&fc.interintra_mode[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob\n"
                     "default_interintra_mode_cdf[BLOCK_SIZE_GROUPS][CDF_SIZE("
                     "INTERINTRA_MODES)]");
Yue Chen's avatar
Yue Chen committed
400

401
  cts_each_dim[0] = BLOCK_SIZES_ALL;
Yue Chen's avatar
Yue Chen committed
402
  cts_each_dim[1] = 2;
403 404 405 406
  optimize_cdf_table(
      &fc.wedge_interintra[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_wedge_interintra_cdf[BLOCK_SIZES_ALL][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
407

Yue Chen's avatar
Yue Chen committed
408
  /* Compound type */
409
  cts_each_dim[0] = BLOCK_SIZES_ALL;
Yue Chen's avatar
Yue Chen committed
410
  cts_each_dim[1] = COMPOUND_TYPES;
411 412 413 414
  optimize_cdf_table(
      &fc.compound_interinter[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_compound_type_cdf[BLOCK_SIZES_ALL][CDF_SIZE(COMPOUND_TYPES)]");
Yue Chen's avatar
Yue Chen committed
415

416
  /* motion_var and warped_motion experiments */
417
  cts_each_dim[0] = BLOCK_SIZES_ALL;
Yue Chen's avatar
Yue Chen committed
418
  cts_each_dim[1] = MOTION_MODES;
419 420 421 422 423
  optimize_cdf_table(
      &fc.motion_mode[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_motion_mode_cdf[BLOCK_SIZES_ALL][CDF_SIZE(MOTION_MODES)]");
  cts_each_dim[0] = BLOCK_SIZES_ALL;
Yue Chen's avatar
Yue Chen committed
424
  cts_each_dim[1] = 2;
425 426 427
  optimize_cdf_table(&fc.obmc[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_obmc_cdf[BLOCK_SIZES_ALL][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
428 429 430 431

  /* Intra/inter flag */
  cts_each_dim[0] = INTRA_INTER_CONTEXTS;
  cts_each_dim[1] = 2;
432 433 434 435
  optimize_cdf_table(
      &fc.intra_inter[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_intra_inter_cdf[INTRA_INTER_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
436 437 438 439

  /* Single/comp ref flag */
  cts_each_dim[0] = COMP_INTER_CONTEXTS;
  cts_each_dim[1] = 2;
440 441 442 443
  optimize_cdf_table(
      &fc.comp_inter[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_comp_inter_cdf[COMP_INTER_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
444 445 446 447 448

/* ext_comp_refs experiment */
#if CONFIG_EXT_COMP_REFS
  cts_each_dim[0] = COMP_REF_TYPE_CONTEXTS;
  cts_each_dim[1] = 2;
449 450 451 452
  optimize_cdf_table(
      &fc.comp_ref_type[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_comp_ref_type_cdf[COMP_REF_TYPE_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
453 454 455 456

  cts_each_dim[0] = UNI_COMP_REF_CONTEXTS;
  cts_each_dim[1] = UNIDIR_COMP_REFS - 1;
  cts_each_dim[2] = 2;
457 458 459 460
  optimize_cdf_table(&fc.uni_comp_ref[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob\n"
                     "default_uni_comp_ref_cdf[UNI_COMP_REF_CONTEXTS][UNIDIR_"
                     "COMP_REFS - 1][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
461 462 463 464 465 466
#endif

  /* Reference frame (single ref) */
  cts_each_dim[0] = REF_CONTEXTS;
  cts_each_dim[1] = SINGLE_REFS - 1;
  cts_each_dim[2] = 2;
467 468 469 470
  optimize_cdf_table(
      &fc.single_ref[0][0][0], probsfile, 3, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_single_ref_cdf[REF_CONTEXTS][SINGLE_REFS - 1][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
471 472

  /* ext_refs experiment */
473
  cts_each_dim[0] = REF_CONTEXTS;
Yue Chen's avatar
Yue Chen committed
474 475
  cts_each_dim[1] = FWD_REFS - 1;
  cts_each_dim[2] = 2;
476 477 478
  optimize_cdf_table(
      &fc.comp_ref[0][0][0], probsfile, 3, cts_each_dim,
      "static const aom_cdf_prob\n"
479
      "default_comp_ref_cdf[REF_CONTEXTS][FWD_REFS - 1][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
480

481
  cts_each_dim[0] = REF_CONTEXTS;
Yue Chen's avatar
Yue Chen committed
482 483
  cts_each_dim[1] = BWD_REFS - 1;
  cts_each_dim[2] = 2;
484 485 486
  optimize_cdf_table(
      &fc.comp_bwdref[0][0][0], probsfile, 3, cts_each_dim,
      "static const aom_cdf_prob\n"
487
      "default_comp_bwdref_cdf[REF_CONTEXTS][BWD_REFS - 1][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
488

489
  /* Transform size */
Yue Chen's avatar
Yue Chen committed
490 491
  cts_each_dim[0] = TXFM_PARTITION_CONTEXTS;
  cts_each_dim[1] = 2;
492 493 494 495
  optimize_cdf_table(
      &fc.txfm_partition[0][0], probsfile, 2, cts_each_dim,
      "static const aom_cdf_prob\n"
      "default_txfm_partition_cdf[TXFM_PARTITION_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
496 497 498 499

  /* Skip flag */
  cts_each_dim[0] = SKIP_CONTEXTS;
  cts_each_dim[1] = 2;
500 501 502
  optimize_cdf_table(&fc.skip[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_skip_cdfs[SKIP_CONTEXTS][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
503 504 505 506

/* intrabc experiment */
#if CONFIG_INTRABC
  cts_each_dim[0] = 2;
507 508 509
  optimize_cdf_table(
      &fc.intrabc[0], probsfile, 1, cts_each_dim,
      "static const aom_cdf_prob default_intrabc_cdf[CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
510 511 512 513
#endif

/* filter_intra experiment */
#if CONFIG_FILTER_INTRA
514
  cts_each_dim[0] = FILTER_INTRA_MODES;
515
  optimize_cdf_table(
516
      &fc.filter_intra_mode[0], probsfile, 1, cts_each_dim,
517
      "static const aom_cdf_prob "
518
      "default_filter_intra_mode_cdf[CDF_SIZE(FILTER_INTRA_MODES)]");
519 520 521 522 523 524

  cts_each_dim[0] = TX_SIZES_ALL;
  cts_each_dim[1] = 2;
  optimize_cdf_table(&fc.filter_intra_tx[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_filter_intra_cdfs[TX_SIZES_ALL][CDF_SIZE(2)]");
Yue Chen's avatar
Yue Chen committed
525 526
#endif

527 528
#if CONFIG_LV_MAP

Jingning Han's avatar
Jingning Han committed
529 530 531 532
  cts_each_dim[0] = TX_SIZES;
  cts_each_dim[1] = TXB_SKIP_CONTEXTS;
  cts_each_dim[2] = 2;
  optimize_entropy_table(
533
      &fc.txb_skip[0][0][0], probsfile, 3, cts_each_dim, 1,
Jingning Han's avatar
Jingning Han committed
534 535 536 537 538 539 540
      "static const aom_prob "
      "default_txk_skip[TX_SIZES][PLANE_TYPES][SIG_COEF_CONTEXTS]");
  optimize_cdf_table(&fc.txb_skip[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_nz_map_cdf[TX_SIZES][PLANE_TYPES][SIG_COEF_"
                     "CONTEXTS][CDF_SIZE(2)]");

541 542 543 544 545
  cts_each_dim[0] = TX_SIZES;
  cts_each_dim[1] = PLANE_TYPES;
  cts_each_dim[2] = EOB_COEF_CONTEXTS;
  cts_each_dim[3] = 2;
  optimize_entropy_table(
546
      &fc.eob_flag[0][0][0][0], probsfile, 4, cts_each_dim, 1,
547 548 549
      "static const aom_prob "
      "default_eob_flag[TX_SIZES][PLANE_TYPES][EOB_COEF_CONTEXTS]");

550 551 552 553
  cts_each_dim[0] = TX_SIZES;
  cts_each_dim[1] = PLANE_TYPES;
  cts_each_dim[2] = EOB_COEF_CONTEXTS;
  cts_each_dim[3] = 2;
554 555 556 557 558
  optimize_cdf_table(
      &fc.eob_extra[0][0][0][0], probsfile, 4, cts_each_dim,
      "static const aom_cdf_prob "
      "default_eob_extra_cdf[TX_SIZES][PLANE_TYPES][EOB_COEF_CONTEXTS]"
      "[CDF_SIZE(2)]");
559

560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607
  cts_each_dim[0] = PLANE_TYPES;
  cts_each_dim[1] = 2;
  cts_each_dim[2] = 5;
  optimize_cdf_table(&fc.eob_multi16[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_eob_multi16[PLANE_TYPES][2][CDF_SIZE(5)]");

  cts_each_dim[0] = PLANE_TYPES;
  cts_each_dim[1] = 2;
  cts_each_dim[2] = 6;
  optimize_cdf_table(&fc.eob_multi32[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_eob_multi32[PLANE_TYPES][2][CDF_SIZE(6)]");

  cts_each_dim[0] = PLANE_TYPES;
  cts_each_dim[1] = 2;
  cts_each_dim[2] = 7;
  optimize_cdf_table(&fc.eob_multi64[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_eob_multi64[PLANE_TYPES][2][CDF_SIZE(7)]");

  cts_each_dim[0] = PLANE_TYPES;
  cts_each_dim[1] = 2;
  cts_each_dim[2] = 8;
  optimize_cdf_table(&fc.eob_multi128[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_eob_multi128[PLANE_TYPES][2][CDF_SIZE(8)]");

  cts_each_dim[0] = PLANE_TYPES;
  cts_each_dim[1] = 2;
  cts_each_dim[2] = 9;
  optimize_cdf_table(&fc.eob_multi256[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_eob_multi256[PLANE_TYPES][2][CDF_SIZE(9)]");

  cts_each_dim[0] = PLANE_TYPES;
  cts_each_dim[1] = 2;
  cts_each_dim[2] = 10;
  optimize_cdf_table(&fc.eob_multi512[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_eob_multi512[PLANE_TYPES][2][CDF_SIZE(10)]");

  cts_each_dim[0] = PLANE_TYPES;
  cts_each_dim[1] = 2;
  cts_each_dim[2] = 11;
  optimize_cdf_table(&fc.eob_multi1024[0][0][0], probsfile, 3, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_eob_multi1024[PLANE_TYPES][2][CDF_SIZE(11)]");
Dake He's avatar
Dake He committed
608

609 610 611 612 613 614
  cts_each_dim[0] = TX_SIZES;
  cts_each_dim[1] = PLANE_TYPES;
  cts_each_dim[2] = BR_CDF_SIZE - 1;
  cts_each_dim[3] = LEVEL_CONTEXTS;
  cts_each_dim[4] = 2;
  optimize_entropy_table(&fc.coeff_lps[0][0][0][0][0], probsfile, 5,
615
                         cts_each_dim, 1,
616 617 618
                         "static const aom_prob "
                         "default_coeff_lps[TX_SIZES][PLANE_TYPES][BR_CDF_SIZE-"
                         "1][LEVEL_CONTEXTS]");
Dake He's avatar
Dake He committed
619 620 621 622 623 624 625 626 627
  cts_each_dim[0] = TX_SIZES;
  cts_each_dim[1] = PLANE_TYPES;
  cts_each_dim[2] = LEVEL_CONTEXTS;
  cts_each_dim[3] = BR_CDF_SIZE;
  optimize_cdf_table(&fc.coeff_lps_multi[0][0][0][0], probsfile, 4,
                     cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_coeff_lps_multi[TX_SIZES][PLANE_TYPES][LEVEL_"
                     "CONTEXTS][CDF_SIZE(BR_CDF_SIZE)]");
Dake He's avatar
Dake He committed
628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647

  cts_each_dim[0] = TX_SIZES;
  cts_each_dim[1] = PLANE_TYPES;
  cts_each_dim[2] = SIG_COEF_CONTEXTS_2D + SIG_COEF_CONTEXTS_1D;
  cts_each_dim[3] = 4;
  optimize_cdf_table(
      &fc.coeff_base_multi[0][0][0][0], probsfile, 4, cts_each_dim,
      "static const aom_cdf_prob "
      "default_coeff_base_multi[TX_SIZES][PLANE_TYPES][SIG_COEF_CONTEXTS]"
      "[CDF_SIZE(NUM_BASE_LEVELS+2)]");

  cts_each_dim[0] = TX_SIZES;
  cts_each_dim[1] = PLANE_TYPES;
  cts_each_dim[2] = SIG_COEF_CONTEXTS_EOB;
  cts_each_dim[3] = 3;
  optimize_cdf_table(
      &fc.coeff_base_eob_multi[0][0][0][0], probsfile, 4, cts_each_dim,
      "static const aom_cdf_prob "
      "default_coeff_base_eob_multi[TX_SIZES][PLANE_TYPES][SIG_COEF_"
      "CONTEXTS_EOB][CDF_SIZE(NUM_BASE_LEVELS+1)]");
648

649
#endif  // CONFIG_LV_MAP
650

651 652 653 654 655
#if CONFIG_EXT_SKIP
  /* Skip mode flag */
  cts_each_dim[0] = SKIP_MODE_CONTEXTS;
  cts_each_dim[1] = 2;
  optimize_entropy_table(
656
      &fc.skip_mode[0][0], probsfile, 2, cts_each_dim, 1,
657 658 659 660 661 662
      "static const aom_prob default_skip_mode_probs[SKIP_MODE_CONTEXTS]");
  optimize_cdf_table(&fc.skip_mode[0][0], probsfile, 2, cts_each_dim,
                     "static const aom_cdf_prob "
                     "default_skip_mode_cdfs[SKIP_MODE_CONTEXTS][CDF_SIZE(2)]");
#endif  // CONFIG_EXT_SKIP

Yue Chen's avatar
Yue Chen committed
663 664 665 666 667 668
  fclose(statsfile);
  fclose(logfile);
  fclose(probsfile);

  return 0;
}