Custom operations

[dsmilkov]

Starting a thread to open the discussion for supporting custom operations.

The ML field is fast moving and model architectures and operations are evolving quickly. In TensorFlow.js, we have around ~200 ops and we still run into issues of missing ops when someone is trying to port a new model to the browser. I believe that the number of built-in ops will be relatively small and will grow very slowly due to standardization.

Thus, it is important to provide a way for library authors to write custom ops that can interop with the built-in neural net ops. That means having high-performance data exchange between custom ops and built-in ops. I stress the importance of high-performance, otherwise lib authors would revert back to implementing all of the ops using lower-level APIs (e.g. WebGPU).

A good way to start is to understand the scope and complexity of the problem. We can look at technical details on how browser-vendors plan to implement built-in ops, which gives us details about where these ops run and where the data lives.

[tomoyukilabs]

I agree to provide a way to create custom ops. More specifically, I guess that developers would like to write custom ops in the similar manner to making custom ops on Python ML frameworks, e.g. using tf.math in TensorFlow, numpy, etc.

The current charter states that basic algebra ops using WebGL, WebGPU shaders and WebAssembly SIMD are out of scope. On the other hand, platform-level neural network APIs seem to support several basic arithmetic ops according to ops mapping table by @huningxin. IMO, we need to survey what ops would be typically needed to write custom ops and what are already provided or not provided by platform-level APIs. Also, it might be important to know how browser vendors will implement built-in ops.

[DanielMazurkiewicz]

@dsmilkov @tomoyukilabs - updated API proposal, please check if this would suit your needs: https://github.com/DanielMazurkiewicz/WebAI

(custom operations and atomic ml operations API described at the end of document)

[cynthia]

For starters, the ergonomics of algebraic operations (especially with custom units) would be problematic as there is no operator overloading. (There have been attempts in the past, but none have been particularly successful.) The pipeline operator proposal (which will also take a decent amount of time) might help with this but only in a limited way; since precedence isn’t nice/straightforward.

The operator list on the document linked above definitely is not enough to do proper custom layers.

I’m under the opinion that this should probably not be in scope for the initial deliverable. It’s a pretty big pond to boil.

[tomoyukilabs]

The current charter says, "The APIs in scope of this group will not be tied to any particular platform and will be implementable on top of existing major platform APIs." In other words, I guess that we may need to consider how operations not included in major platform-level APIs should be provided to framework implementors or web developers.

There may be a couple of possible ways to provide such operations:

• define those operations in the WebNN spec as normative
• provide a way to construct custom operations

The operator list on the document linked above definitely is not enough to do proper custom layers.
I’m under the opinion that this should probably not be in scope for the initial deliverable. It’s a pretty big pond to boil.

I agree with that. However, if custom operations are not supported, available DNN models completely depend on operations supported in WebNN API. We should carefully consider what built-in and custom operations should be, in terms of both application-level use cases and ease of browser vendor's implementation.

[cynthia]

@dsmilkov @tomoyukilabs - updated API proposal, please check if this would suit your needs: https://github.com/DanielMazurkiewicz/WebAI

@dsmilkov emphasized high-performance, falling back into the main thread to do the operation in JS means 1) a tensor copy from accelerator to the CPU over the bus and back, which is expensive and 2) not friendly to vectorization.

[DanielMazurkiewicz]

@dsmilkov @tomoyukilabs - updated API proposal, please check if this would suit your needs: https://github.com/DanielMazurkiewicz/WebAI

@dsmilkov emphasized high-performance, falling back into the main thread to do the operation in JS means 1) a tensor copy from accelerator to the CPU over the bus and back, which is expensive and 2) not friendly to vectorization.

Ok, I don't know what is exactly working on on bottom of TF, assumed it falls back to JS and copies data between VRAM and regular RAM for every operation (or first and last operation in chain). But if not, this API proposal is still valid, domain doesn't force any particular way of operation usage. It can be used to build "programmatically" code for GPU from Javascript (that is only other way that comes to my mind that it works, but correct me if I'm wrong here too please), just appropriate set of operators (or additional operators) should be provided then.

[huningxin]

@dsmilkov wrote:

Thus, it is important to provide a way for library authors to write custom ops that can interop with the built-in neural net ops.

I agree extensibility is an important requirement for the low-level API whose major consumers would be the libraries.

@tomoyukilabs wrote:

The current charter states that basic algebra ops using WebGL, WebGPU shaders and WebAssembly SIMD are out of scope.

Today's libraries implement neural network ops by WebGL and WebAssembly. I suppose they will continue to adopt new features, like WebGPU compute shader and WebAssembly SIMD/Threading, for performance enhancement. So it makes sense to me that this group leaves the implementation of custom ops out of scope.

@dsmilkov worte:

That means having high-performance data exchange between custom ops and built-in ops. I stress the importance of high-performance, otherwise lib authors would revert back to implementing all of the ops using lower-level APIs (e.g. WebGPU).

I agree high-performance is the main reason to create this API, as the current charter states that WebNN is for neural network inference hardware acceleration. To access the hardware acceleration through WebNN, a library can partition a neural network into sub-graphs based on the supported ops of WebNN. Some sub-graphs can be executed by WebNN if their ops are all supported. Other sub-graphs need to be executed by custom ops within the library. The WebNN sub-graphs and custom ops should exchange data (tensor) in a efficient way. Otherwise, it is easy to lose the performance gain of WebNN execution due to memory movement/reordering overhead. We observed this issue when experimenting TensorFlow.js optimization on our WebNN POC earlier.

@dsmilkov wrote:

A good way to start is to understand the scope and complexity of the problem. We can look at technical details on how browser-vendors plan to implement built-in ops, which gives us details about where these ops run and where the data lives.

It sounds like a good approach.

Based on our POC and platform API investigation, WebNN may execute the neural network on different devices, e.g. CPU, GPU or accelerator. The API for custom ops implementation would be related to the WebNN execution device. For instance, a library may select WebAssembly-based custom ops when WebNN executes sub-graphs on CPU and select WebGPU-based custom ops when WebNN executes on GPU.

Given the WebGPU is WIP, I propose to look at the WebAssembly-based custom ops first. Within our POC, the existing graph execution API accepts inputs and outputs in ArrayBufferView where the WebAssembly-based custom ops should be able to read and write. For WebNN CPU execution, we've prototyped backends with BNNS for macOS, MKL-DNN for Linux/Windows and NNAPI for Android (the current version of NNAPI doesn't support hardware selection, we can workaround it by using device that only has CPU implementation of NNAPI). Based on this setup, we can start to investigate the solution and performance for WebAssembly-based custom ops. Later by coordinating with WebGPU @Kangz, we can extend the investigation to WebGPU-based custom ops with WebNN GPU execution, e.g. the MPS backend of our POC.

Any thoughts?

[tomoyukilabs]

LGTM, @huningxin. Thanks for your detailed explanation.

[dsmilkov]

LGTM as well. It comes down to high-performance data exchange between the built-in operations and user code. "Custom operations" follow from that and are the responsibility of user libraries to implement.

[huningxin]

Thanks @tomoyukilabs @dsmilkov !

As mentioned in 14 Feb 2019 call, I will investigate the "custom operations" support on CPU as the first step and hopefully can report back in March call.

[gramalingam]

Doesn't this essentially require standardizing the in-memory representation of tensors?

[huningxin]

I will investigate the "custom operations" support on CPU as the first step and hopefully can report back in March call.

Here is the investigation report. I'll update it in today's CG meeting.

@gramalingam worte:

Doesn't this essentially require standardizing the in-memory representation of tensors?

I agree this is essential requirement. And according to our investigation, we also need to pay attention to memory re-layout overhead, say from plain tensor layout format (e.g. NHWC) to native CPU backends (e.g. MKLDNN uses blocked layout).

[huningxin]

As a case study of custom ops support for frameworks, @pinzhenx and me prototyped a WebNN backend for ONNX.js.

There are some findings:

1. If the framework is able to share the graph info with a backend, it would help the integration of a graph-building API, e.g. WebNN. For example, in our prototype, when ONNX.js loads a model file and builds the graph in memory, the WebNN backend is invoked to transform the graph through SessionHandler.transformGraph interface. The WebNN backend tries to partition a sub-graph whose ops are supported by WebNN and rewrite the graph by replacing the sub-graph with a custom WebNN op node (WebNNGraphOp).
2. If the framework is able to execute ops by multiple backends, it would help the integration of WebNN and custom ops of framework. For example, in our prototype, the WebNN backend only supports 11 ops, however it can fallback to WASM backend and CPU backend thanks to SessionHandler.resolve and ResolveRule interfaces. So when ONNX.js executes the nodes of the graph, executing the WebNNGraphOp node actually executes the corresponding WebNN graph. The other remaining op nodes are executed by WASM backend. If they are still not supported by WASM backend as well, they are fallback to CPU backend.
3. Performance wise, if the WebNN sub-graph is as big as possible, there would be a good speedup for whole graph execution. For example, when offloading the whole graph of Resnet-50, we observed 9X speedup (WebNN-GPU vs. WebGL) for GPU execution and 30X speedup (WebNN-CPU vs. WASM) for CPU execution on a PC.
4. The tensor layout conversions would take significant overhead. It should be minimized as much as possible. For example ONNX.js uses NCHW but the current WebNN foundation API only supports NHWC. In our prototype, when executing a WebNN sub-graph, WebNN backend converts NCHW tensor to NHWC for input tensors and converts back for output tensors in JavaScript code. We may consider to support NCHW in WebNN, so browser can optimize the conversion in native.