From 4e62d106b9651cae15dfbebce9cbee99b617b84a Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Wed, 11 Jan 2023 20:36:53 -0800 Subject: [PATCH 1/9] wip --- tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc | 0 tfjs-backend-wasm/src/cc/pool3d_impl.cc | 101 ++++++++++++++++++ 2 files changed, 101 insertions(+) create mode 100644 tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc create mode 100644 tfjs-backend-wasm/src/cc/pool3d_impl.cc diff --git a/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc b/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tfjs-backend-wasm/src/cc/pool3d_impl.cc b/tfjs-backend-wasm/src/cc/pool3d_impl.cc new file mode 100644 index 00000000000..45ee21420c9 --- /dev/null +++ b/tfjs-backend-wasm/src/cc/pool3d_impl.cc @@ -0,0 +1,101 @@ +#ifdef __EMSCRIPTEN__ +#include +#endif + +#include +#include +#include + +#include "tfjs-backend-wasm/src/cc/backend.h" +#include "tfjs-backend-wasm/src/cc/util.h" + +namespace tfjs::wasm { + +namespace { + +inline int AddUntilNonNegative(int v, int d) { + if (v >= 0) { + return v; + } + return (v % d) + v; +} + +} // namespace + +struct Pool3DInfo { + int stride_batch; + int stride_depth; + int stride_height; + int stride_width; + int dilation_depth; + int dilation_height; + int dilation_width; + int effective_filter_depth; + int effective_filter_height; + int effective_filter_width; + int pad_front; + int pad_top; + int pad_left; + int in_channels; + int in_depth; + int in_height; + int out_depth; +}; +template +inline void Pool3DImpl( + const Pool3DInfo& info, const std::function& filter_init, + const std::function& filter_apply, + const std::function& filter_aggregate) { + for (int batch = 0; batch < info.batch_size; ++batch) { + for (int channel = 0; channel < info.in_channels; ++channel) { + for (int y_depth = 0; y_depth < info.out_depth; ++y_depth) { + int x_depth_corner = y_depth * info.stride_depth - info.pad_front; + int x_depth_min = + AddUntilNonNegative(x_depth_corner, info.dilation_depth); + int x_depth_max = std::min( + info.in_depth, info.effective_filter_depth + info.x_depth_corner); + + for (int y_row = 0; y_row < info.out_height; ++y_row) { + int x_row_corner = y_row * info.stride_height - info.pad_top; + int x_row_min = + AddUntilNonNegative(x_row_corner, info.dilation_height); + int x_row_max = std::min(info.in_height, + info.effective_filter_height + x_row_corner); + for (int y_col = 0; y_col < info.out_width; ++y_col) { + int x_col_corner = y_col * info.info.stride_width - info.pad_left; + int x_col_min = + AddUntilNonNegative(x_col_corner, info.dilation_width); + int x_col_max = std::min( + info.in_width, info.effective_filter_width + x_col_corner); + + // Apply the filter + T filter_data = filter_init(); + for (int x_depth = x_depth_min; x_depth < x_depth_max; + x_depth += info.dilation_depth) { + for (int x_row = x_row_min; x_row < x_row_max; + x_row += info.dilation_height) { + for (int x_col = x_col_min; x_col < x_col_max; + x_col += info.dilation_width) { + int x_offset = + (((batch * info.in_depth + x_depth) * info.in_height + + x_row) * + info.in_width + + x_col) * + info.in_channels + + channel; + filter_apply(filter_data, x_buf[x_offset]); + } + } + } + int out_offset = + + +channel; + out_buf[out_offset] = filter_aggregate(filter_data); + } + } + } + } + } +} + +} // namespace tfjs::wasm From cdbf2c4a178bf76652ad867dcbfca39a7051285b Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Fri, 20 Jan 2023 15:27:05 -0800 Subject: [PATCH 2/9] Add AvgPool3D kernel --- tfjs-backend-wasm/src/cc/BUILD.bazel | 18 +++ tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc | 91 +++++++++++++++ .../src/cc/{pool3d_impl.cc => pool3d_impl.h} | 62 ++++++---- tfjs-backend-wasm/src/kernels/AvgPool3D.ts | 107 ++++++++++++++++++ tfjs-backend-wasm/src/register_all_kernels.ts | 2 + tfjs-backend-wasm/src/setup_test.ts | 3 +- 6 files changed, 256 insertions(+), 27 deletions(-) rename tfjs-backend-wasm/src/cc/{pool3d_impl.cc => pool3d_impl.h} (63%) create mode 100644 tfjs-backend-wasm/src/kernels/AvgPool3D.ts diff --git a/tfjs-backend-wasm/src/cc/BUILD.bazel b/tfjs-backend-wasm/src/cc/BUILD.bazel index d19d41b7fa9..586ea76a6e8 100644 --- a/tfjs-backend-wasm/src/cc/BUILD.bazel +++ b/tfjs-backend-wasm/src/cc/BUILD.bazel @@ -245,6 +245,14 @@ tfjs_cc_library( ], ) +tfjs_cc_library( + name = "pool3d_impl", + hdrs = ["pool3d_impl.h"], + deps = [ + ":backend", + ], +) + tfjs_cc_library( name = "prelu_impl", srcs = ["prelu_impl.cc"], @@ -302,6 +310,7 @@ tfjs_cc_library( ":ArgMax", ":Atan", ":AvgPool", + ":AvgPool3D", ":BatchMatMul", ":Ceil", ":ClipByValue", @@ -491,6 +500,15 @@ tfjs_cc_library( ], ) +tfjs_cc_library( + name = "AvgPool3D", + srcs = ["kernels/AvgPool3D.cc"], + deps = [ + ":backend", + ":pool3d_impl", + ], +) + tfjs_unit_test( name = "AvgPool_test", srcs = ["kernels/AvgPool_test.cc"], diff --git a/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc b/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc index e69de29bb2d..fa344aed26f 100644 --- a/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc +++ b/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc @@ -0,0 +1,91 @@ +/** + * @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/pool3d_impl.h" + +namespace tfjs::wasm { + +// We use C-style API to interface with Javascript. +extern "C" { + +#ifdef __EMSCRIPTEN__ +EMSCRIPTEN_KEEPALIVE +#endif + +// REQUIRES: +// - Tensor `x` and `out` must have dtype float32 (checked in tfjs-core) +// - Tensor `x` and `out` must have data format 'NDHWC' (checked in tfjs-core) +void AvgPool3D(int x_id, int out_id, int batch_size, int in_depth, + int in_height, int in_width, int in_channels, int out_depth, + int out_height, int out_width, int out_channels, + int stride_depth, int stride_height, int stride_width, + int dilation_depth, int dilation_height, int dilation_width, + int effective_filter_depth, int effective_filter_height, + int effective_filter_width, int pad_front, int pad_top, + int pad_left) { + const TensorInfo& x_info = backend::get_tensor_info(x_id); + TensorInfo& out_info = backend::get_tensor_info_out(out_id); + + NDHWCPool3DImpl(x_info.f32(), out_info.f32_write(), + NDHWCPool3DInfo{ + .batch_size = batch_size, + .in_depth = in_depth, + .in_height = in_height, + .in_width = in_width, + .in_channels = in_channels, + .out_depth = out_depth, + .out_height = out_height, + .out_width = out_width, + .out_channels = out_channels, + .stride_depth = stride_depth, + .stride_height = stride_height, + .stride_width = stride_width, + .dilation_depth = dilation_depth, + .dilation_height = dilation_height, + .dilation_width = dilation_width, + .effective_filter_depth = effective_filter_depth, + .effective_filter_height = effective_filter_height, + .effective_filter_width = effective_filter_width, + .pad_front = pad_front, + .pad_top = pad_top, + .pad_left = pad_left, + }, + /*filter_init=*/ + []() -> std::pair { + return {0.0, 0}; + }, + /*filter_apply=*/ + [](std::pair& data, const float& val) { + data.first += val; + ++data.second; + }, + /*filter_aggregate=*/ + [](const std::pair& data) { + return data.first / + static_cast(std::max(data.second, 1)); + }); +} + +} // extern "C" +} // namespace tfjs::wasm diff --git a/tfjs-backend-wasm/src/cc/pool3d_impl.cc b/tfjs-backend-wasm/src/cc/pool3d_impl.h similarity index 63% rename from tfjs-backend-wasm/src/cc/pool3d_impl.cc rename to tfjs-backend-wasm/src/cc/pool3d_impl.h index 45ee21420c9..37c29aca266 100644 --- a/tfjs-backend-wasm/src/cc/pool3d_impl.cc +++ b/tfjs-backend-wasm/src/cc/pool3d_impl.h @@ -2,12 +2,10 @@ #include #endif -#include #include #include #include "tfjs-backend-wasm/src/cc/backend.h" -#include "tfjs-backend-wasm/src/cc/util.h" namespace tfjs::wasm { @@ -17,13 +15,22 @@ inline int AddUntilNonNegative(int v, int d) { if (v >= 0) { return v; } - return (v % d) + v; + return (v % d + v) % d; } } // namespace -struct Pool3DInfo { - int stride_batch; +struct NDHWCPool3DInfo { + int batch_size; + int in_depth; + int in_height; + int in_width; + int in_channels; + int out_depth; + int out_height; + int out_width; + int out_channels; + int stride_depth; int stride_height; int stride_width; @@ -36,16 +43,27 @@ struct Pool3DInfo { int pad_front; int pad_top; int pad_left; - int in_channels; - int in_depth; - int in_height; - int out_depth; + + inline int in_offset(int b, int d, int h, int w, int c) const { + return c + + (w + (h + (d + b * in_depth) * in_height) * in_width) * in_channels; + } + inline int out_offset(int b, int d, int h, int w, int c) const { + return c + (w + (h + (d + b * out_depth) * out_height) * out_width) * + out_channels; + } + inline int in_size() const { + return batch_size * in_depth * in_height * in_width * in_channels; + } + inline int out_size() const { + return batch_size * out_depth * out_height * out_width * out_channels; + } }; -template -inline void Pool3DImpl( - const Pool3DInfo& info, const std::function& filter_init, - const std::function& filter_apply, - const std::function& filter_aggregate) { +template +inline void NDHWCPool3DImpl(const IN* x_buf, OUT* out_buf, + const NDHWCPool3DInfo& info, const FI& filter_init, + const FAP& filter_apply, + const FAG& filter_aggregate) { for (int batch = 0; batch < info.batch_size; ++batch) { for (int channel = 0; channel < info.in_channels; ++channel) { for (int y_depth = 0; y_depth < info.out_depth; ++y_depth) { @@ -53,7 +71,7 @@ inline void Pool3DImpl( int x_depth_min = AddUntilNonNegative(x_depth_corner, info.dilation_depth); int x_depth_max = std::min( - info.in_depth, info.effective_filter_depth + info.x_depth_corner); + info.in_depth, info.effective_filter_depth + x_depth_corner); for (int y_row = 0; y_row < info.out_height; ++y_row) { int x_row_corner = y_row * info.stride_height - info.pad_top; @@ -62,14 +80,14 @@ inline void Pool3DImpl( int x_row_max = std::min(info.in_height, info.effective_filter_height + x_row_corner); for (int y_col = 0; y_col < info.out_width; ++y_col) { - int x_col_corner = y_col * info.info.stride_width - info.pad_left; + int x_col_corner = y_col * info.stride_width - info.pad_left; int x_col_min = AddUntilNonNegative(x_col_corner, info.dilation_width); int x_col_max = std::min( info.in_width, info.effective_filter_width + x_col_corner); // Apply the filter - T filter_data = filter_init(); + auto filter_data = filter_init(); for (int x_depth = x_depth_min; x_depth < x_depth_max; x_depth += info.dilation_depth) { for (int x_row = x_row_min; x_row < x_row_max; @@ -77,19 +95,13 @@ inline void Pool3DImpl( for (int x_col = x_col_min; x_col < x_col_max; x_col += info.dilation_width) { int x_offset = - (((batch * info.in_depth + x_depth) * info.in_height + - x_row) * - info.in_width + - x_col) * - info.in_channels + - channel; + info.in_offset(batch, x_depth, x_row, x_col, channel); filter_apply(filter_data, x_buf[x_offset]); } } } int out_offset = - - +channel; + info.out_offset(batch, y_depth, y_row, y_col, channel); out_buf[out_offset] = filter_aggregate(filter_data); } } diff --git a/tfjs-backend-wasm/src/kernels/AvgPool3D.ts b/tfjs-backend-wasm/src/kernels/AvgPool3D.ts new file mode 100644 index 00000000000..ef6cd93d27e --- /dev/null +++ b/tfjs-backend-wasm/src/kernels/AvgPool3D.ts @@ -0,0 +1,107 @@ +/** + * @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 {AvgPool3D, AvgPool3DAttrs, AvgPool3DInputs, backend_util, KernelConfig, KernelFunc, TensorInfo} from '@tensorflow/tfjs-core'; + +import {BackendWasm} from '../backend_wasm'; + +let wasmAvgPool3D: ( + xId: number, outId: number, batchSize: number, inDepth: number, + inHeight: number, inWidth: number, inChannels: number, outDepth: number, + outHeight: number, outWidth: number, outChannels: number, + strideDepth: number, strideHeight: number, strideWidth: number, + dilationDepth: number, dilationHeight: number, dilationWidth: number, + effectiveFilterDepth: number, effectiveFilterHeight: number, + effectiveFilterWidth: number, padFront: number, padTop: number, + padLeft: number) => void; + +function setup(backend: BackendWasm) { + wasmAvgPool3D = backend.wasm.cwrap('AvgPool3D', null, [ + 'number', // xId + 'number', // outId + 'number', // batchSize + 'number', // inDepth + 'number', // inHeight + 'number', // inWidth + 'number', // inChannels + 'number', // outDepth + 'number', // outHeight + 'number', // outWidth + 'number', // outChannels + 'number', // strideDepth + 'number', // strideHeight + 'number', // strideWidth + 'number', // dilationDepth + 'number', // dilationHeight + 'number', // dilationWidth + 'number', // effectiveFilterDepth + 'number', // effectiveFilterHeight + 'number', // effectiveFilterWidth + 'number', // padFront + 'number', // padTop + 'number', // padLeft + ]); +} + +export function avgPool3D(args: { + inputs: AvgPool3DInputs, + attrs: AvgPool3DAttrs, + backend: BackendWasm, +}): TensorInfo { + const {inputs, backend, attrs} = args; + const {x} = inputs; + const {filterSize, strides, pad, dimRoundingMode, dataFormat} = attrs; + + const convInfo = backend_util.computePool3DInfo( + x.shape as [number, number, number, number, number], filterSize, strides, + 1 /* dilations */, pad, dimRoundingMode, dataFormat); + const out = backend.makeOutput(convInfo.outShape, x.dtype); + + wasmAvgPool3D( + backend.dataIdMap.get(x.dataId).id, + backend.dataIdMap.get(out.dataId).id, + convInfo.batchSize, + convInfo.inDepth, + convInfo.inHeight, + convInfo.inWidth, + convInfo.inChannels, + convInfo.outDepth, + convInfo.outHeight, + convInfo.outWidth, + convInfo.outChannels, + convInfo.strideDepth, + convInfo.strideHeight, + convInfo.strideWidth, + convInfo.dilationDepth, + convInfo.dilationHeight, + convInfo.dilationWidth, + convInfo.effectiveFilterDepth, + convInfo.effectiveFilterHeight, + convInfo.effectiveFilterWidth, + convInfo.padInfo.front, + convInfo.padInfo.top, + convInfo.padInfo.left, + ); + return out; +} + +export const avgPool3DConfig: KernelConfig = { + kernelName: AvgPool3D, + backendName: 'wasm', + setupFunc: setup, + kernelFunc: avgPool3D 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 d9017380aa2..7e5542dad3f 100644 --- a/tfjs-backend-wasm/src/register_all_kernels.ts +++ b/tfjs-backend-wasm/src/register_all_kernels.ts @@ -32,6 +32,7 @@ import {asinConfig} from './kernels/Asin'; import {asinhConfig} from './kernels/Asinh'; import {atanConfig} from './kernels/Atan'; import {avgPoolConfig} from './kernels/AvgPool'; +import {avgPool3DConfig} from './kernels/AvgPool3D'; import {batchMatMulConfig} from './kernels/BatchMatMul'; import {batchToSpaceNDConfig} from './kernels/BatchToSpaceND'; import {castConfig} from './kernels/Cast'; @@ -153,6 +154,7 @@ const kernelConfigs: KernelConfig[] = [ asinhConfig, atanConfig, avgPoolConfig, + avgPool3DConfig, batchMatMulConfig, batchToSpaceNDConfig, castConfig, diff --git a/tfjs-backend-wasm/src/setup_test.ts b/tfjs-backend-wasm/src/setup_test.ts index 0726b2f2bb4..a56493890ef 100644 --- a/tfjs-backend-wasm/src/setup_test.ts +++ b/tfjs-backend-wasm/src/setup_test.ts @@ -66,8 +66,7 @@ const TEST_FILTERS: TestFilter[] = [ { include: 'avgPool', excludes: [ - 'gradient', // Not yet implemented. - 'avgPool3d', // Not yet implemented. + 'gradient', // Not yet implemented. ] }, { From 9bd8458274af0b8b2dc53ae873e5f08e14fca24e Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Fri, 20 Jan 2023 15:38:06 -0800 Subject: [PATCH 3/9] Add MaxPool3D kernel --- tfjs-backend-wasm/src/cc/BUILD.bazel | 10 ++ tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc | 85 ++++++++++++++ tfjs-backend-wasm/src/kernels/MaxPool3D.ts | 107 ++++++++++++++++++ tfjs-backend-wasm/src/register_all_kernels.ts | 2 + 4 files changed, 204 insertions(+) create mode 100644 tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc create mode 100644 tfjs-backend-wasm/src/kernels/MaxPool3D.ts diff --git a/tfjs-backend-wasm/src/cc/BUILD.bazel b/tfjs-backend-wasm/src/cc/BUILD.bazel index 586ea76a6e8..aad27fa673c 100644 --- a/tfjs-backend-wasm/src/cc/BUILD.bazel +++ b/tfjs-backend-wasm/src/cc/BUILD.bazel @@ -341,6 +341,7 @@ tfjs_cc_library( ":LessEqual", ":Max", ":MaxPool", + ":MaxPool3D", ":Maximum", ":Min", ":Minimum", @@ -968,6 +969,15 @@ tfjs_unit_test( ], ) +tfjs_cc_library( + name = "MaxPool3D", + srcs = ["kernels/MaxPool3D.cc"], + deps = [ + ":backend", + ":pool3d_impl", + ], +) + tfjs_cc_library( name = "Min", srcs = ["kernels/Min.cc"], diff --git a/tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc b/tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc new file mode 100644 index 00000000000..a745b71fe98 --- /dev/null +++ b/tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc @@ -0,0 +1,85 @@ +/** + * @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 + +#include "tfjs-backend-wasm/src/cc/backend.h" +#include "tfjs-backend-wasm/src/cc/pool3d_impl.h" + +namespace tfjs::wasm { + +// We use C-style API to interface with Javascript. +extern "C" { + +#ifdef __EMSCRIPTEN__ +EMSCRIPTEN_KEEPALIVE +#endif + +// REQUIRES: +// - Tensor `x` and `out` must have dtype float32 (checked in tfjs-core) +// - Tensor `x` and `out` must have data format 'NDHWC' (checked in tfjs-core) +void MaxPool3D(int x_id, int out_id, int batch_size, int in_depth, + int in_height, int in_width, int in_channels, int out_depth, + int out_height, int out_width, int out_channels, + int stride_depth, int stride_height, int stride_width, + int dilation_depth, int dilation_height, int dilation_width, + int effective_filter_depth, int effective_filter_height, + int effective_filter_width, int pad_front, int pad_top, + int pad_left) { + const TensorInfo& x_info = backend::get_tensor_info(x_id); + TensorInfo& out_info = backend::get_tensor_info_out(out_id); + + NDHWCPool3DImpl( + x_info.f32(), out_info.f32_write(), + NDHWCPool3DInfo{ + .batch_size = batch_size, + .in_depth = in_depth, + .in_height = in_height, + .in_width = in_width, + .in_channels = in_channels, + .out_depth = out_depth, + .out_height = out_height, + .out_width = out_width, + .out_channels = out_channels, + .stride_depth = stride_depth, + .stride_height = stride_height, + .stride_width = stride_width, + .dilation_depth = dilation_depth, + .dilation_height = dilation_height, + .dilation_width = dilation_width, + .effective_filter_depth = effective_filter_depth, + .effective_filter_height = effective_filter_height, + .effective_filter_width = effective_filter_width, + .pad_front = pad_front, + .pad_top = pad_top, + .pad_left = pad_left, + }, + /*filter_init=*/ + []() -> float { return std::numeric_limits::min(); }, + /*filter_apply=*/ + [](float& data, const float& val) { data = std::max(data, val); }, + /*filter_aggregate=*/ + [](const float& data) { return data; }); +} + +} // extern "C" +} // namespace tfjs::wasm diff --git a/tfjs-backend-wasm/src/kernels/MaxPool3D.ts b/tfjs-backend-wasm/src/kernels/MaxPool3D.ts new file mode 100644 index 00000000000..c8af6295491 --- /dev/null +++ b/tfjs-backend-wasm/src/kernels/MaxPool3D.ts @@ -0,0 +1,107 @@ +/** + * @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 {backend_util, KernelConfig, KernelFunc, MaxPool3D, MaxPool3DAttrs, MaxPool3DInputs, TensorInfo} from '@tensorflow/tfjs-core'; + +import {BackendWasm} from '../backend_wasm'; + +let wasmMaxPool3D: ( + xId: number, outId: number, batchSize: number, inDepth: number, + inHeight: number, inWidth: number, inChannels: number, outDepth: number, + outHeight: number, outWidth: number, outChannels: number, + strideDepth: number, strideHeight: number, strideWidth: number, + dilationDepth: number, dilationHeight: number, dilationWidth: number, + effectiveFilterDepth: number, effectiveFilterHeight: number, + effectiveFilterWidth: number, padFront: number, padTop: number, + padLeft: number) => void; + +function setup(backend: BackendWasm) { + wasmMaxPool3D = backend.wasm.cwrap('MaxPool3D', null, [ + 'number', // xId + 'number', // outId + 'number', // batchSize + 'number', // inDepth + 'number', // inHeight + 'number', // inWidth + 'number', // inChannels + 'number', // outDepth + 'number', // outHeight + 'number', // outWidth + 'number', // outChannels + 'number', // strideDepth + 'number', // strideHeight + 'number', // strideWidth + 'number', // dilationDepth + 'number', // dilationHeight + 'number', // dilationWidth + 'number', // effectiveFilterDepth + 'number', // effectiveFilterHeight + 'number', // effectiveFilterWidth + 'number', // padFront + 'number', // padTop + 'number', // padLeft + ]); +} + +export function maxPool3D(args: { + inputs: MaxPool3DInputs, + attrs: MaxPool3DAttrs, + backend: BackendWasm, +}): TensorInfo { + const {inputs, backend, attrs} = args; + const {x} = inputs; + const {filterSize, strides, pad, dimRoundingMode, dataFormat} = attrs; + + const convInfo = backend_util.computePool3DInfo( + x.shape as [number, number, number, number, number], filterSize, strides, + 1 /* dilations */, pad, dimRoundingMode, dataFormat); + const out = backend.makeOutput(convInfo.outShape, x.dtype); + + wasmMaxPool3D( + backend.dataIdMap.get(x.dataId).id, + backend.dataIdMap.get(out.dataId).id, + convInfo.batchSize, + convInfo.inDepth, + convInfo.inHeight, + convInfo.inWidth, + convInfo.inChannels, + convInfo.outDepth, + convInfo.outHeight, + convInfo.outWidth, + convInfo.outChannels, + convInfo.strideDepth, + convInfo.strideHeight, + convInfo.strideWidth, + convInfo.dilationDepth, + convInfo.dilationHeight, + convInfo.dilationWidth, + convInfo.effectiveFilterDepth, + convInfo.effectiveFilterHeight, + convInfo.effectiveFilterWidth, + convInfo.padInfo.front, + convInfo.padInfo.top, + convInfo.padInfo.left, + ); + return out; +} + +export const maxPool3DConfig: KernelConfig = { + kernelName: MaxPool3D, + backendName: 'wasm', + setupFunc: setup, + kernelFunc: maxPool3D 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 7e5542dad3f..41720020cfa 100644 --- a/tfjs-backend-wasm/src/register_all_kernels.ts +++ b/tfjs-backend-wasm/src/register_all_kernels.ts @@ -78,6 +78,7 @@ import {logicalXorConfig} from './kernels/LogicalXor'; import {maxConfig} from './kernels/Max'; import {maximumConfig} from './kernels/Maximum'; import {maxPoolConfig} from './kernels/MaxPool'; +import {maxPool3DConfig} from './kernels/MaxPool3D'; import {meanConfig} from './kernels/Mean'; import {minConfig} from './kernels/Min'; import {minimumConfig} from './kernels/Minimum'; @@ -200,6 +201,7 @@ const kernelConfigs: KernelConfig[] = [ maxConfig, maximumConfig, maxPoolConfig, + maxPool3DConfig, meanConfig, minConfig, minimumConfig, From e25264b5145d6ab5a88b4c7568be1429fdacaab6 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Fri, 20 Jan 2023 15:40:04 -0800 Subject: [PATCH 4/9] fix --- tfjs-backend-wasm/src/kernels/AvgPool3D.ts | 2 +- tfjs-backend-wasm/src/kernels/MaxPool3D.ts | 2 +- tfjs-backend-wasm/src/setup_test.ts | 2 -- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/tfjs-backend-wasm/src/kernels/AvgPool3D.ts b/tfjs-backend-wasm/src/kernels/AvgPool3D.ts index ef6cd93d27e..99f12d0b14a 100644 --- a/tfjs-backend-wasm/src/kernels/AvgPool3D.ts +++ b/tfjs-backend-wasm/src/kernels/AvgPool3D.ts @@ -68,7 +68,7 @@ export function avgPool3D(args: { const convInfo = backend_util.computePool3DInfo( x.shape as [number, number, number, number, number], filterSize, strides, - 1 /* dilations */, pad, dimRoundingMode, dataFormat); + /*dilations=*/1, pad, dimRoundingMode, dataFormat); const out = backend.makeOutput(convInfo.outShape, x.dtype); wasmAvgPool3D( diff --git a/tfjs-backend-wasm/src/kernels/MaxPool3D.ts b/tfjs-backend-wasm/src/kernels/MaxPool3D.ts index c8af6295491..2c25d01bd10 100644 --- a/tfjs-backend-wasm/src/kernels/MaxPool3D.ts +++ b/tfjs-backend-wasm/src/kernels/MaxPool3D.ts @@ -68,7 +68,7 @@ export function maxPool3D(args: { const convInfo = backend_util.computePool3DInfo( x.shape as [number, number, number, number, number], filterSize, strides, - 1 /* dilations */, pad, dimRoundingMode, dataFormat); + /*dilations=*/1, pad, dimRoundingMode, dataFormat); const out = backend.makeOutput(convInfo.outShape, x.dtype); wasmMaxPool3D( diff --git a/tfjs-backend-wasm/src/setup_test.ts b/tfjs-backend-wasm/src/setup_test.ts index a56493890ef..67a97e7b066 100644 --- a/tfjs-backend-wasm/src/setup_test.ts +++ b/tfjs-backend-wasm/src/setup_test.ts @@ -84,11 +84,9 @@ const TEST_FILTERS: TestFilter[] = [ include: 'maxPool', excludes: [ 'maxPoolBackprop', // Not yet implemented. - 'maxPool3d', // Not yet implemented. 'maxPool3dBackprop', // Not yet implemented. 'ignores NaNs', // Actual != expected. 'maxPoolWithArgmax' // Not yet implemented. - ] }, {include: 'cropAndResize'}, From 68550f6952923a37d298350cf73ec7e44d29ce24 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Fri, 20 Jan 2023 15:46:15 -0800 Subject: [PATCH 5/9] Fix template types order --- tfjs-backend-wasm/src/cc/pool3d_impl.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-wasm/src/cc/pool3d_impl.h b/tfjs-backend-wasm/src/cc/pool3d_impl.h index 37c29aca266..aa8dd66c320 100644 --- a/tfjs-backend-wasm/src/cc/pool3d_impl.h +++ b/tfjs-backend-wasm/src/cc/pool3d_impl.h @@ -59,7 +59,7 @@ struct NDHWCPool3DInfo { return batch_size * out_depth * out_height * out_width * out_channels; } }; -template +template inline void NDHWCPool3DImpl(const IN* x_buf, OUT* out_buf, const NDHWCPool3DInfo& info, const FI& filter_init, const FAP& filter_apply, From 217b40e58250248835fb10e1d654b4429c2bb9bc Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Fri, 20 Jan 2023 20:29:45 -0800 Subject: [PATCH 6/9] Update param channel_size --- tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc | 13 +++++------ tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc | 13 +++++------ tfjs-backend-wasm/src/cc/pool3d_impl.h | 15 ++++++------ tfjs-backend-wasm/src/kernels/AvgPool3D.ts | 23 +++++++++---------- tfjs-backend-wasm/src/kernels/MaxPool3D.ts | 23 +++++++++---------- 5 files changed, 42 insertions(+), 45 deletions(-) diff --git a/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc b/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc index fa344aed26f..7aedde7c78a 100644 --- a/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc +++ b/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc @@ -36,11 +36,11 @@ EMSCRIPTEN_KEEPALIVE // REQUIRES: // - Tensor `x` and `out` must have dtype float32 (checked in tfjs-core) // - Tensor `x` and `out` must have data format 'NDHWC' (checked in tfjs-core) -void AvgPool3D(int x_id, int out_id, int batch_size, int in_depth, - int in_height, int in_width, int in_channels, int out_depth, - int out_height, int out_width, int out_channels, - int stride_depth, int stride_height, int stride_width, - int dilation_depth, int dilation_height, int dilation_width, +void AvgPool3D(int x_id, int out_id, int batch_size, int channel_size, + int in_depth, int in_height, int in_width, int out_depth, + int out_height, int out_width, int stride_depth, + int stride_height, int stride_width, int dilation_depth, + int dilation_height, int dilation_width, int effective_filter_depth, int effective_filter_height, int effective_filter_width, int pad_front, int pad_top, int pad_left) { @@ -50,14 +50,13 @@ void AvgPool3D(int x_id, int out_id, int batch_size, int in_depth, NDHWCPool3DImpl(x_info.f32(), out_info.f32_write(), NDHWCPool3DInfo{ .batch_size = batch_size, + .channel_size = channel_size, .in_depth = in_depth, .in_height = in_height, .in_width = in_width, - .in_channels = in_channels, .out_depth = out_depth, .out_height = out_height, .out_width = out_width, - .out_channels = out_channels, .stride_depth = stride_depth, .stride_height = stride_height, .stride_width = stride_width, diff --git a/tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc b/tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc index a745b71fe98..49c841012fe 100644 --- a/tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc +++ b/tfjs-backend-wasm/src/cc/kernels/MaxPool3D.cc @@ -37,11 +37,11 @@ EMSCRIPTEN_KEEPALIVE // REQUIRES: // - Tensor `x` and `out` must have dtype float32 (checked in tfjs-core) // - Tensor `x` and `out` must have data format 'NDHWC' (checked in tfjs-core) -void MaxPool3D(int x_id, int out_id, int batch_size, int in_depth, - int in_height, int in_width, int in_channels, int out_depth, - int out_height, int out_width, int out_channels, - int stride_depth, int stride_height, int stride_width, - int dilation_depth, int dilation_height, int dilation_width, +void MaxPool3D(int x_id, int out_id, int batch_size, int channel_size, + int in_depth, int in_height, int in_width, int out_depth, + int out_height, int out_width, int stride_depth, + int stride_height, int stride_width, int dilation_depth, + int dilation_height, int dilation_width, int effective_filter_depth, int effective_filter_height, int effective_filter_width, int pad_front, int pad_top, int pad_left) { @@ -52,14 +52,13 @@ void MaxPool3D(int x_id, int out_id, int batch_size, int in_depth, x_info.f32(), out_info.f32_write(), NDHWCPool3DInfo{ .batch_size = batch_size, + .channel_size = channel_size, .in_depth = in_depth, .in_height = in_height, .in_width = in_width, - .in_channels = in_channels, .out_depth = out_depth, .out_height = out_height, .out_width = out_width, - .out_channels = out_channels, .stride_depth = stride_depth, .stride_height = stride_height, .stride_width = stride_width, diff --git a/tfjs-backend-wasm/src/cc/pool3d_impl.h b/tfjs-backend-wasm/src/cc/pool3d_impl.h index aa8dd66c320..eac53bf25e7 100644 --- a/tfjs-backend-wasm/src/cc/pool3d_impl.h +++ b/tfjs-backend-wasm/src/cc/pool3d_impl.h @@ -22,14 +22,15 @@ inline int AddUntilNonNegative(int v, int d) { struct NDHWCPool3DInfo { int batch_size; + // Since Pool3D ops (AvgPool3D and MaxPool3D) support 3D filter only, in + // channels should always equal to out channels. + int channel_size; int in_depth; int in_height; int in_width; - int in_channels; int out_depth; int out_height; int out_width; - int out_channels; int stride_depth; int stride_height; @@ -46,17 +47,17 @@ struct NDHWCPool3DInfo { inline int in_offset(int b, int d, int h, int w, int c) const { return c + - (w + (h + (d + b * in_depth) * in_height) * in_width) * in_channels; + (w + (h + (d + b * in_depth) * in_height) * in_width) * channel_size; } inline int out_offset(int b, int d, int h, int w, int c) const { return c + (w + (h + (d + b * out_depth) * out_height) * out_width) * - out_channels; + channel_size; } inline int in_size() const { - return batch_size * in_depth * in_height * in_width * in_channels; + return batch_size * in_depth * in_height * in_width * channel_size; } inline int out_size() const { - return batch_size * out_depth * out_height * out_width * out_channels; + return batch_size * out_depth * out_height * out_width * channel_size; } }; template @@ -65,7 +66,7 @@ inline void NDHWCPool3DImpl(const IN* x_buf, OUT* out_buf, const FAP& filter_apply, const FAG& filter_aggregate) { for (int batch = 0; batch < info.batch_size; ++batch) { - for (int channel = 0; channel < info.in_channels; ++channel) { + for (int channel = 0; channel < info.channel_size; ++channel) { for (int y_depth = 0; y_depth < info.out_depth; ++y_depth) { int x_depth_corner = y_depth * info.stride_depth - info.pad_front; int x_depth_min = diff --git a/tfjs-backend-wasm/src/kernels/AvgPool3D.ts b/tfjs-backend-wasm/src/kernels/AvgPool3D.ts index 99f12d0b14a..34d54ff6032 100644 --- a/tfjs-backend-wasm/src/kernels/AvgPool3D.ts +++ b/tfjs-backend-wasm/src/kernels/AvgPool3D.ts @@ -20,28 +20,26 @@ import {AvgPool3D, AvgPool3DAttrs, AvgPool3DInputs, backend_util, KernelConfig, import {BackendWasm} from '../backend_wasm'; let wasmAvgPool3D: ( - xId: number, outId: number, batchSize: number, inDepth: number, - inHeight: number, inWidth: number, inChannels: number, outDepth: number, - outHeight: number, outWidth: number, outChannels: number, - strideDepth: number, strideHeight: number, strideWidth: number, - dilationDepth: number, dilationHeight: number, dilationWidth: number, - effectiveFilterDepth: number, effectiveFilterHeight: number, - effectiveFilterWidth: number, padFront: number, padTop: number, - padLeft: number) => void; + xId: number, outId: number, batchSize: number, channelSize: number, + inDepth: number, inHeight: number, inWidth: number, outDepth: number, + outHeight: number, outWidth: number, strideDepth: number, + strideHeight: number, strideWidth: number, dilationDepth: number, + dilationHeight: number, dilationWidth: number, effectiveFilterDepth: number, + effectiveFilterHeight: number, effectiveFilterWidth: number, + padFront: number, padTop: number, padLeft: number) => void; function setup(backend: BackendWasm) { wasmAvgPool3D = backend.wasm.cwrap('AvgPool3D', null, [ 'number', // xId 'number', // outId 'number', // batchSize + 'number', // channelSize 'number', // inDepth 'number', // inHeight 'number', // inWidth - 'number', // inChannels 'number', // outDepth 'number', // outHeight 'number', // outWidth - 'number', // outChannels 'number', // strideDepth 'number', // strideHeight 'number', // strideWidth @@ -75,14 +73,15 @@ export function avgPool3D(args: { backend.dataIdMap.get(x.dataId).id, backend.dataIdMap.get(out.dataId).id, convInfo.batchSize, + // Since Pool3D ops (AvgPool3D and MaxPool3D) support 3D filter only, in + // channels should always equal to out channels. + /*channelSize=*/convInfo.inChannels, convInfo.inDepth, convInfo.inHeight, convInfo.inWidth, - convInfo.inChannels, convInfo.outDepth, convInfo.outHeight, convInfo.outWidth, - convInfo.outChannels, convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth, diff --git a/tfjs-backend-wasm/src/kernels/MaxPool3D.ts b/tfjs-backend-wasm/src/kernels/MaxPool3D.ts index 2c25d01bd10..5ab51b9b7b7 100644 --- a/tfjs-backend-wasm/src/kernels/MaxPool3D.ts +++ b/tfjs-backend-wasm/src/kernels/MaxPool3D.ts @@ -20,28 +20,26 @@ import {backend_util, KernelConfig, KernelFunc, MaxPool3D, MaxPool3DAttrs, MaxPo import {BackendWasm} from '../backend_wasm'; let wasmMaxPool3D: ( - xId: number, outId: number, batchSize: number, inDepth: number, - inHeight: number, inWidth: number, inChannels: number, outDepth: number, - outHeight: number, outWidth: number, outChannels: number, - strideDepth: number, strideHeight: number, strideWidth: number, - dilationDepth: number, dilationHeight: number, dilationWidth: number, - effectiveFilterDepth: number, effectiveFilterHeight: number, - effectiveFilterWidth: number, padFront: number, padTop: number, - padLeft: number) => void; + xId: number, outId: number, batchSize: number, channelSize: number, + inDepth: number, inHeight: number, inWidth: number, outDepth: number, + outHeight: number, outWidth: number, strideDepth: number, + strideHeight: number, strideWidth: number, dilationDepth: number, + dilationHeight: number, dilationWidth: number, effectiveFilterDepth: number, + effectiveFilterHeight: number, effectiveFilterWidth: number, + padFront: number, padTop: number, padLeft: number) => void; function setup(backend: BackendWasm) { wasmMaxPool3D = backend.wasm.cwrap('MaxPool3D', null, [ 'number', // xId 'number', // outId 'number', // batchSize + 'number', // channelSize 'number', // inDepth 'number', // inHeight 'number', // inWidth - 'number', // inChannels 'number', // outDepth 'number', // outHeight 'number', // outWidth - 'number', // outChannels 'number', // strideDepth 'number', // strideHeight 'number', // strideWidth @@ -75,14 +73,15 @@ export function maxPool3D(args: { backend.dataIdMap.get(x.dataId).id, backend.dataIdMap.get(out.dataId).id, convInfo.batchSize, + // Since Pool3D ops (AvgPool3D and MaxPool3D) support 3D filter only, in + // channels should always equal to out channels. + /*channelSize=*/convInfo.inChannels, convInfo.inDepth, convInfo.inHeight, convInfo.inWidth, - convInfo.inChannels, convInfo.outDepth, convInfo.outHeight, convInfo.outWidth, - convInfo.outChannels, convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth, From 97932762eac43e5122ef711b288bcfa40ed66e85 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Sun, 22 Jan 2023 20:53:22 -0800 Subject: [PATCH 7/9] Update headers and build deps remove --- tfjs-backend-wasm/src/cc/BUILD.bazel | 3 --- tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc | 1 + tfjs-backend-wasm/src/cc/pool3d_impl.h | 3 --- 3 files changed, 1 insertion(+), 6 deletions(-) diff --git a/tfjs-backend-wasm/src/cc/BUILD.bazel b/tfjs-backend-wasm/src/cc/BUILD.bazel index aad27fa673c..9cf54ca5ff8 100644 --- a/tfjs-backend-wasm/src/cc/BUILD.bazel +++ b/tfjs-backend-wasm/src/cc/BUILD.bazel @@ -248,9 +248,6 @@ tfjs_cc_library( tfjs_cc_library( name = "pool3d_impl", hdrs = ["pool3d_impl.h"], - deps = [ - ":backend", - ], ) tfjs_cc_library( diff --git a/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc b/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc index 7aedde7c78a..44e35047fe7 100644 --- a/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc +++ b/tfjs-backend-wasm/src/cc/kernels/AvgPool3D.cc @@ -20,6 +20,7 @@ #endif #include +#include #include "tfjs-backend-wasm/src/cc/backend.h" #include "tfjs-backend-wasm/src/cc/pool3d_impl.h" diff --git a/tfjs-backend-wasm/src/cc/pool3d_impl.h b/tfjs-backend-wasm/src/cc/pool3d_impl.h index eac53bf25e7..dc17f6d0fba 100644 --- a/tfjs-backend-wasm/src/cc/pool3d_impl.h +++ b/tfjs-backend-wasm/src/cc/pool3d_impl.h @@ -3,9 +3,6 @@ #endif #include -#include - -#include "tfjs-backend-wasm/src/cc/backend.h" namespace tfjs::wasm { From 54525da43ed220cee522889e9516f4395d0beed8 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Mon, 23 Jan 2023 13:56:16 -0800 Subject: [PATCH 8/9] Add license comment --- tfjs-backend-wasm/src/cc/pool3d_impl.h | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tfjs-backend-wasm/src/cc/pool3d_impl.h b/tfjs-backend-wasm/src/cc/pool3d_impl.h index dc17f6d0fba..d4275970398 100644 --- a/tfjs-backend-wasm/src/cc/pool3d_impl.h +++ b/tfjs-backend-wasm/src/cc/pool3d_impl.h @@ -1,3 +1,20 @@ +/** + * @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 From 7b4721f0df8d01aec6c7adf7ac7e4fcf9eab1e5d Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Mon, 23 Jan 2023 14:20:14 -0800 Subject: [PATCH 9/9] Fix non-negative modulo implementation --- tfjs-backend-wasm/src/cc/pool3d_impl.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-wasm/src/cc/pool3d_impl.h b/tfjs-backend-wasm/src/cc/pool3d_impl.h index d4275970398..143574675c2 100644 --- a/tfjs-backend-wasm/src/cc/pool3d_impl.h +++ b/tfjs-backend-wasm/src/cc/pool3d_impl.h @@ -29,7 +29,7 @@ inline int AddUntilNonNegative(int v, int d) { if (v >= 0) { return v; } - return (v % d + v) % d; + return (v % d + d) % d; } } // namespace