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 | // See docs in ../ops/array_ops.cc. |
17 | #include "tensorflow/core/kernels/identity_op.h" |
18 | |
19 | #include "tensorflow/core/framework/op_kernel.h" |
20 | #include "tensorflow/core/framework/register_types.h" |
21 | #include "tensorflow/core/framework/tensor.h" |
22 | #include "tensorflow/core/framework/types.h" |
23 | |
24 | namespace tensorflow { |
25 | |
26 | REGISTER_KERNEL_BUILDER(Name("Identity" ).Device(DEVICE_CPU), IdentityOp); |
27 | REGISTER_KERNEL_BUILDER(Name("Identity" ).Device(DEVICE_TPU_SYSTEM), IdentityOp); |
28 | |
29 | // StopGradient does the same thing as Identity, but has a different |
30 | // gradient registered. |
31 | REGISTER_KERNEL_BUILDER(Name("StopGradient" ).Device(DEVICE_CPU), IdentityOp); |
32 | // PreventGradient does the same thing as Identity, but has a NO |
33 | // gradient registered. |
34 | REGISTER_KERNEL_BUILDER(Name("PreventGradient" ).Device(DEVICE_CPU), IdentityOp); |
35 | |
36 | // PlaceholderWithDefault does the same thing as Identity, but has a |
37 | // different shape function (and constant value function) registered. |
38 | REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault" ).Device(DEVICE_CPU), |
39 | IdentityOp); |
40 | |
41 | // Used executing op-by-op to copy constants to the current device without |
42 | // serializing tensors as TensorProtos, after a host tensor has been |
43 | // created. Same behavior as Identity, but no gradient and potentially relaxed |
44 | // copy semantics. |
45 | REGISTER_KERNEL_BUILDER(Name("_EagerConst" ).Device(DEVICE_CPU), IdentityOp); |
46 | |
47 | REGISTER_KERNEL_BUILDER(Name("RefIdentity" ).Device(DEVICE_CPU), IdentityOp); |
48 | |
49 | // Identity op for gradients debugging in TensorFlow Debugger (hidden op in |
50 | // Python). |
51 | REGISTER_KERNEL_BUILDER(Name("DebugGradientIdentity" ).Device(DEVICE_CPU), |
52 | IdentityOp); |
53 | REGISTER_KERNEL_BUILDER(Name("DebugGradientRefIdentity" ).Device(DEVICE_CPU), |
54 | IdentityOp); |
55 | |
56 | #define REGISTER_GPU_KERNEL(type) \ |
57 | REGISTER_KERNEL_BUILDER( \ |
58 | Name("Identity").Device(DEVICE_GPU).TypeConstraint<type>("T"), \ |
59 | IdentityOp); \ |
60 | REGISTER_KERNEL_BUILDER( \ |
61 | Name("PreventGradient").Device(DEVICE_GPU).TypeConstraint<type>("T"), \ |
62 | IdentityOp); \ |
63 | REGISTER_KERNEL_BUILDER( \ |
64 | Name("RefIdentity").Device(DEVICE_GPU).TypeConstraint<type>("T"), \ |
65 | IdentityOp); \ |
66 | REGISTER_KERNEL_BUILDER( \ |
67 | Name("StopGradient").Device(DEVICE_GPU).TypeConstraint<type>("T"), \ |
68 | IdentityOp); \ |
69 | REGISTER_KERNEL_BUILDER(Name("DebugGradientIdentity") \ |
70 | .Device(DEVICE_GPU) \ |
71 | .TypeConstraint<type>("T"), \ |
72 | IdentityOp); \ |
73 | REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault") \ |
74 | .Device(DEVICE_GPU) \ |
75 | .TypeConstraint<type>("dtype"), \ |
76 | IdentityOp) \ |
77 | REGISTER_KERNEL_BUILDER( \ |
78 | Name("_EagerConst").Device(DEVICE_GPU).TypeConstraint<type>("T"), \ |
79 | IdentityOp) |
80 | |
81 | TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); |
82 | REGISTER_GPU_KERNEL(Variant); |
83 | REGISTER_GPU_KERNEL(bool); |
84 | |
85 | #undef REGISTER_GPU_KERNEL |
86 | |
87 | #define REGISTER_DEFAULT_KERNEL(type) \ |
88 | REGISTER_KERNEL_BUILDER( \ |
89 | Name("Identity").Device(DEVICE_DEFAULT).TypeConstraint<type>("T"), \ |
90 | IdentityOp); \ |
91 | REGISTER_KERNEL_BUILDER(Name("PreventGradient") \ |
92 | .Device(DEVICE_DEFAULT) \ |
93 | .TypeConstraint<type>("T"), \ |
94 | IdentityOp); \ |
95 | REGISTER_KERNEL_BUILDER( \ |
96 | Name("RefIdentity").Device(DEVICE_DEFAULT).TypeConstraint<type>("T"), \ |
97 | IdentityOp); \ |
98 | REGISTER_KERNEL_BUILDER( \ |
99 | Name("StopGradient").Device(DEVICE_DEFAULT).TypeConstraint<type>("T"), \ |
100 | IdentityOp); \ |
101 | REGISTER_KERNEL_BUILDER(Name("DebugGradientIdentity") \ |
102 | .Device(DEVICE_DEFAULT) \ |
103 | .TypeConstraint<type>("T"), \ |
104 | IdentityOp); \ |
105 | REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault") \ |
106 | .Device(DEVICE_DEFAULT) \ |
107 | .TypeConstraint<type>("dtype"), \ |
108 | IdentityOp) \ |
109 | REGISTER_KERNEL_BUILDER( \ |
110 | Name("_EagerConst").Device(DEVICE_DEFAULT).TypeConstraint<type>("T"), \ |
111 | IdentityOp) |
112 | |
113 | TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_DEFAULT_KERNEL); |
114 | REGISTER_DEFAULT_KERNEL(Variant); |
115 | REGISTER_DEFAULT_KERNEL(bool); |
116 | |
117 | #undef REGISTER_DEFAULT_KERNEL |
118 | |
119 | #if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ |
120 | (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) |
121 | // A special GPU kernel for int32 and bool. |
122 | // TODO(b/25387198): Also enable int32 in device memory. This kernel |
123 | // registration requires all int32 inputs and outputs to be in host memory. |
124 | #define REGISTER_GPU_HOST_KERNEL(type) \ |
125 | REGISTER_KERNEL_BUILDER(Name("Identity") \ |
126 | .Device(DEVICE_GPU) \ |
127 | .HostMemory("input") \ |
128 | .HostMemory("output") \ |
129 | .TypeConstraint<type>("T"), \ |
130 | IdentityOp); \ |
131 | REGISTER_KERNEL_BUILDER(Name("RefIdentity") \ |
132 | .Device(DEVICE_GPU) \ |
133 | .HostMemory("input") \ |
134 | .HostMemory("output") \ |
135 | .TypeConstraint<type>("T"), \ |
136 | IdentityOp); \ |
137 | REGISTER_KERNEL_BUILDER(Name("StopGradient") \ |
138 | .Device(DEVICE_GPU) \ |
139 | .HostMemory("input") \ |
140 | .HostMemory("output") \ |
141 | .TypeConstraint<type>("T"), \ |
142 | IdentityOp); \ |
143 | REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault") \ |
144 | .Device(DEVICE_GPU) \ |
145 | .HostMemory("input") \ |
146 | .HostMemory("output") \ |
147 | .TypeConstraint<type>("dtype"), \ |
148 | IdentityOp) \ |
149 | REGISTER_KERNEL_BUILDER(Name("_EagerConst") \ |
150 | .Device(DEVICE_GPU) \ |
151 | .HostMemory("input") \ |
152 | .HostMemory("output") \ |
153 | .TypeConstraint<type>("T"), \ |
154 | IdentityOp); |
155 | |
156 | REGISTER_GPU_HOST_KERNEL(int32); |
157 | REGISTER_GPU_HOST_KERNEL(tstring); |
158 | REGISTER_GPU_HOST_KERNEL(ResourceHandle); |
159 | |
160 | #undef REGISTER_GPU_HOST_KERNEL |
161 | |
162 | #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
163 | |
164 | #define REGISTER_DEFAULT_HOST_KERNEL(type) \ |
165 | REGISTER_KERNEL_BUILDER(Name("Identity") \ |
166 | .Device(DEVICE_DEFAULT) \ |
167 | .HostMemory("input") \ |
168 | .HostMemory("output") \ |
169 | .TypeConstraint<type>("T"), \ |
170 | IdentityOp); \ |
171 | REGISTER_KERNEL_BUILDER(Name("RefIdentity") \ |
172 | .Device(DEVICE_DEFAULT) \ |
173 | .HostMemory("input") \ |
174 | .HostMemory("output") \ |
175 | .TypeConstraint<type>("T"), \ |
176 | IdentityOp); \ |
177 | REGISTER_KERNEL_BUILDER(Name("StopGradient") \ |
178 | .Device(DEVICE_DEFAULT) \ |
179 | .HostMemory("input") \ |
180 | .HostMemory("output") \ |
181 | .TypeConstraint<type>("T"), \ |
182 | IdentityOp); \ |
183 | REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault") \ |
184 | .Device(DEVICE_DEFAULT) \ |
185 | .HostMemory("input") \ |
186 | .HostMemory("output") \ |
187 | .TypeConstraint<type>("dtype"), \ |
188 | IdentityOp) \ |
189 | REGISTER_KERNEL_BUILDER(Name("_EagerConst") \ |
190 | .Device(DEVICE_DEFAULT) \ |
191 | .HostMemory("input") \ |
192 | .HostMemory("output") \ |
193 | .TypeConstraint<type>("T"), \ |
194 | IdentityOp) |
195 | |
196 | REGISTER_DEFAULT_HOST_KERNEL(int32); |
197 | REGISTER_DEFAULT_HOST_KERNEL(tstring); |
198 | REGISTER_DEFAULT_HOST_KERNEL(ResourceHandle); |
199 | |
200 | #undef REGISTER_DEFAULT_HOST_KERNEL |
201 | |
202 | } // namespace tensorflow |
203 | |