| 151 | |
| 152 | template <typename T> |
| 153 | void DnnPoolingOp<T>::Compute(OpKernelContext* context, |
| 154 | se::dnn::PoolingMode pooling_mode, |
| 155 | const std::vector<int32>& size, |
| 156 | const std::vector<int32>& stride, Padding padding, |
| 157 | TensorFormat data_format, const Tensor& tensor_in, |
| 158 | const TensorShape& tensor_out_shape, |
| 159 | bool propagate_nans) { |
| 160 | Tensor* tensor_out = nullptr; |
| 161 | OP_REQUIRES_OK(context, |
| 162 | context->allocate_output(0, tensor_out_shape, &tensor_out)); |
| 163 | if (tensor_in.shape().num_elements() == 0) { |
| 164 | return; |
| 165 | } |
| 166 | |
| 167 | PoolParameters params{context, size, stride, |
| 168 | padding, data_format, tensor_in.shape()}; |
| 169 | if (!context->status().ok()) { |
| 170 | return; |
| 171 | } |
| 172 | |
| 173 | int batch_size = params.tensor_in_batch; |
| 174 | int depth = params.depth; |
| 175 | #if CUDNN_VERSION < 7300 |
| 176 | /// Earlier versions do not support NHWC format, so we need to convert it |
| 177 | /// to NCHW before calling cudnn. We need to get rid of this once it is done |
| 178 | Tensor transformed_input; |
| 179 | if (data_format == FORMAT_NHWC) { |
| 180 | OP_REQUIRES_OK(context, context->allocate_temp( |
| 181 | DataTypeToEnum<T>::value, |
| 182 | ShapeFromFormat(FORMAT_NCHW, tensor_in.shape(), |
| 183 | data_format), |
| 184 | &transformed_input)); |
| 185 | functor::NHWCToNCHW<GPUDevice, T, 4>()(context->eigen_device<Device>(), |
| 186 | tensor_in.tensor<T, 4>(), |
| 187 | transformed_input.tensor<T, 4>()); |
| 188 | } else { |
| 189 | transformed_input = tensor_in; |
| 190 | } |
| 191 | Tensor transformed_output; |
| 192 | if (data_format == FORMAT_NHWC) { |
| 193 | OP_REQUIRES_OK(context, context->allocate_temp( |
| 194 | DataTypeToEnum<T>::value, |
| 195 | ShapeFromFormat(FORMAT_NCHW, tensor_out_shape, |
| 196 | data_format), |
| 197 | &transformed_output)); |
| 198 | } else { |
| 199 | transformed_output = *tensor_out; |
| 200 | } |
| 201 | se::dnn::DataLayout data_layout = se::dnn::DataLayout::kBatchDepthYX; |
| 202 | #else |
| 203 | auto& transformed_input = tensor_in; |
| 204 | auto& transformed_output = *tensor_out; |
| 205 | se::dnn::DataLayout data_layout; |
| 206 | switch (data_format) { |
| 207 | case FORMAT_NHWC: |
| 208 | data_layout = se::dnn::DataLayout::kBatchYXDepth; |
| 209 | break; |
| 210 | case FORMAT_NCHW: |
nothing calls this directly
no test coverage detected