1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#ifndef TENSORFLOW_CORE_KERNELS_AVGPOOLING_OP_H_
17#define TENSORFLOW_CORE_KERNELS_AVGPOOLING_OP_H_
18// Functor definition for AvgPoolingOp, must be compilable by nvcc.
19
20#include "tensorflow/core/framework/tensor_types.h"
21#include "tensorflow/core/kernels/eigen_pooling.h"
22#include "tensorflow/core/platform/types.h"
23
24namespace tensorflow {
25namespace functor {
26
27template <typename Device, typename T>
28struct SpatialAvgPooling {
29 void operator()(const Device& d, typename TTypes<T, 4>::Tensor output,
30 typename TTypes<T, 4>::ConstTensor input, int window_rows,
31 int window_cols, int row_stride, int col_stride,
32 const Eigen::PaddingType& padding) {
33 MaybeWith32BitIndexing<Device>(
34 [&](auto output32, auto input32) {
35 // Because we swap the layout, we swap the row/cols as well.
36 output32.swap_layout().device(d) = Eigen::SpatialAvgPooling(
37 input32.swap_layout(), window_cols, window_rows, col_stride,
38 row_stride, padding);
39 },
40 output, input);
41 }
42};
43
44} // namespace functor
45
46typedef Eigen::GpuDevice GPUDevice;
47
48// Launch a custom GPU kernels from Yanqing for the avgpooling backward
49// operation that works NHWC data formats. Arguments:
50// top_diff: backprop to the output of the pooling layer
51// num: number of input batches
52// height: input height
53// width: input width
54// channels: number of input channels
55// pooled_height: the height of the output to the pooling layer
56// pooled_width: the width of the output to the pooling layer
57// kernel_h: the height of the pooling kernel
58// kernel_w: the width of the pooling kernel
59// stride_h: the height of the vertical stride
60// stride_w: the width of the horizontal stride
61// pad_t: padding size to the top side
62// pad_l: padding size to the left side
63// bottom_diff: backprop to the input of the pooling layer.
64template <typename T>
65bool RunAvePoolBackwardNHWC(const T* const top_diff, const int num,
66 const int height, const int width,
67 const int channels, const int pooled_height,
68 const int pooled_width, const int kernel_h,
69 const int kernel_w, const int stride_h,
70 const int stride_w, const int pad_t,
71 const int pad_l, T* const bottom_diff,
72 const GPUDevice& d);
73
74} // namespace tensorflow
75
76#endif // TENSORFLOW_CORE_KERNELS_AVGPOOLING_OP_H_
77