1 | /* Copyright 2015 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 | #include "tensorflow/core/kernels/reduction_ops_common.h" |
17 | |
18 | namespace tensorflow { |
19 | |
20 | #define REGISTER_CPU_KERNELS(type) \ |
21 | REGISTER_KERNEL_BUILDER( \ |
22 | Name("Min") \ |
23 | .Device(DEVICE_CPU) \ |
24 | .TypeConstraint<type>("T") \ |
25 | .TypeConstraint<int32>("Tidx"), \ |
26 | ReductionOp<CPUDevice, type, int32, \ |
27 | Eigen::internal::MinReducer<type, Eigen::PropagateNaN>>); \ |
28 | REGISTER_KERNEL_BUILDER( \ |
29 | Name("Min") \ |
30 | .Device(DEVICE_CPU) \ |
31 | .TypeConstraint<type>("T") \ |
32 | .TypeConstraint<int64_t>("Tidx"), \ |
33 | ReductionOp<CPUDevice, type, int64, \ |
34 | Eigen::internal::MinReducer<type, Eigen::PropagateNaN>>); |
35 | TF_CALL_REAL_NUMBER_TYPES(REGISTER_CPU_KERNELS); |
36 | #undef REGISTER_CPU_KERNELS |
37 | |
38 | #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
39 | |
40 | #define REGISTER_GPU_KERNELS(type) \ |
41 | REGISTER_KERNEL_BUILDER( \ |
42 | Name("Min") \ |
43 | .Device(DEVICE_GPU) \ |
44 | .TypeConstraint<type>("T") \ |
45 | .TypeConstraint<int32>("Tidx") \ |
46 | .HostMemory("reduction_indices"), \ |
47 | ReductionOp<GPUDevice, type, int32, \ |
48 | Eigen::internal::MinReducer<type, Eigen::PropagateNaN>>); \ |
49 | REGISTER_KERNEL_BUILDER( \ |
50 | Name("Min") \ |
51 | .Device(DEVICE_GPU) \ |
52 | .TypeConstraint<type>("T") \ |
53 | .TypeConstraint<int64_t>("Tidx") \ |
54 | .HostMemory("reduction_indices"), \ |
55 | ReductionOp<GPUDevice, type, int64, \ |
56 | Eigen::internal::MinReducer<type, Eigen::PropagateNaN>>); |
57 | REGISTER_GPU_KERNELS(Eigen::half); |
58 | REGISTER_GPU_KERNELS(float); |
59 | REGISTER_GPU_KERNELS(double); |
60 | |
61 | #undef REGISTER_GPU_KERNELS |
62 | #endif |
63 | |
64 | // A special DEVICE_DEFAULT kernel for int32. |
65 | // TODO(b/25387198): Also enable int32 in device memory. This kernel |
66 | // registration requires all int32 inputs and outputs to be in host memory. |
67 | REGISTER_KERNEL_BUILDER( |
68 | Name("Min" ) |
69 | .Device(DEVICE_DEFAULT) |
70 | .HostMemory("reduction_indices" ) |
71 | .HostMemory("input" ) |
72 | .HostMemory("output" ) |
73 | .TypeConstraint<int32>("T" ) |
74 | .TypeConstraint<int32>("Tidx" ), |
75 | ReductionOp<CPUDevice, int32, int32, Eigen::internal::MinReducer<int32>>); |
76 | REGISTER_KERNEL_BUILDER( |
77 | Name("Min" ) |
78 | .Device(DEVICE_DEFAULT) |
79 | .HostMemory("reduction_indices" ) |
80 | .HostMemory("input" ) |
81 | .HostMemory("output" ) |
82 | .TypeConstraint<int32>("T" ) |
83 | .TypeConstraint<int64_t>("Tidx" ), |
84 | ReductionOp<CPUDevice, int32, int64, Eigen::internal::MinReducer<int32>>); |
85 | |
86 | |
87 | } // namespace tensorflow |
88 | |