webmachinelearning / webnn-polyfill

🧠⚙️ Web Neural Network API polyfill based on TensorFlow.js
https://www.npmjs.com/package/@webmachinelearning/webnn-polyfill
Apache License 2.0
101 stars 18 forks source link

Expose wasm interface #208

Closed BruceDai closed 1 year ago

BruceDai commented 1 year ago

Current WebNN-Polyfill auto select WASM backend depending on environment's supports of TensorFlow Wasm backend, see below codes

// In @tensorflow/tfjs-backend-wasm/dist/backend_wasm.js
/**
 * Initializes the wasm module and creates the js <--> wasm bridge.
 */
 export async function init() {
    const [simdSupported, threadsSupported] = await Promise.all([
        env().getAsync('WASM_HAS_SIMD_SUPPORT'),
        env().getAsync('WASM_HAS_MULTITHREAD_SUPPORT')
    ]);
    return new Promise((resolve, reject) => {
        const factoryConfig = {};
        ... ...
        factoryConfig.locateFile = (path, prefix) => {
            ... ...
            if (path.endsWith('.wasm')) {
                return getPathToWasmBinary(simdSupported, threadsSupported, wasmPathPrefix != null ? wasmPathPrefix : prefix);
            }
            return prefix + path;
        };
   ... ...
 }

/**
 * Returns the path of the WASM binary.
 * @param simdSupported whether SIMD is supported
 * @param threadsSupported whether multithreading is supported
 * @param wasmModuleFolder the directory containing the WASM binaries.
 */
function getPathToWasmBinary(simdSupported, threadsSupported, wasmModuleFolder) {
    ... ...
    let path = 'tfjs-backend-wasm.wasm';
    if (simdSupported && threadsSupported) {
        path = 'tfjs-backend-wasm-threaded-simd.wasm';
    }
    else if (simdSupported) {
        path = 'tfjs-backend-wasm-simd.wasm';
    }
    ... ....
    return wasmModuleFolder + path;
}

This PR is to expose wasm interface for user conveniently configuring threads for WASM backend if environment supported multi-threads, likes

  const wasm = navigator.ml.createContext().wasm;
  // Here to configure threads number
  wasm.setThreadsCount(n);

  const tf = navigator.ml.createContext().tf;
  await tf.setBackend('wasm')

@huningxin @Honry PTAL, thanks.

huningxin commented 1 year ago

Can webnn-polyfill tell which wasm impl is used tfjs-backend-wasm-threaded-simd.wasm or tfjs-backend-wasm-simd.wasm?

BruceDai commented 1 year ago

Can webnn-polyfill tell which wasm impl is used tfjs-backend-wasm-threaded-simd.wasm or tfjs-backend-wasm-simd.wasm?

Addressed it with new printWasmBinary method, please take another look, thanks. @huningxin

BruceDai commented 1 year ago

This updated commit is for user setting Wasm backend with Wasm binary path and threads count configurations by tf and wasm interface from context.

@huningxin @Honry Please take another look, thanks.

BruceDai commented 1 year ago

please also update the version in package.json.

Done, PTAL, thanks.

BruceDai commented 1 year ago

Thanks for your reviewing! Let's merge it.