使用Apache TVM將機器學習編譯為WASM和WebGPU
TLDR
在Apache TVM深度學習編譯器中引入了對WASM和WebGPU的支持。實驗表明，在將模型部署到Web時，TVM的WebGPU后端可以接近本機 GPU的性能。

概述
計算是現代機器學習應用程序的支柱之一。GPU的引入加快了深度學習的工作量，極大地提高了運行速度。部署機器學習的需求不斷增長，瀏覽器已成為部署智能應用程序的自然之所。
TensorFlow.js和ONNX.js將機器學習引入瀏覽器，但Web版本和本機版本之間在性能上仍然存在著不小的差距。許多原因之一是缺乏對Web上GPU的標準訪問和高性能訪問。WebGL缺少高性能著色學習所需的重要功能，例如計算著色器和通用存儲緩沖區(qū)。
WebGPU是下一代Web圖形標準。與最新一代的圖形API（例如Vulkan和Metal）一樣，WebGPU提供了一流的計算著色器支持。
為了探索在瀏覽器中使用WebGPU進行機器學習部署的潛力，增強了深度學習編譯器Apache（incubating）TVM，以WASM（用于計算啟動參數并調用設備啟動的主機代碼）和WebGPU（用于設備）為目標。執(zhí)行初步結果是非常積極的-第一次，可以在Web上部署機器學習應用程序，同時仍能接近GPU的本機性能。
機器學習編譯器

試用WebGPU的自然反應是為深度神經網絡（矩陣乘法和卷積）中的原始算子編寫著色器，然后直接優(yōu)化性能。這是現有框架（例如TensorFlow.js）使用的傳統(tǒng)工作流程。
相反，采用了基于編譯的方法。TVM自動從TensorFlow，Keras，PyTorch，MXNet和ONNX等高級框架中提取模型，使用機器學習驅動的方法自動生成低級代碼，在這種情況下，將以SPIR-V格式計算著色器。然后可以為可部署模塊生成的代碼打包。
編譯的方法的一個重要優(yōu)點是基礎架構的重用。通過重用基礎結構來優(yōu)化CUDA，Metal和OpenCL等本機平臺的GPU內核，能夠輕松地（相對于其它方法）以Web為目標。如果WebGPU API到本機API的映射有效，可以通過很少的工作獲得類似的性能。更重要的是，AutoTVM基礎架構，能夠針對特定模型專門化計算著色器，從而能夠為感興趣的特定模型生成最佳的計算著色器。
構建WASM和WebGPU編譯器
為了構建可以針對WASM和WebGPU的編譯器，需要以下元素：
? 用于計算著色器的SPIR-V生成器。
? 主機程序的WASM生成器。
? 加載和執(zhí)行生成的程序的runtime。
TVM已經有Vulkan的SPIR-V目標，使用LLVM生成主機代碼?？梢詢H將二者的用途重新生成設備和主機程序。
主要挑戰(zhàn)是runtime。需要一個runtime來加載著色器代碼，并使主機代碼對話能夠正確地與著色器通信。TVM具有最低的基于C ++的runtime。構建了一個最小的Web runtime庫，生成的著色器和主機驅動代碼鏈接，生成一個WASM文件。但是，此WASM模塊仍然包含兩個未知的依賴項：
? runtime需要調用系統(tǒng)庫調用（malloc，stderr）。
? wasmruntime需要與WebGPU驅動程序進行交互（在Javascript中，WebGPU API是the first-class citizen）。
WASI是解決第一個問題的標準解決方案。盡管網絡上還沒有成熟的WASI，使用emscripten生成類似WASI的庫，提供這些系統(tǒng)庫。
通過在TVM的JS runtime內部構建WebGPU runtime來解決第二個問題，在調用GPU代碼時，從WASM模塊中回調這些功能。使用TVM runtime系統(tǒng)中的PackedFunc機制，可以通過將JavaScript閉包傳遞到WASM接口，直接公開高級runtime原語。這種方法將大多數runtime代碼保留在JavaScript中，隨著WASI和WASM支持的成熟，可以將更多JS代碼引入WASM runtime。

性能

進行了一個快速實驗，比較了通過TVM的WebGPU后端和使用本地GPU runtime（Metal和OpenCL）的本地目標執(zhí)行完整計算圖的情況。在MobileNet模型上，可以發(fā)現WebGPU可以接近Metal的性能。假設Chrome WebGPU的runtime以MacOS上的Metal（而不是OpenCL）為目標，可以放心地假設以GPU為目標時，性能幾乎沒有損失。
此基準不包括CPU到GPU的數據復制成本，而僅基準GPU的執(zhí)行。從CPU到GPU的數據復制，仍會占用25％的執(zhí)行時間?？梢酝ㄟ^諸如連續(xù)執(zhí)行設置中的雙緩沖之類的方法，進一步攤銷這些成本。
報告的mobilenet的端到端runtime，絕不是最佳選擇，重復使用了GTX 1080 Ti的優(yōu)化程序，這與Intel圖形GPU截然不同。希望通過在目標平臺上使用AutoTVM來進一步提高性能。
展望未來
結果表明，在網絡上進行機器學習有許多有趣的機會。值得注意的是，WebGPU是一個仍在不斷發(fā)展的API，其含義可能會超出Web應用程序。例如，當WebGPU成熟，通過WASI標準化時，可以將其定位為WebGPU的本機API，使用WebGPU的獨立WASM應用程序。
TVM社區(qū)還積極地在基于Rust的runtime上工作，該runtime將提供更強大的WASM支持，wgpu和Rust WASM生態(tài)系統(tǒng)等項目的交互更加輕松。
提出的方法為大多數WASM的應用場景提供了有效的機器學習支持。接近本機的性能，可以釋放瀏覽器上更好的聯(lián)合學習功能。相同的編譯程序包，也應該能夠在本機WASM執(zhí)行程序上運行，為應用程序提供sandbox 。

總結

以上是生活随笔為你收集整理的使用Apache TVM将机器学习编译为WASM和WebGPU的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。