JPEGの展開を速くする
導入
古いJPEGのソースを見ているとIDCT（逆デジタルコサイン変換）で意味不明の謎の足し算をやっているわけである。読んでも理解できないのでそのままコピペされている様だ。しかしWeb Assemblyでこのソースをそのまま使っても速くならない気もするので検証してみる。AP-922の文書（後述）にはそれ以前のアルゴリズムにFeig & Winograd、The Loeffler-Ligtenberg-Moschytz (LLM)、Arai-Agui-Nakajima (AAN) algorithmと言うのものがあるらしい。入手出来ればその実装も確認したいところ。
謎の演算の一部(libjpeg / jidctint.cから）
/* Even part: reverse the even part of the forward DCT. */
/* The rotator is sqrt(2)*c(-6). */
z2 = (INT32) wsptr[2];
z3 = (INT32) wsptr[6];
z1 = MULTIPLY(z2 + z3, FIX_0_541196100);
tmp2 = z1 + MULTIPLY(z3, - FIX_1_847759065);
tmp3 = z1 + MULTIPLY(z2, FIX_0_765366865);
tmp0 = ((INT32) wsptr[0] + (INT32) wsptr[4]) << CONST_BITS;
tmp1 = ((INT32) wsptr[0] - (INT32) wsptr[4]) << CONST_BITS;
tmp10 = tmp0 + tmp3;
tmp13 = tmp0 - tmp3;
tmp11 = tmp1 + tmp2;
tmp12 = tmp1 - tmp2;
/* Odd part per figure 8; the matrix is unitary and hence its
* transpose is its inverse. i0..i3 are y7,y5,y3,y1 respectively.
*/
tmp0 = (INT32) wsptr[7];
tmp1 = (INT32) wsptr[5];
tmp2 = (INT32) wsptr[3];
tmp3 = (INT32) wsptr[1];
z1 = tmp0 + tmp3;
z2 = tmp1 + tmp2;
z3 = tmp0 + tmp2;
z4 = tmp1 + tmp3;
z5 = MULTIPLY(z3 + z4, FIX_1_175875602); /* sqrt(2) * c3 */
tmp0 = MULTIPLY(tmp0, FIX_0_298631336); /* sqrt(2) * (-c1+c3+c5-c7) */
tmp1 = MULTIPLY(tmp1, FIX_2_053119869); /* sqrt(2) * ( c1+c3-c5+c7) */
tmp2 = MULTIPLY(tmp2, FIX_3_072711026); /* sqrt(2) * ( c1+c3+c5-c7) */
tmp3 = MULTIPLY(tmp3, FIX_1_501321110); /* sqrt(2) * ( c1+c3-c5-c7) */
z1 = MULTIPLY(z1, - FIX_0_899976223); /* sqrt(2) * (c7-c3) */
z2 = MULTIPLY(z2, - FIX_2_562915447); /* sqrt(2) * (-c1-c3) */
z3 = MULTIPLY(z3, - FIX_1_961570560); /* sqrt(2) * (-c3-c5) */
z4 = MULTIPLY(z4, - FIX_0_390180644); /* sqrt(2) * (c5-c3) */
z3 += z5;
z4 += z5;
tmp0 += z1 + z3;
tmp1 += z2 + z4;
tmp2 += z2 + z3;
tmp3 += z1 + z4;
/* Final output stage: inputs are tmp10..tmp13, tmp0..tmp3 */
outptr[0] = range_limit[(int) DESCALE(tmp10 + tmp3,
CONST_BITS+PASS1_BITS+3)
& RANGE_MASK];
outptr[7] = range_limit[(int) DESCALE(tmp10 - tmp3,
CONST_BITS+PASS1_BITS+3)
& RANGE_MASK];
outptr[1] = range_limit[(int) DESCALE(tmp11 + tmp2,
CONST_BITS+PASS1_BITS+3)
& RANGE_MASK];
outptr[6] = range_limit[(int) DESCALE(tmp11 - tmp2,
CONST_BITS+PASS1_BITS+3)
& RANGE_MASK];
outptr[2] = range_limit[(int) DESCALE(tmp12 + tmp1,
CONST_BITS+PASS1_BITS+3)
& RANGE_MASK];
outptr[5] = range_limit[(int) DESCALE(tmp12 - tmp1,
CONST_BITS+PASS1_BITS+3)
& RANGE_MASK];
outptr[3] = range_limit[(int) DESCALE(tmp13 + tmp0,
CONST_BITS+PASS1_BITS+3)
& RANGE_MASK];
outptr[4] = range_limit[(int) DESCALE(tmp13 - tmp0,
CONST_BITS+PASS1_BITS+3)
& RANGE_MASK];
wsptr += DCTSIZE; /* advance pointer to next row */
}
}
高速化の手法
素直な演算 IDCT1
そのまま演算すると以下の様になりシンプルだが遅い（しかしuvのループを逆にしたのはなぜだろ……。後で直そう）
let mut val: f32=0.0;
for u in 0..8 {
let cu = if u == 0 {1.0 / 2.0_f32.sqrt()} else {1.0};
for v in 0..8 {
let cv = if v == 0 {1.0_f32 / 2.0_f32.sqrt()} else {1.0};
val += cu * cv * (f[v*8 + u] as f32)
* ((2.0 * x + 1.0) * u as f32 * PI / 16.0_f32).cos()
* ((2.0 * y + 1.0) * v as f32 * PI / 16.0_f32).cos();
}
}
val = val / 4.0
そこで演算テーブルを用意しcosをあらかじめ求めておくとこうなる
let cos_table :[[f32;8];8] =
[[ 1.0, 0.98078528, 0.92387953, 0.83146961, 0.70710678, 0.55557023, 0.38268343, 0.19509032],
[ 1.0, 0.83146961, 0.38268343, -0.19509032, -0.70710678, -0.98078528, -0.92387953, -0.5555702],
[ 1.0, 0.55557023, -0.38268343, -0.98078528, -0.70710678, 0.19509032, 0.92387953, 0.83146961],
[ 1.0, 0.19509032, -0.92387953, -0.55557023, 0.70710678, 0.83146961, -0.38268343, -0.98078528],
[ 1.0, -0.19509032, -0.92387953, 0.55557023, 0.70710678, -0.83146961, -0.38268343, 0.98078528],
[ 1.0, -0.55557023, -0.38268343, 0.98078528, -0.70710678, -0.19509032, 0.92387953, -0.83146961],
[ 1.0, -0.83146961, 0.38268343, 0.19509032, -0.70710678, 0.98078528, -0.92387953, 0.55557023],
[ 1.0, -0.98078528, 0.92387953, -0.83146961, 0.70710678, -0.55557023, 0.38268343, -0.19509032]];
// some code
for u in 0..8 {
for v in 0..8 {
let cucv :f32 = if u == 0 && v ==0 {0.5}
else if u==0 || v==0 {1.0 / 2.0_f32.sqrt()}
else {1.0};
val += cucv * f[v*8 + u] as f32 * cos_table[x][u] * cos_table[y][v];
}
} ちなみにベンチマーク(10万回ループ）では約9000msと2200msと4倍以上変わった（Ryzen7 4900HS + Chrome + Win11 DevToolsを開いた状態) cosの計算のコストが重いことが分かる。
しかし、これでも未だ駄目なのである。
cosのかけ算の正規化 IDCT2
そこでcosの周期性に注目すると、cos[i][j]はこうなる
| 1 cos(0,1) cos(0,2) cos(1,1) cos(0,4) cos(2,1) cos(1,2) cos(3,1)|
| 1 cos(1,1) cos(1,2) -cos(3,1) -cos(0,4) -cos(0,1) -cos(0,2) -cos(2,1)|
| 1 cos(2,1) -cos(0,2) -cos(0,1) -cos(0,4) cos(3,1) cos(0,2) cos(1,1)|
| 1 cos(3,1) -cos(1,2) -cos(2,1) cos(0,4) cos(1,1) -cos(1,2) -cos(0,1)|
| 1 -cos(3,1) -cos(1,2) cos(2,1) cos(0,4) -cos(1,1) -cos(1,2) cos(0,1)|
| 1 -cos(2,1) -cos(0,2) cos(0,1) -cos(0,4) -cos(3,1) cos(0,2) -cos(1,1)|
| 1 -cos(1,1) cos(1,2) cos(3,1) -cos(0,4) cos(0,1) -cos(0,2) cos(2,1)|
| 1 -cos(0,1) cos(0,2) -cos(1,1) cos(0,4) -cos(2,1) cos(1,2) -cos(3,1)|t[n] = cos nπ/16 で置き換えると以下になる
| t[0] t[1] t[2] t[3] t[4] t[5] t[6] t[7] |
| t[0] t[3] t[6] -t[7] -t[4] -t[1] -t[2] -t[5] |
| t[0] t[5] -t[6] -t[1] -t[4] t[7] t[2] t[3] |
| t[0] t[7] -t[2] -t[5] t[4] t[3] -t[6] -t[1] |
| t[0] -t[7] -t[2] t[5] t[4] -t[3] -t[6] t[1] |
| t[0] -t[5] -t[6] t[1] -t[4] -t[7] t[2] -t[3] |
| t[0] -t[3] t[6] t[7] -t[4] t[1] -t[2] t[5] |
| t[0] -t[1] t[2] -t[3] t[4] -t[5] t[6] -t[7] |なんと8つだけ数字を用意すれば良かったのである。cosの解答は8つの解答に±を加えた16だけ。cosの積和の公式cos α cosβ = (cos(α+β) + cos(α-β))/2から
t[i]t[j] = cos i * cos jπ/16 = {cos (i+j)π/16 + cos (i-j)π/16}/2が導出できる。
そうすると (cos (2x+1)uπ/16) * cos ((2y+1)v/π/16) は、
1/2 (cos ((2x+1)u + (2y +1)v)π/16) + cos ((2x+1)u - (2y +1)v)π/16)) になるのであるが、0 ≦ θ ＜ 2π の周期性があるのでそれ以外は丸めればよい。
|(2x+1)u + (2y +1)v| mod 32
|(2x+1)u - (2y +1)v| mod 32
からnが求まれば良い気もするが演算量が減るわけでわけでは無いので最適解ではなさそうである。
t[0] 1.000000
t[1] 0.980800
t[2] 0.923900
t[3] 0.831500
t[4] 0.707100
t[5] 0.555600
t[6] 0.382700
t[7] 0.195100
t[8] 0.000000
t[9] -0.195100
t[10] -0.382700
t[11] -0.555600
t[12] -0.707100
t[13] -0.831500
t[14] -0.923900
t[15] -0.980800
t[16] -1.000000謎の足し算の正体はこれかな……。なお、16<n<32のときはt[32-n]
ここまでやれば固定小数点演算が可能。まず1/2を括りだすと最後に1/8することになる。しかし、計算を256倍して固定小数点演算に変更する。と最後は、256/8 = 64 で割れば良い。CuCvも256倍するので 64*256で割れば良いはず（オーバーフローもしなくなるはず……。）
行列演算の正規化 IDCT3
しかしまだ終わっていない。u=0-7 v=0-7の足し算。足し算64回。これを64回。一つの足し算につきかけ算が4回で、16384回の演算。これをどうサボるかが高速化のポイントなのである。1秒間で100億回計算したとしても61万0351回程度しか計算できないのである。1200x800の画像で0.04秒かかる計算になる（検算してないけど）。動画用途だと遅い……。
今時はsimdやGPUにぶん投げば良いのだろう。ちなみSIMD用にはAP-922と言う実装があるらしい。
https://www.cs.cmu.edu/~barbic/cs-740/ap922.pdf
ちなみにこの方法でも謎の足し算が出てくる。
しかし、これは後で検証するとし、先にCu Cvの正規化を行う。
let cucv :f32 = if u == 0 && v ==0 {0.5}
else if u==0 || v==0 {1.0 / 2.0_f32.sqrt()}
else {1.0}; if分岐をしている部分を、Cv Cuを先にcos_tableにかけてみる。u=0, v=0の時のみ1/√2をかければ良いので最初の一列に1/√2をかけてしまうのだ。
結果、1625msほどに（2-3割減）比較演算と乗算が減ったことによる高速化だろう。
この ときt[n]はこう変化する（ちなみこのテーブル、AP-922と縦横逆なのだが）
| t[4] t[1] t[2] t[3] t[4] t[5] t[6] t[7] |
| t[4] t[3] t[6] -t[7] -t[4] -t[1] -t[2] -t[5] |
| t[4] t[5] -t[6] -t[1] -t[4] t[7] t[2] t[3] |
| t[4] t[7] -t[2] -t[5] t[4] t[3] -t[6] -t[1] |
| t[4] -t[7] -t[2] t[5] t[4] -t[3] -t[6] t[1] |
| t[4] -t[5] -t[6] t[1] -t[4] -t[7] t[2] -t[3] |
| t[4] -t[3] t[6] t[7] -t[4] t[1] -t[2] t[5] |
| t[4] -t[1] t[2] -t[3] t[4] -t[5] t[6] -t[7] |一括計算 IDCT4
c_table[x][u]が何度も出てくるのが分かるだろう。そこで最後の一回にかけることにする。
let mut val = 0.0;
for u in 0..8 {
let mut uval = 0.0;
for v in 0..8 {
uval += f[v*8 + u] as f32 * c_table[y][v];
}
val += uval * c_table[x][u];
}コードはこう変化する。しかしこのアプローチはほとんど速くならなかった。どうやらコンパイラがコードの最適化を行っているようである。
固定小数点化 IDCT5
演算テーブルを整数にして最後でかけた後で変えた数字で割れば、整数演算が可能になる。ちなみに1割程度しか速くならなかった。wasmでは浮動小数点演算が、それほど遅くならないようである（処理系にもよるのだろうが）。速度がクリティカルでなければ敢えて固定小数点化はしなくても良さそう。そのため浮動小数点演算のままでチューニングを進める。
ループの展開 IDCT6
それはともあれ、ループを展開してみた。
let mut val = 0.0;
for u in 0..8 {
let mut uval = 0.0;
uval += f[0*8 + u] as f32 * c_table[y][0];
uval += f[1*8 + u] as f32 * c_table[y][1];
uval += f[2*8 + u] as f32 * c_table[y][2];
uval += f[3*8 + u] as f32 * c_table[y][3];
uval += f[4*8 + u] as f32 * c_table[y][4];
uval += f[5*8 + u] as f32 * c_table[y][5];
uval += f[6*8 + u] as f32 * c_table[y][6];
uval += f[7*8 + u] as f32 * c_table[y][7];
val += uval * c_table[x][u];
}
val = val / 4.0;このベンチは実に739msを叩き出した。IDCT4の2倍を叩き出しており、まずforループがネックになっている様だ。次にテーブル参照がボトルネックになっているようである。そのためテーブル参照も減らす方向で考えてみる。
行列の対称性 IDCT7、8、9
ちなみに、このテーブルは上と下で対象になっているのである。
| a1 a2 a3 a8 |
| b1 b2 b3 b8 |
| c1 c2 c3 c8 |
| d1 d2 d3 d8 |
| d1 -d2 d3 .... -d8 |
| c1 -c2 c3 c8 |
| v1 -b2 b3 b8 |
| a1 -a2 a3 a8 |
謎の足し算はこの行列の対称性を利用するときにでてくるものらしい。しかるに2行同時に計算してしまう方法が考えられる。
取りあえず、xを一度に二回計算することにする。
let mut vals :Vec<u8> = (0..64).map(|_| 0).collect();
for i in 0..32 {
let (x,y) = ((i%4) as usize,(i/4) as usize);
// IDCT from CCITT Rec. T.81 (1992 E) p.27 A3.3
let mut val1 = 0.0;
let mut val2 = 0.0;
let mut plus_minus = 1.0;
for u in 0..8 {
let mut uval1 = 0.0;
uval1 += f[0*8 + u] as f32 * c_table[y][0];
uval1 += f[1*8 + u] as f32 * c_table[y][1];
uval1 += f[2*8 + u] as f32 * c_table[y][2];
uval1 += f[3*8 + u] as f32 * c_table[y][3];
uval1 += f[4*8 + u] as f32 * c_table[y][4];
uval1 += f[5*8 + u] as f32 * c_table[y][5];
uval1 += f[6*8 + u] as f32 * c_table[y][6];
uval1 += f[7*8 + u] as f32 * c_table[y][7];
val1 += uval1 * c_table[x][u];
val2 += uval1 * c_table[x][u] * plus_minus;
plus_minus *= -1.0;
}
val1 = val1 / 4.0;
val2 = val2 / 4.0;
// level shift from CCITT Rec. T.81 (1992 E) p.26 A3.1
let v = val1.round() as isize + 128 ;
vals[y *8 + x] = if v < 0 {0} else if v > 255 {255} else {v as u8};
let v = val2.round() as isize + 128 ;
vals[y *8 + 7-x] = if v < 0 {0} else if v > 255 {255} else {v as u8};
}390msとなり、先のルーチンの倍の速度を叩きだしている。さらにy方向にも適用した（IDCT8）、260msとなり、20-30%程度高速になっている。
更に計算をまとめ行い足し算と引き算を行うと200msになり。おおむね2倍になる。
for i in 0..16 {
let (x,y) = ((i%4) as usize,(i/4) as usize);
// IDCT from CCITT Rec. T.81 (1992 E) p.27 A3.3
let mut val11 = 0.0;
let mut val12 = 0.0;
let mut val21 = 0.0;
let mut val22 = 0.0;
let mut plus_minus = 1.0;
for u in 0..8 {
let temp1 = f[0*8 + u] as f32 * c_table[y][0] + f[2*8 + u] as f32 * c_table[y][2]
+ f[4*8 + u] as f32 * c_table[y][4] + f[6*8 + u] as f32 * c_table[y][6];
let temp2 = f[1*8 + u] as f32 * c_table[y][0] + f[3*8 + u] as f32 * c_table[y][2]
+ f[5*8 + u] as f32 * c_table[y][4] + f[7*8 + u] as f32 * c_table[y][6];
let uval1 = temp1 + temp2;
let uval2 = temp1 - temp2;
val11 += uval1 * c_table[x][u];
val12 += uval1 * c_table[x][u] * plus_minus;
val21 += uval2 * c_table[x][u];
val22 += uval2 * c_table[x][u] * plus_minus;
plus_minus *= -1.0;
}
val11 /= 4.0;
val12 /= 4.0;
val21 /= 4.0;
val22 /= 4.0;
// level shift from CCITT Rec. T.81 (1992 E) p.26 A3.1
let v = val11.round() as isize + 128 ;
vals[y *8 + x] = if v < 0 {0} else if v > 255 {255} else {v as u8};
let v = val12.round() as isize + 128 ;
vals[y *8 + 7-x] = if v < 0 {0} else if v > 255 {255} else {v as u8};
let v = val21.round() as isize + 128 ;
vals[(7 - y) *8 + x] = if v < 0 {0} else if v > 255 {255} else {v as u8};
let v = val22.round() as isize + 128 ;
vals[(7 - y) *8 + 7-x] = if v < 0 {0} else if v > 255 {255} else {v as u8};
}括り出す前の4倍近い速度を叩きだしている。まぁ計算量を64から1/4に減らしたのだから当然なのだが。
ここから先は行列の上下ではなく左右に注目する必要がある。
どうもここまででAP-922の半分の様である。vのループは消滅したので次はuのループを消滅させることにする。道のりは長いぞ……。
さらなるループの展開
下半分を消滅させたので必要な行列はこう変化した。
| t[4] t[1] t[2] t[3] t[4] t[5] t[6] t[7] |
| t[4] t[3] t[6] -t[7] -t[4] -t[1] -t[2] -t[5] |
| t[4] t[5] -t[6] -t[1] -t[4] t[7] t[2] t[3] |
| t[4] t[7] -t[2] -t[5] t[4] t[3] -t[6] -t[1] | ここで、1列目を無視すると4行目を境にして法則がありそう。展開すると謎の足し算が大量に出てきた。しかし、デバッグがしんどいので取りあえず中止……。
AP-922
AP-922の高速化手法
行列を因数分解することにより、不要な演算を取り除き高速化する手法である。数学的に計算されているので、先のような手法と異なり、効果が計算可能なのである。
なおこの情報、Intel公式からは消えているらしい。本語訳らしきものがあるが、AVX用なのでMMX用ではなさそう（MMXは整数実装で、AVX用は浮動小数点演算実装）
アルゴリズムの比較 from AP-922
Algorithm Operations per 8-element 1D DCT Operations per 8x8 2D DCT additions multiplications additions multiplications
要素8 1D DCTの演算数 8x8 2D DCTの演算数
乗算 加算 乗算 加算
Feig & Winograd N/A 454 94 (direct 2D algorithm)
LLM 28 11 448 176
AAN 29 5 464 144
AP-922 320 208AP-922 IDCT10
https://qiita.com/tobira-code/items/91f3578cd7ed5b19c1f9
にPythonの実装コードがあるのでそのままrustに移植してみた。最適化はコンパイラさんにお任せする。ベンチマークにブレが大きくて不明だが、IDCT9のコードの倍の速度が出るようである。最初のコードと比較すると約100倍になる。ただ、もう少しチューニングの余地がありそうである。あと展開後の誤差値±1出ているが定数の設定かな？
ANN
ANNは、細かい情報が落ちてないので困った。取りあえずテスト実装が転がっていたので動かしてみたら、AP-922よりさらに2割ほど速い。ただし、先程のAP-922の実装は最適化されていない。
ANNは以下のサイトにあった。
https://web.stanford.edu/class/ee398a/handouts/lectures/07-TransformCoding.pdf#page=29
LLNは以下のサイトにあった。
https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.463.3353&rep=rep1&type=pdf
https://hal.archives-ouvertes.fr/hal-01797957/document
結局の所行列の因数分解で計算数を減らす手法らしい。
ただ計算数をみるとAP-922よりAANとLLMの方が速くなりそうである。どうもAP-922は、固定小数点演算(i16)に於いての精度の優位性しかないようである。LLMはバグが取れないのでよく分からないがAANとあまり差が出ない模様（メモリ管理やループで変わる範囲ではないかと）
IDCT11のコードは無駄なコードが多いのだが、コンパイラの最適化により書き換えても速度が変わらないのでそのままにしておいた。libjpegは手動で最適化したようである。謎の足し算の正体であった。
初期化しないVecと配列を使った場合も比較。初期化しない方が速いが、64程度ぐらいなら配列でやった安全そう。