• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 = &params(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             &params(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>(&params(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