用 Mojo 编写 GPU 函数：从向量加法到高性能归约的五个实战示例详解

用 Mojo 编写 GPU 函数：从向量加法到高性能归约的五个实战示例详解

【免费下载链接】mojoThe Modular Platform (includes MAX & Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo

本篇技术指南以 Modular 开源仓库 max/examples/gpu-functions 目录为核心，系统讲解如何用 Mojo 语言编写、编译并调度（dispatch）在 GPU 上运行的线程级函数，全程无需编写 CUDA 或任何厂商专有代码。读完本文，你将掌握max.gpu模块的核心 API（DeviceContext、enqueue_create_buffer、enqueue_function、map_to_host）、一维与二维线程网格的调度方式，以及从最简单的向量加法、图像灰度化、朴素矩阵乘法，到 Mandelbrot 分形计算和高性能归约（reduction）内核的完整实战路径。

背景：Mojo 的 GPU 编程定位

Mojo 是一门面向高性能计算的 Python 家族语言。在 GPU 编程场景下，Mojo 允许开发者直接编写自定义的 GPU 算法，而不必依赖 CUDA、ROCm 等厂商专有库——编译期会由 Mojo 编译器将函数编译为面向目标加速器的代码，运行时则由 max.gpu 模块 统一处理硬件相关的细节，包括：

• 在主机与加速器之间分配和迁移内存；
• 把面向 GPU 的函数编译并调度到设备上执行。

[!IMPORTANT] 运行这些示例需要一块兼容的 GPU。示例源码在入口处均通过std.sys的has_accelerator()做了编译期断言，例如 vector_addition.mojo 中的comptime assert has_accelerator(), "This example requires a supported GPU"，在没有加速器的机器上会直接编译失败并给出明确提示。

本目录中的示例与仓库里另一组 custom_ops 示例 互补：后者演示如何把 GPU kernel 封装为可运行在 CPU 和 GPU 之上的自定义图算子（graph op），其中同样包含 Mandelbrot 与向量加法的对应实现（mandelbrot.py、vector_addition.py），适合对比阅读。

目录内共包含五个 Mojo 示例文件：

前三个示例是线程级 GPU 编程最常见的入门起点，与经典教材《Programming Massively Parallel Processors》前三章的内容一一对应，便于读者在阅读教材的同时用 Mojo 亲手复现。后两个示例则展示了更复杂与更追求性能的写法。

环境准备与快速运行

Setup：安装 pixi 并获取仓库

这些示例使用 pixi 作为包与任务管理工具，pixi.toml已锁定完整的依赖环境。按以下步骤准备环境：

1. 确保系统包含兼容的 GPU；
2. 若尚未安装 pixi，执行官方安装脚本：

   curl -fsSL https://pixi.sh/install.sh | sh

3. 克隆本仓库：

   git clone https://gitcode.com/GitHub_Trending/mo/mojo

4. 进入示例目录：

   cd mojo/max/examples/gpu-functions

Quickstart：两种运行方式

方式一：直接用pixi run启动 Mojo 解释器运行指定文件：

pixi run mojo vector_addition.mojo

方式二：先进入 pixi 提供的虚拟 shell，再逐个直接运行：

pixi shell

mojo vector_addition.mojo mojo grayscale.mojo mojo naive_matrix_multiplication.mojo mojo mandelbrot.mojo mojo reduction.mojo

打开目录下的 pixi.toml 可以看到，除了直接运行，还预定义了五个命名任务以及一个聚合的test任务：

[tasks] vector_addition = "mojo run vector_addition.mojo" grayscale = "mojo run grayscale.mojo" naive_matrix_multiplication = "mojo run naive_matrix_multiplication.mojo" mandelbrot = "mojo run mandelbrot.mojo" reduction = "mojo run reduction.mojo" test = { depends-on = ["vector_addition", "grayscale", "naive_matrix_multiplication", "mandelbrot"] }

因此也可以使用pixi run test一次性跑通前四个示例。依赖方面，pixi.toml声明了mojo = "*"与max = "*"两个包，渠道（channels）包含conda-forge与 Modular 的 nightly 源，支持osx-arm64、linux-aarch64、linux-64三种平台。

此外，仓库还提供了 Bazel 构建入口 BUILD.bazel：五个示例均以mojo_binary规则构建，并通过target_compatible_with = ["//:has_gpu"]限定仅在检测到 GPU 的目标平台下编译；每个二进制还对应一个modular_run_binary_test运行测试（打有gpu标签），其中naive_matrix_multiplication因编译问题暂标记为manual（源码中留有 TODO 注释）。有 Bazel 环境的读者可以用 bazel 方式构建与测试。

GPU 编程核心 API：五个关键步骤

无论哪个示例，其主流程（main函数）都遵循同一套模式，以 vector_addition.mojo 为模板可拆解为五步：

1. 定义 GPU 函数：编写一个每线程执行一次的函数（即 kernel），通过thread_idx/global_idx等内置标识确定当前线程负责的数据位置；
2. 获取设备上下文：var ctx = DeviceContext()拿到当前加速器的上下文对象，后续的内存分配、kernel 调度都经由它完成；
3. 分配输入输出缓冲区：用ctx.enqueue_create_bufferdtype在 GPU 地址空间分配缓冲区，再用buffer.enqueue_fill(value)填充初值；
4. 编译并调度函数：ctx.enqueue_functionfn把 GPU 函数按网格（grid，以块为单位划分）和块（block，每块的线程数）调度到设备上，参数按函数签名顺序传入；
5. 取回结果：用with buffer.map_to_host() as host_buffer:上下文管理器把设备内存映射回主机端读取。

其中核心 API 的语义如下：

• DeviceContext：封装对加速器的访问，定义于 max/mojo/max/gpu/host/device_context.mojo；
• enqueue_create_buffer：在设备地址空间分配缓冲区，enqueue_fill以流式（enqueue）方式填充整个缓冲区；
• enqueue_function：编译并提交 kernel 到设备执行队列，其多个重载实现在 max/mojo/max/gpu/host/_device_context_extras.mojo 附近，支持直接传函数符号或已编译的 kernel 对象，并对grid_dim/block_dim的维度合法性做运行时校验（_check_dim）；
• map_to_host：将设备缓冲区映射到主机地址空间，供 CPU 侧读取结果；
• TileTensor与row_major：来自layout模块（对应 Bazel 依赖//max:layout），用指定的 row-major 布局把设备缓冲区包装成可索引的张量视图，MutAnyOrigin标注表示 kernel 内部可写。

从源码结构看，max.mojo层的max/mojo/max/algorithm/backend/gpu/目录下的elementwise.mojo、reduction.mojo、stencil.mojo等高级算子实现同样基于这套enqueue_function调度机制，说明本目录示例演示的就是 MAX 框架底层使用的同一套 GPU 抽象。

示例一：向量加法（GPU 版 "Hello World"）

向量逐元素相加是数据并行编程中最经典的入门示例，对应教材第 2 章。kernel 本身极其简单：每个线程根据自身线程 ID 处理一对元素，把lhs + rhs写入out的对应位置。

以下是仓库中 vector_addition.mojo 的实际实现（注意：README 中的代码片段可能滞后于源码，请以源码文件为准）：

def vector_addition( lhs_tensor: TileTensor[float_dtype, type_of(layout), MutAnyOrigin], rhs_tensor: TileTensor[float_dtype, type_of(layout), MutAnyOrigin], out_tensor: TileTensor[float_dtype, type_of(layout), MutAnyOrigin], size_dev: Int32, ): """The calculation to perform across the vector on the GPU.""" var size = Int(size_dev) var global_tid = global_idx.x if global_tid < size: out_tensor[global_tid] = lhs_tensor[global_tid] + rhs_tensor[global_tid]

这里使用global_idx.x作为全局线程 ID（跨越所有块），并通过if global_tid < size做边界检查，保证即使size不能被块大小整除也不会越界。主流程的调度部分如下：

comptime float_dtype = DType.float32 comptime VECTOR_WIDTH = 10 comptime BLOCK_SIZE = 5 comptime layout = row_major[VECTOR_WIDTH]() # 分配设备端缓冲区 var lhs_buffer = ctx.enqueue_create_bufferfloat_dtype var rhs_buffer = ctx.enqueue_create_bufferfloat_dtype var out_buffer = ctx.enqueue_create_bufferfloat_dtype # 填充初值 lhs_buffer.enqueue_fill(1.25) rhs_buffer.enqueue_fill(2.5) # 用布局包装成张量 var lhs_tensor = TileTensor(lhs_buffer, layout) var rhs_tensor = TileTensor(rhs_buffer, layout) var out_tensor = TileTensor(out_buffer, layout) # 计算需要多少个块才能覆盖整个向量（向上取整） var grid_dim = ceildiv(VECTOR_WIDTH, BLOCK_SIZE) # 编译并调度 GPU kernel ctx.enqueue_functionvector_addition, grid_dim=grid_dim, block_dim=BLOCK_SIZE, ) # 映射回主机读取结果 with out_buffer.map_to_host() as host_buffer: var host_tensor = TileTensor(host_buffer, layout) print("Resulting vector:", host_tensor)

几个值得注意的细节：

• 所有尺寸与布局都是编译期（comptime）常量，layout = row_major[VECTOR_WIDTH]()把一维布局信息编码进类型系统，kernel 签名中的type_of(layout)由此与调度侧严格对齐；
• 网格划分：ceildiv(VECTOR_WIDTH, BLOCK_SIZE)计算覆盖全部元素所需的最小块数（本示例中VECTOR_WIDTH=10、BLOCK_SIZE=5，故grid_dim=1），剩余的边界情况由 kernel 内的global_tid < size兜底；
• 内核在 Mojo 文件编译期即完成 GPU 编译：enqueue_function[vector_addition]的方括号语法在编译期对函数做设备端特化，运行时只需提交调度。

运行方式：

pixi run mojo vector_addition.mojo

预期输出是所有元素均为3.75的向量（因为1.25 + 2.5 = 3.75）。可以尝试修改VECTOR_WIDTH、BLOCK_SIZE、填充值等参数，观察计算在不同规模下的表现。

示例二：彩色图像转灰度

grayscale.mojo 演示把 RGB 彩色图像并行转为灰度图，核心是引入秩 3（rank-3）张量来承载「高度 × 宽度 × 颜色通道」三个维度，并且调度在二维网格上进行。

转换采用经典的加权亮度公式：

gray = 0.21 * red + 0.71 * green + 0.07 * blue

kernel 实现（grayscale.mojo）：

def color_to_grayscale( rgb_tensor: TileTensor[int_dtype, type_of(rgb_layout), MutAnyOrigin], gray_tensor: TileTensor[int_dtype, type_of(gray_layout), MutAnyOrigin], ): var row = global_idx.y var col = global_idx.x if col < WIDTH and row < HEIGHT: var red = rgb_tensor[row, col, 0].cast[float_dtype]() var green = rgb_tensor[row, col, 1].cast[float_dtype]() var blue = rgb_tensor[row, col, 2].cast[float_dtype]() var gray = 0.21 * red + 0.71 * green + 0.07 * blue gray_tensor[row, col] = gray.cast[int_dtype]()

这里通过global_idx.y/global_idx.x分别取行、列坐标；由于三个通道以uint8存储，先.cast[float_dtype]()提升到float32参与加权计算，再cast[int_dtype]()写回uint8的灰度张量。边界检查col < WIDTH and row < HEIGHT保证恰好覆盖图像尺寸。

布局与二维调度（grayscale.mojo 与 L62-L74）：

comptime WIDTH = 5 comptime HEIGHT = 10 comptime NUM_CHANNELS = 3 comptime rgb_layout = row_major[HEIGHT, WIDTH, NUM_CHANNELS]() comptime gray_layout = row_major[HEIGHT, WIDTH]()

comptime BLOCK_SIZE = 16 var num_col_blocks = ceildiv(WIDTH, BLOCK_SIZE) var num_row_blocks = ceildiv(HEIGHT, BLOCK_SIZE) ctx.enqueue_functioncolor_to_grayscale, block_dim=(BLOCK_SIZE, BLOCK_SIZE), )

主流程中还有一个值得学习的模式：先用with rgb_buffer.map_to_host() as host_buffer:在 CPU 侧构造 RGB 图像并填充像素值，再用同一缓冲区的设备端包装发起计算，最后映射回主机打印灰度结果。print_image辅助函数负责把灰度强度格式化输出为对齐的数字网格。

运行命令：

pixi run mojo grayscale.mojo

输出是一张由数字组成的灰度网格。可以修改WIDTH、HEIGHT、像素填充公式或BLOCK_SIZE，观察 GPU 上二维网格并行度的变化。

示例三：朴素矩阵乘法

naive_matrix_multiplication.mojo 实现完全没有硬件优化（无共享内存、无分块 tiling）的朴素矩阵乘法M(I×J) × N(J×K) = P(I×K)，每个线程负责输出矩阵中的一个元素，对j做串行累加。

kernel 实现（naive_matrix_multiplication.mojo）：

def naive_matrix_multiplication( m: TileTensor[float_dtype, type_of(m_layout), MutAnyOrigin], n: TileTensor[float_dtype, type_of(n_layout), MutAnyOrigin], p: TileTensor[float_dtype, type_of(p_layout), MutAnyOrigin], ): var row = global_idx.y var col = global_idx.x var m_dim = Int(p.dim[0]()) var n_dim = Int(p.dim[1]()) var k_dim = Int(m.dim[1]()) if row < m_dim and col < n_dim: for j_index in range(k_dim): p[row, col] = p[row, col] + m[row, j_index] * n[j_index, col]

矩阵维度从张量的运行时dim信息动态读取（p.dim[0]()为输出行数，p.dim[1]()为输出列数，m.dim[1]()为归约维度长度），因此 kernel 并不硬编码尺寸。示例中I=5, J=4, K=6，输入矩阵在主机端用map_to_host填充并打印（M[i][j] = i - j，N[i][j] = i + j），调度配置与灰度示例一致：

comptime BLOCK_SIZE = 16 comptime num_col_blocks = ceildiv(I, BLOCK_SIZE) comptime num_row_blocks = ceildiv(J, BLOCK_SIZE) ctx.enqueue_functionnaive_matrix_multiplication, block_dim=(BLOCK_SIZE, BLOCK_SIZE), )

运行命令：

pixi run mojo naive_matrix_multiplication.mojo

控制台会依次打印M、N两个输入矩阵以及相乘结果P。这个示例最大的价值在于作为正确性基准：后续若要实现基于共享内存 tiling、向量化加载等优化的矩阵乘法，可以先与这份朴素实现的结果对照验证。

示例四：计算 Mandelbrot 集合分形

mandelbrot.mojo 演示一个更有趣的场景：kernel 不接收任何输入张量，只依赖一组编译期标量常量，输出一张记录各位置「逃逸迭代次数」的二维整数矩阵。Mandelbrot 集合的迭代规则是从Z=0出发反复计算Z = Z² + C，直到Z的模超过阈值 4（判定为逃逸），并记录逃逸时的迭代次数；若达到最大迭代次数仍未逃逸，则认为该点属于集合。

kernel 实现（mandelbrot.mojo）：

def mandelbrot( tensor: TileTensor[int_dtype, type_of(layout), MutAnyOrigin], ): var row = global_idx.y var col = global_idx.x comptime SCALE_X = (MAX_X - MIN_X) / GRID_WIDTH comptime SCALE_Y = (MAX_Y - MIN_Y) / GRID_HEIGHT var cx = MIN_X + Float32(col) * SCALE_X var cy = MIN_Y + Float32(row) * SCALE_Y var c = ComplexScalarfloat_dtype var z = ComplexScalarfloat_dtype var iters = Scalarint_dtype var in_set_mask = Scalar.bool for _ in range(MAX_ITERATIONS): if not any(in_set_mask): break in_set_mask = z.squared_norm().le(4) iters = in_set_mask.select(iters + 1, iters) z = z.squared_add(c) tensor[row, col] = iters

要点解析：

• 每个线程由(row, col)映射到复平面上的一点C，随后执行最多MAX_ITERATIONS次的迭代；
• 使用std.complex的ComplexScalar[float_dtype]（float32 标量复数）与Scalar类型；in_set_mask.select(iter + 1, iter)是 SIMD/标量风格的掩码选择：仍在集合内的点迭代次数 +1，其余保持不变；当any(in_set_mask)为假（全部逃逸）时提前break退出循环，避免多余计算；
• 控制观察区域、分辨率与迭代上限的常量均为编译期别名：

comptime GRID_WIDTH = 60 comptime GRID_HEIGHT = 25 comptime MIN_X: Scalar[float_dtype] = -2.0 comptime MAX_X: Scalar[float_dtype] = 0.7 comptime MIN_Y: Scalar[float_dtype] = -1.12 comptime MAX_Y: Scalar[float_dtype] = 1.12 comptime MAX_ITERATIONS = 100

调度与上一示例相同（grid_dim=(COL_BLOCKS, ROW_BLOCKS)、block_dim=(BLOCK_SIZE, BLOCK_SIZE)，BLOCK_SIZE=16），并在ctx.enqueue_function之后显式调用ctx.synchronize()等待 kernel 完成。

运行命令：

pixi run mojo mandelbrot.mojo

结果是在终端以 ASCII 字符绘制的 Mandelbrot 集合图案（字符亮度对应迭代次数，draw_mandelbrot 用一串"....,c8M@jawrpogOQEPGJ"作为灰度色阶，未逃逸点输出空格），例如：

...................................,,,,c@8cc,,,............. ...............................,,,,,,cc8M @Mjc,,,,.......... ............................,,,,,,,ccccM@aQaM8c,,,,,........ ..........................,,,,,,,ccc88g.o. Owg8ccc,,,,...... .......................,,,,,,,,c8888M@j, ,wMM8cccc,,..... .....................,,,,,,cccMQOPjjPrgg, OrwrwMMMjjc,.... ..................,,,,cccccc88MaP @ ,pGa.g8c,... ...............,,cccccccc888MjQp. o@8cc,.. ..........,,,,c8jjMMMMMMMMM@@w. aj8c,,. .....,,,,,,ccc88@QEJwr.wPjjjwG w8c,,. ..,,,,,,,cccccMMjwQ EpQ .8c,,. .,,,,,,cc888MrajwJ MMcc,,, .cc88jMMM@@jaG. oM8cc,,,

修改GRID_WIDTH/GRID_HEIGHT可获得不同分辨率，调整MIN_X/MAX_X/MIN_Y/MAX_Y则可以观察复平面上不同区域的细节。

示例五：高性能归约 kernel

reduction.mojo 是五个示例中性能取向最强的一个：对一个长度SIZE = 1 << 12的int32向量做求和归约，并内置了基准测试。它综合运用了 Mojo GPU 编程的多项进阶技术：

1. 向量化加载 + 跨块网格步长（grid-stride）循环。每个线程不是只处理一个元素，而是通过a.unsafe_loadwidth=batch_size一次加载BATCH_SIZE=8个元素（SIMD 宽度为 8），再用.reduce_add()在寄存器内完成部分归约；for i in range(global_tid, size, threads_in_grid)让所有线程以网格步长遍历，保证任意SIZE都能正确覆盖：

var global_tid = block_idx.x * block_dim.x + thread_idx.x var threads_in_grid = KERNEL_TPB * NUM_BLOCKS var sum: Int32 = 0 for i in range(global_tid, size, threads_in_grid): var idx = i * batch_size if idx < size: sum += a.unsafe_loadwidth=batch_size.reduce_add()

2. 共享内存 + 块内归约。每个线程把局部和写入共享内存数组sums（unsafe_stack_allocation[...address_space=.SHARED]），调用barrier()同步后，通过comptime for展开的树形归约逐轮把活跃线程减半、两两累加，直至只剩一个 warp：

var sums = unsafe_stack_allocation[KERNEL_TPB, Scalar[dtype], address_space=.SHARED]() ... sums[unsafe_offset=tid] = sum barrier() var active_threads = KERNEL_TPB comptime KERNEL_LOG_TPB = log2_floor(KERNEL_TPB) comptime for power in range(1, KERNEL_LOG_TPB - log2_floor(WARP_SIZE) + 1): active_threads >>= 1 if tid < active_threads: sums[unsafe_offset=tid] += sums[unsafe_offset=tid + active_threads] barrier()

3. warp 级归约 + 原子累加。最后只让前WARP_SIZE个线程参与：先读共享内存，用warp.sum做 warp 内归约，再由tid == 0的线程通过Atomic.fetch_add把该块的部分和原子累加到输出指针上，从而支持任意数量的块：

if tid < WARP_SIZE: var warp_sum: Int32 = sums[unsafe_offset=tid][0] warp_sum = warp.sum(warp_sum) if tid == 0: _ = Atomic.fetch_add(output, warp_sum)

4. 内置基准测试与正确性校验。main中先用随机数（randint生成 0–10 之间的整数）填充输入，kernel 结果与 CPU 侧顺序求和的结果通过assert_equal比对；随后用std.benchmark的Bench/Bencher与max.benchmark的bencher_iter_custom对sum_kernel做吞吐量测量（ThroughputMeasure以字节为单位度量处理数据量），并打印格式化的基准表。关键参数集中在文件头部：

comptime TPB = 512 # 每块线程数 comptime BATCH_SIZE = 8 # SIMD 批量宽度，需为 2 的幂 comptime SIZE = 1 << 12 # 向量长度 comptime NUM_BLOCKS = ceildiv(SIZE, TPB * BATCH_SIZE) comptime dtype = DType.int32

注释特别提示：要测量高带宽，应把SIZE增大到较大数值。运行：

pixi run mojo reduction.mojo

进阶方向：从 GPU 函数到自定义图算子

如果想把这里的 GPU 函数封装成 MAX 计算图中的算子（可同时在 CPU 和 GPU 上运行），可以参考仓库中的 max/examples/custom_ops 目录——README 明确说明它与本目录示例互补。其中mandelbrot.py、vector_addition.py、matrix_multiplication.py分别对应本文的 Mandelbrot、向量加法与矩阵乘法示例，但以图算子（custom op）的形式呈现；kernels/子目录内还有对应的 Mojo kernel 实现，适合对照「裸 GPU 函数」与「图算子封装」两种编程范式的差异。

若想继续深入学习 Mojo GPU 编程，仓库内还提供了其他资料：

• max/examples/gpu-block-and-warp：深入 block 与 warp 层级的并行模型；
• max/examples/gpu-functions/BUILD.bazel：了解如何在 Bazel 工程中以mojo_binary+//:has_gpu约束构建 GPU 示例；
• max/mojo/max/gpu/ 与 max/mojo/max/algorithm/backend/gpu/：阅读DeviceContext与 elementwise/reduction/stencil 等高级算子的底层实现，理解enqueue_function之上的抽象层次。

小结

通过gpu-functions目录这五个层层递进的示例，你可以掌握 Mojo GPU 编程的完整链路：从「定义每线程函数 → 获取DeviceContext→ 分配设备缓冲区 → 按网格/块调度 → 映射回主机」的基础五步，到二维网格、秩 3 张量、编译期常量、共享内存、warp 归约与原子操作等进阶技巧。它们既是入门教材的 Mojo 复现，也是通往 MAX 自定义算子与高性能 kernel 开发的坚实起点。

【免费下载链接】mojoThe Modular Platform (includes MAX & Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo