2021.08.22

Juliaコンパイラ開発レポート

門脇宗平

こんにちは。 ATL客員研究員の門脇宗平です。

普段は、MITのJuliaLabという組織の"Research Programmer"という肩書きで、プログラミング言語Juliaのコンパイラやその周辺技術の開発に携わっています。今回は自己紹介も兼ねて、自分の簡単な経歴と今取り組んでいるプロジェクトたちについて紹介したいと思います。

自己紹介

僕は2020年4月に京都大学の総合人間学部というちょっと不思議な名前の学部を卒業しました。総合人間学部はいわゆる教養学部のような場所で、大変自由でのんびりした学部です。そんな時間がたくさんあった学部時代には、海外に住んでみたくてカナダのブリティッシュコロンビア大学に交換留学に行ってみたり、当時の機械学習ブームに乗って人工知能系の授業をつまみ食いしてみたり、休学している間にリクルートのデータサイエンスコースの長期インターン¹に参加して就業経験を積んだりなどしました。卒論のテーマも各自好きに選んでよしとのことだったので、インターンで使用してから好きになったJulia言語について卒論を書きました。

卒業後はリクルート住まいカンパニーに新卒入社しました。住まいカンパニーでは、データサイエンスチーム所属のソフトウェアエンジニアというポジションでのびのびと働くことができ、新しいデータ品質検査フレームワークの設計開発などといった、新規性が高くやりがいのある案件を担当しエンジニアとしての経験を積むことができました。また、入社したときにちょうどコロナ禍に入ってしまったのですが、そのおかげで時間に余裕があったので、仕事のかたわらで卒論の続きの開発を進めてみたり、初めて自分のブログを書いたりもしました。

すると、僕のブログを読んでくれたJuliaLabのPhDの方が声をかけてくれ、ラボ付きの技官のような肩書きで働いてみないかという話になり、今年の6月から今のポジションで働いています。

普段のプロジェクト

それでは僕が今取り組んでいるプロジェクトについて説明したいと思います。

Juliaはオープンソースの文化の中で育った言語なので、JuliaLabのプロジェクトも基本的にパブリックな形で進められています。僕はJuliaコンパイラに直接関わるプロジェクトに携わっているので、特に広くコミュニティと協働する機会が多く、非常にやりがいがあります。

今は以下のプロジェクトに関わっています:

エスケープ解析: オブジェクトのエスケープ情報及びライフタイムを解析し、各種の最適化へ応用する
コンパイラプラグイン: Juliaコンパイラを拡張するインフラを整備し、より応用的なコード生成を可能にする
そのほか色々な開発: Juliaコンパイラの開発、卒論のプロジェクトの続き、などなど

今回は上2つのプロジェクトについて少し掘り下げながらご紹介させていただこうと思います。

エスケープ解析

このプロジェクトは、Juliaコンパイラにオブジェクトの生存範囲やライフタイムを解析する新しいルーチンを追加し、それをメモリ効率の改善などの最適化に応用しようというものです。

Juliaプログラムは、大まかに

(ステージ1) Juliaレベルのプログラム解析と最適化
(ステージ2) LLVMによるプログラム解析と最適化及びコード生成

という2つのステージを経てコンパイルされ実行されます。 LLVMは優秀なコンパイラインフラであり、エスケープ解析を含むさまざまな解析と最適化をステージ2.で行ってくれるため、実はJuliaはすでにLLVMによるメモリ最適化の恩恵を受けています。その一方で、LLVMの段階ではJuliaレベルのセマンティクスを用いた最適化が行えないため、例えば関数呼び出しをまたいだ解析や最適化を行うことはできません。

このプロジェクトではステージ1.のJuliaレベルの段階において、 Juliaのセマンティクスをフルに利用したエスケープ情報の解析を行うことで、特に以下のような最適化を行ってJuliaプログラムのパフォーマンスを向上させることを当面の目標としています。

エスケープしないオブジェクトをヒープ領域ではなくスタック領域にアローケーションする: Juliaプログラム一般についてメモリ効率の改善が期待できます
オブジェクトのライフタイムを解析しfinalizer²をearly callする: 特にGPU計算などメモリプレッシャーの高い計算環境において、より効率的なリソース解放によるパフォーマンスの改善が見込めます

僕たちはこのエスケープ解析を抽象解釈("abstract interpretation")と呼ばれる技術を応用して実装しています。抽象解釈は定式化された理論が与えられている古典的なプログラムの解析技術で、 Juliaコンパイラもプログラムを最適化するための型推論に抽象解釈を応用しています。

エスケープ解析を行う抽象解釈の面白い点は、型推論や定数伝播などを行う抽象解釈とは逆方向に情報が伝播するという点です。つまり、型推論では 変数の定義元から変数の利用場所へ と情報が伝播しますが、一方でエスケープ情報は 変数の利用場所から変数の定義元へ と情報が伝播するのです:

型情報: 定義元(引数型 s::String)から利用場所(関数呼び出し Ref(s::String))へと伝播

エスケープ情報: 利用場所(関数呼び出し broadcast(identity, x))から定義元(オブジェクトアロケーション x = Ref(s))へ伝播

function foo(s) 
    x = Ref(s) # `x`の生存区間はこの「定義元」だけ見てもわからない
    broadcast(identity, x) # 「利用場所」`broadcast(identity, x)`を解析してはじめて`x`の生存区間がわかる
end
foo("42") # => "42"

function foo(s)

x = Ref(s) # `x`の生存区間はこの「定義元」だけ見てもわからない

broadcast(identity, x) # 「利用場所」`broadcast(identity, x)`を解析してはじめて`x`の生存区間がわかる

end

foo("42") # => "42"

さらにこうした逆方向の抽象解釈の面白い点は、順方向の抽象解釈との組み合わせを考えることができるという点です。一般に、異なるプログラム性質を解析する抽象解釈は、別々に行われるよりも同じルーチンの中でフィードバックし合いながら行うことで、一方の解析から得られた情報を用いてもう一方の情報を更新することができ、双方の解析精度を高めることが期待できます。例えば以下のような場合に、エスケープ情報(及びエフェクト情報)により、型情報の改善が見込めます:

e.g. bar(x)がエスケープ及びメモリ書き込み効果を持たないと証明できた場合、x.xを定数で畳込むことができる

mutable struct MyRef{T}
    x::T
end
bar(x) = ...
let
    x = MyRef{Any}("42") # `x::MyRef{Any}`はミュータブルなオブジェクト
    bar(x)               # `bar(x)`内での`x`のエスケープ情報及び`x.x`フィールドへのメモリ書き込みの情報を解析
    x.x                  # `bar(x)`がエスケープ及びメモリ書き込みの効果を持たないと証明できた場合、`x.x::String`というより正確な型情報を得ることができる
end

mutable struct MyRef{T}

x::T

end

bar(x) = ...

let

x = MyRef{Any}("42") # `x::MyRef{Any}`はミュータブルなオブジェクト

bar(x) # `bar(x)`内での`x`のエスケープ情報及び`x.x`フィールドへのメモリ書き込みの情報を解析

x.x # `bar(x)`がエスケープ及びメモリ書き込みの効果を持たないと証明できた場合、`x.x::String`というより正確な型情報を得ることができる

end

ただし、双方向の抽象解釈を組み合わせた場合にいかにして解析の収束性を保つことができるか、という点は自明ではありません。僕の知る限り、双方向の抽象解釈をJITコンパイラで利用している例は聞いたことがなく、実際に僕たちのエスケープ解析も現時点では型推論(順方向解析)とは独立したルーチンとして実装しています。

一方で、僕はこうした双方向の解析を用いることで、Juliaコンパイラを新たな次元に押し上げることができると考えています。そうした意味でこのプロジェクトはメモリ効率の最適化だけにとどまらない可能性を秘めています。

このプロジェクトはGoogle Summer of CodeというGoogleが主催するイベントの一環として、僕がメンターを務めるImplement Escape Analysis in Julia Compilerというプロジェクトとして始まりましたが、今年度のGSoCのタームに限らず、今後も継続して長期的に取り組む予定です。

解析ルーチンの開発はaviatesk/EscapeAnalysis.jlで行っておりますので、興味のある方はぜひ覗いてみて下さい！

コンパイラプラグイン

JuliaはLispのような非常に強力なメタプログラミング機能を持つ言語です。シンタックスレベルでのコード生成を行うことができるマクロはもちろん、 Juliaレベルでの型推論に介入するステージングプログラミングのインターフェースも提供されています。

どちらも非常に強力な言語機能ですが、特にステージングプログラミングの機能を利用することで、ライブラリレベルでユーザコードを自由に書き換えることが可能であり、実際に逆方向自動微分の実装などに使用されています。例えば、Juliaコミュニティで広く利用されているZygote.jlという自動微分ライブラリは、まさにこのステージング機能を用いて実装されています⁴。コンパイラのコード生成に介入してコード変換を行うことにより、ホスト言語(Julia)で書かれたプログラムをそのまま(自動的に)微分することが可能であり、かつ、変換後のコードは(理論的には)通常のプログラム実行と同じように効率的に動作します:

e.g. Zygote.jlによるユーザ定義の累乗計算関数pow(x, n)の微分

julia> function pow(x, n)
           r = 1
           for i = 1:n # コントロールフローなどのホスト言語の機能をサポート
               if i % 10 == 0 # I/Oなど複雑な処理もそのまま微分可能
                   println("itr: $i")
               end
               r *= x
           end
           return r
       end;
julia> using Zygote
julia> @time gradient(x->pow(x,12), 2)
itr: 10
 23.595115 seconds (41.45 M allocations: 2.378 GiB, 4.44% gc time, 99.99% compilation time)
(24576,)
julia> 12*2^11
24576
julia> @time gradient(x->pow(x,20), 2)
itr: 10
itr: 20
  0.041844 seconds (23.41 k allocations: 1.392 MiB, 98.57% compilation time)
(10485760,)

julia> function pow(x, n)

r = 1

for i = 1:n # コントロールフローなどのホスト言語の機能をサポート

if i % 10 == 0 # I/Oなど複雑な処理もそのまま微分可能

println("itr: $i")

end

r *= x

end

return r

end;

julia> using Zygote

julia> @time gradient(x->pow(x,12), 2)

itr: 10

23.595115 seconds (41.45 M allocations: 2.378 GiB, 4.44% gc time, 99.99% compilation time)

(24576,)

julia> 12*2^11

24576

julia> @time gradient(x->pow(x,20), 2)

itr: 10

itr: 20

0.041844 seconds (23.41 k allocations: 1.392 MiB, 98.57% compilation time)

(10485760,)

その一方で、現行のステージングプログラミングのインターフェースにおいては以下のような問題点も存在します:

プログラムの書き換えによるコンパイルタイムのレイテンシが大きい (compile-time latency)
通常のプログラム実行よりも型推論自体の精度が落ちてしまう (limited inference)
型推論直前のステージにしか介入できない (limited customization)

例えば上で紹介したZygote.jlを使用した自動微分では、gradient関数を一番最初に呼び出した時のJITコンパイルの待ち時間(レイテンシ)がかなりかかってしまっており⁵、1つ目の問題点が如実に現れています。

このプロジェクトの目的は、こういった欠点を解消した、さらに強力なステージングプログラミングのインターフェースを用意し、Juliaコンパイラにより高い拡張性を持たせようというものです。コンパイラプラグインにより、機械学習などの分野における計算の基盤となっている自動微分技術が恩恵を受けるのはもちろん、数式処理システムなどの領域へも応用が期待できるほか、ライブラリレベルでの最適化ルーチンの実装が可能になったりします。

このプロジェクトは、Keno FisherさんをはじめとしたJuliaのコンパイラチームの方々からのヘルプを受けながら、広くコミュニティの方々とコミュニケーションをしつつ、僕がプロジェクトリードとして進めています。デザインドキュメントをまとめているほか、すでにWIPのPRも出ているので、もし興味のある方がいらっしゃいましたらぜひ読んでみて下さい。フィードバックや質問なども大歓迎です。 Julia Slackの#compiler-pluginチャンネルでお気軽にコンタクトして下さい。

そのほか色々な開発

上記のプロジェクトの合間に、Juliaに関する色々な開発も行っています。 Juliaの型推論ルーチンの改善や各種最適化機能の充実の他に、僕の卒論のプロジェクトであるJuliaの型推論を用いた静的解析ツールJET.jlの開発などに取り組んでいます。特にJETについては、今年のJuliaConでの発表など(talk, workshop)力を入れており、また別の機会で改めてご紹介させていただければと思っております。

おわりに

今回は、僕の簡単な自己紹介と今取り組んでいるプロジェクトについてご紹介させていただきました。今後もこういった内容のブログを執筆させていただく予定ですので、どうぞよろしくお願いいたします。

リンク

リクルートの技術系のインターンは、実際の開発/解析の現場に入ってかつとても自由に働くことができるとても良い機会なので、ぜひチェックしてみてください。↩︎
オブジェクトのメモリ領域が解放される(つまりGCがオブジェクトを回収する)タイミングで呼ばれるフックで、オブジェクトが利用する外部リソースの解放などに利用されます。↩︎
Zygote.jlについては作者のMike J. Innesさんの論文やハンズオンを読んでみると面白いかもしれません。僕のTwitterでも簡単な解説をしています: https://twitter.com/kdwkshh/status/1414068340335599618 ↩︎
初回の自動微分(gradient(x->pow(x,12), 2))の実行時間23.59秒のうち、なんと99.99%がコンパイルにかかっています(一方で 2回目の呼び出し(gradient(x->pow(x,20), 2))ではキャッシュされたコンパイル済みコードが実行されるのでコンパイルにかかる時間はかからず、高速に動作しているのが確認できると思います)。↩︎

Recruit Tech Blog

タグから見つける