Supercharging the TensorFlow.js WebAssembly backend with SIMD and multi-threading

Posted by , Software Engineers at Google
a new WebAssembly (Wasm) accelerated backend for TensorFlow.js (scroll further down to learn more about Wasm and why this is important). Today we are excited to announce a major performance update: as of TensorFlow.js version 2.3.0, our Wasm backend has become up to 10X faster by leveraging via - a light model with 0.1 million parameters and about 20 million multiply-add operations:

(times listed are milliseconds per inference)
, a medium-sized model with 3.5 million parameters and roughly 300 million multiply-add operations, attain even greater speedups:
. SIMD support in iOS is also still under development.

**Note: Node support for the TF.js multi-threaded Wasm backend is coming soon.

The performance gains from SIMD and multithreading are independent of each other. These benchmarks show that SIMD brings a 1.7-4.5X performance improvement to plain Wasm, and multithreading brings another 1.8-2.9X speedup on top of that.

Usage

SIMD is supported as of TensorFlow.js 2.1.0, and multithreading is supported as of TensorFlow.js 2.3.0.

At runtime we test for SIMD and multithreading support and serve the appropriate Wasm binary. Today we serve a different binary for each of the following cases:

• Default: The runtime does not support SIMD or multithreading
  
• SIMD: The runtime supports SIMD but not multithreading
  
• SIMD + multithreading: The runtime supports SIMD and multithreading
  

Since most runtimes that support multi-threading also support SIMD, we decided to omit the multi-threading-only runtime to keep our bundle size down. This means that if your runtime supports multithreading but not SIMD, you will be served the default binary. There are two ways to use the Wasm backend:

1. With NPM
   
   JAVASCRIPT Kopieren

   // Import @tensorflow/tfjs or @tensorflow/tfjs-core
   const tf = require('@tensorflow/tfjs');
   // Add the WAsm backend to the global backend registry.
   require('@tensorflow/tfjs-backend-wasm');
    
   // Set the backend to WAsm and wait for the module to be ready.
   tf.setBackend('wasm').then(() => main());

   The library expects the Wasm binaries to be located relative to the main JS file. If you’re using a bundler such as parcel or webpack, you may need to manually indicate the location of the Wasm binaries with our setWasmPaths helper:
   
   JAVASCRIPT Kopieren

   import {setWasmPaths} from '@tensorflow/tfjs-backend-wasm';
   setWasmPaths(yourCustomFolder);
   tf.setBackend('wasm').then(() => {...});

   See the To compare against the unoptimized binary, try , and is to see which features are supported by your runtime, including:
   1. SIMD
      SIMD stands for Single Instruction, Multiple Data, which means that SIMD instructions operate on small fixed-size vectors of elements rather than individual scalars. The Wasm SIMD proposal makes the SIMD instructions supported by modern processors usable inside Web browsers, unlocking significant performance gains.
      
      Wasm SIMD is a in Chrome 84-86. This means developers can opt in their websites to Wasm SIMD and all their visitors will enjoy its benefits without needing to explicitly enable the feature in their browser settings. Besides Google Chrome, Firefox Nightly supports Wasm SIMD by default.
      
      
   2. Multi-threading
      Nearly all modern processors have multiple cores, each of which is able to carry out instructions independently and concurrently. WebAssembly programs can spread their work across cores via the object for fast communication between workers.
      
      Wasm threads is a by default since version 74. There is an ongoing cross-browser effort to enable this functionality for mobile devices as well.
      
   To see which browsers support SIMD, threads, and other experimental features, check out the WebAssembly .

   Looking ahead

   We expect the performance of our Wasm backend to keep improving. We’re closely following the progress of several evolving specifications in WebAssembly, including , and for WebAssembly modules. As with SIMD and multithreading, we intend to take advantage of these features as they become available with no implications for TF.js user code.

   More information

   • Check out the WebAssembly roadmap:
     
   • Read more information about the Wasm SIMD proposal:
     
   • Submit feedback and contributions via on GitHub.
     
   • Share what you build with these optimizations on social media with the for future updates.
     

   Acknowledgements

   We would like to thank Daniel Smilkov and Nikhil Thorat for laying the groundwork of the WebAssembly backend and the integration with XNNPACK, Matsvei Zhdanovich for collecting Pixel 4 benchmark numbers, and Frank Barchard for implementing low-level Wasm SIMD optimizations in XNNPACK.
   Vollständiger Original-Bericht

   Ausführliche Details, Code-Beispiele & Hersteller-Stellungnahme auf blog.tensorflow.org.

   ↗ Original-Artikel auf blog.tensorflow.org lesen