1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
24namespace tensorflow {
25
26REGISTER_KERNEL_BUILDER(Name("Identity").Device(DEVICE_CPU), IdentityOp);
27REGISTER_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.
31REGISTER_KERNEL_BUILDER(Name("StopGradient").Device(DEVICE_CPU), IdentityOp);
32// PreventGradient does the same thing as Identity, but has a NO
33// gradient registered.
34REGISTER_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.
38REGISTER_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.
45REGISTER_KERNEL_BUILDER(Name("_EagerConst").Device(DEVICE_CPU), IdentityOp);
46
47REGISTER_KERNEL_BUILDER(Name("RefIdentity").Device(DEVICE_CPU), IdentityOp);
48
49// Identity op for gradients debugging in TensorFlow Debugger (hidden op in
50// Python).
51REGISTER_KERNEL_BUILDER(Name("DebugGradientIdentity").Device(DEVICE_CPU),
52 IdentityOp);
53REGISTER_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
81TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL);
82REGISTER_GPU_KERNEL(Variant);
83REGISTER_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
113TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_DEFAULT_KERNEL);
114REGISTER_DEFAULT_KERNEL(Variant);
115REGISTER_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
156REGISTER_GPU_HOST_KERNEL(int32);
157REGISTER_GPU_HOST_KERNEL(tstring);
158REGISTER_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
196REGISTER_DEFAULT_HOST_KERNEL(int32);
197REGISTER_DEFAULT_HOST_KERNEL(tstring);
198REGISTER_DEFAULT_HOST_KERNEL(ResourceHandle);
199
200#undef REGISTER_DEFAULT_HOST_KERNEL
201
202} // namespace tensorflow
203