实现 Metal 4 机器学习与图形应用程序的完美融合

大家好！ 我是 Apple Metal 框架团队的工程师 Preston Provins 稍后我的同事 Scott 也会加入这个讲座 我将介绍 Metal 的新增功能 这些功能实现了 机器学习与游戏的结合 Scott 将介绍GPU 工具的新增功能 这些功能旨在增强 Metal 4 中机器学习的调试体验 我很高兴能在这个讲座中向大家介绍 如何将 Metal 4 机器学习与图形相结合 如果你有兴趣了解Metal 4 提供的所有功能 请观看 Metal 4 基础知识讲座了解 Metal 4 的其他新功能 机器学习正在通过升频采样、素材压缩、 动画融合和神经网络着色等技术 彻底改变游戏和图形 这些技术推动着创意与沉浸感领域的边界不断拓展 它们能够模拟复杂的现象提高视觉保真度 并简化对全新风格和效果的探索过程 CoreML 非常适合多种机器学习任务 例如分割、分类、生成式 AI 等 它可以帮助你轻松创作机器学习模型 如果你的机器学习应用程序 需要与 GPU 时间线紧密集成那么 Metal 4 可以满足你的需求 在典型的帧中 游戏可能会在计算处理阶段执行顶点蒙皮 在渲染处理阶段对 场景进行光栅化渲染 并在另一个计算处理阶段应用抗锯齿处理 通常采用临时抗锯齿等图像处理技术 来完成抗锯齿处理 前沿技术用机器学习网络取代了 这些传统方法 这个网络能够对图像进行升频采样 让你能够以较低分辨率执行其余的渲染 从而提升性能

在着色器内执行微型神经网络这一方法也变得越来越普遍 例如 传统的片段着色器会对材质纹理进行采样 但突破性技术可使用小型神经网络 即时进行纹理解压缩并实现更高的压缩比 这种神经网络渲染技术可将材质集占用的空间 压缩到块压缩格式所占用空间的 50% 在这个讲座中 我们将了解 MTLTensor 这是 Metal 4 的新资源适用于机器学习工作流程 我们将深入了解全新MTL4MachineLearningCommandEncoder 它在 GPU 时间线上运行整个网络 我们还会深入了解 你的其他绘制和派发操作 我们将介绍 Shader ML 它让你能够将机器学习运算嵌入到自己的着色器中 最后 我们将介绍如何使用 Metal Debugger 将机器学习 无缝集成到你的应用程序中 你已经非常熟悉 MTLBuffer和 MTLTexture 了 今年 Metal 4 推出了 MTLTensor 这一新资源让你能够非常便捷地将机器学习应用于数据 MTLTensor 是一种基本的机器学习数据类型 可用于计算 图形和机器学习场景 机器学习工作负载将广泛使用张量 MTL4MachineLearningCommandEncoder使用 MTLTensor 来表示输入和输出 而 Shader ML使用 MTLTensor 来表示 权重以及输入和输出

MTLTensor 是多维数据容器 由阶和各个阶对应的维数来描述 MTLTensor 能够很好地突破二维限制扩展到更高维度 让你能够灵活描述实际使用机器学习时 所需的任何数据布局 例如 MTLTexture 限制为最多四个通道 并且对依赖于纹理格式的范围存在严格限制 机器学习中常常使用维度超过二维的数据 例如卷积运算 要使用数据的平面表示形式比如 MTLBuffer 就需要为具有多个维度的数据 建立复杂的索引模式 在 MTLTensor 中进行多维数据索引是一种更简单的方法 因为各个阶的步长和维度 会内置在 MTLTensor 对象中 并自动用于索引计算 我们来完成一遍MTLTensor 的创建过程 MTLTensor 的阶描述了它有多少个轴 这个 MTLTensor 的阶为 2 它包含了多行多列数据 维度的范围描述了沿这个轴的数据点个数 dataType 属性定义了MTLTensor 所包装数据的格式 Usage 属性表明将如何利用 MTLTensor 包括表示用于 MTL4MachineLearningCommandEncoder 的 MTLTensorUsageMachineLearning MTLTensorUsageCompute 还有表示在着色器程序内使用的 MTLTensorUsageRender 还可以将多个使用方式组合起来比如多个纹理的 usage 属性 以上都是重要的 MTLTensor 属性 应在 MTLTensorDescriptor对象上填充这些数据 接下来 我们要在代码中 创建一个 MTLTensor 填充描述符属性之后 通过在 MTLDevice 对象上调用newTensorWithDescriptor:offset:error: 来创建新的 MTLTensor MTLTensor 是从 MTLDevice或 MTLBuffer 对象创建的 但是 从设备创建的MTLTensor 可提供最佳性能 和 MTLTexture 的重排方式类似从 MTLDevice 对象创建 MTLTensor 会生成在读写方面进行优化的不透明布局 接下来 我们重点关注如何从预先存在的 MTLBuffer 创建 MTLTensor 不同于从 MTLDevice 创建的 MTLTensor 从 MTLBuffer 创建的 MTLTensor不会被假设为密排 所以你需要指定它的步长 最靠近中心的步长应该始终为 1 第二个步长表明当行索引递增时 跳过的元素数量

源 MTLBuffer 可以包含填充 例如行末尾未使用的列 需要考虑到填充 以便MTLTensor 包装适当的元素 要从底层缓冲区创建 MTLTensor 需要像对设备分配的张量那样设置 dataType 和 usage 属性 然后填写 MTLTensorDescriptor的 strides 属性 因此 生成的 MTLTensor 会适当地包装 MTLBuffer 的内容 最后 在源 MTLBuffer 上使用newTensorWithDescriptor:offset:error: 现在 我们知道了如何分配和创建 MTLTensor 接下来 我们深入了解一下全新机器学习编码器 它可以将机器学习作业 添加到 GPU 时间线 Metal 4 可以分别通过 MTL4ComputeCommandEncoder和 MTL4RenderCommandEncoder 轻松将计算和渲染命令 添加到 GPU 时间线 今年 通过将机器学习作业添加到 GPU 时间线 我们将实现进一步的统一 MTL4MachineLearningCommandEncoder支持并排运行多个完整模型 并与 GPU 上的其他 Metal 命令进行同步 可确保与其他 MTLCommand 无缝集成 这个全新编码器用于编码机器学习命令 它具有与计算编码器和渲染编码器类似的接口 Metal 4 同步原语 也可以使用机器学习命令进行运算 就像使用计算和渲染命令一样 通过进行同步 可以实现对作业编排的控制 并且有助于实现并行化以保持高性能 MTL4MachineLearningCommandEncoder创建工作流程可以分为两部分： 脱机和运行时 这个工作流程的脱机部分在应用程序启动之前进行 而运行时部分会在应用程序生命周期内进行 例如在某一帧中间 我们先来看看这个工作流程的脱机部分 创建 MTLPackage MTLPackage 是一种容器用于一个或多个函数 其中每个函数表示一个机器学习网络 你可以在 Metal 中使用这个网络来执行机器学习作业 这个格式在使用 Metal 进行载入和执行方面进行了优化 要创建 MTLPackage首先需要一个 CoreML 软件包 在这里 我们使用 CoreML 转换器 将在 PyTorch 或 TensorFlow 中 编写的机器学习框架 转换为 CoreML 软件包 这个示例展示了如何使用 Python中的 CoreML Tools 库 导出 PyTorch 模型 只需导入 Tools 库 然后对模型运行 convert 命令即可生成导出 最后 将这个导出另存为 ML 软件包 在这里 我想强调一点 并非每个 CoreML 软件包 都是 ML 程序 受支持的只有 ML 程序 如果 CoreML 软件包是在旧版操作系统上导出的 请查看这篇文章 以进一步了解如何将这些 CoreML 模型文件 导出为 ML 软件包 创建 CoreML 软件包后 接下来的步骤非常简单 只需针对已保存的模型运行 metal-package-builder 命令行以生成 MTLPackage 这会将 CoreML 软件包转换为 可在运行时高效载入的格式 以上就是创建MTLPackage 的整个过程 这个工作流程的脱机部分已经完成 其余部分都在运行时进行 为了对网络进行编译 首先将 MTLPackage 作为 MTLLibrary 打开 使用软件包中表示网络的函数的名称 来创建函数描述符 在这个示例中为 main 函数 为了对网络进行编译 创建MTL4MachineLearningPipelineState 需要使用 MTL4MachineLearningPipelineStateDescriptor 以及函数描述符来完成这个步骤 如果网络具有动态输入 请指定 MTL4MachineLearningPipelineStateDescriptor上每个输入的大小 通过使用 MTL4MachineLearningPipelineStateDescriptor创建 MTL4MachineLearningPipelineState 针对特定设备进行网络编译

这就是创建 MTL4MachineLearningPipelineState 的方法 下一步是创建MTL4MachineLearningCommandEncoder 并对作业进行编码 我们来深入了解一下 如何使用 MTL4MachineLearningCommandEncoder 对象 将作业派发到 GPU 时间线上 只需创建MTL4MachineLearningCommandEncoder 对象 就像创建用于计算或渲染的编码器一样 设置创建好的MTL4MachineLearningPipelineState 对象 并对使用的输入和输出进行绑定 最后 使用 dispatchNetworkWithIntermediatesHeap方法来派发作业

这个机器学习编码器使用堆 来存储多个运算之间的中间数据 而不是创建和释放缓冲区 因此可以将资源重复用于不同的派发 为了创建这个 MTLHeap 需要创建 MTLHeapDescriptor 并将 type 属性设置为MTLHeapTypePlacement 可以通过查询管道的intermediateHeapSize 获取网络的最小堆大小 并将堆的 size 属性设置为大于或等于这个大小 对网络派发进行编码后 结束编码并提交命令以在 GPU 时间线上运行

正如前面提到的 Metal4 同步原语 也可以使用机器学习命令进行运算 就像使用计算和渲染命令一样 如果正确地进行了同步那么不依赖于机器学习输出的作业 可以同时进行

只有消耗网络输出的作业才需要 等待排程机器学习作业运行完毕

为了同步 MTL4MachineLearningCommandEncoder 派发 可以使用标准的 Metal 4 同步原语例如 MTLBarriers 和 MTLFences 全新 MTLStageMachineLearning用于识别屏障中的机器学习工作负载 例如 为了让渲染作业等待网络生成的输出 你可以在适当的渲染阶段和机器学习阶段之间 使用屏障 我们来看看 MTL4MachineLearningCommandEncoder 的实际应用 这个示例使用了MTL4MachineLearningCommandEncoder 来派发全卷积网络 以预测每个像素的遮挡值 进行这一评估需要进行非常细致的同步 深度缓冲区和视图空间法线 是在启动机器学习工作负载之前填充的 在网络处理数据时 渲染器会并行派发与渲染相关的其他任务 并等待神经网络输出的结果然后再合成最终的帧 MTL4MachineLearningCommandEncoder不仅限于处理游戏的全帧信息 还可以用于任何符合 实时预算要求的网络 你可以根据自己的集成需求 以最适合的方式利用 Metal 4 同步原语 因此 借助 Metal 4 的 MTL4MachineLearningCommandEncoder 可以轻松在 GPU 时间线上运行大型机器学习工作负载 总结一下 借助 MTL4MachineLearningCommandEncoder机器学习在 Metal 4 中实现了与计算和渲染的整合 MTL4MachineLearningCommandEncoder 让整个网络能够在 GPU 时间线上运行 资源可以与其他 GPU 命令共享 此外 一组强大的 Metal 4 同步原语 可实现高性能机器学习功能 Metal 4 还引入了 Shader ML 用于将较小的机器学习运算 嵌入到现有的内核和着色器内 一些走在前沿的游戏正在采用机器学习 来取代传统的渲染算法 基于机器学习的技术 可提供全局照明、 材质着色、几何压缩 材料压缩等解决方案 这些技术往往可以提升性能或减少内存占用 作为一个鼓舞人心的例子我们来考虑一下神经网络材质压缩 这种技术与块压缩格式相比可将占用的空间 压缩到后者的 50% 对于传统材质你会对反照率贴图和法线贴图 等材质纹理进行采样 然后使用采样值来执行着色 采用神经网络材质压缩你会对隐式纹理数据进行采样 使用采样值执行推理 然后使用网络的输出来执行着色

将每个步骤拆分到自身管道中是非常低效的做法 因为每个步骤都需要将张量同步到设备内存、 进行运算 然后反过来同步输出 用于后续运算

为了实现最佳性能App 应该将这些步骤 合并到单个着色器派发中 借助 Shader MLMetal 让你能够直接 在你的片段着色器中运行你自己的神经网络 而无需在各个步骤之间遍历设备内存 你可以初始化输入张量、运行自己的网络 然后仅对每帧中必要的像素进行着色 这样可以减少运行时内存占用和游戏占用的磁盘空间 我们来更详细地了解一下神经网络材质评估 对输入 MTLTensor 进行初始化可以分为两部分 载入网络权重以及构建输入特征 MTLTensor 输入特征 MTLTensor是由采样边界纹理 使用片段 UV 坐标生成的

在推理过程中 习得的权重矩阵 对输入特征 MTLTensor 进行变换 以提取特征、计算激活值 并通过不同层传播信息 对多个层重复进行这项评估 从而得到经解压缩的材质 最后 使用经解压缩的材质 来进行片段着色计算

我们来看看如何使用 Shader ML对输入 MTLTensor 进行初始化 首先 声明一个片段着色器它将利用 Shader ML 并传入网络权重 先使用 include 指令添加新的 metal_tensor 头文件 我们将使用 MTLTensor 类型来访问网络权重 MTLTensor 通过缓冲区绑定槽绑定到着色器 也可以使用参数缓冲区 传入 MTLTensor MTLTensor 类型已经模板化 第一个模板参数是 MTLTensor 的 dataType 这些 MTLTensor是在设备内存中创建的 所以我们使用 device地址空间限定符 第二个参数表示MTLTensor 的维度 以及 MTLTensor 的索引要使用的类型 在这个示例中 我们使用 dextents定义一个具有动态范围的二阶张量 这样 我们的片段着色器就设置好了 接下来 我们要实施神经网络材质压缩算法 有了传入的网络权重 我们可以通过对四个隐式纹理进行采样 来创建输入 MTLTensor MTLTensor 不仅仅是一种可以绑定的资源 你也可以直接在着色器中创建内联 MTLTensor！ 创建一个对采样值进行包装的 MTLTensor 并将它用于网络评估 内联 MTLTensor 被假设为密排 因此在创建时无需传递步长 这样 对输入 MTLTensor进行初始化就完成了 我们已经完成了所有准备工作可以通过神经网络推断相应的值了 评估会使用习得的参数对输入进行变换 然后激活这些参数 激活被传递到后续各层 最后一层的激活会形成解压后的材质

今年 Metal 推出了Metal Performance Primitives 使得在着色语言中能够使用MTLTensor 运算 这个库是一组高性能 API 能在 MTLTensor 上实现高性能可移植解决方案 它提供了矩阵乘法和卷积运算 矩阵乘法是神经网络评估的核心 我们将使用Metal Performance Primitives 提供的 matmul2d 实现来实施高性能可移植网络评估例程 首先 使用 include 指令在 Metal 着色器内添加新的 MetalPerformancePrimitives 头文件 矩阵乘法的参数是 使用 matmul2d_descriptor对象进行配置的 第一组模板参数指定了 矩阵乘法的问题大小 下一组模板参数控制着 在执行运算时 矩阵乘法的输入 是否需要进行转置 最后一个模板参数控制着精度要求

除了指定描述符之外还必须使用将参与运算的线程数 指定 matmul2d 运算 这个示例中 我们位于片段着色器内因此将使用 execution_thread 来表明这个线程将执行完整的矩阵乘法 然后 使用这个配置来运行矩阵乘法

最后 使用 ReLU 激活函数 激活矩阵乘法结果的每个元素 对第二层重复这个过程 从而在着色器中进行全面的网络评估 评估完成后 即可使用经解压缩的材质进行着色 输出 MTLTensor存放着通道信息 随后可以像使用任何其他纹理采样值一样使用这些信息 这里实时展示了 神经网络材质压缩与传统材质的比较情况 使用神经网络材质并未造成可察觉的画质损失 尤其是在着色后 这里是分离出来的基本色 我们仍然很难在神经网络材质 与传统材质之间发现任何差异 然而神经网络材质占用的内存和磁盘空间只有传统材质的一半

MTLTensor 运算并非只适用于片段着色器 而是在各种函数内和各个着色器阶段都能使用 如果由整个 SIMD 组或线程组 对相同的数据执行相同的运算 你可以通过选择更大的执行组来发挥硬件的优势 但如果 MTLTensor 运算在数据方面存在分歧 或者在 MTLTensor 运算调用站点表现出不均匀的控制流 就必须使用单线程执行组 其他执行方案假设不存在分歧 并假设执行组的控制流是统一的 总结一下 你现在可以在自己的着色器中 执行一些机器学习运算比如矩阵乘法和卷积 借助 Shader ML 你可以轻松在单个着色器中运行多个机器学习运算 这种方式可以减轻缓存压力减少所需的派发数 并减少内存带宽占用 尤其是在使用较小的网络时 而 Shader ML 可以提供 必要的精细控制帮助你创建自定运算 在 Metal App 中实施前沿机器学习技术的过程从未如此简单 以上介绍了如何使用 Shader ML 将神经网络嵌入到你的着色器程序中 现在有请我的同事 Scott 来介绍 Metal 4 的全新调试工具如何 帮助你轻松调试机器学习工作负载 大家好 我是 GPU 工具团队的软件工程师 Scott Moyers 刚才 Preston 展示了一个应用程序它使用机器学习技术来 来计算环境光遮蔽 这个 App 直接在 它的 Metal 渲染管道中 编码了一个机器学习网络 在帮助开发这个 App 的过程中我遇到了一个问题： 输出存在严重的伪影 我要启用环境光遮蔽通道来突出显示我遇到的问题 物体的角落处应该有阴影 但阴影被大量噪点取而代之而且我们几乎看不到画面结构 我要展示一下我是如何使用新工具找到并解决这个问题的 首先 在 Xcode 中捕获这个 App 的 GPU 追踪情况 为此 我要点按屏幕底部的 Metal 图标 然后点按“Capture”按钮

捕获完成后 我可以在摘要中找到可用的捕获帧 左侧的调试导航器提供了一个列表 其中列出了 应用程序用来构造帧的命令 例如 离屏命令缓冲区包含了许多编码器 其中包括 GBuffer Pass 下一个命令缓冲区包含了MTL4MachineLearningCommandEncoder 通过使用 Metal 4 我可以对同步进行精细控制 虽然我非常细致地设置了 依赖通道之间的屏障和事件 我想知道同步问题有没有可能导致这些问题 为了进行这方面的检查我转到依赖项查看器 这个工具非常实用 能让我大致了解 Metal 应用程序的结构 点按左上角的“Dependencies”图标

通过这个界面 我可以看到应用程序的所有命令 以及任何同步原语 例如屏障和事件 放大命令编码器可以 显示更多详细信息 这样 我的第一个命令缓冲区就完成了

它下面的命令会将法线复制到 MTLTensor 中 然后 MTL4MachineLearningCommandEncoder 后面紧接着就是一个屏障 我要重新缩小显示以便查看整体结构 新的环境光遮蔽通道位于右侧的命令缓冲区中 在我添加这个通道之前应用程序运行良好 因此我可以假设顶部和底部 命令缓冲区中的依赖项正确无误 我要检查一下新的命令缓冲区 其中包含了 MTL4MachineLearningCommandEncoder

命令缓冲区需要等待共享事件信号收到这个信号之后才能启动 然后 命令缓冲区的末尾有一个信号用于解锁下一个缓冲区 因此 不能有任何其他命令与这个命令缓冲区并行运行 此外 在命令缓冲区内前后编码器之间都有屏障 从而确保命令按顺序逐个执行 我相信这个阶段不存在任何同步问题 至少这一帧中不存在同步问题我对此很有把握 排除同步问题后 我决定直接检查MTL4MachineLearningCommandEncoder 点按环境光遮蔽网络对应的派发调用 可以查看它的绑定资源 右侧的辅助编辑器显示了输出 MTLTensor 可以看到它具有与正在运行的应用程序相同的伪影 这显然不正确 我要连按两下输入 MTLTensor将它显示在输出旁边 输入符合我对视图空间法线的预期 朝向不同方向的物体确实具有不同的分量强度 所以问题一定出在我的机器学习网络中 回到绑定资源视图 这次连按两下“Network”以在新的ML Network Debugger 中打开它 这个工具对于了解模型内部正在发生的事情至关重要

这个图表表示环境光遮蔽网络的结构 是我用 PyTorch 编写的 在目标的第一个构建阶段我按照 Preston 之前的建议操作 将它导出为 CoreML 软件包然后转换为 MTLPackage 方框表示运算 方框之间的连接按从左到右的顺序展示了流经模型的数据流 我想找出引入了这些伪影的运算 我知道最终输出很糟糕 而输入没有问题 于是我决定将图表一分为二以缩小范围 我们挑选一个大致位于中间的运算 选择某个运算 就会在右侧显示它的描述 以及它的属性、输入和输出 更重要的是 我能够检查任何运算 输出的中间 MTLTensor 数据 点按预览 以在MTLTensor 查看器中打开 可以看到这里已经存在伪影 所以我要检查之前的运算

这个运算的输出也有伪影 我们来检查一下它的输入

但是这个 MTLTensor 似乎突出显示了画面中的物体边缘 这是意料之中的 我们网络的输入是从深度缓冲区中提取的边缘 因此 网络的这个区域一定出了问题

点按运算左上角的箭头 即可展开这个拼接区域

从这些运算的顺序和类型来看 我认出这是我的SignedSmoothstep 函数 它首先对输入取绝对值 然后将这个值限制在 0 到 1 之间 但随后以结果自身为指数对结果进行幂运算 这看起来不太对 我不记得 SignedSmoothstep函数中有幂运算 我们直接进入 Python 代码弄清楚发生了什么 我要停止调试环节 回到源代码部分

正在运行的模型位于名为 LightUNet 的类中 我要导航到它的 foward 传播函数检查它是否在执行我所预期的运算

它执行的第一个自定义运算是 SignedSmoothstep 这是我在 ML Network Debugger 中看到的拼接区域 我要跳转到它的 forward 传播函数

这应该是一个简单明了的forward smoothstep 运算 其中保持了输入的符号 但是 在这一行 可以看到这个错误 我多键入了一个星号让乘号变成了幂运算符 我们删除一个多余的星号然后再次尝试运行

就是这样 使用 Metal 4 内置的MTL4MachineLearningCommandEncode 实现的神经网络环境光遮蔽可以正常工作了

在这个演示中 我展示了如何使用 Metal Debugger 调试 Metal 4 机器学习应用程序 首先 我借助依赖项查看器对同步情况进行了验证 然后 我使用 MTLTensor 查看器 检查了网络的输入和输出 进而确认问题出在我的网络内部 最后 我使用ML Network Debugger 对网络中的运算进行逐步调试并查明了问题

这些工具属于一个工具系列这个系列囊括了更丰富的工具 可用于调试和优化 Metal App 现在 我们来回顾一下 今天介绍的内容 Metal 4 推出了 MTLTensor 这是专为机器学习数据设计的全新多维资源 MTLTensor 可以灵活处理超出二维的复杂数据布局 并且内置了步长和维度信息从而大大简化了索引编制操作 Metal 4 中的新功能让你能够 将机器学习工作负载整合到你的 Metal 管道中 MTL4MachineLearningCommandEncoder让整个机器学习网络 能够直接在 GPU 时间线上运行 从而能够与你的计算和渲染作业 无缝集成和同步 对于较小的网络 借助 Shader ML 和Metal Performance Primitives 库 你可以将机器学习作业直接嵌入到着色器中 最后 Metal Debugger可提供出色的可见性 让你清楚地了解Metal 4 应用程序中发生的事情 借助全新ML Network Debugger 你可以轻松了解自己的网络以及它在设备上的执行情况 这种洞察信息对于确保正确性 和性能优化至关重要 在后续步骤方面 可以安装最新操作系统和 Xcode然后亲自尝试一下 Metal 4 的 MTL4MachineLearningCommandEncoder和适用于 Xcode 的 Shader ML 要进一步了解 Metal 开发者工具可以如何助你一臂之力 请前往 Apple 开发者网站 为了让你的 Metal 4 应用程序发挥最佳性能 请务必观看有关 Metal 4 的其他讲座 我们非常期待看到大家利用这些新功能打造精彩纷呈的体验 感谢观看