Combine Metal 4 machine learning and graphics
Graphics & Games 进阶 2m

Metal 4 ML 与图形融合：Tensor、ML Encoder 和 Shader ML

Combine Metal 4 machine learning and graphics

2025年6月9日

在 Apple 官方观看视频

一句话判断

Metal 4 把 ML 正式纳入了 GPU 渲染管线——MTLTensor 解决了多维数据的资源管理，ML Command Encoder 让整网推理和渲染/计算命令在同一时间线上同步执行，而 Shader ML 允许你在 fragment shader 内部直接跑神经网络。这套组合拳瞄准的是实时渲染中 ML 最密集的三个场景：超分、材质压缩和光照计算。

这场 Session 讲了什么

分四个部分，由 Preston 和 Scott 两位工程师讲解：

MTLTensor：Metal 4 新引入的多维数据容器。与 MTLBuffer 的线性布局和 MTLTexture 的 4 通道上限不同，MTLTensor 支持任意维度、内建 stride 和维度信息，简化多维数据索引。可以从 MTLDevice（最优性能，不透明布局优化）或 MTLBuffer（需手动指定 stride）创建。Usage 属性支持 MachineLearning / Compute / Render 及其组合。

MTL4MachineLearningCommandEncoder：整网推理的 GPU 时间线编码器。离线流程：CoreML 包通过 metal-package-builder 转为 MTLPackage，加载为 MTLLibrary，编译为 MTL4MachineLearningPipelineState（可指定动态输入尺寸）。运行时：创建 encoder → 设置 pipeline state → 绑定输入/输出 MTLTensor → 用 dispatchNetworkWithIntermediatesHeap 调度。Heap 用 MTLHeapTypePlacement 类型复用中间资源。同步用标准 Metal 4 原语（MTLFence、MTLBarrier），新增 MTLStageMachineLearning 标识 ML 阶段。

Shader ML：在 shader 内部嵌入 ML 操作。核心是 Metal Performance Primitives 库——提供 matmul2d、卷积等高性能 MTLTensor 操作。可以在 fragment shader 中创建 inline MTLTensor（采样潜变量纹理），通过 matmul2d_descriptor 配置矩阵乘法（指定问题规模、转置、精度、执行线程数），执行推理后用输出值做 shading。典型用例是神经材质压缩——在 fragment shader 内完成潜变量采样→网络推理→材质解压→着色，全程不离开 GPU 线程，实现 50% 压缩比。

Metal Debugger ML 工具：Dependency Viewer 验证同步关系；MTLTensor Viewer 检查输入/输出数据；ML Network Debugger 展示网络图结构，可以逐步检查每层操作的中间张量，展开 stitched region 查看子操作。Scott 的调试案例：SignedSmoothstep 函数中多了一个 * 号导致乘法变幂运算，通过二分法定位到具体操作层。

值得深挖的点

MTLTensor 的设计哲学是”为 ML 优化的纹理”。MTLTexture 受限于维度数（最多 4D）和通道数（最多 4），不适合 ML 中常见的高维张量（比如卷积的 batch×channel×height×width）。MTLBuffer 虽然灵活但没有内建的多维索引。MTLTensor 把维度、stride、数据类型封装到对象内部，索引计算自动完成——这是 ML 工作负载最需要的抽象。

Shader ML 的执行组选择是性能关键。每个 fragment thread 独立执行同一份网络但不同数据时用 execution_thread。如果整个 simdgroup 或 threadgroup 操作同一份数据，可以选择更大的执行组利用硬件并行。但如果操作存在数据发散或非均匀控制流，必须用单线程执行组——否则结果未定义。

神经材质压缩的对比效果很直观。Base Color 隔离对比下几乎看不出差异，但神经材质只占 block compressed 格式 50% 的内存和磁盘空间。这个技术对于移动端游戏和 VR 应用的资源优化意义重大——同样的存储预算可以装两倍的材质精度。

调试器的 ML Network Debugger 是新增的独立工具。以前调试 ML 网络需要导出中间张量到 Python 环境分析。现在直接在 Metal Debugger 里点击网络图的任意节点就能查看中间张量数据，展开 stitched region 看子操作。这个工具链闭环大大缩短了”发现问题→定位问题→修复问题”的循环。

代码片段

// Shader ML：在 fragment shader 中做神经材质压缩
#include <metal_tensor>
#include <metal_performance_primitives>

fragment float4 neuralMaterialShader(
    texture2d<float> latentTex0 [[texture(0)]],
    texture2d<float> latentTex1 [[texture(1)]],
    texture2d<float> latentTex2 [[texture(2)]],
    texture2d<float> latentTex3 [[texture(3)]],
    device metal::tensor<float, metal::dextents<int, 2>> &weights [[buffer(0)]],
    float2 uv [[texture_coord]])
{
    // 1. 采样潜变量纹理，创建 inline MTLTensor
    float4 latent = float4(
        latentTex0.sample(sampler, uv).x,
        latentTex1.sample(sampler, uv).x,
        latentTex2.sample(sampler, uv).x,
        latentTex3.sample(sampler, uv).x
    );
    auto input = metal::make_tensor(latent);

    // 2. 配置矩阵乘法（单线程执行）
    auto desc = metal::matmul2d_descriptor<4, 16, false, false, metal::matmul_precision::high>();
    auto result = metal::matmul2d<1>(desc, input, weights);

    // 3. ReLU 激活
    metal::elementwise_relu_inplace(result);

    // 4. 第二层推理...
    // 5. 用解压后的材质值做 shading
    return shade(result);
}
// 创建 MTLPackage 的命令行流程
// Step 1: Python 中导出 CoreML 包
// import coremltools as ct
// model = ct.convert(pytorch_model)
// model.save("MyModel.mlpackage")

// Step 2: 转为 MTLPackage
// metal-package-builder MyModel.mlpackage -o MyModel.metallib

最佳实践

  1. 优先从 MTLDevice 创建 MTLTensor。性能最优，系统会自动选择硬件优化的内存布局。只有在需要复用已有 MTLBuffer 数据时才从 Buffer 创建。

  2. 同步粒度要精确到消费端。不需要 ML 输出的工作可以在 ML 推理期间并行执行。只让真正消费 ML 输出的 render/compute pass 等待 ML 阶段完成。

  3. Shader ML 适用于小网络。整网推理（超分、遮挡计算）用 ML Command Encoder。碎片化的小网络（材质压缩、逐像素光照）用 Shader ML。判断标准：网络是否需要在单次 shader dispatch 中完成且不需要设备内存中间同步。

  4. Shader ML 的执行组选择遵循发散性原则。无数据发散 → 用 simdgroup/threadgroup 最大化并行。有数据发散或非均匀控制流 → 只能用单线程 execution_thread。

  5. 调试 ML 管线先检查同步再检查数据。用 Dependency Viewer 排除同步问题，再用 MTLTensor Viewer 检查输入输出，最后用 ML Network Debugger 定位网络内部问题。

还有什么值得关注

  • 只有 ML Program 格式的 CoreML 包才受支持，旧格式需要先转换
  • MTLTensor 的 Usage 属性可以组合（MachineLearning | Compute | Render），和 MTLTexture 的 usage 设计一致
  • Metal Performance Primitives 不仅提供 matmul2d，还有卷积操作
  • ML Network Debugger 支持展开 stitched region（融合子图），逐层查看操作详情
  • 神经材质压缩的 fragment shader 示例展示了典型的”采样→推理→着色”单 pass 流程
图形与游戏 机器学习