MLX 编译指南使用 mx.compile 合并计算图、融合算子并加速训练【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlxMLXArray framework for Apple silicon在mlx.core中提供了compile这一函数变换function transformation用于编译计算图。函数编译通过合并重复的公共计算、融合特定算子来得到更小的计算图在多数场景下能显著降低运行时间与内存占用。本文以 compile.rst 为主体结合 mlx/compile.cpp 与 python/src/transforms.cpp 的源码实现系统讲解mx.compile的使用方式、缓存与重编译行为、纯函数约束、训练图编译、与其它变换的组合以及 shapeless 编译帮助你写出可安全编译、可复用、可调试的高性能 MLX 代码。一、编译的基本用法从普通函数到编译函数mx.compile的入门用法非常简单把一个以数组为输入、数组为输出的普通 Python 函数包装起来即可。下面的例子演示了普通调用与编译调用的区别def fun(x, y): return mx.exp(-x) y x mx.array(1.0) y mx.array(2.0) # 普通调用不编译 # 输出: array(2.36788, dtypefloat32) print(fun(x, y)) # 编译该函数 compiled_fun mx.compile(fun) # 输出: array(2.36788, dtypefloat32) print(compiled_fun(x, y))普通函数与编译函数的输出在数值精度范围内是一致的——compile是语义保持semantics-preserving的变换只改变执行的组织方式不改变计算结果。从底层看mx.compile返回的编译函数在 mlx/compile.cpp 中实现为一条完整的编译流水线首次调用时会依次执行compile_trace用占位符placeholder输入调用原函数追踪出计算图compile_dfs深度优先遍历计算图构建 tape操作序列与 parents map父节点映射compile_simplify简化 tape——合并相同标量、移除 no-opCopy、StopGradient、多轮合并等价子表达式默认 3 轮见 mlx/compile.cppcompile_fuse将可融合的子图替换为Compiled原语见 mlx/compile.cppcompile_replace把占位符替换为真实数组供后续求值。缓存机制编译一次多次复用编译本身有成本第一次调用编译函数时MLX 需要构建计算图、优化并生成/编译代码这个过程相对较慢。但 MLX 会对编译结果做缓存后续调用不会重新触发编译。因此建议只编译你打算多次调用的函数def fun(x, y): return mx.exp(-x) y x mx.array(1.0) y mx.array(2.0) compiled_fun mx.compile(fun) # 此处发生编译 compiled_fun(x, y) # 不再编译 compiled_fun(x, y) # 不再编译 mx.compile(fun)(x, y)缓存在 mlx/compile.cpp 的CompileCache中实现它以函数地址fun_id为主键逐条缓存 entry命中条件包括默认 stream/设备匹配、shapeless标志一致、输入数组的 shape 与 dtype 相等以及编译常量一致。缓存使用共享锁实现线程安全并且是线程局部的thread_local见 mlx/compile.cpp。何时会触发重新编译以下情况会导致函数被重新编译输入的形状或维度数量发生变化任一输入的类型发生变化函数的输入个数发生变化。其中部分情况只重跑编译栈中的某几层例如仅改变形状而另一些情况例如改变类型会重跑整条编译栈。一般来说应避免过于频繁地编译函数。相应地Python 绑定层在 python/src/transforms.cpp 的compile签名中提供了三个可选参数inputs、outputs与shapeless其中shapeless正是为缓解形状变化导致的重编译而设计的详见后文Shapeless 编译一节。需要警惕的反模式在循环中创建并销毁编译函数另一个容易踩坑的写法是编译那些被频繁创建又销毁的函数例如在循环内编译匿名函数a mx.array(1.0) # 不要这样做每一轮迭代都会重新编译这个 lambda for _ in range(5): mx.compile(lambda x: mx.exp(mx.abs(x)))(a)由于每次迭代都会新建一个 lambda 对象其函数地址缓存键随之变化缓存无法命中从而反复编译。应把编译函数提取到循环之外只创建一次。二、真实加速案例编译 GELUmlx.nn.gelu是 Transformer 类模型中常用的非线性激活函数其实现涉及多个一元unary与二元binary逐元素操作def gelu(x): return x * (1 mx.erf(x / math.sqrt(2))) / 2当输入数组很小时该函数受调用开销overhead-bound限制当输入数组很大时则受内存带宽memory bandwidth-bound限制。而gelu中的所有操作都是可融合的mx.compile可以将其融合进单个 kernel从而在这两种场景下都获得可观加速。下面用计时辅助函数对比普通函数与编译函数的运行时间该辅助函数先做 10 次 warm-up并在计时循环中通过mx.eval完成同步import time def timeit(fun, x): # warm up for _ in range(10): mx.eval(fun(x)) tic time.perf_counter() for _ in range(100): mx.eval(fun(x)) toc time.perf_counter() tpi 1e3 * (toc - tic) / 100 print(fTime per iteration {tpi:.3f} (ms))构造一个大数组并分别计时x mx.random.uniform(shape(32, 1000, 4096)) timeit(gelu, x) timeit(mx.compile(gelu), x)在 M1 Max 上普通gelu约为 15.5 毫秒编译后的gelu约为 3.1 毫秒编译版本快约 5 倍该数值来自文档在 M1 Max 上的实测具体结果随设备与数组规模而异。从源码可以印证哪些算子可融合的判断逻辑在 mlx/compile.cpp 中is_unary覆盖Exp、Erf、Negative、Log、Sigmoid等一元算子is_binary覆盖Add、Multiply、Divide、Subtract等二元算子is_ternary覆盖Selectis_broadcast覆盖Broadcast而is_fusable正是这四类的并集。compile_fuse在反向遍历 tape 时递归收集可融合的算子并受两个常量约束最大融合深度max_compile_depth 11、最多输入数组数max_compile_arrays 24见 mlx/compile.cpp。因此像gelu这种由一元/二元逐元素算子串成的计算链可以整体融合成一个Compiled原语只需在编译阶段为该子图生成一个 kernel 即可。三、调试编译函数占位符、disable_compile 与 MLX_DISABLE_COMPILE编译函数在首次被调用时是用占位符输入进行追踪tracing的。这意味着在编译函数内部不能对数组求值例如打印数组内容否则会崩溃mx.compile def fun(x): z -x print(z) # 崩溃 return mx.exp(z) fun(mx.array(5.0))这是因为占位符数组只用于构建计算图本身不携带数据。需要调试时检查中间数组的内容非常有用方法之一是全局禁用编译使用mx.disable_compile()函数或设置环境变量MLX_DISABLE_COMPILE。例如下面的代码即便fun是编译的也不会崩溃mx.compile def fun(x): z -x print(z) # 正常 return mx.exp(z) mx.disable_compile() fun(mx.array(5.0))disable_compile/enable_compile在 mlx/compile.h 中声明并绑定为mlx.core.disable_compile与mlx.core.enable_compile见 python/src/transforms.cpp。在 mlx/compile.cpp 中compile_mode()首次初始化时会检查环境变量MLX_DISABLE_COMPILE只要该变量被设置即便设为0也会生效属于按存在与否启用的变量编译模式即为disabled而运行时调用enable_compile()可以覆盖该环境变量。环境变量的完整语义记录在 environment_variables.rst布尔型变量用0关闭、非零整数开启但MLX_DISABLE_COMPILE特殊——只要存在即生效所以设置0也会禁用编译mlx.core.enable_compile可以覆盖它。此外编译模式还有更细粒度的控制CompileMode枚举disabled、no_simplify、no_fuse、enabled与set_compile_mode见 mlx/compile.h可供进阶诊断使用例如只跳过简化或只跳过融合以定位性能问题的来源。测试方面python/tests/test_compile.py 的test_enable_disable验证了通过mx.export_to_dot导出计算图并统计节点数禁用编译后节点数明显增多重新启用后恢复为编译时的节点数——这是观察编译合并/融合效果最直观的手段。四、纯函数约束副作用、隐式输入与隐式输出编译函数被设计为纯函数pure即不应该产生副作用。例如下面的代码会出问题state [] mx.compile def fun(x, y): z x y state.append(z) return mx.exp(z) fun(mx.array(1.0), mx.array(2.0)) # 崩溃 print(state)原因在于首次调用fun后state列表中保存的是一个占位符数组。占位符没有真实数据只用于构建计算图打印这样的数组会导致崩溃。针对编译函数内部更新外部容器这一需求文档给出了两种解决方案。方案一把 state 作为返回值输出state [] mx.compile def fun(x, y): z x y state.append(z) return mx.exp(z), state _, state fun(mx.array(1.0), mx.array(2.0)) # 输出 [array(3, dtypefloat32)] print(state)方案二用 outputs 参数捕获隐式输出有些场景下显式返回更新后的 state 很不方便。因此mx.compile提供了outputs参数来捕获隐式输出from functools import partial state [] # 告诉 compile 把 state 捕获为输出 partial(mx.compile, outputsstate) def fun(x, y): z x y state.append(z) return mx.exp(z) fun(mx.array(1.0), mx.array(2.0)) # 输出 [array(3, dtypefloat32)] print(state)这在编译包含更新容器逻辑的函数时特别有用——训练mlx.nn.Module参数时正是典型场景。在 python/src/transforms.cpp 的PyCompiledFun::call_impl中outputs捕获的数组会以tree_flatten方式扁平化后追加到函数输出末尾编译执行后通过tree_fill回写到原容器从而把占位符替换为真实数组。常量参数列表之外的输入被视为常量编译函数会把不在参数列表中的输入当作常量。例如state [mx.array(1.0)] mx.compile def fun(x): return x state[0] # 输出 array(2, dtypefloat32) print(fun(mx.array(1.0))) # 更新 state state[0] mx.array(5.0) # 仍然输出 array(2, dtypefloat32) print(fun(mx.array(1.0)))修改state后输出不变因为首次编译时state[0]的值常量已被固化进编译产物。这也对应 mlx/compile.cpp 中的逻辑compile_fuse会把非输入、无 primitive、大小为 1 的标量标记为constant_ids并参与生成 kernel 名称常量值本身通过constant_hasher哈希后拼入 kernel 名字见 mlx/compile.cpp因此常量变化会触发缓存键变化并重新编译。在 Python 层非数组参数float、int、str、None 以及 list/tuple/dict 的树形结构也会被编码进constants向量参与缓存匹配见 python/src/transforms.cpp。想让state的变化反映到输出同样有两种办法。办法一把 state 作为显式输入传入state [mx.array(1.0)] mx.compile def fun(x, state): return x state[0] # 输出 array(2, dtypefloat32) print(fun(mx.array(1.0), state)) # 更新 state state[0] mx.array(5.0) # 输出 array(6, dtypefloat32) print(fun(mx.array(1.0), state))办法二用 inputs 参数捕获隐式输入from functools import partial state [mx.array(1.0)] # 告诉 compile 把 state 捕获为输入 partial(mx.compile, inputsstate) def fun(x): return x state[0] # 输出 array(2, dtypefloat32) print(fun(mx.array(1.0))) # 更新 state state[0] mx.array(5.0) # 输出 array(6, dtypefloat32) print(fun(mx.array(1.0)))inputs捕获的数组会被追加到实际输入之后参与编译调用时先通过tree_fill将占位符填入捕获容器、执行后再用tree_replace还原见 python/src/transforms.cpp。inputs/outputs都支持 list 或 dict可包含任意嵌套的 list、dict 与 array非 array 的叶子节点会被忽略见 python/src/transforms.cpp 的签名说明。相关行为在 python/tests/test_compile.py 的test_compile_capture中有系统验证。五、编译完整训练图前向 反向 参数更新一步到位本节用一个常见的训练设置示例演示如何用mx.compile编译完整的前向、反向与参数更新流程使用mlx.nn.Module定义模型、mlx.optimizers.Optimizer维护带状态如动量的优化器。先看不编译的版本import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim # 4 个样本每个 10 维特征 x mx.random.uniform(shape(4, 10)) # 0、1 标签 y mx.array([0, 1, 0, 1]) # 简单的线性模型 model nn.Linear(10, 1) # 带动量的 SGD optimizer optim.SGD(learning_rate0.1, momentum0.8) def loss_fn(model, x, y): logits model(x).squeeze() return nn.losses.binary_cross_entropy(logits, y) loss_and_grad_fn nn.value_and_grad(model, loss_fn) # 执行 10 步梯度下降 for it in range(10): loss, grads loss_and_grad_fn(model, x, y) optimizer.update(model, grads) mx.eval(model.parameters(), optimizer.state)要编译更新这一步可以把整个更新过程放进一个函数并用合适的inputs/outputs捕获状态。下面是编译后的相同示例import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim from functools import partial # 4 个样本每个 10 维特征 x mx.random.uniform(shape(4, 10)) # 0、1 标签 y mx.array([0, 1, 0, 1]) # 简单的线性模型 model nn.Linear(10, 1) # 带动量的 SGD optimizer optim.SGD(learning_rate0.1, momentum0.8) def loss_fn(model, x, y): logits model(x).squeeze() return nn.losses.binary_cross_entropy(logits, y) # 将被捕获为输入和输出的状态 state [model.state, optimizer.state] partial(mx.compile, inputsstate, outputsstate) def step(x, y): loss_and_grad_fn nn.value_and_grad(model, loss_fn) loss, grads loss_and_grad_fn(model, x, y) optimizer.update(model, grads) return loss # 执行 10 步梯度下降 for it in range(10): loss step(x, y) # 求值模型与优化器状态 mx.eval(state) print(loss)这里的关键点inputsstate把模型参数与优化器状态含动量缓冲作为隐式输入捕获保证每次step都基于最新状态计算outputsstate把更新后的参数与状态作为隐式输出捕获并写回使model.state、optimizer.state在编译执行后携带真实数据每次迭代后仍需mx.eval(state)触发求值MLX 是惰性求值框架并打印 loss。注意如果使用的模块包含随机采样例如mlx.nn.Dropout请务必把mx.random.state也纳入compile捕获的state中即state [model.state, optimizer.state, mx.random.state]。否则随机数状态不会在编译函数中被更新导致每次调用得到相同的采样结果。这一行为在 python/tests/test_compile.py 的test_compile_rng系列测试中有专门覆盖inputsmx.random.state与outputsmx.random.state必须同时捕获随机状态才能在编译函数间正确传递。提示更多编译完整训练图的示例可参考 MLX 官方示例仓库mlx-examples。本文不展开外部链接训练图编译的通用模式即把 step 整体放进函数用inputs/outputs捕获可变状态。六、与其它函数变换的组合可组合的变换体系MLX 的函数变换是可组合的你可以把任意函数变换应用到任意其它函数变换的输出上。编译经过变换的函数与预期一致grad_fn mx.grad(mx.exp) compiled_grad_fn mx.compile(grad_fn) # 输出: array(2.71828, dtypefloat32) print(grad_fn(mx.array(1.0))) # 同样输出: array(2.71828, dtypefloat32) print(compiled_grad_fn(mx.array(1.0)))一个需要注意的默认行为是对编译函数再施加变换时变换后的函数默认不会被编译这是为了尽可能保留编译产物、避免重复编译。如果需要编译变换后的函数只需把它再传给mx.compile即可。例如test_compile_two_input_gradpython/tests/test_compile.py验证了mx.compile(mx.grad(loss))与直接求梯度结果一致test_vjp_vjp_compiled、test_vmap_compiledpython/tests/test_compile.py也验证了编译函数与vjp、jvp、vmap组合的正确性。你也可以编译那些内部调用了编译函数的函数。最佳实践是编译最外层的函数给compile最大机会去优化整个计算图mx.compile def inner(x): return mx.exp(-mx.abs(x)) def outer(x): inner(inner(x)) # 编译外层函数通常是更好的选择 # 因为即使内层函数已编译外层编译仍可能更快 fun mx.compile(outer)七、Shapeless 编译一次编译多变形状默认情况下编译函数的输入形状一旦改变就会重新编译。通过给mx.compile传入shapelessTrue可以只编译一次然后在任意形状的输入上运行def fun(x, y): return mx.abs(x y) compiled_fun mx.compile(fun, shapelessTrue) x mx.array(1.0) y mx.array(-2.0) # 首次调用触发编译 print(compiled_fun(x, y)) # 换用不同形状再次调用不会重新编译 x mx.array([1.0, -6.0]) y mx.array([-2.0, 3.0]) print(compiled_fun(x, y))从源码看shapeless的影响体现在两个层面缓存匹配在CompileCache::find中shapeless模式下比较输入时跳过 shape 检查只比较ndim与dtype见 mlx/compile.cpp输出形状推断compile_replace在shapeless模式下不再沿用追踪时记录的静态 shape而是调用每个 primitive 的output_shapes(real_inputs)依据真实输入推断输出形状见 mlx/compile.cpp。使用 shapeless 编译的注意事项请谨慎使用 shapeless 编译。由于形状变化不会触发重新编译任何依赖输入形状的条件分支图都不会按预期工作。形状相关的计算很常见而且有时很隐蔽例如def fun(x): return x.reshape(x.shape[0] * x.shape[1], -1) compiled_fun mx.compile(fun, shapelessTrue) x mx.random.uniform(shape(2, 3, 4)) out compiled_fun(x) x mx.random.uniform(shape(5, 5, 3)) # 报错无法将 (5, 5, 3) 变形为 (6, -1) out compiled_fun(x)第二次调用失败的原因是reshape使用了第一次调用时x的静态形状2 * 3 6而 shapeless 模式下图不会按新形状重新构建。解决办法是改用flatten避免在图中硬编码形状def fun(x): return x.flatten(0, 1) compiled_fun mx.compile(fun, shapelessTrue) x mx.random.uniform(shape(2, 3, 4)) out compiled_fun(x) x mx.random.uniform(shape(5, 5, 3)) # 正常 out compiled_fun(x)另外需要留意shapelessTrue并非适用于所有函数尝试编译不支持 shapeless 的函数会抛错并且即便启用shapeless改变输入的维度数ndim或类型仍会触发重新编译见 python/src/transforms.cpp 的签名文档。shapeless的边界情况在 python/tests/test_compile.py 中有大量测试覆盖包括与 broadcast、reduction、gather、full_like、量化矩阵乘等算子的组合可作为判断哪些图可以安全 shapeless 编译的参考。八、小结与最佳实践结合 compile.rst 与源码实现可以把mx.compile的使用要点归纳为编译可复用的函数首次编译成本高但结果会被缓存避免在循环中反复创建编译函数缓存键基于函数地址见 mlx/compile.cpp。警惕重编译触发条件输入 shape/ndim、dtype、输入个数的变化都会导致部分或全部重新编译。保持函数纯净编译函数内部不能对数组求值打印会崩溃有外部可变状态时用inputs/outputs捕获训练时记得把mx.random.state一并捕获。编译最外层让compile有机会融合尽可能多的算子受max_compile_depth11与max_compile_arrays24约束。调试用开关mx.disable_compile()或环境变量MLX_DISABLE_COMPILE可全局关闭编译编译模式在 mlx/compile.h 中定义为disabled / no_simplify / no_fuse / enabled四档enable_compile()可重新开启并覆盖环境变量。shapeless 谨慎用它跳过 shape 相关的重编译但要求计算图本身不依赖静态形状多用flatten而非手写reshape维度。组合变换时显式编译对编译函数再施加变换后默认不编译需要编译时显式再包一层mx.compile。进一步深入时可以参考 python/tests/test_compile.py 中的完整测试集常量、无穷值、闭包捕获、kwargs、多线程编译、随机状态捕获等场景以及 function_transforms.rst 中关于函数变换组合性的系统说明。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 MLX 编译指南：使用 mx.compile 合并计算图、融合算子并加速训练
MLX 编译指南使用 mx.compile 合并计算图、融合算子并加速训练【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlxMLXArray framework for Apple silicon在mlx.core中提供了compile这一函数变换function transformation用于编译计算图。函数编译通过合并重复的公共计算、融合特定算子来得到更小的计算图在多数场景下能显著降低运行时间与内存占用。本文以 compile.rst 为主体结合 mlx/compile.cpp 与 python/src/transforms.cpp 的源码实现系统讲解mx.compile的使用方式、缓存与重编译行为、纯函数约束、训练图编译、与其它变换的组合以及 shapeless 编译帮助你写出可安全编译、可复用、可调试的高性能 MLX 代码。一、编译的基本用法从普通函数到编译函数mx.compile的入门用法非常简单把一个以数组为输入、数组为输出的普通 Python 函数包装起来即可。下面的例子演示了普通调用与编译调用的区别def fun(x, y): return mx.exp(-x) y x mx.array(1.0) y mx.array(2.0) # 普通调用不编译 # 输出: array(2.36788, dtypefloat32) print(fun(x, y)) # 编译该函数 compiled_fun mx.compile(fun) # 输出: array(2.36788, dtypefloat32) print(compiled_fun(x, y))普通函数与编译函数的输出在数值精度范围内是一致的——compile是语义保持semantics-preserving的变换只改变执行的组织方式不改变计算结果。从底层看mx.compile返回的编译函数在 mlx/compile.cpp 中实现为一条完整的编译流水线首次调用时会依次执行compile_trace用占位符placeholder输入调用原函数追踪出计算图compile_dfs深度优先遍历计算图构建 tape操作序列与 parents map父节点映射compile_simplify简化 tape——合并相同标量、移除 no-opCopy、StopGradient、多轮合并等价子表达式默认 3 轮见 mlx/compile.cppcompile_fuse将可融合的子图替换为Compiled原语见 mlx/compile.cppcompile_replace把占位符替换为真实数组供后续求值。缓存机制编译一次多次复用编译本身有成本第一次调用编译函数时MLX 需要构建计算图、优化并生成/编译代码这个过程相对较慢。但 MLX 会对编译结果做缓存后续调用不会重新触发编译。因此建议只编译你打算多次调用的函数def fun(x, y): return mx.exp(-x) y x mx.array(1.0) y mx.array(2.0) compiled_fun mx.compile(fun) # 此处发生编译 compiled_fun(x, y) # 不再编译 compiled_fun(x, y) # 不再编译 mx.compile(fun)(x, y)缓存在 mlx/compile.cpp 的CompileCache中实现它以函数地址fun_id为主键逐条缓存 entry命中条件包括默认 stream/设备匹配、shapeless标志一致、输入数组的 shape 与 dtype 相等以及编译常量一致。缓存使用共享锁实现线程安全并且是线程局部的thread_local见 mlx/compile.cpp。何时会触发重新编译以下情况会导致函数被重新编译输入的形状或维度数量发生变化任一输入的类型发生变化函数的输入个数发生变化。其中部分情况只重跑编译栈中的某几层例如仅改变形状而另一些情况例如改变类型会重跑整条编译栈。一般来说应避免过于频繁地编译函数。相应地Python 绑定层在 python/src/transforms.cpp 的compile签名中提供了三个可选参数inputs、outputs与shapeless其中shapeless正是为缓解形状变化导致的重编译而设计的详见后文Shapeless 编译一节。需要警惕的反模式在循环中创建并销毁编译函数另一个容易踩坑的写法是编译那些被频繁创建又销毁的函数例如在循环内编译匿名函数a mx.array(1.0) # 不要这样做每一轮迭代都会重新编译这个 lambda for _ in range(5): mx.compile(lambda x: mx.exp(mx.abs(x)))(a)由于每次迭代都会新建一个 lambda 对象其函数地址缓存键随之变化缓存无法命中从而反复编译。应把编译函数提取到循环之外只创建一次。二、真实加速案例编译 GELUmlx.nn.gelu是 Transformer 类模型中常用的非线性激活函数其实现涉及多个一元unary与二元binary逐元素操作def gelu(x): return x * (1 mx.erf(x / math.sqrt(2))) / 2当输入数组很小时该函数受调用开销overhead-bound限制当输入数组很大时则受内存带宽memory bandwidth-bound限制。而gelu中的所有操作都是可融合的mx.compile可以将其融合进单个 kernel从而在这两种场景下都获得可观加速。下面用计时辅助函数对比普通函数与编译函数的运行时间该辅助函数先做 10 次 warm-up并在计时循环中通过mx.eval完成同步import time def timeit(fun, x): # warm up for _ in range(10): mx.eval(fun(x)) tic time.perf_counter() for _ in range(100): mx.eval(fun(x)) toc time.perf_counter() tpi 1e3 * (toc - tic) / 100 print(fTime per iteration {tpi:.3f} (ms))构造一个大数组并分别计时x mx.random.uniform(shape(32, 1000, 4096)) timeit(gelu, x) timeit(mx.compile(gelu), x)在 M1 Max 上普通gelu约为 15.5 毫秒编译后的gelu约为 3.1 毫秒编译版本快约 5 倍该数值来自文档在 M1 Max 上的实测具体结果随设备与数组规模而异。从源码可以印证哪些算子可融合的判断逻辑在 mlx/compile.cpp 中is_unary覆盖Exp、Erf、Negative、Log、Sigmoid等一元算子is_binary覆盖Add、Multiply、Divide、Subtract等二元算子is_ternary覆盖Selectis_broadcast覆盖Broadcast而is_fusable正是这四类的并集。compile_fuse在反向遍历 tape 时递归收集可融合的算子并受两个常量约束最大融合深度max_compile_depth 11、最多输入数组数max_compile_arrays 24见 mlx/compile.cpp。因此像gelu这种由一元/二元逐元素算子串成的计算链可以整体融合成一个Compiled原语只需在编译阶段为该子图生成一个 kernel 即可。三、调试编译函数占位符、disable_compile 与 MLX_DISABLE_COMPILE编译函数在首次被调用时是用占位符输入进行追踪tracing的。这意味着在编译函数内部不能对数组求值例如打印数组内容否则会崩溃mx.compile def fun(x): z -x print(z) # 崩溃 return mx.exp(z) fun(mx.array(5.0))这是因为占位符数组只用于构建计算图本身不携带数据。需要调试时检查中间数组的内容非常有用方法之一是全局禁用编译使用mx.disable_compile()函数或设置环境变量MLX_DISABLE_COMPILE。例如下面的代码即便fun是编译的也不会崩溃mx.compile def fun(x): z -x print(z) # 正常 return mx.exp(z) mx.disable_compile() fun(mx.array(5.0))disable_compile/enable_compile在 mlx/compile.h 中声明并绑定为mlx.core.disable_compile与mlx.core.enable_compile见 python/src/transforms.cpp。在 mlx/compile.cpp 中compile_mode()首次初始化时会检查环境变量MLX_DISABLE_COMPILE只要该变量被设置即便设为0也会生效属于按存在与否启用的变量编译模式即为disabled而运行时调用enable_compile()可以覆盖该环境变量。环境变量的完整语义记录在 environment_variables.rst布尔型变量用0关闭、非零整数开启但MLX_DISABLE_COMPILE特殊——只要存在即生效所以设置0也会禁用编译mlx.core.enable_compile可以覆盖它。此外编译模式还有更细粒度的控制CompileMode枚举disabled、no_simplify、no_fuse、enabled与set_compile_mode见 mlx/compile.h可供进阶诊断使用例如只跳过简化或只跳过融合以定位性能问题的来源。测试方面python/tests/test_compile.py 的test_enable_disable验证了通过mx.export_to_dot导出计算图并统计节点数禁用编译后节点数明显增多重新启用后恢复为编译时的节点数——这是观察编译合并/融合效果最直观的手段。四、纯函数约束副作用、隐式输入与隐式输出编译函数被设计为纯函数pure即不应该产生副作用。例如下面的代码会出问题state [] mx.compile def fun(x, y): z x y state.append(z) return mx.exp(z) fun(mx.array(1.0), mx.array(2.0)) # 崩溃 print(state)原因在于首次调用fun后state列表中保存的是一个占位符数组。占位符没有真实数据只用于构建计算图打印这样的数组会导致崩溃。针对编译函数内部更新外部容器这一需求文档给出了两种解决方案。方案一把 state 作为返回值输出state [] mx.compile def fun(x, y): z x y state.append(z) return mx.exp(z), state _, state fun(mx.array(1.0), mx.array(2.0)) # 输出 [array(3, dtypefloat32)] print(state)方案二用 outputs 参数捕获隐式输出有些场景下显式返回更新后的 state 很不方便。因此mx.compile提供了outputs参数来捕获隐式输出from functools import partial state [] # 告诉 compile 把 state 捕获为输出 partial(mx.compile, outputsstate) def fun(x, y): z x y state.append(z) return mx.exp(z) fun(mx.array(1.0), mx.array(2.0)) # 输出 [array(3, dtypefloat32)] print(state)这在编译包含更新容器逻辑的函数时特别有用——训练mlx.nn.Module参数时正是典型场景。在 python/src/transforms.cpp 的PyCompiledFun::call_impl中outputs捕获的数组会以tree_flatten方式扁平化后追加到函数输出末尾编译执行后通过tree_fill回写到原容器从而把占位符替换为真实数组。常量参数列表之外的输入被视为常量编译函数会把不在参数列表中的输入当作常量。例如state [mx.array(1.0)] mx.compile def fun(x): return x state[0] # 输出 array(2, dtypefloat32) print(fun(mx.array(1.0))) # 更新 state state[0] mx.array(5.0) # 仍然输出 array(2, dtypefloat32) print(fun(mx.array(1.0)))修改state后输出不变因为首次编译时state[0]的值常量已被固化进编译产物。这也对应 mlx/compile.cpp 中的逻辑compile_fuse会把非输入、无 primitive、大小为 1 的标量标记为constant_ids并参与生成 kernel 名称常量值本身通过constant_hasher哈希后拼入 kernel 名字见 mlx/compile.cpp因此常量变化会触发缓存键变化并重新编译。在 Python 层非数组参数float、int、str、None 以及 list/tuple/dict 的树形结构也会被编码进constants向量参与缓存匹配见 python/src/transforms.cpp。想让state的变化反映到输出同样有两种办法。办法一把 state 作为显式输入传入state [mx.array(1.0)] mx.compile def fun(x, state): return x state[0] # 输出 array(2, dtypefloat32) print(fun(mx.array(1.0), state)) # 更新 state state[0] mx.array(5.0) # 输出 array(6, dtypefloat32) print(fun(mx.array(1.0), state))办法二用 inputs 参数捕获隐式输入from functools import partial state [mx.array(1.0)] # 告诉 compile 把 state 捕获为输入 partial(mx.compile, inputsstate) def fun(x): return x state[0] # 输出 array(2, dtypefloat32) print(fun(mx.array(1.0))) # 更新 state state[0] mx.array(5.0) # 输出 array(6, dtypefloat32) print(fun(mx.array(1.0)))inputs捕获的数组会被追加到实际输入之后参与编译调用时先通过tree_fill将占位符填入捕获容器、执行后再用tree_replace还原见 python/src/transforms.cpp。inputs/outputs都支持 list 或 dict可包含任意嵌套的 list、dict 与 array非 array 的叶子节点会被忽略见 python/src/transforms.cpp 的签名说明。相关行为在 python/tests/test_compile.py 的test_compile_capture中有系统验证。五、编译完整训练图前向 反向 参数更新一步到位本节用一个常见的训练设置示例演示如何用mx.compile编译完整的前向、反向与参数更新流程使用mlx.nn.Module定义模型、mlx.optimizers.Optimizer维护带状态如动量的优化器。先看不编译的版本import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim # 4 个样本每个 10 维特征 x mx.random.uniform(shape(4, 10)) # 0、1 标签 y mx.array([0, 1, 0, 1]) # 简单的线性模型 model nn.Linear(10, 1) # 带动量的 SGD optimizer optim.SGD(learning_rate0.1, momentum0.8) def loss_fn(model, x, y): logits model(x).squeeze() return nn.losses.binary_cross_entropy(logits, y) loss_and_grad_fn nn.value_and_grad(model, loss_fn) # 执行 10 步梯度下降 for it in range(10): loss, grads loss_and_grad_fn(model, x, y) optimizer.update(model, grads) mx.eval(model.parameters(), optimizer.state)要编译更新这一步可以把整个更新过程放进一个函数并用合适的inputs/outputs捕获状态。下面是编译后的相同示例import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim from functools import partial # 4 个样本每个 10 维特征 x mx.random.uniform(shape(4, 10)) # 0、1 标签 y mx.array([0, 1, 0, 1]) # 简单的线性模型 model nn.Linear(10, 1) # 带动量的 SGD optimizer optim.SGD(learning_rate0.1, momentum0.8) def loss_fn(model, x, y): logits model(x).squeeze() return nn.losses.binary_cross_entropy(logits, y) # 将被捕获为输入和输出的状态 state [model.state, optimizer.state] partial(mx.compile, inputsstate, outputsstate) def step(x, y): loss_and_grad_fn nn.value_and_grad(model, loss_fn) loss, grads loss_and_grad_fn(model, x, y) optimizer.update(model, grads) return loss # 执行 10 步梯度下降 for it in range(10): loss step(x, y) # 求值模型与优化器状态 mx.eval(state) print(loss)这里的关键点inputsstate把模型参数与优化器状态含动量缓冲作为隐式输入捕获保证每次step都基于最新状态计算outputsstate把更新后的参数与状态作为隐式输出捕获并写回使model.state、optimizer.state在编译执行后携带真实数据每次迭代后仍需mx.eval(state)触发求值MLX 是惰性求值框架并打印 loss。注意如果使用的模块包含随机采样例如mlx.nn.Dropout请务必把mx.random.state也纳入compile捕获的state中即state [model.state, optimizer.state, mx.random.state]。否则随机数状态不会在编译函数中被更新导致每次调用得到相同的采样结果。这一行为在 python/tests/test_compile.py 的test_compile_rng系列测试中有专门覆盖inputsmx.random.state与outputsmx.random.state必须同时捕获随机状态才能在编译函数间正确传递。提示更多编译完整训练图的示例可参考 MLX 官方示例仓库mlx-examples。本文不展开外部链接训练图编译的通用模式即把 step 整体放进函数用inputs/outputs捕获可变状态。六、与其它函数变换的组合可组合的变换体系MLX 的函数变换是可组合的你可以把任意函数变换应用到任意其它函数变换的输出上。编译经过变换的函数与预期一致grad_fn mx.grad(mx.exp) compiled_grad_fn mx.compile(grad_fn) # 输出: array(2.71828, dtypefloat32) print(grad_fn(mx.array(1.0))) # 同样输出: array(2.71828, dtypefloat32) print(compiled_grad_fn(mx.array(1.0)))一个需要注意的默认行为是对编译函数再施加变换时变换后的函数默认不会被编译这是为了尽可能保留编译产物、避免重复编译。如果需要编译变换后的函数只需把它再传给mx.compile即可。例如test_compile_two_input_gradpython/tests/test_compile.py验证了mx.compile(mx.grad(loss))与直接求梯度结果一致test_vjp_vjp_compiled、test_vmap_compiledpython/tests/test_compile.py也验证了编译函数与vjp、jvp、vmap组合的正确性。你也可以编译那些内部调用了编译函数的函数。最佳实践是编译最外层的函数给compile最大机会去优化整个计算图mx.compile def inner(x): return mx.exp(-mx.abs(x)) def outer(x): inner(inner(x)) # 编译外层函数通常是更好的选择 # 因为即使内层函数已编译外层编译仍可能更快 fun mx.compile(outer)七、Shapeless 编译一次编译多变形状默认情况下编译函数的输入形状一旦改变就会重新编译。通过给mx.compile传入shapelessTrue可以只编译一次然后在任意形状的输入上运行def fun(x, y): return mx.abs(x y) compiled_fun mx.compile(fun, shapelessTrue) x mx.array(1.0) y mx.array(-2.0) # 首次调用触发编译 print(compiled_fun(x, y)) # 换用不同形状再次调用不会重新编译 x mx.array([1.0, -6.0]) y mx.array([-2.0, 3.0]) print(compiled_fun(x, y))从源码看shapeless的影响体现在两个层面缓存匹配在CompileCache::find中shapeless模式下比较输入时跳过 shape 检查只比较ndim与dtype见 mlx/compile.cpp输出形状推断compile_replace在shapeless模式下不再沿用追踪时记录的静态 shape而是调用每个 primitive 的output_shapes(real_inputs)依据真实输入推断输出形状见 mlx/compile.cpp。使用 shapeless 编译的注意事项请谨慎使用 shapeless 编译。由于形状变化不会触发重新编译任何依赖输入形状的条件分支图都不会按预期工作。形状相关的计算很常见而且有时很隐蔽例如def fun(x): return x.reshape(x.shape[0] * x.shape[1], -1) compiled_fun mx.compile(fun, shapelessTrue) x mx.random.uniform(shape(2, 3, 4)) out compiled_fun(x) x mx.random.uniform(shape(5, 5, 3)) # 报错无法将 (5, 5, 3) 变形为 (6, -1) out compiled_fun(x)第二次调用失败的原因是reshape使用了第一次调用时x的静态形状2 * 3 6而 shapeless 模式下图不会按新形状重新构建。解决办法是改用flatten避免在图中硬编码形状def fun(x): return x.flatten(0, 1) compiled_fun mx.compile(fun, shapelessTrue) x mx.random.uniform(shape(2, 3, 4)) out compiled_fun(x) x mx.random.uniform(shape(5, 5, 3)) # 正常 out compiled_fun(x)另外需要留意shapelessTrue并非适用于所有函数尝试编译不支持 shapeless 的函数会抛错并且即便启用shapeless改变输入的维度数ndim或类型仍会触发重新编译见 python/src/transforms.cpp 的签名文档。shapeless的边界情况在 python/tests/test_compile.py 中有大量测试覆盖包括与 broadcast、reduction、gather、full_like、量化矩阵乘等算子的组合可作为判断哪些图可以安全 shapeless 编译的参考。八、小结与最佳实践结合 compile.rst 与源码实现可以把mx.compile的使用要点归纳为编译可复用的函数首次编译成本高但结果会被缓存避免在循环中反复创建编译函数缓存键基于函数地址见 mlx/compile.cpp。警惕重编译触发条件输入 shape/ndim、dtype、输入个数的变化都会导致部分或全部重新编译。保持函数纯净编译函数内部不能对数组求值打印会崩溃有外部可变状态时用inputs/outputs捕获训练时记得把mx.random.state一并捕获。编译最外层让compile有机会融合尽可能多的算子受max_compile_depth11与max_compile_arrays24约束。调试用开关mx.disable_compile()或环境变量MLX_DISABLE_COMPILE可全局关闭编译编译模式在 mlx/compile.h 中定义为disabled / no_simplify / no_fuse / enabled四档enable_compile()可重新开启并覆盖环境变量。shapeless 谨慎用它跳过 shape 相关的重编译但要求计算图本身不依赖静态形状多用flatten而非手写reshape维度。组合变换时显式编译对编译函数再施加变换后默认不编译需要编译时显式再包一层mx.compile。进一步深入时可以参考 python/tests/test_compile.py 中的完整测试集常量、无穷值、闭包捕获、kwargs、多线程编译、随机状态捕获等场景以及 function_transforms.rst 中关于函数变换组合性的系统说明。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考