GPU 上の 2 つの数値
あなたは Python で c = a + b を実行する方法をすでに知っています。何千回もやってきたでしょう。ここでは、それを PyTorch のテンソルで説明します。テンソルとは、数値の配列にすぎません。それを GPU に置くということは、その配列が通常の CPU メモリではなく GPU のメモリに存在することを意味します。a と b が GPU 上の 2 つのテンソルである場合、その 1 行は十分に高速に終了するため、考える必要すらありません。
では、それを縮小してみましょう。a と b が 2 つの単一の浮動小数点数で、両方とも GPU 上にあるとします。同じ行です。c = a + b。実際にチップ上で実行されるのは何でしょうか？
答えは カーネル です。この世界では、カーネルとは GPU が一部のデータに対して実行する小さなプログラムです。ラップトップが起動する OS カーネルではありません。線形代数の教科書にある数学カーネルでもありません。この言葉は何度も再利用されているので、あなたのせいではありません。GPU の世界では、カーネルは単に次のことを意味します。GPU が今すぐ、受け取ったデータに対して並列に実行するように指示された小さな関数です。
この記事の終わりまでに、あなたは PyTorch のコードの断片を見て、GPU が実行するカーネルの数を数えられるようになります。それは小さな技のように聞こえますが、実際その通りです。しかし、それは「GPU をブラックボックスとして扱う」感覚を振り払うための最初のステップでもあります。あなたのモデルが遅くて理由がわからない、あの感覚です。GPU のパフォーマンスについて尋ねることができるすべての質問は、最終的に「どのカーネルが実行され、何をしていたのか」に戻ります。だから、ここから始めましょう。
最初のカーネル
a と b を少し大きくしてみましょう。今回は長さ 8 のテンソルです。依然として 1 行の PyTorch コード: c = a + b。
これを実行すると、あなたの CPU（実際に Python を実行しているマシン）は GPU に次のように指示します: 「このデータに対してこのプログラムを実行してくれ。」その命令は 起動 (launch) と呼ばれます。起動されるのはカーネル、つまり 1 つのプログラムで、実行する準備ができています。起動自体は安価で、1 つあたりマイクロ秒単位です。起動の 周り（GPU に送られるデータ、戻ってくる結果）に実際のコストがかかります。そして、それが私たちがこれから数えていくものです。
カーネル内部では、実際の作業は スレッド と呼ばれる小さなワーカーによって行われます。GPU は数千ものスレッドを利用できます。長さ 8 の加算の場合、8 つのスレッドが作業を受け持ちます: スレッド 0 は要素 0 を処理し、スレッド 1 は要素 1 を処理し、以下同様にスレッド 7 まで続きます。各スレッドは同じ小さなプログラムを実行します: a の 1 要素を読み取り、対応する b の要素を読み取り、それらを加算し、結果を c に書き込みます。
（実際には、GPU はスレッドをワープと呼ばれる固定サイズのグループで起動し、NVIDIA カードでは常に 32 スレッドであり、配列が均等に分割されない場合には余剰をマスクします。今のところは無視して構いません。）
つまり、1 行の PyTorch コード、1 回の起動、1 つのカーネル、8 つのスレッドが 8 回の加算を行います。では、実際にチップを横断したものを数えてみましょう。加算を行うために、各スレッドは a の要素と b の要素を必要とします。つまり、a からの 8 回の読み取りと b からの 8 回の読み取りです。次に、各スレッドは結果を c に書き込みます。つまり、8 回の書き込みです。
これらの読み取りと書き込みは、GPU チップのすぐ隣にある大きなメモリに行われます。データセンター向けカード（A100、H100）では、そのメモリは HBM（高帯域幅メモリ）と呼ばれます。コンシューマー向けカード（RTX 4090）や、実際に試してみる人が多い Colab スタイルの T4 では、GDDR と呼ばれます。いずれにせよ、それはチップの隣にある高速メモリであり、単に GPU メモリと呼ぶことにします。高速ですが、無料ではありません。そして、そこにアクセスするたびに何らかのコストがかかります。
1 つのカーネル = 1 回の起動 = データに対する 1 回のパス。カーネルがその内部で何をするにせよ、その端での読み取りと書き込み（入力を受け取るための GPU メモリへのアクセス、出力を書き込むための戻りのアクセス）がコストのかかる部分です。これが全体像です。
テンソルが大きくなっても、これは変わりません。同じ PyTorch の行、同じ 1 つのカーネル、単により多くのスレッドが使われるだけです。a と b がそれぞれ 100 万要素の場合、GPU はより大きなスレッド群で同じカーネルを起動します。計算はスケールし、移動するバイト数もスケールしますが、メンタルモデルは変わりません。1 行、1 カーネル。
2 つの操作の間で何が起こるか
1c = (a + b).relu()
あなたは Python を十分に理解しているので、これが 2 つの操作（加算とそれに続く relu）が連鎖していることを知っています。インタプリタでは、これは 2 つの関数呼び出しです。GPU 上では、イージーモードの PyTorch では、これは 2 回のカーネル起動になります: 1 回は加算用、もう 1 回は relu 用です。今のところ、驚くことではありません。
実際に興味深いのは、2 つのカーネルの間で何が起こるかです。
加算が終了すると、その結果を どこかに 置く必要があります。そのどこかは GPU メモリです。加算は完全な中間配列（tmp と呼びます）をメモリに書き込みます。そして、少し後に relu が起動し、その最初の仕事は、同じ tmp 配列をメモリから読み戻すことです。全体を読み取り、各要素に relu を適用し、結果を c に書き込みます。
これら 2 つのカーネルのメモリトラフィックを数えてみましょう:
- 加算:
aの読み取り、bの読み取り、tmpの書き込み。3 回の配列サイズの転送。 - relu:
tmpの読み取り、cの書き込み。さらに 2 回。
合計で 5 回の配列サイズの転送 です。前のセクションの長さ 8 の加算だけの場合（3 回）と比較してください。.relu() をチェーンに追加したことで、relu の計算コストだけでなく、配列の GPU メモリへの完全なラウンドトリップのコストも発生しました。なぜなら、tmp は次のカーネルが読み戻せるように書き出されなければならなかったからです。
何もキャッシュされませんでした。tmp はレジスタや高速なローカルキャッシュに留まることができませんでした。それは GPU メモリ（低速で遠い種類）に送られ、すぐに戻ってきました。2 つのカーネルは互いに無関係です。両方が通信方法を知っている唯一の媒体、つまり GPU メモリを介してデータを受け渡さなければなりません。
なぜ PyTorch はこのようにするのでしょうか？イージーモードでは、a + b と書くと、PyTorch はそれを 今すぐ 実行します。次に .relu() を呼び出すことは知りません。各操作は、その Python の行が実行された瞬間にディスパッチされます。計画も先読みもありません。各操作は独立して、実際の配列を生成し、メモリを介して次に来るものにそれを渡します。
これが覚えておくべきパターンです。 PyTorch コード内のすべての中間値は、物理的に GPU メモリに書き出され、次の操作によって読み戻されます。すべてです。それが「カーネル数」が実際に測定するものです。追加のカーネルはすべて、データが GPU メモリを経由しなければならない別のラウンドトリップです。
融合: 2 つの操作、1 つのカーネル
すべてを一度に行う 1 つのカーネルを想像してください: a の要素を読み取り、b の要素を読み取り、それらを加算し、結果に relu を適用します（すべてカーネル内部で、チップから決して出ることのない小さなスレッドごとのスクラッチスペース上で）、そしてその後に最終値を c に書き込みます。中間の (a + b) は依然として存在しますが、カーネル内部のみ、各スレッドのプライベートなスクラッチスペース内にあります。GPU メモリに書き出されることは決してありません。tmp は実際の配列としてはまったく存在しません。
では、転送を数えてみましょう。a の読み取り: 要素あたり 1 回。b の読み取り: 要素あたり 1 回。c の書き込み: 要素あたり 1 回。3 回の配列サイズの転送。 2 つのカーネル版と同じ計算ですが、ラウンドトリップが 2 回少なくなっています。
長さ 8 では、これは四捨五入誤差です。誰も気にしません。長さ 100 万、または 1 億では、これらの余分なラウンドトリップが実行時間の大部分を占めるようになり、ウォールクロックがそれを反映します。なぜメモリトラフィックがそのように支配的になるのか は、このシリーズの第 2 記事の主題全体なので、ここでは「なぜ」については触れません。今のポイントは単に: 同じ計算、より少ないトリップ、実際に高速化、ということです。
別々のカーネルになるはずだった操作を 1 つのカーネルに結合し、中間値が GPU メモリにアクセスする必要がないようにするこのトリックには、名前があります。それは 融合 (fusion) と呼ばれます。それがすべての言葉です。それがすべてのアイデアです。
さて、厄介な部分です。その結合されたカーネルを手動で書くことは、add + relu の場合には簡単に見えます。2 つの操作。中間に 1 行の「計算」。しかし、実際の PyTorch コードには、数十の操作が連鎖しており、それぞれに独自の形状、データ型、ブロードキャストルールがあります。それらすべてを正しく処理する融合カーネルを書くことは、本格的なエンジニアリング作業です。通常、このような日常的な要素ごとのカーネルを手動で書くことはないでしょう。
良いニュース: PyTorch には、この種のケースに対して自動的に書き換えを行うツールが組み込まれています。それは torch.compile と呼ばれます。
あなたはおそらく、torch.compile が次のように使われているのを見たことがあるでしょう:
1model = torch.compile(model)
1 行です。インターネット上の誰かが、これで処理が速くなると言ったでしょう。これが実際に行うことを平易な英語で説明します: イージーモードのように操作を 1 つずつ実行する代わりに、torch.compile は関数が実行するテンソル操作をキャプチャし、それらを結合する機会を探し、最適化されたコードを生成します。後で同じ前提に一致する呼び出しは、そのコードを再利用できます。
上記で紙面上で行った融合（add と relu を 1 つのカーネルで共有し、tmp がメモリに触れない）は、まさに torch.compile が、操作が十分に単純であれば、自動的にコードに行う種類のものです。人々が torch.compile は「PyTorch を高速化する」と言うとき、これがその大部分を占めています。
torch.compile が独自に融合できない場合（認識できないカスタム操作、異常なリダクション、変なメモリレイアウト）には、誰かが手動でカーネルを書く必要があります。それが Triton や CUDA のようなツールの出番です。別の記事で扱います。
自分で確認する
上記のすべては、紙の上でカーネルを数えるものでした。実際の GPU で数えてみましょう。CUDA GPU を搭載したマシン（ワークステーション、Colab ノートブック、クラウドインスタンス）があれば、これを数分で自分で実行できます。
ツールは torch.profiler です。PyTorch に組み込まれています。コードの実行中に GPU が実際に行ったことを記録し、読み取り可能なテーブルを返すだけです。
ステップ 1: イージーバージョン
2 つの操作の行を関数でラップして、呼び出せるようにします:
1import torch2from torch.profiler import profile, ProfilerActivity34def add_relu(a, b):5 return (a + b).relu()67a = torch.randn(1_000_000, device="cuda")8b = torch.randn(1_000_000, device="cuda")910with profile(activities=[ProfilerActivity.CUDA]) as prof:11 add_relu(a, b)12 torch.cuda.synchronize()1314print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize() は、タイミングを読み取る前に GPU が終了していることを確認するためのものです。GPU の処理は非同期で実行されるため、同期がないと、実際のカーネル作業ではなく、起動のオーバーヘッドを測定することがあります。
ステップ 2: 出力を読む
実際のプロファイラ出力には、予想よりも多くの行が含まれているでしょう。メモリ割り当てや PyTorch の簿記行が多数混ざっています。注目すべき行は、CUDA カーネル、つまり GPU が実行した実際の関数です。名前に kernel を含む行を探してください。2 つの行はおおよそ次のようになります:
1vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us2vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
正確なテンプレート名は PyTorch のバージョンによって変わります（relu はしばしば threshold として表示されます。これは基礎となる操作が threshold であるためです。add は CUDAFunctor_add として表示されることもあります）。全体を解析しようとしないでください。単に行数を数えてください。2 行。2 つのカーネル。 1 つは加算用、もう 1 つは relu 用です。まさに前のセクションで述べた通りです。
ステップ 3: コンパイルバージョン
1 行の変更。関数を torch.compile でラップします:
1compiled = torch.compile(add_relu)
さて、実行する前に知っておくべき落とし穴があります。torch.compile された関数への 最初の 呼び出しは遅く、時には異常に遅くなります。なぜなら、そのときに torch.compile が作業を行うからです: コードを分析し、何を融合するかを決定し、融合カーネルを生成します。最初の呼び出しをプロファイリングすると、カーネルではなくコンパイルステップを測定することになります。したがって、パターンは次のようになります: 1 回呼び出してウォームアップし、結果を破棄し、その後 プロファイリングします。
1compiled(a, b) # ウォームアップ、結果は破棄2torch.cuda.synchronize()34with profile(activities=[ProfilerActivity.CUDA]) as prof:5 compiled(a, b)6 torch.cuda.synchronize()78print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
ステップ 4: 再び出力を読む
今度は:
1triton_poi_fused_add_relu_0 1 14us
1 行。1 つのカーネル。 名前から何をしたかもわかります: fused add and relu。以前と同じ計算で、2 回の起動ではなく 1 回です。
あなたはこの記事で話してきたことを、たった 1 文で実行しました: PyTorch に 2 つの操作を 1 つのカーネルに結合するよう依頼し、プロファイラを確認し、カーネル数が 2 から 1 に減ったことを確認しました。融合を、実際の環境で、あなたのマシンで行ったのです。
もっと劇的に見たい場合は、いくつかの異なるテンサーサイズで試してみてください。長さ 100 では、両方のバージョンが非常に高速に実行されるため、差はノイズに埋もれます。長さ 1000 万または 1 億では、コンパイルバージョンが明らかにリードし始めます。なぜなら、削減したラウンドトリップがその規模では実際の作業の一部だからです。
カーネルを数えることは、もはや抽象的なアドバイスではありません。確認する方法があります。
まとめ
すべてを一気にまとめます。
GPU 上で実行される PyTorch コードは、一連のカーネルに変換されます。各カーネルは 1 回の起動、データに対する 1 回のパス、入力の取得と出力の書き込みのための GPU メモリへの 1 回のラウンドトリップです。単純な操作は 1 つのカーネルになります。操作の連鎖は、デフォルトでは操作ごとに 1 つのカーネルになり、その間で中間値がメモリを往復します。torch.compile は、単純な連鎖を融合して、それらの中間値がメモリに触れないようにすることができます。カーネルが少ないということは、通常、メモリトラフィックが少ないことを意味します。そして、メモリトラフィックが少ないということは、通常、高速化を意味します。
