Show HN: Phobos – タイルDAGサポート付きの小さなスケールフリー・カーネル言語

TLDR: Tritonにインスパイアされた小さなカーネル言語Phobosを開発しました。PTXにコンパイルされ、NVIDIA GPU上で動作します。2080 SUPERでcuBLAS SGEMM GFLOP/sの76%（理論上のピークGFLOP/sの74%）という許容できるパフォーマンスを達成しています。Phobosは、分散タイルDAG設計に自然にマッピングされます。クラスタプロトタイプは単一マシンでのみ検証しており、マルチノードベンチマークはありません。これは、低レベルGPU概念への理解を深めるための個人的な研究プロジェクトでした。 子供の頃から（人工）ニューラルネットワークについて学び始め、人間の脳とAIに魅了され、寝室で貧弱なウェブカメラと1.8GHz CPUを使って学習用の物体認識ソフトウェアを自作するほどでした。学校を終えた後、ビーレフェルト大学で認知情報学を専攻しましたが、パリのスタートアップから仕事のオファーがあり、世界を見てみたかったため、わずか4日で中退しました。 その後、最適化、コンパイラ、その他の機械学習関連のトピックに多くの時間を費やしました。現在も研究をフォローしていますが、最先端のAI研究やテクノロジーからは一歩引いた立場にいます。 ここに、貧弱な2080 SUPERと共に戻ってきました。学ぶ意欲に満ちています。 何かを完全に理解するためには、実際に手を動かして作業することが不可欠だと強く信じています。コンパイラ、最適化、LLVMの経験があるので、GPUをターゲットとする小さな言語とコンパイラを書くことは私にとって手が届く範囲だと考えました。 ソフトウェアエンジニアリングとコンパイラの経験も十分にあるため、慎重に進める必要があることもわかっています。これを学習の演習に留めたいのです。もし間違った方向に進めば、私の利用可能な時間を過剰に消費するプロジェクトに発展する可能性があります。コンパイラは特にトリッキーです。プログラムを最初から最後までコンパイルできなければ、何も見えません。プログラムのコンパイル自体が容易ではないため、いくつかのショートカットを取ることにします。リンクやフェーズモデルのようなものはありません。 目標：NVIDIA GPUでコンパイル・実行できる小さな言語を作成すること。その際、AI計算に適したモデルを選択すること。典型的なAIカーネルの意味論を再構築することなく、最適化を可能にすること。第一級のテンソル型/タイルをサポートしたいと考えています。実際、言語全体がタイルを中心に構成されています。 コンパイルしたい2つのサンプルプログラムを以下に示します。言語については後で詳しく説明します。 ベクトル加算 これは \(\vec{c} = \vec{a} + \vec{b}\) です。 @autotune(BLOCK in [1024]) kernel add(a: tensor<f32>[N], b: tensor<f32>[N], c: tensor<f32>[N]) { let base = program_id(0) * BLOCK c[base :+ BLOCK] = a[base :+ BLOCK] + b[base :+ BLOCK] } 行列乗算 これは \(\mathbf{C}_{M \times N} = \alpha \mathbf{A}_{M \times K} \mathbf{B}_{K \times N} + \beta \mathbf{C}_{M \times N}\) です。 @autotune(TILE_M in [32, 128], TILE_N in [64, 256], TILE_K in [16, 64]) kernel sgemm(A: tensor<f32>[M, K], B: tensor<f32>[K, N], C: tensor<f32>[M, N], alpha: f32, beta: f32) { let pm = program_id(0) let pn = program_id(1) var acc: tile<f32>[TILE_M, TILE_N] = 0.0 for kt in range(0, K, TILE_K) { var a = A[pm * TILE_M :+ TILE_M, kt :+ TILE_K] var b = B[kt :+ TILE_K, pn * TILE_N :+ TILE_N] acc += dot(a, b) } let c_old = C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N] C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N] = alpha * acc + beta * c_old } 余談：カーネル開発に詳しい方なら、はい、これは基本的にミニ・トライトンです。 未知の領域 CUDAとそのターゲットフォーマットを知ることから始めましょう。CUDAリフレッシャー：CUDA入門のようなリソースを見ると、次のようなサンプルに出会います。 /** CUDA kernel device code - CUDA Sample Codes * * Computes the vector addition of A and B into C. * * The three vectors have the same number of elements as numElements. */ __global__ void vectorAdd( float *A, float *B, float *C, int numElements) { int i = blockDim.x * blockIdx.x + threadIdx.x; if (i < numElements) { C[i] = A[i] + B[i]; } } 私たちはこれには興味がありません。これは高レベルコードです。GPUのC++ではなく、GPUのx64が欲しいのです。まずNVIDIA CUDA Compiler (NVCC) のドキュメント1を見ましたが、nvccは実行ファイル全体をリンクするため、非常に複雑です。CUDAドキュメント2自体には、私たちの領域に近いセクションがあります：PTXと低レベルプログラミング。低レベルプログラミングこそが私たちが求めているものです。 PTXは汎用並列プログラミングのための安定したプログラミングモデルと命令セットを提供します。これは、NVIDIA Teslaアーキテクチャで定義された計算機能をサポートするNVIDIA GPUで効率的に動作するように設計されています。CUDAやC/C++のような言語の高レベルコンパイラは、PTX命令を生成し、これはネイティブターゲットアーキテクチャ命令に最適化され、変換されます。 利用できる既存のLLVMバックエンドがあり、NVPTXと呼ばれています。これでパイプラインのアイデアができました。 ソースコード ⟶ レクサー ⟶ パーサー ⟶ コードジェネレータ ⟶ LLVM IR ⟶ NVPTX ⟶ PTX ⟶ GPU PTXはGPUが実行するマシンコードではありませんが、それに最も近いものです。SASSが実際のマシンコード4ですが、文書化されていません。 しかし、それだけではありません。常にあります。機能の肥大化です。それを感じます。MLIRが現在のトレンドです5。LLVM IRを生成する代わりに、それを使用しましょう。実際、カスタムダイアレクトのようなMLIRの機能を使用するつもりはありません。しかし、CUTLASSのようなパフォーマンスを無料で得られるオプションを提供してくれます…夢を見させてください。真面目な話、MLIRは私が必要とする多くの命令を公開しています。 ソースコード ⟶ レクサー ⟶ パーサー ⟶ コードジェネレータ ⟶ MLIR ⟶ PTX ⟶ GPU このプロジェクトでは、LLVMのバインディング（inkwellを使用）とMLIRのバインディング（meliorクレートを使用）を持ち、パターンマッチングをサポートしているため、コンパイラの記述が100倍楽しくなるRustを使用します。MLIRの使用はinkwellの使用と同じくらい「シンプル」になることを期待していました。しかし、残念ながら、melior、mlir-sys、tblgen-rsクレートのパッチ適用やその他の苦難に多くの時間を費やし、ようやく簡単なプログラムをコンパイルできるようになりました。 この小さな冒険を始めたとき、すべてを理解するために、LLVM IRから直接PTXを生成しました。それが私にとって当然の道でした。しかし、すぐにMLIRについて知りました。inkwellを使い始めるのに数時間かかりました。meliorとMLIRを動作させるのに2日かかりました。 右から左へ ここで多くのことが壊れる可能性があるため、右から左へと進みます。最も独立したタスクから始めます：GPU上でPTXコードを実行すること。これは必然的に後でやらなければならないことです。 上記のvectorAddをPTXで記述し、custクレートで実行してみましょう。 まず、GPUのコンピュートキャップを理解する必要があります。私の場合はTuringファミリーで7.5です。はい、VRAMが8GBしかないかなり古いGPUです。残念ながら、float32テンソルコアもありませんが、16ビット精度は利用可能です。それで対応できます。 $ nvidia-smi --query-gpu=name,compute_cap name, compute_cap NVIDIA GeForce RTX 2080 SUPER, 7.5 これでPTXを書き始めることができます。LLVM IRを手書きしたことがあるなら、これは馴染み深いでしょう。少なくとも、すべてをSSA形式で書く必要はありません。いずれにせよ、これは手書きで書く予定だった唯一のPTXコードでした。ptxas add.ptxでアセンブラを直接実行することは、Rustや使用しているCUDAクレートからの問題を分離するための有用な健全性チェックでした。nvcc -ptx in.cu -o out.ptx を使用してカーネルをPTXにコンパイルすることも非常に役立ちます。カーネルを書き、公式ツールチェーンでコンパイルし、生成されたPTX出力を検査できます。 .version 9.0 // 9.3 は動作するはずだが、しない？ .target sm_75 // ターゲットアーキテクチャ（コンピュートキャップ） .address_size 64 .visible .entry add( .param .u64 a, .param .u64 b, .param .u64 c, .param .u32 numElements ) { .reg .pred %p; .reg .f32 %f<4>; .reg .b32 %r<5>; .reg .b64 %rd<11>; ld.param.u64 %rd1, [a]; ld.param.u64 %rd2, [b]; ld.param.u64 %rd3, [c]; ld.param.u32 %r1, [numElements]; // i = blockIdx.x * blockDim.x + threadIdx.x mov.u32 %r2, %ctaid.x; mov.u32 %r3, %ntid.x; mov.u32 %r4