使用 Apache TVM 将机器学习编译到 WASM 和 WebGPU

•


TLDR

我们为 Apache TVM 深度学习编译器引入了对 WASM 和 WebGPU 的支持。我们的实验表明，当将模型部署到 Web 时，TVM 的 WebGPU 后端可以接近原生 GPU 性能。

image

介绍

计算是现代机器学习应用的重要支柱之一。GPU 的引入加速了深度学习工作负载，极大地提高了发展速度。鉴于在任何地方部署机器学习的需求日益增长，浏览器成为部署智能应用的自然场所。

尽管 TensorFlow.js 和 ONNX.js 是将机器学习引入浏览器的现有努力，但 Web 版本和原生版本之间仍然存在明显的性能差距。其中一个重要原因是 Web 上缺乏对 GPU 的标准且高性能的访问。WebGL 缺少高性能深度学习所需的关键特性，例如计算着色器和通用存储缓冲区。

WebGPU 是下一代 Web 图形的新兴标准，它有可能极大地改变这种情况。与最新一代图形 API（如 Vulkan 和 Metal）一样，WebGPU 提供了 first-class 计算着色器支持。

为了探索在浏览器中使用 WebGPU 进行机器学习部署的潜力，我们增强了深度学习编译器 Apache (incubating) TVM，使其以 WASM（用于计算启动参数并调用设备启动的主机代码）和 WebGPU（用于设备执行）为目标。我们的初步结果非常积极——我们首次可以在 Web 上部署机器学习应用程序，同时仍然获得接近原生的 GPU 性能。

机器学习编译器

image

尝试 WebGPU 时，一种自然的反应是为深度神经网络中的原始运算符（矩阵乘法和卷积）编写着色器，然后直接优化它们的性能。这是现有框架（如 TensorFlow.js）使用的传统工作流程。

相反，我们应用了一种基于编译的方法。TVM 自动从高级框架（如 TensorFlow、Keras、PyTorch、MXNet 和 ONNX）中提取模型，并使用机器学习驱动的方法自动生成低级代码，在本例中为 SPIR-V 格式的计算着色器。生成的代码随后可以打包为可部署的模块。

基于编译的方法的一个重要优势是基础设施的重用。通过重用用于优化原生平台（如 CUDA、Metal 和 OpenCL）的 GPU 内核的基础设施，我们可以毫不费力地（相对于其他方法）以 Web 为目标。如果 WebGPU API 到原生 API 的映射是高效的，我们可以期待以很少的工作获得类似的性能。更重要的是，AutoTVM 基础设施允许我们为特定模型定制计算着色器，从而为我们感兴趣的特定模型生成最佳计算着色器。

构建 WASM 和 WebGPU 编译器

为了构建一个可以以 WASM 和 WebGPU 为目标的编译器，我们需要以下要素

  • 用于计算着色器的 SPIR-V 生成器。
  • 用于主机程序的 WASM 生成器。
  • 用于加载和执行生成程序的运行时。

幸运的是，TVM 已经有用于 Vulkan 的 SPIR-V 目标，并使用 LLVM 进行主机代码生成。因此，我们可以重新利用这两者来生成设备和主机程序。

主要挑战是运行时。我们需要一个运行时来加载着色器代码，并使主机代码能够与着色器正确通信。TVM 有一个最小的基于 C++ 的运行时。我们构建了一个最小的 Web 运行时库，并将其与生成的着色器和主机驱动代码链接，生成一个 WASM 文件。但是，这个 WASM 模块仍然包含两个未知的依赖项

  • 运行时需要调用系统库调用 (malloc, stderr)。
  • wasm 运行时需要与 WebGPU 驱动程序交互（在 javascript 中，WebGPU API 是 first-class citizen）。

WASI 是解决第一个问题的标准解决方案。虽然 Web 上还没有成熟的 WASI，但我们可以使用 emscripten 生成一个类似 WASI 的库（参见此处的讨论）来提供这些系统库。

我们通过在 TVM 的 JS 运行时中构建 WebGPU 运行时来解决第二个问题，并在调用 GPU 代码时从 WASM 模块回调这些函数。使用 TVM 运行时系统中的 PackedFunc 机制，我们可以通过将 JavaScript 闭包传递给 WASM 接口来直接公开高级运行时原语。这种方法将大部分运行时代码保留在 JavaScript 中，随着 WASI 和 WASM 支持的成熟，我们可以将更多的 JS 代码引入 WASM 运行时。

image

性能

image

我们进行了一个快速实验，比较了通过 TVM 的 WebGPU 后端和使用原生 GPU 运行时（Metal 和 OpenCL）的原生目标执行完整计算图的情况。在 MobileNet 模型上，我们发现 WebGPU 可以接近 Metal 的性能。假设 Chrome WebGPU 的运行时以 Metal 而不是 MacOS 上的 OpenCL 为目标，我们可以安全地假设在以 GPU 为目标时几乎没有性能损失。

此基准测试排除了 CPU 到 GPU 的数据复制成本，仅对 GPU 执行进行基准测试。目前，从 CPU 到 GPU 的数据复制仍然可能占用 25% 的执行时间；但是，可以通过诸如在连续执行设置中使用双缓冲等方法进一步分摊这些成本。

我们报告的 mobilenet 端到端运行时间绝非最佳，因为我们只是重用了来自 GTX 1080 Ti 的调优程序，这与 Intel 图形 GPU 非常不同。我们期望通过在感兴趣的目标平台上使用 AutoTVM 来进一步提高性能。

展望未来

我们的结果表明，Web 上的机器学习有很多有趣的机会。值得注意的是，WebGPU 是一个仍在发展中的 API，其影响可能超出 Web 应用程序。例如，当 WebGPU 成熟并通过 WASI 标准化时，人们可以以 WebGPU 的原生 API 为目标，从而实现利用 WebGPU 的独立 WASM 应用程序。

TVM 社区也在积极开发 基于 Rust 的运行时，这将实现更强大的 WASM 支持，并使与 wgpu 和 Rust WASM 生态系统等项目的交互更加容易。作为一个开源项目，我们正在寻找能够带来新想法并帮助将项目推向这些令人兴奋的方向的贡献者。

提出的方法为大多数 WASM 的应用场景提供了有效的机器学习支持。接近原生的性能可以解锁浏览器上更好的联邦学习能力。相同的编译包也应该能够在原生 WASM 执行器上运行，为应用程序提供沙箱。

示例代码

致谢

我们要感谢 emscripten 项目提供 WASM 编译基础设施以及 Web 上的 JS 库支持。我们还要感谢 WebGPU 社区的各种有益讨论。感谢 Fletcher Haynes 对本文提供的宝贵反馈。