Commit 331c56b1ca for aom
commit 331c56b1cafa306552834e0ffbd4f516ec3e61fd
Author: Ihor Buhaiov <freedbrt@gmail.com>
Date: Mon Oct 26 12:49:22 2026 +0100
Keep the range coder state in registers
Every symbol the bitstream writer codes goes through an out of line
od_ec_encode_cdf_q15() or od_ec_encode_bool_q15(). Each call loads low,
rng and cnt from the od_ec_enc struct, branches on the symbol value,
renormalizes in od_ec_enc_normalize() and stores the state back. The
next symbol needs that state, so the store and the reload sit on the
critical chain, and the branches on coefficient levels and signs
mispredict often. In real time screen encodes 61% to 88% of all symbols
come from av1_write_coeffs_txb().
Add od_ec_enc_state with low, rng and cnt, and inline symbol and bool
coders that work on it. A caller takes the state with
od_ec_enc_get_state(), codes many symbols and puts it back with
od_ec_enc_set_state() before any other call on the encoder. The byte
output part of the old normalize moves to od_ec_enc_flush(), which stays
out of line and runs about once every 50 to 75 symbols. The new symbol
coder computes both ends of the range and selects, so it has no branch
on the symbol. The plain od_ec_encode_cdf_q15() and
od_ec_encode_bool_q15() become wrappers over the same code.
av1_write_coeffs_txb() keeps one state from the eob symbol to the last
sign. It collects the nonzero coefficients without branches, so the sign
loop has no data dependent branch, and writes the base range and golomb
codes with inline helpers. pack_map_tokens() keeps the state over the
palette color indices. update_cdf() now computes both updates and
selects one. The decoder shares it and runs at the same speed.
The output doesn't change: the arithmetic and the symbol order are the
same. Checked with stream hashes on Arm and x86, with MSVC and clang,
and with the new EC_TEST.state_functions_match and
AV1.TestWriteSymbolState tests. After a failed buffer realloc the state
now keeps changing where the old code froze it; od_ec_enc_done() still
returns NULL in both cases.
RT speed 11, screen content, 3 screen recordings, fewer encoder cycles
than before, median of interleaved runs:
Apple M3 Pro, clang 4:2:0 4:4:4
1 thread 2.1% to 4.5% 2.3% to 6.1%
4 threads 2.4% to 3.2% 2.5% to 5.6%
AMD Ryzen 9 5900X, MSVC 4:2:0 4:4:4
1 thread 3.0% to 4.5% 2.9% to 5.6%
4 threads 2.0% to 4.1% 2.1% to 5.8%
The writer itself takes 28% to 32% fewer cycles on the M3 Pro. On a
32 bit Arm phone the coder replay of recorded symbol streams is 1.37x
to 1.40x faster on a Cortex A75 and 1.67x to 1.69x on a Cortex A55,
even though the 64 bit low takes two registers there.
Change-Id: Ib7db1b7a9e6081d9038e2ede9e679f9ff4a2b584
diff --git a/aom_dsp/aom_dsp_common.h b/aom_dsp/aom_dsp_common.h
index ea6a2cecdd..7acd41f6bf 100644
--- a/aom_dsp/aom_dsp_common.h
+++ b/aom_dsp/aom_dsp_common.h
@@ -25,8 +25,10 @@ extern "C" {
#if defined(_MSC_VER)
#define AOM_FORCE_INLINE __forceinline
+#define AOM_NOINLINE __declspec(noinline)
#else
#define AOM_FORCE_INLINE __inline__ __attribute__((always_inline))
+#define AOM_NOINLINE __attribute__((noinline))
#endif
#define PI 3.141592653589793238462643383279502884
diff --git a/aom_dsp/bitwriter.h b/aom_dsp/bitwriter.h
index 35f82daefc..d017da79d6 100644
--- a/aom_dsp/bitwriter.h
+++ b/aom_dsp/bitwriter.h
@@ -104,6 +104,29 @@ static inline void aom_write_symbol(aom_writer *w, int symb, aom_cdf_prob *cdf,
if (w->allow_update_cdf) update_cdf(cdf, symb, nsymbs);
}
+// Same as aom_write_symbol() and aom_write_bit(), but use the coder state in
+// *s. See od_ec_enc_state.
+static AOM_FORCE_INLINE void aom_write_symbol_state(aom_writer *w,
+ od_ec_enc_state *s,
+ int symb, aom_cdf_prob *cdf,
+ int nsymbs) {
+#if CONFIG_BITSTREAM_DEBUG
+ bitstream_queue_push(symb, cdf, nsymbs);
+#endif
+ od_ec_encode_cdf_q15_state(&w->ec, s, symb, cdf, nsymbs);
+ if (w->allow_update_cdf) update_cdf(cdf, symb, nsymbs);
+}
+
+static AOM_FORCE_INLINE void aom_write_bit_state(aom_writer *w,
+ od_ec_enc_state *s, int bit) {
+ const int p = CDF_PROB_TOP / 2;
+#if CONFIG_BITSTREAM_DEBUG
+ aom_cdf_prob cdf[2] = { (aom_cdf_prob)p, 32767 };
+ bitstream_queue_push(bit, cdf, 2);
+#endif
+ od_ec_encode_bool_q15_state(&w->ec, s, bit, p);
+}
+
#ifdef __cplusplus
} // extern "C"
#endif
diff --git a/aom_dsp/entenc.c b/aom_dsp/entenc.c
index a56fb89c49..f717aa5579 100644
--- a/aom_dsp/entenc.c
+++ b/aom_dsp/entenc.c
@@ -11,18 +11,10 @@
#include <stdlib.h>
#include <string.h>
-#include <math.h>
#include <assert.h>
#include "aom_dsp/entenc.h"
#include "aom_dsp/prob.h"
-#if OD_MEASURE_EC_OVERHEAD
-#if !defined(M_LOG2E)
-#define M_LOG2E (1.4426950408889634073599246810019)
-#endif
-#define OD_LOG2(x) (M_LOG2E * log(x))
-#endif // OD_MEASURE_EC_OVERHEAD
-
/*A range encoder.
See entdec.c and the references for implementation details \cite{Mar79,MNW98}.
@@ -48,78 +40,65 @@
URL="http://researchcommons.waikato.ac.nz/bitstream/handle/10289/78/content.pdf"
}*/
-/*Takes updated low and range values, renormalizes them so that
- 32768 <= rng < 65536 (flushing bytes from low to the output buffer if
- necessary), and stores them back in the encoder context.
- low: The new value of low.
- rng: The new value of the range.*/
-static void od_ec_enc_normalize(od_ec_enc *enc, od_ec_enc_window low,
- unsigned rng) {
- int d;
- int c;
- int s;
- if (enc->error) return;
- c = enc->cnt;
- assert(rng <= 65535U);
- /*The number of leading zeros in the 16-bit binary representation of rng.*/
- d = 16 - OD_ILOG_NZ(rng);
- s = c + d;
-
- /* We flush every time "low" cannot safely and efficiently accommodate any
- more data. Overall, c must not exceed 63 at the time of byte flush out. To
- facilitate this, "s" cannot exceed 56-bits because we have to keep 1 byte
- for carry. Also, we need to subtract 16 because we want to keep room for
- the next symbol worth "d"-bits (max 15). An alternate condition would be if
- (e < d), where e = number of leading zeros in "low", indicating there is
- not enough rooom to accommodate "rng" worth of "d"-bits in "low". However,
- this approach needs additional computations: (i) compute "e", (ii) push
- the leading 0x00's as a special case.
- */
- if (s >= 40) { // 56 - 16
- unsigned char *out = enc->buf;
- uint32_t storage = enc->storage;
- uint32_t offs = enc->offs;
- if (offs + 8 > storage) {
- storage = 2 * storage + 8;
- out = (unsigned char *)realloc(out, sizeof(*out) * storage);
- if (out == NULL) {
- enc->error = -1;
- return;
- }
- enc->buf = out;
- enc->storage = storage;
+/*Writes the bytes of low that are ready to the output buffer.
+ od_ec_enc_normalize_state() calls this when s, the number of bits of data
+ in low after the normalization shift, reaches 40.
+ We flush every time "low" cannot safely and efficiently accommodate any
+ more data. Overall, c must not exceed 63 at the time of byte flush out. To
+ facilitate this, "s" cannot exceed 56-bits because we have to keep 1 byte
+ for carry. Also, we need to subtract 16 because we want to keep room for
+ the next symbol worth "d"-bits (max 15). An alternate condition would be if
+ (e < d), where e = number of leading zeros in "low", indicating there is
+ not enough rooom to accommodate "rng" worth of "d"-bits in "low". However,
+ this approach needs additional computations: (i) compute "e", (ii) push
+ the leading 0x00's as a special case.
+ It is kept out of line, so the common path of the callers stays small.
+ low: The new value of low, before the normalization shift.
+ c: The number of bits of data in low before the symbol.
+ s: c plus the normalization shift.
+ Return: low without the bits that were written out.*/
+AOM_NOINLINE od_ec_enc_window od_ec_enc_flush(od_ec_enc *enc,
+ od_ec_enc_window low, int c,
+ int s) {
+ /*After an error the buffer may be too small, so write nothing more.*/
+ if (enc->error) return low;
+ unsigned char *out = enc->buf;
+ uint32_t storage = enc->storage;
+ uint32_t offs = enc->offs;
+ if (offs + 8 > storage) {
+ storage = 2 * storage + 8;
+ out = (unsigned char *)realloc(out, sizeof(*out) * storage);
+ if (out == NULL) {
+ enc->error = -1;
+ return low;
}
- // Need to add 1 byte here since enc->cnt always counts 1 byte less
- // (enc->cnt = -9) to ensure correct operation
- uint8_t num_bytes_ready = (s >> 3) + 1;
-
- // Update "c" to contain the number of non-ready bits in "low". Since "low"
- // has 64-bit capacity, we need to add the (64 - 40) cushion bits and take
- // off the number of ready bits.
- c += 24 - (num_bytes_ready << 3);
+ enc->buf = out;
+ enc->storage = storage;
+ }
+ // Need to add 1 byte here since enc->cnt always counts 1 byte less
+ // (enc->cnt = -9) to ensure correct operation
+ uint8_t num_bytes_ready = (s >> 3) + 1;
- // Prepare "output" and update "low"
- uint64_t output = low >> c;
- low = low & (((uint64_t)1 << c) - 1);
+ // Update "c" to contain the number of non-ready bits in "low". Since "low"
+ // has 64-bit capacity, we need to add the (64 - 40) cushion bits and take
+ // off the number of ready bits.
+ c += 24 - (num_bytes_ready << 3);
- // Prepare data and carry mask
- uint64_t mask = (uint64_t)1 << (num_bytes_ready << 3);
- uint64_t carry = output & mask;
+ // Prepare "output" and update "low"
+ uint64_t output = low >> c;
+ low = low & (((uint64_t)1 << c) - 1);
- mask = mask - 0x01;
- output = output & mask;
+ // Prepare data and carry mask
+ uint64_t mask = (uint64_t)1 << (num_bytes_ready << 3);
+ uint64_t carry = output & mask;
- // Write data in a single operation
- write_enc_data_to_out_buf(out, offs, output, carry, &enc->offs,
- num_bytes_ready);
+ mask = mask - 0x01;
+ output = output & mask;
- // Update state of the encoder: enc->cnt to contain the number of residual
- // bits
- s = c + d - 24;
- }
- enc->low = low << d;
- enc->rng = rng << d;
- enc->cnt = s;
+ // Write data in a single operation
+ write_enc_data_to_out_buf(out, offs, output, carry, &enc->offs,
+ num_bytes_ready);
+ return low;
}
/*Initializes the encoder.
@@ -152,63 +131,13 @@ void od_ec_enc_reset(od_ec_enc *enc) {
/*Frees the buffers used by the encoder.*/
void od_ec_enc_clear(od_ec_enc *enc) { free(enc->buf); }
-/*Encodes a symbol given its frequency in Q15.
- fl: CDF_PROB_TOP minus the cumulative frequency of all symbols that come
- before the one to be encoded.
- fh: CDF_PROB_TOP minus the cumulative frequency of all symbols up to and
- including the one to be encoded.*/
-static void od_ec_encode_q15(od_ec_enc *enc, unsigned fl, unsigned fh, int s,
- int nsyms) {
- od_ec_enc_window l;
- unsigned r;
- unsigned u;
- unsigned v;
- l = enc->low;
- r = enc->rng;
- assert(32768U <= r);
- assert(fh <= fl);
- assert(fl <= 32768U);
- assert(7 - EC_PROB_SHIFT >= 0);
- const int N = nsyms - 1;
- if (fl < CDF_PROB_TOP) {
- u = ((r >> 8) * (uint32_t)(fl >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT)) +
- EC_MIN_PROB * (N - (s - 1));
- v = ((r >> 8) * (uint32_t)(fh >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT)) +
- EC_MIN_PROB * (N - (s + 0));
- l += r - u;
- r = u - v;
- } else {
- r -= ((r >> 8) * (uint32_t)(fh >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT)) +
- EC_MIN_PROB * (N - (s + 0));
- }
- od_ec_enc_normalize(enc, l, r);
-#if OD_MEASURE_EC_OVERHEAD
- enc->entropy -= OD_LOG2((double)(OD_ICDF(fh) - OD_ICDF(fl)) / CDF_PROB_TOP.);
- enc->nb_symbols++;
-#endif
-}
-
/*Encode a single binary value.
val: The value to encode (0 or 1).
f: The probability that the val is one, scaled by 32768.*/
void od_ec_encode_bool_q15(od_ec_enc *enc, int val, unsigned f) {
- od_ec_enc_window l;
- unsigned r;
- unsigned v;
- assert(0 < f);
- assert(f < 32768U);
- l = enc->low;
- r = enc->rng;
- assert(32768U <= r);
- v = ((r >> 8) * (uint32_t)(f >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT));
- v += EC_MIN_PROB;
- if (val) l += r - v;
- r = val ? v : r - v;
- od_ec_enc_normalize(enc, l, r);
-#if OD_MEASURE_EC_OVERHEAD
- enc->entropy -= OD_LOG2((double)(val ? f : (32768 - f)) / 32768.);
- enc->nb_symbols++;
-#endif
+ od_ec_enc_state st = od_ec_enc_get_state(enc);
+ od_ec_encode_bool_q15_state(enc, &st, val, f);
+ od_ec_enc_set_state(enc, &st);
}
/*Encodes a symbol given a cumulative distribution function (CDF) table in Q15.
@@ -221,11 +150,9 @@ void od_ec_encode_bool_q15(od_ec_enc *enc, int val, unsigned f) {
This should be at most 16.*/
void od_ec_encode_cdf_q15(od_ec_enc *enc, int s, const uint16_t *icdf,
int nsyms) {
- (void)nsyms;
- assert(s >= 0);
- assert(s < nsyms);
- assert(icdf[nsyms - 1] == OD_ICDF(CDF_PROB_TOP));
- od_ec_encode_q15(enc, s > 0 ? icdf[s - 1] : OD_ICDF(0), icdf[s], s, nsyms);
+ od_ec_enc_state st = od_ec_enc_get_state(enc);
+ od_ec_encode_cdf_q15_state(enc, &st, s, icdf, nsyms);
+ od_ec_enc_set_state(enc, &st);
}
#if OD_MEASURE_EC_OVERHEAD
diff --git a/aom_dsp/entenc.h b/aom_dsp/entenc.h
index dffe32dae5..3976b3ad83 100644
--- a/aom_dsp/entenc.h
+++ b/aom_dsp/entenc.h
@@ -25,6 +25,14 @@ typedef struct od_ec_enc od_ec_enc;
#define OD_MEASURE_EC_OVERHEAD (0)
+#if OD_MEASURE_EC_OVERHEAD
+#include <math.h>
+#if !defined(M_LOG2E)
+#define M_LOG2E (1.4426950408889634073599246810019)
+#endif
+#define OD_LOG2(x) (M_LOG2E * log(x))
+#endif // OD_MEASURE_EC_OVERHEAD
+
/*The entropy encoder context.*/
struct od_ec_enc {
/*Buffered output.
@@ -49,6 +57,20 @@ struct od_ec_enc {
#endif
};
+/*The part of the encoder context that changes with every symbol.
+ A caller that writes many symbols in a row can keep a copy of it in local
+ variables, so it stays in registers: take it with od_ec_enc_get_state(),
+ use the od_ec_encode_*_state() functions and put it back with
+ od_ec_enc_set_state() before any other call on the encoder.*/
+typedef struct od_ec_enc_state {
+ /*The low end of the current range.*/
+ od_ec_enc_window low;
+ /*The number of values in the current range.*/
+ unsigned rng;
+ /*The number of bits of data in the current value.*/
+ int cnt;
+} od_ec_enc_state;
+
/*See entenc.c for further documentation.*/
void od_ec_enc_init(od_ec_enc *enc, uint32_t size) OD_ARG_NONNULL(1);
@@ -98,6 +120,120 @@ static inline void write_enc_data_to_out_buf(unsigned char *out, uint32_t offs,
*enc_offs = offs + num_bytes_ready;
}
+OD_WARN_UNUSED_RESULT od_ec_enc_window od_ec_enc_flush(od_ec_enc *enc,
+ od_ec_enc_window low,
+ int c, int s)
+ OD_ARG_NONNULL(1);
+
+static inline od_ec_enc_state od_ec_enc_get_state(const od_ec_enc *enc) {
+ od_ec_enc_state st;
+ st.low = enc->low;
+ st.rng = enc->rng;
+ st.cnt = enc->cnt;
+ return st;
+}
+
+static inline void od_ec_enc_set_state(od_ec_enc *enc,
+ const od_ec_enc_state *st) {
+ enc->low = st->low;
+ enc->rng = (uint16_t)st->rng;
+ enc->cnt = (int16_t)st->cnt;
+}
+
+/*Takes updated low and range values, renormalizes them so that
+ 32768 <= rng < 65536 (flushing bytes from low to the output buffer if
+ necessary), and stores them back in the state.
+ low: The new value of low.
+ rng: The new value of the range.*/
+static AOM_FORCE_INLINE void od_ec_enc_normalize_state(od_ec_enc *enc,
+ od_ec_enc_state *st,
+ od_ec_enc_window low,
+ unsigned rng) {
+ assert(rng > 0 && rng <= 65535U);
+ /*The number of leading zeros in the 16-bit binary representation of rng.
+ get_msb(rng) is at most 15, so the xor is 15 - get_msb(rng).*/
+ const int d = 15 ^ get_msb(rng);
+ int s = st->cnt + d;
+ /*The limit is explained in od_ec_enc_flush().*/
+ if (UNLIKELY(s >= 40)) {
+ low = od_ec_enc_flush(enc, low, st->cnt, s);
+ /*The flush leaves s - 8 * ((s >> 3) + 1) bits.*/
+ s = (s & 7) - 8;
+ }
+ st->low = low << d;
+ st->rng = rng << d;
+ st->cnt = s;
+}
+
+/*Encode a single binary value.
+ val: The value to encode (0 or 1).
+ f: The probability that the val is one, scaled by 32768.*/
+static AOM_FORCE_INLINE void od_ec_encode_bool_q15_state(od_ec_enc *enc,
+ od_ec_enc_state *st,
+ int val, unsigned f) {
+ od_ec_enc_window l;
+ unsigned r;
+ unsigned v;
+ assert(0 < f);
+ assert(f < 32768U);
+ l = st->low;
+ r = st->rng;
+ assert(32768U <= r);
+ v = ((r >> 8) * (uint32_t)(f >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT));
+ v += EC_MIN_PROB;
+ if (val) l += r - v;
+ r = val ? v : r - v;
+ od_ec_enc_normalize_state(enc, st, l, r);
+#if OD_MEASURE_EC_OVERHEAD
+ enc->entropy -= OD_LOG2((double)(val ? f : (32768 - f)) / 32768.);
+ enc->nb_symbols++;
+#endif
+}
+
+/*Encodes a symbol given a cumulative distribution function (CDF) table in Q15,
+ like od_ec_encode_cdf_q15().
+ No branch on s, because coefficient levels are hard to predict.
+ s: The index of the symbol to encode.
+ icdf: 32768 minus the CDF, such that symbol s falls in the range
+ [s > 0 ? (32768 - icdf[s - 1]) : 0, 32768 - icdf[s]).
+ The values must be monotonically decreasing, and icdf[nsyms - 1] must
+ be 0.
+ nsyms: The number of symbols in the alphabet.
+ This should be at most 16.*/
+static AOM_FORCE_INLINE void od_ec_encode_cdf_q15_state(od_ec_enc *enc,
+ od_ec_enc_state *st,
+ int s,
+ const uint16_t *icdf,
+ int nsyms) {
+ const unsigned r = st->rng;
+ assert(s >= 0);
+ assert(s < nsyms);
+ assert(icdf[nsyms - 1] == OD_ICDF(CDF_PROB_TOP));
+ assert(32768U <= r);
+ assert(7 - EC_PROB_SHIFT >= 0);
+ const int N = nsyms - 1;
+ /*icdf[s - 1] for s > 0, and a valid entry to read for s == 0.*/
+ const unsigned fl = icdf[s - (s > 0)];
+ const unsigned fh = icdf[s];
+ const unsigned r8 = r >> 8;
+ const unsigned pu =
+ r8 * (uint32_t)(fl >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT);
+ const unsigned pv =
+ r8 * (uint32_t)(fh >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT);
+ /*s spans [u, v) below rng, with u = pu + EC_MIN_PROB * (N - s + 1) and
+ v = pv + EC_MIN_PROB * (N - s). um is u without their common term. For
+ s == 0 (or fl == CDF_PROB_TOP) u is rng, so low stays as it is. Doing
+ the select first keeps the chain from rng to the new rng short.*/
+ const unsigned rk = r - EC_MIN_PROB * (N - s);
+ const unsigned um = (s > 0 && fl < CDF_PROB_TOP) ? pu + EC_MIN_PROB : rk;
+ od_ec_enc_normalize_state(enc, st, st->low + (rk - um), um - pv);
+#if OD_MEASURE_EC_OVERHEAD
+ enc->entropy -= OD_LOG2(
+ (double)(OD_ICDF(fh) - OD_ICDF(s > 0 ? fl : OD_ICDF(0))) / CDF_PROB_TOP);
+ enc->nb_symbols++;
+#endif
+}
+
#ifdef __cplusplus
} // extern "C"
#endif
diff --git a/aom_dsp/prob.h b/aom_dsp/prob.h
index 8bec067ca8..5d6903c5b7 100644
--- a/aom_dsp/prob.h
+++ b/aom_dsp/prob.h
@@ -128,11 +128,10 @@ static inline void update_cdf(aom_cdf_prob *cdf, int8_t val, int nsymbs) {
int i = 0;
do {
- if (i < val) {
- cdf[i] += (CDF_PROB_TOP - cdf[i]) >> rate;
- } else {
- cdf[i] -= cdf[i] >> rate;
- }
+ // Compute both and select, a branch on val is hard to predict.
+ const int up = cdf[i] + ((CDF_PROB_TOP - cdf[i]) >> rate);
+ const int down = cdf[i] - (cdf[i] >> rate);
+ cdf[i] = (aom_cdf_prob)(i < val ? up : down);
} while (++i < nsymbs - 1);
cdf[nsymbs] += (count < 32);
}
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index efac97a036..991ba72590 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -357,13 +357,15 @@ static inline void pack_map_tokens(aom_writer *w, const TokenExtra **tp, int n,
write_uniform(w, n, p->token); // The first color index.
++p;
--num;
+ od_ec_enc_state st = od_ec_enc_get_state(&w->ec);
for (int i = 0; i < num; ++i) {
assert((p->color_ctx >= 0) &&
(p->color_ctx < PALETTE_COLOR_INDEX_CONTEXTS));
aom_cdf_prob *color_map_cdf = map_pb_cdf[palette_size_idx][p->color_ctx];
- aom_write_symbol(w, p->token, color_map_cdf, n);
+ aom_write_symbol_state(w, &st, p->token, color_map_cdf, n);
++p;
}
+ od_ec_enc_set_state(&w->ec, &st);
*tp = p;
}
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index baedd1320d..41246d6973 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -115,7 +115,8 @@ void av1_free_txb_buf(AV1_COMP *cpi) {
coeff_buf_pool->entropy_ctx = NULL;
}
-static void write_golomb(aom_writer *w, int level) {
+static AOM_FORCE_INLINE void write_golomb(aom_writer *w, od_ec_enc_state *s,
+ int level) {
int x = level + 1;
int i = x;
int length = 0;
@@ -126,9 +127,20 @@ static void write_golomb(aom_writer *w, int level) {
}
assert(length > 0);
- for (i = 0; i < length - 1; ++i) aom_write_bit(w, 0);
+ for (i = 0; i < length - 1; ++i) aom_write_bit_state(w, s, 0);
- for (i = length - 1; i >= 0; --i) aom_write_bit(w, (x >> i) & 0x01);
+ for (i = length - 1; i >= 0; --i) aom_write_bit_state(w, s, (x >> i) & 0x01);
+}
+
+// Writes the base range symbols of a level above NUM_BASE_LEVELS.
+static AOM_FORCE_INLINE void write_coeff_br(aom_writer *w, od_ec_enc_state *s,
+ aom_cdf_prob *cdf, int level) {
+ const int base_range = level - 1 - NUM_BASE_LEVELS;
+ for (int idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
+ const int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);
+ aom_write_symbol_state(w, s, k, cdf, BR_CDF_SIZE);
+ if (k < BR_CDF_SIZE - 1) break;
+ }
}
int av1_get_eob_pos_token(const int eob, int *const extra) {
@@ -306,118 +318,135 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *const x,
const int eob_multi_size = txsize_log2_minus4[tx_size];
const TX_CLASS tx_class = tx_type_to_class[tx_type];
const int eob_multi_ctx = (tx_class == TX_CLASS_2D) ? 0 : 1;
+
+ const int width = get_txb_wide(tx_size);
+ const int height = get_txb_high(tx_size);
+ uint8_t levels_buf[TX_PAD_2D];
+ uint8_t *const levels = set_levels(levels_buf, height);
+ const tran_low_t *tcoeff_txb =
+ cb_coef_buff->tcoeff[plane] + x->mbmi_ext_frame->cb_offset[plane_type];
+ const tran_low_t *tcoeff = tcoeff_txb + BLOCK_OFFSET(block);
+ const SCAN_ORDER *const scan_order = get_scan(tx_size, tx_type);
+ const int16_t *const scan = scan_order->scan;
+ DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
+ if (eob > 1) {
+ av1_txb_init_levels(tcoeff, width, height, levels);
+ av1_get_nz_map_contexts(levels, scan, eob, tx_size, tx_class,
+ coeff_contexts);
+ } else {
+ coeff_contexts[0] = 0;
+ }
+
+ // The rest of the block only writes symbols, so keep the range coder state
+ // in registers.
+ od_ec_enc_state st = od_ec_enc_get_state(&w->ec);
+ aom_cdf_prob *eob_cdf;
+ int eob_nsymbs;
switch (eob_multi_size) {
case 0:
- aom_write_symbol(w, eob_pt - 1,
- ec_ctx->eob_flag_cdf16[plane_type][eob_multi_ctx], 5);
+ eob_cdf = ec_ctx->eob_flag_cdf16[plane_type][eob_multi_ctx];
+ eob_nsymbs = 5;
break;
case 1:
- aom_write_symbol(w, eob_pt - 1,
- ec_ctx->eob_flag_cdf32[plane_type][eob_multi_ctx], 6);
+ eob_cdf = ec_ctx->eob_flag_cdf32[plane_type][eob_multi_ctx];
+ eob_nsymbs = 6;
break;
case 2:
- aom_write_symbol(w, eob_pt - 1,
- ec_ctx->eob_flag_cdf64[plane_type][eob_multi_ctx], 7);
+ eob_cdf = ec_ctx->eob_flag_cdf64[plane_type][eob_multi_ctx];
+ eob_nsymbs = 7;
break;
case 3:
- aom_write_symbol(w, eob_pt - 1,
- ec_ctx->eob_flag_cdf128[plane_type][eob_multi_ctx], 8);
+ eob_cdf = ec_ctx->eob_flag_cdf128[plane_type][eob_multi_ctx];
+ eob_nsymbs = 8;
break;
case 4:
- aom_write_symbol(w, eob_pt - 1,
- ec_ctx->eob_flag_cdf256[plane_type][eob_multi_ctx], 9);
+ eob_cdf = ec_ctx->eob_flag_cdf256[plane_type][eob_multi_ctx];
+ eob_nsymbs = 9;
break;
case 5:
- aom_write_symbol(w, eob_pt - 1,
- ec_ctx->eob_flag_cdf512[plane_type][eob_multi_ctx], 10);
+ eob_cdf = ec_ctx->eob_flag_cdf512[plane_type][eob_multi_ctx];
+ eob_nsymbs = 10;
break;
default:
- aom_write_symbol(w, eob_pt - 1,
- ec_ctx->eob_flag_cdf1024[plane_type][eob_multi_ctx], 11);
+ eob_cdf = ec_ctx->eob_flag_cdf1024[plane_type][eob_multi_ctx];
+ eob_nsymbs = 11;
break;
}
+ aom_write_symbol_state(w, &st, eob_pt - 1, eob_cdf, eob_nsymbs);
const int eob_offset_bits = av1_eob_offset_bits[eob_pt];
if (eob_offset_bits > 0) {
const int eob_ctx = eob_pt - 3;
int eob_shift = eob_offset_bits - 1;
int bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
- aom_write_symbol(w, bit,
- ec_ctx->eob_extra_cdf[txs_ctx][plane_type][eob_ctx], 2);
+ aom_write_symbol_state(
+ w, &st, bit, ec_ctx->eob_extra_cdf[txs_ctx][plane_type][eob_ctx], 2);
for (int i = 1; i < eob_offset_bits; i++) {
eob_shift = eob_offset_bits - 1 - i;
bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
- aom_write_bit(w, bit);
+ aom_write_bit_state(w, &st, bit);
}
}
- const int width = get_txb_wide(tx_size);
- const int height = get_txb_high(tx_size);
- uint8_t levels_buf[TX_PAD_2D];
- uint8_t *const levels = set_levels(levels_buf, height);
- const tran_low_t *tcoeff_txb =
- cb_coef_buff->tcoeff[plane] + x->mbmi_ext_frame->cb_offset[plane_type];
- const tran_low_t *tcoeff = tcoeff_txb + BLOCK_OFFSET(block);
- const SCAN_ORDER *const scan_order = get_scan(tx_size, tx_type);
- const int16_t *const scan = scan_order->scan;
- DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
- if (eob > 1) {
- av1_txb_init_levels(tcoeff, width, height, levels);
- av1_get_nz_map_contexts(levels, scan, eob, tx_size, tx_class,
- coeff_contexts);
- } else {
- coeff_contexts[0] = 0;
- }
-
const int bhl = get_txb_bhl(tx_size);
- for (int c = eob - 1; c >= 0; --c) {
- const int pos = scan[c];
- const int coeff_ctx = coeff_contexts[pos];
+ aom_cdf_prob(*const base_cdf)[CDF_SIZE(4)] =
+ ec_ctx->coeff_base_cdf[txs_ctx][plane_type];
+ aom_cdf_prob(*const br_cdf)[CDF_SIZE(BR_CDF_SIZE)] =
+ ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_32X32)][plane_type];
+ // The nonzero coefficients from c = eob - 1 down, for the sign loop.
+ tran_low_t nz_coeffs[32 * 32];
+ assert(eob <= av1_get_max_eob(tx_size));
+
+ int num_nz = 0;
+ {
+ // The last coefficient is never 0 and has its own cdf.
+ const int pos = scan[eob - 1];
const tran_low_t v = tcoeff[pos];
const tran_low_t level = abs(v);
-
- if (c == eob - 1) {
- aom_write_symbol(
- w, AOMMIN(level, 3) - 1,
- ec_ctx->coeff_base_eob_cdf[txs_ctx][plane_type][coeff_ctx], 3);
- } else {
- aom_write_symbol(w, AOMMIN(level, 3),
- ec_ctx->coeff_base_cdf[txs_ctx][plane_type][coeff_ctx],
- 4);
+ nz_coeffs[num_nz++] = v;
+ aom_write_symbol_state(
+ w, &st, AOMMIN(level, 3) - 1,
+ ec_ctx->coeff_base_eob_cdf[txs_ctx][plane_type][coeff_contexts[pos]],
+ 3);
+ if (level > NUM_BASE_LEVELS) {
+ write_coeff_br(w, &st, br_cdf[get_br_ctx_eob(pos, bhl, tx_class)], level);
}
+ }
+ for (int c = eob - 2; c >= 0; --c) {
+ const int pos = scan[c];
+ const tran_low_t v = tcoeff[pos];
+ const tran_low_t level = abs(v);
+ nz_coeffs[num_nz] = v;
+ num_nz += (v != 0);
+ aom_write_symbol_state(w, &st, AOMMIN(level, 3),
+ base_cdf[coeff_contexts[pos]], 4);
if (level > NUM_BASE_LEVELS) {
- // level is above 1.
- const int base_range = level - 1 - NUM_BASE_LEVELS;
- const int br_ctx = (eob > 1) ? get_br_ctx(levels, pos, bhl, tx_class) : 0;
- aom_cdf_prob *cdf =
- ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_32X32)][plane_type][br_ctx];
- for (int idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
- const int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);
- aom_write_symbol(w, k, cdf, BR_CDF_SIZE);
- if (k < BR_CDF_SIZE - 1) break;
- }
+ write_coeff_br(w, &st, br_cdf[get_br_ctx(levels, pos, bhl, tx_class)],
+ level);
}
}
- // Loop to code all signs in the transform block,
- // starting with the sign of DC (if applicable)
- for (int c = 0; c < eob; ++c) {
- const tran_low_t v = tcoeff[scan[c]];
+ // Code all signs in the transform block, starting with the sign of DC (if
+ // applicable). nz_coeffs holds them in the reverse order.
+ int i = num_nz - 1;
+ if (tcoeff[scan[0]] != 0) {
+ const tran_low_t v = nz_coeffs[i--];
const tran_low_t level = abs(v);
- const int sign = (v < 0) ? 1 : 0;
- if (level) {
- if (c == 0) {
- const int dc_sign_ctx =
- (entropy_ctx[block] >> DC_SIGN_CTX_SHIFT) & DC_SIGN_CTX_MASK;
- aom_write_symbol(w, sign, ec_ctx->dc_sign_cdf[plane_type][dc_sign_ctx],
- 2);
- } else {
- aom_write_bit(w, sign);
- }
- if (level > COEFF_BASE_RANGE + NUM_BASE_LEVELS)
- write_golomb(w, level - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
- }
+ const int dc_sign_ctx =
+ (entropy_ctx[block] >> DC_SIGN_CTX_SHIFT) & DC_SIGN_CTX_MASK;
+ aom_write_symbol_state(w, &st, v < 0,
+ ec_ctx->dc_sign_cdf[plane_type][dc_sign_ctx], 2);
+ if (level > COEFF_BASE_RANGE + NUM_BASE_LEVELS)
+ write_golomb(w, &st, level - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
+ }
+ for (; i >= 0; --i) {
+ const tran_low_t v = nz_coeffs[i];
+ const tran_low_t level = abs(v);
+ aom_write_bit_state(w, &st, v < 0);
+ if (level > COEFF_BASE_RANGE + NUM_BASE_LEVELS)
+ write_golomb(w, &st, level - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
}
+ od_ec_enc_set_state(&w->ec, &st);
}
void av1_write_intra_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
diff --git a/test/boolcoder_test.cc b/test/boolcoder_test.cc
index 3629bd7140..6e5c8db6b5 100644
--- a/test/boolcoder_test.cc
+++ b/test/boolcoder_test.cc
@@ -13,6 +13,8 @@
#include <stdlib.h>
#include <string.h>
+#include <vector>
+
#include "gtest/gtest.h"
#include "test/acm_random.h"
@@ -171,3 +173,100 @@ TEST(AV1, TestHasOverflowed) {
ASSERT_TRUE(aom_reader_has_overflowed(&br));
}
}
+
+// Fills cdf with a random CDF of nsymbs symbols and a zero counter.
+static void RandomCdf(ACMRandom *rnd, int nsymbs, aom_cdf_prob *cdf) {
+ int prev = 0;
+ for (int i = 0; i < nsymbs - 1; ++i) {
+ // Room for the symbols left, each with at least one unit.
+ const int left = nsymbs - 1 - i;
+ const int span = CDF_PROB_TOP - prev - left;
+ prev += 1 + (*rnd)(span);
+ cdf[i] = AOM_ICDF(prev);
+ }
+ cdf[nsymbs - 1] = AOM_ICDF(CDF_PROB_TOP);
+ cdf[nsymbs] = 0;
+}
+
+// aom_write_symbol_state() and aom_write_bit_state() keep the range coder
+// state in a local copy. They must write the same bytes as
+// aom_write_symbol() and aom_write_bit() and adapt the CDFs the same way.
+TEST(AV1, TestWriteSymbolState) {
+ ACMRandom rnd(ACMRandom::DeterministicSeed());
+ const int kSymbols = 20000;
+ const int kMaxSymbs = 16;
+ const int kNumCdfs = kMaxSymbs - 1;
+ const int kBufferSize = 3 * kSymbols;
+ for (int n = 0; n < num_tests; ++n) {
+ aom_cdf_prob cdf_a[kNumCdfs][CDF_SIZE(kMaxSymbs)];
+ aom_cdf_prob cdf_b[kNumCdfs][CDF_SIZE(kMaxSymbs)];
+ aom_cdf_prob cdf_r[kNumCdfs][CDF_SIZE(kMaxSymbs)];
+ for (int k = 0; k < kNumCdfs; ++k) RandomCdf(&rnd, k + 2, cdf_a[k]);
+ memcpy(cdf_b, cdf_a, sizeof(cdf_a));
+ memcpy(cdf_r, cdf_a, sizeof(cdf_a));
+ const uint8_t allow_update_cdf = n & 1;
+ std::vector<int> kinds(kSymbols);
+ std::vector<int> values(kSymbols);
+ std::vector<uint8_t> buf_a(kBufferSize);
+ std::vector<uint8_t> buf_b(kBufferSize);
+ aom_writer wa;
+ aom_writer wb;
+ aom_start_encode(&wa, buf_a.data(), kBufferSize);
+ aom_start_encode(&wb, buf_b.data(), kBufferSize);
+ wa.allow_update_cdf = allow_update_cdf;
+ wb.allow_update_cdf = allow_update_cdf;
+ od_ec_enc_state st = od_ec_enc_get_state(&wb.ec);
+ for (int i = 0; i < kSymbols; ++i) {
+ // kind 0 is a bit, kind k > 0 a symbol coded with cdf_*[k - 1].
+ const int kind = rnd(kNumCdfs + 5) - 4;
+ kinds[i] = kind > 0 ? kind : 0;
+ // Mix in a write through the plain path now and then. It works on the
+ // state stored in the writer.
+ const bool plain = rnd(100) == 0;
+ if (plain) od_ec_enc_set_state(&wb.ec, &st);
+ if (kinds[i] == 0) {
+ values[i] = rnd(2);
+ aom_write_bit(&wa, values[i]);
+ if (plain) {
+ aom_write_bit(&wb, values[i]);
+ } else {
+ aom_write_bit_state(&wb, &st, values[i]);
+ }
+ } else {
+ const int k = kinds[i] - 1;
+ values[i] = rnd(k + 2);
+ aom_write_symbol(&wa, values[i], cdf_a[k], k + 2);
+ if (plain) {
+ aom_write_symbol(&wb, values[i], cdf_b[k], k + 2);
+ } else {
+ aom_write_symbol_state(&wb, &st, values[i], cdf_b[k], k + 2);
+ }
+ }
+ if (plain) {
+ ASSERT_EQ(aom_tell_size(&wa), aom_tell_size(&wb)) << "symbol " << i;
+ st = od_ec_enc_get_state(&wb.ec);
+ }
+ }
+ od_ec_enc_set_state(&wb.ec, &st);
+ GTEST_ASSERT_GE(aom_stop_encode(&wa), 0);
+ GTEST_ASSERT_GE(aom_stop_encode(&wb), 0);
+ ASSERT_EQ(wa.pos, wb.pos);
+ EXPECT_EQ(memcmp(buf_a.data(), buf_b.data(), wa.pos), 0);
+ EXPECT_EQ(memcmp(cdf_a, cdf_b, sizeof(cdf_a)), 0);
+
+ aom_reader br;
+ aom_reader_init(&br, buf_b.data(), wb.pos);
+ br.allow_update_cdf = allow_update_cdf;
+ for (int i = 0; i < kSymbols; ++i) {
+ if (kinds[i] == 0) {
+ GTEST_ASSERT_EQ(aom_read_bit(&br, nullptr), values[i]) << "pos " << i;
+ } else {
+ const int k = kinds[i] - 1;
+ GTEST_ASSERT_EQ(aom_read_symbol(&br, cdf_r[k], k + 2, nullptr),
+ values[i])
+ << "pos " << i;
+ }
+ }
+ EXPECT_EQ(memcmp(cdf_a, cdf_r, sizeof(cdf_a)), 0);
+ }
+}
diff --git a/test/ec_test.cc b/test/ec_test.cc
index 5586a60aae..65311cbbde 100644
--- a/test/ec_test.cc
+++ b/test/ec_test.cc
@@ -12,6 +12,7 @@
#include "gtest/gtest.h"
#include <cstdlib>
+#include <cstring>
#include <memory>
#include <new>
@@ -143,3 +144,63 @@ TEST(EC_TEST, random_ec_test) {
od_ec_enc_clear(&enc);
EXPECT_EQ(ret, 0);
}
+
+// The od_ec_encode_*_state() functions keep the coder state in a local copy.
+// They must write the same bytes as the plain functions, also when the output
+// buffer has to grow.
+TEST(EC_TEST, state_functions_match) {
+ const unsigned int seed = 0xdaa1b;
+ srand(seed);
+ od_ec_enc enc_a;
+ od_ec_enc enc_b;
+ od_ec_enc_init(&enc_a, 1);
+ od_ec_enc_init(&enc_b, 1);
+ for (int i = 0; i < 400; i++) {
+ od_ec_enc_reset(&enc_a);
+ od_ec_enc_reset(&enc_b);
+ od_ec_enc_state st = od_ec_enc_get_state(&enc_b);
+ const int sz = rand() % 4000;
+ for (int j = 0; j < sz; j++) {
+ if (rand() % 3 == 0) {
+ const int val = rand() & 1;
+ const unsigned f = 1 + rand() % 32767;
+ od_ec_encode_bool_q15(&enc_a, val, f);
+ od_ec_encode_bool_q15_state(&enc_b, &st, val, f);
+ } else {
+ // A random CDF of 2 to 16 symbols. Sometimes the first symbols have
+ // a zero probability, so icdf[s - 1] can be CDF_PROB_TOP.
+ const int nsyms = 2 + rand() % 15;
+ const int zeros = rand() % 8 == 0 ? rand() % nsyms : 0;
+ uint16_t icdf[16];
+ int prev = 0;
+ for (int k = 0; k < nsyms - 1; k++) {
+ if (k >= zeros) {
+ prev += 1 + rand() % (CDF_PROB_TOP - prev - (nsyms - 1 - k));
+ }
+ icdf[k] = OD_ICDF(prev);
+ }
+ icdf[nsyms - 1] = OD_ICDF(CDF_PROB_TOP);
+ const int s = zeros + rand() % (nsyms - zeros);
+ od_ec_encode_cdf_q15(&enc_a, s, icdf, nsyms);
+ od_ec_encode_cdf_q15_state(&enc_b, &st, s, icdf, nsyms);
+ }
+ if (rand() % 64 == 0) {
+ od_ec_enc_set_state(&enc_b, &st);
+ ASSERT_EQ(od_ec_enc_tell_frac(&enc_a), od_ec_enc_tell_frac(&enc_b))
+ << "symbol " << j << " (Random seed: " << seed << ")";
+ st = od_ec_enc_get_state(&enc_b);
+ }
+ }
+ od_ec_enc_set_state(&enc_b, &st);
+ uint32_t sz_a;
+ uint32_t sz_b;
+ unsigned char *ptr_a = od_ec_enc_done(&enc_a, &sz_a);
+ unsigned char *ptr_b = od_ec_enc_done(&enc_b, &sz_b);
+ ASSERT_NE(ptr_a, nullptr);
+ ASSERT_NE(ptr_b, nullptr);
+ ASSERT_EQ(sz_a, sz_b) << "(Random seed: " << seed << ")";
+ EXPECT_EQ(memcmp(ptr_a, ptr_b, sz_a), 0) << "(Random seed: " << seed << ")";
+ }
+ od_ec_enc_clear(&enc_a);
+ od_ec_enc_clear(&enc_b);
+}