pvq: fix typo
[daala.git] / src / pvq_encoder.c
1 /*Daala video codec
2 Copyright (c) 2012 Daala project contributors.  All rights reserved.
3
4 Redistribution and use in source and binary forms, with or without
5 modification, are permitted provided that the following conditions are met:
6
7 - Redistributions of source code must retain the above copyright notice, this
8   list of conditions and the following disclaimer.
9
10 - Redistributions in binary form must reproduce the above copyright notice,
11   this list of conditions and the following disclaimer in the documentation
12   and/or other materials provided with the distribution.
13
14 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
15 AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
16 IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
17 DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
18 FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
19 DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
20 SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
21 CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
22 OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
23 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.*/
24
25 #ifdef HAVE_CONFIG_H
26 # include "config.h"
27 #endif
28
29 #include <stdlib.h>
30 #include <stdio.h>
31 #include "internal.h"
32 #include "block_size.h"
33 #include "entenc.h"
34 #include "entcode.h"
35 #include "filter.h"
36 #include "pvq_encoder.h"
37 #include "partition.h"
38
39 #define OD_PVQ_RATE_APPROX (0)
40 /*Shift to ensure that the upper bound (i.e. for the max blocksize) of the
41    dot-product of the 1st band of chroma with the luma ref doesn't overflow.*/
42 #define OD_CFL_FLIP_SHIFT (OD_LIMIT_BSIZE_MAX + 0)
43
44 static void od_encode_pvq_codeword(od_ec_enc *ec, od_pvq_codeword_ctx *adapt,
45  const od_coeff *in, int n, int k) {
46   int i;
47   od_encode_band_pvq_splits(ec, adapt, in, n, k, 0);
48   for (i = 0; i < n; i++) if (in[i]) od_ec_enc_bits(ec, in[i] < 0, 1);
49 }
50
51 /* Computes 1/sqrt(i) using a table for small values. */
52 static double od_rsqrt_table(int i) {
53   static double table[16] = {
54     1.000000, 0.707107, 0.577350, 0.500000,
55     0.447214, 0.408248, 0.377964, 0.353553,
56     0.333333, 0.316228, 0.301511, 0.288675,
57     0.277350, 0.267261, 0.258199, 0.250000};
58   if (i <= 16) return table[i-1];
59   else return 1./sqrt(i);
60 }
61
62 /*Computes 1/sqrt(start+2*i+1) using a lookup table containing the results
63    where 0 <= i < table_size.*/
64 static double od_custom_rsqrt_dynamic_table(const double* table,
65  const int table_size, const double start, const int i) {
66   if (i < table_size) return table[i];
67   else return od_rsqrt_table(start + 2*i + 1);
68 }
69
70 /*Fills tables used in od_custom_rsqrt_dynamic_table for a given start.*/
71 static void od_fill_dynamic_rsqrt_table(double *table, const int table_size,
72  const double start) {
73   int i;
74   for (i = 0; i < table_size; i++)
75     table[i] = od_rsqrt_table(start + 2*i + 1);
76 }
77
78 /** Find the codepoint on the given PSphere closest to the desired
79  * vector. Double-precision PVQ search just to make sure our tests
80  * aren't limited by numerical accuracy.
81  *
82  * @param [in]      xcoeff  input vector to quantize (x in the math doc)
83  * @param [in]      n       number of dimensions
84  * @param [in]      k       number of pulses
85  * @param [out]     ypulse  optimal codevector found (y in the math doc)
86  * @param [out]     g2      multiplier for the distortion (typically squared
87  *                          gain units)
88  * @param [in] pvq_norm_lambda enc->pvq_norm_lambda for quantized RDO
89  * @param [in]      prev_k  number of pulses already in ypulse that we should
90  *                          reuse for the search (or 0 for a new search)
91  * @return                  cosine distance between x and y (between 0 and 1)
92  */
93 static double pvq_search_rdo_double(const od_val16 *xcoeff, int n, int k,
94  od_coeff *ypulse, double g2, double pvq_norm_lambda, int prev_k) {
95   int i, j;
96   double xy;
97   double yy;
98   /* TODO - This blows our 8kB stack space budget and should be fixed when
99    converting PVQ to fixed point. */
100   double x[MAXN];
101   double xx;
102   double lambda;
103   double norm_1;
104   int rdo_pulses;
105   double delta_rate;
106   xx = xy = yy = 0;
107   for (j = 0; j < n; j++) {
108     x[j] = fabs((float)xcoeff[j]);
109     xx += x[j]*x[j];
110   }
111   norm_1 = 1./sqrt(1e-30 + xx);
112   lambda = pvq_norm_lambda/(1e-30 + g2);
113   i = 0;
114   if (prev_k > 0 && prev_k <= k) {
115     /* We reuse pulses from a previous search so we don't have to search them
116        again. */
117     for (j = 0; j < n; j++) {
118       ypulse[j] = abs(ypulse[j]);
119       xy += x[j]*ypulse[j];
120       yy += ypulse[j]*ypulse[j];
121       i += ypulse[j];
122     }
123   }
124   else if (k > 2) {
125     double l1_norm;
126     double l1_inv;
127     l1_norm = 0;
128     for (j = 0; j < n; j++) l1_norm += x[j];
129     l1_inv = 1./OD_MAXF(l1_norm, 1e-100);
130     for (j = 0; j < n; j++) {
131       double tmp;
132       tmp = k*x[j]*l1_inv;
133       ypulse[j] = OD_MAXI(0, (int)floor(tmp));
134       xy += x[j]*ypulse[j];
135       yy += ypulse[j]*ypulse[j];
136       i += ypulse[j];
137     }
138   }
139   else OD_CLEAR(ypulse, n);
140
141   /* Only use RDO on the last few pulses. This not only saves CPU, but using
142      RDO on all pulses actually makes the results worse for reasons I don't
143      fully understand. */
144   rdo_pulses = 1 + k/4;
145   /* Rough assumption for now, the last position costs about 3 bits more than
146      the first. */
147   delta_rate = 3./n;
148   /* Search one pulse at a time */
149   for (; i < k - rdo_pulses; i++) {
150     int pos;
151     double best_xy;
152     double best_yy;
153     pos = 0;
154     best_xy = -10;
155     best_yy = 1;
156     for (j = 0; j < n; j++) {
157       double tmp_xy;
158       double tmp_yy;
159       tmp_xy = xy + x[j];
160       tmp_yy = yy + 2*ypulse[j] + 1;
161       tmp_xy *= tmp_xy;
162       if (j == 0 || tmp_xy*best_yy > best_xy*tmp_yy) {
163         best_xy = tmp_xy;
164         best_yy = tmp_yy;
165         pos = j;
166       }
167     }
168     xy = xy + x[pos];
169     yy = yy + 2*ypulse[pos] + 1;
170     ypulse[pos]++;
171   }
172   /* Search last pulses with RDO. Distortion is D = (x-y)^2 = x^2 - 2*x*y + y^2
173      and since x^2 and y^2 are constant, we just maximize x*y, plus a
174      lambda*rate term. Note that since x and y aren't normalized here,
175      we need to divide by sqrt(x^2)*sqrt(y^2). */
176   for (; i < k; i++) {
177     double rsqrt_table[4];
178     int rsqrt_table_size = 4;
179     int pos;
180     double best_cost;
181     pos = 0;
182     best_cost = -1e5;
183     /*Fill the small rsqrt lookup table with inputs relative to yy.
184       Specifically, the table of n values is filled with
185        rsqrt(yy + 1), rsqrt(yy + 2 + 1) .. rsqrt(yy + 2*(n-1) + 1).*/
186     od_fill_dynamic_rsqrt_table(rsqrt_table, rsqrt_table_size, yy);
187     for (j = 0; j < n; j++) {
188       double tmp_xy;
189       double tmp_yy;
190       tmp_xy = xy + x[j];
191       /*Calculate rsqrt(yy + 2*ypulse[j] + 1) using an optimized method.*/
192       tmp_yy = od_custom_rsqrt_dynamic_table(rsqrt_table, rsqrt_table_size,
193        yy, ypulse[j]);
194       tmp_xy = 2*tmp_xy*norm_1*tmp_yy - lambda*j*delta_rate;
195       if (j == 0 || tmp_xy > best_cost) {
196         best_cost = tmp_xy;
197         pos = j;
198       }
199     }
200     xy = xy + x[pos];
201     yy = yy + 2*ypulse[pos] + 1;
202     ypulse[pos]++;
203   }
204   for (i = 0; i < n; i++) {
205     if (xcoeff[i] < 0) ypulse[i] = -ypulse[i];
206   }
207   return xy/(1e-100 + sqrt(xx*yy));
208 }
209
210 /** Encodes the gain so that the return value increases with the
211  * distance |x-ref|, so that we can encode a zero when x=ref. The
212  * value x=0 is not covered because it is only allowed in the noref
213  * case.
214  *
215  * @param [in]      x      quantized gain to encode
216  * @param [in]      ref    quantized gain of the reference
217  * @return                 interleave-encoded quantized gain value
218  */
219 static int neg_interleave(int x, int ref) {
220   if (x < ref) return -2*(x - ref) - 1;
221   else if (x < 2*ref) return 2*(x - ref);
222   else return x-1;
223 }
224
225 int od_vector_is_null(const od_coeff *x, int len) {
226   int i;
227   for (i = 0; i < len; i++) if (x[i]) return 0;
228   return 1;
229 }
230
231 static double od_pvq_rate(int qg, int icgr, int theta, int ts,
232  const od_adapt_ctx *adapt, const od_coeff *y0, int k, int n,
233  int is_keyframe, int pli, int speed) {
234   double rate;
235   if (k == 0) rate = 0;
236   else if (speed > 0) {
237     int i;
238     int sum;
239     double f;
240     /* Compute "center of mass" of the pulse vector. */
241     sum = 0;
242     for (i = 0; i < n - (theta != -1); i++) sum += i*abs(y0[i]);
243     f = sum/(double)(k*n);
244     /* Estimates the number of bits it will cost to encode K pulses in
245        N dimensions based on hand-tuned fit for bitrate vs K, N and
246        "center of mass". */
247     rate = (1 + .4*f)*n*OD_LOG2(1 + OD_MAXF(0, log(n*2*(1*f + .025))*k/n)) + 3;
248   }
249   else {
250     od_ec_enc ec;
251     od_pvq_codeword_ctx cd;
252     int tell;
253     od_ec_enc_init(&ec, 1000);
254     OD_COPY(&cd, &adapt->pvq.pvq_codeword_ctx, 1);
255     tell = od_ec_enc_tell_frac(&ec);
256     od_encode_pvq_codeword(&ec, &cd, y0, n - (theta != -1), k);
257     rate = (od_ec_enc_tell_frac(&ec)-tell)/8.;
258     od_ec_enc_clear(&ec);
259   }
260   if (qg > 0 && theta >= 0) {
261     /* Approximate cost of entropy-coding theta */
262     rate += .9*OD_LOG2(ts);
263     /* Adding a cost to using the H/V pred because it's going to be off
264        most of the time. Cost is optimized on subset1, while making
265        sure we don't hurt the checkerboard image too much.
266        FIXME: Do real RDO instead of this arbitrary cost. */
267     if (is_keyframe && pli == 0) rate += 6;
268     if (qg == icgr) rate -= .5;
269   }
270   return rate;
271 }
272
273 #define MAX_PVQ_ITEMS (20)
274 /* This stores the information about a PVQ search candidate, so we can sort
275    based on K. */
276 typedef struct {
277   int gain;
278   int k;
279   od_val32 qtheta;
280   int theta;
281   int ts;
282   od_val32 qcg;
283 } pvq_search_item;
284
285 int items_compare(pvq_search_item *a, pvq_search_item *b) {
286   return a->k - b->k;
287 }
288
289 /** Perform PVQ quantization with prediction, trying several
290  * possible gains and angles. See draft-valin-videocodec-pvq and
291  * http://jmvalin.ca/slides/pvq.pdf for more details.
292  *
293  * @param [out]    out       coefficients after quantization
294  * @param [in]     x0        coefficients before quantization
295  * @param [in]     r0        reference, aka predicted coefficients
296  * @param [in]     n         number of dimensions
297  * @param [in]     q0        quantization step size
298  * @param [out]    y         pulse vector (i.e. selected PVQ codevector)
299  * @param [out]    itheta    angle between input and reference (-1 if noref)
300  * @param [out]    max_theta maximum value of itheta that could have been
301  * @param [out]    vk        total number of pulses
302  * @param [in]     beta      per-band activity masking beta param
303  * @param [out]    skip_diff distortion cost of skipping this block
304  *                           (accumulated)
305  * @param [in]     robust    make stream robust to error in the reference
306  * @param [in]     is_keyframe whether we're encoding a keyframe
307  * @param [in]     pli       plane index
308  * @param [in]     adapt     probability adaptation context
309  * @param [in]     qm        QM with magnitude compensation
310  * @param [in]     qm_inv    Inverse of QM with magnitude compensation
311  * @param [in] pvq_norm_lambda enc->pvq_norm_lambda for quantized RDO
312  * @param [in]     speed     Make search faster by making approximations
313  * @return         gain      index of the quatized gain
314 */
315 static int pvq_theta(od_coeff *out, const od_coeff *x0, const od_coeff *r0,
316  int n, int q0, od_coeff *y, int *itheta, int *max_theta, int *vk,
317  od_val16 beta, double *skip_diff, int robust, int is_keyframe, int pli,
318  const od_adapt_ctx *adapt, const int16_t *qm,
319  const int16_t *qm_inv, double pvq_norm_lambda, int speed) {
320   od_val32 g;
321   od_val32 gr;
322   od_coeff y_tmp[MAXN];
323   int i;
324   /* Number of pulses. */
325   int k;
326   /* Companded gain of x and reference, normalized to q. */
327   od_val32 cg;
328   od_val32 cgr;
329   int icgr;
330   int qg;
331   /* Best RDO cost (D + lamdba*R) so far. */
332   double best_cost;
333   double dist0;
334   /* Distortion (D) that corresponds to the best RDO cost. */
335   double best_dist;
336   double dist;
337   /* Sign of Householder reflection. */
338   int s;
339   /* Dimension on which Householder reflects. */
340   int m;
341   od_val32 theta;
342   double corr;
343   int best_k;
344   od_val32 best_qtheta;
345   od_val32 gain_offset;
346   int noref;
347   double skip_dist;
348   int cfl_enabled;
349   int skip;
350   double gain_weight;
351   od_val16 x16[MAXN];
352   od_val16 r16[MAXN];
353   int xshift;
354   int rshift;
355   /* Give more weight to gain error when calculating the total distortion. */
356   gain_weight = 1.4;
357   OD_ASSERT(n > 1);
358   corr = 0;
359 #if !defined(OD_FLOAT_PVQ)
360   /* Shift needed to make x fit in 16 bits even after rotation.
361      This shift value is not normative (it can be changed without breaking
362      the bitstream) */
363   xshift = OD_MAXI(0, od_vector_log_mag(x0, n) - 15);
364   /* Shift needed to make the reference fit in 15 bits, so that the Householder
365      vector can fit in 16 bits.
366      This shift value *is* normative, and has to match the decoder. */
367   rshift = OD_MAXI(0, od_vector_log_mag(r0, n) - 14);
368 #else
369   xshift = 0;
370   rshift = 0;
371 #endif
372   for (i = 0; i < n; i++) {
373 #if defined(OD_FLOAT_PVQ)
374     /*This is slightly different from the original float PVQ code,
375        where the qm was applied in the accumulation in od_pvq_compute_gain and
376        the vectors were od_coeffs, not od_val16 (i.e. double).*/
377     x16[i] = x0[i]*(double)qm[i]*OD_QM_SCALE_1;
378     r16[i] = r0[i]*(double)qm[i]*OD_QM_SCALE_1;
379 #else
380     x16[i] = OD_SHR_ROUND(x0[i]*qm[i], OD_QM_SHIFT + xshift);
381     r16[i] = OD_SHR_ROUND(r0[i]*qm[i], OD_QM_SHIFT + rshift);
382 #endif
383     corr += OD_MULT16_16(x16[i], r16[i]);
384   }
385   cfl_enabled = is_keyframe && pli != 0 && !OD_DISABLE_CFL;
386   cg  = od_pvq_compute_gain(x16, n, q0, &g, beta, xshift);
387   cgr = od_pvq_compute_gain(r16, n, q0, &gr, beta, rshift);
388   if (cfl_enabled) cgr = OD_CGAIN_SCALE;
389   /* gain_offset is meant to make sure one of the quantized gains has
390      exactly the same gain as the reference. */
391 #if defined(OD_FLOAT_PVQ)
392   icgr = (int)floor(.5 + cgr);
393 #else
394   icgr = OD_SHR_ROUND(cgr, OD_CGAIN_SHIFT);
395 #endif
396   gain_offset = cgr - OD_SHL(icgr, OD_CGAIN_SHIFT);
397   /* Start search with null case: gain=0, no pulse. */
398   qg = 0;
399   dist = gain_weight*cg*cg*OD_CGAIN_SCALE_2;
400   best_dist = dist;
401   best_cost = dist + pvq_norm_lambda*od_pvq_rate(0, 0, -1, 0, adapt, NULL, 0,
402    n, is_keyframe, pli, speed);
403   noref = 1;
404   best_k = 0;
405   *itheta = -1;
406   *max_theta = 0;
407   OD_CLEAR(y, n);
408   best_qtheta = 0;
409   m = 0;
410   s = 1;
411   corr = corr/(1e-100 + g*(double)gr/OD_SHL(1, xshift + rshift));
412   corr = OD_MAXF(OD_MINF(corr, 1.), -1.);
413   if (is_keyframe) skip_dist = gain_weight*cg*cg*OD_CGAIN_SCALE_2;
414   else {
415     skip_dist = gain_weight*(cg - cgr)*(cg - cgr)
416      + cgr*(double)cg*(2 - 2*corr);
417     skip_dist *= OD_CGAIN_SCALE_2;
418   }
419   if (!is_keyframe) {
420     /* noref, gain=0 isn't allowed, but skip is allowed. */
421     od_val32 scgr;
422     scgr = OD_MAXF(0,gain_offset);
423     if (icgr == 0) {
424       best_dist = gain_weight*(cg - scgr)*(cg - scgr)
425        + scgr*(double)cg*(2 - 2*corr);
426       best_dist *= OD_CGAIN_SCALE_2;
427     }
428     best_cost = best_dist + pvq_norm_lambda*od_pvq_rate(0, icgr, 0, 0, adapt,
429      NULL, 0, n, is_keyframe, pli, speed);
430     best_qtheta = 0;
431     *itheta = 0;
432     *max_theta = 0;
433     noref = 0;
434   }
435   dist0 = best_dist;
436   if (n <= OD_MAX_PVQ_SIZE && !od_vector_is_null(r0, n) && corr > 0) {
437     od_val16 xr[MAXN];
438     int gain_bound;
439     int prev_k;
440     pvq_search_item items[MAX_PVQ_ITEMS];
441     int idx;
442     int nitems;
443     double cos_dist;
444     idx = 0;
445     gain_bound = OD_SHR(cg - gain_offset, OD_CGAIN_SHIFT);
446     /* Perform theta search only if prediction is useful. */
447     theta = OD_ROUND32(OD_THETA_SCALE*acos(corr));
448     m = od_compute_householder(r16, n, gr, &s, rshift);
449     od_apply_householder(xr, x16, r16, n);
450     prev_k = 0;
451     for (i = m; i < n - 1; i++) xr[i] = xr[i + 1];
452     /* Compute all candidate PVQ searches within a reasonable range of gain
453        and theta. */
454     for (i = OD_MAXI(1, gain_bound - 1); i <= gain_bound + 1; i++) {
455       int j;
456       od_val32 qcg;
457       int ts;
458       int theta_lower;
459       int theta_upper;
460       /* Quantized companded gain */
461       qcg = OD_SHL(i, OD_CGAIN_SHIFT) + gain_offset;
462       /* Set angular resolution (in ra) to match the encoded gain */
463       ts = od_pvq_compute_max_theta(qcg, beta);
464       theta_lower = OD_MAXI(0, (int)floor(.5 +
465        theta*OD_THETA_SCALE_1*2/M_PI*ts) - 2);
466       theta_upper = OD_MINI(ts - 1, (int)ceil(theta*OD_THETA_SCALE_1*2/M_PI*ts));
467       /* Include the angles within a reasonable range. */
468       for (j = theta_lower; j <= theta_upper; j++) {
469         od_val32 qtheta;
470         qtheta = od_pvq_compute_theta(j, ts);
471         k = od_pvq_compute_k(qcg, j, qtheta, 0, n, beta, robust || is_keyframe);
472         items[idx].gain = i;
473         items[idx].theta = j;
474         items[idx].k = k;
475         items[idx].qcg = qcg;
476         items[idx].qtheta = qtheta;
477         items[idx].ts = ts;
478         idx++;
479         OD_ASSERT(idx < MAX_PVQ_ITEMS);
480       }
481     }
482     nitems = idx;
483     cos_dist = 0;
484     /* Sort PVQ search candidates in ascending order of pulses K so that
485        we can reuse all the previously searched pulses across searches. */
486     qsort(items, nitems, sizeof(items[0]),
487      (int (*)(const void *, const void *))items_compare);
488     /* Search for the best gain/theta in order. */
489     for (idx = 0; idx < nitems; idx++) {
490       int j;
491       od_val32 qcg;
492       int ts;
493       double cost;
494       double dist_theta;
495       double sin_prod;
496       od_val32 qtheta;
497       /* Quantized companded gain */
498       qcg = items[idx].qcg;
499       i = items[idx].gain;
500       j = items[idx].theta;
501       /* Set angular resolution (in ra) to match the encoded gain */
502       ts = items[idx].ts;
503       /* Search for the best angle within a reasonable range. */
504       qtheta = items[idx].qtheta;
505       k = items[idx].k;
506       /* Compute the minimal possible distortion by not taking the PVQ
507          cos_dist into account. */
508       dist_theta = 2 - 2.*od_pvq_cos(theta - qtheta)*OD_TRIG_SCALE_1;
509       dist = gain_weight*(qcg - cg)*(qcg - cg) + qcg*(double)cg*dist_theta;
510       dist *= OD_CGAIN_SCALE_2;
511       /* If we have no hope of beating skip (including a 1-bit worst-case
512          penalty), stop now. */
513       if (dist > dist0 + 1.0*pvq_norm_lambda && k != 0) continue;
514       sin_prod = od_pvq_sin(theta)*OD_TRIG_SCALE_1*od_pvq_sin(qtheta)*
515        OD_TRIG_SCALE_1;
516       /* PVQ search, using a gain of qcg*cg*sin(theta)*sin(qtheta) since
517          that's the factor by which cos_dist is multiplied to get the
518          distortion metric. */
519       if (k == 0) {
520         cos_dist = 0;
521         OD_CLEAR(y_tmp, n-1);
522       }
523       else if (k != prev_k) {
524         cos_dist = pvq_search_rdo_double(xr, n - 1, k, y_tmp,
525          qcg*(double)cg*sin_prod*OD_CGAIN_SCALE_2, pvq_norm_lambda, prev_k);
526       }
527       prev_k = k;
528       /* See Jmspeex' Journal of Dubious Theoretical Results. */
529       dist_theta = 2 - 2.*od_pvq_cos(theta - qtheta)*OD_TRIG_SCALE_1
530        + sin_prod*(2 - 2*cos_dist);
531       dist = gain_weight*(qcg - cg)*(qcg - cg) + qcg*(double)cg*dist_theta;
532       dist *= OD_CGAIN_SCALE_2;
533       /* Do approximate RDO. */
534       cost = dist + pvq_norm_lambda*od_pvq_rate(i, icgr, j, ts, adapt, y_tmp,
535        k, n, is_keyframe, pli, speed);
536       if (cost < best_cost) {
537         best_cost = cost;
538         best_dist = dist;
539         qg = i;
540         best_k = k;
541         best_qtheta = qtheta;
542         *itheta = j;
543         *max_theta = ts;
544         noref = 0;
545         OD_COPY(y, y_tmp, n - 1);
546       }
547     }
548   }
549   /* Don't bother with no-reference version if there's a reasonable
550      correlation. The only exception is luma on a keyframe because
551      H/V prediction is unreliable. */
552   if (n <= OD_MAX_PVQ_SIZE &&
553    ((is_keyframe && pli == 0) || corr < .5
554    || cg < (od_val32)(OD_SHL(2, OD_CGAIN_SHIFT)))) {
555     int gain_bound;
556     int prev_k;
557     gain_bound = OD_SHR(cg, OD_CGAIN_SHIFT);
558     prev_k = 0;
559     /* Search for the best gain (haven't determined reasonable range yet). */
560     for (i = OD_MAXI(1, gain_bound); i <= gain_bound + 1; i++) {
561       double cos_dist;
562       double cost;
563       od_val32 qcg;
564       qcg = OD_SHL(i, OD_CGAIN_SHIFT);
565       k = od_pvq_compute_k(qcg, -1, -1, 1, n, beta, robust || is_keyframe);
566       /* Compute the minimal possible distortion by not taking the PVQ
567          cos_dist into account. */
568       dist = gain_weight*(qcg - cg)*(qcg - cg);
569       dist *= OD_CGAIN_SCALE_2;
570       if (dist > dist0 && k != 0) continue;
571       cos_dist = pvq_search_rdo_double(x16, n, k, y_tmp,
572        qcg*(double)cg*OD_CGAIN_SCALE_2, pvq_norm_lambda, prev_k);
573       prev_k = k;
574       /* See Jmspeex' Journal of Dubious Theoretical Results. */
575       dist = gain_weight*(qcg - cg)*(qcg - cg)
576        + qcg*(double)cg*(2 - 2*cos_dist);
577       dist *= OD_CGAIN_SCALE_2;
578       /* Do approximate RDO. */
579       cost = dist + pvq_norm_lambda*od_pvq_rate(i, 0, -1, 0, adapt, y_tmp, k,
580        n, is_keyframe, pli, speed);
581       if (cost <= best_cost) {
582         best_cost = cost;
583         best_dist = dist;
584         qg = i;
585         noref = 1;
586         best_k = k;
587         *itheta = -1;
588         *max_theta = 0;
589         OD_COPY(y, y_tmp, n);
590       }
591     }
592   }
593   k = best_k;
594   theta = best_qtheta;
595   skip = 0;
596   if (noref) {
597     if (qg == 0) skip = OD_PVQ_SKIP_ZERO;
598   }
599   else {
600     if (!is_keyframe && qg == 0) {
601       skip = (icgr ? OD_PVQ_SKIP_ZERO : OD_PVQ_SKIP_COPY);
602     }
603     if (qg == icgr && *itheta == 0 && !cfl_enabled) skip = OD_PVQ_SKIP_COPY;
604   }
605   /* Synthesize like the decoder would. */
606   if (skip) {
607     if (skip == OD_PVQ_SKIP_COPY) OD_COPY(out, r0, n);
608     else OD_CLEAR(out, n);
609   }
610   else {
611     if (noref) gain_offset = 0;
612     g = od_gain_expand(OD_SHL(qg, OD_CGAIN_SHIFT) + gain_offset, q0, beta);
613     od_pvq_synthesis_partial(out, y, r16, n, noref, g, theta, m, s,
614      qm_inv);
615   }
616   *vk = k;
617   *skip_diff += skip_dist - best_dist;
618   /* Encode gain differently depending on whether we use prediction or not.
619      Special encoding on inter frames where qg=0 is allowed for noref=0
620      but not noref=1.*/
621   if (is_keyframe) return noref ? qg : neg_interleave(qg, icgr);
622   else return noref ? qg - 1 : neg_interleave(qg + 1, icgr + 1);
623 }
624
625 /** Encodes a single vector of integers (eg, a partition within a
626  *  coefficient block) using PVQ
627  *
628  * @param [in,out] ec         range encoder
629  * @param [in]     qg         quantized gain
630  * @param [in]     theta      quantized post-prediction theta
631  * @param [in]     max_theta  maximum possible quantized theta value
632  * @param [in]     in         coefficient vector to code
633  * @param [in]     n          number of coefficients in partition
634  * @param [in]     k          number of pulses in partition
635  * @param [in,out] model      entropy encoder state
636  * @param [in,out] adapt      adaptation context
637  * @param [in,out] exg        ExQ16 expectation of gain value
638  * @param [in,out] ext        ExQ16 expectation of theta value
639  * @param [in]     nodesync   do not use info that depend on the reference
640  * @param [in]     cdf_ctx    selects which cdf context to use
641  * @param [in]     is_keyframe whether we're encoding a keyframe
642  * @param [in]     code_skip  whether the "skip rest" flag is allowed
643  * @param [in]     skip_rest  when set, we skip all higher bands
644  * @param [in]     encode_flip whether we need to encode the CfL flip flag now
645  * @param [in]     flip       value of the CfL flip flag
646  */
647 static void pvq_encode_partition(od_ec_enc *ec,
648                                  int qg,
649                                  int theta,
650                                  int max_theta,
651                                  const od_coeff *in,
652                                  int n,
653                                  int k,
654                                  generic_encoder model[3],
655                                  od_adapt_ctx *adapt,
656                                  int *exg,
657                                  int *ext,
658                                  int nodesync,
659                                  int cdf_ctx,
660                                  int is_keyframe,
661                                  int code_skip,
662                                  int skip_rest,
663                                  int encode_flip,
664                                  int flip) {
665   int noref;
666   int id;
667   noref = (theta == -1);
668   id = (qg > 0) + 2*OD_MINI(theta + 1,3) + 8*code_skip*skip_rest;
669   if (is_keyframe) {
670     OD_ASSERT(id != 8);
671     if (id >= 8) id--;
672   }
673   else {
674     OD_ASSERT(id != 10);
675     if (id >= 10) id--;
676   }
677   /* Jointly code gain, theta and noref for small values. Then we handle
678      larger gain and theta values. For noref, theta = -1. */
679   od_encode_cdf_adapt(ec, id, &adapt->pvq.pvq_gaintheta_cdf[cdf_ctx][0],
680    8 + 7*code_skip, adapt->pvq.pvq_gaintheta_increment);
681   if (encode_flip) {
682     /* We could eventually do some smarter entropy coding here, but it would
683        have to be good enough to overcome the overhead of the entropy coder.
684        An early attempt using a "toogle" flag with simple adaptation wasn't
685        worth the trouble. */
686     od_ec_enc_bits(ec, flip, 1);
687   }
688   if (qg > 0) {
689     int tmp;
690     tmp = *exg;
691     generic_encode(ec, &model[!noref], qg - 1, -1, &tmp, 2);
692     OD_IIR_DIADIC(*exg, qg << 16, 2);
693   }
694   if (theta > 1 && (nodesync || max_theta > 3)) {
695     int tmp;
696     tmp = *ext;
697     generic_encode(ec, &model[2], theta - 2, nodesync ? -1 : max_theta - 3,
698      &tmp, 2);
699     OD_IIR_DIADIC(*ext, theta << 16, 2);
700   }
701   od_encode_pvq_codeword(ec, &adapt->pvq.pvq_codeword_ctx, in,
702    n - (theta != -1), k);
703 }
704
705 /** Quantizes a scalar with rate-distortion optimization (RDO)
706  * @param [in] x      unquantized value
707  * @param [in] q      quantization step size
708  * @param [in] delta0 rate increase for encoding a 1 instead of a 0
709  * @param [in] pvq_norm_lambda enc->pvq_norm_lambda for quantized RDO
710  * @retval quantized value
711  */
712 int od_rdo_quant(od_coeff x, int q, double delta0, double pvq_norm_lambda) {
713   int threshold;
714   /* Optimal quantization threshold is 1/2 + lambda*delta_rate/2. See
715      Jmspeex' Journal of Dubious Theoretical Results for details. */
716   threshold = 128 + OD_CLAMPI(0, (int)(256*pvq_norm_lambda*delta0/2), 128);
717   if (abs(x) < q*threshold/256) {
718     return 0;
719   }
720   else {
721     return OD_DIV_R0(x, q);
722   }
723 }
724
725 #if OD_SIGNAL_Q_SCALING
726 void od_encode_quantizer_scaling(daala_enc_ctx *enc, int q_scaling,
727  int sbx, int sby, int skip) {
728   int nhsb;
729   OD_ASSERT(skip == !!skip);
730   nhsb = enc->state.nhsb;
731   OD_ASSERT(sbx < nhsb);
732   OD_ASSERT(sby < enc->state.nvsb);
733   OD_ASSERT(!skip || q_scaling == 0);
734   enc->state.sb_q_scaling[sby*nhsb + sbx] = q_scaling;
735   if (!skip) {
736     int above;
737     int left;
738     /* use value from neighbour if possible, otherwise use 0 */
739     above = sby > 0 ? enc->state.sb_q_scaling[(sby - 1)*enc->state.nhsb + sbx]
740      : 0;
741     left = sbx > 0 ? enc->state.sb_q_scaling[sby*enc->state.nhsb + (sbx - 1)]
742      : 0;
743     od_encode_cdf_adapt(&enc->ec, q_scaling,
744      enc->state.adapt.q_cdf[above + left*4], 4,
745      enc->state.adapt.q_increment);
746   }
747 }
748 #endif
749
750 /** Encode a coefficient block (excepting DC) using PVQ
751  *
752  * @param [in,out] enc     daala encoder context
753  * @param [in]     ref     'reference' (prediction) vector
754  * @param [in]     in      coefficient block to quantize and encode
755  * @param [out]    out     quantized coefficient block
756  * @param [in]     q0      scale/quantizer
757  * @param [in]     pli     plane index
758  * @param [in]     bs      log of the block size minus two
759  * @param [in]     beta    per-band activity masking beta param
760  * @param [in]     robust  make stream robust to error in the reference
761  * @param [in]     is_keyframe whether we're encoding a keyframe
762  * @param [in]     q_scaling scaling factor to apply to quantizer
763  * @param [in]     bx      x-coordinate of this block
764  * @param [in]     by      y-coordinate of this block
765  * @param [in]     qm      QM with magnitude compensation
766  * @param [in]     qm_inv  Inverse of QM with magnitude compensation
767  * @param [in]     speed   Make search faster by making approximations
768  * @return         Returns 1 if both DC and AC coefficients are skipped,
769  *                 zero otherwise
770  */
771 int od_pvq_encode(daala_enc_ctx *enc,
772                    od_coeff *ref,
773                    const od_coeff *in,
774                    od_coeff *out,
775                    int q0,
776                    int pli,
777                    int bs,
778                    const od_val16 *beta,
779                    int robust,
780                    int is_keyframe,
781                    int q_scaling,
782                    int bx,
783                    int by,
784                    const int16_t *qm,
785                    const int16_t *qm_inv,
786                    int speed){
787   int theta[PVQ_MAX_PARTITIONS];
788   int max_theta[PVQ_MAX_PARTITIONS];
789   int qg[PVQ_MAX_PARTITIONS];
790   int k[PVQ_MAX_PARTITIONS];
791   od_coeff y[OD_BSIZE_MAX*OD_BSIZE_MAX];
792   int *exg;
793   int *ext;
794   int nb_bands;
795   int i;
796   const int *off;
797   int size[PVQ_MAX_PARTITIONS];
798   generic_encoder *model;
799   double skip_diff;
800   int tell;
801   uint16_t *skip_cdf;
802   od_rollback_buffer buf;
803   int dc_quant;
804   int flip;
805   int cfl_encoded;
806   int skip_rest;
807   int skip_dir;
808   int skip_theta_value;
809   const unsigned char *pvq_qm;
810   double dc_rate;
811 #if !OD_SIGNAL_Q_SCALING
812   OD_UNUSED(q_scaling);
813   OD_UNUSED(bx);
814   OD_UNUSED(by);
815 #endif
816   pvq_qm = &enc->state.pvq_qm_q4[pli][0];
817   exg = &enc->state.adapt.pvq.pvq_exg[pli][bs][0];
818   ext = enc->state.adapt.pvq.pvq_ext + bs*PVQ_MAX_PARTITIONS;
819   skip_cdf = enc->state.adapt.skip_cdf[2*bs + (pli != 0)];
820   model = enc->state.adapt.pvq.pvq_param_model;
821   nb_bands = OD_BAND_OFFSETS[bs][0];
822   off = &OD_BAND_OFFSETS[bs][1];
823   dc_quant = OD_MAXI(1, q0*pvq_qm[od_qm_get_index(bs, 0)] >> 4);
824   tell = 0;
825   for (i = 0; i < nb_bands; i++) size[i] = off[i+1] - off[i];
826   skip_diff = 0;
827   flip = 0;
828   /*If we are coding a chroma block of a keyframe, we are doing CfL.*/
829   if (pli != 0 && is_keyframe) {
830     od_val32 xy;
831     xy = 0;
832     /*Compute the dot-product of the first band of chroma with the luma ref.*/
833     for (i = off[0]; i < off[1]; i++) {
834 #if defined(OD_FLOAT_PVQ)
835       xy += ref[i]*(double)qm[i]*OD_QM_SCALE_1*
836        (double)in[i]*(double)qm[i]*OD_QM_SCALE_1;
837 #else
838       od_val32 rq;
839       od_val32 inq;
840       rq = ref[i]*qm[i];
841       inq = in[i]*qm[i];
842       xy += OD_SHR(rq*(int64_t)inq, OD_SHL(OD_QM_SHIFT + OD_CFL_FLIP_SHIFT,
843        1));
844 #endif
845     }
846     /*If cos(theta) < 0, then |theta| > pi/2 and we should negate the ref.*/
847     if (xy < 0) {
848       flip = 1;
849       for(i = off[0]; i < off[nb_bands]; i++) ref[i] = -ref[i];
850     }
851   }
852   for (i = 0; i < nb_bands; i++) {
853     int q;
854     q = OD_MAXI(1, q0*pvq_qm[od_qm_get_index(bs, i + 1)] >> 4);
855     qg[i] = pvq_theta(out + off[i], in + off[i], ref + off[i], size[i],
856      q, y + off[i], &theta[i], &max_theta[i],
857      &k[i], beta[i], &skip_diff, robust, is_keyframe, pli,
858      &enc->state.adapt, qm + off[i], qm_inv + off[i], enc->pvq_norm_lambda,
859      speed);
860   }
861   od_encode_checkpoint(enc, &buf);
862   if (is_keyframe) out[0] = 0;
863   else {
864     dc_rate = -OD_LOG2((double)(skip_cdf[3] - skip_cdf[2])/
865      (double)(skip_cdf[2] - skip_cdf[1]));
866     out[0] = od_rdo_quant(in[0] - ref[0], dc_quant, dc_rate,
867      enc->pvq_norm_lambda);
868   }
869   tell = od_ec_enc_tell_frac(&enc->ec);
870   /* Code as if we're not skipping. */
871   od_encode_cdf_adapt(&enc->ec, 2 + (out[0] != 0), skip_cdf,
872    4 + (pli == 0 && bs > 0), enc->state.adapt.skip_increment);
873 #if OD_SIGNAL_Q_SCALING
874   if (bs == OD_NBSIZES - 1 && pli == 0) {
875     od_encode_quantizer_scaling(enc, q_scaling, bx >> (OD_NBSIZES - 1),
876      by >> (OD_NBSIZES - 1), 0);
877   }
878 #endif
879   cfl_encoded = 0;
880   skip_rest = 1;
881   skip_theta_value = is_keyframe ? -1 : 0;
882   for (i = 1; i < nb_bands; i++) {
883     if (theta[i] != skip_theta_value || qg[i]) skip_rest = 0;
884   }
885   skip_dir = 0;
886   if (nb_bands > 1) {
887     for (i = 0; i < 3; i++) {
888       int j;
889       int tmp;
890       tmp = 1;
891       for (j = i + 1; j < nb_bands; j += 3) {
892         if (theta[j] != skip_theta_value || qg[j]) tmp = 0;
893       }
894       skip_dir |= tmp << i;
895     }
896   }
897   if (theta[0] == skip_theta_value && qg[0] == 0 && skip_rest) nb_bands = 0;
898   for (i = 0; i < nb_bands; i++) {
899     int encode_flip;
900     /* Encode CFL flip bit just after the first time it's used. */
901     encode_flip = pli != 0 && is_keyframe && theta[i] != -1 && !cfl_encoded;
902     if (i == 0 || (!skip_rest && !(skip_dir & (1 << ((i - 1)%3))))) {
903       pvq_encode_partition(&enc->ec, qg[i], theta[i], max_theta[i], y + off[i],
904        size[i], k[i], model, &enc->state.adapt, exg + i, ext + i,
905        robust || is_keyframe, (pli != 0)*OD_NBSIZES*PVQ_MAX_PARTITIONS
906        + bs*PVQ_MAX_PARTITIONS + i, is_keyframe, i == 0 && (i < nb_bands - 1),
907        skip_rest, encode_flip, flip);
908     }
909     if (i == 0 && !skip_rest && bs > 0) {
910       od_encode_cdf_adapt(&enc->ec, skip_dir,
911        &enc->state.adapt.pvq.pvq_skip_dir_cdf[(pli != 0) + 2*(bs - 1)][0], 7,
912        enc->state.adapt.pvq.pvq_skip_dir_increment);
913     }
914     if (encode_flip) cfl_encoded = 1;
915   }
916   tell = od_ec_enc_tell_frac(&enc->ec) - tell;
917   /* Account for the rate of skipping the AC, based on the same DC decision
918      we made when trying to not skip AC. */
919   {
920     double skip_rate;
921     if (out[0] != 0) {
922       skip_rate = -OD_LOG2((skip_cdf[1] - skip_cdf[0])/
923      (double)skip_cdf[3 + (pli == 0 && bs > 0)]);
924     }
925     else {
926       skip_rate = -OD_LOG2(skip_cdf[0]/
927      (double)skip_cdf[3 + (pli == 0 && bs > 0)]);
928     }
929     tell -= (int)floor(.5+8*skip_rate);
930   }
931   if (nb_bands == 0 || skip_diff <= enc->pvq_norm_lambda/8*tell) {
932     if (is_keyframe) out[0] = 0;
933     else {
934       dc_rate = -OD_LOG2((double)(skip_cdf[1] - skip_cdf[0])/
935        (double)skip_cdf[0]);
936       out[0] = od_rdo_quant(in[0] - ref[0], dc_quant, dc_rate,
937        enc->pvq_norm_lambda);
938     }
939     /* We decide to skip, roll back everything as it was before. */
940     od_encode_rollback(enc, &buf);
941     od_encode_cdf_adapt(&enc->ec, out[0] != 0, skip_cdf,
942      4 + (pli == 0 && bs > 0), enc->state.adapt.skip_increment);
943 #if OD_SIGNAL_Q_SCALING
944     if (bs == OD_NBSIZES - 1 && pli == 0) {
945       int skip;
946       skip = out[0] == 0;
947       if (skip) {
948         q_scaling = 0;
949       }
950       od_encode_quantizer_scaling(enc, q_scaling, bx >> (OD_NBSIZES - 1),
951        by >> (OD_NBSIZES - 1), skip);
952     }
953 #endif
954     if (is_keyframe) for (i = 1; i < 1 << (2*bs + 4); i++) out[i] = 0;
955     else for (i = 1; i < 1 << (2*bs + 4); i++) out[i] = ref[i];
956     if (out[0] == 0) return 1;
957   }
958   return 0;
959 }