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生態系統等項目進行更輕松的互動。
源碼