Pruning dead exception handlers

原文はこちら。
The original article was written by Jorn Vernee (Principal Member Of Technical Staff at Oracle).
https://jornvernee.github.io/hotspot/jit/2024/02/16/prune-dead-exception-handlers.html

JDK 22では、しばらくの間FFM APIを悩ませていたパフォーマンスの問題を解決しました。新しく作成されたArenatry-with-resourcesを使用すると、過剰なアロケーションが発生していました。これは、例外ハンドラ(catchブロック)のデッドコードがC2 JITコンパイラによって削除されないために発生していました。これをJDK 22で解決しました。

8267532: C2: Profile and prune untaken exception handlers #16416
https://github.com/openjdk/jdk/pull/16416

この最適化の効果は、try-with-resourcesを使用するJavaコード、または単に使われないcatchブロックに広く適用できます。そこで、このエントリでこの問題について議論し、その解決方法を説明するのが面白いと思いました。

ちなみに、タイトルにある”pruning”とは、ガーデニングの際に行う剪定のことで、枝などの植物の一部を切り落とすことを指します。これは、JITコンパイル中にコードの枯れ枝を「切り落とす」方法と似ているからです。

Why are my allocations escaping?

この問題に遭遇したケースは、FFM APIを使用する際にかなり一般的なものです。

try (Arena arena = Arena.ofConfined()) {
    MemorySegment segment = arena.allocateFrom("Hello!");
    func(segment);
}

このコードは、新しいarenaを作成し、そのarenaにいくつかのデータを割り当て、そのデータへのポインタをネイティブ関数(func)に渡しています。funcの実装は、 java.lang.foreign.Linker::downcallHandle が生成するネイティブメソッドハンドルを使用して呼び出しを転送するだけだとします。

このコードにはいくつかのアロケーションがありますが、その中でもArenaMemorySegmentが最も顕著なものです。しかし、funcの実装はネイティブ関数にプリミティブアドレスを渡すだけなので、アロケートされたオブジェクトはどれもより広いJavaコードにエスケープされません。理論的には、C2はこれらのオブジェクトをスカラー置換し、割り当てを完全に回避できるはずです。

しかし、私の別の投稿で説明したテクニックを使ってこれを検証すると、エスケープしている割り当てがいくつかあることが判明しました。

Tracking down escaping objects
https://jornvernee.github.io/hotspot/jit/2023/08/18/debugging-jit.html#5-tracking-down-escaping-objects

JDK 21では、このコードには以下のエスケープ割り当てがあります。

JavaObject(38) allocation in: MemorySessionImpl::createConfined @ bci:0 (line 145)
  -> Field(63)
  -> JavaObject(40)
  -> LocalVar(117)
  -> LocalVar(157)
  Reason: Escapes as argument to call to: jdk.internal.foreign.MemorySessionImpl$1::close void ( jdk/internal/foreign/MemorySessionImpl$1 (java/lang/AutoCloseable,java/lang/foreign/Arena,java/lang/foreign/SegmentAllocator):NotNull * ) TestArena::payload @ bci:36 (line 25)


JavaObject(39) allocation in: ConfinedSession::<init> @ bci:2 (line 55)
  -> Field(45)
  -> JavaObject(38)
  -> Field(63)
  -> JavaObject(40)
  -> LocalVar(117)
  -> LocalVar(157)
  Reason: Escapes as argument to call to: jdk.internal.foreign.MemorySessionImpl$1::close void ( jdk/internal/foreign/MemorySessionImpl$1 (java/lang/AutoCloseable,java/lang/foreign/Arena,java/lang/foreign/SegmentAllocator):NotNull * ) TestArena::payload @ bci:36 (line 25)


JavaObject(40) allocation in: MemorySessionImpl::asArena @ bci:0 (line 80)
  -> LocalVar(117)
  -> LocalVar(157)
  Reason: Escapes as argument to call to: jdk.internal.foreign.MemorySessionImpl$1::close void ( jdk/internal/foreign/MemorySessionImpl$1 (java/lang/AutoCloseable,java/lang/foreign/Arena,java/lang/foreign/SegmentAllocator):NotNull * ) TestArena::payload @ bci:36 (line 25)


JavaObject(41) allocation in: NativeMemorySegmentImpl::makeNativeSegment @ bci:112 (line 136)
  -> Field(67)
  -> JavaObject(39)
  -> Field(45)
  -> JavaObject(38)
  -> Field(63)
  -> JavaObject(40)
  -> LocalVar(117)
  -> LocalVar(157)
  Reason: Escapes as argument to call to: jdk.internal.foreign.MemorySessionImpl$1::close void ( jdk/internal/foreign/MemorySessionImpl$1 (java/lang/AutoCloseable,java/lang/foreign/Arena,java/lang/foreign/SegmentAllocator):NotNull * ) TestArena::payload @ bci:36 (line 25)

これらのオブジェクトのエスケープ・ルートを見ると、以下のことに気づくことでしょう。

  • JavaObject(38)がエスケープするのはJavaObject(40)がエスケープするから
  • JavaObject(39)がエスケープするのはJavaObject(38)がエスケープするから
  • JavaObject(41)がエスケープするのはJavaObject(39)がエスケープするから

言い換えれば、オブジェクト・グラフ全体が、JavaObject(40)と一緒にエスケープします。JavaObject(40)は、jdk.internal.foreign.MemorySessionImpl$1::closeの引数ならびにアウトオブライン呼び出しとしてエスケープされます。

Why is this call not being inlined?

では、なぜここでアウトオブライン呼び出しがあり、オブジェクト・グラフをエスケープしているのでしょうか?インライン化トレースを見ると、呼び出しがインライン化されていることがわかります。

Printing inlining traces
https://jornvernee.github.io/hotspot/jit/2023/08/18/debugging-jit.html#3-printing-inlining-traces

TestArena::payload (53 bytes)
@ 22   jdk.internal.foreign.MemorySessionImpl$1::close (8 bytes)   inline (hot)
  @ 4   jdk.internal.foreign.MemorySessionImpl::close (12 bytes)   inline (hot)
    @ 1   jdk.internal.foreign.ConfinedSession::justClose (52 bytes)   inline (hot)
     ...

いや、ちょっと待ってください。違うのかな。。

@ 36   java.lang.foreign.Arena::close (0 bytes)   virtual call
@ 36   jdk.internal.foreign.MemorySessionImpl$1::close (8 bytes)   low call site frequency

なるほど!例外が発生していないパスと、例外が発生しているパス(catchブロック内)の2つのclose呼び出しがあります。これはjavacがJavaコード用に生成したバイトコードです。

Code:
 0: invokestatic  #12                 // InterfaceMethod java/lang/foreign/Arena.ofConfined:()Ljava/lang/foreign/Arena;
 3: astore_0
 4: aload_0
 5: ldc           #18                 // String Hello!
 7: invokeinterface #20,  2           // InterfaceMethod java/lang/foreign/Arena.allocateUtf8String:(Ljava/lang/String;)Ljava/lang/foreign/MemorySegment;
12: astore_1
13: aload_1
14: invokestatic  #24                 // Method func:(Ljava/lang/foreign/MemorySegment;)V
17: aload_0
18: ifnull        52
21: aload_0
22: invokeinterface #28,  1           // InterfaceMethod java/lang/foreign/Arena.close:()V
27: goto          52
30: astore_1
31: aload_0
32: ifnull        50
35: aload_0
36: invokeinterface #28,  1           // InterfaceMethod java/lang/foreign/Arena.close:()V
41: goto          50
44: astore_2
45: aload_1
46: aload_2
47: invokevirtual #33                 // Method java/lang/Throwable.addSuppressed:(Ljava/lang/Throwable;)V
50: aload_1
51: athrow
52: return
Exception table:
   from    to  target type
       4    17    30   Class java/lang/Throwable
      35    41    44   Class java/lang/Throwable

バイトコード・インデックス (bci) の2236closeが2回呼び出されています。これはjavacfinallyブロックを翻訳した結果です。finallyブロック内のコードは、基本的に非例外パスと例外パス(例外ハンドラ内)に沿ってコピーペーストされます。例外テーブルを見ると、bci 36は例外ハンドラの中にあります。ここまではすべてチェックアウトされています。

そこで、インライン化トレースを見返してみると、通常の、例外のないパス(@ 22)に沿ったcloseへの呼び出しは期待通りにインライン化されていますが、例外ハンドラ内の close への呼び出し(@ 36)は「コールサイトの頻度が低い(low call site frequency)」ためにインライン化されていません。エスケープされたアロケーションをふり返ると、オブジェクト・グラフがエスケープされるのは、bci 36の例外ハンドラのこの呼び出しであることがわかります。

Reason: Escapes as argument to call to: jdk.internal.foreign.MemorySessionImpl$1::close ... TestArena::payload @ bci:36 (line 25)

C2は、ほとんど到達しない(頻度が低い)コールをインライン化しません。なぜなら、これらのコールはホットパス上にない可能性が高いため、インライン化によるメリットが少ないからです。私たちの具体的なケースでは、例外がスローされることがないため、例外ハンドラのcloseのコールサイトに到達することはありません。しかし、このdeadコードは、スカラー置換のような他の最適化をまだ妨害しています。というのも、例外が発生する可能性はまだ残っているため、C2が生成するコードはその可能性を考慮しなければならないからです。

しかし、C2には、実際には決して到達しないコードに対処する方法(uncommon trap)もあります。これは、(プロファイリング情報に基づいて)必要になる可能性が非常に低いコードの一部を、コードの最適化を解除してインタプリタ内で実行し続けるトラップに置き換えます。C2 は本質的に、このコードは不要であり、最適化する価値がないことを意味します。これが可能であることが、コード・インタプリタとJITコンパイルの両方が可能なmixed mode VMの強みの1つです。uncommon trapは、例えば、if/else文のありそうもない分岐を削除するために使われたりもします。

また、uncommon trapはスカラー置換されたオブジェクトを戻すことができるので、uncommon trapが割り当てられたオブジェクトを「使用」することがあっても、スカラー置換の最適化を妨げることはありません。したがって、(我々が例外ハンドラに入ることはないため)理論的には、C2がtry-with-resourcesブロックの例外ハンドラをuncommon trapに置き換えることは可能なはずです。そうすれば、オブジェクト・グラフをスカラー置換できるはずです。

Why is there no uncommon trap?

if文の分岐のような通常の分岐では、VMは分岐回数をカウントして分岐をプロファイリングします。JITはこの回数と、それを囲むメソッドの呼び出し回数を用いて、分岐の頻度を決定します。もしある分岐がヒューリスティックに「めったに」行われないと判断された場合、その分岐はuncommon trapに置き換えられます。

例外ハンドラは分岐の一種であり、例外ハンドラがカバーするコード内のどのポイントからでも分岐に入ることができます。であれば、例外ハンドラもプロファイリングできるはずですよね?ほとんどのプロファイリングは、実行中のバイトコードに基づいて行われます。例えば、if文にはgotoバイトコードが使われます。gotoバイトコードが実行されると、分岐プロファイリングも行われます。これは簡単で高速です。しかし、例外ハンドラは任意のバイトコードで開始できるので、同じ戦略を適用できません。例えばiconst_0バイトコードを実行した場合、それが例外ハンドラの最初のバイトコードであるかどうかは、バイトコードを見ただけではわかりません。おそらく、実行中のメソッドに含まれるすべての例外ハンドラの最初のバイトコードインデックスのテーブルを保持し、実行したバイトコードがそのうちの1つであれば、プロファイリングできるでしょう。しかし、これでは例外的な、つまりめったに起こらないはずのことのために、すべてのバイトコードの実行が遅くなってしまいます。これはあまり良いこととは思えません。これが例外ハンドラのプロファイリングがすぐに行われなかった主な理由の1つなのではないか、と思います。

しかし、例外がスローされると、例外ハンドラを検索するためにランタイムコール(C++コードへのアウトオブラインコール)を行います。実際に例外ハンドラーを実行するときだけ例外ハンドラーを検索すればいいという前提で、プロファイリング・コードをそのランタイム・コールの中に入れればいいのです。結局そのように実施しています。例外ハンドラを検索するたびに、そのハンドラをenteredとしてマークします。これが我々のプロファイリングです。C2が例外ハンドラを解析する際、その例外ハンドラに入ったか(entered)どうかをチェックし、もし入っていなければ、例外ハンドラの代わりにuncommon trapを挿入します。それに加えて、このuncommon trapを使って最適化を解除するときに、例外ハンドラをenteredとしてマークする必要があります。そうすることで、結局例外がスローされた場合、(例外が発生する可能性があることは明らかなので)次にコードがコンパイルされるときに別のuncommon trapを挿入しようとしなくなります。

Success!

オブジェクトがエスケープするアウトオブライン呼び出しの問題は、その呼び出しがある分岐をuncommon trapで置き換えることで回避したところ、ほとんどのオブジェクトがエスケープしなくなりました。

JavaObject(31) allocation in: SegmentFactories::allocateSegment @ bci:100 (line 158)
  Reason: MergedWithObject[other=JavaObject(1) [ [ ]]    128  ConP  === 0  [[ 216 210 209 213 643 212 59 208 214 1178 167 211 151 503 165 508 166 138 ]]  #null]

MemorySegmentだけがエスケープされるのは、別の理由によるものです(これについても修正する可能性があります)。

このちょっとしたコードサンプルのような使用例では、これで割り当てバイト数が半分になり、実行速度が2倍速くなる可能性があります。例えば、リンク先のプルリクエストに含まれるベンチマークの数字をご覧ください。

Before:

Benchmark                                                      Mode  Cnt      Score     Error   Units
ResourceScopeCloseMin.confined_close                           avgt   30     10.458 ±   0.070   ns/op
ResourceScopeCloseMin.confined_close:gc.alloc.rate.norm        avgt   30    104.000 ±   0.001    B/op

After:

Benchmark                                                      Mode  Cnt      Score     Error   Units
ResourceScopeCloseMin.confined_close                           avgt   30      4.563 ±   0.043   ns/op
ResourceScopeCloseMin.confined_close:gc.alloc.rate.norm        avgt   30     56.000 ±   0.001    B/op

この件が素晴らしいのは、これがFFM APIに役立つだけでなく、try-with-resourcescatchなどの例外ハンドラを使用するすべてのコードに役立つ可能性があるということです。従って、コードのホットパスで使われない例外ハンドラがあるようなユースケースがあれば、JDK 22でどれぐらいパフォーマンスが改善するか注目してください!

コメントを残す

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください