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

TLDR

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

概述

计算是现代机器学习应用程序的支柱之一。GPU的引入加快了深度学习的工作量，极大地提高了运行速度。部署机器学习的需求不断增长，浏览器已成为部署智能应用程序的自然之所。

TensorFlow.js和ONNX.js将机器学习引入浏览器，但Web版本和本机版本之间在性能上仍然存在着不小的差距。许多原因之一是缺乏对Web上GPU的标准访问和高性能访问。WebGL缺少高性能着色学习所需的重要功能，例如计算着色器和通用存储缓冲区。

WebGPU是下一代Web图形标准。与最新一代的图形API（例如Vulkan和Metal）一样，WebGPU提供了一流的计算着色器支持。

为了探索在浏览器中使用WebGPU进行机器学习部署的潜力，增强了深度学习编译器Apache（incubating）TVM，以WASM（用于计算启动参数并调用设备启动的主机代码）和WebGPU（用于设备）为目标。执行初步结果是非常积极的-第一次，可以在Web上部署机器学习应用程序，同时仍能接近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具有最低的基于C ++的runtime。构建了一个最小的Web runtime库，生成的着色器和主机驱动代码链接，生成一个WASM文件。但是，此WASM模块仍然包含两个未知的依赖项：

• runtime需要调用系统库调用（malloc，stderr）。
• wasmruntime需要与WebGPU驱动程序进行交互（在Javascript中，WebGPU API是the first-class citizen）。

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

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

性能

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

此基准不包括CPU到GPU的数据复制成本，而仅基准GPU的执行。从CPU到GPU的数据复制，仍会占用25％的执行时间。可以通过诸如连续执行设置中的双缓冲之类的方法，进一步摊销这些成本。

报告的mobilenet的端到端runtime，绝不是最佳选择，重复使用了GTX 1080 Ti的优化程序，这与Intel图形GPU截然不同。希望通过在目标平台上使用AutoTVM来进一步提高性能。

展望未来

结果表明，在网络上进行机器学习有许​​多有趣的机会。值得注意的是，WebGPU是一个仍在不断发展的API，其含义可能会超出Web应用程序。例如，当WebGPU成熟，通过WASI标准化时，可以将其定位为WebGPU的本机API，使用WebGPU的独立WASM应用程序。

TVM社区还积极地在基于Rust的runtime上工作，该runtime将提供更强大的WASM支持，wgpu和Rust WASM生态系统等项目的交互更加轻松。

提出的方法为大多数WASM的应用场景提供了有效的机器学习支持。接近本机的性能，可以释放浏览器上更好的联合学习功能。相同的编译程序包，也应该能够在本机WASM执行程序上运行，为应用程序提供sandbox 。