OpenAI、ニューラルネットワーク向けGPUプログラミング言語「Triton 1.0」公開

 AIを研究する非営利団体OpenAIは7月28日、GPUを使った深層学習プログラムを作成できるPythonライクなプログラミング言語「Triton 1.0」を公開した。

 TritonはCUDA体験のない研究者が効率の良いGPUコードを作成できることを目指して開発したプログラミング言語およびコンパイラ。既存のDSL以上の柔軟性も目指す。OSはLinux、ハードウェアはNVIDIA GPU(Compute Capability 7.0+)をサポートし、現在AMD GPU、CPUもサポートに向けて作業中という。Tritonライブラリを使うことで、研究者はPythonでコードを作成し、JITコンパイルしてGPU上で動かすことができる。Pythonはバージョン3.6以上をサポートする。

 開発チームによると、プログラミングモデルで最も近いものはNumbaで、カーネルをデコレートされたPython関数として定義し、インスタンス上で並列的に異なるprogram_idで起動する。だが、SIMT実行モデルは取らず、ブロック上のオペレーションでイントラインスタンス並列をエクスポーズするという。これらの仕組みにより、比較的少ない労力でハードウェア性能のピークに到達できるという。例えば、cuBLASの性能に匹敵するFP16行列乗算カーネルを25行以下のコードで作成できるとのこと。すでに内部の開発者はTritonを使って、同等のTorch実装と比較して最大2倍の効率でカーネルを生成できているという。

 TritonはプロジェクトのWebサイトより入手できる。

Triton
https://github.com/openai/triton