[WASM] Add Concat and Transpose kernels

.vscode/settings.json

@@ -127,6 +127,7 @@
     "queue": "cpp",
     "stack": "cpp",
     "*.ipp": "cpp",
-    "forward_list": "cpp"
+    "forward_list": "cpp",
+    "hash_map": "cpp"
   }
 }

tfjs-backend-wasm/scripts/cpplint.js

@@ -29,4 +29,4 @@ console.log(result);
 const cwd = process.cwd() + '/' + CC_FILEPATH;
 
 const filenameArgument = result.join(' ');
-exec(`python tools/cpplint.py --root ${cwd} ${filenameArgument}`);
+exec(`python2 tools/cpplint.py --root ${cwd} ${filenameArgument}`);

tfjs-backend-wasm/src/backend_wasm.ts

@@ -124,16 +124,18 @@ export class BackendWasm extends KernelBackend {
     return {dataId, shape, dtype};
   }
 
-  typedArrayFromHeap(offset: number, dtype: DataType, size: number):
+  typedArrayFromHeap({shape, dtype, dataId}: TensorInfo):
       backend_util.TypedArray {
     const buffer = this.wasm.HEAPU8.buffer;
+    const {memoryOffset} = this.dataIdMap.get(dataId);
+    const size = util.sizeFromShape(shape);
     switch (dtype) {
       case 'float32':
-        return new Float32Array(buffer, offset, size);
+        return new Float32Array(buffer, memoryOffset, size);
       case 'int32':
-        return new Int32Array(buffer, offset, size);
+        return new Int32Array(buffer, memoryOffset, size);
       case 'bool':
-        return new Uint8Array(buffer, offset, size);
+        return new Uint8Array(buffer, memoryOffset, size);
       default:
         throw new Error(`Uknown dtype ${dtype}`);
     }

tfjs-backend-wasm/src/cc/BUILD

@@ -69,6 +69,7 @@ tfjs_cc_library(
     ":Min",
     ":Sigmoid",
     ":Sub",
+    ":Transpose",
   ]
 )
 
@@ -185,20 +186,31 @@ tfjs_cc_library(
   ],
 )
 
+tfjs_cc_library(
+  name = "Transpose",
+  srcs = ["kernels/Transpose.cc"],
+  deps = [
+    ":backend",
+    ":util",
+  ],
+)
+
 tfjs_cc_library(
   name = "util",
-  srcs = ["util.h"],
+  hdrs = ["util.h"],
+  srcs = ["util.cc"],
 )
 
 test_suite(
   name = "cc_tests"
 )
 
 tfjs_unit_test(
-  name = "backend_test",
-  srcs = ["backend_test.cc"],
+  name = "backend_tests",
+  srcs = glob(["*_test.cc"]),
   deps = [
     ":backend",
+    ":util",
     ":Prelu",
   ]
 )

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

@@ -37,7 +37,7 @@ std::unordered_map<int, std::vector<tfjs::backend::DisposeFunction>>
 
 namespace tfjs {
 namespace backend {
-TensorInfo get_tensor_info(int tensor_id) { return data.at(tensor_id); }
+TensorInfo &get_tensor_info(int tensor_id) { return data.at(tensor_id); }
 
 int xnn_operator_count = 0;
 

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

@@ -45,7 +45,7 @@ namespace tfjs {
 namespace backend {
 // Returns the tensor information object associated with a given tensor_id
 // bucket.
-TensorInfo get_tensor_info(int tensor_id);
+TensorInfo &get_tensor_info(int tensor_id);
 
 // Registers a function callback to be called when a tensor with a given ID is
 // disposed.

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

@@ -51,7 +51,7 @@ void BatchMatMul(int a_id, int b_id, int shared_dim, int left_dim,
   // Zero out the output buffer because it might have been used before.
   std::fill(out_buf, out_buf + batch_dim * size, 0);
 
-  for (int b = 0; b < batch_dim; b++) {
+  for (int b = 0; b < batch_dim; ++b) {
     for (int i0 = 0; i0 < left_dim; i0 += kBlockSize) {
       for (int j0 = 0; j0 < right_dim; j0 += kBlockSize) {
         for (int k0 = 0; k0 < shared_dim; k0 += kBlockSize) {
@@ -60,11 +60,11 @@ void BatchMatMul(int a_id, int b_id, int shared_dim, int left_dim,
           int j_block = std::min(j0 + kBlockSize, right_dim);
           int k_block = std::min(k0 + kBlockSize, shared_dim);
 
-          for (int i = i0; i < i_block; i++) {
-            for (int j = j0; j < j_block; j++) {
+          for (int i = i0; i < i_block; ++i) {
+            for (int j = j0; j < j_block; ++j) {
               float sum = 0.0;
 
-              for (int k = k0; k < k_block; k++) {
+              for (int k = k0; k < k_block; ++k) {
                 sum +=
                     a_buf[b * a_batch + i * a_outer_step + k * a_inner_step] *
                     b_buf[k * b_inner_step + j * b_outer_step + b * b_batch];