[wasm] Add XNNPACK operators for unary ops (Square, Neg, Abs).

tfjs-backend-wasm/src/cc/BUILD

@@ -172,7 +172,8 @@ tfjs_cc_library(
 
 tfjs_cc_library(
     name = "unary",
-    srcs = ["unary.h"],
+    srcs = ["unary.cc"],
+    hdrs = ["unary.h"],
     deps = [
         ":backend",
     ],

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

@@ -15,9 +15,7 @@
 #ifdef __EMSCRIPTEN__
 #include <emscripten.h>
 #endif
-
-#include <cmath>
-#include <cstddef>
+#include <xnnpack.h>
 
 #include "src/cc/backend.h"
 #include "src/cc/unary.h"
@@ -31,7 +29,7 @@ extern "C" {
 EMSCRIPTEN_KEEPALIVE
 #endif
 void Abs(const size_t x_id, const size_t out_id) {
-  unary(x_id, out_id, std::abs);
+  unary_xnn_f32(x_id, out_id, xnn_create_abs_nc_f32, xnn_setup_abs_nc_f32);
 }
 
 }  // extern "C"

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

@@ -15,16 +15,11 @@
 #ifdef __EMSCRIPTEN__
 #include <emscripten.h>
 #endif
-
-#include <cmath>
+#include <xnnpack.h>
 
 #include "src/cc/backend.h"
 #include "src/cc/unary.h"
 
-namespace {
-inline float negate(const float val) { return -val; }
-}  // namespace
-
 namespace tfjs {
 namespace wasm {
 // We use C-style API to interface with Javascript.
@@ -33,7 +28,10 @@ extern "C" {
 #ifdef __EMSCRIPTEN__
 EMSCRIPTEN_KEEPALIVE
 #endif
-void Negate(const int x_id, const int out_id) { unary(x_id, out_id, negate); }
+void Negate(const int x_id, const int out_id) {
+  unary_xnn_f32(x_id, out_id, xnn_create_negate_nc_f32,
+                xnn_setup_negate_nc_f32);
+}
 
 }  // extern "C"
 }  // namespace wasm

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

@@ -15,16 +15,11 @@
 #ifdef __EMSCRIPTEN__
 #include <emscripten.h>
 #endif
-
-#include <cstddef>
+#include <xnnpack.h>
 
 #include "src/cc/backend.h"
 #include "src/cc/unary.h"
 
-namespace {
-inline float square(const float val) { return val * val; }
-}  // namespace
-
 namespace tfjs {
 namespace wasm {
 // We use C-style API to interface with Javascript.
@@ -34,7 +29,8 @@ extern "C" {
 EMSCRIPTEN_KEEPALIVE
 #endif
 void Square(const size_t x_id, const size_t out_id) {
-  unary(x_id, out_id, square);
+  unary_xnn_f32(x_id, out_id, xnn_create_square_nc_f32,
+                xnn_setup_square_nc_f32);
 }
 
 }  // extern "C"

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

@@ -0,0 +1,74 @@
+/* Copyright 2020 Google LLC. 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/unary.h"
+
+#include <xnnpack.h>
+#include <cstddef>
+#include <limits>
+#include <unordered_map>
+
+#include "src/cc/backend.h"
+#include "src/cc/util.h"
+
+namespace {
+// Maps an `xnn_create_*_nc_f32` function pointer to an instantiated operator.
+std::unordered_map<tfjs::wasm::xnn_create_unary_op, xnn_operator_t> op_cache;
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+
+void unary_xnn_f32(const size_t x_id, const size_t out_id,
+                   xnn_create_unary_op create_op, xnn_setup_unary_op setup_op) {
+  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 unary_op = nullptr;
+
+  auto cache_result = op_cache.find(create_op);
+  if (cache_result == op_cache.end()) {
+    const size_t channels = 1, input_stride = 1, output_stride = 1;
+    const uint32_t flags = 1;
+    xnn_status status =
+        create_op(channels, input_stride, output_stride, flags, &unary_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, unary_op});
+    backend::xnn_operator_count++;
+  } else {
+    unary_op = cache_result->second;
+  }
+  const size_t batch_size = out_info.size;
+  xnn_status status =
+      setup_op(unary_op, batch_size, x_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(unary_op, nullptr /* thread pool */);
+}
+
+}  // namespace wasm
+}  // namespace tfjs

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

@@ -15,8 +15,11 @@
 #ifndef UNARY_H_
 #define UNARY_H_
 
+#include <xnnpack.h>
 #include <cstddef>
 
+#include "src/cc/backend.h"
+
 namespace tfjs {
 namespace wasm {
 
@@ -33,6 +36,14 @@ inline void unary(const size_t x_id, const size_t out_id,
   }
 }
 
+typedef xnn_status (*xnn_create_unary_op)(size_t, size_t, size_t, uint32_t,
+                                          xnn_operator_t*);
+typedef xnn_status (*xnn_setup_unary_op)(xnn_operator_t, size_t, const float*,
+                                         float*, pthreadpool_t);
+
+void unary_xnn_f32(const size_t x_id, const size_t out_id,
+                   xnn_create_unary_op create_op, xnn_setup_unary_op setup_op);
+
 }  // namespace wasm
 }  // namespace tfjs