12. ポインタ・GPU・レイアウト
12.1. この章で学ぶこと
ポインタの基本と
unsafeの考え方GPU 実行モデルの入口
LayoutとLayoutTensorの役割
この章のテーマは、「安全な高級 API を優先し、必要なときだけ低水準へ降りる」 という設計の軸を理解することです。
ポインタは
Pointer(借用)やOwnedPointer(所有)から始め、UnsafePointerは最後の手段GPU も
DeviceContext/LayoutTensorの高級抽象から入るLayout は shape と stride を型で表した抽象で、直接バイト列を触る手前の層
今すぐ GPU や数値計算が必要でなければ、最初は見出しだけ追っても大丈夫です。
12.2. Pointers
この節でいちばん大事なのは、Mojo ではポインタも用途ごとに型を分けて扱う ことです。
12.2.1. どんなポインタがあるか
Mojo には、たとえば次のような種類があります。
Pointer:借用OwnedPointer:所有ArcPointer:共有UnsafePointer:低水準操作用
ここで大事なのは、ただの生アドレスとして扱うのではなく、意味を型で分ける ことです。アライメント・ストライドの話とも、後述の GPU や layout ともつながります。
12.2.2. 基本の考え方
まずは安全な API を優先する
生ポインタに近い操作は
unsafeに閉じ込める借用・所有・共有を意識して使い分ける
つまり、便利だからすぐ UnsafePointer を使う、という流れではない ということです。
12.2.3. 最小の例
次の例は、UnsafePointer に近い操作を最小限で示したものです。実アプリでは安全なラッパーを優先し、こうした操作は unsafe ブロックやレビュー体制のもとに置く、という前提で読んでください。
def main():
var p = alloc[Int](1)
p.init_pointee_copy(42)
print(p[])
p.free()
リスト-1: pointers_unsafe_minimal.mojo
この例では、
alloc[Int](1)で 1 要素分の領域を確保するinit_pointee_copyで値を書き込む[]で読み出す最後に
freeで解放する
という流れになっています。
つまり、確保して、使って、最後に返す という順番です。Pointer / OwnedPointer など別の型は、マニュアルの API に沿って より制約のはっきりした借用・所有を表せます。
このプログラムをコンパイルすると以下のようになります。
pointers_unsafe_minimal.o:
(__TEXT,__text) section
_pointers_unsafe_minimal::main():
; def main():
0: ff 03 06 d1 sub sp, sp, #0x180
4: fc 6f 16 a9 stp x28, x27, [sp, #0x160]
8: fd 7b 17 a9 stp x29, x30, [sp, #0x170]
; var p = alloc[Int](1)
c: 1f 20 03 d5 nop
10: eb a3 00 91 add x11, sp, #0x28
14: c8 00 80 52 mov w8, #0x6
18: e8 1b 00 f9 str x8, [sp, #0x30]
1c: 08 00 00 90 adrp x8, _static_string_aa50425c37b3053a@PAGE
20: 08 01 00 91 add x8, x8, _static_string_aa50425c37b3053a@PAGEOFF
24: e9 03 08 aa mov x9, x8
28: e9 23 00 f9 str x9, [sp, #0x40]
2c: e9 03 08 aa mov x9, x8
30: e9 27 00 f9 str x9, [sp, #0x48]
34: e8 17 00 f9 str x8, [sp, #0x28]
38: 08 00 e4 d2 mov x8, #0x2000000000000000
3c: e8 1f 00 f9 str x8, [sp, #0x38]
40: e9 43 01 91 add x9, sp, #0x50
44: 0a 04 80 52 mov w10, #0x20
48: ea 2f 00 f9 str x10, [sp, #0x58]
4c: 0a 00 00 90 adrp x10, _static_string_d03f408500ce58ad@PAGE
50: 4a 01 00 91 add x10, x10, _static_string_d03f408500ce58ad@PAGEOFF
54: ec 03 0a aa mov x12, x10
58: ec 37 00 f9 str x12, [sp, #0x68]
5c: ec 03 0a aa mov x12, x10
60: ec 3b 00 f9 str x12, [sp, #0x70]
64: ea 2b 00 f9 str x10, [sp, #0x50]
68: e8 33 00 f9 str x8, [sp, #0x60]
6c: 0d 00 00 90 adrp x13, _static_string_bb8091a8286e9abe@PAGE
70: ad 01 00 91 add x13, x13, _static_string_bb8091a8286e9abe@PAGEOFF
74: e8 03 0d aa mov x8, x13
78: e8 3f 00 f9 str x8, [sp, #0x78]
7c: 68 00 80 52 mov w8, #0x3
80: ec 03 08 aa mov x12, x8
84: ec 43 00 f9 str x12, [sp, #0x80]
88: 28 00 80 52 mov w8, #0x1
8c: ea 03 08 aa mov x10, x8
90: ea 47 00 f9 str x10, [sp, #0x88]
94: 1f 20 03 d5 nop
98: e8 43 02 91 add x8, sp, #0x90
9c: ea 4b 00 f9 str x10, [sp, #0x90]
a0: ea 63 02 91 add x10, sp, #0x98
a4: ed 4f 00 f9 str x13, [sp, #0x98]
a8: ec 53 00 f9 str x12, [sp, #0xa0]
ac: eb 7b 00 f9 str x11, [sp, #0xf0]
b0: ea 7f 00 f9 str x10, [sp, #0xf8]
b4: e9 83 00 f9 str x9, [sp, #0x100]
b8: e8 87 00 f9 str x8, [sp, #0x108]
bc: e8 1f 40 f9 ldr x8, [sp, #0x38]
c0: 88 05 f0 b6 tbz x8, #0x3e, 0x170
c4: 01 00 00 14 b 0xc8
c8: e8 17 40 f9 ldr x8, [sp, #0x28]
cc: 1f 20 03 d5 nop
d0: e8 8b 00 f9 str x8, [sp, #0x110]
d4: 09 01 80 52 mov w9, #0x8
d8: e9 8f 00 f9 str x9, [sp, #0x118]
dc: 08 21 00 f1 subs x8, x8, #0x8
e0: e8 13 00 f9 str x8, [sp, #0x20]
e4: e9 03 08 aa mov x9, x8
e8: e9 57 00 f9 str x9, [sp, #0xa8]
ec: e8 5b 00 f9 str x8, [sp, #0xb0]
f0: 1f 20 03 d5 nop
f4: 01 00 00 14 b 0xf8
f8: e9 13 40 f9 ldr x9, [sp, #0x20]
fc: 28 00 80 52 mov w8, #0x1
100: e8 93 00 f9 str x8, [sp, #0x120]
104: ea 03 09 aa mov x10, x9
108: ea 5f 00 f9 str x10, [sp, #0xb8]
10c: ea 03 09 aa mov x10, x9
110: ea 63 00 f9 str x10, [sp, #0xc0]
114: e8 03 08 cb neg x8, x8
118: .long 0xf8680128
11c: e8 97 00 f9 str x8, [sp, #0x128]
120: e8 0f 00 f9 str x8, [sp, #0x18]
124: 01 00 00 14 b 0x128
128: e8 0f 40 f9 ldr x8, [sp, #0x18]
12c: 08 05 00 f1 subs x8, x8, #0x1
130: c1 01 00 54 b.ne 0x168
134: 01 00 00 14 b 0x138
138: 01 00 00 14 b 0x13c
13c: bf 39 03 d5 dmb ishld
140: 01 00 00 14 b 0x144
144: e0 13 40 f9 ldr x0, [sp, #0x20]
148: 1f 20 03 d5 nop
14c: e8 03 00 aa mov x8, x0
150: e8 9b 00 f9 str x8, [sp, #0x130]
154: 1f 20 03 d5 nop
158: e8 03 00 aa mov x8, x0
15c: e8 9f 00 f9 str x8, [sp, #0x138]
160: 00 00 00 94 bl _KGEN_CompilerRT_AlignedFree
164: 02 00 00 14 b 0x16c
168: 01 00 00 14 b 0x16c
16c: 02 00 00 14 b 0x174
170: 01 00 00 14 b 0x174
174: e8 33 40 f9 ldr x8, [sp, #0x60]
178: 88 05 f0 b6 tbz x8, #0x3e, 0x228
17c: 01 00 00 14 b 0x180
180: e8 2b 40 f9 ldr x8, [sp, #0x50]
184: 1f 20 03 d5 nop
188: e8 8b 00 f9 str x8, [sp, #0x110]
18c: 09 01 80 52 mov w9, #0x8
190: e9 8f 00 f9 str x9, [sp, #0x118]
194: 08 21 00 f1 subs x8, x8, #0x8
198: e8 0b 00 f9 str x8, [sp, #0x10]
19c: e9 03 08 aa mov x9, x8
1a0: e9 67 00 f9 str x9, [sp, #0xc8]
1a4: e8 6b 00 f9 str x8, [sp, #0xd0]
1a8: 1f 20 03 d5 nop
1ac: 01 00 00 14 b 0x1b0
1b0: e9 0b 40 f9 ldr x9, [sp, #0x10]
1b4: 28 00 80 52 mov w8, #0x1
1b8: e8 93 00 f9 str x8, [sp, #0x120]
1bc: ea 03 09 aa mov x10, x9
1c0: ea 6f 00 f9 str x10, [sp, #0xd8]
1c4: ea 03 09 aa mov x10, x9
1c8: ea 73 00 f9 str x10, [sp, #0xe0]
1cc: e8 03 08 cb neg x8, x8
1d0: .long 0xf8680128
1d4: e8 97 00 f9 str x8, [sp, #0x128]
1d8: e8 07 00 f9 str x8, [sp, #0x8]
1dc: 01 00 00 14 b 0x1e0
1e0: e8 07 40 f9 ldr x8, [sp, #0x8]
1e4: 08 05 00 f1 subs x8, x8, #0x1
1e8: c1 01 00 54 b.ne 0x220
1ec: 01 00 00 14 b 0x1f0
1f0: 01 00 00 14 b 0x1f4
1f4: bf 39 03 d5 dmb ishld
1f8: 01 00 00 14 b 0x1fc
1fc: e0 0b 40 f9 ldr x0, [sp, #0x10]
200: 1f 20 03 d5 nop
204: e8 03 00 aa mov x8, x0
208: e8 9b 00 f9 str x8, [sp, #0x130]
20c: 1f 20 03 d5 nop
210: e8 03 00 aa mov x8, x0
214: e8 9f 00 f9 str x8, [sp, #0x138]
218: 00 00 00 94 bl _KGEN_CompilerRT_AlignedFree
21c: 02 00 00 14 b 0x224
220: 01 00 00 14 b 0x224
224: 02 00 00 14 b 0x22c
228: 01 00 00 14 b 0x22c
22c: 1f 20 03 d5 nop
230: 08 01 80 52 mov w8, #0x8
234: e1 03 08 aa mov x1, x8
238: e0 03 01 aa mov x0, x1
23c: 00 00 00 94 bl _KGEN_CompilerRT_AlignedAlloc
240: e8 03 00 aa mov x8, x0
244: e8 03 00 f9 str x8, [sp]
248: e0 03 08 aa mov x0, x8
24c: e0 77 00 f9 str x0, [sp, #0xe8]
250: e0 03 08 aa mov x0, x8
254: e0 a3 00 f9 str x0, [sp, #0x140]
; p.init_pointee_copy(42)
258: 1f 20 03 d5 nop
25c: e0 03 08 aa mov x0, x8
260: e0 a7 00 f9 str x0, [sp, #0x148]
264: 1f 20 03 d5 nop
268: 49 05 80 52 mov w9, #0x2a
26c: e0 03 09 aa mov x0, x9
270: 00 01 00 f9 str x0, [x8]
; print(p[])
274: 01 00 00 90 adrp x1, _static_string_a8d4ace0dc8d360e@PAGE
278: 21 00 00 91 add x1, x1, _static_string_a8d4ace0dc8d360e@PAGEOFF
27c: 28 00 80 52 mov w8, #0x1
280: e6 03 08 aa mov x6, x8
284: e2 03 06 aa mov x2, x6
288: 03 00 00 90 adrp x3, _static_string_bbe01a6a523daf15@PAGE
28c: 63 00 00 91 add x3, x3, _static_string_bbe01a6a523daf15@PAGEOFF
290: e4 03 06 aa mov x4, x6
294: 08 00 80 52 mov w8, #0x0
298: 05 01 00 12 and w5, w8, #0x1
29c: 00 00 00 94 bl "_std::io::io::print[*::Writable](*$0,sep:::StringSlice[::Bool(False), StaticConstantOrigin, *?],end:::StringSlice[::Bool(False), StaticConstantOrigin, *?],flush:::Bool,file:::FileDescriptor$),Ts=[[typevalue<#kgen.instref<\"std::builtin::int::Int\">>, index]]"
2a0: e0 03 40 f9 ldr x0, [sp]
; p.free()
2a4: 1f 20 03 d5 nop
2a8: e8 03 00 aa mov x8, x0
2ac: e8 ab 00 f9 str x8, [sp, #0x150]
2b0: 1f 20 03 d5 nop
2b4: e8 03 00 aa mov x8, x0
2b8: e8 af 00 f9 str x8, [sp, #0x158]
2bc: 00 00 00 94 bl _KGEN_CompilerRT_AlignedFree
; def main():
2c0: fd 7b 57 a9 ldp x29, x30, [sp, #0x170]
2c4: fc 6f 56 a9 ldp x28, x27, [sp, #0x160]
2c8: ff 03 06 91 add sp, sp, #0x180
2cc: c0 03 5f d6 ret
リスト-2: pointers_unsafe_minimal.asm
詳細: Pointers
補足: C に近い感覚がありますが、Mojo では意味ごとにポインタ型がわかれています。
12.3. Unsafe pointers
この節は、本当に低水準なメモリ操作をする場面 の話です。
12.3.1. 何をするのか
UnsafePointer では、
alloc/freeで領域を管理するinit_pointeeなどで中身を初期化する必要なら自分で後始末する
という流れを自分で扱います。
12.3.2. なぜ注意が必要か
この領域では、コンパイラの安全保証が弱くなります。
そのため、unsafe ブロックで ここは自分で責任を持つ操作だ と明示します。
つまり、unsafe は「速そうだから使う」ものではなく、
安全な方法では足りないときの最後の手段 です。FFI・低水準バッファで必要になる 最後の手段として位置づけます。
12.3.3. 例
複数要素ぶんを連続確保し、要素ごとに初期化してから読み、destroy_pointee と free で後始末する流れです。+ i でポインタを進めている部分が、配列のように見えるメモリを C 風に扱うイメージです。
次の例では、複数要素ぶんの領域を連続で確保して、要素ごとに初期化してから読み出しています。
def main():
var n = 3
var p = alloc[Int](n)
for i in range(n):
(p + i).init_pointee_copy(i * 10)
var sum = 0
for i in range(n):
sum += (p + i)[]
for i in range(n):
(p + i).destroy_pointee()
p.free()
print(sum)
リスト-3: unsafe_buffer_three_ints.mojo
ここで大事なのは、
初期化する
使う
destroy_pointeeとfreeで片づける
という対応関係です。init_pointee_copy と destroy_pointee を対にすることで、型のデストラクタが必要なオブジェクトを置く場合にも拡張しやすいパターンです。
実際の unsafe ブロックで囲むかどうかは、呼び出しコンテキストとマニュアルの推薦に沿ってください。
詳細: Unsafe pointers
出典: Mojo Manual — unsafe-pointers
補足: この領域は、実務ではレビュー前提で扱うと考えるのが安全です。
12.4. GPU architecture
ここからは GPU の基本です。
この節でまず押さえたいのは、GPU はたくさんのスレッドを同時に動かす前提で設計されている ことです。
12.4.1. よく出てくる単語
SM
warp
SIMT
grid
block
thread
最初は全部を厳密に覚えなくて大丈夫です。 まずは、thread が集まって block になり、block が集まって grid になる と捉えるとよいでしょう。
12.4.2. 何が大事か
GPU では、多くの処理を小さな単位に分けて、一気に並列実行します。 そのとき、自分が何番目のスレッドなのか がよく重要になります。
1 次元では、たとえば次のように考えます。
block_id * block_dim + thread_in_block
これは、グローバルなスレッド番号の典型的な考え方です。
12.4.3. 例
次の例は、GPU 上で実行するコードではなく、 その番号計算の考え方だけを CPU 上で確認するためのものです。GPU では スレッドが大量に並列に動き、ブロックや グリッドといった単位でまとめられます(実カーネルやデバイス API は含みません。環境依存が大きいため、実行モデルの数式のイメージ用です)。
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)
リスト-4: gpu_linear_thread_index.mojo
多次元グリッドやワープ単位の動きは、マニュアルの図とセットで見るのが確実です。
詳細: GPU architecture
出典: Mojo Manual — gpu/architecture
補足: 初学者は、ここはマニュアルの図と併せて確認すると理解が深まります。CUDA 経験者はスキップしてかまいません。
12.5. GPU fundamentals
この節では、GPU を使うときの大まかな流れ を見ます。
12.5.1. 基本の流れ
CPU 側でデータを用意する
GPU 用のバッファへ渡す
カーネルを実行する
結果を読み戻す
ここでいう カーネル は、GPU 上で並列に動く関数です(戻り値はなく バッファへ書き込みます)。
12.5.2. よく出てくる名前
DeviceContextcompile_functionenqueue_functionDeviceBufferHostBuffer
名前は少し多いですが、 要するに GPU を使う準備・実行・データ受け渡し のための部品です。
12.5.3. 例
次の例は、本物の HostBuffer ではなく、
CPU 側で値をまとめて持つイメージを List で簡単に示したものです。ホスト側で 入力データをまとめておき、カーネルに渡すバッファへコピーする流れのうち、ホスト上で値を保持する部分のイメージに使えます(HostBuffer そのものではありません)。
def main():
var host: List[Float32] = [1.0, 2.0, 3.0, 4.0]
var acc: Float32 = 0.0
for x in host:
acc += x
print(acc)
リスト-5: gpu_host_buffer_list_stand_in.mojo
実際の DeviceContext・DeviceBuffer による転送や同期は、インストール環境とマニュアルの API に沿ってください(名前や手順はバージョンで変わり得ます)。
詳細: GPU fundamentals
出典: Mojo Manual — gpu/fundamentals
補足: GPU 関連の API 名や細かな手順は、バージョンで変わることがあります。実際に試すときは公式を確認するのが安全です。
12.6. GPU block and warp
この節は、GPU の中でスレッド同士がどう協力するか の話です。
12.6.1. まず押さえること
block 内では同期が必要になることがある(
barrier等)warp 単位で効率よく通信できる場合がある(shuffle など)
大きな処理を小さなタイルに分ける考え方が重要(タイル化(tiling)、共有メモリ・協調ロードはマニュアル本編)
12.6.2. タイル化とは
タイル化は、大きな配列や行列をそのまま処理せず、 小さなかたまりに分けて順番に処理する 考え方です。
これは GPU でとてもよく出てきます。 なぜなら、共有メモリや協調処理と相性がよいからです。
12.6.3. 例
次の例は、GPU カーネルそのものではなく、 CPU 上の二重ループでタイルの進み方だけを確認する簡単な形です。タイル化は、大きな配列を 小さなブロック(タイル)に分けて順に処理する考え方のイメージです。GPU の共有メモリや協調ロードの説明はマニュアルに譲ります。
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)
リスト-6: gpu_tile_loop_nest.mojo
16×16 を 4 ずつ進めると、タイル数は 4×4 = 16 になります。実カーネルでのタイルとバリアは GPU block and warp のコード例を参照してください。
出典: Mojo Manual — gpu/block-and-warp
補足: ここは性能チューニングの入口です。実際の同期や共有メモリのコードは、まず公式例を見るのがおすすめです。
12.7. Layouts
この節で大事なのは、同じ配列でも、メモリ上の並び方を明示できる ことです。
12.7.1. Layout とは
Layout は、
データの形と並び方を表すための仕組み です。
ここでよく出てくるのが次の2つです。
shape:形
stride:並び方の間隔
つまり、論理的にどう見えるか と 実際にどう並んでいるか を結びつける仕組みです。コンパイル時に IntTuple 等で表現し、メモリ上の並びと 論理インデックスの対応を 型レベルで固定します。tile などで 部分レイアウトを切り出し、カーネルに渡す使い方にもつながります。
12.7.2. row-major の考え方
Layout.row_major は、行優先の並びです。
これは、最も右のインデックス側が連続したメモリになる並び方です。
12.7.3. 例
次の例では、IntTuple で形状を与え、ランクを確認しています。
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())
リスト-7: layout_row_major_shape.mojo
ストライドや tile による部分ビューは、マニュアルで数値例と合わせて確認すると理解しやすくなります。
詳細: Layouts
出典: Mojo Manual — layout/layouts
補足: レイアウトは、次の
LayoutTensorを理解するための前提です。
12.8. Layout tensors
最後に、レイアウト情報つきの多次元データ を見ます。
12.8.1. LayoutTensor とは
LayoutTensor は、
Layout で決めた並び方の上に、実際のデータを載せた多次元ビュー です。
つまり、ただの配列ではなく、 どう並んでいるかまでわかった状態で扱うテンソル です。
12.8.2. 何ができるか
tileで分割するvectorizeでベクトル化を考えるdistributeで並列処理の単位へ分ける
このあたりは、CPU SIMD と GPU の両方に関わってきます。GPU カーネルと CPU SIMD の両方で 同じ抽象を使う方向性です(詳細はマニュアルを参照します)。
12.8.3. 例
次の例は、CPU 上の InlineArray を使った小さな 2×3 テンソルです。
二重インデックスで要素にアクセスしています。
from layout import Layout, LayoutTensor
def main():
comptime layout = Layout.row_major(2, 3)
var storage = InlineArray[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])
リスト-8: layout_tensor_small.mojo
tile や GPU バッファとの組み合わせは、ワークロードに応じてマニュアル本編のパターンを参照してください。
詳細: Layout tensors
出典: Mojo Manual — layout/tensors
補足: AI や数値計算では重要ですが、最初は「並び方つきのテンソル」と理解すれば十分です。
12.9. この章を一文で言うと
この章は、低水準なメモリ操作、GPU の並列実行、データの並び方を、Mojo でどう扱うかを見る章です。
12.10. まとめ
Mojo ではポインタも用途ごとに型がわかれている
unsafeは最後の手段として使うGPU では thread・block・grid の階層を意識する
カーネル実行では、データ転送と同期の流れが重要になる
Layout は shape と stride を使って並び方を表す
LayoutTensorはレイアウトつきの多次元データを扱う仕組み
この章は、今すぐ全部を使わなくても大丈夫です。 必要になったときに、「メモリ」「並列実行」「並び方」 の3つに分けて読み返すと整理しやすいです。