💻

CUDA Tileを解説: cuTile Pythonで行列積を実装してみよう

に公開

1. はじめに

2025年12月はじめ、NVIDIAはCUDA 13.1とともにCUDA Tileという新しいプログラミングモデルを発表しました。NVIDIAはテックブログにて、CUDA Tileのことを「2006年のCUDAプラットフォーム発明以来、最大の進歩」と位置づけています。

最大の進歩と言われてしまうと、どんなものなのか気になってしまいますね。そこで本記事では、CUDA Tileについて公式情報に基づき概要を整理しつつ、Python APIであるcuTile Pythonの行列積サンプルを紐解き、Tileプログラミングに触れてみようと思います。

想定読者

  • CUDAプログラミングの基礎知識がある方(以前の記事等を適宜参照ください)
  • Pythonプログラミング知識は前提とします

注意事項

2025年12月時点で、cuTile PythonはBlackwell GPU(Compute Capability 10.x / 12.x)のみ対応しており、現時点では実行できる環境は限られていますのでご注意ください。ただ、実行できなくてもTileプログラミングの雰囲気は掴んでもらえると考えています。

2. CUDA Tileとは何か

CUDA Tileについて、従来のCUDAプログラミングと比較しながら解説します。

2.1 SIMTモデルの振り返り

従来のCUDAプログラミングはSIMT(Single Instruction, Multiple Thread)モデルに基づいています。このモデルでは、プログラマは各スレッドの気持ちになり、スレッドがどの要素にアクセスし、どのような演算を行うかを明示的に記載します。

// 従来のCUDA: スレッド単位で記述
__global__ void vector_add(float *a, float *b, float *c, int n) {
    // どのインデックスの計算を行うか特定
    int idx = blockIdx.x * blockDim.x + threadIdx.x;  
    // 境界チェックを実施のうえ、上で特定したインデックスの要素を加算
    if (idx < n) {
        c[idx] = a[idx] + b[idx];
    }
}

このアプローチには以下の課題があります:

  • ハードウェア依存: スレッド数、ブロックサイズ、共有メモリの使い方がGPUアーキテクチャに依存
  • Tensor Coreの複雑さ: wmma等で行列の乗算・積算を明示的にプログラミング
  • 移植性の低さ: 新しいGPU世代ごとに再チューニングが必要

上記の課題は、この画像をみると明らかだと思います。左側がカリカリにチューニング(Speed-Of-Light)したSIMTプログラミングでの行列乗算カーネル、右側が本記事で解説するcuTileで記載した行列乗算カーネルです。 (DeepDive: How to Use cuTile Python (Youtube)1:00〜あたり)

特にTensorcoreをフル活用しようとすると、上記画像の左側のようなコード量を記載しないと性能が発揮されないです。さらに、アーキテクチャが異なるとそのまま移行できず(Non-Portable)、アーキテクチャ毎にチューニングが必要となります。やってられないですよね。

なお、上記のような従来のCUDAプログラミングは、これまでは「CUDAプログラミング」や「CUDA Kernelプログラミング」と呼ばれていました。しかし今回のCUDA Tileの登場により、レトロニウムとして従来のCUDAプログラミングのことを「SIMTプログラミング」(もしくはシンプルに「SIMT」)と呼ぶようになったようです。本記事でもその記載方法に倣います。

2.2 Tensor Core

CUDA Tileは、Tensor Coreをプログラマの負担を下げつつ、最大限活用することを念頭に開発されたものと考えられます。Tensor Coreは、Volta世代以降NVIDIA GPUに搭載された行列演算に特化したハードウェアユニットです。通常のCUDA Coreが条件分岐等を含む汎用的な計算を行えるのと異なり、行列積和演算(MMA: Matrix Multiply-Accumulate)に特化した計算コアとなっています。

Tensor Coreが実行するMMA演算は以下のとおりです:

D = A \times B + C

MMA演算は、行列積を効率的にTensor Coreで計算するために用いられます。例えば4 \times 4 行列X, Yについて積Z=X\times Y(1, 2)成分(Z_{1,2})を求める場合を考えます。計算式は以下となります。
Z_{1,2}= X_{1,0}*Y_{0,2}+X_{1,1}*Y_{1,2}+X_{1,2}*Y_{2,2}+X_{1,3}*Y_{3,2}

この計算は、以下のように逐次的に足し算を行うことで求めることができます。
S_{1,2}=0

S_{1,2}=X_{1,0}*Y_{0,2}+S_{1,2}

S_{1,2}=X_{1,1}*Y_{1,2}+S_{1,2}

S_{1,2}=X_{1,2}*Y_{2,2}+S_{1,2}

S_{1,2}=X_{1,3}*Y_{3,2}+S_{1,2}

Z_{1,2}=S_{1,2}

結局、行列積は2つの数字を掛け算し、それをひたすら積算していくという処理に帰着されます。そのため、Tensor Coreは下の図のような掛け算→積算に特化した構成となっています。以下はVolta世代のTensor Coreの処理イメージです。
Figure 9: Volta GV100 Tensor Core operation.
一番左の二つのFP16がXおよびYの特定の成分の入力に相当します。そしてFP32のaccumulator(積算器)がSの特定の成分にあたります。ここで地味なようで非常に大切なポイントが、積を計算する二つの数値はFP16(単精度)をインプットするのに、掛け算の出力およびその加算はFP32(倍精度)で行なっている点です。これは、インプットの精度を下げグローバルメモリからの読み込み等を高速化しつつも、乗算の結果を積算するさいの丸め誤差を抑える仕組みとなっており、計算速度と精度の両立を図っています。
上記はVolta世代のTensor Coreのイメージですが、それ以降の世代も基本的に構造は変わっていないと思われます。Volta以降の世代では、量子化したディープラーニングモデルに対応するため、FP8やBF16などより低精度の浮動小数点に対応したり、Tensor Core専用のオンチップメモリ(TMEM: Tensor Memory)が追加されたりしています。

2.3 CUDA Tile誕生の背景

従来、CUDAはPTX(Parallel Thread Execution)というISAを介することで、異なるGPUアーキテクチャ間でのコードの互換性を実現してきました。開発者はCUDAコードを記述し、それがPTXに変換され、実行時にドライバがJIT(Just-In-Time)コンパイルによって実際のGPU用のマシンコードに変換します。この仕組みにより、古いコードが新しいGPU上でも動作することが保証されていました。

しかし、Tensor Coreの登場およびその最大限の活用を目指して、上記のPTXの互換性が一部破られ、特定のCompute Capability(CC)に限定したコンパイル方法が登場しました。Hopper以降で導入されたArchitecture-specific features(9.0や12.0など特定のCCのみに特化)と、Blackwellで導入されたFamily-specific features(10.xや12.xなど特定の世代に特化)です。

これらの機能は主にTensor Coreに関連する高度な最適化を可能にしますが、PTXの通常の抽象化メカニズムでは表現しきれません。これは、PTXによる抽象化に限界が訪れていたことを示しています。

2.4 Tile IRとは

CUDA Tileの核心はCUDA Tile IRという新しいISA(命令セットアーキテクチャ)です。前節のPTXの限界が、Tile IRを新規に開発した主要なモチベーションであったと推察しています。Tile IRはPTXとは異なる抽象化を提供することで:

  • アーキテクチャ固有の最適化を自動化:Tensor Coreの世代間の違いをコンパイラが吸収
  • 将来互換性の向上:コード変更なしで新アーキテクチャの恩恵を受けられる
  • 開発者の負担軽減:Compute Capabilityを意識せずに高性能コードを記述

Tile IRは、PTXと並存する新しいコンパイル方法として位置づけられており、特にTensor Coreを活用する行列演算において、SIMTプログラミングの複雑さを大幅に軽減することを目指しています。TileプログラミングではcuTileで記載したコードをTile IRに変換します。Tile IRではコンパイル時にアーキテクチャにあわせた最適化が可能なため、完璧ではないものの最新アーキの恩恵をコード変更無しで受けられます。

cuTileやTile IRは、下図のとおり従来のCUDAやPTXと並列な位置付けとなっています。

A diagram of how the Tile path of compilation fits into the full software stack, adjacent to the SIMT path. The SIMT path includes NVVM/LVVM and PTX, whereas the tile path includes Tile IR.

繰り返しになりますが、TileはSIMTを置き換えるものではありません。Tileは行列を含む配列の演算に特化した構成となっていることから、上記画像のとおり汎用的なSIMTと今後も共存するものであることに留意してください。

3. cuTile Pythonの基本

ここまででCUDA Tileの概要について解説しました。ここからは実際にCUDA Tileを動かしてみましょう。本記事執筆時点(2025年12月)では、cuTileを利用する唯一の方法はPythonベースのDomain Specific Language(DSL)であるcuTile Pythonとなっています。

3.1 プログラミングモデル

Tileプログラミングの流れは大まかに以下のとおりとなります。
①ブロックにグローバルメモリ上のデータのタイルを読み込ませる
②ブロックに読み込んだタイルの計算処理を行わせる
③ブロックの処理結果をグローバルメモリに返す

上記でわかるとおり、cuTile Pythonでは開発者はスレッドに対する明示的な指示を行いません。代わりに

  • どのブロックが
  • どのタイルに対し
  • どのような処理をするのか

を意識すればよく、スレッドの割り当てやメモリ管理はコンパイラが自動的に行います。

3.2 基本的なカーネル構造

cuTile Pythonでベクトル加算を実装すると、以下のようになります:

import cuda.tile as ct

@ct.kernel
def vector_add(a, b, c, tile_size: ct.Constant[int]):
    # ブロックIDを取得
    bid = ct.bid(0)
    
    # グローバルメモリに配置した配列からタイルをロード
    a_tile = ct.load(a, index=(bid,), shape=(tile_size,))
    b_tile = ct.load(b, index=(bid,), shape=(tile_size,))
    
    # 演算(タイル全体に対して行われる)
    result = a_tile + b_tile
    
    # 結果をストア
    ct.store(c, index=(bid,), tile=result)

このコードには、cuTileカーネルに共通する構造が現れています:

  1. @ct.kernelでカーネルであることを宣言
  2. ct.bid()でブロックIDを取得
  3. ct.load()でグローバルメモリからタイルを読み込み
  4. タイルに対して演算を実行
  5. ct.store()で結果をグローバルメモリに書き戻し

注目すべき点として、ct.Constant[int]で宣言された引数はコンパイル時定数として埋め込まれます。この変数はConstantと名乗ってはいますが少しややこしく、あくまでカーネル実行中は変更できないというだけで、ホストコード内では変更可能です。ただしこの値が変更になった場合は、再度JITコンパイルが実行されます。

4. 環境セットアップ・動作確認

cuTile PythonのQuickstartを概説します。更新があるかもしれませんので、うまくいかない場合は適宜元ページを参照ください。

4.1 動作要件

動作要件は記事執筆時点(2025年12月)で以下となっています

項目 要件
GPU(CC) 10.x / 12.x(Blackwellのみ)
Driver R580以降
CUDA Toolkit 13.1以降
Python 3.10, 3.11, 3.12, 3.13
OS Linux x86_64, Linux aarch64, Windows x86_64

DriverやCUDA Toolkitのインストール/アップデートについてはこちらの記事等を参照ください。Linux環境の場合のDriverアップデートのさいは、アップデート前に古いものを削除することを強くお勧めします。

4.2 インストール

cuTile Pythonのインストールは、公式の案内に従いpipで行います。適宜venvを作成し実施してください。

# cuTile Pythonのインストール
pip install cuda-tile

そのほか、サンプルコードの実行にあたりいくつか追加インストールが指定されていますので、素直に従います。

# サンプル実行に必要なパッケージ
pip install cupy-cuda13x
pip install pytest numpy

インストール後、以下のベクトル加算のTileコードを実行してみてください。
https://docs.nvidia.com/cuda/cutile-python/quickstart.html#example-code

5. MatMul.py解説

cuTile Pythonを体験するために、行列積(matmul)のサンプルコードMatMul.pyを詳しく見ていきましょう。5.2で解説するswizzleというテクニックを使っているため少々複雑ですが、行列積の演算に関わる部分は非常にシンプルとなっており、本記事で解説するシンプルな記載だけで、高度に最適化されたPyTorchによる演算と比較できるレベルの性能を発揮することができます。

5.1 ブロックに処理を割り当てる

Tileプログラミングでは、以下のような流れで指示を行います。

  1. タイルの大きさを定義
  2. ブロックをどのように並べるか(ディメンジョン)をグリッドとして定義
  3. 各ブロックが「どのタイル」を「どう処理する」かをカーネルとして定義
  4. 以下をランチャに与えてカーネルを起動
    • グリッド
    • カーネル
    • インプット
    • アウトプット出力先
    • タイルのサイズ

それぞれのステップについて、MatMul.pyに基づき概説します。まずタイルの大きさですが、MatMul.pyでは以下のように与えられています。

if A.dtype.itemsize == 2:  # Likely torch.float16 or torch.bfloat16
	tm, tn, tk = 128, 256, 64  # Larger tiles for Tensor Core friendly types
else:  # Likely torch.float32 or other
	tm, tn, tk = 32, 32, 32   # Smaller, more general tiles

このサイズの決定方法についての解説は見つけられませんでしたが、FP16やBF16などの16bit(2byte)の型については比較的大きなタイルサイズ、FP32などの32bitの型については比較的小さいタイルサイズを採用すると良さそうです。このオーダー感を参照すると良いでしょう。

なお、タイルサイズは現時点で2の累乗(2, 4, 16, 32, ...)である必要があるという制限がありますので留意してください。

次にグリッドの定義です。前のステップで定義したタイルを処理対象に敷き詰めて、その個数に応じてグリッドのサイズを定義することになります。行列積の場合はSIMTと同様に、出力Cに対してタイルを敷き詰めて、どのタイルを処理するかを指示するためにグリッドを定義します。出力Cに対してタイルを敷き詰めるイメージを示します。グレーのセルが行列Cであり、セル1つが行列の1要素を示しています。そして、黒い太線の四角それぞれがタイルです。

なお、行列の縦横の幅は必ずしもタイルの縦横の整数倍ではありません。そのため、オレンジ・紫・赤のようなタイルが行列からはみ出る部分が出てきます。この部分の処理については5.3節のpaddingの処理内容の紹介にて解説します。

グリッドは三次元まで定義でき、行列は二次元の配列なので、素直にグリッドも二次元にするのが良さそうです。しかしMatMul.pyでは、次の節で解説するSwizzleという効率的なメモリアクセス手法を使用するために、以下のとおりグリッドとして一次元を採用しています。

# 行列サイズを取得
m, k_a = A.shape
k_b, n = B.shape

# グリッドサイズを計算
grid_x = ceil(m / tm) #タイルをCのM方向に敷き詰める
grid_y = ceil(n / tn) #タイルをCのN方向に敷き詰める
grid = (grid_x * grid_y, 1, 1) #1次元のグリッドを定義

カーネルの定義については後述します。カーネルの定義後はランチャによりカーネルを起動します。

# カーネルを起動
ct.launch(torch.cuda.current_stream(), grid, matmul_kernel, (A, B, C, tm, tn, tk))

最初の引数のtorch.cuda.current_stream()は、デフォルトのCUDAストリームで実行することを示します。複数のカーネルを並列に起動したい場合はここでCUDAストリームを指定することになります。その他の引数は上で解説したとおりとなります。

補足

タイルサイズが2の累乗である必要があるとしました。この制限についてはNVIDIA公式のcuTile Python解説動画にて「将来的に緩和予定」とされています。

以下は個人的な推測なのですが、2の累乗の制限というのは、bitのshift演算を用いた効率的な処理を利用していることによると思われます(こちらの記事等を参照)。そのため、緩和するというのはあくまで「2の累乗以外のタイルサイズも利用できる」ようになるだけで、パフォーマンスは犠牲になり、結局「タイルサイズは原則的に2の累乗を採用すべき」となるのではないかなと推測しています。

5.2 swizzleとは何か?

ここではswizzleというテクニックを解説します。swizzleはキャッシュのヒット率を向上し、グローバルメモリからの読み取り時間を短縮するために用います。まず、行列積でキャッシュがどのように役立つのかを解説します。A\times B=Cの行列積を演算するさい、以下の式のようにA3 \times 4のタイルに、B4 \times 3のタイルに分割して処理する場合を考えます。

A=\begin{pmatrix} A_{0,0} &A_{0,1} & A_{0,2} & A_{0,3}\\ A_{1,0} &A_{1,1} & A_{1,2} & A_{1,3}\\ A_{2,0} &A_{2,1} & A_{2,2} & A_{2,3}\\ \end{pmatrix}, \quad B=\begin{pmatrix} B_{0,0} &B_{0,1} & B_{0,2}\\ B_{1,0} &B_{1,1} & B_{1,2}\\ B_{2,0} &B_{2,1} & B_{2,2}\\ B_{3,0} &B_{3,1} & B_{3,2}\\ \end{pmatrix}

この場合、C_{0,0}の計算は、行列成分のMMAと同じような要領で以下のとおり計算されます(ただし、今回は各要素自体が分割した行列であることに留意してください)

C_{0,0}= A_{0,0}*B_{0,0}+A_{0,1}*B_{1,0}+A_{0,2}*B_{2,0}+A_{0,3}*B_{3,0}

同様に、右隣のC_{0,1}の計算は以下となります。
C_{0,1}= A_{0,0}*B_{0,1}+A_{0,1}*B_{1,1}+A_{0,2}*B_{2,1}+A_{0,3}*B_{3,1}

ここで、C_{0,0}C_{0,1}の計算を見比べると、両方ともA_{0,0}A_{0,3}を用いていることがわかります。つまりC_{0,0}C_{0,1}の演算をできるだけ時間的に近接して行った場合、キャッシュヒットする可能性が高くなります。

同様に、C_{0,0}のすぐ下の成分C_{1,0}の計算は以下となります。

C_{1,0}= A_{1,0}*B_{0,0}+A_{1,1}*B_{1,0}+A_{1,2}*B_{2,0}+A_{1,3}*B_{3,0}

ここでも、C_{0,0}C_{1,0}の計算を見比べると、今度は両方ともB_{0,0}B_{0,3}を用いていることがわかり、これらの計算も時間的近接して行わせたいところです。

しかし、この横方向と縦方向の時間的近接は、ナイーブに実装してしまうと実現しません。具体的に、横方向に10成分あるタイルを考えてみましょう。

上記の赤矢印は、ナイーブな実装を行う場合のタイルの並び順を示します。この図からわかるように、横方向の計算の時間的近接はナイーブに実装しても実現しますが、縦方向については横方向に該当行を全て処理した後に次の行に移ることから、時間的近接性が実現しません。実際、GPUで扱う行列は巨大な場合がほとんどのため、上記で0番目のタイルで読み込んだ行列Bの特定のタイルは、10番目のタイル実行時にはキャッシュは忘れているでしょう。

そこで、横方向および縦方向にバランスよく演算を時間的近接させるため、下の画像のようにSwizzleというタイルの実行順の並べ替えを行います。

このような実行順にすることで、横方向に行を全て並べてから次の行に移るのではなく、縦方向と横方向にギザギザと処理を並べ、縦横方向の時間的近接性のバランスをとります。このような並べ替えをswizzleと呼びます。パラメータGROUP_SIZEは「配列を何個並べたら折り返して次の列に移るか」の設定値であり、上記の例ではGROUP_SIZE = 4としています。なお、サンプルコードでは上記の例とは異なり、より大きい値であるGROUP_SIZE = 8を採用していることに留意してください。

1次元のブロックid(bid)を上記の要領でswizzleするコードは以下となります。やや複雑ですが、上記の図のイメージを持っていれば読み解けると思います。

def swizzle_2d_from_bid(M, N, tm, tn, GROUP_SIZE_M, bid):
    """1DのブロックIDから2Dの(bidx, bidy)を計算"""
    num_bid_m = ct.cdiv(M, tm) #ceil(M/tm)と同等
    num_bid_n = ct.cdiv(N, tn) #ceil(N/tn)と同等
    num_bid_in_group = GROUP_SIZE_M * num_bid_n
    group_id = bid // num_bid_in_group
    first_bid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_bid_m - first_bid_m, GROUP_SIZE_M)
    bid_m = first_bid_m + (bid % group_size_m)
    bid_n = (bid % num_bid_in_group) // group_size_m
    return bid_m, bid_n

なお、返り値はブロックのM方向およびN方向のidとなっており、この値のペアにより、各ブロックがCのどのタイルを処理すれば良いか特定できます。

なお、前節でグリッドを1次元にした理由は、上記から理解できると思います。上の図のギザギザした順序は、2次元で順番を表現しようとすると、分岐処理が必要になりかえって複雑になります。上記の矢印の流れのように、順序の軌跡を1次元の変数で表現する方がシンプルです。

5.3 Tile Kernel

ついにカーネルに辿り着きました!といってもcuTileのカーネルはSIMTと比較すると非常にシンプルなので、あまり身構えず読み進めてもらえると思います。

まずは処理の概要を解説します。タイルプログラミングでの行列積は、これまで解説してきているように行列をタイルに分割して計算を行います。具体例として、A\times B=Cの行列積を演算するさい、以下の式のようにA3 \times 4のタイルに、B4 \times 3のタイルに分割して処理する場合を考えます。

A=\begin{pmatrix} A_{0,0} &A_{0,1} & A_{0,2} & A_{0,3}\\ A_{1,0} &A_{1,1} & A_{1,2} & A_{1,3}\\ A_{2,0} &A_{2,1} & A_{2,2} & A_{2,3}\\ \end{pmatrix}, \quad B=\begin{pmatrix} B_{0,0} &B_{0,1} & B_{0,2}\\ B_{1,0} &B_{1,1} & B_{1,2}\\ B_{2,0} &B_{2,1} & B_{2,2}\\ B_{3,0} &B_{3,1} & B_{3,2}\\ \end{pmatrix}

この場合、C_{0,0}の計算は、行列成分のMMAと同じような要領で以下のとおり計算されます。

C_{0,0}= A_{0,0}*B_{0,0}+A_{0,1}*B_{1,0}+A_{0,2}*B_{2,0}+A_{0,3}*B_{3,0}

念の為、この式のA_{0,0}等の成分(タイル)自体も行列であることに留意してください。

上記の処理をループで記載すると、C_{i,j}タイルの計算はK方向のタイル数をnum_tiles_kとして以下のようになります。

C[i,j] = [0で初期化]

for k in range(num_tiles_k):
    a = [Aからタイル(i,k)をロード]
    b = [Bからタイル(k,j)をロード]
    accumulator = [aとbのMMAを実施し、accumulatorに加算]

C[i,j] = accumulator

上記の式を実装するさいには、A_{0,0}*B_{0,0}等の個別のタイルの積の計算に対して、Tensor CoreによるMMAを行います。

以下がCの(bidx, bidy)成分のKernelの記載内容となっています(bidはblock idのことです)。上記の疑似コードを意識しながらであれば、読み取りはそれほど難しくないと思います。

# K次元のタイル数を計算
num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk))

# 累積器を初期化(float32で高精度を維持)
accumulator = ct.full((tm, tn), 0, dtype=ct.float32)
zero_pad = ct.PaddingMode.ZERO #行列の下端および右端ではタイルがはみ出るため、ゼロ埋め

# dtypeを取得(FP32の場合はTF32に変換)
dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype

# Kループ: K次元に沿って部分積を累積
for k in range(num_tiles_k):
    # 行列Aからタイルをロード
    a = ct.load(A, index=(bidx, k), shape=(tm, tk), padding_mode=zero_pad).astype(dtype)
    
    # 行列Bからタイルをロード
    b = ct.load(B, index=(k, bidy), shape=(tk, tn), padding_mode=zero_pad).astype(dtype)
    
    # 行列積を計算し累積
    accumulator = ct.mma(a, b, accumulator)

# 累積結果を出力データ型に変換
accumulator = ct.astype(accumulator, C.dtype)

# グローバルメモリにストア
ct.store(C, index=(bidx, bidy), tile=accumulator)

ct.num_tiles()は、指定した形状でタイル分割した場合のタイル数を返します。従来はceil(K / tk)のように計算が必要でしたが、cuTileではこのAPIで宣言的に取得できます。

以下、上記のカーネルに現れるAPI等の補足を行います。

ct.mma(Matrix Multiply-Accumulate)

ct.mmacuda.tile.mma)はcuTileの根幹をなすAPIです。ct.mmaは2.2節で解説したような行列積和(a @ b) + accumulatorを行います(@は行列積)。

accumulator = ct.mma(a, b, accumulator) 

演算はaccumulatorの精度で実行されます。そのため、accumulatorを高精度にしておくことで、丸め誤差の影響を緩和できます。

なお、単純な行列積のみが必要な場合はct.matmul(a, b)cuda.tile.matmul)を使用します。

tfloat32への変換

サンプルコードでは、入力行列がfloat32の場合、tfloat32に変換しています:

dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype
a = ct.load(...).astype(dtype)

tfloat32(TensorFloat-32)は、FP32と比較して仮数部のビット数を大幅に圧縮するかわりに、大幅なパフォーマンス向上が得られます。これはこのサンプルのみでなく、GPUでのモデル訓練等で一般的な戦略となっています。

タイルのロードおよびPadding

グローバルメモリに配置した行列から、指定した位置(index)および次元(shape)のタイルを読み込みます。

zero_pad = ct.PaddingMode.ZERO
a = ct.load(A, index=(bidx, k), shape=(tm, tk), padding_mode=zero_pad)

上記で特に補足すべきはpadding処理でしょう。二次元データの処理においては、5.1節で触れたように、一定のサイズのタイルをデータに敷き詰めると、右端および下端のブロックがデータからはみ出てしまい、何らか適切な処理が必要です。このような処理を境界チェックと呼びます。

SIMTの場合、データからはみ出たスレッドを実行してしまうと、行列の範囲外のデータにアクセスしてしまったり、プログラムの管理外のデータを上書きしてしまったりして非常に危険です。そのため、SIMTの場合はif文により、インデックスがデータの範囲からはみ出た場合は処理しないようにカーネルを記述する必要があり、やや面倒でした。

一方、cuTileでは範囲外のデータをパディングすることにより、境界チェックを不要にします。これは、Tensor Coreではif文を処理できないので、無害な0を流し込むことで計算処理を変更せず、処理結果を維持する戦略を取るためです。

アキュムレータによる累積、結果のストア

Cの(bidx, bidy)タイルを計算するためのアキュムレータを定義のうえ、ループ処理でMMAを実行し、結果を出力行列に書き込みます。

# 累積器を初期化(float32で高精度を維持)
accumulator = ct.full((tm, tn), 0, dtype=ct.float32)



# 累積結果を出力データ型に変換
accumulator = ct.astype(accumulator, C.dtype)

# グローバルメモリにストア
ct.store(C, index=(bidx, bidy), tile=accumulator)

ct.fullは第一引数で指定したサイズのタイルを、第二引数で指定した値および第三引数で指定した型で埋めます。
なお、ここでaccumulatorの精度指定にてct.float32をしていることに注目してください。ABがFP16やBF16等の低精度の場合でも、2.2節で解説したように累積器はFP32として丸め誤差の影響を抑えており、Tensor Coreの基本的な戦略に沿ったとなっています。

最後のct.storeはCの(bidx, bidy)タイルに対し、計算結果を書き込む操作です。これでCの(bidx, bidy)タイルの計算が完了です、めでたしめでたし。

5.4 ホストコード

ホストからカーネルを実行するコードは以下のとおりです。

def cutile_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    # 入力検証
    if A.shape[1] != B.shape[0]:
        raise ValueError(f"Incompatible matrices: K dimension mismatch")
    if not A.is_cuda or not B.is_cuda:
        raise ValueError("Input tensors must be on a CUDA device.")

    # データ型に応じてタイルサイズを決定
    if A.dtype.itemsize == 2:  # float16/bfloat16
        tm, tn, tk = 128, 256, 64  # Tensor Core向けの大きなタイル
    else:  # float32
        tm, tn, tk = 32, 32, 32    # より小さなタイル

    # 行列サイズを取得
    m, k_a = A.shape
    k_b, n = B.shape

    # グリッドサイズを計算
    grid_x = ceil(m / tm)
    grid_y = ceil(n / tn)
    grid = (grid_x * grid_y, 1, 1)

    # 出力テンソルを作成
    C = torch.empty((m, n), device=A.device, dtype=A.dtype)

    # カーネルを起動
    ct.launch(torch.cuda.current_stream(), grid, matmul_kernel, (A, B, C, tm, tn, tk))

    return C

PyTorch TensorをそのままcuTileカーネルに渡せる点が便利です。そのほか、基本的にはこれまで解説した内容になっています。

6. サンプルの実行

MatMul.pyのサンプルのうち、テストケース(TC)1〜3の3ケースを実行し、実行し時間をPyTorchrと比較してみました。各ケースは以下の内容となっており、次節で詳細に説明します。

  • TC1: 入力行列がFP16かつタイルの整数倍のサイズ
  • TC2: 入力行列がFP32かつタイルの整数倍のサイズ
  • TC3: 入力行列がFP32かつタイルの整数倍でないサイズ

なお、TC4はPersistent(常駐)な行列積を扱い複雑なため、本記事ではスキップします。

6.1 テストケースの条件

TC1〜3の条件を示します。

テストケース1

TC1は、以下のサイズのランダムなFP16をAおよびBにインプットします。

M_dim, N_dim, K_dim = 512, 512, 768

なお、上記の行列のディメンジョンはいずれも以下の各軸方向のタイルサイズ(FP16用)の倍数となっています。

tm, tn, tk = 128, 256, 64

テストケース2

TC2は、TC1のFP32番となります。FP32用の各軸方向のタイルサイズは以下のとおりとなっており、こちらも行列のディメンジョンはいずれも各軸方向のタイルサイズの倍数となります。

tm, tn, tk = 32, 32, 32

テストケース3

TC3は、以下のサイズのランダムなFP32をAおよびBにインプットします。

M_dim, N_dim, K_dim = 1_000, 500, 700

FP32用のタイルサイズはTC2と同様tm,tn,tk全て32ですが、上記の行列のディメンジョンはいずれも32の倍数ではないため、はみ出たタイル部分についてゼロ埋めを行います。

6.2 サンプルコードの修正内容

MatMul.pyでは処理時間の計測を行っていません。そのため、計測のための追記を行います。加えて、cuTileおよびPyTorchではJITコンパイルを採用していることから、特定のカーネルの初回実行時には追加のオーバーヘッドが発生します。これは一般的なGPUワークロードでは無視できる内容であることから、このオーバーヘッドの影響をウォームアップにより除去します。

上記をふまえ、TC1のcuTileカーネルの実行時間計測を例とし、修正内容を示します。「#追記: 〜」を付与している行が追記した行です。

# Perform matrix multiplication using the cuTile wrapper function.
_ = cutile_matmul(A_fp16, B_fp16) #追記: ウォームアップの実施
torch.cuda.synchronize() #追記: 計測開始前に同期
start_cuTile = time.time() #追記: 計測開始
C_fp16_cutile = cutile_matmul(A_fp16, B_fp16)
torch.cuda.synchronize()  #追記: 実行完了を待つ
time_cuTile = time.time() - start_cuTile #追記: カーネル実行時間を取得
print(f"cuTile execution time: {time_cuTile*1000:.3f} ms") #追記: カーネル実行時間を取得
print(f"cuTile Output C shape: {C_fp16_cutile.shape}, dtype: {C_fp16_cutile.dtype}")

6.3 テストケースの実行結果

TC1〜3に対し、6.2節の修正内容をcuTile実行およびPyTorch実行の両方に反映のうえ、手持ちのRTX5060Ti 16GBでの実行時間を計算してみました。

TC cuTile [ms] PyTorch [ms]
1 315 87
2 131 96
3 205 65

上記のとおり、残念ながらPyTorchの実行時間にはかないませんでした。ただし、PyTorchはcuBLASやcuTLASを利用しており、NVIDIAのNinjaたちによる各アーキテクチャに特化した最適化の恩恵を受けています。それと比較すると、この記事内で解説できるレベルの記述により、ここまで比較できるレベルの結果を得られていることは驚くべきことです(2.1節のSIMTとTileのコード記載量の比較画像を思い出してください)。特にTC2ではPyTorchの実行時間の1.4倍程度であり、かなり良い線を行っています。

上記の結果で少し気になるのが、cuTileの実行時間のうち、全く同じ行列のサイズを扱うTC1(FP16)およびTC2(FP32)で、より軽量なデータを扱うTC1の方が2倍以上遅くなっている点です。対象的に、PyTorchはTC2の方が少し遅く、直感と整合します。この原因はよくわかりませんが、何らか不具合があるのかもしれません。

まだcuTileは発表されたばかりであることをふまえると、今後の伸び代が大きいと思われます。特にタイルサイズの自動最適化に今後取り組む予定となっており、パフォーマンス向上に期待できます。

7. まとめ

本記事では、「2006年のCUDA発明以来最大の進歩」と2025年12月にNVIDIAが発表し、鳴物入りで登場したCUDA Tile、およびそのPython実装であるcuTile Pythonについて解説しました。

本記事で解説した行列積のサンプルでは、わずか数十行のシンプルなコードで、高度に最適化されたPyTorchと比較できるレベルのパフォーマンスを得られました。さらに良いことに、Tile IRはハードウェアの詳細を抽象化するため、新しいGPUアーキテクチャが登場してもコード変更なしに恩恵を受けられるとされています。

今回紐解いた行列積(MMA)はcuBLASやcuTLAS等のライブラリでも実行可能なため、cuTile Pythonでゼロから実装する意義は薄いかもしれません。しかし、GPUによる並列処理を実装したいが、cuBLAS等では対応できない処理内容の場合は、自分でカーネルを記述しなければならず、そのような場合に第一に検討すべきオプションとなるでしょう。

何か誤り等がありましたら、コメントでご指摘いただけますと幸いです。

Discussion