第12章 ポインタ・GPU・レイアウト
12.1 この章で学ぶこと
この章で学ぶ内容は次のとおりです。
学ぶこと |
ポイント |
|---|---|
ポインタの種類と選び方 |
|
アセンブリで見る alloc / init / free |
|
GPU 実行モデル |
thread → block → grid の 3 層階層 |
Layout：shape と stride の型化 |
|
12.2 Pointers：ポインタの種類と選び方
Mojo では「生アドレス」を直接使う前に、用途に合ったポインタ型を選ぶのが設計の原則です。
型 |
種類 |
主な用途 |
|---|---|---|
|
借用ポインタ |
既存の値への一時的なエイリアス。所有権を持たない |
|
所有ポインタ |
ヒープ上の値を 1 オーナーで管理。スコープ終了で自動 free |
|
参照カウントポインタ |
複数オーナーが必要な場合。カウントが 0 になったら自動 free |
|
生ポインタとしての操作 |
FFI、低レベルバッファ操作用。管理はプログラマが行う |
コラム: なぜポインタを型で分けるか
C では「ポインタはただのアドレス」です。 借用なのか所有なのか、解放責任があるのかは、コメントや慣習でしか表現できません。
Mojo はポインタの意図を型で表現します。
OwnedPointer を使えば「スコープを抜けたら自動解放」が型として保証されます。
かつては手動管理専用の UnsafePointer[T] という別の型が存在しましたが、Mojo 1.0 で Pointer[T] に統合されました。
型を分けるのではなく、unsafe_offset や unsafe_write のように 危険な操作それぞれに unsafe_ 接頭辞を付けることで、「この部分は手動管理」であることを操作単位で明示する設計に変わっています。
12.3 Pointer[T]：借用ポインタ
Pointer[T] は既存の値へのアドレスを保持するだけの借用ポインタです。
値の所有権を持たず、ヒープを確保しません。
スタック上の変数や構造体のフィールドへの参照を、ポインタとして扱いたいときに使います。
用途の典型例は次のとおりです。
大きな struct を関数に渡すとき、コピーを避けてアドレスだけ渡す
配列の先頭アドレスを保持して、ポインタ演算で要素を走査する
RAII struct の内部で、外部バッファを参照するフィールドとして持つ
def main():
var x: Int = 42
var p = Pointer(to=x) # x への借用ポインタ（所有権なし）
print(p[]) # 42
p[] = 100 # x の値をポインタ越しに書き換え
print(x) # 100
リスト12-1: pointer_borrow.mojo
Pointer(to=x) は x のアドレスを取得します。
p[] で間接参照（デリファレンス）して値を読み書きします。
p[] = 100 は x の値をポインタ越しに書き換えます。
アセンブリで何が起きているか確認します。
; var x: Int = 42 → スタックに配置
4: movq $42, 8(%rsp)
; var p = Pointer(to=x) → x のアドレスをレジスタに取得（ヒープ確保なし）
d: leaq 8(%rsp), %rax
12: movq %rax, 16(%rsp)
; print(p[]) → スタック上の x を読み出して print に渡す
17: movq 8(%rsp), %rdi
40: callq print ; 42 を表示
; p[] = 100 → スタック上の x を直接書き換え
45: movq $100, 8(%rsp)
77: callq print ; 100 を表示
操作 |
アセンブリでの意味 |
|---|---|
|
|
|
|
|
|
|
呼ばれない。ヒープ確保ゼロ |
Pointer はアドレスを保持するだけで、ヒープ確保も解放も行いません。
アセンブリには callq AlignedAlloc が存在せず、x はスタック上（8(%rsp)）に留まり続けます。
p はそのアドレスへの参照に過ぎません。
ref との関係
Pointer[T] と予約語 ref は同じ機械レベルの操作（アドレス取得 leaq）を、異なる抽象レベルで表現したものです。
項目 |
|
|
|---|---|---|
何か |
言語レベルの参照構文 |
|
使える場所 |
関数引数と戻り値の型注釈 |
変数、struct フィールド、戻り値 |
値として保存 |
できない（引数の alias） |
できる（first-class value） |
アセンブリ |
|
同じ |
origin 追跡 |
あり（借用チェッカーが検査） |
あり（ |
Pointer(to=x) は内部で x への ref を取り、それを値型として包みます。
どちらも同じ leaq 命令にコンパイルされます。
var x: Int = 42
# ref は言語構文の参照（引数の alias として使う）
# var p = Pointer(to=x) は ref を値型でラップしたもの
var p = Pointer(to=x)
ref だけでは足りない場面が Pointer[T] の存在理由です。
ref は関数の引数やローカルの alias として使えますが、struct のフィールドに保存したり別の変数に「格納」したりできません。
Pointer[T] はその制約を外し、参照を first-class の値として持ち回れるようにします。
unsafe_offset などの unsafe 操作とは異なり origin（ライフタイム）を型パラメータで追跡するため、参照元より長生きするとコンパイルエラーになります。
図12-3: Pointer[T]（借用ポインタ）のまとめ
12.4 OwnedPointer[T]：所有ポインタ
OwnedPointer[T] はヒープに値を 1 つ確保し、唯一のオーナーとして管理するポインタです。
スコープを抜けると __deinit__ が自動で呼ばれ、ヒープを解放します。
unsafe_alloc/unsafe_free を手動でペアにする代わりに、RAII を型として保証します。
OwnedPointer を選ぶ典型的な状況は次のとおりです。
スタックに置くには大きすぎる値をヒープに確保したい
関数をまたいで値を渡したいが、所有者は常に 1 つに限定したい
ヒープの手動管理をなくして確実に解放したい
from std.memory import OwnedPointer
def main():
var p = OwnedPointer[Int](42) # ヒープに 42 を確保、p が唯一のオーナー
print(p[]) # 42
# スコープを抜けると p.__del__ が自動でヒープを解放
リスト12-2: pointer_owned.mojo
OwnedPointer[Int](42) は Int 1 個分のヒープ領域を確保し、42 で初期化します。
p[] で値を読み出せます。
変数 p がスコープを抜けると、コンパイラが __deinit__ を自動で挿入し、ヒープを解放します。
アセンブリで動作を確認します。
; OwnedPointer[Int](42) → ヒープ確保 + 42 で初期化
4: movl $42, %edi
9: callq AlignedAlloc ; 8 バイト確保、戻り値 = ヒープアドレス
e: movq %rax, 16(%rsp) ; ポインタをスタックに保存
; p[] → ヒープから値を読み出して print へ
1d: movq (%rax), %rdi ; ヒープ上の 42 を読む
44: callq print
; スコープ終了 → __deinit__ が自動で解放
49: movq 16(%rsp), %rdi
4e: callq AlignedFree ; ヒープを自動解放
操作 |
アセンブリでの意味 |
|---|---|
|
|
|
|
スコープ終了時の自動解放 |
|
所有権の移転 |
|
手動での unsafe_alloc/unsafe_free との最大の違いは「解放を書き忘れられない」点です。
OwnedPointer のスコープが終わると、コンパイラが callq AlignedFree を確実に挿入します。
alloc/free の対応が型で保証されます。
図12-4: OwnedPointer[T]（所有ポインタ）のまとめ
12.5 ArcPointer[T]：参照カウントポインタ
ArcPointer[T]（Arc は Atomic Reference Count）は複数のオーナーで同じヒープ値を共有するポインタです。
ポインタをコピーするたびに内部の参照カウントがインクリメントされ、スコープを抜けるたびにデクリメントされます。
カウントが 0 になった時点で自動解放されます。
ArcPointer を選ぶ典型的な状況は次のとおりです。
同じデータを複数の変数やスレッドから参照したい
「最後に使い終わった側が解放する」という所有権を型で表現したい
解放タイミングをコンパイル時に固定できない場合
from std.memory import ArcPointer
def main():
var p1 = ArcPointer[Int](42) # 参照カウント = 1
var p2 = p1 # 参照カウント = 2（コピーで共有）
print(p1[], p2[]) # 42 42
# p1, p2 両方がスコープを抜けるとカウント 0 → 自動解放
リスト12-3: pointer_arc.mojo
var p1 = ArcPointer[Int](42) でカウント 1 のオブジェクトをヒープに確保します。
var p2 = p1 はポインタをコピーし、カウントを 2 に増やします。
p1 と p2 両方がスコープを抜けると、カウントは 1 → 0 となり、その時点でヒープが解放されます。
アセンブリで動作を確認します。
; ArcPointer[Int](42) → ヒープ確保（カウンタ領域 + 値）、カウント = 1
4: movl $42, %edi
9: callq AlignedAlloc ; カウンタ + Int のヒープ領域を確保
; var p2 = p1 → 参照カウント +1（コピー）
e: movq %rax, 24(%rsp) ; p1 をスタックへ
13: movq 24(%rsp), %rdi
18: callq <Arc::clone> ; カウントをインクリメントして p2 を作成
1d: movq %rax, 32(%rsp) ; p2 をスタックへ
; p1[], p2[] → ヒープから値を読み出して print へ
40: movq (%rcx), %rdi ; p1 が指す値（42）
43: movq (%rax), %rsi ; p2 が指す値（42）
6e: callq print
; スコープ終了 → p1.__deinit__（カウント 2→1）、p2.__deinit__（カウント 1→0 → 解放）
73: movq 24(%rsp), %rdi
78: callq <Arc::drop> ; カウント -1（まだ残る）
7d: movq 32(%rsp), %rdi
82: callq <Arc::drop> ; カウント -1 → 0 → AlignedFree
操作 |
アセンブリでの意味 |
|---|---|
|
|
|
|
|
|
スコープ終了 |
|
OwnedPointer との違いは「コピーできる」点です。
OwnedPointer はムーブしかできず、所有者は常に 1 つです。
ArcPointer はコピーのたびにカウントが増え、複数の場所から同じ値を参照できます。
その代わり、カウントの増減に原子操作（atomic）を使うため、OwnedPointer より若干コストがかかります。
図12-5: ArcPointer[T]（参照カウントポインタ）のまとめ
12.6 Pointer の unsafe 操作：最小の例とアセンブリ
Mojo 1.0 より前は、手動管理専用の UnsafePointer[T] という別の型が存在しました。
1.0 ではこの型は Pointer[T] に統合され、ヒープの確保・初期化・解放を手動で行う操作は unsafe_ 接頭辞を持つメソッド群として提供されます。
型システムによる制約がほぼないため、FFI（他言語との連携）や低レベルバッファ操作など、他のポインタでは表現しづらい処理を書くことができます。
その反面、unsafe_alloc と unsafe_free のペアが崩れるとメモリリークや二重解放が起き、Mojo コンパイラはこれを検出できません。
選択順は最後で、OwnedPointer や ArcPointer で表現できる場合はそちらを優先するべきです。
from std.memory.alloc import unsafe_alloc
def main():
var p = unsafe_alloc[Int](1)
p.unsafe_write(42)
print(p[])
p.unsafe_free()
リスト12-4: pointers_unsafe_minimal.mojo
このコードは Pointer の unsafe 操作を使った最小の例です。
4 つの操作で構成されています。
unsafe_alloc[Int](1)：Int1 個分（8 バイト）のヒープ領域を確保し、そのアドレスを返す。この時点でメモリの中身は未初期化です。p.unsafe_write(42)：未初期化のメモリに42を書き込んで初期化します。p[] = 42ではなくunsafe_writeを使う理由は、p[] = 42が「既存の値の上書き」（古い値のデストラクタを呼んでから代入）を意味するのに対し、unsafe_writeは「何もない場所への初回書き込み」だからです。未初期化領域にp[] = 42を使うと未定義動作になります。p[]：ポインタを間接参照して値を読み出します（デリファレンス）。p.unsafe_free()：ヒープ領域を解放します。unsafe_allocとunsafe_freeは必ずペアで書く必要があります。解放を忘れるとメモリリークになります。
注意：未初期化領域に p[] = value を使ってはいけない
p[] = value は「既存の値を上書きする」操作です。
内部では次の順序で動きます。
*p（現在そこにある値）のデストラクタ__deinit__を呼ぶ新しい値をコピーして書き込む
unsafe_alloc 直後のメモリは未初期化で、ゴミの bit パターンが入っています。
ステップ 1 でそのゴミを String や List[T] として解釈して __deinit__ を呼ぶと、ランダムなアドレスを解放しようとしてクラッシュまたはメモリ破壊が起きます。
Int のような trivial な型（デストラクタを持たない型）では p[] = 42 も unsafe_write(42) も同じ結果になりますが、非 trivial な型では必ず unsafe_write を使ってください。
コンパイラはこの誤りを検出しません。
アセンブリで実際に何が起きているか確認します。
注釈
実際のリスト12-4（pointers_unsafe_minimal.asm）は 900 行を超えます。
これは unsafe_alloc や print の呼び出しごとに、デフォルト引数（呼び出し位置を記録する SourceLocation など）の組み立てや、値の有無を表す Optional/Variant の判定処理がインライン展開されているためです。
この節では、4 つの操作それぞれに対応する本質的な命令だけを抜き出します。
; var p = unsafe_alloc[Int](1) → AlignedAlloc で 8 バイト確保
502: callq _KGEN_CompilerRT_AlignedAlloc
; 戻り値（ヒープアドレス）は %rax に入り、スタック上のスロットへ保存される
; p.unsafe_write(42) → ヒープ先頭に 0x2a = 42 を書き込む
14d0: movq $0x2a, (%rax)
; (%rax) = "%rax が指すアドレス" に直接書き込む
; print(p[]) → ポインタを読み出して print に渡す
159d: movq (%rax), %rdi
; (%rax) から値を読み出し、引数レジスタ %rdi に移す
; p.unsafe_free() → AlignedFree で解放
15e0: callq _KGEN_CompilerRT_AlignedFree
操作 |
アセンブリでの意味 |
|---|---|
|
|
|
|
|
|
|
|
AlignedAlloc が malloc ではなくアライメント付き確保（引数はどちらも 8：サイズとアライメント）を使うのは、SIMD 命令がアドレスの境界整合を要求するためです。
Int であれば 8 バイトアライメントが保証されます。
コラム: なぜ 4 行のコードが 900 行超のアセンブリになるのか
Mojo 1.0.0 では unsafe_alloc や print のようなジェネリック関数の呼び出し 1 回ごとに、呼び出し位置を記録する SourceLocation の構築と、それを保持するかどうかを表す Optional/Variant 型の判定処理が、最適化なしのビルドではインライン展開されて残ります。
これは Mojo 1.0 で強化されたエラーメッセージ・診断情報の充実（実行時エラーがどこで起きたかをより正確に特定できるようにする設計）の一環です。
--no-optimization を外した通常ビルドでは、これらの命令の多くはコンパイラの最適化によって削除されます。
ポインタとヒープのメモリ関係を図示します。
図12-6a: Pointer の unsafe 操作におけるスタックとヒープの関係
p 自体はスタック上にあり、ヒープのアドレスを保持するだけです。
p[]（movq (%rax), %rdi）は「%rax が指すアドレスの値を読む」という間接参照です。
スタック上の p を読んでいるのではなく、p が指す先のヒープを読んでいる点が Pointer[T]（借用ポインタとしての用法）との根本的な違いです。
図12-6: Pointer の unsafe 操作（最小の例とアセンブリ）のまとめ
12.7 複数要素の Pointer（配列風）
12.6 では 1 要素のヒープ確保を扱いました。
Pointer は unsafe_offset によって複数要素の連続領域を配列のように扱うこともできます。
C の malloc(n * sizeof(T)) + ポインタ演算と同じ考え方であり、Mojo の List や Tensor のような動的コレクションも内部ではこのパターンで実装されています。
from std.memory.alloc import unsafe_alloc
def main():
var n = 3
var p = unsafe_alloc[Int](n)
for i in range(n):
p.unsafe_offset(i).unsafe_write(i * 10)
var sum = 0
for i in range(n):
sum += p.unsafe_offset(i)[]
for i in range(n):
p.unsafe_offset(i).unsafe_deinit_pointee()
p.unsafe_free()
print(sum)
リスト12-5: unsafe_buffer_three_ints.mojo
まず unsafe_alloc[Int](n) で n × sizeof(Int) バイト（ここでは 3 × 8 = 24 バイト）の連続ブロックをヒープに確保します。
p が指すのは先頭アドレスだけです。
C で int *p = malloc(n * sizeof(int)) と書くのと本質的に同じです。
p.unsafe_offset(i) は「先頭アドレスから i 要素分だけ進んだ位置のポインタ」を返します。
Int は 8 バイトなので、p.unsafe_offset(1) は p のアドレスに 8 バイト加算したアドレスを指します。
これはアセンブリレベルでは単純な加算命令（leaq 8(%rax), %rdx など）に翻訳されます。
p.unsafe_offset(2) なら 16 バイト加算です。
かつては p + i というポインタ演算（__add__）で同じことができましたが、Mojo 1.0 では非推奨になり、unsafe_offset という名前のメソッドに置き換わりました。
「危険な操作であること」がメソッド名から一目で分かるようにする、という 1.0 全体の設計方針の一環です。
図12-7a: 連続ヒープとポインタ演算（unsafe_alloc[Int](3) = 24 バイト）
初期化には unsafe_write を使います。
12.6 で述べたとおり、確保直後のメモリは未初期化領域であり、p[] = value を使うと既存値の __deinit__ が呼ばれてクラッシュします。
ループで各要素を順に初期化しているのはこのためです。
読み出しは p.unsafe_offset(i)[] です。
unsafe_offset(i) でインデックス i の要素を指すポインタを作り、[] でデリファレンスして値を取り出します。
合計を求める sum += p.unsafe_offset(i)[] は、3 要素（0, 10, 20）を足して 30 を返します。
解放フェーズでは必ず unsafe_deinit_pointee → unsafe_free の順で行います。
unsafe_deinit_pointee は各要素の型のデストラクタ（__deinit__）を呼びます。
Int は trivial 型なので実質的に no-op ですが、String のようにヒープ上の文字列バッファを持つ型では unsafe_deinit_pointee を省略するとメモリリークになります。
その後 p.unsafe_free() でブロック全体を返却します。
コード |
意味 |
|---|---|
|
|
|
先頭から |
|
未初期化領域を値 |
|
|
|
|
|
連続ブロック全体をヒープに返却（ |
unsafe_write と unsafe_deinit_pointee のペアは、所有権を持つ型（String など）を手動管理する際の基本作法です。
unsafe_write で構築し、unsafe_deinit_pointee で破棄します。
この順序を守ることで、メモリリークや二重解放を防ぎます。
p[i] と p.unsafe_offset(i)[] の使い分け
初期化済みの要素を読み書きするだけなら、キーワード引数を使って p[unsafe_offset=i] と書くこともできます。
p[unsafe_offset=i] は p.unsafe_offset(i)[] と等価であり、C の p[i] ≡ *(p + i) と同じ関係です。
Mojo 1.0 では、位置引数での p[i] という書き方（positional __getitem__）は非推奨になり、unsafe_offset= というキーワード引数で「これは unsafe なオフセットアクセスだ」と明示する形に変わりました。
操作 |
|
|
|---|---|---|
読み出し |
|
|
書き換え（初期化済み） |
|
|
初回初期化（未初期化領域） |
|
|
デストラクタ呼び出し |
|
—（ |
unsafe_write と unsafe_deinit_pointee はポインタ自身のメソッドであるため、値を返す p[unsafe_offset=i] 構文では呼び出せません。
リスト12-5 では初回初期化と破棄が必要なため unsafe_offset(i) 形式で統一していますが、既に初期化済みの領域を読み書きするだけであれば p[unsafe_offset=i] の方が簡潔です。
図12-7: 複数要素の Pointer（配列風）のまとめ
12.7.1 RAII パターン：__deinit__ で自動解放
手動の unsafe_alloc/unsafe_free はペアが崩れるとメモリリークや二重解放につながります。
実用では __deinit__（デストラクタ）を持つ struct に Pointer の unsafe 操作をラップすることで、スコープを抜けたときに自動解放を保証できます。
from std.memory.alloc import unsafe_alloc
struct HeapInts(Writable):
var data: Pointer[Int, MutUntrackedOrigin]
var size: Int
def __init__(out self, *values: Int):
self.size = len(values)
self.data = unsafe_alloc[Int](self.size)
for i in range(self.size):
self.data.unsafe_offset(i).unsafe_write(values[i])
def __deinit__(deinit self):
for i in range(self.size):
self.data.unsafe_offset(i).unsafe_deinit_pointee()
self.data.unsafe_free()
def get(self, i: Int) -> Int:
return self.data.unsafe_offset(i)[]
def main():
var h = HeapInts(1, 2, 3)
print(h.get(0), h.get(2))
リスト12-6: reference_heapints_managed.mojo
HeapInts struct の各メソッドの役割を整理します。
__init__：unsafe_allocでヒープを確保し、可変長引数*valuesを順に書き込む。確保と初期化を__init__にまとめることで「確保済みだが未初期化」という中間状態が外から見えない__deinit__：struct がスコープを抜けるとコンパイラが自動で呼ぶデストラクタ。unsafe_deinit_pointeeで各要素のデストラクタを呼んでからunsafe_freeするget：unsafe_offset+ デリファレンス（[]）でインデックスアクセス
def main():
var h = HeapInts(1, 2, 3)
print(h.get(0), h.get(2)) # 1 3
# ここで h のスコープが終わり __deinit__ が自動呼び出し
# → unsafe_deinit_pointee × 3 → unsafe_free の順で解放される
unsafe_free を明示的に呼ぶ必要がなく、例外が発生した場合もデストラクタは必ず実行されます。
これは C++ の RAII（Resource Acquisition Is Initialization）と同じパターンです。
Pointer の unsafe 操作を直接使う場面では、このラップ構造を基本形として採用することを推奨します。
12.7.2 外部ポインタ（foreign pointer）：他言語との連携
ここまでは Mojo 自身が unsafe_alloc で確保したメモリを扱ってきました。
C や Python など他の言語とデータを交換するときは、外部コードが確保したメモリのアドレスを Pointer として受け取ることがあります。
これを 外部ポインタ（foreign pointer） と呼びます。
代表的なのは NumPy 配列との連携です。
NumPy 配列の ctypes.data は配列先頭の生アドレス（Python 側が確保したメモリ）で、unsafe_get_as_pointer で Pointer に変換すると、配列の中身を Mojo から直接読めます。
# 外部ポインタ（foreign pointer）の例。
# NumPy が確保した配列のメモリを、Mojo の Pointer 経由で直接読む。
from std.python import Python
def main() raises:
var np = Python.import_module("numpy")
var arr = np.array(Python.list(1, 2, 3, 4, 5, 6, 7, 8, 9))
# arr.ctypes.data は配列先頭の生アドレス（Python 側が所有するメモリ）。
# unsafe_get_as_pointer で Pointer に変換し、Mojo から直接読む。
var ptr = arr.ctypes.data.unsafe_get_as_pointer[DType.int64]()
for i in range(9):
print(ptr[unsafe_offset=i], end=", ")
print()
# このメモリは NumPy が所有する。Mojo 側で unsafe_free() してはいけない。
リスト12-7: foreign_pointer_numpy.mojo
実行すると、NumPy 配列の要素がそのまま表示されます。
1, 2, 3, 4, 5, 6, 7, 8, 9,
C ライブラリと連携する場合も同様で、external_call などで受け取った生アドレスを Pointer として扱います。
注意：外部で確保したメモリを unsafe_free してはいけない
外部ポインタを扱うときは、そのメモリを誰が所有しているかを意識する必要があります。
他の言語やライブラリが確保したメモリに対して unsafe_free() を呼ぶと未定義動作になります。
解放の責任は確保した側（NumPy や C ライブラリなど）にあり、Mojo 側は借りて読み書きするだけです。
自分で unsafe_alloc したメモリだけを unsafe_free する、と覚えておくと安全です。
12.8 GPU アーキテクチャ：実行モデルの 3 層
GPU とは何か
GPU（Graphics Processing Unit）はもともと画像レンダリングのために設計されたプロセッサです。 ピクセルの色計算は「同じ演算を大量のデータに対して独立して適用する」という構造を持つため、GPU は数千〜数万のコアを並列動作させるアーキテクチャに進化しました。
CPU との違いは、コアの性質にあります。 CPU は少数（4〜128 コア程度）の高性能コアを持ち、複雑な制御フロー、分岐予測、投機的実行を備えて 1 スレッドの処理を高速化することに特化しています。 対して GPU は数千の単純なコアを持ち、同じ命令を大量のデータに対して同時実行する（SIMT: Single Instruction, Multiple Threads）ことに特化しています。
図12-8: CPU と GPU のアーキテクチャの違い
この特性から、GPU は深層学習の行列演算、物理シミュレーション、画像処理など「同じ計算を大量のデータに適用する」タスクで CPU を大きく凌駕します。 Mojo の GPU サポートもこの並列性を活かすことを主目的としています。
実行モデルの階層
GPU は大量のスレッドを階層的に管理します。
単位 |
まとめ方 |
一般的なサイズ |
|---|---|---|
thread |
最小単位 |
1 スレッド = 1 要素を処理 |
warp（NVIDIA）/ wavefront（AMD） |
32 thread を束ねた SIMT 実行単位 |
同じ命令を 32 スレッドが同時実行 |
block |
複数の warp の集合。shared memory を共有 |
典型的に 64〜1024 thread |
grid |
複数の block の集合。GPU 全体に分配 |
カーネル起動時に指定 |
図12-8b: GPU 実行モデルの階層（grid → block → warp → thread）
12.9 GPU スレッド番号の計算
12.8 で見たように、GPU のスレッドは block と grid という階層で管理されています。
GPU カーネルを実行するとき、各スレッドは「自分が配列のどの要素を担当するか」を自分自身で計算する必要があります。
その計算に使うのがグローバルスレッド番号（global_idx）です。
グローバルスレッド番号は次の式で求めます。
global_idx = block_id × block_size + thread_in_block
たとえば block_size = 4 のとき、Block 0 のスレッドは 0〜3、Block 1 のスレッドは 4〜7 を担当します。
この番号を配列インデックスとして使うことで、各スレッドが重複なく分担して要素を処理できます。
def main():
var block_dim = 128
var block_id = 2
var thread_in_block = 7
var global_linear = block_id * block_dim + thread_in_block
print(global_linear)
リスト12-8: gpu_linear_thread_index.mojo
このコードは GPU カーネルそのものではなく、スレッド番号の計算ロジックを CPU 上で確認するものです。
実際のカーネルでは block_id と thread_in_block が GPU ハードウェアから取得されます。
変数 |
意味 |
|---|---|
|
Grid の中で何番目の Block か |
|
1 Block に含まれる thread 数 |
|
Block の中で何番目の thread か |
|
Grid 全体での通し番号（この値で配列要素を選択する） |
12.10 GPU fundamentals：CPU→GPU データ転送の流れ
CPU と GPU は別々のメモリ空間を持ちます。 CPU は主メモリ（DRAM）を、GPU は VRAM（HBM/GDDR）をそれぞれ独立して使い、互いのメモリを直接参照できません。 この 2 つを橋渡しするのが PCIe バス経由の DMA 転送であり、データのやり取りには明示的なコピー操作が必要です。
以下は GPU の転送サイクルを示す例です。
注釈
このコードは GPU の概念を示す例です。 実行には NVIDIA CUDA 対応の GPU 環境が必要で、本書の標準セットアップ（CPU / Apple Silicon の Metal）では実行検証していません。 GPU 関連の API は発展途上で、バージョンによってシグネチャが異なる場合があります。
from gpu.host import DeviceContext
from gpu import thread_idx, block_idx, block_dim
def double_fn(
out: Pointer[Float32],
inp: Pointer[Float32],
n: Int,
):
var i = Int(block_idx.x) * Int(block_dim.x) + Int(thread_idx.x)
if i < n:
out[i] = inp[i] * 2.0
def main():
alias N = 4
var ctx = DeviceContext() # ① GPU デバイスを初期化
var h_in = ctx.alloc_host_buffer[Float32, N]() # ② ピンメモリ確保
h_in[0] = 1.0; h_in[1] = 2.0
h_in[2] = 3.0; h_in[3] = 4.0
var d_in = ctx.alloc_device_buffer[Float32, N]() # ③ GPU メモリ確保
var d_out = ctx.alloc_device_buffer[Float32, N]()
ctx.enqueue_copy(d_in, h_in) # ④ CPU → GPU 転送
ctx.enqueue_function[double_fn]( # ⑤ カーネル起動
N, d_out.unsafe_ptr(), d_in.unsafe_ptr(), N
)
var h_out = ctx.alloc_host_buffer[Float32, N]()
ctx.enqueue_copy(h_out, d_out) # ⑥ GPU → CPU 転送
ctx.synchronize() # ⑦ 完了待機
for i in range(N):
print(h_out[i]) # 2.0 4.0 6.0 8.0
各ステップを順に見ていきます。
① DeviceContext()：GPU コンテキストの初期化
DeviceContext() は内部で CUDA Driver API を順に呼び出し、GPU との通信チャネルを確立します。
CPU 側ではこれらは通常の関数呼び出し命令（callq）に翻訳されます。
callq cuInit ; CUDA ドライバ初期化
callq cuDeviceGet ; デバイス 0 のハンドルを取得
callq cuCtxCreate ; GPU コンテキスト作成
movq %rax, ctx(%rip) ; コンテキストハンドルをスタックへ保存
② / ③ HostBuffer と DeviceBuffer：2 種類のメモリ
alloc_host_buffer と alloc_device_buffer はメモリの確保場所が異なります。
図12-10a: HostBuffer と DeviceBuffer：2 種類のメモリ
alloc_host_buffer は内部で cuMemAllocHost を呼びます。
通常の malloc と異なり、確保したページを OS がスワップアウトしないようロック（ピン）します。
DMA コントローラが物理アドレスを直接参照できるため、転送速度が向上します。
alloc_device_buffer は cuMemAlloc で GPU の VRAM を確保します。
CPU からは直接アクセスできず、ポインタとして扱えるのは GPU カーネル内だけです。
; h_in = ctx.alloc_host_buffer[Float32, 4]()
movl $16, %edi ; 4 × sizeof(Float32) = 16 バイト
callq cuMemAllocHost ; ページロックされたピンメモリ確保
movq %rax, h_in(%rip)
; d_in = ctx.alloc_device_buffer[Float32, 4]()
movl $16, %edi
callq cuMemAlloc ; GPU VRAM 確保
movq %rax, d_in(%rip)
④ enqueue_copy（PCIe DMA 転送）
ctx.enqueue_copy(d_in, h_in) は内部で cuMemcpyHtoDAsync を呼びます。
「Async」の通り、この呼び出しは即座に返ります。
実際の転送は GPU コマンドキューに積まれ、DMA エンジンが PCIe バスを経由して非同期に実行します。
; ctx.enqueue_copy(d_in, h_in)
movq d_in(%rip), %rdi ; 転送先（GPU アドレス）
movq h_in(%rip), %rsi ; 転送元（ピンメモリ）
movl $16, %edx ; バイト数
movq stream(%rip), %rcx ; GPU コマンドストリーム
callq cuMemcpyHtoDAsync ; キューに積んで即返る
PCIe Gen4 の帯域幅はおよそ 32 GB/s です。 16 バイトはマイクロ秒以下で完了しますが、大規模モデルの重みを転送する場合はこのレイテンシがボトルネックになります。
⑤ enqueue_function：カーネル起動と PTX
ctx.enqueue_function[double_fn](...) は内部で cuLaunchKernel を呼び出し、double_fn を GPU スレッド上で実行するよう命令します。
; ctx.enqueue_function[double_fn](N, ...)
movq kernel_handle(%rip), %rdi
movl $1, %esi ; gridDimX
movl $1, %edx ; gridDimY
movl $1, %ecx ; gridDimZ
movl $4, %r8d ; blockDimX（= N = 4 スレッド）
movl $1, %r9d ; blockDimY
callq cuLaunchKernel
double_fn 関数本体は Mojo コンパイラが PTX（NVIDIA の GPU 中間アセンブリ）に変換します。
PTX は x86 アセンブリと同様の形式ですが、GPU の命令セットを対象としています。
; double_fn の PTX（Mojo コンパイラ生成、概略）
.visible .entry double_fn(
.param .u64 out, .param .u64 inp, .param .u64 n)
{
; グローバルスレッド番号 i = block_idx.x * block_dim.x + thread_idx.x
mov.u32 %r_bid, %ctaid.x ; block_idx.x
mov.u32 %r_bdim, %ntid.x ; block_dim.x
mov.u32 %r_tid, %tid.x ; thread_idx.x
mad.lo.s32 %r_i, %r_bid, %r_bdim, %r_tid
; if i >= n: 終了
setp.ge.s32 %p0, %r_i, %r_n
@%p0 bra $done
; inp[i] を VRAM から読む
mul.wide.u32 %rd_off, %r_i, 4 ; i × sizeof(Float32)
add.u64 %rd_src, %rd_inp, %rd_off
ld.global.f32 %f_val, [%rd_src] ; GPU VRAM からロード
; val × 2.0 を計算して out[i] に書く
mul.f32 %f_out, %f_val, 0f40000000 ; 0f40000000 = 2.0（IEEE 754）
add.u64 %rd_dst, %rd_out, %rd_off
st.global.f32 [%rd_dst], %f_out ; GPU VRAM にストア
$done:
ret;
}
ld.global.f32 と st.global.f32 が GPU の VRAM へのアクセス命令です。
x86 の movq (%rax) に相当しますが、アクセス先が CPU の DRAM ではなく GPU の VRAM（HBM/GDDR）です。
4 スレッドが同じ PTX 命令を同時に実行し（SIMT）、それぞれ異なる %rd_off（オフセット）で別々の要素を処理します。
⑥ / ⑦ 結果の読み戻しと synchronize
enqueue_copy(h_out, d_out) は cuMemcpyDtoHAsync でキューに積まれ、synchronize() で CPU がブロックして全完了を待ちます。
; ctx.synchronize()
callq cuStreamSynchronize ; GPU のコマンドキューが空になるまで CPU をブロック
synchronize() が返った後にはじめて h_out の値が確定します。
Mojo API |
CUDA Driver API |
動作 |
|---|---|---|
|
|
GPU コンテキストを初期化 |
|
|
ページロック済みピンメモリを CPU 側に確保 |
|
|
GPU の VRAM を確保 |
|
|
PCIe DMA でホスト→デバイス転送（非同期） |
|
|
GPU カーネルをコマンドキューに投入（非同期） |
|
|
PCIe DMA でデバイス→ホスト転送（非同期） |
|
|
CPU がブロックしてキューの全完了を待機 |
enqueue_ 系の呼び出しはすべて即座に返る非同期操作です。
GPU はコマンドキューを自律的に処理し、CPU は synchronize() に達して初めてその完了を待ちます。
これにより CPU 側の処理と GPU 側の実行を重ねることができます。
図12-10: GPU fundamentals：CPU→GPU データ転送の流れ
コラム: Apple Silicon（M1/M2）における GPU プログラミング
NVIDIA GPU（PCIe 接続）と Apple Silicon（M1/M2）では、CPU↔GPU のデータ転送コストが根本的に異なります。 最大の違いは UMA（Unified Memory Architecture）です。
PCIe 接続型 GPU（NVIDIA 等）
CPU の DRAM と GPU の VRAM は物理的に別チップ
データは PCIe バス（Gen4: 〜32 GB/s）を経由してコピーが必要
cuMemcpyHtoDAsyncが実際に DMA 転送を発生させる
Apple Silicon（M1/M2 等）
CPU と GPU が同じシリコンダイ上にあり、同じ物理メモリプールを共有
「転送」は TLB エントリの追加に過ぎず、実データのコピーが不要
.storageModeSharedバッファは CPU からも GPU からも直接アクセス可能
Apple Silicon では HostBuffer/DeviceBuffer の区別がなく、共有バッファ 1 本で済みます。
GPU フレームワークは CUDA ではなく Metal です。
カーネルは Metal Shading Language（C++ ベース）で記述します。
// double.metal
#include <metal_stdlib>
using namespace metal;
kernel void double_fn(
device float* out [[buffer(0)]],
device const float* inp [[buffer(1)]],
uint i [[thread_position_in_grid]])
{
out[i] = inp[i] * 2.0;
}
PTX の ld.global.f32 に相当するのが device float* です。
[[thread_position_in_grid]] が CUDA の blockIdx.x * blockDim.x + threadIdx.x に相当するグローバルスレッド番号です。
ホスト側は Swift + Metal API で記述します。
import Metal
let device = MTLCreateSystemDefaultDevice()!
let queue = device.makeCommandQueue()!
// 共有バッファ（CPU・GPU 両方から直接アクセス可能）
var data: [Float] = [1.0, 2.0, 3.0, 4.0]
let buf = device.makeBuffer(bytes: &data, length: 16,
options: .storageModeShared)!
let lib = device.makeDefaultLibrary()!
let fn = try! lib.makeFunction(name: "double_fn")!
let pipeline = try! device.makeComputePipelineState(function: fn)
let cmd = queue.makeCommandBuffer()!
let enc = cmd.makeComputeCommandEncoder()!
enc.setComputePipelineState(pipeline)
enc.setBuffer(buf, offset: 0, index: 1) // inp
enc.setBuffer(buf, offset: 0, index: 0) // out（in-place 上書き）
enc.dispatchThreads(MTLSize(width: 4, height: 1, depth: 1),
threadsPerThreadgroup: MTLSize(width: 4, height: 1, depth: 1))
enc.endEncoding()
cmd.commit()
cmd.waitUntilCompleted()
// buf のメモリをそのまま CPU から読める
let result = buf.contents().bindMemory(to: Float.self, capacity: 4)
print(result[0], result[1], result[2], result[3]) // 2.0 4.0 6.0 8.0
コラム（続き）: CUDA と Metal API の対応
CUDA と Metal の対応関係は以下のとおりです。
NVIDIA / CUDA |
Apple Silicon / Metal |
備考 |
|---|---|---|
|
不要（UMA） |
データコピー自体が不要 |
|
|
共有バッファ 1 本 |
|
不要（UMA） |
ポインタ共有のみ |
|
|
カーネル投入 |
|
|
CPU 側同期 |
PTX（中間表現） |
AIR（Apple IR → Metal bytecode） |
コンパイル後の中間形式 |
Mojo 1.0.0 時点では GPU バックエンドは主に NVIDIA（CUDA/PTX）を対象としており、Apple Silicon GPU への直接サポートは将来のロードマップに含まれています。 現時点で M1/M2 上の GPU を使う場合は Metal API を直接呼び出すか、MAX Engine の CPU パスを利用します。
12.11 GPU block and warp：タイル化
なぜタイル化が必要か
GPU のメモリには速度と容量が異なる 3 つの層があります。
図12-11a: GPU メモリ階層（速い順）
問題は、単純に「各スレッドが VRAM から直接読む」と同じデータを何度も読むことになる点です。
たとえば行列乗算 C = A × B では、A の 1 行を B の全列が共通して参照します。
スレッドごとに VRAM から読み直すと、n スレッド分の読み込みが発生します。
タイル化（Tiling）はこれを解消します。 block 内の全スレッドが協調して VRAM から小さなかたまり（タイル）を一度だけ shared memory にロードし、その後は高速な shared memory から読みます。
図12-11b: タイル化なし vs タイル化あり：VRAM の重複読み込みを Shared Memory で解消
タイル分割のしくみ
大きな行列を tile_size × tile_size の小さなかたまりに分割します。
図12-11c: タイル分割：16×16 行列をタイルサイズ 4 で分割
CPU 側でこの分割ループの構造を確認します。
def main():
var n = 16
var tile = 4
var count = 0
for bi in range(0, n, tile):
for bj in range(0, n, tile):
var _ = bi + bj
count += 1
print(count)
リスト12-9: gpu_tile_loop_nest.mojo
range(0, n, tile) で 0, 4, 8, 12 と tile ステップで進み、タイルの左上角 (bi, bj) を列挙します。
count は 4 × 4 = 16 になります。
実際の GPU カーネルでは、この各ループ反復が 1 つの block に対応します。
GPU カーネルでの実装
実際の GPU カーネルでのタイル化は、shared memory への協調ロードと barrier の 2 ステップで構成されます。
注釈
以下のカーネルは GPU 上のタイル化の概念を示す例です。 実行には NVIDIA CUDA 対応の GPU 環境が必要で、本書の標準セットアップでは実行検証していません。 GPU 関連の API は発展途上で、バージョンによってシグネチャが異なる場合があります。
# 行列乗算カーネルのタイル化（Mojo 1.0.0・概略）
alias TILE = 16
def matmul_tiled(
C: Pointer[Float32],
A: Pointer[Float32],
B: Pointer[Float32],
N: Int,
):
# shared memory：block 内の全スレッドが共有する高速領域
var tile_a = stack_allocation[TILE * TILE, Float32]()
var tile_b = stack_allocation[TILE * TILE, Float32]()
var row = Int(block_idx.x) * TILE + Int(thread_idx.x)
var col = Int(block_idx.y) * TILE + Int(thread_idx.y)
var local_r = Int(thread_idx.x)
var local_c = Int(thread_idx.y)
var acc: Float32 = 0.0
# タイル数だけ反復（K 方向を TILE ずつ分割）
for t in range(0, N // TILE):
# ① VRAM → shared memory への協調ロード
tile_a[local_r * TILE + local_c] = A[row * N + t * TILE + local_c]
tile_b[local_r * TILE + local_c] = B[(t * TILE + local_r) * N + col]
# ② barrier：全スレッドのロード完了を待つ
barrier()
# ③ shared memory から積和演算（高速）
for k in range(TILE):
acc += tile_a[local_r * TILE + k] * tile_b[k * TILE + local_c]
# ④ 次タイルのロード前に再度 barrier
barrier()
C[row * N + col] = acc
barrier の役割
barrier()（CUDA では __syncthreads()）は block 内の全スレッドが同じ地点に達するまで待機させる命令です。
図12-11d: barrier が 2 回必要な理由
barrier なしでは、まだロードが終わっていない tile_a を別のスレッドが読み出し、誤った値で計算してしまいます。
PTX ではこの命令は bar.sync 0; と表現されます。
; PTX の barrier 命令
bar.sync 0; ; block 内の全スレッドがここに到達するまでストール
概念 |
意味 |
|---|---|
タイル化 |
大きな配列を |
shared memory |
block 内の全スレッドが共有する高速メモリ（48KB 程度）。VRAM の約 100 倍速い |
協調ロード |
block 内の各スレッドが 1 要素ずつ担当して tile を shared memory へ転送 |
barrier |
|
タイルサイズの選択 |
典型的に 16×16 または 32×32。shared memory に収まり、warp（32 スレッド）の倍数が望ましい |
12.12 Layouts：データの並び方を型で表す
なぜ Layout が必要か
多次元配列はメモリ上では 1 次元のフラットな列として格納されます。
A[i][j] という論理インデックスをメモリオフセットに変換する計算は、従来は手動で行っていました。
# Layout なし：手動でオフセット計算（cols=3 を知っていないと書けない）
var A = unsafe_alloc[Float32](4 * 3)
A.unsafe_offset(i * 3 + j).unsafe_write(value)
この方法では cols=3 という知識がコードに散らばり、行列のサイズが変わるたびに全箇所を修正する必要があります。
Layout はこの「論理インデックス → メモリオフセット」の写像を型として閉じ込めます。
# Layout あり：オフセット計算が Layout に封じ込められている
comptime layout = Layout.row_major(4, 3)
# layout(i, j) が i×3 + j を返す
shape と stride：2 つの数列が Layout を決める
Layout は shape（各次元のサイズ）と stride（次の要素へ進む要素数）の 2 つの数列で定義されます。
shape = (行数, 列数) 例: (4, 3)
stride = (行方向の歩幅, 列方向の歩幅)
オフセット計算式:
offset = index[0] × stride[0] + index[1] × stride[1]
= i × stride_row + j × stride_col
stride の値がメモリ配置を決めます。
同じ shape = (4, 3) でも stride が違えばメモリ上の並びが変わります。
row-major と col-major
row-major（行優先）は C 言語、NumPy、Mojo のデフォルト配置です。 行の要素がメモリ上で連続します。
図12-12a: row-major（行優先）：shape=(4,3), stride=(3,1)
col-major（列優先）は Fortran、BLAS、OpenGL のデフォルト配置です。 列の要素がメモリ上で連続します。
図12-12b: col-major（列優先）：shape=(4,3), stride=(1,4)
同じ論理行列でも stride が違えば同じ [1,2] が指す物理アドレスが異なります。
BLAS ライブラリに行列を渡すとき、row-major か col-major かを正しく伝えないと計算結果が壊れるのはこのためです。
stride トリック：コピーなしの転置
stride を入れ替えるだけで、データを一切コピーせず転置ビューを作れます。
図12-12c: stride トリック：コピーなしの転置
転置の offset(i, j) = i × 1 + j × 3 になります。
NumPy の A.T もこの仕組みです。
コードで確認
from layout import Layout
from layout.int_tuple import IntTuple
def main():
var shape = IntTuple(3, 4)
var l = Layout.row_major(shape)
print(l.rank())
リスト12-10: layout_row_major_shape.mojo
Layout.row_major(shape) は shape = (3, 4) から stride を自動計算して stride = (4, 1) の Layout を構築します。
rank() は次元数（ここでは 2）を返します。
Layout 自体はデータを持たず、オフセット計算のルールだけを保持します。
概念 |
定義 |
具体例（4×3 行列） |
|---|---|---|
|
各次元の要素数 |
|
|
|
|
|
|
|
|
行要素が連続。C、NumPy、Mojo のデフォルト |
GPU で行方向アクセス時に coalesced（効率的） |
|
列要素が連続。Fortran、BLAS の慣例 |
GPU で列方向アクセス時に coalesced |
stride トリック |
stride を変えるだけで転置ビューを作れる |
コピーなし。 |
12.13 LayoutTensor：レイアウト情報つきの多次元ビュー
Layout と LayoutTensor の違い
12.12 で見た Layout はオフセット計算のルールだけを持ち、実データを持ちません。
LayoutTensor はそこに実際のデータ領域（ストレージ）を結合したものです。
Layout = shape + stride（オフセット計算のルール）
LayoutTensor = Layout + 実データ（Pointer または Array）
─────────────────────────────────────────────────────
「どう並ぶか」 + 「実際のデータ」= 多次元ビュー
Python の NumPy ndarray に相当しますが、layout がコンパイル時定数（comptime）であれば、インデックス→オフセットの計算がコンパイル時に解決されるゼロコスト抽象になります。
コードを読む
from layout import Layout, LayoutTensor
def main():
comptime layout = Layout.row_major(2, 3)
var storage = Array[Float32, 2 * 3](uninitialized=True)
var t = LayoutTensor[DType.float32, layout](storage)
t[0, 0] = 1.0
t[1, 2] = 4.0
print(t[0, 0], t[1, 2])
リスト12-11: layout_tensor_small.mojo
各行を順に追います。
① comptime layout = Layout.row_major(2, 3)
comptime キーワードにより、layout はコンパイル時定数になります。
Layout.row_major(2, 3) は shape = (2, 3)、stride = (3, 1) の Layout を構築します。
② var storage = Array[Float32, 2 * 3](uninitialized=True)
スタック上に Float32 6 個分（24 バイト）の連続領域を確保します。
ヒープ確保は行われません。
uninitialized=True により初期値のゼロ埋めをスキップします。
③ var t = LayoutTensor[DType.float32, layout](storage)
storage を layout の規則で解釈する多次元ビューを作ります。
LayoutTensor 自体はデータのコピーを行わず、storage へのポインタと layout の情報だけを保持します。
④ t[0, 0] = 1.0 と t[1, 2] = 4.0
論理インデックスをオフセットに変換して書き込みます。
layout が comptime なので、オフセット計算はコンパイル時に済んでいます。
t[0, 0] → offset = 0×3 + 0×1 = 0 → storage[0]（= 1.0）
t[1, 2] → offset = 1×3 + 2×1 = 5 → storage[5]（= 4.0）
storage のメモリレイアウト（アドレス順）：
[0] [1] [2] [3] [4] [5]
1.0 0.0 0.0 0.0 0.0 4.0
↑ ↑
t[0,0] t[1,2]
アセンブリではオフセット計算が定数に畳み込まれます。
; t[1, 2] = 4.0 の場合（layout が comptime のため offset=5 は定数）
; address = storage_base + 5 × sizeof(Float32) = storage_base + 20
movss %xmm0, 20(%rsp) ; スタック上の storage+20 に直接ストア
i × stride_row + j × stride_col の乗算、加算がアセンブリから消えています。
comptime layout によるゼロコストの証拠です。
tile / vectorize / distribute
LayoutTensor が提供する主要な変換操作を説明します。
tile(size)：テンソルをタイルサイズで切り出したサブビューを返します。
データのコピーなしに「この LayoutTensor の中の (row, col) 番目のタイル」という参照を作ります。
12.11 のタイル化を LayoutTensor レベルで表現するための操作です。
# 8×8 テンソルを 4×4 タイルに分割して処理
for tile_row in range(0, 8, 4):
for tile_col in range(0, 8, 4):
var sub = t.tile[4, 4](tile_row, tile_col)
# sub は 4×4 の LayoutTensor（コピーなし）
vectorize(width)：SIMD 幅を指定して水平方向の要素を一括ロード／ストアします。
CPU の AVX2 では width=8（Float32 × 8 = 256 bit）が典型的です。
# SIMD 幅 8 でロード：8 要素を一命令で読む
var vec = t.vectorize[8](row, col_start) # SIMD レジスタに 8 個ロード
distribute(threads)：スレッド数で分割し、各スレッドが担当するサブビューを返します。
GPU カーネル内でスレッド番号に応じた担当範囲を得る操作です。
# 4 スレッドで 8 要素を分担（各スレッドが 2 要素を担当）
var my_part = t.distribute[4](thread_id)
要素 / 操作 |
説明 |
|---|---|
|
Layout とストレージを結合した多次元ビュー。コピーなし |
|
インデックス→オフセット計算をコンパイル時解決。実行時コストゼロ |
|
|
|
指定サイズのサブビューを取得（データコピーなし） |
|
SIMD 幅で水平方向を一括ロードとストア |
|
スレッド数で分割し、担当サブビューを返す |
コラム: CPU SIMD と GPU の共通抽象
LayoutTensor は CPU の SIMD 演算と GPU カーネルの両方で同じ抽象として使えます。
「データがどう並んでいるか」を Layout で型に閉じ込めることで、演算コードから並び方の詳細が切り離されます。
これにより、「row_major で CPU SIMD ループ」と「row_major で GPU カーネル」を同じ tile / vectorize インターフェースで表現できます。
これは第 3 部で実装する MicroGPT の行列演算の基盤でもあります。
図12-13: LayoutTensor：レイアウト情報つきの多次元ビュー
12.14 まとめ
本章で扱った内容を以下にまとめます。
要素 |
押さえるべきこと |
|---|---|
ポインタの選択順 |
|
|
|
GPU の 3 層 |
thread → block → grid。 |
タイル化 |
大きな行列を shared memory に収まる単位に分割して処理 |
Layout |
shape（形）と stride（ステップ）を型で表す |
LayoutTensor |
Layout つきの多次元ビュー。 |