Mojo 的 GPU 编程模型是如何组织的？

1. 主机与设备的分工

Mojo 通过 MAX 加速器库（MAX accelerator library）提供 GPU 编程支持，该库提供设备管理、GPU 抽象与张量运算等 Mojo API。CPU 作为主机负责管理程序流程、协调 GPU 操作，GPU 作为设备负责在大量线程上并行执行计算。主机与设备各自拥有独立内存空间，二者之间的数据交换需要显式管理。

2. 一次典型 GPU 程序的四个步骤

1. 在主机内存中初始化数据
2. 分配设备内存，并把数据从主机传输到设备
3. 在 GPU 上执行 kernel 函数处理数据
4. 把结果从设备传回主机内存

整个过程通常是异步的：GPU 计算期间 CPU 可以并行处理其他任务，CPU 在需要确认 GPU 已完成操作（例如拷贝结果之前）时，必须显式调用 synchronize() 进行同步。

3. 线程的组织层次

线程的最顶层组织是网格（grid），网格由若干线程块（thread block）构成，二者都可以按一维、二维或三维组织。流式多处理器会把线程块进一步拆分为 warp，其大小依架构通常为 32 或 64 个线程；若线程块线程数不能整除 warp 大小，剩余线程仍会占满一个 warp 的资源，编码时需留意越界访问。

kernel 内部可直接访问以下索引与维度，均为 Int 类型：

• grid_dim：网格各维度大小
• block_dim：线程块各维度大小
• block_idx：当前线程块在网格中的索引
• thread_idx：当前线程在线程块中的索引
• global_idx：线程的全局偏移，例如 global_idx.x = block_dim.x * block_idx.x + thread_idx.x

4. kernel 的编写约束

• kernel 必须是不抛错（non-raising）的函数，如需表达错误，可使用不分配堆内存的自定义类型
• 参数必须符合 DevicePassable trait，如 SIMD、DeviceBuffer（设备侧为 Pointer）、TileTensor 以及 Int32、Int64 等定宽类型；Int 与 UInt 不能直接传入
• kernel 不能有返回值，结果必须写入传入的内存缓冲区
• 需要自行做越界检查，例如通过 if global_idx.x < size 判断

5. 关键 API 与启动流程

主机侧主要通过 max.gpu.host 包工作，核心类型包括表示 GPU 逻辑实例的 DeviceContext、设备与主机内存缓冲区 DeviceBuffer 与 HostBuffer，以及维度描述结构 Dim。典型流程为：创建 DeviceContext，通过 enqueue 系列接口创建缓冲区并完成主机与设备间拷贝，调用 compile_function 编译 kernel 或直接用 enqueue_function 一步完成编译并入队，最后调用 synchronize() 等待执行完成。

6. 硬件厂商支持

目前通过 api 参数与探测函数明确支持三家厂商：

上述检测函数位于标准库 sys 包中，通常在编译期用 comptime if 求值，以便在无 GPU 环境下只编译对应的 CPU 分支。在多 GPU 系统中，设备以从 0 开始的整数索引标识，可通过 DeviceContext.number_of_devices() 查询数量。

需要说明的是，GPU 编程原语并不在标准库内，而是由 MAX 加速器库的 max.gpu 包提供，其下包含 host（主机与设备的交互）、memory（GPU 内存操作）、sync（同步原语）、primitives（warp、block 与 cluster 操作）以及 compute（MMA 与 tensor core 操作）等子包；上文的 grid_dim、block_dim、block_idx、thread_idx、global_idx 等索引别名同样从 max.gpu 导入。