誤差無し Barrett Reduction

$\mathrm{mod}\ M (\lt 2 ^ {31})$ 上で様々な計算を行うことを考える。このとき、よくある状況として、たくさんの $0\leq a,b\lt M $ なる整数組 $(a,b)$ に対して $a \star b \bmod M\ (\star\in\lbrace +,-,\times\rbrace)$ を計算することになる。

$\star\in\lbrace +,-\rbrace$ の場合は次のようにして容易に剰余演算を回避できる。

// 1 <= M < 2^31
uint32_t M = ...;
// 0 <= a,b < M の仮定下で a+b mod M を計算
uint32_t add(uint32_t x, uint32_t y) {
    uint32_t c = a + b;
    return c >= M ? c - M : c;
}
// 0 <= a,b < M の仮定下で a-b mod M を計算
uint32_t sub(uint32_t x, uint32_t y) {
    uint32_t c = a + (M - b);
    return c >= M ? c - M : c;
}

$a\times b \bmod M $ を剰余演算無しで計算するのは容易ではないが、Barrett Reduction はこれを実現する方法の $1$ つとして知られている。1 2

状況を少し一般化して、ある正整数 $M $ が固定されており、たくさんの非負整数 $x$ に対して $x \bmod M $ を計算することを考える。

また $x$ の上界 $X\ (\geq M)$ が定まっているものとする。つまり、以下では $0 \leq x\leq X$ を仮定する。$a\times b\bmod M $ の例では $X$ として $(M - 1) ^ 2$ 以上の値を与えることになる。

Barrett Reduction では、初めに $M,X$ のみに依存する (剰余演算を含む) 簡単な前計算を行うことで、各クエリを剰余演算を用いずに処理する。

本記事では、Barrett Reduction における商の近似値の誤差を $0$ にすることで条件分岐を解消する手法について説明する。なお、この手法が高速化に寄与するかは不明である (誰か実験してほしい)。

記号

定数と変数の区別を明確にするため、以下では $x$ に依存しない値を大文字で定義し、$x$ に依存する値を小文字で定義する。

通常の (?) Barrett Reduction

目標は $r\coloneqq x\bmod M $ を計算することである。$r = x - M\cdot \lfloor x/M \rfloor$ より、$q\coloneqq \lfloor x/M \rfloor$ を求められれば十分である。

Barrett Reduction では次のような方針で $q$ を計算する。

  1. $q$ の近似値 $q'$ を計算する。
  2. $q'$ を微調整して $q$ を得る。

Step 1. 近似値 $q'$ の計算

適切に正整数 $K$ を固定して $L\coloneqq \lceil 2 ^ K / M\rceil$ とする 3。$L$ は二進法における $1/M $ の小数第 $K + 1$ 位以下を切り上げて $2 ^ K$ 倍した値である。$K$ は (大文字で定義していることからも分かるように) $x$ に依存しない値であり、その決め方は後述する。$L$ も $x$ に依存しない値であることに注意する。

$q'\coloneqq \lfloor xL / 2 ^ K \rfloor$ を $q$ の近似値として、その誤差を評価する。任意の実数 $a$ に対して $a\leq \lceil a\rceil \lt a + 1$ および $a - 1\lt \lfloor a\rfloor \leq a$ が成り立つので、次が従う。

$$\begin{alignedat}{3} 2 ^ K / M & {}\leq{} & L & {}\lt{} & 2 ^ K / M + 1, \\ \lfloor xL/2 ^ K \rfloor & {}={} & q' & {}\leq{} & xL/2 ^ K. \end{alignedat}$$

これらを併せて次の式 $(1)$ を得る ($L,x,M $ は全て $0$ 以上であることに注意)。

$$q = \lfloor x/M \rfloor \leq q' \lt x/M + x/2 ^ K.\tag{1}$$

従って、$x\leq 2 ^ K$ ならば $q'\in \lbrace q,q+1\rbrace$ が成り立つ。つまり、$K \coloneqq \lceil \log _ 2 X \rceil$ と定めれば必ず $q'\in \lbrace q,q+1\rbrace$ が成り立つ。

実装に当たっては以下の点に注意せよ:

  • $L$ を計算するためには除算を行う必要があるが、これは $x$ に依存しないので前計算して使いまわすこと。クエリ毎に毎回 $L$ を除算で計算すると Barrett Reduction を用いる意味が無くなるので注意。
  • $q'$ の計算に当たって、$2 ^ K$ での切り捨て除算はシフト演算を用いること。具体的には q_prime = (x * L) >> K と計算すること。これも除算を使うと Barrett Reduction を用いる意味が無くなってしまう。(最適化でシフト演算に置き換えてくれるかもしれないけど)

Step 2. $q$ の計算

$q'\in \lbrace q, q+1\rbrace$ であるから、$r' \coloneqq x - q' M $ とすれば次のように $q$ を計算できる。同時に $r$ も計算できる。

  • $r'\lt 0$ の場合: $q = q' - 1,\ r = r' + M $ である。
  • $r' \geq 0$ の場合: $q = q',\ r = r'$ である。

$r'$ の値によって条件分岐が生じることに注意する。

計算途中に現れる値の最大値

オーバーフロー回避のために、計算途中に現れうる値の最大値 $U$ を見積もっておく。これは $U = XL = X\lceil 2 ^ K / M\rceil$ である。

例えば $X\lt 2 ^ {63}$ の場合は $K = 64$ とすれば $K \geq \lceil \log _ 2 X\rceil$ を満たす。このとき $U\lt 2 ^ {127}$ であるから、計算途中の値を $128$ bit 整数型で持てばオーバーフローは起こらない。

$r'$ は負の値を取り得るので符号無し整数型を用いて実装する場合は十分注意すること。なお、$x$ と $q'M $ の大小比較で条件分岐すればこの問題を避けることができる。

(本当は負方向のオーバーフロー回避のために最小値も確認すべきだが、省略する)

誤差無し Barrett Reduction

実は、上記アルゴリズムにおいて $K$ をさらに大きく取ることで $q'$ は $q$ と必ず一致するようにできる。これによって Step 2 を省略でき、即ちクエリ処理時に条件分岐が一切現れないようにできる。

$q'$ と $q$ が必ず一致するような $K$ の十分条件を考える。式 $(1)$ より、全ての $0\leq x\leq X$ について $x / M + x/2 ^ K \leq \lfloor x/ M \rfloor + 1$ を満たすことが十分である。

$x / M $ の小数部分 $t\coloneqq x/M - \lfloor x/M\rfloor$ について $0\leq t\leq 1 - 1/M $ が成り立つ。従って、$X / 2 ^ K \leq 1/M $ 即ち $K \geq \lceil \log _ 2 (MX)\rceil$ は十分条件である。

modint

冒頭の話に戻り、Barrett Reduction を modint 実装に用いることを考える。

固定された $M\ (\lt 2 ^ {31})$ とたくさんの $0\leq a,b\lt M $ なる $a,b$ に対して $a\times b\bmod M $ を計算したいのであった。$0\leq a\times b \leq (M - 1) ^ 2$ であるから $X = (M - 1) ^ 2\ (\lt 2 ^ {62})$ とすれば全ての入力をカバーすることができる。

また、$K = 96$ と定めると、$M \lt 2 ^ {31},\ X\lt 2 ^ {62}$ の下で誤差無し Barrett Reduction の条件 $K \geq \lceil \log _ 2 (MX)\rceil$ を満たす。

計算途中に現れる値の最大値は $U = (M - 1) ^ 2\lceil 2 ^ {96} / M\rceil$ であり、次が成り立つ。

$$\begin{aligned} U & \lt{} (M - 1) ^ 2 (2 ^ {96} / M + 1) &\quad &(\because \lceil 2 ^ {96} / M\rceil \lt 2 ^ {96} / M + 1)\\ & ={} 2 ^ {96} M - 2\times 2 ^ {96} + 2 ^ {96} / M + (M - 1) ^ 2 \\ & \lt{} 2 ^ {127}. &\quad &(\because 1\leq M\lt 2 ^ {31}) \end{aligned}$$

つまり、計算途中の値は $128$ bit 整数型で持つことができる。

modint 実装例

以上を踏まえて modint を実装すると次のようになる (※動作確認はしてません。雰囲気を伝えるだけのものと解釈してください。)

class modint {
    struct barrett {
        static constexpr int K = 96;
        uint32_t M;
        __uint128_t L;
        barrett(uint32_t M) : M(M), L(((__uint128_t(1) << K) - 1) / M + 1) {}
        // a*b mod M
        uint32_t mul(uint32_t a, uint32_t b) const {
            uint64_t c = uint64_t(a) * b;
            uint64_t q = (c * L) >> K;
            return c - q * M;
        }
    };
    static inline barrett _bt{1};
    uint32_t _val;
public:
    // ... (コンストラクタやメソッドなど)

    static void set_mod(int M) {
        assert(1 <= M);
        _bt = barrett(M);
    }

    modint& operator*=(const modint & b) {
        _val = _bt.mul(_val, b._val);
        return *this;
    }
};


  1. $M $ がコンパイル時定数である場合は a * b % M などと書いても最適化によって剰余演算を回避するバイナリが吐かれる (と思われる)。
  2. 他の方法としては、例えば Montgomery 乗算など
  3. バレットの還元アルゴリズム - Wikipedia のように $L\coloneqq \lfloor 2 ^ K / M\rfloor$ とした場合に誤差無し Barrett Reduction を作ろうとするとどうなるかを考えると面白いかもしれません。