wasm-micro-runtime icon indicating copy to clipboard operation
wasm-micro-runtime copied to clipboard

[wasi-nn] add more backends into `libraries/wasi-nn`

Open lum1n0us opened this issue 1 month ago • 0 comments

Background

We are planning to add more backends into WAMR to enrich its capability in ML/LLM/AI field. Currently the only backend(in libraries/wasi-nn) is tensorflow-lite.

The code in core/iwasm/libraries/wasi-nn can be separated into two parts:

  • wasi-nn-general. It provides host implementation of wasi-nn APIs. It is the unchanged part for every backends. It will
    1. accept parameters from the Wasm world
    2. translate parameters from Wasm world to host world
    3. call appropriate backend to finish requests based on graph_encoding
    4. translate results from host world to Wasm world
    5. send results to the Wasm world
  • wasi-nn-backends. For now, it is only wasi_nn_tensorflowlite.cpp. It will use backends' APIs and libraries to finish requests. Usually, variant backends lead to variant implementation.

Every backend contains huge libraries. For example, libtensorflow-lite.a for the tensorflowlite backend is 651M. It is not possible to pack all backends, or even one backend, into the iwasm binary. A nature way is to create a shared library for each backend.

Proposal

There will be libwasi_nn_backend_tflite.so, libwasi_nn_backend_openvino.so, and so on. Users of iwasm need to pass the right backend w/ --native-lib, like --native-lib=libwasi_nn_backend_tflite.so or --naive-lib=libwasi_nn_backend_pytorch.so or something else, before running a wasi-nn wasm application.

Each of them contains

  • a shared wasi-nn host implementation(wasi-nn-general). Match wasi_ephemeral_nn imports requirements.
  • a separated backend in variant framework. like w/ tflite, w/ openvino, w/ pytorch, w/ onnx, and so on.

Choice: use vmlib.a or iwasm.so

Since wasi-nn-general(wasi_nn.c) using APIs of WAMR(like wasm_runtime_malloc(), bh_hash_map_xxx(), ...), the final libwasi_nn_backend_xx.so not only depends on ML framework libraries but also WAMR libraries, especially libvmlib.a or libiwasm.so. If realized iwasm also depends on the same library, a problem is raised.

If, iwasm and libwasi_nn_backend_xxx.so, choose static library(vmlib.a), both needs to maintain the right runtime status. Because two static libraries causes separated sets of status. For example, even though iwasm will load_and_register_native_lib() after wasm_runtime_full_init(), wasm_runtime_malloc() in libwasi_nn_backend_xxx.so will not work unless calling wasm_runtime_full_init() by itself. The reason is there are two global variables memory_mode in space for two components.

A quick and simple solution is let iwasm and libwasi_nn_backend_xxx.so use the dynamic library(libiwasm.so). But it may change the distribution form of iwasm. ~~So, we have to stick at a static library~~.

The reason of not using another way is it requires a lot effort to resolve incoming problems, like:

~~## Choice: cross platform APIs~~

~~Another way to sync runtime status in iwasm and libwasi_nn_backend_xxx is avoiding two status at the beginning. It means to left wasi_nn.c in iwasm. Only pack ML framework(tensorflow-lite.a) and integration code(wasi_nn_tensorflowlite.cpp) together. Plus, need to choose WAMR APIs used in integration code wisely. Majority of those WAMR APIs are about cross platforms. If consider the very limited number of supported platforms of ML frameworks, maybe we don't need to involve any WAMR APIs at all. Handle it by creating a separate cross-platform layer for backend or using POSIX for now~~

~~## Choice: work with libwasi_nn_backend_xxx.so~~

~~If not contains wasi_nn.c in, the wasi-nn backend can't work via iwasm --native-lib=libwasi_nn_backend_xxx.so. It turns to a "dlopen()/dlsym()" case.~~

Choice: distribution of iwasm and iwasm-nn

Usually, the release of iwasm is just one executable binary. iwasm. Since the wasi-nn backends plan requires dynamic libraries, iwasm-nn becomes iwasm-nn + libiwasm.so + libwasi_nn_backends_xxx.so(multiple libraries for multiple backends).

By default,

  • iwasm, with static library (vmlib.a), can't execute wasi-nn related wasm modules
  • when -DWAMR_BUILD_WASI_NN=1, iwasm uses shared libraries automatically.

  • [ ] make WASI NN backend as dynamic libraries by default. Currently, tensorflowlite is statically linked to iwasm. And WASI_NN_SHARED is optional.
  • [x] add wasi-nn-examples into integration test cases
  • [ ] sync up with latest wasi-nn APIs
  • [ ] add openvino as a new backend
  • [ ] add Intel open source framework

lum1n0us avatar Jun 03 '24 09:06 lum1n0us