TVM编译机器学习到 WASM 和 WebGPU

TLDR

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

引论

计算是现代机器学习应用的支柱之一。引入 GPU 以加快深度学习工作量，大大提高了进步速度。鉴于部署机器学习无处不在的需求日益增长，浏览器成为部署智能应用程序的自然场所。

虽然 TensorFlow .js 和 ONNX .js是将机器学习引入浏览器的现有努力，但 Web 版本和本地版本在性能上仍然存在非同小的差距。众多原因之一是缺乏对 Web 上的 GPU 的标准和执行访问。WebGL 缺乏重要的功能，如计算着色器和通用存储缓冲器，这些功能是高性能深度学习所必需的。

WebGPU 是下一代 Web 图形的即将推出的标准，有可能显著改变这种状况。与最新一代图形 API（如 Vulkan 和 Metal）一样，WebGPU 提供一流的计算着色器支持。

为了探索在浏览器中使用 WebGPU 进行机器学习部署的潜力，增强了深度学习编译器 Apache（孵化）TVM，以针对 WASM（用于计算启动参数和调用进入设备启动的主机代码）和 WebGPU（用于设备执行）。初步结果是相当积极的-第一次，可以部署机器学习应用程序在网络上，同时仍然接近本地性能的GPU。

机器学习编译器

在尝试 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 生成器。
加载和执行生成程序的runtime。

幸运的是，TVM已经为Vulkan制定了SPIR-V目标，并且使用LLVM生成主机代码。因此，可以重新调整两者的用途，以生成设备和主机程序。

主要的挑战是runtime。需要一个runtime来加载着色器代码，并使主机代码通话能够正确地与着色器通信。TVM 的runtime最少C++。构建了一个最低的 Web runtime库，与生成的着色器和主机驱动代码链接，生成单个 WASM 文件。WASM 模块仍然包含两个未知的依赖关系：

runtime需要调用到系统库calls (malloc, stderr)。
runtime需要与 WebGPU 驱动程序（在 javascript中，WebGPU API is the first-class citizenWebGPU API 是一流公民）进行交互。

WASI 是解决第一个问题的标准解决方案。虽然网络上还没有成熟的 WASI，但可以使用 emscript 生成类似 WASI 的库，以提供这些系统库。

通过在 TVM 的 JS runtime内构建 WebGPU runtime以及在调用 GPU 代码时从 WASM 模块调用这些功能来解决第二个问题。使用 TVM runtime系统中的打包机制，可以通过将 JavaScript 关闭传递到 WASM 接口来直接输出高级runtime原始。此方法保留了 JavaScript 中的大部分runtime代码，随着 WASI 和 WASM 支持的成熟，可以将更多的 JS 代码引入 WASM runtime。

性能

运行了一个快速实验，比较了通过 TVM 的 WebGPU 后端和使用本地 GPU runtime （金属和 OpenCL）的本地目标执行完整计算图的执行情况。在移动网络模型上，可以发现 WebGPU 可以接近metal性能。假设 Chrome WebGPU 的runtime目标为metal，而不是 MacOS 上的 OpenCL，可以放心地假设，在面对 GPU 时，几乎没有性能损失。

此基准不包括 CPU 到 GPU 数据复制成本，仅对 GPU 执行进行基准。目前，从 CPU 到 GPU 的数据副本，仍可能需要 25% 的执行时间：这些成本可以通过连续执行设置中的双缓冲等方法进一步摊销。

报告的移动网端到端runtime绝不是最佳的，因为只是重复使用 GTX 1080 Ti 的调度，这与英特尔图形 GPU 非常不同。期望通过在感兴趣的目标平台上使用AutoTVM来进一步提升性能。

展望

结果表明，在网络上机器学习有许多有趣的机会。值得注意的是，WebGPU 是一种仍在不断发展的 API，其影响可能超越 Web 应用程序。例如，随着 WebGPU 的成熟并通过 WASI 实现标准化，可以针对 WebGPU 的原生 API，从而支持使用 WebGPU 的独立 WASM 应用程序。

TVM 社区还积极致力于基于 Rust 的runtime，这将提供更强大的 WASM 支持，并能够与wgpu和Rust WASM生态系统等项目进行更轻松的互动。

源码