Metal Debugger

性能分析是性能优化的关键步骤。你可以使用 MLX_METAL_DEBUG 选项构建 MLX，以改进 Metal 调试和优化工作流程。MLX_METAL_DEBUG 调试选项

• 在 Metal 编译期间记录源代码，以便在调试时进行检查。

• 标记 Metal 对象（如命令队列），提高捕获的可读性。

要在 Python 中启用调试构建，请在构建调用前加上 CMAKE_ARGS="-DMLX_METAL_DEBUG=ON"。

函数 metal.start_capture() 启动对所有 MLX GPU 工作的捕获。

注意

要捕获 GPU 跟踪，必须在运行应用程序时设置环境变量 MTL_CAPTURE_ENABLED=1。

import mlx.core as mx

a = mx.random.uniform(shape=(512, 512))
b = mx.random.uniform(shape=(512, 512))
mx.eval(a, b)

trace_file = "mlx_trace.gputrace"

# Make sure to run with MTL_CAPTURE_ENABLED=1 and
# that the path trace_file does not already exist.
mx.metal.start_capture(trace_file)

for _ in range(10):
  mx.eval(mx.add(a, b))

mx.metal.stop_capture()

你可以在 Xcode 中打开并回放 GPU 跟踪。Dependencies 视图提供了所有操作的良好概览。有关更多信息，请查阅 Metal 调试器文档。

Xcode 工作流程

你可以通过在 Xcode 中运行来跳过保存到路径。首先，使用 CMake 生成一个 Xcode 项目。

mkdir build && cd build
cmake .. -DMLX_METAL_DEBUG=ON -G Xcode
open mlx.xcodeproj

选择 metal_capture 示例方案并运行。