plsql一直正在编译_面向下一代Web的深度学习编译：WebAssembly和WebGPU初探

最新推荐文章于 2024-07-29 00:00:00 发布

qhumble

最新推荐文章于 2024-07-29 00:00:00 发布

阅读量396

点赞数

文章标签： plsql一直正在编译

本文链接：https://blog.csdn.net/weixin_31451231/article/details/112530888

版权

本文探讨了WebAssembly和WebGPU在浏览器中运行机器学习模型的潜力，通过TVM编译框架实现了WebGPU后端，实验结果显示在MacOS上WebGPU模块的效率接近本地Metal。文章介绍了实现过程、面临挑战及实验结果，指出未来可能的发展方向。

摘要由CSDN通过智能技术生成

在深度学习的大潮下，算力始终是其中最重要的一环。GPU加速使得我们可以更快地部署训练好的的模型。我们可以看到越来越多部署深度学习到各个环境的需求。浏览器环境本身也自然地成为了其中部署目标之一。

比较遗憾的是，虽然目前有很多已有框架（如tensorflowjs, onnxjs)的工作在尝试支持机器学习模型的浏览器部署。在浏览器中运行机器学习始终存在着性能问题。而这个问题的根本原因是因为目前的浏览器端编程模型无法充分利用GPU资源。虽然WebGL允许我们可以通过图形渲染的方式去访问GPU，我们依然无法引入shared memory，generic storage buffer这些计算特有的概念去优化浏览器端的程序（新版本的OpenGL部分解决了这个问题，但是目前WebGL依然还是基于旧的OpenGL标准）。

当然标准也在不断演化，最近Web端两个重要的新元素-- WebAssembly和WebGPU给了解决浏览器端机器学习一个新的希望。WebGPU是下一代互联网的图形学接口，目前已经进入了实现阶段，主要浏览器的nightly版本已经加入了WebGPU的支持。从API上，WebGPU支持了compute shader，使得更加极致优化浏览器端的算子成为可能。

为了探索这个可能性，TVM社区最近加入了WebAssembly和WebGPU后端的支持。通过已有的架构生成嵌入WebGPU compute shader的wasm模块。我们在Chrome预览版本上的测试结果展示了很大的潜力 -- tvm生成的WebGPU模块在MacOS上可以获得和直接本地运行native metal几乎一样的效率。

编译到WebAssembly和WebGPU

提到浏览器编程，传统的解决方法自然是手写矩阵算子代码。这也是tensorflow.js和其它已有框架干的事情。TVM采用编译和自动代码优化的手段绕过了手写优化的步骤。在这个任务下，我们的一个明显优势是我们不需要重新为WebGPU手写算子，而是直接通过给TVM加后端代码生成到WebGPU支持的输入格式就可以了。

为了支持WebAssembly和WebGPU，我们需要额外完成一下功能。

一个SPIRV的代码生成器来生成WebGPU支持的shader格式
可以直接代码生成到WebAssembly
一个runtime来运行我们生成的程序

TVM本身已经有面向vulkan的SPIRV代码生成，所以我们可以直接复用这份代码，并且我们可以通过LLVM的wasm代码生成功能来解决WebAssembly的代码生成问题。

所以剩下的主要挑战是runtime。我们需要一个runtime来读入生成的compute shader，已经让生成的wasm和compute shader进行通信。TVM本身有一个极简化的c++ runtime，我们直接用emscripten把这个runtime编译成了webassembly。但是为了运行这个wasm，我们还需要满足两个条件：

因为代码本身会依赖于各种系统调用（malloc， stderr），我们需要解决这些系统库的实现问题
如何让wasm和js端的WebGPU API进行交互。

最近提出的WASI（WebAssembly System Interface）解决了第一个问题。WASI的主要想法是标准化wasm里面可能存在的系统调用，并且允许各个环境提供它们自己的WASI实现。虽然在浏览器端暂时没有标准的WASI支持，我们改写了emscripten生成的js来模拟了一个WASI的实现。

我们直接基于TVM的PackedFunc的机制解决了第二个问题。我们先采用了Typescript在js端完成了一个GPU runtime，然后通过closure的方式把runtime函数调用传给了wasm端，使得wasm端可以调用对应的GPU接口。

实验结果

为了测试webgpu的效果，我们直接在Macbook 13inch上面进行了测试了mobilenet的运行效率。并且比较了采用TVM Metal和OpenCL后端的结果。

我们可以发现WebGPU的运行时间几乎和Metal的运行时间差不多。我们应该可以假设在MacOS上Chrome应该是直接把GPU计算offload给了Metal，所以虽然在这个情况下OpenCL效率还是比metal快，随着metal的优化的提升，WebGPU的效率也会越来越好。

当然我们的这个benchmark只测试了GPU运行开销。CPU到GPU的内存拷贝开销依然会占据25%左右的运行时间。当然如果我们可以进一步做类似double buffer之类的优化，应该可以进一步降低拷贝开销。为了快速进行这个这个实验，我们直接用了OpenCL 1080ti的AutoTVM历史记录。相信进一步通过AutoTVM的专门优化，我们可以在这个平台上获得更高的效率。

虽然只是初步实验，这一结果本身是非常有趣的。和本地平台几乎一样的效率意味着我们可以真正地把一些机器学习任务搬到前端。WebGPU本身还在演化阶段，除了javascript端的代码，也有一些native api的proposal。在未来我们或许可以把GPU runtime的部分也搬运到wasm端。

另外TVM社区最近也一直在推进rust runtime的建设。因为rust本身对于webassembly和webgpu的支持都非常不错，相信基于rust的wasm runtime可以推动许多有趣的新方向，也欢迎有兴趣的同学加入进来。