Add Prod kernel to WASM backend. (#4138)

FEATURE
Co-authored-by: Ann Yuan <[email protected]>

Author: miaowzhang

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

@@ -0,0 +1,85 @@
+/* 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 <cstddef>
+
+#include "src/cc/backend.h"
+#include "src/cc/util.h"
+
+namespace {
+
+template <typename T>
+void prod(const size_t x_id, const size_t reduce_size, const size_t out_id) {
+  auto& x_info = tfjs::backend::get_tensor_info(x_id);
+  auto& out_info = tfjs::backend::get_tensor_info_out(out_id);
+
+  const T* x_buf = reinterpret_cast<const T *>(x_info.memory_offset);
+  const size_t x_size = x_info.size;
+
+  T* out_buf = reinterpret_cast<T *>(out_info.memory_offset);
+  const size_t out_size = out_info.size;
+
+  const T* x_offset = x_buf;
+
+  for (size_t i = 0; i < out_size; ++i) {
+    const size_t offset = i * reduce_size;
+    T product = 1;
+
+    const T* x_iter_end = x_offset + reduce_size;
+
+    for (const T* x = x_offset; x < x_iter_end; ++x) {
+      product *= (*x);
+    }
+
+    x_offset += reduce_size;
+    out_buf[i] = product;
+  }
+}
+
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void Prod(const size_t x_id, const size_t reduce_size,
+          const DType dtype, const size_t out_id) {
+  switch (dtype) {
+    case DType::float32:
+      prod<float>(x_id, reduce_size, out_id);
+      break;
+    case DType::int32:
+      prod<int32_t>(x_id, reduce_size, out_id);
+      break;
+    case DType::boolean:
+      prod<bool>(x_id, reduce_size, out_id);
+      break;
+    default:
+      util::warn("Prod failed. Unknown dtype %d", dtype);
+  }
+}
+
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

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

@@ -0,0 +1,98 @@
+/**
+ * @license
+ * 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.
+ * =============================================================================
+ */
+
+import {backend_util, KernelConfig, KernelFunc, Prod, ProdAttrs, ProdInputs, TensorInfo, util} from '@tensorflow/tfjs-core';
+
+import {BackendWasm} from '../backend_wasm';
+
+import {permuteAxesAndTranspose} from './kernel_utils';
+
+import {CppDType} from './types';
+
+let wasmProd: (
+    xId: number, reduceSize: number,
+    dtype: number, outId: number) => void;
+
+function setup(backend: BackendWasm): void {
+  wasmProd = backend.wasm.cwrap(Prod, null /*void*/, [
+    'number',
+    'number',
+    'number',
+    'number'
+  ]);
+}
+
+function prod(args: {
+  backend: BackendWasm,
+  inputs: ProdInputs,
+  attrs: ProdAttrs
+}): TensorInfo {
+  const {backend, inputs, attrs} = args;
+  const {axis, keepDims} = attrs;
+  const {x} = inputs;
+  const xId = backend.dataIdMap.get(x.dataId).id;
+  let inputId = xId;
+  let input = x;
+
+  const {transposed, axes, originalAxes, inputWasTransposed} =
+      permuteAxesAndTranspose(x, axis, backend);
+
+  let reductionAxes = axes;
+  if (inputWasTransposed) {
+    const transposedId = backend.dataIdMap.get(transposed.dataId).id;
+    if (transposedId !== xId) {
+      // transpose was not a no-op. We will need to dispose of this
+      // once we are done.
+      input = transposed;
+      inputId = transposedId;
+      reductionAxes = backend_util.getInnerMostAxes(
+          reductionAxes.length, input.shape.length);
+    }
+  }
+
+  backend_util.assertAxesAreInnerMostDims(
+      'prod', reductionAxes, input.shape.length);
+  const [outShape, reduceShape] =
+      backend_util.computeOutAndReduceShapes(input.shape, reductionAxes);
+  const reduceSize = util.sizeFromShape(reduceShape);
+
+  const out = backend.makeOutput(outShape, input.dtype);
+  if (util.sizeFromShape(input.shape) !== 0) {
+    const outId = backend.dataIdMap.get(out.dataId).id;
+    wasmProd(inputId, reduceSize, CppDType[out.dtype], outId);
+  }
+
+  if (inputWasTransposed) {
+    // dispose of the transposed tensor.
+    backend.disposeData(transposed.dataId);
+  }
+
+  if (keepDims) {
+    // reshape
+    const newShape = backend_util.expandShapeToKeepDim(out.shape, originalAxes);
+    out.shape = newShape;
+  }
+
+  return out;
+}
+
+export const prodConfig: KernelConfig = {
+  kernelName: Prod,
+  backendName: 'wasm',
+  setupFunc: setup,
+  kernelFunc: prod as {} as KernelFunc
+};

tfjs-backend-wasm/src/register_all_kernels.ts

@@ -70,6 +70,7 @@ import {onesLikeConfig} from './kernels/OnesLike';
 import {padV2Config} from './kernels/PadV2';
 import {powConfig} from './kernels/Pow';
 import {preluConfig} from './kernels/Prelu';
+import {prodConfig} from './kernels/Prod';
 import {reluConfig} from './kernels/Relu';
 import {relu6Config} from './kernels/Relu6';
 import {reshapeConfig} from './kernels/Reshape';
@@ -149,6 +150,7 @@ const kernelConfigs: KernelConfig[] = [
   padV2Config,
   powConfig,
   preluConfig,
+  prodConfig,
   reluConfig,
   relu6Config,
   reshapeConfig,

tfjs-backend-wasm/src/setup_test.ts

@@ -376,7 +376,8 @@ const TEST_FILTERS: TestFilter[] = [
     startsWith: 'onesLike',
     // Complex numbers not supported yet.
     excludes: ['complex'],
-  }
+  },
+  {include: 'prod'},
 ];
 
 const customInclude = (testName: string) => {