LLVMのk乗和の最適化手法について

clang-O1以上の最適化オプションをつけると、以下のコード(単純に考えるとΘ(n)命令必要そう)からループを取り除き、Θ(1)命令のコードに最適化します。

uint64_t sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0; i < n; ++i ) {
        ret += i;
    }
    return ret;
}

これはclangの驚異的な最適化として挙げられることが多い例ですが、どのような仕組みで行われているのか深く調べたことがなかったので、調べてみました。

環境

以下では、clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04を使っていきます。 なお、この最適化はclangが独立したCコンパイラとなったclang 3.0の時点ですでに実装されていたようです(clang+llvm-3.0-x86_64-linux-Ubuntu-11_04/bin/clangで確認)。

実験には、以下のソースコードを使います。

// square_sum.c
#include <stdio.h>

unsigned long long n;
int main() {
        unsigned long long i, sum = 0;
        for( i = 0; i < n; ++i ) {
                sum += i * i;
        }
        printf( "%llu\n", sum );
}

clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/clang -O2 -S -masm=intel square_sum.cとしてコンパイルすると、ループ部分は以下の機械語コードになります。

        mov     rdi, qword ptr [rip + n]
        test    rdi, rdi
        je      .LBB0_1
# %bb.2:
        lea     rax, [rdi - 1]
        lea     rcx, [rdi - 2]
        mul     rcx
        mov     r8, rax
        mov     rsi, rdx
        lea     rcx, [rdi - 3]
        mul     rcx
                                        # kill: def $ecx killed $ecx killed $rcx
        imul    ecx, esi
        add     edx, ecx
        shld    rdx, rax, 63
        movabs  rax, 6148914691236517206
        imul    rax, rdx
        add     rax, rdi
        shld    rsi, r8, 63
        lea     rcx, [rsi + 2*rsi]
        lea     rsi, [rax + rcx]
        add     rsi, -1
        jmp     .LBB0_3
.LBB0_1:
        xor     esi, esi
.LBB0_3:
        # printfの第二引数はrsiで渡すことになっているのでrsiが計算結果のはず

ちゃんとループが消えていることがわかります。

どこで最適化が行われているのかを特定

前提知識

clangC言語から機械語を生成しますが、その裏側では複数のソフトウェアが動いています。 具体的には、以下のようにLLVMに含まれるいくつかのソフトウェアが順に使われます。

  1. C言語フロントエンドとしてのclangで、C言語ソースコードLLVM-IRに変換する
    • 主に言語に依存する最適化に関係します
  2. ミドルエンドであるoptで、LLVM-IRを最適化されたLLVM-IRに変換する
    • 主に言語にも機械語にも依存しない最適化に関係します
  3. バックエンドであるllcで、LLVM-IRを機械語に変換する
    • 主に機械語に依存する最適化に関係します

バックエンドではない

3.の直前のLLVM-IRを手に入れるためには、clang -S -emit-llvmとすればよいです。 clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/clang -O2 -S -emit-llvm square_sum.cとすると、main関数は以下のLLVM-IRになっていました。

; Function Attrs: nofree nounwind uwtable
define dso_local i32 @main() local_unnamed_addr #0 {
  %1 = load i64, i64* @n, align 8, !tbaa !2
  %2 = icmp eq i64 %1, 0
  br i1 %2, label %21, label %3

3:                                                ; preds = %0
  %4 = add i64 %1, -1
  %5 = zext i64 %4 to i65
  %6 = add i64 %1, -2
  %7 = zext i64 %6 to i65
  %8 = mul i65 %5, %7
  %9 = add i64 %1, -3
  %10 = zext i64 %9 to i65
  %11 = mul i65 %8, %10
  %12 = lshr i65 %11, 1
  %13 = trunc i65 %12 to i64
  %14 = mul i64 %13, 6148914691236517206
  %15 = add i64 %1, %14
  %16 = lshr i65 %8, 1
  %17 = trunc i65 %16 to i64
  %18 = mul i64 %17, 3
  %19 = add i64 %15, %18
  %20 = add i64 %19, -1
  br label %21

21:                                               ; preds = %3, %0
  %22 = phi i64 [ 0, %0 ], [ %20, %3 ]
  %23 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([6 x i8], [6 x i8]* @.str, i64 0, i64 0), i64 %22)
  ret i32 0
}

この時点でループが消失しており、x86特有の最適化というわけではなさそうです。

ミドルエンドで行われている

clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/clang -O0 -S -emit-llvm square_sum.cとして、フロントエンドでは最適化しないようにします。

; Function Attrs: noinline nounwind optnone uwtable
define dso_local i32 @main() #0 {
  %1 = alloca i32, align 4
  %2 = alloca i64, align 8
  %3 = alloca i64, align 8
  store i32 0, i32* %1, align 4
  store i64 0, i64* %3, align 8
  store i64 0, i64* %2, align 8
  br label %4

4:                                                ; preds = %14, %0
  %5 = load i64, i64* %2, align 8
  %6 = load i64, i64* @n, align 8
  %7 = icmp ult i64 %5, %6
  br i1 %7, label %8, label %17

8:                                                ; preds = %4
  %9 = load i64, i64* %2, align 8
  %10 = load i64, i64* %2, align 8
  %11 = mul i64 %9, %10
  %12 = load i64, i64* %3, align 8
  %13 = add i64 %12, %11
  store i64 %13, i64* %3, align 8
  br label %14

14:                                               ; preds = %8
  %15 = load i64, i64* %2, align 8
  %16 = add i64 %15, 1
  store i64 %16, i64* %2, align 8
  br label %4, !llvm.loop !2

17:                                               ; preds = %4
  %18 = load i64, i64* %3, align 8
  %19 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([6 x i8], [6 x i8]* @.str, i64 0, i64 0), i64 %18)
  %20 = load i32, i32* %1, align 4
  ret i32 %20
}

この時点ではまだループが残っています。

ここで、このLLVM-IRファイル(square_sum.ll)からoptnoneを消します。 optnoneが最適化を強制的に妨げるためです。 二つありますが、一つ目はコメントの中なので、消しても消さなくても同じです。 二つ目を消すのが大事です。

その後、clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt -O2 -S square_sum.llとしてoptコマンドを適用します。 すると、以下のようにループが消えます。

; Function Attrs: nofree noinline nounwind uwtable
define dso_local i32 @main() local_unnamed_addr #0 {
  %1 = load i64, ptr @n, align 8
  %.not = icmp eq i64 %1, 0
  br i1 %.not, label %._crit_edge, label %.lr.ph.preheader

.lr.ph.preheader:                                 ; preds = %0
  %2 = add i64 %1, -1
  %3 = zext i64 %2 to i65
  %4 = add i64 %1, -2
  %5 = zext i64 %4 to i65
  %6 = mul i65 %3, %5
  %7 = add i64 %1, -3
  %8 = zext i64 %7 to i65
  %9 = mul i65 %6, %8
  %10 = lshr i65 %9, 1
  %11 = trunc i65 %10 to i64
  %12 = mul i64 %11, 6148914691236517206
  %13 = add i64 %1, %12
  %14 = lshr i65 %6, 1
  %15 = trunc i65 %14 to i64
  %16 = mul i64 %15, 3
  %17 = add i64 %13, %16
  %18 = add i64 %17, -1
  br label %._crit_edge

._crit_edge:                                      ; preds = %.lr.ph.preheader, %0
  %.0.lcssa = phi i64 [ 0, %0 ], [ %18, %.lr.ph.preheader ]
  %19 = tail call i32 (ptr, ...) @printf(ptr noundef nonnull dereferenceable(1) @.str, i64 noundef %.0.lcssa)
  ret i32 0
}

したがって、この最適化はミドルエンドで行われているようです。

Induction Variable Simplificationで行われている

さて、-O2はいくつかの最適化パスの集合体です。 具体的にどのような最適化が行われているかを調べるには、--print-after-allをつけて最適化パスの名前とその時点でのLLVM-IRを出力すればよさそうです(llvm opt optionsでググって出てきたoptの概要 - nothingcosmos wikiに書いてありました)。 これにより、*** IR Dump After Induction Variable Simplification ***の後に出力されるLLVM-IRではループが消失していることがわかります。

Induction Variable Simplificationだけではダメ

clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt --help | grep 'Induction Variable Simplification'とすると、Induction Variable Simplificationは--indvarsというオプションを指定すれば実行されるということがわかります。 clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt --indvars -S square_sum.llとすればよいようです。 ※llvm-16.0.0ではオプションの文法が違い、The `opt -passname` syntax for the new pass manager is not supported,と怒られます。--passes=indvarsなどとすればよいようです。

しかし、これは残念ながらループのままです。

Induction Variable Simplificationの前提条件

最適化パスは他の最適化パスによりLLVM-IRが整えられていることを前提としていることがあります。 どの最適化パスがInduction Variable Simplificationの前提条件となっているかを調べていきます。

少なくとも、-O2と同じ順番でInduction Variable Simplificationまで最適化パスを実行すれば同じ結果が得られるはずです。 実行されているパスの一覧は、以下のようにして得られます。

clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt -S -O2 --print-after-all square_sum.ll 2>&1 | grep '\*\*\*'`

以下のような出力を得ます。

*** IR Dump After Module Verifier ***
*** IR Dump After Instrument function entry/exit with calls to e.g. mcount() (pre inlining) ***
*** IR Dump After Simplify the CFG ***
*** IR Dump After SROA ***
(以下略。計121行)

これに対するオプションを調べていきます。 対応するオプションを以下のようにして得ます。

for i in $(seq 121); do clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt --help 2>&1 | grep -- "- $(clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt square_sum.ll -S -O2 --print-after-all 2>&1 | grep '\*\*\*' | sed 's/.*After \(.*\) \*\*\*/\1/' | head -n $i | tail -n 1)\$" | awk '{print$1}'; done

llvm-16.0.0では末尾に親切な情報がついてくるため、このコマンドではうまくいきません。

--verify
--ee-instrument
--simplifycfg
--early-cse
--lower-expect
--annotation2metadata
--forceattrs
--inferattrs
--ipsccp
--called-value-propagation
--globalopt
--mem2reg
--deadargelim
--instcombine
--simplifycfg
--prune-eh
--inline
--openmpopt
--function-attrs
--early-cse-memssa
--jump-threading
--correlated-propagation
--simplifycfg
--instcombine
--libcalls-shrinkwrap
--tailcallelim
--simplifycfg
--reassociate
--loop-simplify
--lcssa-verification
--lcssa
--loop-rotate
--licm
--loop-unswitch
--simplifycfg
--instcombine
--loop-simplify
--lcssa-verification
--lcssa
--loop-idiom
--indvars
(以下略)

後はこれを適当に消していきます。すると、以下のように二つの最適化パスだけでうまくいくことがわかります。

clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt --licm --indvars -S square_sum.ll

--licmは、Loop Invariant Code Motionです。 Loop Invariantはループ不変条件を意味しますが、そんなに大層なことはやっておらず、単にループ中で変わらない変数の計算をループの外に追い出す最適化がかかるだけのようです。 元のコードはループの終了条件にグローバル変数が使われており、毎回ロードが必要だったのが、Induction Variable Simplificationで最適化がかからなかった原因のようです。

実際、以下のコードに変えてみると、clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt --sroa --indvars -S square_sum_with_end.llでうまくいくことがわかります(optnoneを消すのをお忘れなく)。 ここで、--sroaはローカル変数endをスタックではなくレジスタに割り付けるために指定しています。

// square_sum_with_end.c
#include <stdio.h>

unsigned long long n;
int main() {
        unsigned long long i, sum = 0, end = n;
        for( i = 0; i < end; ++i ) {
                sum += i * i;
        }
        printf( "%llu\n", sum );
}

これらのことから、ループの終了条件がループ中で変わらないとすぐにわかる(=レジスタに割り付けられた)変数であることが前提条件となっていることが推測できます。

Induction Variable Simplificationの仕組み

ソースコード上の場所を特定

ここからはLLVMソースコードを見ていきます。 かなり昔から行われている最適化なので、かなり安定していると考えられます。 今のバージョンのソースコードを見てもそう違いはないでしょう。

GitHub - llvm/llvm-projectでindvarsと検索すると、llvm-project/llvm/lib/Transforms/Scalar/IndVarSimplify.cppが最適化パスのソースコードのようです。 2000行以上あって読むのが大変ですが、bool IndVarSimplify::run(Loop *L)が本体なので、コメントをたよりにそこを見ていきます。 そうすると、rewriteLoopExitValuesというのがいかにも怪しそうだとわかります。

rewriteLoopExitValuesググるLoopUtils.cppの中にある関数のようです。 1315行目からを眺めていくと、1425行目ExitValue = AddRec->evaluateAtIteration(ExitCount, *SE);でループが終わるときの値を計算しているとみて良さそうです。

evaluateAtIterationググるllvm::SCEVAddRecExprというクラスメンバ関数のようです。 ドキュメントを見てみると、isAffineisQuadratic*1などの誘導変数最適化関連っぽいメンバ関数があり、最適化がここで行われている期待が持てます。 肝心のevaluateAtIteration関数は978行目にあって、BinomialCoefficient関数を呼び出しています。 BinomialCoefficient(It, K, SE, ResultTy)関数は、863行目のコメントにあるように、 \frac{It(It-1)\cdots(It-K+1)}{K!}を計算するコードを生成します。 これは、コンビネーションの記法を用いれば、 {}_{It}C_{K}です。

LLVM-IRを確認

clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt --sroa --indvars -S square_sum_with_end.llをやってみると、以下のようになります。

define dso_local i32 @main() #0 {
  %1 = load i64, i64* @n, align 8
  %2 = zext i64 %1 to i65
  %3 = add i64 %1, -1
  %4 = zext i64 %3 to i65
  %5 = mul i65 %2, %4
  %6 = add i64 %1, -2
  %7 = zext i64 %6 to i65
  %8 = mul i65 %5, %7
  %9 = lshr i65 %8, 1
  %10 = trunc i65 %9 to i64
  %11 = mul i65 %2, %4
  %12 = lshr i65 %11, 1
  %13 = trunc i65 %12 to i64
  br label %14

14:                                               ; preds = %16, %0
  br i1 false, label %15, label %17

15:                                               ; preds = %14
  br label %16

16:                                               ; preds = %15
  br label %14, !llvm.loop !2

17:                                               ; preds = %14
  %18 = mul i64 %10, 6148914691236517206
  %19 = add i64 %18, %13
  %20 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([6 x i8], [6 x i8]* @.str, i64 0, i64 0), i64 %19)
  ret i32 0
}

計算過程を平易に記述してみると、

  1. nをロードする(%1)
  2. n - 1を計算する(%3)
  3. n * (n - 1)を65bit整数として計算する(%5)
  4. n - 2を計算する(%7)
  5. n * (n - 1) * (n - 2)を65bit整数として計算する(%8)
  6. n * (n - 1) * (n - 2) / 2を計算する。64bit整数に収まるはず(%10)
  7. n * (n - 1)を65bit整数として計算する(%11。%5と同じなのにもう一度計算しているのは、共通部分式最適化を施す前の自動生成されたコードそのままであるため)
  8. n * (n - 1) / 2を計算する。64bit整数に収まるはず(%13)
  9. n * (n - 1) * (n - 2) / 3を定数乗算テクニックで計算する(%18)
    • 3の倍数を2/3倍するには、0x5555555555555556を掛けます。この定数は \frac{2^{64}}3+\frac23で、被乗数が3の倍数であれば最初の項は \mod 2^{64}で消えるので、2/3倍できます。
  10. n * (n - 1) * (n - 2) / 3 + n * (n - 1) / 2を計算する。これが求めるものである(%19)

となります。 つまり、 2\,{}_{n}C_{3} + {}_{n}C_{2}を計算するという非常に単純な仕組みです。 この係数21がどこから出てきたのかを、次に説明します。

ループ誘導変数最適化の仕組み

ループ誘導変数最適化の基礎

ループ誘導変数最適化は、典型的な演算強度低減最適化です。 例えば、以下のコードを考えます。

uint32_t sum( uint32_t* arr, size_t n ) {
    uint32_t ret = 0;
    for( size_t i = 0; i < n; ++i ) {
        ret += arr[i];
    }
    return ret;
}

これを素直にコンパイルすると、以下のようになるはずです。

uint32_t sum( uint32_t* arr, size_t n ) {
    uint32_t ret = 0;
    for( size_t i = 0; i < n; ++i ) {
        char* p = (char*)arr + i * 4;
        uint32_t val = *(uint32_t*)p;
        ret += val;
    }
    return ret;
}

乗算が必要になっている点に注意します。 x86ではlea命令があるのであまり気になりませんが、乗算は一般に高コストであるので、可能であれば避けたいです。 そこで、コンパイラは以下のように最適化します。

uint32_t sum( uint32_t* arr, size_t n ) {
    uint32_t ret = 0;
    char* p = (char*)arr;
    for( size_t i = 0; i < n; ++i, p += 4 ) {
        uint32_t val = *(uint32_t*)p;
        ret += val;
    }
    return ret;
}

乗算がなくなり、+= 4という加算になりました。 このような最適化を(ループ誘導変数に関連する)演算強度低減最適化と言います。

発展的なループ誘導変数最適化

この考え方を発展させます。 以下のコードについて考えます。

uint64_t square_sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0; i < n; ++i ) {
        ret += i * i;
    }
    return ret;
}

これを誘導変数最適化すると以下のようになります。

uint64_t sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0, square = 0; i < n; square += 2 * i + 1, ++i ) {
        ret += square;
    }
    return ret;
}

本物の乗算から定数乗算へと、演算強度が低減されたことがわかります。

ここに現れる定数21こそが、 2\,{}_{n}C_{3} + {}_{n}C_{2}の係数だったのです。 この係数はどう求めるかというと、ループ誘導変数同士の乗算の公式を使っています(ScalarEvolution.cppの3311行目から3314行目)。 Clang の k 乗和の最適化を眺める - えびちゃんの日記では競プロerらしく(?)第二種 Stirling 数との関連を指摘&高速に計算する方法を考察していますが、そんなに高尚なものではないです。 愚直な計算なのでオーダーはおそらくΘ(k3)ですが、そんなに大きな次数の多項式が現れることはありえなさそうです。 そもそも、そういう場合は係数がオーバーフローしてあきらめるので、遅さが問題となることはありません。 逆に、この公式を使うと有理数演算が出てこないのがうれしいのだと思います。

ループ誘導変数同士の乗算の公式の説明の前に、まずループ誘導変数を定式化します。

ループ誘導変数の定式化

三乗和のコードを手で最適化して考えてみる

実際のコードで出てくることはほぼありませんが、二乗だとまだ一般化が見えづらいため、以下のように三乗が出てくるループを考えます。

uint64_t cube_sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0, square = 0; i < n; ++i ) {
        ret += i * i * i;
    }
    return ret;
}

 (i+1)^3 = i^3 + 3i^2 + 3i +1であることから、これを以下のように変形できます。

uint64_t cube_sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0, cube = 0; i < n; cube += 3 * i * i + 3 * i + 1, ++i ) {
        ret += cube;
    }
    return ret;
}

さて、ここでi * iが出てきました。これに対しても再帰的に適用することで、以下のようになります。

uint64_t cub_sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0, square = 0, cube = 0; i < n; cube += 3 * square + 3 * i + 1, square += 2 * i + 1, ++i ) {
        ret += cube;
    }
    return ret;
}

3 * square3 * iに対して演算強度低減最適化をしましょう。

uint64_t cube_sum(uint64_t n) {
    uint64_t ret = 0;
    for( uint64_t i = 0, three_i = 0, three_square = 0, cube = 0; i < n; cube += three_square + three_i + 1, three_square += 6 * i + 3, three_i += 3, ++i ) {
        ret += cube;
    }
    return ret;
}

three_square + three_iはまとめられそうです。

uint64_t cube_sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0, tmp = 0, cube = 0; i < n; cube += tmp + 1, tmp += 6 * i + 6, ++i ) {
        ret += cube;
    }
    return ret;
}

6 * iも演算強度低減最適化しましょう。

uint64_t cube_sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0, six_i = 0, tmp = 0, cube = 0; i < n; cube += tmp + 1, tmp += six_i + 6, six_i += 6, ++i ) {
        ret += cube;
    }
    return ret;
}

ここで現れる係数(1, 6, 6)が求めるものです。 つまり、retの最終値 {}_nC_2 + 6\,{}_nC_3 + 6\,{}_nC_4だということです。 実際、optの出力はこれを計算するものになっていました。 この式を展開すると  \left(-\frac12n+\frac12n^2\right)+\left(2n-3n^2+n^3\right)+\left(-\frac32n+\frac{11}4n^2-\frac32n^3+\frac14n^4\right)になっていて、  \sum_{i=0}^{n-1} i^3= \frac14n^2 - \frac12n^3 + \frac14n^4と一致しています。

ここで出てきた誘導変数は、cubetmpsix_iでした。 これらのいずれも、cube += tmp + 1tmp += six_i + 6six_i += 6のように、更新の式が(他の誘導変数または0)+(定数)を加算する形になっています。 そこで、Aが誘導変数であるということを、次のように定義しましょう。

  • Aが誘導変数であるとは、他の誘導変数Bを用いて、更新の式がA += B + cのようになることである
    • ここで、増分cはループ不変式である必要がある
    • B0でもよい

※これは正確には加法的な誘導変数に関する定義です。指数関数を乗算に落とす演算強度最適化を行うと、更新の式がA *= B * cB1でもよい)のような乗法的な誘導変数が出現します。LLVMはこれを取り扱えず、 \sum_{i=0}^{n-1}2^iを計算するコードはSIMDを駆使して頑張る命令列になります。

このように定義すると、Aを記述するのに必要な情報は、Aの初期値と、各誘導変数の更新式に現れる増分cだけであり、簡潔に記述することができます。 LLVMの流儀(SCEVAddRecExprの出力フォーマット)では、iのことを{0,+,1}(初期値が0で1ずつ増える)、six_iのことを{0,+,6}(初期値が0で6ずつ増える)、tmpのことを{0,+,6,+,6}(初期値が0で、6+「初期値が0で6ずつ増える値」ずつ増える)、cubeのことを{0,+,1,+,6,+,6}(初期値が0で、1+『初期値が0で6+「初期値が0で6ずつ増える値」ずつ増える値』ずつ増える)、retのことを{0,+,0,+,1,+,6,+,6}(初期値が0で、0+【初期値が0で1+『初期値が0で6+「初期値が0で6ずつ増える値」ずつ増える値』ずつ増える値】ずつ増える)と表記しているようです。 なお、途中に出てくる誘導変数の初期値はすべて0に正規化出来ます。一つ外側の誘導変数の増分cを変えればよいからです。

実際の計算手順(ループ誘導変数同士の乗算の公式を使う方法)

上記やり方は、「three_square + three_iはまとめられそう」のようにあまり系統立った方法ではありませんでした。 実はもう一通りやり方があって、それはまずi * iだけを誘導変数とみなす方法です。 この方法で最適化すると、まず一回目では以下のようになります。

uint64_t cube_sum( uint64_t n ) {
    uint64_t ret = 0;
    for( uint64_t i = 0, square = 0, cube = 0; i < n; square += 2 * i + 1, ++i ) {
        ret += square * i;
    }
    return ret;
}

ここで、squareLLVMの記法でいうと{0,+,1,+,2}で、i{0,+,1}です。 問題は、{0,+,1,+,2}{0,+,1}の乗算をどう取り扱うかです。 ようするに{a,+,b,+,c} a\,{}_iC_0 + b\,{}_iC_1 + c\,{}_iC_2なので、この形同士の乗算を考えればいいわけですが、これには公式があるようです。 ソースコードの方が信頼できるのでScalarEvolution.cppの3340行目から3364行目を眺めると、以下のような計算をしているようです。

// X = { 0, 1, 2 };
// Y = { 0, 1 };
ret = {};
for( x = 0; x < X.size() + Y.size() - 1; ++x ) {
  sum = 0;
  for( y = x; y < 2*x + 1; ++y ) {
    Coeff1 = Choose(x, 2*x - y);
    for( z = max(y-x, y-X.size()+1), z < min(x+1, Y.size()); ++z )
      Coeff2 = Choose(2*x - y, x - z);
      sum += Coeff1*Coeff2*X[y-z]*Y[z];
    }
  }
  ret.push_back(sum);
}

ということはScalarEvolution.cppの3311行目から3314行目choose(x, 2x)*choose(2x-y, x-z)と書いてあるところは間違い(choose(x-1, 2x-y-1)*choose(2x-y-1, x-z)が正しい?)のようです。

これを動かすと、以下のようになります。

x = 0
  y = 0
    z = 0
      Choose(0,0)Choose(0,0)X[0]Y[0]
x = 1
  y = 1
    z = 0
      Choose(1,1)Choose(1,1)X[1]Y[0]
    z = 1
      Choose(1,1)Choose(1,0)X[0]Y[1]
  y = 2
    z = 1
      Choose(1,0)Choose(0,0)X[1]Y[1]
x = 2
  y = 2
    z = 0
      Choose(2,2)Choose(2,2)X[2]Y[0]
    z = 1
      Choose(2,2)Choose(2,1)X[1]Y[1]
  y = 3
    z = 1
      Choose(2,1)Choose(1,1)X[2]Y[1]
  y = 4
    満たすzなし
x = 3
  y = 3
    z = 1
      Choose(3,3)Choose(3,2)X[2]Y[1]
  y = 4
    満たすzなし
  y = 5
    満たすzなし

これを計算すると、一つ目の係数は0、二つ目の係数はChoose(1,0)Choose(0,0)X[1]Y[1]=1、三つ目の係数はChoose(2,2)Choose(2,1)X[1]Y[1]+Choose(2,1)Choose(1,1)X[2]Y[1]=6、四つ目の係数はChoose(3,3)Choose(3,2)X[2]Y[1]=6、となり、無事に{0,+,1,+,6,+,6}が導出できました。

retcubeの総和なので、もう一レベル上がって{0,+,0,+,1,+,6,+,6}になります。これは  0\times{}_iC_0 + 0\times{}_iC_1 + 1\times{}_iC_2+ 6\times{}_iC_3+ 6\times{}_iC_4です。 ループ脱出時にはi == nになっているので、  0\times{}_nC_0 + 0\times{}_nC_1 + 1\times{}_nC_2+ 6\times{}_nC_3+ 6\times{}_nC_4になっているはずです。 これで、結論に達することができました。

clangの出力が汚い理由

上では、clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt --sroa --indvars -S square_sum_with_end.llという、本質部分だけを抜き出した最適化により得られたきれいなコードを確認しました。 しかし普通にclang -O2としたときに出力される機械語コードのアルゴリズムは、これより汚いです。 同様にオプションを一つずつ削ることで、アルゴリズムが汚くなる原因が--loop-rotateにあることがわかります。 この最適化パスは、while文をdo文に変えるような最適化を行うようです。 つまり、

// whileっぽい方式
    unsigned long long sum = 0, i = 0, end = n;
loop:
    if( i >= end ) { goto loop_end; }
    sum += i;
    ++i;
    goto loop;
loop_end:
    printf( "%llu\n", sum );

とするとループ一周当たりに分岐が二回(最後以外成立しないif文と、loopに戻るgoto文)ですが、

// doっぽい方式
    unsigned long long sum = 0, i = 0, end = n;
    if( end == 0 ) { goto loop_end; }
loop:
    sum += i * i;
    ++i;
    if( i < end ) { goto loop; }
loop_end:
    printf( "%llu\n", sum );

のようにすれば、ループ脱出の条件分岐とloopに戻る分岐を兼ねられて、ループ一周当たりに分岐が一回になります。 この最適化をLoop Rotation最適化というようです。

if( i >= end ) goto loop_end; sum += i * i; ++i;                               goto loop;

                              sum += i * i; ++i; if( i >= end ) goto loop_end; goto loop;

に変わる(ループ本体が文単位で見て左ローテートになっている)あたりが、Rotationと呼ばれる理由なのでしょう。

この最適化は分岐命令の実行数を減らす点では有用ですが、どうもInduction Variable Simplificationとの相性が悪いようです。 clang+llvm-12.0.1-x86_64-linux-gnu-ubuntu-16.04/bin/opt --loop-rotate --sroa --indvars -S square_sum_with_end.llとすると、条件がfalseの分岐が取り残されます。 これを見るに、必ず一回は実行されるループ構造になっているのが原因なのか、最後の一周が取り残されているような雰囲気があります。 実際、optの出力が計算しているのは  {}_{n-1}C_{1} + 3\,{}_{n-1}C_{2} + 2\,{}_{n-1}C_{3}で、一周分少なくなっていることがわかります。 おそらく、{0,+,0,+,1,+,2}(ループの最終周が始まるときのretの値)と{0,+,1,+,2}(ループの最終周におけるi * iの値)を足して{0,+,1,+,3,+,2}になったのでしょう(誘導変数の加算は要素ごとの加算でできます)。

まとめ

  • clangは誘導変数最適化の枠組みで、 \sum_{i=0}^{n-1} i = \frac{n(n-1)}2のような最適化を行っている
    • 具体的には、optの中のInduction Variable Simplification最適化パスで計算されている
  • clangは、SCEVAddRecExprという形で加法的な誘導変数を高度に解析している
    • 加法的な誘導変数であれば、それが k次の誘導変数であってもループの各回や終了時点の値を直接算出できる
      • cube += tmp + 1tmp += six_i + 6six_i += 6のように、誘導変数の連鎖で表す
      •  {}_{i}C_{\kappa} (0\le\kappa\le k)を基底とした表現で持っているのと実質的に同じ
    • 誘導変数同士の演算結果がどのような誘導変数になるかが計算できる
      • 加算は自明(係数の要素ごと加算でよい)
      • 乗算は有理数演算が不要な公式がある
  • 普通に-O2などの最適化オプションを使う場合、Rotate Loop最適化パスが先に実行されるためにきれいなコードとならない
    •  {}_{n-1}C_{\kappa} (0\le\kappa\le k)を基底とした表現で計算されてしまう

関連情報

*1:i*i<5みたいな条件式の時、二次方程式を解けばiの範囲を特定することができます。isQuadratic関数は、このために使われるようです。