ClPy meets libTooling (1): Ultima

ClPyチームの今泉です。 ClPyがβ版になった際のブログ記事に、以下の記述がありました。

これらの機能（特に2つ目と3つ目）のために、LLVM/Clang(libTooling)を用いて、「生成元になるC++のソースコードを抽象構文木(AST)経由で意味解析をしてから適切な変換を施す」という技術を用いています。これについては、ここで解説すると長くなりますので、近い内に詳細を解説する記事を執筆し紹介したいと思います。

1年半ほど経ってますが「近い内」ですので、ClPyのインストール手順にClangのインストールが含まれている理由、ClPyにおけるlibToolingの利用事例の1つについて述べていきたいと思います。 なお、本記事中においてChainerはv3.3.0、CuPyはv2.1.0を指すものとします。

α版までの問題点

α版のClPyは、Chainerを完全にサポートできていませんでした。 ChainerのexampleにあるPTBやword2vecを動作させるためには、Chainerにいくつかパッチを適用する必要があったのです。 これは何故かというと、端的に言えばCuPyが扱うのはCUDA C++であり、ClPyが扱うのはOpenCL Cであることに起因します。

CuPyにはElementwiseKernelというPythonのクラスがあります。 これはざっくり言うと「入出力とCUDA C++のコード片を渡すといい感じのカーネル関数を作ってオンラインコンパイルする」というもので、 __call__() メソッドに ndarray を渡すと内部でいい感じに記述したカーネル関数を呼び出してくれるようになっています。 またCuPyの数学関数は概ねこの ElementwiseKernel を使って実装されており、いわばCuPyの肝となる存在です(ufunc とか ReductionKernel とかについては同様の状況に対して同様の対処を行っているので本稿では割愛します)。 以下に ElementwiseKernel の使用例を示します。

>>> import cupy
>>> f = cupy.ElementwiseKernel(
...     in_params='raw T x, T y',
...     out_params='T z',
...     operation='''
...         const ptrdiff_t idx[] = { i / x.shape()[1], i % x.shape()[1] };
...         z = x[idx] + y
...         ''',
...     name=u'add2d')
>>> x = cupy.arange(4).reshape(2, 2).astype('f')
>>> y = cupy.array([[2, 4], [6, 8]]).astype('f')
>>> f(x, y)
array([[  2.,   5.],
       [  8.,  11.]], dtype=float32)

typedef float T;
extern "C" __global__ void add2d(CArray<float, 2> x,
                                 CArray<float, 1> _raw_y,
                                 CArray<float, 1> _raw_z,
                                 CIndexer<1> _ind) {
  for(ptrdiff_t i = blockIdx.x * blockDim.x + threadIdx.x;
                i < _ind.size();
                i += blockDim.x * gridDim.x) {
    _ind.set(i);
    T &y = _raw_y[_ind.get()];
    T &z = _raw_z[_ind.get()];
    const ptrdiff_t idx[] = { i / x.shape()[1], i % x.shape()[1] };
    z = x[idx] + y;
  }
}

このように ElementwiseKernel は外部に公開されたインターフェースなので、CuPyの内部実装に限らず自由に独自のカーネル関数を実装することができます。 そしてChainerでは ElementwiseKernel を使ってDNN用の関数を自前でいくつも作っており、その中に以下のようなコードが含まれています。

• template
• CUDAの予約語(__device__)
• 変数 half
• ptrdiff_t の配列を渡すべきところでint の配列を渡している
  • CuPyではtemplate を用いた演算子オーバーロードによって実装されている部分なので問題なく動きます
  • ClPyではα時代は ptrdiff_t 決め打ちだったため、 sizeof(ptrdiff_t) != sizeof(int) な環境でメモリアクセスに問題が生じてしまっていました

これらはいずれもOpenCL Cでは許容されない記述です。 そりゃそうです、Cなんだから template も演算子オーバーロードもありませんし、OpenCLはCUDAじゃないので予約語も異なります。 half が変数名として使えないのは、OpenCL Cにおいて half は型名の予約語だからです。 また、上記3つ目の配列を引数に取る ndarray の operator[] によるメモリアクセスについて、α版のClPyでは単純な文字列検索を使って ident[ の形と後続の ] を探し、インデックスを計算する関数呼び出しに文字列置換という実装になっていました。 当たり前ですが、この実装だと入れ子になったときに対応できませんし、そもそも識別子と [ の間にスペースがあると検索にひっかからないという問題もありました(尤も、当時サポートを表明していた範囲のChainerやCuPyにはそのようなコードは含まれていなかったため、OpenCLでChainerを動かす、という観点では現実的に大きな問題にはなりませんでしたが)。

以上のように、実はα版のClPyはパッチを当てないとChainerのexampleのいくつかは動かない上、入力できるコードにかなり制約がありました。

Ultima

上記の問題を解決すべくClPyチームが取った手段は、 「Clangで構文解析して得たASTを使ってCUDA C++をOpenCL Cに変換しちゃおうぜ」でした。 構文解析を行えば上記の文字列置換による実装の入れ子や空白の問題は自明に解決できます。 C++の構文解析をフルスクラッチで実装するのは現実的ではないため、Clangを用いてASTを取得します。 またASTが取れているため、 template や関数オーバーロード、CUDAの予約語や変数名としての half については、後述する方法で変換していきます。 なお念の為明記しておきますが、これはCUDA C++のフルセットからOpenCL Cのフルセットへの変換は目指していません。 当然ですがソースコードレベルでC++からCに完全に変換し切ることは技術的に非常に難しいことですし(例外やデストラクタの扱い、コンパイル時処理については考えたくありません…)、ClPyの一部としてそこまでの機能を提供するのはオーバースペックなためです。 あくまでもChainerのコードが問題なく動かせる程度、CUDA C++のサブセット(Chainerが使うくらい)からOpenCL Cのサブセット(ClPyで実現したいくらい)への変換となります。

上述のコード変換器はUltimaという名前でClPyに含まれています。 ベースとなるのは実はただのpretty printer(ソースコードを読み込んで、構文解析を行った後、解析で得た構文木を基に入力と同内容のコードを整形して出力するだけのプログラム)で、そこに「CUDAの予約語だったら出力しない」などの変更を複数加えることでCUDA C++(のサブセット)をOpenCL C(のサブセット)に変換するプログラムにしたのがUltimaです。 以下ではまず技術選択としてlibToolingを採用するに至った経緯を述べ、続いて主要な変更点として、 template や関数オーバーロード、C++のクラスとコンストラクタ、 ndarrayへの要素アクセスの対応、 half 識別子への対応の4つについて述べます。

libToolingを選択した理由

Clangを用いた構文解析を行うところまではよいとして、Clang ASTを扱うプログラムを書く方法は3つあります。

• libClang
• Clangのプラグインとして実装
• libTooling

ClPyはPythonモジュールですから、libClangのPython bindingsを使用したほうが良いのでは？とお考えの読者の方もいらっしゃるかもしれません。 私もそう思っていた時期がありました。 しかし、C++の演算子オーバーロードがまともに扱えない(include先でオーバーロードされている演算子は普通の演算子として扱う)上に、構文木上の演算子のノードにアクセスしても「こいつは二項演算子」というレベルの情報しか取り出せず、具体的な演算子がなんなのかは直接取得できません¹。 技術選択で真っ先に検討しましたが、即ボツになりました。

次にClangのプラグインとして実装する場合ですが、これはClangのコンパイル前にASTに対してなにか処理を行うための選択肢です。 ClangベースのOpenCL Cコンパイラに組み込む形で実装する場合には検討の余地がありますが、今回はそうではありませんので、これも無しです。

というわけで、消去法的にlibToolingしか選択肢が残りませんでした…。 libToolingは使ってみたところ、必要な機能が揃っていることが確認できたため採用しました。

template 関数と関数オーバーロード

C++では関数オーバーロードができます。 関数オーバーロードは、内部ではマングリング(name mangling)という処理で引数の型情報などが関数シンボル名にくっつけられており、 これによってリンカは呼び出し先の関数を一意に特定してリンクできます。 もうおわかりですね？Ultimaはソースコードレベルで関数名をマングリングします。

Ultimaを使うと、

void f(){}
void f(int arg){}

void f() 
{
}
void f__left_paren__int__right_paren__(int arg) 
{
}

template も同様です。全てのインスタンス化された関数をマングリングしてやれば解決します。 なお、インスタンス化されたASTはClangで作ってくれるので、Ultimaとしては名前だけ気をつけて出力するだけです。

なお、入力として f__left_paren__int__right_paren__ という名前の関数を別途定義した場合、OpenCL Cのコンパイラでエラーとなるでしょう。 しかし、よく見てください。 _ が2つ連続していますね。 つまりこれは(CUDA)C++としては予約語となるため、そんな関数を宣言した時点で未定義動作です。 未定義動作のコードを書いて問題が起きても未定義動作のコードが悪いので、Ultimaには何ら非はありません。

コンストラクタ

C++にはクラスがあり、コンストラクタがあります。 この際メンバアクセスは全部public扱いで構造体で代替するとして(CUDA C++時点でアクセス制御に問題がないコードなら全部public扱いでも問題ないため)、コンストラクタは困ります。 Cでは構造体の中で関数定義はできません。 しかし、CuPyの内部実装では1引数コンストラクタによる暗黙の型変換を前提としたコードが存在し、これに対処するためにはコンストラクタの実装が必要でした。

そこで、 struct 、 class 、 union について、コンストラクタ T(Args... args) を void constructor_T(T*const this, Args... args) へと書き換え、コンストラクタ呼び出し T t(args...); を T t;constructor_T(&t, args...); へと書き換えて対処しました。 この際、関数定義を一度退避しておき、構造体定義が終了してから関数定義を出力するようにすることで、Cコンパイラでも問題なく処理できるコードにすることができます。 なおコンストラクタが複数ある場合は、例によってマングリングされます。

ndarrayへの要素アクセス

上記の ElementwiseKernel の例の x のように、引数定義に raw と付けておくと明示的にインデックスアクセスを行えるようになります。 この際、引数には以下の2つの型の値を渡すことができます。

• 整数型
  • n次元配列を1次元に直列化した際のインデックスアクセスを行います。
  • 上記例のように、 raw でないndarrayでも内部実装に使用されています。
• 整数配列型
  • n次元のndarrayにn次元の整数配列を与えることでも指定要素にアクセスが可能です。

このインデックスアクセス、CuPyにおいては CArray というテンプレートクラスの operator[] のオーバーロードによって実現されており、templateを使って値を取るので数値型なら何でも取れたり、整数値か配列かによって処理を変えたりすることができるのですが、ClPyではクラスが使えないのでα版においては文字列検索と置換で適切な関数呼び出しを挿入してなんとかしていたのでした。 Ultimaでは、以下のように処理していきます。

1. CuPyのような CArray クラスを用意しておきます。
   • これによりUltimaがオーバーロードされた演算子の呼び出しであることを型情報つきで得られます。
2. 引数の CArray<T, N> クラスをOpenCL C用の CArray_N 構造体とバッファの先頭アドレス T* に分解します。
   • OpenCLだとホスト側でGPU上のメモリのGPUにおけるアドレスを取り出して構造体に格納する、といった前処理が不可能なためです。
3. operator[] のオペランドの型に応じて適切な get_CArrayIndex 関数の呼び出しを挿入します。
   • 特に整数配列型でのアクセスの際、 get_CArrayIndexRaw_N 関数をテンプレート関数にしておくことで上記のインスタンス化・マングリングが行われ、適切なサイズでの配列アクセスが実現できます。
4. CArray クラス自体は除去します。
   • これによりOpenCL Cコンパイラに渡す際にも問題なくコンパイルが通るようになります。

変数名 half

OpenCL Cにおいては half という識別子は半精度浮動小数点数型の型名として予約されているため、識別子として使うことができません。 しかし、Chainerのコード内では変数 half が登場します。 これに対処するために、Ultimaをかける前に以下のコードを挿入します。

#define half __clpy__half

おわりに

というわけで、ClPyがどのようにlibToolingを使用しているかについての最初の記事でした。 ClPyがChainerをサポートするために、このような ゴリ押し黒魔術 コンパイラ技術が使われています。 Ultimaについては複数のClangのバージョンに対応させるのに苦慮した話²などがありますが、これはまた別の機会に。

気づけばChainerは開発終了となってしまいましたが、ClPyはCuPy同様引き続き開発を継続しています。 ClPyの現状についてですが、RC1以降実行時間の改善に取り組んできました。 例としてMNISTの1イテレーションあたりの実行時間で言えば、CuPyに比べて37%程度遅かったものが8%程度の速度低下まで改善されました。 また、UltimaのテストやRNGの追加、CUDA 10やClang 8/9の対応なども行われました。 また、ユーザーからのサポート対応も数件ありました。 これまでと同様GitHub上で開発しておりますので、リポジトリのwatch、Issueへのフィードバックなどなどお待ちしております。

[1]: ソースコード全体の文字列の中でその演算子が構成する式の開始・終了位置のインデックスは取れるので、式の文字列を得ることはできます。また、同様に子要素となるオペランドのインデックスも取れるので、これを用いて式の文字列からオペランドの文字列を差し引けば演算子を取得すること自体は可能です。ただし、マクロの中で使われた演算子の場合インデックスアクセスしても演算子の文字列は取れないため、この手法は万能ではありません。

[2]: これはC++のライブラリを開発している人たちには有名な話ですが、Clangのバージョンを __clang_major__ で取得して #if で選り分けてはいけない、と公式ドキュメントに記載があるため、複数のClangのバージョンに対応させるためにはどのClangのバージョンでもコンパイルが通るコードを書くしかありません。ところで、LLVMやClangのAPIを叩いてる人たちには有名な話ですが、LLVMやClangのAPIはメジャーバージョンごとにメチャクチャ壊れます。あとはご想像ください

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