12. ポインタ・GPU・レイアウト

12.1. この章で学ぶこと

  • ポインタの基本と unsafe の考え方

  • GPU 実行モデルの入口

  • LayoutLayoutTensor の役割

この章のテーマは、「安全な高級 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

出典: Mojo Manual — 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_pointeefree で後始末する流れです。+ 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_pointeefree で片づける

という対応関係です。init_pointee_copydestroy_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. よく出てくる名前

  • DeviceContext

  • compile_function

  • enqueue_function

  • DeviceBuffer

  • HostBuffer

名前は少し多いですが、 要するに 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

実際の DeviceContextDeviceBuffer による転送や同期は、インストール環境とマニュアルの 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×164 ずつ進めると、タイル数は 4×4 = 16 になります。実カーネルでのタイルとバリアGPU block and warp のコード例を参照してください。

詳細: 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つに分けて読み返すと整理しやすいです。