1 /* Copyright 2016 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 #ifndef TENSORFLOW_CORE_KERNELS_GATHER_FUNCTOR_H_
17 #define TENSORFLOW_CORE_KERNELS_GATHER_FUNCTOR_H_
18
19 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
20
21 #include "tensorflow/core/framework/bounds_check.h"
22 #include "tensorflow/core/framework/op_kernel.h"
23 #include "tensorflow/core/framework/tensor_types.h"
24 #include "tensorflow/core/framework/type_traits.h"
25 #include "tensorflow/core/framework/variant.h"
26 #include "tensorflow/core/platform/prefetch.h"
27 #include "tensorflow/core/platform/types.h"
28 #include "tensorflow/core/util/work_sharder.h"
29
30 namespace tensorflow {
31 typedef Eigen::ThreadPoolDevice CPUDevice;
32 typedef Eigen::GpuDevice GPUDevice;
33
34 namespace functor {
35
36 // Helper method to copy using memcpy.
37 template <typename T, typename Index, typename SliceIndex,
38 SliceIndex static_slice_elems>
HandleCopies(OpKernelContext * ctx,typename TTypes<T,3>::ConstTensor params,typename TTypes<Index>::ConstFlat indices,SliceIndex slice_elems,typename TTypes<T,3>::Tensor out)39 SliceIndex HandleCopies(OpKernelContext* ctx,
40 typename TTypes<T, 3>::ConstTensor params,
41 typename TTypes<Index>::ConstFlat indices,
42 SliceIndex slice_elems,
43 typename TTypes<T, 3>::Tensor out) {
44 const SliceIndex indices_size = static_cast<SliceIndex>(indices.dimension(0));
45 const SliceIndex batch_size = static_cast<SliceIndex>(params.dimension(0));
46 const Index limit = static_cast<Index>(params.dimension(1));
47 T* out_base = &out(0, 0, 0);
48 const T* params_base = ¶ms(0, 0, 0);
49 if (static_slice_elems >= 0) {
50 // Give compiler static knowledge of the number of elements/bytes
51 slice_elems = static_slice_elems;
52 }
53 // Compute slice_bytes here so that static knowledge is available
54 const size_t slice_bytes = slice_elems * sizeof(T);
55 auto* worker_threads = ctx->device()->tensorflow_cpu_worker_threads();
56 mutex mu;
57 // Store the value of invalidate index for printing error information, it's a
58 // shared variable.
59 SliceIndex result = -1;
60 auto work = [&](int64 start, int64 end) {
61 SliceIndex batch_idx = static_cast<SliceIndex>(start / indices_size);
62 SliceIndex indices_idx = static_cast<SliceIndex>(start % indices_size);
63 SliceIndex batch_idx_end = static_cast<SliceIndex>(end / indices_size);
64 SliceIndex indices_idx_end = static_cast<SliceIndex>(end % indices_size);
65
66 while ((batch_idx < batch_idx_end) ||
67 (batch_idx == batch_idx_end && indices_idx < indices_idx_end)) {
68 SliceIndex i_next = indices_idx + 1;
69 SliceIndex b_next = batch_idx + 1;
70 if ((batch_idx == batch_idx_end && i_next < indices_idx_end) ||
71 (i_next < indices_size)) {
72 port::prefetch<port::PREFETCH_HINT_T0>(
73 ¶ms(batch_idx, indices(i_next), 0));
74 port::prefetch<port::PREFETCH_HINT_T0>(&out(batch_idx, i_next, 0));
75 b_next = batch_idx;
76 } else if (b_next <= batch_idx_end) {
77 port::prefetch<port::PREFETCH_HINT_T0>(¶ms(b_next, indices(0), 0));
78 port::prefetch<port::PREFETCH_HINT_T0>(&out(b_next, 0, 0));
79 i_next = 0;
80 }
81 const Index index = internal::SubtleMustCopy(indices(indices_idx));
82 if (!FastBoundsCheck(index, limit)) {
83 mutex_lock l(mu);
84 result = indices_idx;
85 return;
86 }
87 // Copy using memcpy if possible, otherwise an Eigen loop
88 // TODO(cwhipkey): avoid linking to framework to get Allocator (to improve
89 // ahead-of-time compilation binary size).
90 if (is_simple_type<T>::value) {
91 // Avoid auto-promotion to Index from SliceIndex by casting.
92 memcpy(
93 out_base + (batch_idx * indices_size + indices_idx) * slice_elems,
94 params_base + (batch_idx * static_cast<SliceIndex>(limit) +
95 static_cast<SliceIndex>(index)) *
96 slice_elems,
97 slice_bytes);
98 } else {
99 // For non-"simple" types (e.g. strings).
100 out.template chip<1>(indices_idx) = params.template chip<1>(index);
101 }
102 indices_idx = i_next;
103 batch_idx = b_next;
104 }
105 };
106
107 Shard(worker_threads->num_threads, worker_threads->workers,
108 batch_size * indices_size, slice_elems * sizeof(T), work);
109 return result;
110 }
111
112 template <typename T, typename Index>
113 struct GatherFunctorCPU {
operatorGatherFunctorCPU114 int64 operator()(OpKernelContext* ctx,
115 typename TTypes<T, 3>::ConstTensor params,
116 typename TTypes<Index>::ConstFlat indices,
117 typename TTypes<T, 3>::Tensor out) {
118 const int64 N = indices.size();
119 const int64 slice_size = out.dimension(2);
120 int64 bad_i;
121
122 bool use_large = (slice_size > std::numeric_limits<int32>::max() ||
123 params.size() > std::numeric_limits<int32>::max() ||
124 N > std::numeric_limits<int32>::max());
125 #define CALL(elems) \
126 do { \
127 if (use_large) { \
128 bad_i = HandleCopies<T, Index, int64, elems>(ctx, params, indices, \
129 slice_size, out); \
130 } else { \
131 const int32 small_slice = static_cast<int32>(slice_size); \
132 bad_i = HandleCopies<T, Index, int32, elems>(ctx, params, indices, \
133 small_slice, out); \
134 } \
135 } while (0)
136
137 if (slice_size == 10)
138 CALL(10);
139 else if (slice_size == 20)
140 CALL(20);
141 else
142 CALL(-1);
143 #undef CALL
144
145 return bad_i;
146 }
147 };
148
149 template <typename Device, typename T, typename Index>
150 struct GatherFunctor {
151 int64 operator()(OpKernelContext* ctx,
152 typename TTypes<T, 3>::ConstTensor params,
153 typename TTypes<Index>::ConstFlat indices,
154 typename TTypes<T, 3>::Tensor out);
155 };
156
157 template <typename T, typename Index>
158 struct GatherFunctor<CPUDevice, T, Index> {
159 int64 operator()(OpKernelContext* ctx,
160 typename TTypes<T, 3>::ConstTensor params,
161 typename TTypes<Index>::ConstFlat indices,
162 typename TTypes<T, 3>::Tensor out) {
163 return GatherFunctorCPU<T, Index>()(ctx, params, indices, out);
164 }
165 };
166
167 template <typename Index>
168 struct GatherFunctor<GPUDevice, Variant, Index> {
169 int64 operator()(OpKernelContext* ctx,
170 typename TTypes<Variant, 3>::ConstTensor params,
171 typename TTypes<Index>::ConstFlat indices,
172 typename TTypes<Variant, 3>::Tensor out) {
173 return GatherFunctorCPU<Variant, Index>()(ctx, params, indices, out);
174 }
175 };
176
177 } // namespace functor
178 } // namespace tensorflow
179
180 #endif // TENSORFLOW_CORE_KERNELS_GATHER_FUNCTOR_H_
181