[wasm] Add ArgMax kernel

tfjs-backend-wasm/src/cc/BUILD

@@ -91,6 +91,7 @@ tfjs_cc_library(
     ":Add",
     ":AvgPool",
     ":AddN",
+    ":ArgMax",
     ":BatchMatMul",
     ":MaxPool",
     ":ClipByValue",
@@ -189,6 +190,15 @@ tfjs_cc_library(
   ],
 )
 
+tfjs_cc_library(
+  name = "ArgMax",
+  srcs = ["kernels/ArgMax.cc"],
+  deps = [
+    ":backend",
+    ":util",
+  ],
+)
+
 tfjs_cc_library(
   name = "BatchMatMul",
   srcs = ["kernels/BatchMatMul.cc"],

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

@@ -0,0 +1,75 @@
+/* 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include "src/cc/backend.h"
+#include "src/cc/util.h"
+
+namespace {
+
+template <typename T>
+void argmax(const T* x, const int outer_size, const int inner_size,
+            int* out_buf) {
+  for (int i = 0; i < outer_size; ++i) {
+    const int offset = i * inner_size;
+    T max = x[offset];
+    int max_index = 0;
+    for (int j = 1; j < inner_size; ++j) {
+      const T val = x[offset + j];
+      if (val > max) {
+        max = val;
+        max_index = j;
+      }
+    }
+    out_buf[i] = max_index;
+  }
+}
+
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void ArgMax(const int x_id, const DType dtype, const int outer_size,
+            const int inner_size, const int out_id) {
+  auto& x_info = backend::get_tensor_info(x_id);
+  auto& out_info = backend::get_tensor_info_out(out_id);
+  int* out_buf = out_info.i32_write();
+
+  switch (dtype) {
+    case DType::float32:
+      argmax<float>(x_info.f32(), outer_size, inner_size, out_buf);
+      break;
+    case DType::int32:
+      argmax<int>(x_info.i32(), outer_size, inner_size, out_buf);
+      break;
+    case DType::boolean:
+      argmax<bool>(x_info.b(), outer_size, inner_size, out_buf);
+      break;
+    default:
+      util::warn("Argmax failed. Unknown dtype %d", dtype);
+  }
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

tfjs-backend-wasm/src/kernels/ArgMax.ts

@@ -0,0 +1,64 @@
+/**
+ * @license
+ * 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.
+ * =============================================================================
+ */
+
+import {KernelFunc, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+
+import {BackendWasm} from '../backend_wasm';
+
+import {CppDType} from './types';
+
+interface ArgMaxInputs {
+  x: TensorInfo;
+}
+
+interface ArgMaxAttrs {
+  axis: number;
+}
+
+let wasmFunc: (
+    xId: number, dtype: number, outerSize: number, innerSize: number,
+    outId: number) => void;
+
+function setup(backend: BackendWasm) {
+  wasmFunc = backend.wasm.cwrap('ArgMax', null /* void */, [
+    'number',  // x_id
+    'number',  // dtype
+    'number',  // outer_size
+    'number',  // inner_size
+    'number'   // out_id
+  ]);
+}
+
+function argmax(
+    args: {inputs: ArgMaxInputs, backend: BackendWasm, attrs: ArgMaxAttrs}) {
+  const {inputs: {x}, backend, attrs: {axis}} = args;
+  const outShape = x.shape.slice(0, -1);
+  const out = backend.makeOutput(outShape, 'int32');
+  const xId = backend.dataIdMap.get(x.dataId).id;
+  const outId = backend.dataIdMap.get(out.dataId).id;
+  const outerSize = util.sizeFromShape(out.shape);
+  const innerSize = x.shape[axis];
+  wasmFunc(xId, CppDType[x.dtype], outerSize, innerSize, outId);
+  return out;
+}
+
+registerKernel({
+  kernelName: 'ArgMax',
+  backendName: 'wasm',
+  kernelFunc: argmax as {} as KernelFunc,
+  setupFunc: setup
+});

tfjs-backend-wasm/src/kernels/all_kernels.ts

@@ -20,8 +20,9 @@
 // the contents of this file and import only the kernels that are needed.
 import './Abs';
 import './Add';
-import './AvgPool';
 import './AddN';
+import './ArgMax';
+import './AvgPool';
 import './BatchMatMul';
 import './Cast';
 import './ClipByValue';

tfjs-backend-wasm/src/setup_test.ts

@@ -190,6 +190,7 @@ const TEST_FILTERS: TestFilter[] = [
   {include: 'clip', excludes: ['gradient']},
   {include: 'addN'},
   {include: 'nonMaxSuppression'},
+  {include: 'argmax', excludes: ['gradient']},
 ];
 
 const customInclude = (testName: string) => {

tfjs-core/src/ops/reduction_ops.ts

@@ -477,13 +477,15 @@ function argMax_<T extends Tensor>(x: Tensor|TensorLike, axis = 0): T {
   }
   const grad = (dy: T, saved: Tensor[]) => {
     const [$x] = saved;
-    return {$x: () => zerosLike($x)};
+    return {x: () => zerosLike($x)};
   };
+  const attrs = {axis: axes[0]};
+  const inputsToSave = [$x];
   return ENGINE.runKernelFunc((backend, save) => {
     const res = backend.argMax($x, axes[0]);
     save([$x]);
     return res;
-  }, {$x}, grad) as T;
+  }, {x: $x}, grad, 'ArgMax', attrs, inputsToSave) as T;
 }
 
 /**