From d7ccaa4a626b2bb921d6ec5fcaa54dcb53f06fa3 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 12 Jan 2023 10:42:35 -0800 Subject: [PATCH 1/5] Add DenseBincount kernel --- tfjs-backend-wasm/src/cc/BUILD.bazel | 9 ++ .../src/cc/kernels/DenseBincount.cc | 136 ++++++++++++++++++ .../src/kernels/DenseBincount.ts | 84 +++++++++++ tfjs-backend-wasm/src/register_all_kernels.ts | 2 + tfjs-backend-wasm/src/setup_test.ts | 1 + 5 files changed, 232 insertions(+) create mode 100644 tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc create mode 100644 tfjs-backend-wasm/src/kernels/DenseBincount.ts diff --git a/tfjs-backend-wasm/src/cc/BUILD.bazel b/tfjs-backend-wasm/src/cc/BUILD.bazel index 48b423f7754..7f336bf6692 100644 --- a/tfjs-backend-wasm/src/cc/BUILD.bazel +++ b/tfjs-backend-wasm/src/cc/BUILD.bazel @@ -633,6 +633,15 @@ tfjs_cc_library( ], ) +tfjs_cc_library( + name = "DenseBincount", + srcs = ["kernels/DenseBincount.cc"], + deps = [ + ":backend", + ":util", + ], +) + tfjs_cc_library( name = "DepthToSpace", srcs = ["kernels/DepthToSpace.cc"], diff --git a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc new file mode 100644 index 00000000000..ac2af463984 --- /dev/null +++ b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc @@ -0,0 +1,136 @@ +/** + * @license + * Copyright 2023 Google LLC. + * 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 +#endif + +#include +#include "tfjs-backend-wasm/src/cc/backend.h" +#include "tfjs-backend-wasm/src/cc/util.h" + +namespace tfjs { +namespace wasm { + +namespace { + +template +inline void Bincount1D(const int32_t* x_buf, int32_t x_len, int32_t size, + const T* weight_buf, bool binary_output, T* out_buf) { + std::fill(out_buf, out_buf + size, 0); + for (int32_t i = 0; i < x_len; ++i) { + int32_t value = x_buf[i]; + if (value < 0) { + util::warn("DenseBincount error: input x must be non-negative."); + return; + } + if (value >= size) { + continue; + } + + if (binary_output) { + out_buf[value] = 1; + } else if (weight_buf == nullptr) { + out_buf[value] += 1; + } else { + out_buf[value] += weight_buf[i]; + } + } +} + +template +inline void Bincount2D(const int32_t* x_buf, int32_t x_shape_0, + int32_t x_shape_1, int32_t size, const T* weight_buf, + bool binary_output, T* out_buf) { + for (int32_t i = 0; i < x_shape_0; ++i) { + Bincount1D(x_buf + i * x_shape_1, x_shape_1, size, + weight_buf != nullptr ? weight_buf + i * x_shape_1 : nullptr, + binary_output, out_buf + i * size); + } +} + +} // namespace + +extern "C" { + +#ifdef __EMSCRIPTEN__ +EMSCRIPTEN_KEEPALIVE +#endif + +// REQUIRES: +// - Tensor `x` must have dtype int32 +// - Tensor `x` must be 1D or 2D +// - If has_weights is true, tensor `weights` must have the same shape as `x` +// - Tensor `out` must have shape [x.shape[0], size] or [size] +// - Tensor `out` must have the same dtype as weights +void DenseBincount(const int32_t x_id, const int32_t* x_shape_ptr, + const int32_t x_shape_len, const int32_t size, + const bool has_weights, const int32_t weights_id, + const DType weights_dtype, const bool binary_output, + const int32_t out_id) { + const TensorInfo& x_info = backend::get_tensor_info(x_id); + const TensorInfo* weights_info = + has_weights ? &backend::get_tensor_info(weights_id) : nullptr; + TensorInfo& out_info = backend::get_tensor_info_out(out_id); + + const int32_t* x_buf = x_info.i32(); + if (x_shape_len == 1) { + switch (weights_dtype) { + case DType::float32: + Bincount1D(x_buf, x_shape_ptr[0], size, + weights_info ? weights_info->f32() : nullptr, binary_output, + out_info.f32_write()); + break; + case DType::int32: + Bincount1D(x_buf, x_shape_ptr[0], size, + weights_info ? weights_info->i32() : nullptr, binary_output, + out_info.i32_write()); + break; + default: + util::warn( + "DenseBincount with weights tensor id %d failed. Unsupported " + "weights " + "dtype %d", + weights_id, weights_dtype); + } + return; + } + + // x_shape_len == 2 + switch (weights_dtype) { + case DType::float32: + Bincount2D(x_buf, x_shape_ptr[0], x_shape_ptr[1], size, + weights_info ? weights_info->f32() : nullptr, binary_output, + out_info.f32_write()); + break; + case DType::int32: + Bincount2D(x_buf, x_shape_ptr[0], x_shape_ptr[1], size, + weights_info ? weights_info->i32() : nullptr, binary_output, + out_info.i32_write()); + break; + default: + util::warn( + "DenseBincount with weights tensor id %d failed. Unsupported " + "weights " + "dtype %d", + weights_id, weights_dtype); + } +} + +} // extern "C" +} // namespace wasm +} // namespace tfjs diff --git a/tfjs-backend-wasm/src/kernels/DenseBincount.ts b/tfjs-backend-wasm/src/kernels/DenseBincount.ts new file mode 100644 index 00000000000..125deeec820 --- /dev/null +++ b/tfjs-backend-wasm/src/kernels/DenseBincount.ts @@ -0,0 +1,84 @@ +/** + * @license + * Copyright 2023 Google LLC. + * 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 {DenseBincount, DenseBincountAttrs, DenseBincountInputs, KernelConfig, KernelFunc, TensorInfo} from '@tensorflow/tfjs-core'; + +import {BackendWasm} from '../backend_wasm'; + +import {CppDType} from './types'; + +let wasmDenseBincount: ( + xId: number, xShape: Uint8Array, xShapeLen: number, size: number, + hasWeights: boolean, weightsId: number, weightsDType: CppDType, + binaryOutput: boolean, outId: number) => void; + +function setup(backend: BackendWasm): void { + wasmDenseBincount = backend.wasm.cwrap('DenseBincount', null /*void*/, [ + 'number', // xId + 'array', // xShape + 'number', // xShapeLen + 'number', // size + 'boolean', // hasWeights + 'number', // weightsId + 'number', // weightsDType + 'boolean', // binaryOutput + 'number', // outId + ]); +} + +function denseBincount(args: { + backend: BackendWasm, + inputs: DenseBincountInputs, + attrs: DenseBincountAttrs +}): TensorInfo { + const {backend, inputs, attrs} = args; + const {x, weights} = inputs; + const {size, binaryOutput} = attrs; + if (x.dtype !== 'int32') { + throw new Error('DenseBincount: Tensor x must dtype int32'); + } + + if (x.shape.length !== 1 && x.shape.length !== 2) { + throw new Error('DenseBincount: Tensor x must be 1D or 2D'); + } + const hasWeights = weights.shape.reduce((p, v) => p * v, 1) !== 0; + if (hasWeights && + (x.shape.length !== weights.shape.length || + x.shape.some((v, i) => v !== weights.shape[i]))) { + throw new Error( + 'DenseBincount: Tensor weights must be empty or have the same shape as tensor x'); + } + + const outShape = x.shape.length === 1 ? [size] : [x.shape[0], size]; + const out = backend.makeOutput(outShape, weights.dtype); + + function tensorId(x: TensorInfo) { + return backend.dataIdMap.get(x.dataId).id; + } + wasmDenseBincount( + tensorId(x), new Uint8Array(new Int32Array(x.shape).buffer), + x.shape.length, size, hasWeights, tensorId(weights), + CppDType[weights.dtype], binaryOutput, tensorId(out)); + + return out; +} + +export const denseBincountConfig: KernelConfig = { + kernelName: DenseBincount, + backendName: 'wasm', + setupFunc: setup, + kernelFunc: denseBincount as unknown as KernelFunc +}; diff --git a/tfjs-backend-wasm/src/register_all_kernels.ts b/tfjs-backend-wasm/src/register_all_kernels.ts index cc519fce755..9e877190c03 100644 --- a/tfjs-backend-wasm/src/register_all_kernels.ts +++ b/tfjs-backend-wasm/src/register_all_kernels.ts @@ -45,6 +45,7 @@ import {coshConfig} from './kernels/Cosh'; import {cropAndResizeConfig} from './kernels/CropAndResize'; import {cumprodConfig} from './kernels/Cumprod'; import {cumsumConfig} from './kernels/Cumsum'; +import {denseBincountConfig} from './kernels/DenseBincount'; import {depthToSpaceConfig} from './kernels/DepthToSpace'; import {depthwiseConv2dNativeConfig} from './kernels/DepthwiseConv2dNative'; import {eluConfig} from './kernels/Elu'; @@ -164,6 +165,7 @@ const kernelConfigs: KernelConfig[] = [ cropAndResizeConfig, cumprodConfig, cumsumConfig, + denseBincountConfig, depthToSpaceConfig, depthwiseConv2dNativeConfig, eluConfig, diff --git a/tfjs-backend-wasm/src/setup_test.ts b/tfjs-backend-wasm/src/setup_test.ts index 34e975ebea2..9d7653fc2c5 100644 --- a/tfjs-backend-wasm/src/setup_test.ts +++ b/tfjs-backend-wasm/src/setup_test.ts @@ -407,6 +407,7 @@ const TEST_FILTERS: TestFilter[] = [ {include: 'acosh '}, {include: 'asin '}, {include: 'asinh '}, + {include: 'denseBincount '}, ]; const customInclude = (testName: string) => { From bef8a6102ca78690b54930fb06451ad01233c80f Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 12 Jan 2023 11:52:27 -0800 Subject: [PATCH 2/5] Remove redundant checks --- tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc | 5 +++-- tfjs-backend-wasm/src/kernels/DenseBincount.ts | 13 ------------- 2 files changed, 3 insertions(+), 15 deletions(-) diff --git a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc index ac2af463984..a5dc374272e 100644 --- a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc +++ b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc @@ -72,9 +72,10 @@ EMSCRIPTEN_KEEPALIVE #endif // REQUIRES: -// - Tensor `x` must have dtype int32 -// - Tensor `x` must be 1D or 2D +// - Tensor `x` must have dtype int32 (checked in tfjs-core) +// - Tensor `x` must be 1D or 2D (checked in tfjs-core) // - If has_weights is true, tensor `weights` must have the same shape as `x` +// (checked in tfjs-core) // - Tensor `out` must have shape [x.shape[0], size] or [size] // - Tensor `out` must have the same dtype as weights void DenseBincount(const int32_t x_id, const int32_t* x_shape_ptr, diff --git a/tfjs-backend-wasm/src/kernels/DenseBincount.ts b/tfjs-backend-wasm/src/kernels/DenseBincount.ts index 125deeec820..f38e8bcb4c4 100644 --- a/tfjs-backend-wasm/src/kernels/DenseBincount.ts +++ b/tfjs-backend-wasm/src/kernels/DenseBincount.ts @@ -47,21 +47,8 @@ function denseBincount(args: { const {backend, inputs, attrs} = args; const {x, weights} = inputs; const {size, binaryOutput} = attrs; - if (x.dtype !== 'int32') { - throw new Error('DenseBincount: Tensor x must dtype int32'); - } - if (x.shape.length !== 1 && x.shape.length !== 2) { - throw new Error('DenseBincount: Tensor x must be 1D or 2D'); - } const hasWeights = weights.shape.reduce((p, v) => p * v, 1) !== 0; - if (hasWeights && - (x.shape.length !== weights.shape.length || - x.shape.some((v, i) => v !== weights.shape[i]))) { - throw new Error( - 'DenseBincount: Tensor weights must be empty or have the same shape as tensor x'); - } - const outShape = x.shape.length === 1 ? [size] : [x.shape[0], size]; const out = backend.makeOutput(outShape, weights.dtype); From eb9c5a2e37f52a1a7843149e0e36acdcc7ef8c49 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 12 Jan 2023 13:05:44 -0800 Subject: [PATCH 3/5] Reduce buf reset calls --- tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc index a5dc374272e..51a96a972bb 100644 --- a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc +++ b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc @@ -28,10 +28,12 @@ namespace wasm { namespace { -template +template inline void Bincount1D(const int32_t* x_buf, int32_t x_len, int32_t size, const T* weight_buf, bool binary_output, T* out_buf) { - std::fill(out_buf, out_buf + size, 0); + if (reset_out_buf) { + std::fill(out_buf, out_buf + size, 0); + } for (int32_t i = 0; i < x_len; ++i) { int32_t value = x_buf[i]; if (value < 0) { @@ -56,10 +58,12 @@ template inline void Bincount2D(const int32_t* x_buf, int32_t x_shape_0, int32_t x_shape_1, int32_t size, const T* weight_buf, bool binary_output, T* out_buf) { + std::fill(out_buf, out_buf + (x_shape_0 * size), 0); for (int32_t i = 0; i < x_shape_0; ++i) { - Bincount1D(x_buf + i * x_shape_1, x_shape_1, size, - weight_buf != nullptr ? weight_buf + i * x_shape_1 : nullptr, - binary_output, out_buf + i * size); + Bincount1D( + x_buf + i * x_shape_1, x_shape_1, size, + weight_buf != nullptr ? weight_buf + i * x_shape_1 : nullptr, + binary_output, out_buf + i * size); } } From 8f32e709c371567628cb854e45dce78f80be72b1 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 12 Jan 2023 11:12:18 -0800 Subject: [PATCH 4/5] Update to C++17 --- .bazelrc | 2 +- tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc | 6 ++---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/.bazelrc b/.bazelrc index 514b7e4285d..a5ca9ab41f0 100644 --- a/.bazelrc +++ b/.bazelrc @@ -8,7 +8,7 @@ build --symlink_prefix=dist/ # These compile flags are active no matter which build mode we are in # (dbg vs opt). For flags specific to build mode, see cc_toolchain_config.bzl. -build --cxxopt="-std=c++11" +build --cxxopt="-std=c++17" build --cxxopt="-fno-rtti" build --cxxopt="-fno-exceptions" build --cxxopt="-fomit-frame-pointer" diff --git a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc index 51a96a972bb..e095f39ed5f 100644 --- a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc +++ b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc @@ -23,8 +23,7 @@ #include "tfjs-backend-wasm/src/cc/backend.h" #include "tfjs-backend-wasm/src/cc/util.h" -namespace tfjs { -namespace wasm { +namespace tfjs::wasm { namespace { @@ -137,5 +136,4 @@ void DenseBincount(const int32_t x_id, const int32_t* x_shape_ptr, } } // extern "C" -} // namespace wasm -} // namespace tfjs +} // namespace tfjs::wasm From be1f152c2094de788c40328caf382e63dc666727 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 19 Jan 2023 10:04:22 -0800 Subject: [PATCH 5/5] fix --- tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc index 51a96a972bb..e095f39ed5f 100644 --- a/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc +++ b/tfjs-backend-wasm/src/cc/kernels/DenseBincount.cc @@ -23,8 +23,7 @@ #include "tfjs-backend-wasm/src/cc/backend.h" #include "tfjs-backend-wasm/src/cc/util.h" -namespace tfjs { -namespace wasm { +namespace tfjs::wasm { namespace { @@ -137,5 +136,4 @@ void DenseBincount(const int32_t x_id, const int32_t* x_shape_ptr, } } // extern "C" -} // namespace wasm -} // namespace tfjs +} // namespace tfjs::wasm