自作言語でGPUを少しだけ動かしてみた

この記事は 三重大学 計算研究会 Advent Calendar 2020 20日目の記事です．

背景・概要

以前にRustでGPGPUしたり自作言語でセルフホストしたりしていて，知識としては揃ったので自作言語でGPUを動かしたいとなりました．
x86_64のアセンブリ・PTXコードの生成と，GPUでの実行までできたので，必要だった知識や感想などを書きたいと思います．

リポジトリは https://gitlab.com/spica314/neco2 です．
コンパイラの実装言語はRustです．今回はC言語にしようかと思ったけど#[derive(Debug)]的なのがないのはつらすぎた．

できたものとしては，Ray Tracing in One Weekend の chapter 2 のImage1をGPUを使って計算できました．わーい．(リンク: 自作言語で書かれたソースコード)

...本当は chapter 6 ぐらいまで動かしたかったのですが，GPU向けのコード生成で，条件分岐もループも構造体を(PTX上の)レジスタに割り当てるのもまだ実装できていないので，とりあえずここまで．

自作言語の概要

Hello, worldは...と書こうと思いましたが文字列リテラルが未実装だった

構文的にはRustから関数呼び出しの括弧を消し去って，謎を入れたっぽい言語です．
なおHaskellの$的な存在をまだ実装していないので結局括弧だらけです．

今回は適当な型推論を実装したので型注釈なしでletが使えるようになりました．

全然まとまってないテスト用のプログラム集: https://gitlab.com/spica314/neco2/-/tree/0dbbc599ec0afa91db134242db6feb3b0bbc02ea/examples

exit 42

fn main () = {
    __mov __rax 60;
    __mov __rdi 42;
    __syscall ();
}

インラインアセンブリ的なのが文字列ではないです．
そのうちunsafeブロックっぽいasmブロックを生やす予定です．

exit_group 42 (GPU ver.)

#[ptx]
fn f (x: &[i32; 32]) = {
    let thread_id = __ctaid_x * __ntid_x + __tid_x;
    (*x)[thread_id] = 6 * thread_id;
}

fn main () = {
    __init_cuda ();
    let xs = [0i32; 32];
    __call_cuda f &xs 1 1 1 32 1 1;
    let x = xs[7];
    __mov __rax 231;
    __mov __rdi x;
    __syscall ();
}

GPU版?です．配列のi番目に6*iが入るようにして7番目で42です．関数fはGPU上で実行されます．

コンパイルの流れ

1. GPU向けにPTXコードを生成します (PTXコードについても後述)
2. x86_64のアセンブリを生成します
   • 生成したPTXコードは文字列としてアセンブリに含めておきます
   • GPUでの実行はCUDA Driver APIを使います．詳細は後述
3. 生成されたアセンブリを gcc を使ってアセンブルとリンクをします．
   • e.g. gcc -o a.out t.s /opt/cuda/lib64/stubs/libcuda.so

先に 2. から書きます．

2. x86_64のアセンブリの生成

この部分はいつも通り(?)です．

主に参考にしたのは，

• 低レイヤを知りたい人のためのCコンパイラ作成入門, https://www.sigbus.info/compilerbook
  • 特にスタックとスタックポインタ周りでこんがらがったとき
• LINUX SYSTEM CALL TABLE FOR X86 64, https://blog.rchapman.org/posts/Linux_System_Call_Table_for_x86_64/
  • システムコールの番号とレジスタの使い方

の2つです．

また，今回はRustの syn crate を参考にパーサ部分を書きました．
参考にしていて，かっこいいなと思ったのは，https://docs.rs/syn/1.0.56/syn/parse/index.html にある，

impl Parse for ItemStruct {
    fn parse(input: ParseStream) -> Result<Self> {
        let content;
        Ok(ItemStruct {
            struct_token: input.parse()?,
            ident: input.parse()?,
            brace_token: braced!(content in input),
            fields: content.parse_terminated(Field::parse_named)?,
        })
    }
}

のような書き方です．parse()のところで型を書かなくていいのが楽なのと，(たぶん)途中で失敗したときなどの処理を(おそらく内部でParse traitのparse()を呼ぶ)ParseStream::parse()側で処理しているんだろうなぁ(想像)になりました．
前回，?を使いたいけどもとに戻すの書けないしなぁになってごちゃごちゃしましたが，今回はこれっぽい書き方にしてプログラム見た目がだいぶ良くなりました．たぶん．

CUDA Driver APIの話は 1. の後で書きます．

1. PTXコードの生成

PTXは，NVIDIAのGPU向けのISAです．個人的には中間言語と言ったほうがしっくりきます．

例えば，CUDAのプログラム

extern "C" {
  __global__ void copy_array(int *xs, int *ys) {
    int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
    xs[thread_id] = ys[thread_id];
  }
}

をnvcc --ptx hoge.ptxでPTXコードにコンパイルすると，次のような記述になります．

.visible .entry copy_array(
	.param .u64 copy_array_param_0,
	.param .u64 copy_array_param_1
)
{
	.reg .b32 	%r<6>;
	.reg .b64 	%rd<8>;


	ld.param.u64 	%rd1, [copy_array_param_0];
	ld.param.u64 	%rd2, [copy_array_param_1];
	cvta.to.global.u64 	%rd3, %rd1;
	cvta.to.global.u64 	%rd4, %rd2;
	mov.u32 	%r1, %ctaid.x;
	mov.u32 	%r2, %ntid.x;
	mov.u32 	%r3, %tid.x;
	mad.lo.s32 	%r4, %r1, %r2, %r3;
	mul.wide.s32 	%rd5, %r4, 4;
	add.s64 	%rd6, %rd4, %rd5;
	ld.global.u32 	%r5, [%rd6];
	add.s64 	%rd7, %rd3, %rd5;
	st.global.u32 	[%rd7], %r5;
	ret;

}

一応生成できそうな見た目をしていると思います．

ちなみに，.regは最初にまとめる必要はなく，分けて書いても動きます(<6>とかは1~6をまとめて定義，別々でもいい)．また，SSA(静的単一代入)っぽい見た目をしていますが，その必要もないようです．

少し罠だと思ったのは，

__global__ void f(int *xs, int *ys) {
  // ...
}

のPTXコードは

__global__ void f(int **xs, int **ys) {
  // ...
}

っぽく見えます(実際にどうなのかは自分はよくわかっていない)．
まぁあまり気にする必要はなくて(?)，ld.param.u64とcvta.to.global.u64をテンプレ的に生成すれば動くと思います．たぶん．

生成するときに参考にしたのは，

• 小さめのCUDAプログラムを書いてnvccでコンパイルしたもの
• PARALLEL THREAD EXECUTION ISA v7.0, https://docs.nvidia.com/pdf/ptx_isa_7.0.pdf

の2つです．

デバッグの際には，生成したPTXコードをptxasコマンドに渡すと親切めのエラーメッセージが得られます．

2.(続き) CUDA Driver APIを使ったPTXコードの実行

生成したPTXコードをGPUで動かすには，CUDA Driver APIを使います．
CUDA Driver API のドキュメントは https://docs.nvidia.com/cuda/cuda-driver-api/index.html です．

使った関数名を書いておくと，

• cuInit
• cuDeviceGet
• cuCtxCreate_v2
• cuModuleLoadData
  • ここで生成したPTXコードを渡す
• cuModuleGetFunction
• cuMemAlloc_v2
• cuMemcpyHtoD_v2
• cuLaunchKernel
• cuMemcpyDtoH_v2

です．この順番で呼び出すとPTXコードの実行(とデータの転送など)ができます．

現時点では自作コンパイラがSystem VのABIに対応していないので，初期化や呼び出し部分はアセンブリを直打ちしたものを出力するようにしています．

共有ライブラリの関数を呼び出すには，関数名に@PLTをつけておくといいようです．

ちなみに，CUDA Driver APIを使うときには，60番のexitシステムコールではなく，231番のexit_groupシステムコールを使わないと，動作が固まるようです(Destroy系の関数の呼び出しをさぼっているせいかもしれないです)．

3. アセンブル・リンク

最後に生成したアセンブリをgccでアセンブルとリンクをします．

GPU周りを除けば，生成されるアセンブリはlibcなどに依存していなくて，本当はasとldで処理したかったのですが，libcuda.so が C Runtime を要求するっぽいので gcc でいいかとなりました．

(コンパイルの流れの話はここまで)

プログラムを書き直した話など

今回のプログラムは，前回セルフホストしたプログラムを捨てて，1から書き直しました．

理由としては，以前書いたコンパイラに多少まともなIR(中間表現)を入れようとしたのですが，書き直し量が大変なことになったので，1から書き直すことにしました．IRを入れるなら初めから入れるつもりで書いておいたほうがいいかもしれません．

また，前回はセルフホストを目標にしていましたが，今回はセルフホストは目標にしませんでした．
理由としては，機能を追加するときに，ブートストラップ用のコードとセルフホスト用のコードの2つがあって，セルフホスト用のほうは言語機能があんまりそろってないので，ブートストラップ用のほうに書きます．その後，追加した機能を使うようにセルフホストのほうを書き換えて，機能の追加して...と二度手間になりました．
言語機能がかなり揃ってからセルフホストしようと思います(いつになるやら...)．

今後?

• とりあえずRay Tracing in One Weekendを自作言語でGPUで最後まで動かしたい
• 標準ライブラリ的なのにはいままで手をつけれていないので，そのあたり?
• アーキ?的には高位合成とかDPU(できるんか?)とかWasmとかOpenQASMとか...

さいごに

• ありとあらゆる計算を自作言語で実装したい(野望)
• コンパイラたのしい