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 */
  }
}

copy

高速化の手法

素直な演算 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

copy

 そこで演算テーブルを用意し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];
            }
        }

copy

 ちなみにベンチマーク(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)|

copy

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] |

copy

 なんと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

copy

 が導出できる。

 そうすると (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

copy

 謎の足し算の正体はこれかな……。なお、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回程度しか計算できないのである。1200×800の画像で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};

copy

 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] |

copy

一括計算 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];
        }

copy

 コードはこう変化する。しかしこのアプローチはほとんど速くならなかった。どうやらコンパイラがコードの最適化を行っているようである。

固定小数点化 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;

copy

 このベンチは実に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 |

copy

 謎の足し算はこの行列の対称性を利用するときにでてくるものらしい。しかるに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};
    }

copy

 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};
    }

copy

 括り出す前の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] | 

copy

ここで、1列目を無視すると4行目を境にして法則がありそう。展開すると謎の足し算が大量に出てきた。しかし、デバッグがしんどいので取りあえず中止……。

AP-922

AP-922の高速化手法

 行列を因数分解することにより、不要な演算を取り除き高速化する手法である。数学的に計算されているので、先のような手法と異なり、効果が計算可能なのである。

 なおこの情報、Intel公式からは消えているらしい。本語訳らしきものがあるが、AVX用なのでMMX用ではなさそう(MMXは整数実装で、AVX用は浮動小数点演算実装)

インテル® AVX を使用した逆離散コサイン変換の実装 | iSUSこの記事は、インテル® ソフトウェア・ネットワークに掲載されている「Using Intel® Advanced Vectowww.isus.jp

アルゴリズムの比較 from AP-922
Algorithm Operations per 8-element 1D DCT Operations per 8×8 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   208

copy

AP-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程度ぐらいなら配列でやった安全そう。