Modularize the WASM kernels into separate build targets. (#2173)

- Each kernel gets its own cc_library and file.
- A new cc_library encapsulates all of them: ":all_kernels"
- backend.cc contains global state and backend utilities.
- Add linkstatic=1 so the compiler yells before emscripten fails at linking (it doesn't support dynamic linking, e.g. extern globals).

DEV

Author: Nikhil Thorat

.vscode/settings.json

@@ -35,6 +35,97 @@
   "typescript.tsdk": "${workspaceRoot}/node_modules/typescript/lib",
   "clang-format.executable": "${workspaceRoot}/node_modules/.bin/clang-format",
   "files.associations": {
-    "memory": "cpp"
+    "memory": "cpp",
+    "any": "cpp",
+    "array": "cpp",
+    "atomic": "cpp",
+    "strstream": "cpp",
+    "*.tcc": "cpp",
+    "bitset": "cpp",
+    "cctype": "cpp",
+    "chrono": "cpp",
+    "cinttypes": "cpp",
+    "clocale": "cpp",
+    "cmath": "cpp",
+    "codecvt": "cpp",
+    "complex": "cpp",
+    "condition_variable": "cpp",
+    "cstdarg": "cpp",
+    "cstddef": "cpp",
+    "cstdint": "cpp",
+    "cstdio": "cpp",
+    "cstdlib": "cpp",
+    "cstring": "cpp",
+    "ctime": "cpp",
+    "cwchar": "cpp",
+    "cwctype": "cpp",
+    "deque": "cpp",
+    "list": "cpp",
+    "unordered_map": "cpp",
+    "unordered_set": "cpp",
+    "vector": "cpp",
+    "exception": "cpp",
+    "algorithm": "cpp",
+    "filesystem": "cpp",
+    "functional": "cpp",
+    "iterator": "cpp",
+    "map": "cpp",
+    "memory_resource": "cpp",
+    "numeric": "cpp",
+    "optional": "cpp",
+    "random": "cpp",
+    "ratio": "cpp",
+    "regex": "cpp",
+    "set": "cpp",
+    "string": "cpp",
+    "string_view": "cpp",
+    "system_error": "cpp",
+    "tuple": "cpp",
+    "type_traits": "cpp",
+    "utility": "cpp",
+    "fstream": "cpp",
+    "future": "cpp",
+    "initializer_list": "cpp",
+    "iomanip": "cpp",
+    "iosfwd": "cpp",
+    "iostream": "cpp",
+    "istream": "cpp",
+    "limits": "cpp",
+    "mutex": "cpp",
+    "new": "cpp",
+    "ostream": "cpp",
+    "shared_mutex": "cpp",
+    "sstream": "cpp",
+    "stdexcept": "cpp",
+    "streambuf": "cpp",
+    "thread": "cpp",
+    "cfenv": "cpp",
+    "typeinfo": "cpp",
+    "valarray": "cpp",
+    "variant": "cpp",
+    "charconv": "cpp",
+    "__bit_reference": "cpp",
+    "__config": "cpp",
+    "__debug": "cpp",
+    "__errc": "cpp",
+    "__functional_base": "cpp",
+    "__hash_table": "cpp",
+    "__locale": "cpp",
+    "__mutex_base": "cpp",
+    "__node_handle": "cpp",
+    "__nullptr": "cpp",
+    "__split_buffer": "cpp",
+    "__sso_allocator": "cpp",
+    "__std_stream": "cpp",
+    "__string": "cpp",
+    "__threading_support": "cpp",
+    "__tree": "cpp",
+    "__tuple": "cpp",
+    "bit": "cpp",
+    "ios": "cpp",
+    "locale": "cpp",
+    "queue": "cpp",
+    "stack": "cpp",
+    "*.ipp": "cpp"
   }
 }

tfjs-backend-wasm/src/backend_wasm.ts

@@ -212,7 +212,7 @@ async function init(): Promise<{wasm: BackendWasmModule}> {
       dispose: wasm.cwrap('dispose', voidReturnType, []),
       add: wasm.cwrap('add', voidReturnType, ['number', 'number', 'number']),
       batchMatMul: wasm.cwrap(
-          'batchMatMul', voidReturnType,
+          'batch_matmul', voidReturnType,
           [
             'number', 'number', 'number', 'number', 'number', 'number',
             'number', 'number', 'number', 'number', 'number', 'number', 'number'

tfjs-backend-wasm/src/cc/BUILD

@@ -1,6 +1,5 @@
 cc_binary(
     name = "tfjs-backend-wasm.js",
-    srcs = ["lib.cc"],
     linkopts = [
         "-s ALLOW_MEMORY_GROWTH=1",
         "-s DEFAULT_LIBRARY_FUNCS_TO_INCLUDE=[]",
@@ -15,19 +14,53 @@ cc_binary(
         "-s MALLOC=emmalloc",
     ],
     deps = [
-        ":kernels",
-        ":util",
+        ":backend",
+        ":all_kernels",
     ],
 )
 
 cc_library(
-    name = "kernels",
-    srcs = ["kernels.cc"],
-    hdrs = ["kernels.h"],
-    deps = ["@xnnpack//:XNNPACK"],
+  name = "backend",
+  srcs = ["backend.cc"],
+  hdrs = ["backend.h"],
+  linkstatic = 1,
+  deps = [
+    ":util",
+    "@xnnpack//:XNNPACK",
+  ],
 )
 
 cc_library(
-    name = "util",
-    srcs = ["util.h"],
+  name = "all_kernels",
+  linkstatic = 1,
+  deps = [
+    ":add",
+    ":batch_matmul",
+  ]
+)
+
+cc_library(
+  name = "add",
+  srcs = ["add.cc"],
+  linkstatic = 1,
+  deps = [
+    ":backend",
+    ":util",
+  ],
+)
+
+cc_library(
+  name = "batch_matmul",
+  srcs = ["batch_matmul.cc"],
+  linkstatic = 1,
+  deps = [
+    ":backend",
+    ":util",
+  ],
+)
+
+cc_library(
+  name = "util",
+  linkstatic = 1,
+  srcs = ["util.h"],
 )

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

@@ -0,0 +1,68 @@
+/* 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 <emscripten.h>
+#include <math.h>
+#include <cstdio>
+#include <map>
+#include <vector>
+
+#include "src/cc/backend.h"
+#include "src/cc/util.h"
+
+template <class T>
+void add_impl(T* a_buf, int a_size, T* b_buf, int b_size, T* out_buf) {
+  int size = std::max(a_size, b_size);
+  for (int i = 0; i < size; ++i) {
+    out_buf[i] = a_buf[i % a_size] + b_buf[i % b_size];
+  }
+}
+// Templates need explicit instantiation when implemented in a .cc file.
+template void add_impl<float>(float* a_buf, int a_size, float* b_buf,
+                              int b_size, float* out_buf);
+template void add_impl<int>(int* a_buf, int a_size, int* b_buf, int b_size,
+                            int* out_buf);
+template void add_impl<bool>(bool* a_buf, int a_size, bool* b_buf, int b_size,
+                             bool* out_buf);
+
+namespace tfjs {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+EMSCRIPTEN_KEEPALIVE
+void add(int a_id, int b_id, int out_id) {
+  const auto a_info = backend::get_tensor_info(a_id);
+  const auto b_info = backend::get_tensor_info(b_id);
+  const auto out_info = backend::get_tensor_info(out_id);
+  switch (a_info.dtype) {
+    case DType::float32:
+      add_impl(a_info.buf.f32, a_info.size, b_info.buf.f32, b_info.size,
+               out_info.buf.f32);
+      break;
+    case DType::int32:
+      add_impl(a_info.buf.i32, a_info.size, b_info.buf.i32, b_info.size,
+               out_info.buf.i32);
+      break;
+    case DType::boolean:
+      add_impl(a_info.buf.b, a_info.size, b_info.buf.b, b_info.size,
+               out_info.buf.b);
+      break;
+    default:
+      util::warn("Add for tensor ids %d and %d failed. Unknown dtype %d", a_id,
+                 b_id, a_info.dtype);
+  }
+}
+
+}  // extern "C"
+}  // namespace tfjs

tfjs-backend-wasm/src/cc/lib.cc renamed to tfjs-backend-wasm/src/cc/backend.cc

@@ -13,44 +13,24 @@
  * ===========================================================================*/
 
 #include <emscripten.h>
-#include <math.h>
 #include <xnnpack.h>
 #include <cstdio>
 #include <map>
 #include <vector>
 
-#include "src/cc/kernels.h"
+#include "src/cc/backend.h"
 #include "src/cc/util.h"
 
-namespace tfjs {
-
-enum DType {
-  float32 = 0,
-  int32 = 1,
-  boolean = 2,
-};
-
-// A union of pointers that points to memory for a given tensor.
-union DataPtrUnion {
-  float *f32;
-  int *i32;
-  bool *b;
-};
-
-// Holds information about a tensor such as dtype, shape and pointer to its data
-// in memory.
-struct TensorInfo {
-  // Pointer to the bytes where the data is allocated.
-  DataPtrUnion buf;
-  DType dtype;
-  std::vector<int> shape;
-  // Total number of elements.
-  int size;
-};
-
+namespace {
 // Maps a unique tensor id to info about that tensor. The map owns all of its
 // entries.
 std::map<int, TensorInfo> data;
+}  // namespace
+
+namespace tfjs {
+namespace backend {
+TensorInfo get_tensor_info(int tensor_id) { return data.at(tensor_id); }
+}  // namespace backend
 
 // We use C-style API to interface with Javascript.
 extern "C" {
@@ -103,52 +83,6 @@ void dispose_data(int data_id) {
   data.erase(data_id);
 }
 
-EMSCRIPTEN_KEEPALIVE
-void add(int a_id, int b_id, int out_id) {
-  const auto a_info = data.at(a_id);
-  const auto b_info = data.at(b_id);
-  const auto out_info = data.at(out_id);
-  switch (a_info.dtype) {
-    case DType::float32:
-      kernels::add(a_info.buf.f32, a_info.size, b_info.buf.f32, b_info.size,
-                   out_info.buf.f32);
-      break;
-    case DType::int32:
-      kernels::add(a_info.buf.i32, a_info.size, b_info.buf.i32, b_info.size,
-                   out_info.buf.i32);
-      break;
-    case DType::boolean:
-      kernels::add(a_info.buf.b, a_info.size, b_info.buf.b, b_info.size,
-                   out_info.buf.b);
-      break;
-    default:
-      util::warn("Add for tensor ids %d and %d failed. Unknown dtype %d", a_id,
-                 b_id, a_info.dtype);
-  }
-}
-
-EMSCRIPTEN_KEEPALIVE
-void batchMatMul(int a_id, int b_id, int shared_dim, int left_dim,
-                 int right_dim, int batch_dim, int a_batch, int a_outer_step,
-                 int a_inner_step, int b_batch, int b_outer_step,
-                 int b_inner_step, int out_id) {
-  const auto a_info = data.at(a_id);
-  const auto b_info = data.at(b_id);
-  const auto out_info = data.at(out_id);
-  switch (a_info.dtype) {
-    case DType::float32:
-      kernels::batchMatMul(a_info.buf.f32, b_info.buf.f32, shared_dim, left_dim,
-                           right_dim, batch_dim, a_batch, a_outer_step,
-                           a_inner_step, b_batch, b_outer_step, b_inner_step,
-                           out_info.buf.f32);
-      break;
-    default:
-      util::warn(
-          "batchMatMul for tensor ids %d and %d failed. Unknown dtype %d", a_id,
-          b_id, a_info.dtype);
-  }
-}
-
 EMSCRIPTEN_KEEPALIVE
 void dispose() {
   for (auto const &element : data) {