diff options
Diffstat (limited to 'tensorflow/core/kernels/reduction_ops_prod.cc')
| -rw-r--r-- | tensorflow/core/kernels/reduction_ops_prod.cc | 68 |
1 files changed, 44 insertions, 24 deletions
diff --git a/tensorflow/core/kernels/reduction_ops_prod.cc b/tensorflow/core/kernels/reduction_ops_prod.cc index 33f6ae6bae..e9b23df746 100644 --- a/tensorflow/core/kernels/reduction_ops_prod.cc +++ b/tensorflow/core/kernels/reduction_ops_prod.cc @@ -17,26 +17,39 @@ limitations under the License. namespace tensorflow { -#define REGISTER_CPU_KERNELS(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("Prod") \ - .Device(DEVICE_CPU) \ - .TypeConstraint<type>("T") \ - .TypeConstraint<int32>("Tidx"), \ - ReductionOp<CPUDevice, type, Eigen::internal::ProdReducer<type>>); +#define REGISTER_CPU_KERNELS(type) \ + REGISTER_KERNEL_BUILDER(Name("Prod") \ + .Device(DEVICE_CPU) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int32>("Tidx"), \ + ReductionOp<CPUDevice, type, int32, \ + Eigen::internal::ProdReducer<type>>); \ + REGISTER_KERNEL_BUILDER(Name("Prod") \ + .Device(DEVICE_CPU) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int64>("Tidx"), \ + ReductionOp<CPUDevice, type, int64, \ + Eigen::internal::ProdReducer<type>>); TF_CALL_NUMBER_TYPES(REGISTER_CPU_KERNELS); #undef REGISTER_CPU_KERNELS #if GOOGLE_CUDA -#define REGISTER_GPU_KERNELS(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("Prod") \ - .Device(DEVICE_GPU) \ - .TypeConstraint<type>("T") \ - .TypeConstraint<int32>("Tidx") \ - .HostMemory("reduction_indices"), \ - ReductionOp<GPUDevice, type, Eigen::internal::ProdReducer<type>>); +#define REGISTER_GPU_KERNELS(type) \ + REGISTER_KERNEL_BUILDER(Name("Prod") \ + .Device(DEVICE_GPU) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int32>("Tidx") \ + .HostMemory("reduction_indices"), \ + ReductionOp<GPUDevice, type, int32, \ + Eigen::internal::ProdReducer<type>>); \ + REGISTER_KERNEL_BUILDER(Name("Prod") \ + .Device(DEVICE_GPU) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int64>("Tidx") \ + .HostMemory("reduction_indices"), \ + ReductionOp<GPUDevice, type, int64, \ + Eigen::internal::ProdReducer<type>>); TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS); TF_CALL_int32(REGISTER_GPU_KERNELS); TF_CALL_complex64(REGISTER_GPU_KERNELS); @@ -46,18 +59,25 @@ TF_CALL_complex128(REGISTER_GPU_KERNELS); #endif #ifdef TENSORFLOW_USE_SYCL -#define REGISTER_SYCL_KERNELS(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("Prod") \ - .Device(DEVICE_SYCL) \ - .TypeConstraint<type>("T") \ - .TypeConstraint<int32>("Tidx") \ - .HostMemory("reduction_indices"), \ - ReductionOp<SYCLDevice, type, Eigen::internal::ProdReducer<type>>); +#define REGISTER_SYCL_KERNELS(type) \ + REGISTER_KERNEL_BUILDER(Name("Prod") \ + .Device(DEVICE_SYCL) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int32>("Tidx") \ + .HostMemory("reduction_indices"), \ + ReductionOp<SYCLDevice, type, int32, \ + Eigen::internal::ProdReducer<type>>); \ + REGISTER_KERNEL_BUILDER(Name("Prod") \ + .Device(DEVICE_SYCL) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int64>("Tidx") \ + .HostMemory("reduction_indices"), \ + ReductionOp<SYCLDevice, type, int64, \ + Eigen::internal::ProdReducer<type>>); REGISTER_SYCL_KERNELS(int32); REGISTER_SYCL_KERNELS(float); REGISTER_SYCL_KERNELS(double); #undef REGISTER_SYCL_KERNELS -#endif // TENSORFLOW_USE_SYCL +#endif // TENSORFLOW_USE_SYCL } // namespace tensorflow |