[wasm] Add Softmax, Neg, NotEqual.

tfjs-backend-wasm/.bazelrc

@@ -13,8 +13,6 @@ build:wasm --cxxopt="-std=c++11"
 build:wasm --cxxopt="-fno-rtti"
 build:wasm --cxxopt="-fno-exceptions"
 build:wasm --cxxopt="-fomit-frame-pointer"
-build:wasm --cxxopt="-ffast-math"
-build:wasm --copt="-ffast-math"
 
 # Disable sandbox environment because emsdk caches files by writing to
 # home directory.

tfjs-backend-wasm/WORKSPACE

@@ -8,9 +8,9 @@ emsdk_configure(name = "emsdk")
 
 git_repository(
     name = "xnnpack",
-    commit = "3a77ea7bbe30b2411591f2ab15f9c5032a25f688",
+    commit = "7278a95e3cfae6eac73f363c4fda5db53e1b2a87",
     remote = "https://github.com/google/XNNPACK.git",
-    shallow_since = "1577131863 -0800",
+    shallow_since = "1580796377 -0800",
 )
 
 # The libraries below are transitive dependencies of XNNPACK that we need to
@@ -64,6 +64,7 @@ http_archive(
 http_archive(
     name = "cpuinfo",
     build_file = "@xnnpack//third_party:cpuinfo.BUILD",
+    patches = ["@xnnpack//third_party:cpuinfo.patch"],
     sha256 = "3f2dc1970f397a0e59db72f9fca6ff144b216895c1d606f6c94a507c1e53a025",
     strip_prefix = "cpuinfo-d5e37adf1406cf899d7d9ec1d317c47506ccb970",
     urls = [

tfjs-backend-wasm/package.json

@@ -32,15 +32,15 @@
     "path": false
   },
   "peerDependencies": {
-    "@tensorflow/tfjs-core": "1.5.2"
+    "@tensorflow/tfjs-core": "link:../tfjs-core"
   },
   "dependencies": {
     "@types/emscripten": "~0.0.34"
   },
   "devDependencies": {
     "@bazel/bazel": "^0.28.0",
     "@bazel/buildifier": "0.29.0",
-    "@tensorflow/tfjs-core": "1.5.2",
+    "@tensorflow/tfjs-core": "link:../tfjs-core",
     "@types/jasmine": "~2.8.6",
     "clang-format": "~1.2.4",
     "jasmine": "~3.1.0",

tfjs-backend-wasm/src/cc/BUILD

@@ -177,15 +177,18 @@ tfjs_cc_library(
         ":Min",
         ":Minimum",
         ":Mul",
+        ":Neg",
         ":NonMaxSuppressionV3",
         ":NonMaxSuppressionV5",
+        ":NotEqual",
         ":PadV2",
         ":Prelu",
         ":Relu",
         ":Relu6",
         ":ResizeBilinear",
         ":ScatterNd",
         ":Sigmoid",
+        ":Softmax",
         ":Sub",
         ":Tile",
         ":Transpose",
@@ -457,6 +460,16 @@ tfjs_cc_library(
     ],
 )
 
+tfjs_cc_library(
+    name = "NotEqual",
+    srcs = ["kernels/NotEqual.cc"],
+    deps = [
+        ":backend",
+        ":binary",
+        ":util",
+    ],
+)
+
 tfjs_cc_library(
     name = "LogicalAnd",
     srcs = ["kernels/LogicalAnd.cc"],
@@ -542,6 +555,15 @@ tfjs_cc_library(
     ],
 )
 
+tfjs_cc_library(
+    name = "Neg",
+    srcs = ["kernels/Neg.cc"],
+    deps = [
+        ":backend",
+        ":unary",
+    ],
+)
+
 tfjs_cc_library(
     name = "NonMaxSuppressionV3",
     srcs = ["kernels/NonMaxSuppressionV3.cc"],
@@ -647,6 +669,16 @@ tfjs_cc_library(
     ],
 )
 
+tfjs_cc_library(
+    name = "Softmax",
+    srcs = ["kernels/Softmax.cc"],
+    hdrs = ["kernels/Softmax.h"],
+    deps = [
+        ":backend",
+        ":unary",
+    ],
+)
+
 tfjs_unit_test(
     name = "Sigmoid_test",
     srcs = ["kernels/Sigmoid_test.cc"],
@@ -655,6 +687,14 @@ tfjs_unit_test(
     ],
 )
 
+tfjs_unit_test(
+    name = "Softmax_test",
+    srcs = ["kernels/Softmax_test.cc"],
+    deps = [
+        ":Softmax",
+    ],
+)
+
 tfjs_cc_library(
     name = "Square",
     srcs = ["kernels/Square.cc"],

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

@@ -0,0 +1,40 @@
+/* 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include <cmath>
+
+#include "src/cc/backend.h"
+#include "src/cc/unary.h"
+
+namespace {
+inline float neg(const float val) { return -val; }
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void Neg(const int x_id, const int out_id) { unary(x_id, out_id, neg); }
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

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

@@ -0,0 +1,60 @@
+/* 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include "src/cc/binary.h"
+#include "src/cc/util.h"
+
+namespace {
+template <class T>
+inline bool notEqual(T a, T b) {
+  return a != b;
+}
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void NotEqual(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 input_type, const int out_id) {
+  switch (input_type) {
+    case DType::float32:
+      compare_f32(a_id, b_id, out_id, notEqual<float>);
+      break;
+    case DType::int32:
+      compare_i32(a_id, b_id, out_id, notEqual<int>);
+      break;
+    case DType::boolean:
+      compare_bool(a_id, b_id, out_id, notEqual<bool>);
+      break;
+    default:
+      util::warn(
+          "NotEqual for tensor ids %d and %d failed."
+          "Unsupported input_type %d",
+          a_id, b_id, input_type);
+  }
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

tfjs-backend-wasm/src/cc/kernels/Sigmoid.h

@@ -22,7 +22,6 @@ namespace wasm {
 extern "C" {
 
 void Sigmoid(const size_t x_id, const size_t out_id);
-
 }
 
 }  // namespace wasm

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

@@ -0,0 +1,99 @@
+/* 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include <xnnpack.h>
+#include <cmath>
+#include <cstddef>
+#include <map>
+#include <tuple>
+
+#include "src/cc/backend.h"
+#include "src/cc/kernels/Softmax.h"
+#include "src/cc/util.h"
+
+namespace {
+// We use std::tuple as the cache key as it implements the compare operator
+// needed for std::map.
+typedef std::tuple<size_t> OperatorCacheKey;
+
+// The operator cache maps the weights id to the xnn_operator_t instantiated for
+// this set of weights.
+std::map<OperatorCacheKey, xnn_operator_t> operator_cache;
+
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+
+void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
+             const size_t batch) {
+  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 softmax_op = nullptr;
+
+  OperatorCacheKey cache_key = {channels};
+
+  auto operator_cache_idx = operator_cache.find(cache_key);
+  if (operator_cache_idx == operator_cache.end()) {
+    const size_t input_stride = channels;
+    const size_t output_stride = channels;
+    const uint32_t flags = 0;
+
+    xnn_status status = xnn_create_softmax_nc_f32(
+        channels, input_stride, output_stride, flags, &softmax_op);
+    if (status != xnn_status_success) {
+      tfjs::util::warn(
+          "XNN status for xnn_create_softmax_nc_f32 is not "
+          "successful. Got status %d. Use -c dbg to see XNN logs.",
+          status);
+      return;
+    }
+
+    operator_cache.insert({cache_key, softmax_op});
+
+    tfjs::backend::xnn_operator_count++;
+  } else {
+    softmax_op = operator_cache_idx->second;
+  }
+
+  xnn_status status = xnn_setup_softmax_nc_f32(
+      softmax_op, batch, x_buf, out_buf, nullptr /* thread pool */);
+  if (status != xnn_status_success) {
+    tfjs::util::warn(
+        "XNN status for xnn_setup_softmax_nc_f32 is not "
+        "successful. Got status %d. Use -c dbg to see XNN logs.",
+        status);
+    return;
+  }
+
+  xnn_run_operator(softmax_op, nullptr /* thread pool */);
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

tfjs-backend-wasm/src/cc/kernels/Softmax.h

@@ -0,0 +1,30 @@
+/* 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.
+ * ===========================================================================*/
+
+#ifndef KERNELS_SOFTMAX_H_
+#define KERNELS_SOFTMAX_H_
+
+#include <cstddef>
+
+namespace tfjs {
+namespace wasm {
+extern "C" {
+void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
+             const size_t batch);
+}
+
+}  // namespace wasm
+}  // namespace tfjs
+
+#endif  // KERNELS_SOFTMAX_H_