Triton 3.9.0は、GPUカーネル開発者向けに、他の処理が割り込まないメモリ操作であるアトミック演算とコンパイル時の検査を強化した。カーネルのメモリ動作を直接制御でき、一部の不正な入力は実行前のコンパイル段階で検出される。コードを読み取って動作を模倣するデバッグ用インタープリターも、コンパイル済みカーネルとより一致するよう修正され、結果の確認に役立つ。
1. アトミック演算でメモリの読み書きを制御
今回のバージョンでは、メモリの値を分割できない1回の動作で読み書きするアトミック演算が追加された。tl.atomic_loadとtl.atomic_storeが新たに導入され、Gluonとインタープリターでも利用できる。
tl.atomic_pollは、指定したアドレスの値が期待値と一致するまで待機する。制限時間内に一致しなければ偽を返すため、値の準備ができたかを確認する待機処理に終了条件を設けられる。
読み書きが観測される順序であるメモリ順序の処理をブロック全体に広げる利点には、同期処理が追加されるコストが伴う。読み書きの順序を指定したアトミック演算には、同じ実行ブロック内の流れをそろえるため、一時停止して同期する地点であるバリアが入る。このバリアを避けるにはsem="relaxed"を指定できるため、開発者は必要な順序保証と同期の負担を併せて検討する必要がある。
2. 不正な入力を実行前に発見
複数の値をまとめたタプルを展開して受け取る項目数が一致しない場合、コンパイル時に失敗するようになった。上位の値を選ぶtl.topkの選択数が無効な場合や、配列を連結するtl.catの形状が合わない場合も、同じ段階でエラーになる。
度数分布を計算するtl.histogramに32ビットではない入力を渡した場合も、実行前にエラーになる。入力の誤りの検出をコンパイル段階に前倒しし、カーネルの実行結果から異常を探す代わりに、入力と呼び出し箇所から調べられるようにした変更だ。
コンパイル時間も改善された。リリースノートによると、ネストの深さが8の深いループのコンパイルは約10倍高速化した。これは、そのネストの深さについて示された結果だ。
3. インタープリターとコンパイル済みカーネルの結果を一致させる
デバッグ用インタープリターはコードを解釈して実行するもので、今回の修正は、その結果をコンパイル済みカーネルと一致させることに重点を置いた。整数が表現可能な範囲を超えるオーバーフローと、整数の積の上位ビット、つまり整数表現の高い桁のビットを求める演算の動作を修正した。
乗算と加算を1つにまとめる融合演算と、ヒストグラム計算も一致するよう修正した。開発者がこれらの演算をインタープリターで点検する際、実行方式の違いによる結果の不一致が減る。エラーを早期に見つけるコンパイル時の検査とともに、計算結果を確認する段階も整備した形だ。
4. 既存カーネルの型変換を点検
32ビットより幅の狭い浮動小数点値が、除算や剰余演算の入力値であるオペランドに含まれると、結果はfloat32になる。Pythonのスカラーと併用する場合も該当する。既存の型の結果を維持するには、演算結果を目的の型に変換する必要がある。
テンソルの数学メソッドにも、関数と同じ型検査が適用される。例えば、指数関数メソッドは16ビット浮動小数点の入力でコンパイルに失敗する。リリースノートでは、先に32ビット浮動小数点へ変換する次の例を示している。
x.to(tl.float32).exp()こうした変更により、型を自動的に処理していた既存コードにも修正が必要な箇所が生じる。アップグレード時は、新たに発生するコンパイルエラーの修正とともに、演算結果の保存先の型を確認する必要がある。
5. 次のバージョンの融合のデフォルト設定に備える
乗算と加算を自動的にまとめる暗黙の浮動小数点融合は、3.9では引き続き有効だ。ただし、enable_fp_fusionを指定していないカーネルをコンパイルすると、プロセスごとに1回警告する。3.10では、この融合はデフォルトで無効になる予定だ。
リリースノートでは、数値結果が融合に依存するカーネルに対し、tl.fmaを明示的に呼び出すよう案内している。現在のデフォルト設定に頼るのではなく、計算の意図をコードに残し、次のバージョンのデフォルト設定の変更に備える方法だ。
Triton 3.9.0の要点は、カーネル作成者がアトミックなメモリ動作を直接扱いながらエラーをより早く発見し、変更された順序と型の規則を次のバージョンに先立って点検できるようにした点だ。