2026/9/12 16:56:22
FEATURED · 精选文章 
用 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模块的核心 APIDeviceContext、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 示例文件文件主题对应教材章节vector_addition.mojo向量逐元素相加GPU 编程版 Hello WorldProgramming Massively Parallel Processors第 2 章grayscale.mojoRGB 图像并行转灰度同书第 3 章naive_matrix_multiplication.mojo无硬件优化的朴素矩阵乘法同书第 3 章mandelbrot.mojo并行计算 Mandelbrot 集合逃逸迭代次数自定义计算无输入张量、仅标量参数reduction.mojo高性能归约 kernel含基准测试进阶优化示例前三个示例是线程级 GPU 编程最常见的入门起点与经典教材《Programming Massively Parallel Processors》前三章的内容一一对应便于读者在阅读教材的同时用 Mojo 亲手复现。后两个示例则展示了更复杂与更追求性能的写法。环境准备与快速运行Setup安装 pixi 并获取仓库这些示例使用 pixi 作为包与任务管理工具pixi.toml已锁定完整的依赖环境。按以下步骤准备环境确保系统包含兼容的 GPU若尚未安装 pixi执行官方安装脚本curl -fsSL https://pixi.sh/install.sh | sh克隆本仓库git clone https://gitcode.com/GitHub_Trending/mo/mojo进入示例目录cd mojo/max/examples/gpu-functionsQuickstart两种运行方式方式一直接用pixi run启动 Mojo 解释器运行指定文件pixi run mojo vector_addition.mojo方式二先进入 pixi 提供的虚拟 shell再逐个直接运行pixi shellmojo 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 为模板可拆解为五步定义 GPU 函数编写一个每线程执行一次的函数即 kernel通过thread_idx/global_idx等内置标识确定当前线程负责的数据位置获取设备上下文var ctx DeviceContext()拿到当前加速器的上下文对象后续的内存分配、kernel 调度都经由它完成分配输入输出缓冲区用ctx.enqueue_create_bufferdtype在 GPU 地址空间分配缓冲区再用buffer.enqueue_fill(value)填充初值编译并调度函数ctx.enqueue_functionfn把 GPU 函数按网格grid以块为单位划分和块block每块的线程数调度到设备上参数按函数签名顺序传入取回结果用with buffer.map_to_host() as host_buffer:上下文管理器把设备内存映射回主机端读取。其中核心 API 的语义如下DeviceContext封装对加速器的访问定义于 max/mojo/max/gpu/host/device_context.mojoenqueue_create_buffer在设备地址空间分配缓冲区enqueue_fill以流式enqueue方式填充整个缓冲区enqueue_function编译并提交 kernel 到设备执行队列其多个重载实现在 max/mojo/max/gpu/host/_device_context_extras.mojo 附近支持直接传函数符号或已编译的 kernel 对象并对grid_dim/block_dim的维度合法性做运行时校验_check_dimmap_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_dimgrid_dim, block_dimBLOCK_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_WIDTH10、BLOCK_SIZE5故grid_dim1剩余的边界情况由 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 彩色图像并行转为灰度图核心是引入秩 3rank-3张量来承载「高度 × 宽度 × 颜色通道」三个维度并且调度在二维网格上进行。转换采用经典的加权亮度公式gray 0.21 * red 0.71 * green 0.07 * bluekernel 实现grayscale.mojodef 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-L74comptime 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.mojodef 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 并不硬编码尺寸。示例中I5, J4, K6输入矩阵在主机端用map_to_host填充并打印M[i][j] i - jN[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 集合的迭代规则是从Z0出发反复计算Z Z² C直到Z的模超过阈值 4判定为逃逸并记录逃逸时的迭代次数若达到最大迭代次数仍未逃逸则认为该点属于集合。kernel 实现mandelbrot.mojodef 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_SIZE16并在ctx.enqueue_function之后显式调用ctx.synchronize()等待 kernel 完成。运行命令pixi run mojo mandelbrot.mojo结果是在终端以 ASCII 字符绘制的 Mandelbrot 集合图案字符亮度对应迭代次数draw_mandelbrot 用一串....,c8MjawrpogOQEPGJ作为灰度色阶未逃逸点输出空格例如...................................,,,,c8cc,,,............. ...............................,,,,,,cc8M Mjc,,,,.......... ............................,,,,,,,ccccMaQaM8c,,,,,........ ..........................,,,,,,,ccc88g.o. Owg8ccc,,,,...... .......................,,,,,,,,c8888Mj, ,wMM8cccc,,..... .....................,,,,,,cccMQOPjjPrgg, OrwrwMMMjjc,.... ..................,,,,cccccc88MaP ,pGa.g8c,... ...............,,cccccccc888MjQp. o8cc,.. ..........,,,,c8jjMMMMMMMMMw. aj8c,,. .....,,,,,,ccc88QEJwr.wPjjjwG w8c,,. ..,,,,,,,cccccMMjwQ EpQ .8c,,. .,,,,,,cc888MrajwJ MMcc,,, .cc88jMMMjaG. oM8cc,,,修改GRID_WIDTH/GRID_HEIGHT可获得不同分辨率调整MIN_X/MAX_X/MIN_Y/MAX_Y则可以观察复平面上不同区域的细节。示例五高性能归约 kernelreduction.mojo 是五个示例中性能取向最强的一个对一个长度SIZE 1 12的int32向量做求和归约并内置了基准测试。它综合运用了 Mojo GPU 编程的多项进阶技术1. 向量化加载 跨块网格步长grid-stride循环。每个线程不是只处理一个元素而是通过a.unsafe_loadwidthbatch_size一次加载BATCH_SIZE8个元素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_loadwidthbatch_size.reduce_add()2. 共享内存 块内归约。每个线程把局部和写入共享内存数组sumsunsafe_stack_allocation[...address_space.SHARED]调用barrier()同步后通过comptime for展开的树形归约逐轮把活跃线程减半、两两累加直至只剩一个 warpvar sums unsafe_stack_allocation[KERNEL_TPB, Scalar[dtype], address_space.SHARED]() ... sums[unsafe_offsettid] 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_offsettid] sums[unsafe_offsettid 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_offsettid][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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
用 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模块的核心 APIDeviceContext、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 示例文件文件主题对应教材章节vector_addition.mojo向量逐元素相加GPU 编程版 Hello WorldProgramming Massively Parallel Processors第 2 章grayscale.mojoRGB 图像并行转灰度同书第 3 章naive_matrix_multiplication.mojo无硬件优化的朴素矩阵乘法同书第 3 章mandelbrot.mojo并行计算 Mandelbrot 集合逃逸迭代次数自定义计算无输入张量、仅标量参数reduction.mojo高性能归约 kernel含基准测试进阶优化示例前三个示例是线程级 GPU 编程最常见的入门起点与经典教材《Programming Massively Parallel Processors》前三章的内容一一对应便于读者在阅读教材的同时用 Mojo 亲手复现。后两个示例则展示了更复杂与更追求性能的写法。环境准备与快速运行Setup安装 pixi 并获取仓库这些示例使用 pixi 作为包与任务管理工具pixi.toml已锁定完整的依赖环境。按以下步骤准备环境确保系统包含兼容的 GPU若尚未安装 pixi执行官方安装脚本curl -fsSL https://pixi.sh/install.sh | sh克隆本仓库git clone https://gitcode.com/GitHub_Trending/mo/mojo进入示例目录cd mojo/max/examples/gpu-functionsQuickstart两种运行方式方式一直接用pixi run启动 Mojo 解释器运行指定文件pixi run mojo vector_addition.mojo方式二先进入 pixi 提供的虚拟 shell再逐个直接运行pixi shellmojo 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 为模板可拆解为五步定义 GPU 函数编写一个每线程执行一次的函数即 kernel通过thread_idx/global_idx等内置标识确定当前线程负责的数据位置获取设备上下文var ctx DeviceContext()拿到当前加速器的上下文对象后续的内存分配、kernel 调度都经由它完成分配输入输出缓冲区用ctx.enqueue_create_bufferdtype在 GPU 地址空间分配缓冲区再用buffer.enqueue_fill(value)填充初值编译并调度函数ctx.enqueue_functionfn把 GPU 函数按网格grid以块为单位划分和块block每块的线程数调度到设备上参数按函数签名顺序传入取回结果用with buffer.map_to_host() as host_buffer:上下文管理器把设备内存映射回主机端读取。其中核心 API 的语义如下DeviceContext封装对加速器的访问定义于 max/mojo/max/gpu/host/device_context.mojoenqueue_create_buffer在设备地址空间分配缓冲区enqueue_fill以流式enqueue方式填充整个缓冲区enqueue_function编译并提交 kernel 到设备执行队列其多个重载实现在 max/mojo/max/gpu/host/_device_context_extras.mojo 附近支持直接传函数符号或已编译的 kernel 对象并对grid_dim/block_dim的维度合法性做运行时校验_check_dimmap_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_dimgrid_dim, block_dimBLOCK_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_WIDTH10、BLOCK_SIZE5故grid_dim1剩余的边界情况由 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 彩色图像并行转为灰度图核心是引入秩 3rank-3张量来承载「高度 × 宽度 × 颜色通道」三个维度并且调度在二维网格上进行。转换采用经典的加权亮度公式gray 0.21 * red 0.71 * green 0.07 * bluekernel 实现grayscale.mojodef 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-L74comptime 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.mojodef 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 并不硬编码尺寸。示例中I5, J4, K6输入矩阵在主机端用map_to_host填充并打印M[i][j] i - jN[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 集合的迭代规则是从Z0出发反复计算Z Z² C直到Z的模超过阈值 4判定为逃逸并记录逃逸时的迭代次数若达到最大迭代次数仍未逃逸则认为该点属于集合。kernel 实现mandelbrot.mojodef 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_SIZE16并在ctx.enqueue_function之后显式调用ctx.synchronize()等待 kernel 完成。运行命令pixi run mojo mandelbrot.mojo结果是在终端以 ASCII 字符绘制的 Mandelbrot 集合图案字符亮度对应迭代次数draw_mandelbrot 用一串....,c8MjawrpogOQEPGJ作为灰度色阶未逃逸点输出空格例如...................................,,,,c8cc,,,............. ...............................,,,,,,cc8M Mjc,,,,.......... ............................,,,,,,,ccccMaQaM8c,,,,,........ ..........................,,,,,,,ccc88g.o. Owg8ccc,,,,...... .......................,,,,,,,,c8888Mj, ,wMM8cccc,,..... .....................,,,,,,cccMQOPjjPrgg, OrwrwMMMjjc,.... ..................,,,,cccccc88MaP ,pGa.g8c,... ...............,,cccccccc888MjQp. o8cc,.. ..........,,,,c8jjMMMMMMMMMw. aj8c,,. .....,,,,,,ccc88QEJwr.wPjjjwG w8c,,. ..,,,,,,,cccccMMjwQ EpQ .8c,,. .,,,,,,cc888MrajwJ MMcc,,, .cc88jMMMjaG. oM8cc,,,修改GRID_WIDTH/GRID_HEIGHT可获得不同分辨率调整MIN_X/MAX_X/MIN_Y/MAX_Y则可以观察复平面上不同区域的细节。示例五高性能归约 kernelreduction.mojo 是五个示例中性能取向最强的一个对一个长度SIZE 1 12的int32向量做求和归约并内置了基准测试。它综合运用了 Mojo GPU 编程的多项进阶技术1. 向量化加载 跨块网格步长grid-stride循环。每个线程不是只处理一个元素而是通过a.unsafe_loadwidthbatch_size一次加载BATCH_SIZE8个元素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_loadwidthbatch_size.reduce_add()2. 共享内存 块内归约。每个线程把局部和写入共享内存数组sumsunsafe_stack_allocation[...address_space.SHARED]调用barrier()同步后通过comptime for展开的树形归约逐轮把活跃线程减半、两两累加直至只剩一个 warpvar sums unsafe_stack_allocation[KERNEL_TPB, Scalar[dtype], address_space.SHARED]() ... sums[unsafe_offsettid] 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_offsettid] sums[unsafe_offsettid 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_offsettid][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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
)
