2024年07月12日

LLVM18の驚異的な最適化はどのように実装されているのか調べてみました

ずっと問題なく動いていたコードが LLVM 18 から動かなくなってしまったことをきっかけに始まった調査でしたが、気が付けばけっこうな分量になりました。

SoftFloatの未定義動作バグ(1)signedのunsignedな絶対値を求める際にINT_MIN
SoftFloatの未定義動作バグ(2)RISC-VのRV64Iではunsignedの32bit即値でも64bitレジスタの上位32bitが0とは限らない
RV64Iでunsignedの32bit値が符号拡張されないで関数にレジスタ渡しされるのはどういう時か?
SoftFloatの未定義動作バグ(3)そもそも単精度浮動小数点数演算をソフトウェアエミュレーションする関数の仮引数はfloatにするべき

そして、コンパイラの最適化のバグではなく、コードにもともとあった問題が顕在化したという結論になりました。詳細は上記リンクを参照してください。

その調査の過程で LLVM 18 の驚異的な最適化を目の当たりにしました。以下のように float を uint32_t として扱い、符号ビットを比較するコードが…
    aSign = a >> 31;
    bSign = b >> 31;
    if ( aSign != bSign ) { ... }
RV64I ターゲットでは uint32_t であっても常に 64 bit に符号拡張されるという RISC-V の呼び出し規約を利用し、(64 bit の)xor を取って符号ビット(bit 31 ではなく bit 63)が 1、つまり負の数ならば符号ビットが異なるので次にジャンプ(Branch Less Than Zero; bltz)というコードが生成されました。
  xor         a2, a1, a0
  bltz        a2, 0x1c
LLVM 17 までは、以下のように下位 32 bit を 31 bit 論理右シフト(Shift Right Logical Immediate Word; srliw)して比較(Branch Not Equal; bne)という素直なコードが出ていました。そのため問題のあるコードでも期待通り動作していたのです。
  srliw       a2, a0, 0x1f
  srliw       a3, a1, 0x1f
  bne         a2, a3, 0x38
こんな最適化、いったいどうやったら実現できるのでしょうか?調べてみました。



まず最初に、一切ライブラリを使用せず、ビルドした Clang コンパイラ単体でコンパイルできるテストコードを作成したのですが、その過程で以下のように簡単にしすぎると LLVM 17 でも最適化できてしまうことがわかりました。
typedef unsigned int float32;
typedef int flag;

static inline flag extractFloat32Sign( float32 a )
{
    return a>>31;
}

flag is_same_sign( float32 a, float32 b )
{
    flag aSign = extractFloat32Sign( a );
    flag bSign = extractFloat32Sign( b );
    if ( aSign != bSign ) return 0;
    return 1;
}
先ほどと同様にシフトして比較が xor に最適化され、Set Less Than Immediate で負の場合(a0 < 0)は 1 をセットして、 最後に 1 と xor して反転しています。
is_same_sign:                           # @is_same_sign
# %bb.0:
	xor	a0, a1, a0
	slti	a0, a0, 0
	xori	a0, a0, 1
	ret
LLVM の -print-after-all オプションで、どこの最適化パスでシフトが xor に最適化されるのか調べてみます。
$ clang -v
clang version 17.0.4
$ clang --target=riscv64-unknown-elf -march=rv64ima -mabi=lp64 -Wall -Wextra -O2 -S -mllvm -print-after-all is_same_sign.c 
すると InstCombinePass 終了後に lshr(Logical SHift Right)が xor に変わりました。おそらくこのパスが LLVM 18 で強化されたのではないか?と予想できます。
*** IR Dump After SimplifyCFGPass on is_same_sign ***
; Function Attrs: nounwind
define dso_local signext i32 @is_same_sign(i32 noundef signext %0, i32 noundef signext %1) local_unnamed_addr #0 {
  %3 = lshr i32 %0, 31
  %4 = lshr i32 %1, 31
  %5 = icmp eq i32 %3, %4
  %6 = select i1 %5, i32 1, i32 0
  ret i32 %6
}
*** IR Dump After InstCombinePass on is_same_sign ***
; Function Attrs: nounwind
define dso_local signext i32 @is_same_sign(i32 noundef signext %0, i32 noundef signext %1) local_unnamed_addr #0 {
  %3 = xor i32 %0, %1
  %4 = icmp sgt i32 %3, -1
  %5 = zext i1 %4 to i32
  ret i32 %5
}
今度はもう少しテストコードを複雑にして、LLVM 17 では xor 最適化できず、LLVM 18 では最適化できる例を作ります。とはいえ、もともとのコードが十分簡単かつ条件を満たしているので、そのまま使います。
typedef unsigned int float32;
typedef int flag;
typedef unsigned int bits32;

static inline flag extractFloat32Sign( float32 a )
{
    return a>>31;
}

flag float32_lt( float32 a, float32 b )
{
    flag aSign = extractFloat32Sign( a );
    flag bSign = extractFloat32Sign( b );
    if ( aSign != bSign ) return aSign && ( (bits32) ( ( a | b )<<1 ) != 0 );
    return ( a != b ) && ( aSign ^ ( a < b ) );
}
LLVM 17 では最適化できず、最後まで lshr のまま。
$ clang -v
clang version 17.0.4
$ clang --target=riscv64-unknown-elf -march=rv64ima -mabi=lp64 -Wall -Wextra -O2 -S -mllvm -print-after-all float32_lt.c

float32_lt:                             # @float32_lt
# %bb.0:
	srliw	a2, a0, 31
	srliw	a3, a1, 31
	bne	a2, a3, .LBB0_2
LLVM 18 では、やはり InstCombinePass で最適化されています。
$ clang -v
clang version 18.1.7
$ clang --target=riscv64-unknown-elf -march=rv64ima -mabi=lp64 -Wall -Wextra -O2 -S -mllvm -print-after-all float32_lt.c

float32_lt:                             # @float32_lt
# %bb.0:
	xor	a2, a1, a0
	bltz	a2, .LBB0_2

; *** IR Dump After SimplifyCFGPass on float32_lt ***
; Function Attrs: nounwind
define dso_local signext i32 @float32_lt(i32 noundef signext %0, i32 noundef signext %1) local_unnamed_addr #0 {
  %3 = lshr i32 %0, 31
  %4 = lshr i32 %1, 31
  %5 = icmp eq i32 %3, %4
  br i1 %5, label %13, label %6
(省略)
}
; *** IR Dump After InstCombinePass on float32_lt ***
; Function Attrs: nounwind
define dso_local signext i32 @float32_lt(i32 noundef signext %0, i32 noundef signext %1) local_unnamed_addr #0 {
  %3 = xor i32 %0, %1
  %4 = icmp sgt i32 %3, -1
  br i1 %4, label %11, label %5
(省略)
}
だいぶ絞り込まれてきました。続いて LLVM 17(の最新 17.0.6)と LLVM 18(の最新 18.1.8)の llvm/lib/Transforms/InstCombine/ 以下の diff を取り、怪しそうなところをコメントアウトしてビルドし、生成コードの変化を調べました。その結果、LLVM 18 で追加された以下のコードをコメントアウトすると最適化が行われないことがわかりました。また、ここをコメントアウトしても LLVM 17 と同様の最適化は行われたので、2 つは異なる最適化であることがわかりました。
llvm-18.1.8\llvm\lib\Transforms\InstCombine\InstCombineCompares.cpp

7163 行目

    // Signbit test folds
    // Fold (X u>> BitWidth - 1 Pred ZExt(i1))  -->  X s< 0 Pred i1
    // Fold (X s>> BitWidth - 1 Pred SExt(i1))  -->  X s< 0 Pred i1
    Instruction *ExtI;
    if ((I.isUnsigned() || I.isEquality()) &&
        match(Op1,
              m_CombineAnd(m_Instruction(ExtI), m_ZExtOrSExt(m_Value(Y)))) &&
        Y->getType()->getScalarSizeInBits() == 1 &&
        (Op0->hasOneUse() || Op1->hasOneUse())) {
      unsigned OpWidth = Op0->getType()->getScalarSizeInBits();
      Instruction *ShiftI;
      if (match(Op0, m_CombineAnd(m_Instruction(ShiftI),
                                  m_Shr(m_Value(X), m_SpecificIntAllowUndef(
                                                        OpWidth - 1))))) {
        unsigned ExtOpc = ExtI->getOpcode();
        unsigned ShiftOpc = ShiftI->getOpcode();
        if ((ExtOpc == Instruction::ZExt && ShiftOpc == Instruction::LShr) ||
            (ExtOpc == Instruction::SExt && ShiftOpc == Instruction::AShr)) {
          Value *SLTZero =
              Builder.CreateICmpSLT(X, Constant::getNullValue(X->getType()));
          Value *Cmp = Builder.CreateICmp(Pred, SLTZero, Y, I.getName());
          return replaceInstUsesWith(I, Cmp);
        }
      }
    }
git blame すると
b22917e6e2a0a (XChy                     2023-10-13 22:02:57 +0800 7476)     // Signbit test folds

    [InstCombine] Fold Ext(i1) Pred shr(A, BW - 1) => i1 Pred A s< 0 (#68244)
    
    Resolves #67916 .
    This patch folds `Ext(icmp (A, xxx)) Pred shr(A, BW - 1)` into `i1 Pred
    A s< 0`.
    [Alive2](https://alive2.llvm.org/ce/z/k53Xwa).
もともとは !a == (a < 0) が最適化されないという issue に対するパッチのようです。

https://github.com/llvm/llvm-project/issues/67916
https://github.com/llvm/llvm-project/pull/68244

過去記事
これ、コンパイラが unsigned でも 32 bit 値の 31 bit 目の比較は bltz(if a2 < 0 branch 0x1c)を使った符号判定に置き換えられるという知識を持ってるってことなんですかね?
とか書いてましたが、まさか本当にそのまま予想通りのルールが埋め込まれていたとは。

コードは難しいですが、コメントを参考にすると、このルールは直接 xor に変換するものではなく、これによって別の最適化が可能になってる感じですかね。

では、シフトを xor に変換しているルールはズバリどこなのでしょうか?

これは diff などの絞り込むための手がかりが無いためなかなか難しく、諦めかけたのですが、ふと上記のパッチを眺めている時に
// (A >> C) == (B >> C) --> (A^B) u< (1 << C)
というコメントを見つけました。そしてこれが当たりでした。以下のコードをコメントアウトすると xor に変換されなくなります。そしてこのコードは LLVM 17 から変わってないです。
llvm-18.1.8\llvm\lib\Transforms\InstCombine\InstCombineCompares.cpp

5349 行目

  // (A >> C) == (B >> C) --> (A^B) u< (1 << C)
  // For lshr and ashr pairs.
  const APInt *AP1, *AP2;
  if ((match(Op0, m_OneUse(m_LShr(m_Value(A), m_APIntAllowUndef(AP1)))) &&
       match(Op1, m_OneUse(m_LShr(m_Value(B), m_APIntAllowUndef(AP2))))) ||
      (match(Op0, m_OneUse(m_AShr(m_Value(A), m_APIntAllowUndef(AP1)))) &&
       match(Op1, m_OneUse(m_AShr(m_Value(B), m_APIntAllowUndef(AP2)))))) {
    if (AP1 != AP2)
      return nullptr;
    unsigned TypeBits = AP1->getBitWidth();
    unsigned ShAmt = AP1->getLimitedValue(TypeBits);
    if (ShAmt < TypeBits && ShAmt != 0) {
      ICmpInst::Predicate NewPred =
          Pred == ICmpInst::ICMP_NE ? ICmpInst::ICMP_UGE : ICmpInst::ICMP_ULT;
      Value *Xor = Builder.CreateXor(A, B, I.getName() + ".unshifted");
      APInt CmpVal = APInt::getOneBitSet(TypeBits, ShAmt);
      return new ICmpInst(NewPred, Xor, ConstantInt::get(A->getType(), CmpVal));
    }
  }
このコードを release-17.x ブランチで git blame すると
10494b268239b (Sanjay Patel             2016-09-16 16:10:22 +0000 5205)   // (A >> C) == (B >> C) --> (A^B) u< (1 << C)
10494b268239b (Sanjay Patel             2016-09-16 16:10:22 +0000 5206)   // For lshr and ashr pairs.
0eff6c6ba81c4 (Chenbing Zheng           2022-06-20 10:55:47 +0800 5207)   const APInt *AP1, *AP2;
2016 年ごろからあるようです。

コメントアウトした場合の生成コードはそれぞれ以下のようになります。
is_same_sign:                           # @is_same_sign
# %bb.0:
        srliw   a0, a0, 31
        srliw   a1, a1, 31
        xor     a0, a0, a1
        seqz    a0, a0
        ret
float32_lt:                             # @float32_lt
# %bb.0:
        srliw   a2, a0, 31
        srliw   a3, a1, 31
        bne     a2, a3, .LBB0_2
最適化ルールの実装は突き止められましたが、なぜ LLVM 17 ではこのルールが float32_lt には適用できないのかはよくわかりませんでした。

そもそも当初の目的は、シフトが xor になるなんて、いったいどんな最適化をすれば実現できるんだろう…という疑問だったのですが、実は「(A >> C) == (B >> C) --> (A^B) u< (1 << C)」というルール一発だったというオチでした。

issue トラッカーを見るに、LLVM のオプティマイザの開発者は一つ一つの最適化ルールの正しさを専用の検証ツールで証明して添付している感じなんですね。こんな膨大な最適化ルールの集合体がちゃんと動いているってすごいな…と改めて思いました。

kmckk at 18:59コメント(0)Clang | 若槻 

コメントする

名前
 
  絵文字
 
 
記事検索
最新コメント
アクセスカウンター
  • 今日:
  • 昨日:
  • 累計:

QRコード
QRコード