diff options
Diffstat (limited to 'tensorflow/core/kernels/reduction_ops_max.cc')
| -rw-r--r-- | tensorflow/core/kernels/reduction_ops_max.cc | 90 |
1 files changed, 64 insertions, 26 deletions
diff --git a/tensorflow/core/kernels/reduction_ops_max.cc b/tensorflow/core/kernels/reduction_ops_max.cc index 4ca5c11a48..9cf953f4bf 100644 --- a/tensorflow/core/kernels/reduction_ops_max.cc +++ b/tensorflow/core/kernels/reduction_ops_max.cc @@ -17,26 +17,39 @@ limitations under the License. namespace tensorflow { -#define REGISTER_CPU_KERNELS(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("Max") \ - .Device(DEVICE_CPU) \ - .TypeConstraint<type>("T") \ - .TypeConstraint<int32>("Tidx"), \ - ReductionOp<CPUDevice, type, Eigen::internal::MaxReducer<type>>); +#define REGISTER_CPU_KERNELS(type) \ + REGISTER_KERNEL_BUILDER( \ + Name("Max") \ + .Device(DEVICE_CPU) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int32>("Tidx"), \ + ReductionOp<CPUDevice, type, int32, Eigen::internal::MaxReducer<type>>); \ + REGISTER_KERNEL_BUILDER( \ + Name("Max") \ + .Device(DEVICE_CPU) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int64>("Tidx"), \ + ReductionOp<CPUDevice, type, int64, Eigen::internal::MaxReducer<type>>); TF_CALL_REAL_NUMBER_TYPES(REGISTER_CPU_KERNELS); #undef REGISTER_CPU_KERNELS #if GOOGLE_CUDA -#define REGISTER_GPU_KERNELS(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("Max") \ - .Device(DEVICE_GPU) \ - .TypeConstraint<type>("T") \ - .TypeConstraint<int32>("Tidx") \ - .HostMemory("reduction_indices"), \ - ReductionOp<GPUDevice, type, Eigen::internal::MaxReducer<type>>); +#define REGISTER_GPU_KERNELS(type) \ + REGISTER_KERNEL_BUILDER( \ + Name("Max") \ + .Device(DEVICE_GPU) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int32>("Tidx") \ + .HostMemory("reduction_indices"), \ + ReductionOp<GPUDevice, type, int32, Eigen::internal::MaxReducer<type>>); \ + REGISTER_KERNEL_BUILDER( \ + Name("Max") \ + .Device(DEVICE_GPU) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int64>("Tidx") \ + .HostMemory("reduction_indices"), \ + ReductionOp<GPUDevice, type, int64, Eigen::internal::MaxReducer<type>>); REGISTER_GPU_KERNELS(float); REGISTER_GPU_KERNELS(double); REGISTER_GPU_KERNELS(int64); @@ -52,21 +65,37 @@ REGISTER_KERNEL_BUILDER( .HostMemory("output") .TypeConstraint<int32>("T") .TypeConstraint<int32>("Tidx"), - ReductionOp<CPUDevice, int32, Eigen::internal::MaxReducer<int32>>); + ReductionOp<CPUDevice, int32, int32, Eigen::internal::MaxReducer<int32>>); +REGISTER_KERNEL_BUILDER( + Name("Max") + .Device(DEVICE_GPU) + .HostMemory("reduction_indices") + .HostMemory("input") + .HostMemory("output") + .TypeConstraint<int32>("T") + .TypeConstraint<int64>("Tidx"), + ReductionOp<CPUDevice, int32, int64, Eigen::internal::MaxReducer<int32>>); #undef REGISTER_GPU_KERNELS #endif #ifdef TENSORFLOW_USE_SYCL -#define REGISTER_SYCL_KERNELS(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("Max") \ - .Device(DEVICE_SYCL) \ - .TypeConstraint<type>("T") \ - .TypeConstraint<int32>("Tidx") \ - .HostMemory("reduction_indices"), \ - ReductionOp<SYCLDevice, type, Eigen::internal::MaxReducer<type>>); +#define REGISTER_SYCL_KERNELS(type) \ + REGISTER_KERNEL_BUILDER(Name("Max") \ + .Device(DEVICE_SYCL) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int32>("Tidx") \ + .HostMemory("reduction_indices"), \ + ReductionOp<SYCLDevice, type, int32, \ + Eigen::internal::MaxReducer<type>>); \ + REGISTER_KERNEL_BUILDER(Name("Max") \ + .Device(DEVICE_SYCL) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int64>("Tidx") \ + .HostMemory("reduction_indices"), \ + ReductionOp<SYCLDevice, type, int64, \ + Eigen::internal::MaxReducer<type>>); REGISTER_SYCL_KERNELS(float); REGISTER_SYCL_KERNELS(double); @@ -78,8 +107,17 @@ REGISTER_KERNEL_BUILDER( .HostMemory("output") .TypeConstraint<int32>("T") .TypeConstraint<int32>("Tidx"), - ReductionOp<CPUDevice, int32, Eigen::internal::MaxReducer<int32>>); + ReductionOp<CPUDevice, int32, int32, Eigen::internal::MaxReducer<int32>>); +REGISTER_KERNEL_BUILDER( + Name("Max") + .Device(DEVICE_SYCL) + .HostMemory("reduction_indices") + .HostMemory("input") + .HostMemory("output") + .TypeConstraint<int32>("T") + .TypeConstraint<int64>("Tidx"), + ReductionOp<CPUDevice, int32, int64, Eigen::internal::MaxReducer<int32>>); #undef REGISTER_SYCL_KERNELS -#endif // TENSORFLOW_USE_SYCL +#endif // TENSORFLOW_USE_SYCL } // namespace tensorflow |