1 /* Copyright 2016 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 // See docs in ../ops/nn_ops.cc. 17 18 #define EIGEN_USE_THREADS 19 20 #include "tensorflow/core/kernels/mirror_pad_op.h" 21 #include <string> 22 23 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 24 25 #include "tensorflow/core/framework/op.h" 26 #include "tensorflow/core/framework/op_kernel.h" 27 #include "tensorflow/core/framework/register_types.h" 28 #include "tensorflow/core/framework/tensor.h" 29 #include "tensorflow/core/framework/tensor_shape.h" 30 #include "tensorflow/core/framework/tensor_types.h" 31 #include "tensorflow/core/framework/types.h" 32 #include "tensorflow/core/platform/logging.h" 33 #include "tensorflow/core/platform/types.h" 34 #include "tensorflow/core/util/mirror_pad_mode.h" 35 36 namespace tensorflow { 37 38 template <typename Device, typename T, typename Tpaddings> 39 class MirrorPadOp : public OpKernel { 40 public: MirrorPadOp(OpKernelConstruction * context)41 explicit MirrorPadOp(OpKernelConstruction* context) : OpKernel(context) { 42 MirrorPadMode mode; 43 OP_REQUIRES_OK(context, context->GetAttr("mode", &mode)); 44 45 switch (mode) { 46 case MirrorPadMode::SYMMETRIC: { 47 offset_ = 0; 48 break; 49 } 50 case MirrorPadMode::REFLECT: { 51 offset_ = 1; 52 break; 53 } 54 default: 55 OP_REQUIRES(context, false, 56 errors::InvalidArgument( 57 "mode must be either REFLECT or SYMMETRIC.")); 58 } 59 } 60 61 ~MirrorPadOp() override = default; 62 Compute(OpKernelContext * context)63 void Compute(OpKernelContext* context) override { 64 const Tensor& in0 = context->input(0); 65 const Tensor& in1 = context->input(1); 66 const int dims = in0.dims(); 67 constexpr int kMinDims = 0; 68 constexpr int kMaxDims = 5; 69 OP_REQUIRES(context, kMinDims <= dims && dims <= kMaxDims, 70 errors::Unimplemented("inputs rank not in [", kMinDims, ",", 71 kMaxDims, "]: ", dims)); 72 OP_REQUIRES( 73 context, 74 TensorShapeUtils::IsMatrix(in1.shape()) && in1.dim_size(1) == 2, 75 errors::InvalidArgument("paddings must be a matrix with 2 columns: ", 76 in1.shape().DebugString())); 77 OP_REQUIRES( 78 context, dims == in1.dim_size(0), 79 errors::InvalidArgument( 80 "The first dimension of paddings must be the rank of inputs", 81 in1.shape().DebugString(), ", ", in0.shape().DebugString())); 82 83 // Compute the shape of the output tensor, and allocate it. 84 TensorShape output_shape; 85 typename TTypes<Tpaddings>::ConstMatrix paddings = in1.matrix<Tpaddings>(); 86 for (int d = 0; d < dims; ++d) { 87 const Tpaddings before = paddings(d, 0); // Pad before existing elements. 88 const Tpaddings after = paddings(d, 1); // Pad after existing elements. 89 OP_REQUIRES(context, before >= 0 && after >= 0, 90 errors::InvalidArgument( 91 "paddings must be non-negative: ", before, " ", after)); 92 if (offset_ == 0) { // SYMMETRIC mode. 93 OP_REQUIRES(context, 94 before <= in0.dim_size(d) && after <= in0.dim_size(d), 95 errors::InvalidArgument("paddings must be no greater " 96 "than the dimension size: ", 97 before, ", ", after, 98 " greater than ", in0.dim_size(d))); 99 } else if (offset_ == 1) { // REFLECT mode. 100 OP_REQUIRES( 101 context, before < in0.dim_size(d) && after < in0.dim_size(d), 102 errors::InvalidArgument("paddings must be less than" 103 " the dimension size: ", 104 before, ", ", after, " not less than ", 105 in0.dim_size(d))); 106 } 107 108 output_shape.AddDim(before + in0.dim_size(d) + after); 109 } 110 111 if (output_shape.num_elements() == in0.NumElements()) { 112 // When num_elements == 0, shape may have changed. 113 Tensor out; 114 CHECK(out.CopyFrom(in0, output_shape)); 115 context->set_output(0, out); 116 return; 117 } 118 119 Tensor* output = nullptr; 120 OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); 121 122 #define MIRROR_PAD_CASE(i) \ 123 case i: { \ 124 functor::MirrorPad<Device, T, Tpaddings, i>()( \ 125 context->eigen_device<Device>(), To32Bit(output->tensor<T, i>()), \ 126 To32Bit(in0.tensor<T, i>()), paddings, offset_); \ 127 break; \ 128 } 129 130 // Invoke the dims-specific implementation. 131 switch (dims) { 132 MIRROR_PAD_CASE(1) 133 MIRROR_PAD_CASE(2) 134 MIRROR_PAD_CASE(3) 135 MIRROR_PAD_CASE(4) 136 MIRROR_PAD_CASE(5) 137 default: 138 OP_REQUIRES(context, false, 139 errors::InvalidArgument("Unsupported rank: ", 140 in0.shape().DebugString())); 141 } 142 #undef MIRROR_PAD_CASE 143 } 144 145 private: 146 int offset_; 147 }; 148 149 using CpuDevice = Eigen::ThreadPoolDevice; 150 using GpuDevice = Eigen::GpuDevice; 151 152 namespace functor { 153 // Forward declarations of the functor specializations defined in the sharded 154 // files. 155 #define DECLARE_CPU_SPEC(T, Tpaddings, i) \ 156 template <> \ 157 void MirrorPad<CpuDevice, T, Tpaddings, i>::operator()( \ 158 const CpuDevice&, typename TTypes<T, i, int32>::Tensor, \ 159 typename TTypes<T, i, int32>::ConstTensor, \ 160 TTypes<Tpaddings>::ConstMatrix, int); \ 161 extern template struct MirrorPad<CpuDevice, T, Tpaddings, i>; 162 163 #define DECLARE_CPU_SPECS(T) \ 164 DECLARE_CPU_SPEC(T, int32, 1); \ 165 DECLARE_CPU_SPEC(T, int32, 2); \ 166 DECLARE_CPU_SPEC(T, int32, 3); \ 167 DECLARE_CPU_SPEC(T, int32, 4); \ 168 DECLARE_CPU_SPEC(T, int32, 5); \ 169 DECLARE_CPU_SPEC(T, int64, 1); \ 170 DECLARE_CPU_SPEC(T, int64, 2); \ 171 DECLARE_CPU_SPEC(T, int64, 3); \ 172 DECLARE_CPU_SPEC(T, int64, 4); \ 173 DECLARE_CPU_SPEC(T, int64, 5); 174 175 TF_CALL_POD_TYPES(DECLARE_CPU_SPECS); 176 TF_CALL_string(DECLARE_CPU_SPECS); 177 178 #undef DECLARE_CPU_SPEC 179 #undef DECLARE_CPU_SPECS 180 } // namespace functor 181 182 #define REGISTER_KERNEL(type) \ 183 REGISTER_KERNEL_BUILDER(Name("MirrorPad") \ 184 .Device(DEVICE_CPU) \ 185 .TypeConstraint<type>("T") \ 186 .TypeConstraint<int32>("Tpaddings") \ 187 .HostMemory("paddings"), \ 188 MirrorPadOp<CpuDevice, type, int32>); \ 189 REGISTER_KERNEL_BUILDER(Name("MirrorPad") \ 190 .Device(DEVICE_CPU) \ 191 .TypeConstraint<type>("T") \ 192 .TypeConstraint<int64>("Tpaddings") \ 193 .HostMemory("paddings"), \ 194 MirrorPadOp<CpuDevice, type, int64>); 195 196 // Note that we do register for bool type, but not in the gradient op. 197 TF_CALL_POD_TYPES(REGISTER_KERNEL); 198 TF_CALL_string(REGISTER_KERNEL); 199 #undef REGISTER_KERNEL 200 201 #if GOOGLE_CUDA 202 namespace functor { 203 // Forward declarations of the functor specializations for GPU. 204 #define DECLARE_GPU_SPEC(T, Tpaddings, i) \ 205 template <> \ 206 void MirrorPad<GpuDevice, T, Tpaddings, i>::operator()( \ 207 const GpuDevice&, typename TTypes<T, i, int32>::Tensor, \ 208 typename TTypes<T, i, int32>::ConstTensor, \ 209 TTypes<Tpaddings>::ConstMatrix, int); \ 210 extern template struct MirrorPad<GpuDevice, T, Tpaddings, i>; 211 212 #define DECLARE_GPU_SPECS(T) \ 213 DECLARE_GPU_SPEC(T, int32, 1); \ 214 DECLARE_GPU_SPEC(T, int32, 2); \ 215 DECLARE_GPU_SPEC(T, int32, 3); \ 216 DECLARE_GPU_SPEC(T, int32, 4); \ 217 DECLARE_GPU_SPEC(T, int32, 5); \ 218 DECLARE_GPU_SPEC(T, int64, 1); \ 219 DECLARE_GPU_SPEC(T, int64, 2); \ 220 DECLARE_GPU_SPEC(T, int64, 3); \ 221 DECLARE_GPU_SPEC(T, int64, 4); \ 222 DECLARE_GPU_SPEC(T, int64, 5); 223 224 TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPECS); 225 #undef DECLARE_GPU_SPECS 226 #undef DECLARE_GPU_SPEC 227 } // namespace functor 228 229 // Registration of the GPU implementations. 230 #define REGISTER_GPU_KERNEL(T) \ 231 REGISTER_KERNEL_BUILDER(Name("MirrorPad") \ 232 .Device(DEVICE_GPU) \ 233 .TypeConstraint<T>("T") \ 234 .TypeConstraint<int32>("Tpaddings") \ 235 .HostMemory("paddings"), \ 236 MirrorPadOp<GpuDevice, T, int32>); \ 237 REGISTER_KERNEL_BUILDER(Name("MirrorPad") \ 238 .Device(DEVICE_GPU) \ 239 .TypeConstraint<T>("T") \ 240 .TypeConstraint<int64>("Tpaddings") \ 241 .HostMemory("paddings"), \ 242 MirrorPadOp<GpuDevice, T, int64>); 243 244 TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNEL); 245 #undef REGISTER_GPU_KERNEL 246 #endif // GOOGLE_CUDA 247 248 // Gradient op. 249 template <typename Device, typename T, typename Tpaddings> 250 class MirrorPadGradOp : public OpKernel { 251 public: MirrorPadGradOp(OpKernelConstruction * context)252 explicit MirrorPadGradOp(OpKernelConstruction* context) : OpKernel(context) { 253 MirrorPadMode mode; 254 OP_REQUIRES_OK(context, context->GetAttr("mode", &mode)); 255 256 switch (mode) { 257 case MirrorPadMode::SYMMETRIC: { 258 offset_ = 0; 259 break; 260 } 261 case MirrorPadMode::REFLECT: { 262 offset_ = 1; 263 break; 264 } 265 default: 266 OP_REQUIRES(context, false, 267 errors::InvalidArgument( 268 "mode must be either REFLECT or SYMMETRIC.")); 269 } 270 } 271 272 ~MirrorPadGradOp() override = default; 273 Compute(OpKernelContext * context)274 void Compute(OpKernelContext* context) override { 275 const Tensor& in0 = context->input(0); 276 const Tensor& in1 = context->input(1); 277 const int dims = in0.dims(); 278 constexpr int kMinDims = 0; 279 constexpr int kMaxDims = 5; 280 OP_REQUIRES(context, kMinDims <= dims && dims <= kMaxDims, 281 errors::Unimplemented("inputs rank not in [", kMinDims, ",", 282 kMaxDims, "]: ", dims)); 283 OP_REQUIRES( 284 context, 285 TensorShapeUtils::IsMatrix(in1.shape()) && in1.dim_size(1) == 2, 286 errors::InvalidArgument("paddings must be a matrix with 2 columns: ", 287 in1.shape().DebugString())); 288 OP_REQUIRES( 289 context, dims == in1.dim_size(0), 290 errors::InvalidArgument( 291 "The first dimension of paddings must be the rank of inputs", 292 in1.shape().DebugString(), " ", in0.shape().DebugString())); 293 294 // Compute the shape of the output tensor, and allocate it. 295 TensorShape output_shape; 296 typename TTypes<Tpaddings>::ConstMatrix paddings = in1.matrix<Tpaddings>(); 297 for (int d = 0; d < dims; ++d) { 298 const Tpaddings before = paddings(d, 0); // Pad before existing elements. 299 const Tpaddings after = paddings(d, 1); // Pad after existing elements. 300 OP_REQUIRES(context, before >= 0 && after >= 0, 301 errors::InvalidArgument( 302 "Paddings must be non-negative: ", before, ", ", after)); 303 304 const int64 out_size = in0.dim_size(d) - (before + after); 305 if (offset_ == 0) { // SYMMETRIC mode. 306 OP_REQUIRES(context, before <= out_size && after <= out_size, 307 errors::InvalidArgument("paddings must be no greater " 308 "than the output dimension size: ", 309 before, ", ", after, 310 " greater than ", out_size)); 311 } else if (offset_ == 1) { // REFLECT mode. 312 OP_REQUIRES(context, before < out_size && after < out_size, 313 errors::InvalidArgument("paddings must be less than" 314 " the output dimension size: ", 315 before, ", ", after, 316 " not less than ", out_size)); 317 } 318 output_shape.AddDim(out_size); 319 } 320 321 if (output_shape == in0.shape()) { 322 context->set_output(0, in0); 323 return; 324 } 325 326 Tensor scratch; 327 OP_REQUIRES_OK(context, context->allocate_temp(DataTypeToEnum<T>::value, 328 in0.shape(), &scratch)); 329 330 Tensor* output = nullptr; 331 OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); 332 333 #define MIRROR_PAD_GRAD_CASE(k) \ 334 case k: { \ 335 functor::MirrorPadGrad<Device, T, Tpaddings, k>()( \ 336 context->eigen_device<Device>(), To32Bit(output->tensor<T, k>()), \ 337 To32Bit(in0.tensor<T, k>()), paddings, offset_, \ 338 To32Bit(scratch.tensor<T, k>())); \ 339 break; \ 340 } 341 342 // Invoke the dims-specific implementation. 343 switch (dims) { 344 MIRROR_PAD_GRAD_CASE(1); 345 MIRROR_PAD_GRAD_CASE(2); 346 MIRROR_PAD_GRAD_CASE(3); 347 MIRROR_PAD_GRAD_CASE(4); 348 MIRROR_PAD_GRAD_CASE(5); 349 default: 350 OP_REQUIRES(context, false, 351 errors::InvalidArgument("Unsupported rank: ", 352 in0.shape().DebugString())); 353 } 354 #undef MIRROR_PAD_GRAD_CASE 355 } 356 357 private: 358 int offset_; 359 }; 360 361 namespace functor { 362 // Forward declarations of the functor specializations defined in the sharded 363 // files. 364 #define DECLARE_CPU_SPEC(T, Tpaddings, k) \ 365 template <> \ 366 void MirrorPadGrad<CpuDevice, T, Tpaddings, k>::operator()( \ 367 const CpuDevice&, typename TTypes<T, k, int32>::Tensor, \ 368 typename TTypes<T, k, int32>::ConstTensor, \ 369 TTypes<Tpaddings>::ConstMatrix, int, \ 370 typename TTypes<T, k, int32>::Tensor); \ 371 extern template struct MirrorPadGrad<CpuDevice, T, Tpaddings, k>; 372 373 #define DECLARE_CPU_SPECS(T) \ 374 DECLARE_CPU_SPEC(T, int32, 1); \ 375 DECLARE_CPU_SPEC(T, int32, 2); \ 376 DECLARE_CPU_SPEC(T, int32, 3); \ 377 DECLARE_CPU_SPEC(T, int32, 4); \ 378 DECLARE_CPU_SPEC(T, int32, 5); \ 379 DECLARE_CPU_SPEC(T, int64, 1); \ 380 DECLARE_CPU_SPEC(T, int64, 2); \ 381 DECLARE_CPU_SPEC(T, int64, 3); \ 382 DECLARE_CPU_SPEC(T, int64, 4); \ 383 DECLARE_CPU_SPEC(T, int64, 5); 384 385 TF_CALL_NUMBER_TYPES(DECLARE_CPU_SPECS); 386 #undef DECLARE_CPU_SPECS 387 #undef DECLARE_CPU_SPEC 388 } // namespace functor 389 390 #define REGISTER_KERNEL(type) \ 391 REGISTER_KERNEL_BUILDER(Name("MirrorPadGrad") \ 392 .Device(DEVICE_CPU) \ 393 .TypeConstraint<type>("T") \ 394 .TypeConstraint<int32>("Tpaddings") \ 395 .HostMemory("paddings"), \ 396 MirrorPadGradOp<CpuDevice, type, int32>); \ 397 REGISTER_KERNEL_BUILDER(Name("MirrorPadGrad") \ 398 .Device(DEVICE_CPU) \ 399 .TypeConstraint<type>("T") \ 400 .TypeConstraint<int64>("Tpaddings") \ 401 .HostMemory("paddings"), \ 402 MirrorPadGradOp<CpuDevice, type, int64>); 403 404 TF_CALL_NUMBER_TYPES(REGISTER_KERNEL); 405 #undef REGISTER_KERNEL 406 407 #if GOOGLE_CUDA 408 namespace functor { 409 // Forward declarations of the functor specializations for GPU. 410 #define DECLARE_GPU_SPEC(T, Tpaddings, k) \ 411 template <> \ 412 void MirrorPadGrad<GpuDevice, T, Tpaddings, k>::operator()( \ 413 const GpuDevice&, typename TTypes<T, k, int32>::Tensor, \ 414 typename TTypes<T, k, int32>::ConstTensor, \ 415 TTypes<Tpaddings>::ConstMatrix, int, \ 416 typename TTypes<T, k, int32>::Tensor); \ 417 extern template struct MirrorPadGrad<GpuDevice, T, Tpaddings, k>; 418 419 #define DECLARE_GPU_SPECS(T) \ 420 DECLARE_GPU_SPEC(T, int32, 1); \ 421 DECLARE_GPU_SPEC(T, int32, 2); \ 422 DECLARE_GPU_SPEC(T, int32, 3); \ 423 DECLARE_GPU_SPEC(T, int32, 4); \ 424 DECLARE_GPU_SPEC(T, int32, 5); \ 425 DECLARE_GPU_SPEC(T, int64, 1); \ 426 DECLARE_GPU_SPEC(T, int64, 2); \ 427 DECLARE_GPU_SPEC(T, int64, 3); \ 428 DECLARE_GPU_SPEC(T, int64, 4); \ 429 DECLARE_GPU_SPEC(T, int64, 5); 430 431 TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPECS); 432 #undef DECLARE_GPU_SPECS 433 #undef DECLARE_GPU_SPEC 434 } // namespace functor 435 436 // Registration of the GPU implementations. 437 #define REGISTER_GPU_KERNEL(T) \ 438 REGISTER_KERNEL_BUILDER(Name("MirrorPadGrad") \ 439 .Device(DEVICE_GPU) \ 440 .TypeConstraint<T>("T") \ 441 .TypeConstraint<int32>("Tpaddings") \ 442 .HostMemory("paddings"), \ 443 MirrorPadGradOp<GpuDevice, T, int32>); \ 444 REGISTER_KERNEL_BUILDER(Name("MirrorPadGrad") \ 445 .Device(DEVICE_GPU) \ 446 .TypeConstraint<T>("T") \ 447 .TypeConstraint<int64>("Tpaddings") \ 448 .HostMemory("paddings"), \ 449 MirrorPadGradOp<GpuDevice, T, int64>); 450 451 TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNEL); 452 #undef REGISTER_GPU_KERNEL 453 #endif // GOOGLE_CUDA 454 455 } // namespace tensorflow 456