[js/web] WebGPU backend via JSEP by fs-eire · Pull Request #14579 · microsoft/on...

[js/web] WebGPU backend via JSEP by fs-eire · Pull Request #14579 · microsoft/onnxruntime · GitHub

Description

This change introduced the following new components into ONNX Runtime Web:

• JavaScript Execution Provider (JSEP)
  • Asynchronized inferencing execution powered by Emscripten's Asyncify
• WebGPU backend implemented in TypeScript
  • initial implementation of kernels:
    • elementwise operators (22)
    • binary operators (5)
    • tensor: Shape, Reshape, Transpose, Gemm
    • nn: Conv, {Global}Maxpool, {Global}AveragePool

Code need to be polished. still working on it.

What is JSEP?

JSEP, aka JavaScript Execution Provider, is a new ONNXRuntime execution provider that specifically works on Web environment (browsers). JSEP allows JavaScript code to kick in from various places when ONNX Runtime inferences a model.

Why JSEP?

JSEP is a hybrid mode EP that contains both C/C++ and TypeScript/JavaScript implementation. There are 2 strong reasons why we introduces JSEP:

1. the C/C++ part helps JSEP to leverage ONNX Runtime's capabilities as much as possible including graph transformer, optimizers and also the capabilities to fallback to CPU EP. TypeScript/JavaScript helps JSEP to develop and debug much easier in the browser for the kernel implementation.
2. the requirement of asynchronized execution from JavaScript API (eg. buffer.mapAsync()) makes it impossible to run OrtRun() in a synchronized context (see "async problem" section below). This is done by using Emscripten's Asyncify.

What is WebGPU?

WebGPU is the new GPU API that available in browser. It's one of the only 2 APIs that currently available to access the GPU from browser (the other is WebGL).
WebGPU is designed with more advanced and stronger features comparing to WebGL and is potentially solution that offer the best GPU performance for model inferencing that currently available.

What is the async problem and why we have the problem?

The "async problem" is a problem that you cannot call an async function in a synchronous context. Think about the following C++ code:

// C-style declarations (API)
typedef void (*ON_COMPLETE)(PVOID state, DATA *data);
void read_data_from_file(FILEHANDLE file, ON_COMPLETE on_complete);

// implementation
DATA * my_impl_read_data_from_file_sync(FILEHANDLE file) {
  // how to implement?
}

The answer is, it's impossible to implement this function. Usually we try to find a sync version API, or launch a thread to call the async function and sync-wait on the main thread. Unfortunately, in browser environment, neither is possible.

WebGPU does not offer any synchronized API for data downloading (GPU to CPU). This is the only operation that MUST be async. As OrtRun() will eventually call into DataTransfer for copy data from GPU to CPU, and OrtRun() is a synchronized function, this cannot be done in normal way.

What is Emscripten? How is the Asyncify feature resolved the problem?

Emscripten is the C/C++ compiler for WebAssembly. It's what we use to compile ORT and generates the WebAssembly artifacts which runs on browsers.

Asyncify is a compiler feature that allows calling async functions from a synchronized context. In short, it generates code to unwind and rewind call stack to emulate async execution. With this feature, we are able to call the async function inside OrtRun() call.

Design Overview

Inter-op

JSEP is doing pretty much same thing to just another EP. It exposes an interface for inter-op with JavaScript, which is defined in onnxruntime/wasm/js_internal_api.js:

// init JSEP
Module["jsepInit"] = function (backend, alloc, free, copy, copyAsync, createKernel, releaseKernel, run) {
    Module.jsepBackend = backend;
    Module.jsepAlloc = alloc;
    Module.jsepFree = free;
    Module.jsepCopy = copy;
    Module.jsepCopyAsync = copyAsync;
    Module.jsepCreateKernel = createKernel;
    Module.jsepReleaseKernel = releaseKernel;
    Module.jsepRun = run;
};

This simple JavaScript snippet defines all language barrier level functions that requires by JSEP to achieve implementing kernels and data transfers using JavaScript inside ONNX Runtime:

• jsepBackend: assign the singleton object to webassembly module
• jsepAlloc and jsepFree: implementation of data transfer's Alloc() and Free()
• jsepCopy: synchronized copy ( GPU to GPU, CPU to GPU)
• jsepCopyAsync: asynchronized copy ( GPU to CPU)
• jsepCreateKernel and jsepReleaseKernel: a corresponding object that maintained in JS to match lifecycle of Kernel in ORT
• jsepRun: OpKernel::Compute() should call into this

The abstraction above allows to tie as little as possible connections and dependencies between C/C++ and TypeScript/JavaScript.

Resource Management

Lifecycle of tensor data and kernels are managed by ORT(C/C++) but the implementation are left to JavaScript. JavaScript code are responsible to implement the callbacks correctly.

For WebGPU, the GPU data is managed by JavaScript using a singleton map (tensot_data_id => GPUBuffer). GPU pipeline is managed as singleton. Shaders are managed using a singletonmap (shader_key => gpu_program), while shader_key is generated by cache_key (OP specific, including attributes) and input shapes.

about data transfer
js::DataTransfer::CopyTensor implemented to call either synchronized or asynchronized copy callback, depending on the destination is GPU or not. Emscripten's macro EM_ASYNC_JS is used to wrap the async function to be called in the synchronized context.

run kernel in JS

Kernel class constructor calls once jsepCreateKernel() with an optional per-kernel specific serialization to pass attributes into JavaScript.

Compute() are implemented in a way that a metadata serialization is performed in a base class and JavaScript code can access the data using the Emscripten specific builtin macro EM_ASM_*.

disabled features
memory pattern is force disabled, because the WebGPU data is not presented by a general memory model (a buffer can be represented by offset + size).
concurrent run support is disabled. WebGPU is stateful and it also has async function call. To support concurrent run will significantly increase the complexity and we don't get any real benefit from it.

prefer channels last
JSEP prefers channels last and returns DataLayout::NHWC in method GetPreferredLayout(). This will let the graph transformers to preprocess the graph into a channels last form so that a more optimized WebGPU shader can be used.

Testing code
It's impossible to test JSEP directly because JSEP itself does not contain any kernel implementation. However, it has the kernel registration which need to work together with the corresponding JavaScript code. There are unit tests that run onnx models from JavaScript API.