えびちゃんの日記

えびちゃん（競プロ）の日記です。

Montgomery 乗算について

Montgomery(モンゴメリ) 乗算の話をします。

Barrett reduction について前に 紹介したことがある のですが、それと同じように除算命令を避けて剰余演算をする手法の一つです。

前提

以下では、法を $N$ とします。また、以下の条件を満たすような $R$ を用意します。

  • $R \gt N$ である。
  • $\gcd(R, N) = 1$ である。
  • $R$ による除算・剰余演算は簡単に行える。

ここでは、$N$ は奇数であるとし $R = 2^{32}$ としてみます。 $\floor{x / R} = x \gg 32$ とできるほか、$x$ を符号なし 64-bit 整数型 x で表せば、$\floor{x / R}$ は x の上位 32-bit、$x\bmod R$ は x の下位 32-bit に対応します。

このとき、法 $N$ での $R$ の乗法逆元 $R^{-1}$ が存在します。すなわち、ある $0\lt R^{-1} \lt N$ と $0\lt N' \lt R$ が存在して $RR^{-1}-NN' = 1$ が成り立ちます。 除算を介さずに $N'$ を求める方法は後述しますが、最初に一度求めれば使い回せるので、適当に互除法などで（除算を用いつつ）求めてもよいでしょう。

Algorithm REDC

まず、$0\le T\lt RN$ から $TR^{-1}\bmod N$ を求めるアルゴリズムを紹介します。ここで、$RN \gt N^2$ に注意しておくとよいかもしれません。

  • function $\text{\small REDC}(T)$
    • $m \gets ((T \bmod R)\cdot N') \bmod R$
      • note: $0\le m\lt R$。
      • note: $m \equiv TN' \pmod{R}$。
    • $t \gets (T + mN) / R$
      • note: $mN \equiv TNN' \equiv T\cdot(RR^{-1}-1) \equiv -T \pmod{R}$。
      • note: なので $T + mN \equiv 0 \pmod{R}$ であり $t$ は整数となる。
      • note: また $tR \equiv T \pmod{N}$ なので $t \equiv TR^{-1} \pmod{N}$。
      • note: $0\le T\lt RN$ と $0\le mN\lt RN$ より $0\le t\lt 2N$。
    • return if $t \lt N$ then $t$ else $t-N$

値の表現・各演算

ここで、$i$ ($0\le i\lt N$) を使って $i$ の属する剰余類を表すのではなく、$i$ を使って $i R^{-1} \bmod N$ の属する剰余類を表すことを考えます。 あるいは、$i$ の属する剰余類を表すために $i$ ではなく $iR\bmod N$ を用いると見ることもできます。

このとき、$(a\pm b)R \equiv aR\pm bR \pmod{N}$ なので、加減算については普段通りに行うことができます。 $0\le a\pm b\lt N$ の範囲を超えた場合は、REDC 同様に $N$ を足し引きして調整します。

また、乗算については、$aR\cdot bR \equiv abR^2 \pmod{N}$ であるため、$(ab)R \equiv \text{\small REDC}(aR\cdot bR) \pmod{N}$ として求められます。

さらに、$aR \equiv \text{\small REDC}(aR^2) \pmod{N}$ なので、予め $R^2 \pmod{N}$ を求めておくことで、以降は除算なしで初期化できます。

各値の前計算

$R^2$ の計算の際は、$(R \bmod N)^2 \bmod N$ としてもよいですが、$R^2 \equiv R^2-N \pmod{N}$ であり、符号なし 64-bit 整数を使うと $R^2-N$ を $-N$ と書けるので、それを利用するのもよいでしょう。

$N'$ については、$N' \equiv -N^{-1} \pmod{R}$ であるので、$N^{-1} \pmod{R}$ を求めることを考えます。 これは、Newton 法を用いると除算を避けることができます。

$n' \gets N$ で初期化し、$n' \gets (n' \cdot (2 - n\cdot n') ) \bmod R$ で更新するのを 4 回程度繰り返すと $n' = N'$ に収束します。 $3$ 以上 $2^{30}$ 未満の各奇数について、これでうまくいくのを 1 秒程度で確認できました。高速ですごいですね。

法が $R$ じゃなくて一般のもののときは Newton 法だとうまくいかないように見えたのですが、うまくいく条件はなんなのでしょう。ねむくて考えられません。

→ 法が \(m\) での逆元を得ているときに、法が $m^2$ での逆元を得るのは上記の Newton 法でできます。ところで法が $4$ のときは（奇数であれば）自分自身が逆元となるため、そこから 4 回行うことで $2^{32}$ を法としたときの逆元が求められるわけですね。以下のスレッドが勉強になります。

twitter.com

並列化

たとえば、メモリ上で連続した位置にある 8 個の mod int をまとめて乗算したいような状況はそれなりにあります。 こうしたとき、SIMD の命令をうまく使うことで、REDC などを並列して行うことができ、高速化が期待できます。 FFT を高速化したいときなどに有効そうです。

下位ワードをまるっと捨ててよいのとかは、複雑なことを気にしなくてよいがちなのでうれしいですね。

正規化のタイミング

REDC 内で $t$ の範囲によって $t-N$ にしたりしなかったりする処理がありますが、これをしなくても $2N$ 未満にはなってくれることは前述の通りです。

そこで、乗算のための REDC で正規化をしないことにすると、加減算でも $0\le i \lt N$ に収めなくてよくなり、$0\le i\lt 2N$ に収まるように調整してもよい感じにできます。

ただし、値の比較をする際に注意しないと $x \not\equiv x+N$ と判定してしまいがちなので、気をつけましょう。

注意

各剰余計算ごとに REDC で変換するとオーバーヘッドが大きそうです。 計算の最初に $i → iR \bmod N$ に変換し、計算はずっとその表現方法のまま行い、計算が終わって出力する段階で $i → iR^{-1}\bmod N$ に戻すのがよいでしょう。

参考文献

  • Montgomery, Peter L. “Modular multiplication without trial division.” Mathematics of computation 44, no. 170 (1985): 519-521. [PDF]

おわり

朝になりました。なんで？