[GPU] CUDA Rustを試してみる

[GPU] CUDA Rustを試してみる

[AI要約] NVIDIAが2026年9月に発表したCUDA Rustの解説。cuda-oxideとcutile-rsを使ってみる。
2026.09.16

Introduction

2026年9月8日、NVIDIA Technical Blog に Introducing CUDA Rust: Two Tracks for Writing GPU Kernels が公開されました。

これまでも Rust から CUDA を呼ぶことはできましたが、
GPU カーネルそのものは CUDA C++ や PTX (後述) で別に用意する構成が一般的で、
NVIDIA も
「Rust からカーネルを起動できても、カーネル本体は別言語で書かなければならないことが多かった」
と記事で言ってます。

CUDA Rust は、そこを埋めるものです。GPU カーネルを Rust で書き、
SIMT トラックでは PTX、Tile トラックでは CUDA Tile IR を経由して
GPUコードへコンパイルできます。

「Two Tracks」というのは、
CUDA 自体が持つ2つのプログラミングモデルに対応した2つのプロジェクトを指します。

  • cuda-oxide — 従来型の SIMT モデル。「1スレッドが何をするか」を書く
  • cutile-rs — 新しい Tile モデル。「1つのタイル(データのかたまり)に何をするか」を書く

SIMTとTileのモデルについては後述するAbout GPU kernel / CUDA で解説します。
本記事では、GPU カーネル・CUDA の概要からはじめて、
CUDA Rust の2トラックを両方解説し、実際に動かすところまでやってみます。

本記事で確認した結論は以下。

  • Tile(cutile-rs)は高水準でシンプル。stable Rust で入り、unsafe もスレッド管理も要らない
  • SIMT(cuda-oxide)は低水準で自由度が高いが、固定 nightly が必要で shared memory には unsafe が要る
  • 素直に書くと Tile のほうが速かった。プロファイラで原因を調べ、SIMT のスレッドへの要素割り当てを見直すと、書き出しの unsafe を増やさずに上回れた
  • NVIDIA は「まず Tile を選び、制御が必要なときに SIMT を」と言っており、実際の感触もそう
  • 両プロジェクトとも early-stage で production-ready ではない。現状は実験・学習・PoC 向け
  • GPU が無い環境でも、ビルド(cuda-oxide なら PTX の確認)まではできる。実行と測定だけクラウド GPU に出す方法で検証した

本文は流れを追うのに必要な内容を書いてます。
※セットアップでつまずいた点や検証ログは GitHub に置きました(記事末尾にリンク)

About GPU kernel / CUDA

そもそも GPU とは

まずは GPU というハードウェアの性質から見ていきます。
これがわかっていると、後で出てくる
「なぜ32スレッド単位なのか」「なぜメモリの読み方で速度がそこまで変わるのか」
についても分かります。

グラボから汎用計算へ

GPU と聞いて私が思い浮かべるのはグラボです。出自もそのとおりで、
GPU(Graphics Processing Unit)は画面描画の専用ハードウェアとして生まれました。

今では画像処理に限らず、行列計算や LLM の推論にも GPU が広く使われています。
その理由は、3D 描画というタスクの性質にあります。

たとえば画面に200万ピクセルあれば、各ピクセルの色はそれぞれ独立に計算できます。
つまり描画とは「同じ計算を大量のデータに独立して適用する」処理で、
この形はほかの多くの計算にも当てはまります。
描画専用の並列計算機を汎用の計算に使うという発想が
GPGPU(General-Purpose computing on GPU)で、
それをプログラムできるようにしたのが2007年に登場した CUDA です。

現在の NVIDIA 製品は、GeForce RTX のようなコンシューマ向け(ゲームやクリエイティブ用途、映像出力あり)と、
L4 や A100、H100 のようなデータセンター向け(サーバ搭載前提)に分かれます。
今回使う L4 は後者で、AWS では g6 インスタンスに載っています。
NVIDIA は L4 を動画・AI・グラフィックス向けの汎用アクセラレータと位置づけており、
特に AI 推論を大規模に回す用途に最適化されているとしています
最大消費電力は 72 W で、ハイエンドのゲーム向けカードが数百 W を使うのと比べると、
電力あたりの効率を重視した設計です。

GPU と CPU

GPU と CPU は設計の目標が違います。

CPU は少数(数個〜数十個)の高性能なコアを持ち、1つの処理を速く終わらせることに最適化されています。
分岐予測やアウトオブオーダー実行など、逐次処理を速くする仕組みに多くのトランジスタを使っています。

GPU はCPUのコアより単純な演算器を多数持ちます。
1つ1つの処理は速くないですが、同じ処理を大量のデータへ一斉に適用することに向いています。

CPU: [とても速いコア] × 少数   → 1つの仕事を速く終わらせる
GPU: [単純な演算器]  × 多数    → 同じ仕事を大量に同時進行させる

演算器1個とスレッド1個が対応するわけではない点に注意してください。
NVIDIA GPU では、SM(Streaming Multiprocessor)がスレッドを32本ずつの束(warp)にまとめ、
warp 単位でスケジューリングして実行します(CUDA Programming Guide 3.2.2)。
各スレッドは自分のレジスタとプログラムカウンタを持ちますが、
いつ・どの実行ユニットで動かすかはハードウェアが決めます。

そして GPU では、数十万〜数百万という多数のスレッドを起動するのが普通です。
ただし、これらが一斉に動くわけではありません。
SM に載せられる warp の数には限りがあり、実行可能な warp を順に処理していきます。

メモリからデータが届くまでには時間がかかり、その間 warp は先に進めません。
この待ち時間を埋める仕組みがレイテンシ隠蔽(latency hiding)です。
warp がメモリからデータが届くのを待っている間、SM はその warp を待たせたまま、
すでに実行可能な別の warp を選んで命令を発行します。

warp A: メモリ要求 → [待ち] ──────────→ 計算再開
warp B:              計算 → メモリ要求 → [待ち]
warp C:                      計算 → ...
        ↑ Aが待っている間にB、Bが待っている間にCを走らせる

CPU にも待ち時間を減らす工夫はありますし、GPU にもキャッシュはあります。
違いは重点で、CPU は1つ1つの処理を待たせないこと、
GPU は待っている warp があれば別の warp に切り替えて演算器を遊ばせないことに力を入れています。

そのためには、各 SM に切り替え先となる warp が十分に載っている必要があります。
さらにブロックを十分な数だけ用意しておけば、先に終わったブロックのあとにも
次の仕事を SM に渡し続けられます。
GPU で膨大な数のスレッドを起動するのはこのためです。

本記事で後ほど自作する softmax カーネルも、
4096ブロック × 256スレッド = 約100万スレッドを起動します。

以降で出てくる warp やメモリ階層の話は、すべてこの設計から来ています。

GPU カーネルとは何か

GPU カーネル(kernel) とは、CPU 側のプログラムから GPU に投入され、
GPU 上の非常に多数のスレッドによって同時に実行される関数のことです。

CUDA の用語では、CPU と CPU 側のコードを host、
GPU と GPU 側のコードを device と呼びます。
一般的な外付け GPU では両者は別のメモリ空間を持つので、
典型的な処理は次の流れになります。
※Unified Memory を使う場合や統合 GPU の場合はこの限りではない

1. host でデータを用意する
2. device メモリを確保し、host → device へコピーする
3. カーネルを起動する(grid / block の形を指定する)
4. device → host へ結果をコピーバックする
5. 同期して結果を検証する

例えば、CUDA C++ でベクトル加算を書くと以下のようになります。

// device コード:GPU 上で動くカーネル
__global__ void vecAdd(const float* a, const float* b, float* c, int n) {
    // 自分が担当する要素のインデックスを計算する
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {  // 範囲外アクセスを防ぐ
        c[idx] = a[idx] + b[idx];
    }
}

// host 側は device メモリを確保してコピーし、triple chevron で起動する
// vecAdd<<<blocks, threadsPerBlock>>>(d_a, d_b, d_c, n);

idx = blockIdx.x * blockDim.x + threadIdx.x の箇所は、
各スレッドが「自分は配列の何番目を担当するのか」を
ブロック ID とスレッド ID から自分で計算します。
1024要素なら 256スレッド × 4ブロック = 1024スレッドが起動し、
それぞれが1要素ずつ並列に加算します。

この「1スレッドが何をするかを書いて、それを何千個も起動する」
という書き方がSIMT(Single Instruction, Multiple Threads) モデルです。

スレッド階層

CUDA のスレッドには階層があります(CUDA Programming Guide)。

  • Thread:最小の処理単位。1スレッドがカーネルの1インスタンスを実行する
  • Warp:32スレッドの集まり。ハードウェアが自動的にまとめ、並列で実行する
  • Block:協調して動くスレッドの集まり。同じ SM 上で実行され、shared memory を共有できる
  • Grid:1回のカーネル起動で生成される全ブロックの集合

ハードウェア側では、GPU は多数の SM(Streaming Multiprocessor) から構成され、
スケジューラが各ブロックを SM に割り当てます。
プログラマは grid と block の次元だけを指定し、
どの SM がどのブロックをどの順で実行するかはハードウェアが決めます。
この分離があるおかげで、同じプログラミングモデルのまま
さまざまなGPUに対応しやすくなっています。

warp が32スレッド単位というのは性能に直結します。
warp 内でスレッドごとに異なる分岐を取ると、
分岐の両側を順番に処理することになり(その間、該当しないスレッドが停止)、
実行効率が低下します。(warp divergence)
ブロックのスレッド数を32の倍数にするのが基本、というのはこれが理由らしいです。

メモリ階層

GPU では、CPU 以上にメモリ階層が性能に影響します。

メモリ スコープ 特徴
Register 1スレッド SM 上。最速。スレッドローカル変数が置かれる
Shared memory 1ブロック SM 上の明示管理スクラッチパッド。ブロック内のスレッドで共有できる
Global memory GPU 全体 GPU の DRAM。大容量だがレイテンシが大きい。host との転送もここを経由
Constant memory grid 全体から read read-only、専用キャッシュがある

このうち Global memory は大容量ですがレイテンシが大きく、カーネルの入力も出力もほぼここを通ります。
そのため、Global memory をどう読み書きするかが性能を大きく左右します。

特に重要なのが、warp 内のアクセスをまとめる coalesced memory access です。
1つの warp の連続したスレッドが連続したアドレスを読む場合、
ハードウェアはそれを少数のメモリトランザクションにまとめられ効率的です。
逆にスレッドごとに離れたアドレスを読むと、大量のトランザクションが必要になり帯域を消費します。

NVIDIA の CUDA C++ Best Practices Guide も、
Global memory へのアクセスを coalescing させることを "High Priority" の項目に挙げています。

thread i: C[i] = A[i] + B[i]   → 隣接スレッドが隣接要素を読む。coalescing しやすい
thread i: C[i] = A[i * 1024]   → 大きな stride。帯域を浪費しやすい

GPU の最適化では、演算そのものを減らすより、
データをどの順番で読むかが重要なこともあるようです。

2つのプログラミングモデル:SIMT と Tile

ここまで見てきたのは SIMT モデルの書き方ですが、CUDA にはもう1つの考え方があります。
冒頭の「Two Tracks」はこの2つに対応しています。
どちらも Rust 固有の話ではなく、CUDA C++ や Python でも使えるモデルです。

SIMT モデルは、これまで見てきたとおり「1スレッドが何をするか」を書きます。

// スレッド i がやること
idx = blockIdx.x * blockDim.x + threadIdx.x;
C[idx] = A[idx] + B[idx];

これを1024スレッド起動して、1024要素を処理します。
プログラマはスレッド番号を計算し、範囲外を自分で弾き、
shared memory を使うなら確保と同期も自分で書きます。

Tile モデルは、「1つのタイル(データのかたまり)に何をするか」を書きます。

// 128要素のタイル1枚がやること
z_tile = x_tile + y_tile

これを8タイル分(1024 ÷ 128)用意すれば同じ計算になります。
タイルを実際の何本のスレッドで処理するか、shared memory をどう使うかは
Tile IR コンパイラが決めます。

SIMTとTileの比較は以下。

  • SIMT
    • 1スレッドが何をするかを書く
    • スレッド番号の計算、範囲外アクセスの判定、shared memory の確保と同期は自分で書く
    • 細かい制御や特殊なアクセスパターンが得意
    • CUDA C++ や numba-cuda と同じモデル
  • Tile
    • 1タイルに何をするかを書く
    • スレッド番号は出てこない。shared memory はコンパイラが管理し、範囲外はタイルの境界として扱われる
    • 定型的な行列・ベクトル演算が得意
    • CUDA Tile(C++ / Python / Rust)と同じモデル

感覚としては、SIMT が「for ループの中身を1要素分だけ書く」のに対し、
Tile は「NumPy や PyTorch のように配列同士の演算を書く」のに近いです。
このあと書く softmax の Tile 版は reduce_maxbroadcast といった
NumPyライクな書き方で完成します。

これは、どちらが優れているという話ではなく、抽象度と制御のトレードオフです。
Tile は書く量が減り間違いにくくなりますが、
スレッドや shared memory を直接触る自由は手放します。
NVIDIA はこの選択について
「まず Tile を選び、制御が必要になったら SIMT を選ぶ」
と言ってます(後述)。

CUDA と PTX

CUDAは、NVIDIA GPU で汎用計算を行うためのプラットフォームです。
中身は主に以下の4つです。

  • 言語拡張: C++ に __global__<<<>>> を足した CUDA C++
  • コンパイラ: nvcc
  • ランタイム / ドライバ API: cudaMalloccuLaunchKernel など
  • ライブラリ群: cuBLAS(線形代数)、cuFFT、cuDNN(深層学習)など

そして本記事で何度も出てくるのが PTX です。
CUDA Rust の主張の中心にあるので、ここで押さえておきます。

PTX(Parallel Thread Execution)は、NVIDIA GPU 向けの中間表現です。
GPU のコードは2段階でコンパイルされます。

ソースコード(Rust / C++)
      ↓  コンパイラ(cuda-oxide、nvcc など)
    PTX          ← アーキテクチャ非依存の中間表現。人が読めるテキスト
      ↓  CUDA ドライバが実行時に JIT コンパイル
    SASS         ← 実際に GPU が実行する機械語。世代ごとに異なる

PTXはJava のバイトコードみたいなものです。(特定ハードウェア用の機械語ではない)
実機の機械語は SASS と呼ばれる別のもので、
Ampere 用と Blackwell 用では中身が変わります。

この中間コードを挟む方式には利点があります。
コンパイル済みバイナリ(cubin)は対象アーキテクチャへの依存が強く、
メジャー世代をまたぐと読み込めませんが、PTX は
対象 compute capability 以上の GPU へドライバが
JIT コンパイルできるため、後の世代の GPU への互換性を確保しやすくなります。
※そのかわり初回ロード時に変換の時間がかかる

なお PTX も無条件に将来の GPU で動くわけではなく、
特定世代専用の機能を使った PTX には例外があります。
cutile-rs のドキュメントも、この理由から
「JIT のオーバーヘッドを消したい場合を除いて PTX を選ぶべし」
と案内しています。

PTX が重要なのは、CUDA Rust の特徴につながります。
NVIDIA 公式ブログはこう書いています。

compiled natively to PTX, rather than a wrapper around code from somewhere else

「別言語のコードを包んでいるのではなく、Rust そのものがネイティブに PTX へコンパイルされる」
(厳密には、PTX を出すのは SIMT トラックの cuda-oxide です。Tile トラックの cutile-rs は CUDA Tile IR を経由して GPU コードを生成します)
これは、PTX を実際に出力させれば目で確認できます。
しかも PTX を出すところまでは CPU だけで完結するので、GPU が無くても検証できます。
これは Try 1-1 で実際に確かめます。

GPUまとめ

冒頭で見たとおり、GPU の性能を引き出しやすいのは、次のような条件がある場合です。

  • 大量の独立した処理があり、数千以上のスレッドに分割できる(レイテンシ隠蔽がうまく働く)
  • メモリアクセスが規則的で coalescing できる
  • GPU に載せたデータを複数のカーネルで再利用できる

深層学習の学習・推論は、この形に当てはまります。
行列積、要素ごとの演算、attention、正規化のいずれも、
大量のデータに同じ処理を適用する部分が多く、GPU 向きです。

多くの開発者は cuBLAS や cuDNN のような既製ライブラリを呼ぶだけで済みます。
それでも自分でカーネルを書く必要が出てくるのは、
ライブラリにない独自演算を実装したいときや
複数の演算を1カーネルに融合(kernel fusion)して
メモリ往復を減らしたいときなどです。
LLM 推論エンジンがよく独自カーネルを持っているのはこのためです。

そのような推論エンジンが Rust で書かれる例が増えてきました。
NVIDIA 自身も、Nova Linux ドライバや NVIDIA Dynamo で Rust を使っています。
ホスト側が Rust なのにカーネルだけ別言語という状態を解消しよう、
というのが CUDA Rust の目的です。

About CUDA Rust

Two Tracks

CUDA Rust の GitHub リポジトリでは、
さきほど見た SIMT と Tile という2つのプログラミングモデルが、
そのまま2つのプロジェクトに対応しています。

NVIDIA 公式ブログは、これを「CUDA 自体が持つ2つのトラックに合わせたもの」と説明しています。

There are two tracks to use Rust, matching the two tracks CUDA itself has.

そして、どちらを選ぶかについてもはっきり書いています。

When you are picking one to build on, reach for Tile first. The compiler decides how tiles map onto each architecture, so your source doesn't encode architecture-specific choices, and you drop to SIMT when you need that control or want to manage memory and threads yourself.

さきほども言ったように、
「まず Tile を選択する。コンパイラがタイルとアーキテクチャの対応を決めるので
ソースにアーキテクチャ固有の選択が埋め込まれずに済む。
制御が必要になったら SIMT を使う」
という指針です。
この推奨が実際どうなのかは、Try 2 の自作カーネルで確かめます。

cuda-oxide(SIMT トラック)

cuda-oxide はカスタムの rustc codegen backend です。
rustcは、ソースを型チェックや借用チェックにかけたあと MIR という中間表現に変換し、
最後に codegen backend が MIR を機械語に変換します。
普段は LLVM を使う標準のバックエンドが動いていますが、
rustc にはこの最終段を差し替える仕組みがあり、
cuda-oxide はそこを GPU 向けに作り替えています。

コンパイルに割り込んで #[kernel] を付けた関数だけを GPU 向けの処理に回し、
それ以外は標準のバックエンドに渡します。

Rust source → Rust MIR → Pliron IR → LLVM IR → PTX

中間の Pliron はコミュニティ製の IR フレームワークで、
その上に載る GPU dialect は NVIDIA 製です。
dialect も変換パスも、標準の LLVM バックエンドに渡るまでは
Rustで書かれています。

バックエンドの差し替えは rustc の内部 API に依存していて、その API は安定版として公開されていません。
そのため cuda-oxide には、特定の nightly への固定と、
rustc の内部ライブラリである rustc-dev が必要です。
初回のビルドに時間がかかるのも、このバックエンド自体をビルドするためです。
一方の cutile-rs はバックエンドを差し替えず、マクロでカーネルを埋め込んで実行時にコンパイルする方式なので、stable Rust で動きます。

実際のコードです。
↓はNVIDIA 公式ブログに載っている vecadd で、
cargo oxide new が生成するテンプレートです。

use cuda_device::{kernel, launch_bounds, launch_contract, thread, DisjointSlice};
use cuda_host::cuda_module;
use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig1D};

// === DEVICE CODE - everything in here is compiled to PTX ===
#[cuda_module]
mod kernels {
    use super::*;

    #[kernel] // GPU entry point
    #[launch_bounds(256)] // max threads per block; lets the compiler budget registers
    #[launch_contract(domain = 1, block = (256, 1, 1))] // indexes in 1-D, 256-thread blocks
    pub fn vecadd(a: &[f32], b: &[f32], mut c: DisjointSlice<f32>) {
        let idx = thread::index_1d();
        let idx_raw = idx.get(); // the plain usize, for reading the inputs
        if let Some(c_elem) = c.get_mut(idx) {
            *c_elem = a[idx_raw] + b[idx_raw];
        }
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // === HOST SETUP - device, stream, and buffers ===
    let ctx = CudaContext::new(0)?;
    let stream = ctx.default_stream();

    const N: usize = 1024;   // a_host / b_host を用意して転送する
    let a_dev = DeviceBuffer::from_host(&stream, &a_host)?;
    let b_dev = DeviceBuffer::from_host(&stream, &b_host)?;
    let mut c_dev = DeviceBuffer::<f32>::zeroed(&stream, N)?;

    // === LOAD, PREPARE, LAUNCH ===
    let module = unsafe { kernels::load(&ctx)? };
    let prepared = module.prepare_vecadd(LaunchConfig1D::new((N as u32).div_ceil(256), 256, 0))?;
    module.vecadd(&stream, &prepared, &a_dev, &b_dev, &mut c_dev)?;

    // === READ BACK AND VERIFY ===
    let c_host = c_dev.to_host_vec(&stream)?;   // ... CPU の結果と突き合わせて PASSED を表示 ...
    Ok(())
}

CUDA C++ と比べたときの違いは3つ。

  1. host コードと device コードが1ファイルに存在
    CUDA C++ でも .cu に両方書いて nvcc example.cu -o example
    の1コマンドでビルドできるので、ここは CUDA Rust だけの特徴ではありません。
    違いは、それが Rust のソースと Cargo のワークフローにそのまま乗ることです。

  2. 出力バッファの型が &mut [f32] ではなく DisjointSlice<f32>
    ここが cuda-oxide の安全性の中心なので、小見出しを分けて詳しく見ます。

  3. 起動時に設定が検証されます。
    #[launch_contract] が「1次元・256スレッドブロックでインデックスする」と宣言し、
    prepare_vecadd が渡された LaunchConfig1D
    その宣言と実機のデバイス制限に照らして検証します。
    検証を通った証(token)を安全な vecadd メソッドが要求する、という構造です。
    contract を持たないカーネルは unsafe な起動メソッドしか公開しません。

DisjointSlice:Rust の借用ルールを GPU に持ち込む

公式ブログは、この vecadd の説明をこう始めています。

Read the kernel signature first, because it carries the whole safety argument.

「安全性の主張はすべてカーネルのシグネチャに詰まっている」ということで、
その中心が出力側の DisjointSlice<f32> です。
入力の ab は全スレッドが読むだけなので、
普通の共有スライス &[f32] で問題ありません。

Rust には
「書き換え可能な参照(&mut)は、同じデータに対して同時に1つまで」
というルールがあります。
ところがカーネルの本体は何千ものスレッドで同時に実行されるので、
出力を &mut [f32] にすると次の状態になります。

スレッド0 ─ &mut [配列全体]
スレッド1 ─ &mut [配列全体]   ← 同じ配列への &mut が同時に何千個も
スレッド2 ─ &mut [配列全体]
   ...

これはルール違反なので、Rust は拒否します。
ただ実際には、各スレッドが書き込むのは自分の担当の1要素だけです。
それなら配列全体への &mut は要らず、重ならない1要素ずつに分けて渡せば安全です。

スレッド0 ─ &mut c[0]
スレッド1 ─ &mut c[1]   ← 互いに重ならない(disjoint)ので、同時にあっても安全
スレッド2 ─ &mut c[2]

DisjointSlice は、この分け方を型で表したものです。
CPU 側の Rust で split_at_mutchunks_mut を使い、
1つの &mut [T] を重ならない断片に分けるのと同じ考え方を、
GPU のスレッド単位でやっています。

分けた断片を取り違えないための仕組みもあります。
thread::index_1d() は生の整数ではなくインデックス型を返し、
c.get_mut(idx) はその型しか受け取りません。
戻り値が Option なので、範囲外は「あとで見つかるメモリエラー」ではなく
「自分で処理する分岐」になります。

後半の Try でも、この型は2か所で結果に関わってきます。
Try 1-2 では、出力バッファを入力にも渡すとコンパイル時に止まることを確かめます。
Try 2 では、1スレッドに2要素を書かせようとして
get_mut(thread::index_1d()) では表現できずに一度 unsafe に降り、
LinearTiles という別の分け方で安全に戻す、という形で再び登場します。

cutile-rs(Tile トラック)

この仕組みは rustc のバックエンド差し替えではなく、マクロと JIT です。
#[cutile::module] マクロがカーネルの AST をホストバイナリに埋め込み、
そのカーネルが最初に必要になった時点でCUDA Tile IR 経由で JIT コンパイルして
GPU cubin を生成します(cutile-rs README)。

同じベクトル加算を Tile で書くとこうなります(公式のコード)。

use cutile::prelude::*;

#[cutile::module]
mod kernel {
    use cutile::core::*;

    #[cutile::entry()]
    fn add<const B: i32>(
        // B is the tile width, a static dimension.
        z: &mut Tensor<f32, { [B] }>, // exclusive output, one sub-tensor of B elements
        x: &Tensor<f32, { [-1] }>,    // shared input; -1 is a dynamic dimension, resolved at launch
        y: &Tensor<f32, { [-1] }>,
    ) {
        // This body runs once per mut sub-tensor, as a single logical thread.
        let tx = load_tile_like(x, z); // the slice of x lining up with this sub-tensor of z
        let ty = load_tile_like(y, z);
        z.store(tx + ty); // elementwise across the whole tile
    }
}

fn main() -> Result<(), Error> {
    let device = Device::new(0)?;
    let stream = device.new_stream()?;

    // These are lazy. Nothing has touched the GPU yet.
    let x = api::ones::<f32>(&[1024]);
    let y = api::ones::<f32>(&[1024]);

    // partition は3つを同時に決める: 各タイルが自分の128要素を排他的に所有すること、
    // グリッドが 1024/128 = 8 タイルであること、そして B の値。
    let z = api::zeros::<f32>(&[1024]).partition([128]);

    let c: Vec<f32> = kernel::add(z, x, y) // takes ownership of all three tensors
        .first()                           // ...and returns them; pick the output back out
        .unpartition()                     // drop the host-side partition wrapper; no data moves
        .to_host_vec()                     // record the copy back
        .sync_on(&stream)?;                // and only now does any of it run
    Ok(())
}

SIMT 版と違い、thread::index_1d() に相当するものが出てきません。
もう1つの大きな特徴が遅延実行です。
api::ones()partition()to_host_vec() も、
その場では GPU に触れません。
公式ドキュメントの用語集は DeviceOp をこう定義しています。

A lazy description of GPU work — allocation, kernel launch, or data transfer — that is not executed until either .sync_on(&stream), .sync(), or .await is invoked.

GPU 上の処理をいったんグラフとして組み立て、最後に一度だけ同期する設計です。
公式ガイドは、22層・1層あたり6カーネルの Transformer を素直に書くと
1トークンあたり132回の同期ギャップが発生する、
という例でこの設計の理由を説明しています。
実行のトリガは .sync() / .sync_on(&stream) / .await / .graph()
の4通りです。

この遅延実行は、後半の Try 2 でも結果に関わってきます。
最初は毎回同期するベンチマークを書いていたのですが、
最後にまとめて同期する書き方より約1.6倍遅い数字が出ていました。

所有権で GPU のバグを防げる、という主張

CUDA Rust のセールスポイントは、カーネルの引数や所有権の誤りを
Rust の型システムでコンパイル時に弾けることです。
GPU では何千ものスレッドが順序保証なく同じバッファに到達するので、
2つのスレッドが同じアドレスに触れて片方が書き込めばデータ競合になり、再現しにくいバグになります。

ただし範囲は限られます。コンパイル時に捕まるのは、
主に「同じバッファを入力と出力の両方に渡した」といった起動時の引数の重なりと、
Tile 側ではテンソルの所有権(GPU が使っている最中のバッファを使い回す等)です。
デバイス側では DisjointSlice や Tile のパーティションが
各スレッドの書き込み領域を分けますが、unsafe で書いた領域、
shared memory の同期の書き忘れ、複数カーネル間の順序といった問題までは関わりません。

SIMT 側では、出力バッファを入力にも渡すと error[E0502] でコンパイルが止まります。
Tile 側で同じことをするとこうなります。

let z = api::zeros::<f32>(&[1024]);
kernel::add(z.partition([128]), z, y)
error[E0382]: use of moved value: `z`

どちらもエイリアシングのミスをコンパイル時に検出しますが、
公式ブログは、Tile 側のほうが安全を保証できる範囲が広いと書いています。
cuda-oxide がチェックするのは、起動の呼び出し1回ごとです。
一方 cutile-rs では、テンソルの所有権が起動の境界をまたいでついて回ります。

cuda-oxide checks each launch call. cutile-rs's ownership follows the tensors across the launch boundary, which is the stronger of the two claims.

Tile 側がここまで安全にできるのは、shared memory の使い方も、
どのスレッドがどのデータを担当するか(スレッドインデックス)も、コンパイラが決めるからです。
プログラマはどちらにも触れないので、そこで間違えることもありません。
その代わり、これらを自分で細かく調整することはできません。

SIMT 側はこの調整をプログラマに残していますが、
shared memory を扱うには現状 unsafe が必要です。
ここを安全に書けるようにする作業は進行中とのことです。

SIMT keeps that control, and today shared memory there requires unsafe. Shared memory is the bedrock of fast SIMT kernels, so making that path safe is active work.

性能は?

公式ブログ本文に具体的な数値はなく、リンク先の論文 Fearless Concurrency on the GPU(arXiv:2606.15991)に載っています。
論文では cutile-rs を cuTile Rust と呼んでいます。
NVIDIA のデータセンター向け GPU である B200 で測った結果は、次のとおりです。

  • 行列積(GEMM)で、NVIDIA のライブラリ cuBLAS の 96% の性能
  • 要素ごとの演算(配列どうしの足し算など)で 7 TB/s
  • どちらも Python 版の cuTile と測定誤差の範囲で同じ

ただし測定対象は cuTile Rust 0.2.0 で、本記事で使った 0.3.1 ではありません。
本記事では L4 上で、自作したカーネルを使って実際に比べます。

現在の成熟度

公式ブログによると以下。

  • 両方とも early-stage。どちらも production-ready ではない
  • cuda-oxide は early alpha
  • cutile-rs はより先行しており crates.io 公開済み。NVIDIA 外でも Grout(HuggingFace の推論エンジン)と mistral.rs で使われている
  • カバレッジは不完全&API は変わる可能性あり

CUDA C++ と CUDA Python を成熟した enterprise-grade のツールチェーンと位置づけたうえで、
CUDA Rust は 2027年以降にかけて育てていく(into 2027 and beyond)としています。
ロードマップには、CUDA C++ / Python との言語間相互運用、
SIMT の shared memory の safe 化、nightly 固定依存の解消が挙がっています。
既存の Rust-CUDA などを置き換えるのではなく、並走する立場です。

Environment

  • MacBook Pro (M3, Apple Silicon) / macOS
  • OrbStack 2.2.3(Ubuntu 24.04 arm64 マシンを作成)
  • AWS GPU インスタンス(g6.xlarge / NVIDIA L4、us-east-1)

測定は NVIDIA ドライバ 595.91.07、CUDA Toolkit 13.2(DLAMI 同梱)、cuda-oxide cargo-oxide v0.2.1、cutile-rs 0.3.1 で行いました。
バージョンの一覧は測定環境に記載しています。

M3 Mac では実行できない

理由は3つあります。

  1. M3 に載っているのは Apple 製 GPU で、NVIDIA GPU ではない。CUDA は NVIDIA GPU 専用です
  2. macOS の CUDA サポートは終了済み。CUDA Toolkit 10.2 が最後で、CUDA 11.0 のリリースノートにも "CUDA 11.0 does not support macOS for developing and running CUDA applications." とあります
  3. CUDA Rust 自体が Linux 前提。cuda-oxide も cutile-rs も要件に Linux を挙げています

OrbStack も GPU を Linux 側へ渡せません(FAQ の「What can't you do?」に Linux graphical apps and GPU とあります)。仮に渡せても Apple GPU なので CUDA は動きません。

ただし、カーネルを実行しないところまでなら話は別です。cuda-oxide は GPU が無い環境でのビルドを公式に想定しています(詳細は B. GPU なし環境での検証)。そこで本記事では、OrbStack の Linux で cuda-oxide の「Rust → PTX」までを確認し、実行と測定だけ AWS の GPU インスタンスで行う、という分担にしました。

GPU の選定:compute capability 8.0 以上が必須

CUDA Rust は両トラックとも compute capability 8.0 以上を要求します。
安易に g4dn を選ぶと動きません。

NVIDIA 公式の compute capability 表AWS の各インスタンスページをチェックすると以下。

インスタンス GPU Compute Capability CUDA Rust us-east-1(USD/時間) ap-northeast-1(USD/時間)
g4dn.xlarge NVIDIA T4 7.5 ❌ 不可 0.526 0.71
g5.xlarge NVIDIA A10G ※後述 1.006 1.459
g6.xlarge NVIDIA L4 8.9 ✅ 可(最安) 0.8048 1.1672
g6e.xlarge NVIDIA L40S 8.9 ✅ 可
p4d.24xlarge NVIDIA A100 ×8 8.0 ✅ 可
p5.48xlarge NVIDIA H100 ×8 9.0 ✅ 可

価格は AWS の料金データフィード(Linux / 共有テナンシー、2026年9月9日時点)の値です。
東京リージョンのプレミアムは一律ではなく、g5/g6 が約1.45倍、g4dn が約1.35倍でした。

※ g5 の A10G は、NVIDIA 公式の compute capability 表に名指しで載っていません(データセンター版の A10 は 8.6)。確実なのは L4(8.9)の g6 です。

というわけで、掲載した候補の中では、執筆時点の us-east-1 で
g6.xlarge が最も安価な対応インスタンスでした。
以降はこれを前提にします。

Setup

公式ブログのコマンドどおりには進まない箇所がいくつかありました。
実際にパスした手順を記述します(理由は A. セットアップの補足)。

1. OrbStack で Linux 環境を作る(ビルド確認用・GPU 不要)

brew install orbstack
orb create ubuntu:24.04 cuda-build
orb -m cuda-build uname -m     # aarch64

作った Linux マシンの中に、次を入れます(全コマンドはセットアップ手順に置きました)。

  • ビルド依存(build-essential、pkg-config など)
  • Rust と固定 nightly(rustup component add rust-src rustc-dev llvm-tools
  • LLVM / Clang 21。Ubuntu 24.04 の標準 clang は 18 なので足りません。llvm.sh 21 では llc を含む llvm-21 が入らないので明示的に追加し、bindgen 用に update-alternatives で無印の clang も 21 に向けます
  • CUDA Toolkit。Apple Silicon 上の Linux は aarch64 なので、x86_64 ではなく sbsa(NVIDIA の ARM サーバ向け)のリポジトリを使います

GPU が無いので nvidia-smi は動きませんが、ビルドには問題ありません。

2. AWS GPU インスタンスを起動する

AMI は AWS 公式の
Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 24.04)
を使います。NVIDIA ドライバ(595.91.07)と
CUDA Toolkit(12.8 / 12.9 / 13.0 / 13.2、既定は 13.2)が入っているので、
そこの手間が省けます(DLAMI リリースノート)。

# 最新の AMI ID を SSM パラメータから取得
aws ssm get-parameters \
  --names /aws/service/deeplearning/ami/x86_64/base-oss-nvidia-driver-gpu-ubuntu-24.04/latest/ami-id \
  --region us-east-1 --query 'Parameters[0].Value' --output text

aws ec2 run-instances \
  --region us-east-1 \
  --image-id <上で取得した AMI ID> \
  --instance-type g6.xlarge \
  --key-name <YOUR_KEY> \
  --security-group-ids <YOUR_SG> \
  --block-device-mappings '[{"DeviceName":"/dev/sda1","Ebs":{"VolumeSize":100,"VolumeType":"gp3"}}]' \
  --tag-specifications 'ResourceType=instance,Tags=[{Key=Name,Value=cuda-rust-demo}]'

EBS は 100GB 以上を推奨。
cuda-oxide は初回に codegen バックエンドをビルドするので、ディスクも時間もそれなりに必要です。
セキュリティグループは SSH(22番ポート)を自分の IP だけに開けましょう。

起動したら SSH して GPU を確認します。

nvidia-smi --query-gpu=name,compute_cap,driver_version,memory.total --format=csv
name, compute_cap, driver_version, memory.total [MiB]
NVIDIA L4, 8.9, 595.91.07, 23034 MiB

compute_cap 8.9 は CUDA Rust の要件(8.0 以上)を、
ドライバ 595.91.07 は cuda-oxide の要件(580 以上)を満たしています。

3. cuda-oxide のセットアップ

GPU インスタンス側では、Rust の固定 nightly と LLVM/Clang 21 を足します。
ドライバと CUDA は DLAMI 同梱のものでOK。

# GPU インスタンス上で
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
source "$HOME/.cargo/env"

sudo apt-get update
sudo apt-get install -y build-essential pkg-config git lsb-release software-properties-common gnupg

# LLVM / Clang 21 と update-alternatives は手順1と同じ(bindgen がバージョンなしの clang を探すため)
llc-21 --version | grep -i nvptx    # NVPTX バックエンドが入っていることを確認

# cargo-oxide のインストール(固定 nightly を使う)
rustup toolchain install nightly-2026-04-03
rustup component add rust-src rustc-dev llvm-tools --toolchain nightly-2026-04-03
cargo +nightly-2026-04-03 install --git https://github.com/NVlabs/cuda-oxide.git cargo-oxide

注意点は以下です(理由は A. セットアップの補足)。

  • llvm-tools コンポーネントを入れる。これが PTX 生成に使う llc を提供する
  • Clang は 21 が必要で、無印の clang が 21 を指すよう update-alternatives を張る
  • 公式ブログのコマンドは nightly-2026-04-03 だが、プロジェクトは別の nightly(2026-08-28)を pin しており、初回ビルド時に自動でダウンロードされる

4. cutile-rs のセットアップ

cutile-rs の要件は cuda-oxide より軽く、stable Rust(1.89 以上)で動きます。
CUDA は README 上「13.3 recommended」ですが、
L4(sm_89)なら DLAMI 同梱の 13.2 で動きました(A-6)。

rustup default stable
rustc --version   # 1.89 以上であることを確認

export CUDA_TOOLKIT_PATH=/usr/local/cuda-13.2
export PATH="$CUDA_TOOLKIT_PATH/bin:$PATH"

git clone https://github.com/NVlabs/cutile-rs.git

自分のプロジェクトに組み込むなら、
0.3.1 以降は cargo add cutile の1行で済みます。

Try

1. Hello world レベルの動作確認

1-1. OrbStack(GPU なし)で PTX まで作成

まず Mac 上の OrbStack で Rust → PTX が通るかを確認します。

cargo oxide new vecadd_demo
cd vecadd_demo
cargo oxide doctor

cargo oxide doctor は、Rust ツールチェーン・CUDA Toolkit・LLVM・clang/libclang・GPU ドライバを順に検証します。
GPU が無い M3 Mac 上の Linux での結果です。

CUDA toolkit (nvcc)... ✓ Cuda compilation tools, release 13.3, V13.3.73
libNVVM (libnvvm.so)... ✓ libNVVM 2.0
llc (LLVM)... ✓ LLVM version 23.1.0-rust-1.100.0-nightly
clang / libclang resource dir... ✓ /usr/lib/llvm-21/lib/clang/21
NVIDIA driver / GPU... - no NVIDIA driver detected
  Only `cargo oxide run` (kernel execution) needs the driver;
  `cargo oxide build` and `pipeline` work without one.

✅ Environment looks good!

GPU が無いのに ✅ Environment looks good! です。
ドライバの項目に付いている一文のとおり、GPU 無しでビルドまで進めるのは公式に想定された使い方です。

ビルドして、生成された PTX を見ます。

cargo oxide build      # 初回は codegen バックエンドのビルドを含むので2〜3分かかる
cargo oxide inspect    # 生成された PTX を表示

NVIDIA GPU が無いマシンで生成された PTX から、要点だけ解説。

// Generated by LLVM NVPTX Back-End

.version 7.0
.target sm_80

.visible .entry vecadd( ... )
.reqntid 256, 1, 1
{
	mov.u32 	%r1, %ctaid.x;
	mov.u32 	%r2, %ntid.x;
	mov.u32 	%r3, %tid.x;
	mul.wide.u32 	%rd12, %r1, %r2;
	add.s64 	%rd3, %rd12, %rd13;
	setp.ge.u64 	%p1, %rd3, %rd11;

%ctaid.x × %ntid.x + %tid.x
先ほど見た CUDA C++ のblockIdx.x * blockDim.x + threadIdx.x です。
Rust では thread::index_1d() としか書いていないのに、
PTX レベルでは、CUDA C++ のインデックス計算に対応する命令列になっています。
「ラッパーではなくネイティブに PTX へコンパイルする」という
NVIDIA の主張のとおり。

.reqntid 256, 1, 1 は、#[launch_contract(block = (256, 1, 1))]
をコンパイラが PTX の指示に変換したものです。
「このカーネルは必ず 1ブロック 256×1×1 スレッドで起動する」
という約束を、PTX 側に書き込んでいます。
#[launch_bounds(256)] のほうは上限を示す .maxntid に対応しますが、
両方を付けたカーネルは .reqntid だけを出します。
setp.ge.u64 の比較は c.get_mut(idx) の境界チェックです。
.target sm_80 は、GPU を検出できないときの既定のターゲットです。
※実機がある環境では自動検出

念のため cargo oxide run も叩くと、ビルドとリンクは通ったうえで、
実行時に libcuda.so が見つからずにエラーになります。

     Running `target/release/vecadd_demo`
Error: DriverError(3, "CUDA driver library unavailable: failed to load any of
  [\"libcuda.so.1\", \"libcuda.so\"]: libcuda.so: cannot open shared object file:
  No such file or directory")

OrbStack でできるのはここまでです。
cargo oxide doctor / build / inspect / pipeline
は GPU がなくても動きますが、カーネルを実行する
cargo oxide run には GPU が必要です。

1-2. cuda-oxide の vecadd を GPU で動かす

ここから AWS の g6.xlarge に移ります。

cargo oxide new vecadd_demo
cd vecadd_demo
cargo oxide run
Detected GPU arch: sm_89 (via nvidia-smi)
...
     Running `target/release/vecadd_demo`
PASSED: all 1024 elements correct

Rust で書いた GPU カーネルが動きました。
GPU があれば nvidia-smi からアーキテクチャ(sm_89)を自動検出します。

次に、わざと間違った呼び出しをしてみます。
入力 a の位置に、出力用のバッファ c_dev を渡します。

module.vecadd(&stream, &prepared, &c_dev, &b_dev, &mut c_dev)?;
error[E0502]: cannot borrow `c_dev` as mutable because it is also borrowed as immutable
  --> src/main.rs:37:55
   |
37 |     module.vecadd(&stream, &prepared, &c_dev, &b_dev, &mut c_dev)?;
   |            ------                     ------          ^^^^^^^^^^ mutable borrow occurs here
   |            |                          |
   |            |                          immutable borrow occurs here
   |            immutable borrow later used by call

コンパイル時点で止まります。
矢印の図で、どこが共有借用でどこが可変借用かまで示されます。

Rust が止めたのは、データ競合を見つけたわけではなく、
同じ変数を読み取り用と書き込み用に同時に借用しようとしたからです。

vecadd では各スレッドが自分の担当する要素だけを読み書きするので、
入力と出力が同じバッファでも、実際には競合は起きません。
公式ブログも、実際に競合するかどうかに関係なくコンパイルは通らない、と書いています。

Passing the SIMT kernel's output buffer as one of its own inputs does not compile, whether or not that kernel would actually race

ここで言っているのは Rust の普通の借用ルールです。
(CUDA C++ では同じ呼び出しがそのままコンパイルできてしまう)

1-3. cutile-rs の hello_world を動かす

cd ~/cutile-rs
cargo run -p cutile-examples --example hello_world

このサンプルは tile プログラム自身に ID を表示させるだけのものです。

#[cutile::entry(print_ir = true)]
fn hello_world_kernel() {
    let (pid0, pid1, pid2) = (program_id(0), program_id(1), program_id(2));
    let (n0, n1, n2) = (num_programs(0), num_programs(1), num_programs(2));
    cuda_tile_print!(
        "Hello, I am program <{}, {}, {}> in a kernel with <{}, {}, {}> programs.\n",
        pid0, pid1, pid2, n0, n1, n2
    );
}

実行結果。

Hello, I am program <0, 0, 0> in a kernel with <1, 1, 1> programs.

DLAMI(AWS Deep Learning AMIs)同梱の CUDA 13.2 のまま動きました。
README には期待出力が Hello, I am tile ... と書かれていますが、
実際の出力は program です(0.3.1 で Triton 由来の語彙に変わった際に、直し忘れた?)。

このサンプルは print_ir = true が付いているので、生成された Tile IR も表示されます。

cuda_tile.module @hello_world_module {
  entry @hello_world_kernel_entry() di_name = "hello_world_kernel" {
    %1, %2, %3 = get_tile_block_id : tile<i32>
    %4 = assume bounded<0, ?>, %1 : tile<i32>
    ...
    %78 = print_tko "Hello, I am program <%i, %i, %i> ...", %4, %15, %29, ... : tile<i32>, ...
  }
}

すべての値の型が tile<i32> で、スカラではなくタイル型になっています。
SIMT 側で inspect すると PTX が出てきたのに対し、
Tile 側はタイル粒度の中間表現が出てきます。

2. softmax を両トラックで自作して比べる

ここからは同じ処理を両トラックで自分で書いて、コードと性能を比較します。
2つのトラックの使い分けを確かめます。

softmax とは?

softmax は、任意の数値の並びを「合計が 1 になる確率」に変換する関数です。
深層学習モデルにおいてマルチクラス分類の出力層でよく使われるやつです。

入力: [-1.5, -1.0, -0.5,  0.0,  0.5,  1.0,  1.5,  2.0]   ← ただの数値
出力: [0.012, 0.020, 0.033, 0.054, 0.089, 0.147, 0.243, 0.401]   ← 合計 ≒ 1.0

大きい値ほど大きな確率になり、合計が(理論上)1 になります。
※実際の計算は浮動小数点なのでずれるが

計算式は「各要素を exp して、その総和で割る」だけです。
exp はすぐ桁あふれするので、実装は先に行の最大値を引く

softmaxは主に以下の箇所で使われます。

  1. LLM が次のトークンを選ぶとき。 モデルは語彙すべてにスコアを出すが、そのままでは扱いづらいので softmax で確率に直し、そこからサンプリングして1トークンを決める。文章生成中、1トークンを生成するたびに実行される処理。

  2. Attention の重み計算。 Transformer の中核で、「この単語は他のどの単語にどれだけ注目すべきか」を決める。1回の推論で何回も実行される。

名前に「row-wise」(行ごと)と付くのは、実際には数値の並びが1本だけ来るのではなく、
たくさんの並びを積み重ねた行列が来て、その1行ずつに softmax をかけるからです。

Attention を例にすると、トークン1つにつき
「他のトークンそれぞれにどれだけ注目するか」が1行できます。
これがヘッド(注目の観点)ごとに作られるので、
たとえばヘッド数 32、文章の長さ 1,024 トークンなら、
32 × 1,024 = 32,768 行、
1行の長さは 1,024 になります。

どの行も他の行とは無関係に計算できるので、
行ごとに別のブロックへ割り当てて同時に処理できます。
行が数千〜数万あれば GPU 全体に仕事が行き渡るので、GPU 向きの処理です。
本記事では、4,096 行 × 256 列の行列を使います。

2トラックの比較に向く理由

softmax は GPU カーネルの2つの要素を両方含んでいます。
1行の softmax は、次の4ステップです。

  1. 行の最大値を求める(reduction:多数の値を1つに畳みこむ)
  2. 各要素から最大値を引いて exp する(elementwise:要素ごとに独立)
  3. その総和を求める(reduction)
  4. 各要素を総和で割る(elementwise)

elementwise は簡単です。
スレッドが自分の担当を処理して終わりで、他のスレッドとやりとりする必要がありません。

reductionの場合、「256個のスレッドがそれぞれ持っている値の最大値」を求めるには、
スレッド同士が値をやり取りしないといけません。
GPU にはそのための仕組みがあります。

  • warp shuffle: 高速。同じ warp(32スレッド)の中なら、レジスタの値を直接やり取りできる。
  • shared memory: warp をまたぐには、ブロック共有のメモリに一度書き出し、sync_threads() で全て書き込み完了を待ってから読む

SIMT ではこの2つの処理を自分で実装します。
※Tile では reduce_max(tile, 1) の1行で済む
GEMM だと両トラックともそれなりに書けてしまうのですが、
reduction は差がはっきりと出ます。

SIMT 版(cuda-oxide)

cuda-oxide には softmax の example が存在しないので完全に自作です。
設計は「1行を1ブロック、1要素を1スレッド」にしました。
こうすると thread::index_1d()
「行 × 列数 + 列」というフラットな添字になり、
入出力の添字に直接使えます。

const COLS: usize = 256;        // 1行の長さ = ブロックのスレッド数
const WARPS: usize = COLS / 32; // ブロック内の warp 数 = 8

#[kernel]
#[launch_bounds(256)]
pub fn softmax_rows(x: &[f32], mut y: DisjointSlice<f32>) {
    // ブロック内 reduction の中継地点。warp ごとに1スロット使う。
    static mut REDUCE: SharedArray<f32, WARPS> = SharedArray::UNINIT;

    let gid = thread::index_1d();        // 行 * COLS + 列
    let lane = warp::lane_id();          // warp 内での位置(0..32)
    let wid = warp::warp_id() as usize;  // ブロック内での warp 番号(0..8)

    let v = x[gid.get()];

    // ---- 1. 行の最大値 ----
    let wmax = warp::reduce_max_f32(v);      // まず warp 内32レーンを集約
    if lane == 0 {
        unsafe { REDUCE[wid] = wmax };       // 代表が shared memory に置く
    }
    thread::sync_threads();                  // 全 warp の書き込み完了を待つ

    let mut row_max = unsafe { REDUCE[0] };  // 8個の warp 最大値から全体の最大値
    for i in 1..WARPS {
        let m = unsafe { REDUCE[i] };
        if m > row_max { row_max = m; }
    }
    thread::sync_threads();                  // 全員が読み終わるまで上書きさせない

    // ---- 2. 最大値を引いて exp ----
    let e = (v - row_max).exp();

    // ---- 3. 行の総和(1と同じ2段構え)----
    let wsum = warp::reduce_sum_f32(e);
    if lane == 0 {
        unsafe { REDUCE[wid] = wsum };
    }
    thread::sync_threads();

    let mut row_sum = 0.0f32;
    for i in 0..WARPS {
        row_sum += unsafe { REDUCE[i] };
    }

    // ---- 4. 正規化して書き出し ----
    if let Some(dst) = y.get_mut(gid) {
        *dst = e / row_sum;
    }
}

sync_threads() の位置に注意。
2回目の reduction で REDUCE を再利用しているので、
全スレッドが最大値を読み終わる前に総和で上書きしてはいけません。
そのため最大値を読んだ直後にもう1回同期を入れています。
ここを忘れると、たまに間違った答えが出るというバグになります。
※これはSIMTでreductionを書くときのあるあるらしい

unsafe は shared memory アクセスの5箇所に必要でした。
記事の前半で触れた「SIMT の shared memory は現状 unsafe」が
そのまま形になっています。

このコードは GPU が無い OrbStack でコンパイルしてからGPU(AWSインスタンス)に送りました。
cargo oxide inspect で見ると、意図したとおりの命令になっています。

.visible .shared .align 4 .b8 __shared_mem_0[32];   ← f32 × 8 warp = 32 バイト
shfl.sync.bfly.b32  ...                              ← 10個(butterfly 5段 × 2回)
bar.sync 0;                                          ← 3個(sync_threads() 3回)
ex2.approx.ftz.f32  ...                              ← exp がハードウェア命令に

warp::reduce_max_f32 が PTX でどう展開されるかは D-1 に記述

Tile 版(cutile-rs)

同じ処理を Tile で書きます。

#[cutile::entry()]
fn softmax_rows<const BM: i32, const BN: i32>(
    y: &mut Tensor<f32, { [BM, BN] }>,  // 排他的な出力タイル
    x: &Tensor<f32, { [-1, -1] }>,      // 共有入力
) {
    let tx: Tile<f32, { [BM, BN] }> = x.load_like(y);

    let m: Tile<f32, { [BM] }> = reduce_max(tx, 1i32);                  // ← reduction が1行
    let m: Tile<f32, { [BM, BN] }> = m.reshape(shape![BM, 1]).broadcast(y.shape());

    let num: Tile<f32, { [BM, BN] }> = exp(tx - m);

    let den: Tile<f32, { [BM] }> = reduce_sum(num, 1i32);               // ← reduction が1行
    let den: Tile<f32, { [BM, BN] }> = den.reshape(shape![BM, 1]).broadcast(y.shape());

    y.store(num / den);
}

スレッドもshared memory もsync_threads()unsafe も出てきません。
数式の softmax がほぼそのまま書けています。
SIMT 版で神経を使った「2回目の reduction の前にもう1回同期する」という話も、
そもそも同期という概念がないので考える必要がありません。

Tile 版は cutile-rs 同梱の softmax example とほぼ同じコードになりました。
※SIMT 版のほうは、ブロックの割り当て方も reduction の組み方も何通りもある

結果

L4(g6.xlarge)で、4096行 × 256列に対して測定しました。
入力は両方とも x[i] = i * 0.01 で揃え、
ウォームアップ1回のあと200回の平均を取っています。

Tile は BM=1, BN=256 です。

SIMT Tile 結果
batched 14.2 µs / 588.8 GB/s 10.0 µs / 843.1 GB/s Tile が 1.42倍速い
per-iteration sync 18.9 µs / 443.3 GB/s 15.9 µs / 526.9 GB/s Tile が 1.19倍速い
max │rowsum - 1.0│ 4.768e-7 4.768e-7 一致
カーネル本体の行数 約40行 約10行
unsafe の箇所 5 0
sync_threads() 3 0

どちらも特に最適化せずに書いたものですが、
Tile 版のほうが速い結果になりました(batched で 1.42倍)。
Tile 版はコードが短く、unsafesync_threads() も要りません。

計算結果はどちらも正しく出ています。
各行の合計は理論上ちょうど 1 になるはずで、
実際のずれは両方とも約 0.0000005(4.768e-7)と、浮動小数点の誤差の範囲でした。
SIMT 版は CPU で計算した答えとも比較し、ほぼ一致しました(差は最大 5.6e-9)。

batched は200回まとめて投げて最後に1回だけ同期した場合、
per-iteration sync は毎回同期した場合です。
毎回同期すると SIMT 版で1反復あたり 4.7 µs 増えており、
14.2 µs のカーネルに対して無視できない比率になります。

なお実効帯域(論理的な読み書き量 8 MB ÷ 時間)が
L4 のメモリ帯域(300 GB/s)を超えていますが、
これはL4 の L2 キャッシュは 48 MB あり(Ada アーキテクチャ白書 Appendix D)、今回のワーキングセット 8 MB は L2 に収まるうえ、同じバッファを200回連続で処理しているため、
DRAM をほとんど触らないからです。
D-2 のプロファイルでも DRAM 転送は 218 KB

補足:Tile 版の数字を出すまでの2つの失敗

上の Tile の数字は、2つの失敗を直したあとのものです。
どちらも cutile-rs のドキュメント(performance.md の Common Pitfalls)に、
よくある間違いとしてそのまま載っていました。

1つめはタイルの形です。最初はタイル1枚に8行を持たせる BM = 8 にしていて、
毎回同期する条件で 1371.8 µs かかっていました。
1行ずつの BM = 1 にすると 15.9 µs で、変えたのは BM だけなのに約86倍の差です。
ドキュメントは「タイルが大きすぎるとレジスタ圧で occupancy が下がる、あるいはスピルする」と説明しています。
プロファイラで確かめたわけではないので、この説明が当てはまる可能性が高い、というところまでです。
NVIDIA の言う「コンパイラがタイルをアーキテクチャに対応させる」は
タイル→スレッドの割り当ての話で、タイルをどう切るかはプログラマの設計判断として残ります。
cutile-rs にはその探索用の autotuning 機能(cutile::tune)も用意されています。

2つめは同期の仕方です。最初は1回ごとに .sync_on(&stream) でブロックしていましたが、
公式ベンチマークは .async_on(&stream) で投げっぱなしにし、最後に1回だけ同期していました。
書き換えると 15.9 µs → 10.0 µs(1.59倍)になりました。

  • Excessive synchronization: .sync() after every operation creates CPU/GPU gaps.

前半で見たとおり、DeviceOp は「遅延して組み立て、最後に一度だけ同期する」ために存在します。
毎回同期する書き方は、その設計を無効にしてしまいます。
cutile-rs を触るなら、Common Pitfalls には目を通しておくとよいです。

なぜ SIMT のほうが遅かったか

Nsight Compute(ncu。CUDA Toolkit に同梱)で調べました。

sudo ncu -k regex:softmax --launch-count 1 --metrics <上の表のメトリク> ./target/release/softmax_simt
メトリクス SIMT(自作) Tile(自作)
grid 4096 blocks 4096 blocks
block size 256 threads 128 threads
スレッドあたりのレジスタ数 20 21
ブロックあたりの shared memory 32 B 16 B
achieved occupancy 89.19% 89.38%
実行命令数 4,030,464 2,506,752
DRAM 転送量 4.51 MB 4.50 MB

occupancy もレジスタ使用量も DRAM 転送量もほぼ同じで、違うのは命令数だけです。
SIMT 版が 1.61倍多く実行しており、実際に測った速度差(1.42倍)と合っています。

命令数を warp 数で割ると理由がわかります。
smsp__inst_executed は warp 単位の命令数

ブロックサイズ warp 数 warp あたり命令数 総命令数
SIMT 256 スレッド 4096 × 8 = 32,768 123 4,030,464
Tile 128 スレッド 4096 × 4 = 16,384 153 2,506,752

Tile コンパイラは「1スレッドが2要素を担当」という割り当てを選び、
warp 数を半分にしていました。
warp あたりの命令数は Tile のほうが多い(153 vs 123)のですが、
reduction の固定コスト(warp shuffle、shared memory への書き出し、sync_threads())が
2要素分に薄まるため、総命令数では SIMT のほうが多くなります。

私は「1要素 = 1スレッド」という最も素直な対応付けを選んでいて、
それが命令数を増やしていたと思われます。

確かめる&SIMT を書き直す(v2)

SIMT 版を、1ブロック128スレッド・1スレッド2要素に変えてみます。
各スレッドが TPB 間隔で2要素を担当するので、
隣接するスレッドが隣接要素を読む形は保たれます。

const TPB: usize = 128;         // 1ブロックのスレッド数
const PER: usize = COLS / TPB;  // 1スレッドが担当する要素数 = 2

let mut v = [0.0f32; PER];
for k in 0..PER {
    v[k] = x[base + k * TPB];
}

// reduction に渡す前に、まず自分の担当分の中で畳んでおく
let mut local = v[0];
for k in 1..PER {
    if v[k] > local { local = v[k]; }
}
let wmax = warp::reduce_max_f32(local);

結果は以下。

SIMT softmax (2要素/スレッド)  rows=4096 cols=256 tpb=128
  batched (1 sync)     = 8.2 us  (1018.8 GB/s)
  per-iteration sync   = 12.5 us  (668.6 GB/s)

14.2 µs → 8.2 µs。1.73倍速くなり、Tile(10.0 µs)を上回りました。
命令数も 4,030,464 → 2,064,384 と、
warp 数を Tile と同じ 16,384 に揃えたうえで
Tile(2,506,752)を下回っています。

ただし、このままでは書き出しに unsafe が必要です。
それまで使っていた get_mut(thread::index_1d())
1スレッドが1要素だけを触る API なので、
1スレッドが複数要素を書く形を表現できないためです。

for k in 0..PER {
    let idx = base + k * TPB;
    unsafe { *y.get_unchecked_mut(idx) = e[k] / row_sum };
}

unsafe使わない(v3)

DisjointSlice には、複数要素を安全に扱う API も別に用意されています。
DisjointSlice<f32, LinearTiles<2>> と宣言すると、
スレッド tt*2 .. t*2+2 の連続領域を所有し、
tile_thread32 が境界チェックを1回だけ行ってチェック不要のビューを返します。

pub fn softmax_rows(x: &[f32], mut y: DisjointSlice<f32, LinearTiles<PER>>) {
    let thread_index = thread::index_1d_u32(launch_context);
    // ...
    let Some(mut run) = y.tile_thread32(thread_index) else { return; };
    // ...
    // 書き出し。unsafe も実行時の境界チェックもない
    run.at_const::<0>().write(e[0] / row_sum);
    run.at_const::<1>().write(e[1] / row_sum);
}

at_const::<I>() は「ビューの何番目か」をコンパイル時にチェックするので、
実行時のチェックも unsafe も要りません。

SIMT softmax v3 (安全API / 2要素連続)  rows=4096 cols=256 tpb=128
  all finite           = true
  max |gpu - cpu|      = 5.588e-9
  batched (1 sync)     = 9.4 us  (890.4 GB/s)

書き出しに unsafe を使わなくても、Tile(10.0 µs)を上回りました。
v2(8.2 µs)よりは約15%長いものの、v2 と v3 はアクセスの配置も違います
※v2 は各スレッドが128要素おきの2要素、v3 は連続した2要素

v3 のほうが遅い原因は、1つに絞れません。
プロファイラで見ると、v3 はメモリアクセスの効率が v2 の半分でした。
※L1/TEX のセクタアクセス数が v2 の2倍
v3 では1回の読み書きで warp 内の32スレッドが1要素おきの位置に触れるため、
1命令で取ってきたデータの半分しかその命令では使われないからです。
(coalescing の話)

一方で、実行した命令の数も7%増えていました。
配置と API の両方が変わっているので、15%の差がすべて配置のせいとは言えませんし、
「安全な API にしたコスト」とも言い切れません(詳細は D-3)。

比較結果まとめ

最終的な数字は以下。

実装 時間 実効帯域 行数 unsafe
SIMT v1(素直に書いた) 14.2 µs 588.8 GB/s 約40 5
Tile 10.0 µs 843.1 GB/s 約10 0
SIMT v3(安全 API・連続配置) 9.4 µs 890.4 GB/s 約45 5
SIMT v2(unsafe・ストライド配置) 8.2 µs 1018.8 GB/s 約45 6

unsafe の残り5箇所はいずれも shared memory(SIMT の既知の制約)

  • 書きやすさは明確に Tile。行数が1/4、unsafe ゼロ、同期の位置を考えなくてよい
  • 素直に書く限り、性能も Tile。Tile コンパイラは何も指定しなくても、素直な手書き SIMT より良いスレッド割り当てを選んでいた
  • 突き詰めれば SIMT が勝つが、「warp 数を減らす」という勘所と、プロファイラで確かめる手間が必要。書き出しに安全な API を使ったままでも Tile は上回れた
  • Tile でも、タイルの切り方は自分で決める必要がある。タイル1枚が受け持つ行数(BM)を 8 から 1 に変えただけで、実行時間が大きく変わった

NVIDIA の「まず Tile を選び、制御が必要なときに SIMT へ」という指針のとおり。

開発フロー:型チェックはローカル、GPU 実行はクラウドで

今回の検証ではOrbStack とクラウド GPU を以下のように使いました。

  • OrbStack(GPU なし):コードを書いて型チェック&PTX の確認(cuda-oxide)
  • g6.xlarge(L4):カーネル実行、性能測定、プロファイル

cuda-oxide の初回ビルドは、OrbStack(M3)で 2分23秒、
g6.xlarge で 2分26秒とほぼ同じでした。
ビルドは CPU 律速なので、GPU インスタンスの時間でコンパイラを待つ意味はありません。
本記事の自作カーネルも、すべて OrbStack で型エラーと PTX を確認してから送り、
GPU 側でのビルド回数を減らしました。
※送ったのはソースで、バイナリではない

OrbStack は aarch64、g6.xlarge は x86_64 なので、
GPU 側でもビルドし直しています。
ローカルでは以下を実施し、ビルド作業そのものはクラウドで実施。

  • 型エラーや API の実装ミスがないか確認
  • PTX を確認

後片付け

検証が終わったらちゃんとインスタンスを止めましょう。

aws ec2 stop-instances --region us-east-1 --instance-ids <INSTANCE_ID>
# もう使わないなら
aws ec2 terminate-instances --region us-east-1 --instance-ids <INSTANCE_ID>

Summary

本記事では、CUDA Rust を、実際に動かし、
softmax を両トラックで書いて比較してみました。

CUDA Rustが出てきたことで、「Rust から CUDA を呼ぶ」から
「Rust を CUDA のカーネル言語として正式に育てる」ように
ステージがかわりつつあります。

ただ、 NVIDIA 自身が両プロジェクトとも production-ready ではないと明言しているので、
本番採用にはまだ慎重さが必要でしょう。

2つのトラックの使い分けをまとめると以下。

  • cutile-rs(Tile)— NVIDIA が「まずこっち」と言っているほう。stable Rust で入り、softmax が数式に近い形で書け、スレッドや shared memory はコンパイラが持つ。普通に実装すれば SIMT より速かった。ただしタイルの形は自分で決める必要あり
  • cuda-oxide(SIMT)— thread / block / warp / shared memory という CUDA の本質をそのまま Rust で扱える。スレッドへの要素割り当てをしっかり考慮すれば Tile を上回れる。固定 nightly が必要で、shared memory はまだ unsafe

安全性の面では、起動時の引数の重なりはホスト側の借用規則が E0502 で止めてくれますが、
shared memory は unsafe のままで、どこまで減らせるかは実装方法とデータの配置次第です。
同梱の gemm_views では safe 版と unsafe 版の結果がビット単位で一致し、時間差も誤差の範囲でした(C-1)。

M3 のローカルではカーネルを実行できませんが、OrbStack の Linux で「Rust → PTX」までは確認でき、
実行と測定だけクラウドのインスタンスで行いました。

本番で使えるようになるには、まだ以下の条件がクリアされる必要がありそうです。

  • cuda-oxide が nightly 依存でなくなる
  • SIMT の shared memory が unsafe でなくなる
  • API が安定し、crates.io に安定版がリリース
  • 対応 CUDA / GPU アーキテクチャの正式なサポートマトリクスが公開

References

そのほかの参考リンクは本文中にインラインで貼っています。全リストは検証メモにもまとめてあります。

Appendix:検証メモと自作カーネル

セットアップでつまずいた点、GPU なし環境での検証、同梱サンプルの実行結果、プロファイラの数値、
そして自作した softmax カーネル4本のソースは GitHub に置きました。

nakamura-shuta/examples の cuda-rust

この記事をシェアする

関連記事