1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifdef INTEL_MKL 17 18 #include "tensorflow/core/kernels/mkl_pooling_ops_common.h" 19 #include <limits> 20 #include <vector> 21 #include "tensorflow/core/common_runtime/device.h" 22 #include "tensorflow/core/framework/common_shape_fns.h" 23 #include "tensorflow/core/kernels/bounds_check.h" 24 25 namespace tensorflow { 26 27 // Initialization for TensorFlow format 28 void MklPoolParameters::Init(OpKernelContext* context, 29 const std::vector<int32>& ksize, 30 const std::vector<int32>& stride, Padding padding, 31 TensorFormat data_format, 32 const TensorShape& tensor_in_shape) { 33 // For maxpooling, tensor_in should have 4 dimensions. 34 OP_REQUIRES(context, tensor_in_shape.dims() == 4, 35 errors::InvalidArgument("tensor_in must be 4-dimensional")); 36 37 depth = GetTensorDim(tensor_in_shape, data_format, 'C'); 38 tensor_in_cols = GetTensorDim(tensor_in_shape, data_format, 'W'); 39 tensor_in_rows = GetTensorDim(tensor_in_shape, data_format, 'H'); 40 tensor_in_batch = GetTensorDim(tensor_in_shape, data_format, 'N'); 41 42 Init(context, ksize, stride, padding, data_format); 43 } 44 45 #ifdef INTEL_MKL_ML 46 // Initialization for MKL format 47 void MklPoolParameters::Init(OpKernelContext* context, 48 const std::vector<int32>& ksize, 49 const std::vector<int32>& stride, Padding padding, 50 TensorFormat data_format, 51 const MklShape* mklInputShape) { 52 // Get the input sizes 53 depth = mklInputShape->GetSizes()[2]; 54 tensor_in_cols = mklInputShape->GetSizes()[0]; 55 tensor_in_rows = mklInputShape->GetSizes()[1]; 56 tensor_in_batch = mklInputShape->GetSizes()[3]; 57 58 Init(context, ksize, stride, padding, data_format); 59 } 60 #else 61 // Initialization for MKL format 62 void MklPoolParameters::Init(OpKernelContext* context, 63 const std::vector<int32>& ksize, 64 const std::vector<int32>& stride, Padding padding, 65 TensorFormat data_format, 66 const MklDnnShape* mklInputShape) { 67 // Get the input sizes 68 depth = mklInputShape->GetDimension('C'); 69 tensor_in_cols = mklInputShape->GetDimension('W'); 70 tensor_in_rows = mklInputShape->GetDimension('H'); 71 tensor_in_batch = mklInputShape->GetDimension('N'); 72 73 Init(context, ksize, stride, padding, data_format); 74 } 75 #endif // INTEL_MKL_ML 76 // Common Initialization for TensorFlow and MKL formats 77 void MklPoolParameters::Init(OpKernelContext* context, 78 const std::vector<int32>& ksize, 79 const std::vector<int32>& stride, Padding padding, 80 TensorFormat data_format) { 81 // Get the data format 82 this->data_format = data_format; 83 84 // Get the output sizes 85 window_rows = GetTensorDim(ksize, data_format, 'H'); 86 window_cols = GetTensorDim(ksize, data_format, 'W'); 87 depth_window = GetTensorDim(ksize, data_format, 'C'); 88 89 // Get the strides 90 row_stride = GetTensorDim(stride, data_format, 'H'); 91 col_stride = GetTensorDim(stride, data_format, 'W'); 92 depth_stride = GetTensorDim(stride, data_format, 'C'); 93 94 // We only support 2D pooling across width/height and depthwise 95 // pooling, not a combination. 96 OP_REQUIRES(context, 97 (depth_window == 1 || (window_rows == 1 && window_cols == 1)), 98 errors::Unimplemented( 99 "MaxPooling supports exactly one of pooling across depth " 100 "or pooling across width/height.")); 101 102 if (depth_window == 1) { // we are pooling in the H and W 103 OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose( 104 tensor_in_rows, window_rows, row_stride, 105 padding, &out_height, &pad_top, &pad_bottom)); 106 107 OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose( 108 tensor_in_cols, window_cols, col_stride, 109 padding, &out_width, &pad_left, &pad_right)); 110 #ifndef INTEL_MKL_ML 111 // TF can work with int64, but mkldnn only supports int32 112 // Fail if the height or width are greater than MAX_INT 113 114 OP_REQUIRES(context, 115 FastBoundsCheck(out_height, std::numeric_limits<int>::max()), 116 errors::InvalidArgument("output height is too large")); 117 118 OP_REQUIRES(context, 119 FastBoundsCheck(out_width, std::numeric_limits<int>::max()), 120 errors::InvalidArgument("output width is too large")); 121 122 #endif 123 out_depth = depth; // output will have the same depth as the input 124 } else { // we are pooling in the depth dimension 125 // Our current version of depthwise max pooling does not support 126 // any padding, and expects the depth_window to equal the depth 127 // stride (no overlapping). 128 OP_REQUIRES(context, depth % depth_window == 0, 129 errors::Unimplemented("Depthwise max pooling requires the" 130 " depth window to evenly divide the" 131 " input depth")); 132 OP_REQUIRES(context, depth_stride == depth_window, 133 errors::Unimplemented("Depthwise max pooling requires the" 134 " depth window to equal the depth" 135 " stride")); 136 137 // The current version of depthwise max is only implemented on CPU. 138 OP_REQUIRES(context, 139 (DeviceType(static_cast<Device*>(context->device()) 140 ->attributes() 141 .device_type()) == DeviceType(DEVICE_CPU)), 142 errors::Unimplemented("Depthwise max pooling is currently " 143 "only implemented for CPU devices.")); 144 145 out_depth = depth / depth_window; 146 } 147 } 148 149 // Transfers the right parameters for pooling to the op parameters 150 // Updates context->status if there is an invalid input. 151 void ExtractMklOpParams(OpKernelContext* context, TensorFormat data_format, 152 const MklPoolParameters& params, 153 MklPoolingOpParams* mkl_params) { 154 mkl_params->in_sizes[0] = params.tensor_in_cols; 155 mkl_params->in_sizes[1] = params.tensor_in_rows; 156 mkl_params->in_sizes[2] = params.depth; 157 mkl_params->in_sizes[3] = params.tensor_in_batch; 158 159 GetStridesFromSizes(data_format, mkl_params->in_strides, 160 mkl_params->in_sizes); 161 162 mkl_params->out_sizes[0] = params.out_width; 163 mkl_params->out_sizes[1] = params.out_height; 164 mkl_params->out_sizes[2] = params.depth; 165 mkl_params->out_sizes[3] = params.tensor_in_batch; 166 167 GetStridesFromSizes(data_format, mkl_params->out_strides, 168 mkl_params->out_sizes); 169 170 mkl_params->in_offset[0] = -params.pad_left; 171 mkl_params->in_offset[1] = -params.pad_top; 172 mkl_params->in_offset[2] = -params.pad_right; 173 mkl_params->in_offset[3] = -params.pad_bottom; 174 175 mkl_params->kernel_stride[0] = params.col_stride; 176 mkl_params->kernel_stride[1] = params.row_stride; 177 178 mkl_params->kernel_size[0] = params.window_cols; 179 mkl_params->kernel_size[1] = params.window_rows; 180 } 181 } // namespace tensorflow 182 #endif // INTEL_MKL 183
