[wasm] Use xnnpack for Add/Sub/Mul/Relu/Relu6 (#2506)

Call into xnnpack for Add/Sub/Mul/Relu/Relu6. This provides general axis broadcast support for Add, Sub and Mul.

Perf improvements (average of 200 runs on Macbook Pro 15 2018):
- Detector (15.8ms --> 13.9ms) (~14%)
- Mesh (9.2ms --> 8.2ms) (~12%)
- PoseNet (197ms --> 179ms) (~10%)
- MobileNet (103ms --> 101ms) (~0%)

PERF

Author: dsmilkov

.vscode/settings.json

@@ -130,6 +130,7 @@
     "forward_list": "cpp",
     "typeindex": "cpp",
     "*.inc": "cpp",
-    "hash_map": "cpp"
+    "hash_map": "cpp",
+    "__refstring": "cpp"
   }
 }

tfjs-backend-wasm/scripts/cpplint.js

@@ -26,27 +26,25 @@ const ignoreCode = true;
 const commandOpts = null;
 
 let pythonVersion = exec('python --version', commandOpts, ignoreCode);
-if(pythonVersion['stderr'].includes('Python 2')) {
+if (pythonVersion['stderr'].includes('Python 2')) {
   python2Cmd = 'python';
 } else {
   pythonVersion = exec('python2 --version', commandOpts, ignoreCode);
-  if(pythonVersion.code === 0) {
+  if (pythonVersion.code === 0) {
     python2Cmd = 'python2';
   }
 }
 
-if(python2Cmd != null) {
+if (python2Cmd != null) {
   const result = shell.find('src/cc').filter(
-    fileName => fileName.endsWith('.cc') || fileName.endsWith('.h'));
-
-  console.log(`C++ linting files:`);
-  console.log(result);
+      fileName => fileName.endsWith('.cc') || fileName.endsWith('.h'));
 
   const cwd = process.cwd() + '/' + CC_FILEPATH;
   const filenameArgument = result.join(' ');
 
   exec(`${python2Cmd} tools/cpplint.py --root ${cwd} ${filenameArgument}`);
 } else {
-  console.warn('No python2.x version found - please install python2. ' +
-  'cpplint.py only works correctly with python 2.x.');
+  console.warn(
+      'No python2.x version found - please install python2. ' +
+      'cpplint.py only works correctly with python 2.x.');
 }

tfjs-backend-wasm/src/cc/BUILD

@@ -49,7 +49,8 @@ tfjs_cc_library(
 
 tfjs_cc_library(
   name = "binary",
-  srcs = ["binary.h"],
+  hdrs = ["binary.h"],
+  srcs = ["binary.cc"],
   deps = [
     ":backend",
   ],
@@ -94,6 +95,16 @@ tfjs_cc_library(
   ],
 )
 
+tfjs_cc_library(
+  name = "clamp_impl",
+  hdrs = ["clamp_impl.h"],
+  srcs = ["clamp_impl.cc"],
+  deps = [
+    ":backend",
+    ":util"
+  ],
+)
+
 tfjs_cc_library(
   name = "all_kernels",
   deps = [
@@ -360,6 +371,7 @@ tfjs_cc_library(
   srcs = ["kernels/Relu.cc"],
   deps = [
     ":backend",
+    ":clamp_impl",
     ":unary",
   ],
 )
@@ -369,6 +381,7 @@ tfjs_cc_library(
   srcs = ["kernels/Relu6.cc"],
   deps = [
     ":backend",
+    ":clamp_impl",
     ":unary",
   ],
 )

tfjs-backend-wasm/src/cc/binary.cc

@@ -0,0 +1,79 @@
+/* Copyright 2019 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * ===========================================================================*/
+
+#include "src/cc/binary.h"
+
+#include <xnnpack.h>
+#include <limits>
+#include <unordered_map>
+
+#include "src/cc/backend.h"
+#include "src/cc/util.h"
+
+namespace {
+// Maps an `xnn_create_*_nd_f32` function pointer to an instantiated operator.
+std::unordered_map<tfjs::wasm::xnn_create_binary_op, xnn_operator_t> op_cache;
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+
+void binary_xnn_f32(const int a_id, const size_t* a_shape_ptr,
+                    const int a_shape_len, const int b_id,
+                    const size_t* b_shape_ptr, const int b_shape_len,
+                    const int out_id, xnn_create_binary_op create_op,
+                    xnn_setup_binary_op setup_op) {
+  auto& a_info = backend::get_tensor_info(a_id);
+  auto& b_info = backend::get_tensor_info(b_id);
+  auto& out_info = backend::get_tensor_info_out(out_id);
+  const float* a_buf = a_info.f32();
+  const float* b_buf = b_info.f32();
+  float* out_buf = out_info.f32_write();
+
+  xnn_operator_t binary_op = nullptr;
+
+  auto cache_result = op_cache.find(create_op);
+  if (cache_result == op_cache.end()) {
+    const float sum_min = -std::numeric_limits<float>::infinity(),
+                sum_max = std::numeric_limits<float>::infinity();
+    const int flags = 0;
+    xnn_status status = create_op(sum_min, sum_max, flags, &binary_op);
+    if (status != xnn_status_success) {
+      util::warn(
+          "XNN status for xnn_create_*_nd_f32 is not successful. Got "
+          "status %d. Use -c dbg to see XNN logs.");
+      return;
+    }
+    op_cache.insert({create_op, binary_op});
+    backend::xnn_operator_count++;
+  } else {
+    binary_op = cache_result->second;
+  }
+  const int batch_size = out_info.size;
+  xnn_status status =
+      setup_op(binary_op, a_shape_len, a_shape_ptr, b_shape_len, b_shape_ptr,
+               a_buf, b_buf, out_buf, nullptr /* thread pool */);
+  if (status != xnn_status_success) {
+    util::warn(
+        "XNN status for xnn_setup_*_nd_f32 is not successful. Got "
+        "status %d. Use -c dbg to see XNN logs.",
+        status);
+    return;
+  }
+
+  xnn_run_operator(binary_op, nullptr /* thread pool */);
+}
+
+}  // namespace wasm
+}  // namespace tfjs

tfjs-backend-wasm/src/cc/binary.h

@@ -15,6 +15,7 @@
 #ifndef BINARY_H_
 #define BINARY_H_
 
+#include <xnnpack.h>
 #include <algorithm>
 
 #include "src/cc/backend.h"
@@ -58,6 +59,18 @@ inline void binary_bool(const int a_id, const int b_id, const int out_id,
                     out_info.b_write(), operation);
 }
 
+typedef xnn_status (*xnn_create_binary_op)(float, float, uint32_t,
+                                           xnn_operator_t*);
+typedef xnn_status (*xnn_setup_binary_op)(xnn_operator_t, size_t, const size_t*,
+                                          size_t, const size_t*, const float*,
+                                          const float*, float*, pthreadpool_t);
+
+void binary_xnn_f32(const int a_id, const size_t* a_shape_ptr,
+                    const int a_shape_len, const int b_id,
+                    const size_t* b_shape_ptr, const int b_shape_len,
+                    const int out_id, xnn_create_binary_op create_op,
+                    xnn_setup_binary_op setup_op);
+
 }  // namespace wasm
 }  // namespace tfjs
 

tfjs-backend-wasm/src/cc/clamp_impl.cc

@@ -0,0 +1,77 @@
+/* Copyright 2019 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * ===========================================================================*/
+
+#include "src/cc/clamp_impl.h"
+
+#include <xnnpack.h>
+#include <map>
+#include <tuple>
+
+#include "src/cc/backend.h"
+#include "src/cc/util.h"
+
+namespace {
+// These values are keys to creating the xnn clamp operator. We use
+// std::tuple since it implements the compare operator needed for std::map.
+typedef std::tuple<float, float> CacheKey;
+// The operator cache maps the params of xnn_create_clamp_nc_f32 to an operator.
+std::map<CacheKey, xnn_operator_t> op_cache;
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+
+void xnn_clamp(const int x_id, const int out_id, const float min,
+               const float max) {
+  auto& x_info = backend::get_tensor_info(x_id);
+  auto& out_info = backend::get_tensor_info_out(out_id);
+  const float* x_buf = x_info.f32();
+  float* out_buf = out_info.f32_write();
+
+  xnn_operator_t op = nullptr;
+  CacheKey cache_key = {min, max};
+  const auto& cache_result = op_cache.find(cache_key);
+  if (cache_result == op_cache.end()) {
+    const size_t channels = 1, input_stride = 1, output_stride = 1, flags = 1;
+    xnn_status status = xnn_create_clamp_nc_f32(
+        channels, input_stride, output_stride, min, max, flags, &op);
+    if (status != xnn_status_success) {
+      util::warn(
+          "XNN status for xnn_create_clamp_nc_f32 is not successful. "
+          "Got status %d. Use -c dbg to see XNN logs.",
+          status);
+      return;
+    }
+    op_cache.emplace(cache_key, op);
+    backend::xnn_operator_count++;
+  } else {
+    op = cache_result->second;
+  }
+
+  const size_t batch_size = out_info.size;
+  xnn_status status = xnn_setup_clamp_nc_f32(op, batch_size, x_buf, out_buf,
+                                             nullptr /* thread pool */);
+  if (status != xnn_status_success) {
+    util::warn(
+        "XNN status for xnn_setup_clamp_nc_f32 is not successful. "
+        "Got status %d. Use -c dbg to see XNN logs.",
+        status);
+    return;
+  }
+
+  xnn_run_operator(op, nullptr /* thread pool */);
+}
+
+}  // namespace wasm
+}  // namespace tfjs

tfjs-backend-wasm/src/cc/clamp_impl.h

@@ -0,0 +1,27 @@
+/* Copyright 2019 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * ===========================================================================*/
+
+#ifndef CLAMP_IMPL_H_
+#define CLAMP_IMPL_H_
+
+namespace tfjs {
+namespace wasm {
+
+void xnn_clamp(const int x_id, const int out_id, const float min,
+               const float max);
+
+}  // namespace wasm
+}  // namespace tfjs
+
+#endif  // CLAMP_IMPL_H_

tfjs-backend-wasm/src/cc/kernels/Add.cc

@@ -15,8 +15,8 @@
 #ifdef __EMSCRIPTEN__
 #include <emscripten.h>
 #endif
+#include <xnnpack.h>
 
-#include "src/cc/backend.h"
 #include "src/cc/binary.h"
 #include "src/cc/util.h"
 
@@ -35,11 +35,14 @@ extern "C" {
 #ifdef __EMSCRIPTEN__
 EMSCRIPTEN_KEEPALIVE
 #endif
-void Add(const int a_id, const int b_id, const DType dtype, const int out_id) {
-  auto& a_info = backend::get_tensor_info(a_id);
+void Add(const int a_id, const size_t* a_shape_ptr, const int a_shape_len,
+         const int b_id, const size_t* b_shape_ptr, const int b_shape_len,
+         const DType dtype, const int out_id) {
   switch (dtype) {
     case DType::float32:
-      binary_f32(a_id, b_id, out_id, add<float>);
+      binary_xnn_f32(a_id, a_shape_ptr, a_shape_len, b_id, b_shape_ptr,
+                     b_shape_len, out_id, xnn_create_add_nd_f32,
+                     xnn_setup_add_nd_f32);
       break;
     case DType::int32:
       binary_i32(a_id, b_id, out_id, add<int>);

tfjs-backend-wasm/src/cc/kernels/Div.cc

@@ -35,7 +35,9 @@ extern "C" {
 #ifdef __EMSCRIPTEN__
 EMSCRIPTEN_KEEPALIVE
 #endif
-void Div(const int a_id, const int b_id, const DType dtype, const int out_id) {
+void Div(const int a_id, const size_t* a_shape_ptr, const int a_shape_len,
+         const int b_id, const size_t* b_shape_ptr, const int b_shape_len,
+         const DType dtype, const int out_id) {
   auto& a_info = backend::get_tensor_info(a_id);
   switch (dtype) {
     case DType::float32:

tfjs-backend-wasm/src/cc/kernels/FloorDiv.cc

@@ -30,8 +30,9 @@ extern "C" {
 #ifdef __EMSCRIPTEN__
 EMSCRIPTEN_KEEPALIVE
 #endif
-void FloorDiv(const int a_id, const int b_id, const DType dtype,
-              const int out_id) {
+void FloorDiv(const int a_id, const size_t* a_shape_ptr, const int a_shape_len,
+              const int b_id, const size_t* b_shape_ptr, const int b_shape_len,
+              const DType dtype, const int out_id) {
   auto& a_info = backend::get_tensor_info(a_id);
   switch (dtype) {
     case DType::float32:

tfjs-backend-wasm/src/cc/kernels/Mul.cc

@@ -15,8 +15,8 @@
 #ifdef __EMSCRIPTEN__
 #include <emscripten.h>
 #endif
+#include <xnnpack.h>
 
-#include "src/cc/backend.h"
 #include "src/cc/binary.h"
 #include "src/cc/util.h"
 
@@ -35,11 +35,14 @@ extern "C" {
 #ifdef __EMSCRIPTEN__
 EMSCRIPTEN_KEEPALIVE
 #endif
-void Mul(const int a_id, const int b_id, const DType dtype, const int out_id) {
-  auto& a_info = backend::get_tensor_info(a_id);
+void Mul(const int a_id, const size_t* a_shape_ptr, const int a_shape_len,
+         const int b_id, const size_t* b_shape_ptr, const int b_shape_len,
+         const DType dtype, const int out_id) {
   switch (dtype) {
     case DType::float32:
-      binary_f32(a_id, b_id, out_id, mul<float>);
+      binary_xnn_f32(a_id, a_shape_ptr, a_shape_len, b_id, b_shape_ptr,
+                     b_shape_len, out_id, xnn_create_multiply_nd_f32,
+                     xnn_setup_multiply_nd_f32);
       break;
     case DType::int32:
       binary_i32(a_id, b_id, out_id, mul<int>);