Add a control to set the ANS window size

Change-Id: I3d64ec4bbc72143b30a094ece7a6c711d6b479cd

aom/aomcx.h

@@ -471,6 +471,18 @@ enum aome_enc_control_id {
    * Experiment: TILE_GROUPS
    */
   AV1E_SET_MTU,
+
+  /*!\brief Codec control function to set the number of symbols in an ANS data
+   * window.
+   *
+   * The number of ANS symbols (both boolean and non-booleans alphabets) in an
+   * ANS data window is set to 1 << value.
+   *
+   * \note Valid range: [8, 23]
+   *
+   * Experiment: ANS
+   */
+  AV1E_SET_ANS_WINDOW_SIZE_LOG2,
 };
 
 /*!\brief aom 1-D scaling mode
@@ -664,6 +676,9 @@ AOM_CTRL_USE_TYPE(AV1E_SET_TARGET_LEVEL, unsigned int)
 
 AOM_CTRL_USE_TYPE(AV1E_GET_LEVEL, int *)
 #define AOM_CTRL_AV1E_GET_LEVEL
+
+AOM_CTRL_USE_TYPE(AV1E_SET_ANS_WINDOW_SIZE_LOG2, unsigned int)
+#define AOM_CTRL_AV1E_SET_ANS_WINDOW_SIZE_LOG2
 /*!\endcond */
 /*! @} - end defgroup aom_encoder */
 #ifdef __cplusplus

aom_dsp/ans.h

@@ -23,7 +23,8 @@
 extern "C" {
 #endif  // __cplusplus
 
-#define ANS_MAX_SYMBOLS (1 << 24)
+// Use windowed ANS, size is passed in at initialization
+#define ANS_MAX_SYMBOLS 1
 #define ANS_REVERSE 1
 
 typedef uint8_t AnsP8;

aom_dsp/ansreader.h

@@ -177,6 +177,7 @@ static INLINE int ans_read_init(struct AnsDecoder *const ans,
   ans->state += L_BASE;
   if (ans->state >= L_BASE * IO_BASE) return 1;
 #if ANS_MAX_SYMBOLS
+  assert(ans->window_size > 1);
   ans->symbols_left = ans->window_size;
 #endif
   return 0;

aom_dsp/buf_ans.c

@@ -19,6 +19,7 @@ void aom_buf_ans_alloc(struct BufAnsCoder *c,
                        struct aom_internal_error_info *error, int size) {
   c->error = error;
   c->size = size;
+  assert(c->size > 1);
   AOM_CHECK_MEM_ERROR(error, c->buf, aom_malloc(c->size * sizeof(*c->buf)));
   // Initialize to overfull to trigger the assert in write.
   c->offset = c->size + 1;

av1/av1_cx_iface.c

@@ -66,6 +66,9 @@ struct av1_extracfg {
   int render_width;
   int render_height;
   aom_superblock_size_t superblock_size;
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+  int ans_window_size_log2;
+#endif
 };
 
 static struct av1_extracfg default_extra_cfg = {
@@ -106,16 +109,19 @@ static struct av1_extracfg default_extra_cfg = {
   1,  // max number of tile groups
   0,  // mtu_size
 #endif
-  1,                           // frame_parallel_decoding_mode
-  NO_AQ,                       // aq_mode
-  0,                           // frame_periodic_delta_q
-  AOM_BITS_8,                  // Bit depth
-  AOM_CONTENT_DEFAULT,         // content
-  AOM_CS_UNKNOWN,              // color space
-  0,                           // color range
-  0,                           // render width
-  0,                           // render height
-  AOM_SUPERBLOCK_SIZE_DYNAMIC  // superblock_size
+  1,                            // frame_parallel_decoding_mode
+  NO_AQ,                        // aq_mode
+  0,                            // frame_periodic_delta_q
+  AOM_BITS_8,                   // Bit depth
+  AOM_CONTENT_DEFAULT,          // content
+  AOM_CS_UNKNOWN,               // color space
+  0,                            // color range
+  0,                            // render width
+  0,                            // render height
+  AOM_SUPERBLOCK_SIZE_DYNAMIC,  // superblock_size
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+  23,  // ans_window_size_log2
+#endif
 };
 
 struct aom_codec_alg_priv {
@@ -310,6 +316,9 @@ static aom_codec_err_t validate_config(aom_codec_alg_priv_t *ctx,
   }
   RANGE_CHECK(extra_cfg, color_space, AOM_CS_UNKNOWN, AOM_CS_SRGB);
   RANGE_CHECK(extra_cfg, color_range, 0, 1);
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+  RANGE_CHECK(extra_cfg, ans_window_size_log2, 8, 23);
+#endif
   return AOM_CODEC_OK;
 }
 
@@ -475,6 +484,9 @@ static aom_codec_err_t set_encoder_config(
 #if CONFIG_EXT_PARTITION
   oxcf->superblock_size = extra_cfg->superblock_size;
 #endif  // CONFIG_EXT_PARTITION
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+  oxcf->ans_window_size_log2 = extra_cfg->ans_window_size_log2;
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
 
 #if CONFIG_EXT_TILE
   {
@@ -1334,6 +1346,15 @@ static aom_codec_err_t ctrl_set_superblock_size(aom_codec_alg_priv_t *ctx,
   return update_extra_cfg(ctx, &extra_cfg);
 }
 
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+static aom_codec_err_t ctrl_set_ans_window_size_log2(aom_codec_alg_priv_t *ctx,
+                                                     va_list args) {
+  struct av1_extracfg extra_cfg = ctx->extra_cfg;
+  extra_cfg.ans_window_size_log2 = CAST(AV1E_SET_ANS_WINDOW_SIZE_LOG2, args);
+  return update_extra_cfg(ctx, &extra_cfg);
+}
+#endif
+
 static aom_codec_ctrl_fn_map_t encoder_ctrl_maps[] = {
   { AOM_COPY_REFERENCE, ctrl_copy_reference },
   { AOME_USE_REFERENCE, ctrl_use_reference },
@@ -1384,6 +1405,9 @@ static aom_codec_ctrl_fn_map_t encoder_ctrl_maps[] = {
   { AV1E_SET_MAX_GF_INTERVAL, ctrl_set_max_gf_interval },
   { AV1E_SET_RENDER_SIZE, ctrl_set_render_size },
   { AV1E_SET_SUPERBLOCK_SIZE, ctrl_set_superblock_size },
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+  { AV1E_SET_ANS_WINDOW_SIZE_LOG2, ctrl_set_ans_window_size_log2 },
+#endif
 
   // Getters
   { AOME_GET_LAST_QUANTIZER, ctrl_get_quantizer },

av1/common/onyxc_int.h

@@ -16,6 +16,9 @@
 #include "./av1_rtcd.h"
 #include "aom/internal/aom_codec_internal.h"
 #include "aom_util/aom_thread.h"
+#if CONFIG_ANS
+#include "aom_dsp/ans.h"
+#endif
 #include "av1/common/alloccommon.h"
 #include "av1/common/entropy.h"
 #include "av1/common/entropymode.h"
@@ -422,6 +425,9 @@ typedef struct AV1Common {
   int refresh_mask;
   int invalid_delta_frame_id_minus1;
 #endif
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+  int ans_window_size_log2;
+#endif
 } AV1_COMMON;
 
 #if CONFIG_REFERENCE_BUFFER

av1/decoder/decodeframe.c

@@ -2348,8 +2348,11 @@ static void decode_partition(AV1Decoder *const pbi, MACROBLOCKD *const xd,
 static void setup_bool_decoder(const uint8_t *data, const uint8_t *data_end,
                                const size_t read_size,
                                struct aom_internal_error_info *error_info,
-                               aom_reader *r, aom_decrypt_cb decrypt_cb,
-                               void *decrypt_state) {
+                               aom_reader *r,
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+                               int window_size,
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
+                               aom_decrypt_cb decrypt_cb, void *decrypt_state) {
   // Validate the calculated partition length. If the buffer
   // described by the partition can't be fully read, then restrict
   // it to the portion that can be (for EC mode) or throw an error.
@@ -2358,7 +2361,7 @@ static void setup_bool_decoder(const uint8_t *data, const uint8_t *data_end,
                        "Truncated packet or corrupt tile length");
 
 #if CONFIG_ANS && ANS_MAX_SYMBOLS
-  r->window_size = ANS_MAX_SYMBOLS;
+  r->window_size = window_size;
 #endif
   if (aom_reader_init(r, data, read_size, decrypt_cb, decrypt_state))
     aom_internal_error(error_info, AOM_CODEC_MEM_ERROR,
@@ -3393,7 +3396,11 @@ static const uint8_t *decode_tiles(AV1Decoder *pbi, const uint8_t *data,
 #endif
       av1_tile_init(&td->xd.tile, td->cm, tile_row, tile_col);
       setup_bool_decoder(buf->data, data_end, buf->size, &cm->error,
-                         &td->bit_reader, pbi->decrypt_cb, pbi->decrypt_state);
+                         &td->bit_reader,
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+                         1 << cm->ans_window_size_log2,
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
+                         pbi->decrypt_cb, pbi->decrypt_state);
 #if CONFIG_ACCOUNTING
       if (pbi->acct_enabled) {
         td->bit_reader.accounting = &pbi->accounting;
@@ -3746,8 +3753,11 @@ static const uint8_t *decode_tiles_mt(AV1Decoder *pbi, const uint8_t *data,
         av1_tile_init(tile_info, cm, tile_row, buf->col);
         av1_tile_init(&twd->xd.tile, cm, tile_row, buf->col);
         setup_bool_decoder(buf->data, data_end, buf->size, &cm->error,
-                           &twd->bit_reader, pbi->decrypt_cb,
-                           pbi->decrypt_state);
+                           &twd->bit_reader,
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+                           1 << cm->ans_window_size_log2,
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
+                           pbi->decrypt_cb, pbi->decrypt_state);
         av1_init_macroblockd(cm, &twd->xd,
 #if CONFIG_PVQ
                              twd->pvq_ref_coeff,
@@ -4010,6 +4020,9 @@ static size_t read_uncompressed_header(AV1Decoder *pbi,
       memset(&cm->ref_frame_map, -1, sizeof(cm->ref_frame_map));
       pbi->need_resync = 0;
     }
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+    cm->ans_window_size_log2 = aom_rb_read_literal(rb, 4) + 8;
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
 #if CONFIG_PALETTE
     cm->allow_screen_content_tools = aom_rb_read_bit(rb);
 #endif  // CONFIG_PALETTE
@@ -4049,6 +4062,9 @@ static size_t read_uncompressed_header(AV1Decoder *pbi,
         memset(&cm->ref_frame_map, -1, sizeof(cm->ref_frame_map));
         pbi->need_resync = 0;
       }
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+      cm->ans_window_size_log2 = aom_rb_read_literal(rb, 4) + 8;
+#endif
     } else if (pbi->need_resync != 1) { /* Skip if need resync */
       pbi->refresh_frame_flags = aom_rb_read_literal(rb, REF_FRAMES);
 
@@ -4380,7 +4396,7 @@ static int read_compressed_header(AV1Decoder *pbi, const uint8_t *data,
 #endif
 
 #if CONFIG_ANS && ANS_MAX_SYMBOLS
-  r.window_size = ANS_MAX_SYMBOLS;
+  r.window_size = 1 << cm->ans_window_size_log2;
 #endif
   if (aom_reader_init(&r, data, partition_size, pbi->decrypt_cb,
                       pbi->decrypt_state))

av1/encoder/bitstream.c

@@ -4124,6 +4124,11 @@ static void write_uncompressed_header(AV1_COMP *cpi,
     write_sync_code(wb);
     write_bitdepth_colorspace_sampling(cm, wb);
     write_frame_size(cm, wb);
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+    assert(cpi->common.ans_window_size_log2 >= 8);
+    assert(cpi->common.ans_window_size_log2 < 24);
+    aom_wb_write_literal(wb, cpi->common.ans_window_size_log2 - 8, 4);
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
 #if CONFIG_PALETTE
     aom_wb_write_bit(wb, cm->allow_screen_content_tools);
 #endif  // CONFIG_PALETTE
@@ -4159,6 +4164,12 @@ static void write_uncompressed_header(AV1_COMP *cpi,
       aom_wb_write_literal(wb, get_refresh_mask(cpi), REF_FRAMES);
 #endif  // CONFIG_EXT_REFS
       write_frame_size(cm, wb);
+
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+      assert(cpi->common.ans_window_size_log2 >= 8);
+      assert(cpi->common.ans_window_size_log2 < 24);
+      aom_wb_write_literal(wb, cpi->common.ans_window_size_log2 - 8, 4);
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
     } else {
       MV_REFERENCE_FRAME ref_frame;
 

av1/encoder/encoder.c

@@ -799,9 +799,8 @@ void av1_alloc_compressor_data(AV1_COMP *cpi) {
     unsigned int tokens = get_token_alloc(cm->mb_rows, cm->mb_cols);
     CHECK_MEM_ERROR(cm, cpi->tile_tok[0][0],
                     aom_calloc(tokens, sizeof(*cpi->tile_tok[0][0])));
-#if CONFIG_ANS
-    aom_buf_ans_alloc(&cpi->buf_ans, &cm->error,
-                      ANS_MAX_SYMBOLS ? ANS_MAX_SYMBOLS : tokens);
+#if CONFIG_ANS && !ANS_MAX_SYMBOLS
+    aom_buf_ans_alloc(&cpi->buf_ans, &cm->error, (int)tokens);
 #endif  // CONFIG_ANS
   }
 
@@ -2034,6 +2033,15 @@ void av1_change_config(struct AV1_COMP *cpi, const AV1EncoderConfig *oxcf) {
 #if CONFIG_AOM_HIGHBITDEPTH
   highbd_set_var_fns(cpi);
 #endif
+
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+  cpi->common.ans_window_size_log2 = cpi->oxcf.ans_window_size_log2;
+  if (cpi->buf_ans.size != (1 << cpi->common.ans_window_size_log2)) {
+    aom_buf_ans_free(&cpi->buf_ans);
+    aom_buf_ans_alloc(&cpi->buf_ans, &cpi->common.error,
+                      1 << cpi->common.ans_window_size_log2);
+  }
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
 }
 
 #ifndef M_LOG2_E

av1/encoder/encoder.h

@@ -23,6 +23,7 @@
 #include "av1/common/onyxc_int.h"
 #include "av1/encoder/aq_cyclicrefresh.h"
 #if CONFIG_ANS
+#include "aom_dsp/ans.h"
 #include "aom_dsp/buf_ans.h"
 #endif
 #include "av1/encoder/context_tree.h"
@@ -264,6 +265,9 @@ typedef struct AV1EncoderConfig {
 #if CONFIG_EXT_PARTITION
   aom_superblock_size_t superblock_size;
 #endif  // CONFIG_EXT_PARTITION
+#if CONFIG_ANS && ANS_MAX_SYMBOLS
+  int ans_window_size_log2;
+#endif  // CONFIG_ANS && ANS_MAX_SYMBOLS
 } AV1EncoderConfig;
 
 static INLINE int is_lossless_requested(const AV1EncoderConfig *cfg) {

test/ans_test.cc

@@ -26,9 +26,8 @@ namespace {
 typedef std::vector<std::pair<uint8_t, bool> > PvVec;
 
 const int kPrintStats = 0;
-// When ANS is windowed use the window size, otherwise use a small value to
-// exercise the buffer growth code
-const int kBufAnsSize = ANS_MAX_SYMBOLS ? ANS_MAX_SYMBOLS : 100;
+// Use a small buffer size to exercise ANS window spills or buffer growth
+const int kBufAnsSize = 1 << 8;
 
 PvVec abs_encode_build_vals(int iters) {
   PvVec ret;