TVM编译机器学习到WASM和WebGPU

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

引论
计算是现代机器学习应用的支柱之一。引入 GPU 以加快深度学习工作量&＃xff0c;大大提高了进步速度。鉴于部署机器学习无处不在的需求日益增长&＃xff0c;浏览器成为部署智能应用程序的自然场所。
虽然 TensorFlow .js 和 ONNX .js是将机器学习引入浏览器的现有努力&＃xff0c;但 Web 版本和本地版本在性能上仍然存在非同小的差距。众多原因之一是缺乏对 Web 上的 GPU 的标准和执行访问。WebGL 缺乏重要的功能&＃xff0c;如计算着色器和通用存储缓冲器&＃xff0c;这些功能是高性能深度学习所必需的。
WebGPU 是下一代 Web 图形的即将推出的标准&＃xff0c;有可能显著改变这种状况。与最新一代图形 API&＃xff08;如 Vulkan 和 Metal&＃xff09;一样&＃xff0c;WebGPU 提供一流的计算着色器支持。
为了探索在浏览器中使用 WebGPU 进行机器学习部署的潜力&＃xff0c;增强了深度学习编译器 Apache&＃xff08;孵化&＃xff09;TVM&＃xff0c;以针对 WASM&＃xff08;用于计算启动参数和调用进入设备启动的主机代码&＃xff09;和 WebGPU&＃xff08;用于设备执行&＃xff09;。初步结果是相当积极的-第一次&＃xff0c;可以部署机器学习应用程序在网络上&＃xff0c;同时仍然接近本地性能的GPU。
机器学习编译器

在尝试 WebGPU 时&＃xff0c;一个自然反应是为深神经网络中的原始算子编写着色器&＃xff08;矩阵乘法和卷积&＃xff09;&＃xff0c;然后直接优化性能。这是现有框架&＃xff08;如 TensorFlow&＃xff09;使用的传统工作流程.js。
相反&＃xff0c;采用基于编译的方法。TVM 会自动从高级框架&＃xff08;如 TensorFlow、Keras、PyTorch、MXNet 和 ONNX&＃xff09;中获取模型&＃xff0c;并使用机器学习驱动方法自动生成低级别代码&＃xff0c;在这种情况下&＃xff0c;以 SPIR-V 格式计算着色器。然后&＃xff0c;生成的代码可以打包为可部署模块。
基于编译的方法的一个重要优势是基础设施的再利用。能够毫不费力地&＃xff08;相对于其它方法&＃xff09;通过重新利用基础架构&＃xff0c;优化本地平台&＃xff08;如 CUDA、metal和 OpenCL&＃xff09;的 GPU 内核来定位 Web。如果 WebGPU API 与本地 API 的映射效率很高&＃xff0c;则可以期望类似的性能&＃xff0c;但工作很少。更重要的是&＃xff0c;AutoTVM基础架构能够专门计算特定型号的点着色器&＃xff0c;从而能够生成针对特定兴趣模型的最佳计算着色器。
构建 WASM 和 WebGPU 编译器
为了构建一个可以针对 WASM 和 WebGPU 的编译器&＃xff0c;需要以下元素&＃xff1a;
• 用于计算着色器的 SPIR-V 生成器。
• 主机程序的 WASM 生成器。
• 加载和执行生成程序的runtime。
幸运的是&＃xff0c;TVM已经为Vulkan制定了SPIR-V目标&＃xff0c;并且使用LLVM生成主机代码。因此&＃xff0c;可以重新调整两者的用途&＃xff0c;以生成设备和主机程序。
主要的挑战是runtime。需要一个runtime来加载着色器代码&＃xff0c;并使主机代码通话能够正确地与着色器通信。TVM 的runtime最少C&＃43;&＃43;。构建了一个最低的 Web runtime库&＃xff0c;与生成的着色器和主机驱动代码链接&＃xff0c;生成单个 WASM 文件。WASM 模块仍然包含两个未知的依赖关系&＃xff1a;
• runtime需要调用到系统库calls (malloc, stderr)。
• runtime需要与 WebGPU 驱动程序&＃xff08;在 Javascript中&＃xff0c;WebGPU API is the first-class citizenWebGPU API 是一流公民&＃xff09;进行交互。
WASI 是解决第一个问题的标准解决方案。虽然网络上还没有成熟的 WASI&＃xff0c;但可以使用 emscript 生成类似 WASI 的库&＃xff0c;以提供这些系统库。
通过在 TVM 的 JS runtime内构建 WebGPU runtime以及在调用 GPU 代码时从 WASM 模块调用这些功能来解决第二个问题。使用 TVM runtime系统中的打包机制&＃xff0c;可以通过将 Javascript 关闭传递到 WASM 接口来直接输出高级runtime原始。此方法保留了 Javascript 中的大部分runtime代码&＃xff0c;随着 WASI 和 WASM 支持的成熟&＃xff0c;可以将更多的 JS 代码引入 WASM runtime。

性能

运行了一个快速实验&＃xff0c;比较了通过 TVM 的 WebGPU 后端和使用本地 GPU runtime &＃xff08;金属和 OpenCL&＃xff09; 的本地目标执行完整计算图的执行情况。在移动网络模型上&＃xff0c;可以发现 WebGPU 可以接近metal性能。假设 Chrome WebGPU 的runtime目标为metal&＃xff0c;而不是 MacOS 上的 OpenCL&＃xff0c;可以放心地假设&＃xff0c;在面对 GPU 时&＃xff0c;几乎没有性能损失。
此基准不包括 CPU 到 GPU 数据复制成本&＃xff0c;仅对 GPU 执行进行基准。目前&＃xff0c;从 CPU 到 GPU 的数据副本&＃xff0c;仍可能需要 25% 的执行时间&＃xff1a;这些成本可以通过连续执行设置中的双缓冲等方法进一步摊销。
报告的移动网端到端runtime绝不是最佳的&＃xff0c;因为只是重复使用 GTX 1080 Ti 的调度&＃xff0c;这与英特尔图形 GPU 非常不同。期望通过在感兴趣的目标平台上使用AutoTVM来进一步提升性能。
展望
结果表明&＃xff0c;在网络上机器学习有许多有趣的机会。值得注意的是&＃xff0c;WebGPU 是一种仍在不断发展的 API&＃xff0c;其影响可能超越 Web 应用程序。例如&＃xff0c;随着 WebGPU 的成熟并通过 WASI 实现标准化&＃xff0c;可以针对 WebGPU 的原生 API&＃xff0c;从而支持使用 WebGPU 的独立 WASM 应用程序。
TVM 社区还积极致力于基于 Rust 的runtime&＃xff0c;这将提供更强大的 WASM 支持&＃xff0c;并能够与wgpu和Rust WASM生态系统等项目进行更轻松的互动。
源码
• Example project for image classification
• Apache TVM on github