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 | #ifndef TENSORFLOW_CORE_UTIL_TENSOR_SLICE_UTIL_H_ |
17 | #define TENSORFLOW_CORE_UTIL_TENSOR_SLICE_UTIL_H_ |
18 | |
19 | #include "tensorflow/core/framework/tensor_shape.h" |
20 | #include "tensorflow/core/framework/tensor_slice.h" |
21 | #include "tensorflow/core/platform/logging.h" |
22 | |
23 | namespace tensorflow { |
24 | |
25 | namespace { |
26 | |
27 | // Some hackery to invoke eigen tensor to copy over tensor slices with variable |
28 | // dimension tensors. |
29 | // TODO(yangke): get rid of that once the variable dimension tensor support is |
30 | // in. |
31 | static const int kTensorSliceMaxRank = 8; |
32 | |
33 | // Create a tensor map with the given shape: we support up to 8 dimensions. If |
34 | // the shape has less than 8 dimensions, we pad the remaining dimension with 1. |
35 | template <typename T> |
36 | Eigen::TensorMap<Eigen::Tensor<T, kTensorSliceMaxRank, Eigen::RowMajor>> |
37 | GetEigenTensorMapFromTensorShape(const TensorShape& shape, T* data) { |
38 | Eigen::DSizes<Eigen::DenseIndex, kTensorSliceMaxRank> dsizes = |
39 | shape.AsEigenDSizesWithPadding<kTensorSliceMaxRank>(); |
40 | Eigen::TensorMap<Eigen::Tensor<T, kTensorSliceMaxRank, Eigen::RowMajor>> eig( |
41 | data, dsizes); |
42 | return eig; |
43 | } |
44 | |
45 | // For everything except string, a standard Eigen cast and assignment works |
46 | template <typename DstT> |
47 | struct CopyThatWorksWithStringPointer { |
48 | template <typename SrcTensor, typename DstTensor, typename Shape> |
49 | static void Copy(const SrcTensor& s, Shape s_start, Shape len, DstTensor& d, |
50 | Shape d_start) { |
51 | d.slice(d_start, len) = s.slice(s_start, len).template cast<DstT>(); |
52 | } |
53 | }; |
54 | |
55 | // Eigen makes it extremely difficult to dereference a tensor of string* into |
56 | // string, so we roll our own loop instead. |
57 | template <> |
58 | struct CopyThatWorksWithStringPointer<tstring> { |
59 | template <typename SrcTensor, typename DstTensor, typename Shape> |
60 | static void Copy(const SrcTensor& s, Shape s_start, Shape len, DstTensor& d, |
61 | Shape d_start) { |
62 | typedef typename SrcTensor::Index Index; |
63 | static_assert(kTensorSliceMaxRank == 8, |
64 | "If kTensorSliceMaxRank changes, modify the loop below." ); |
65 | for (Index i0 = 0; i0 < len[0]; i0++) { |
66 | for (Index i1 = 0; i1 < len[1]; i1++) { |
67 | for (Index i2 = 0; i2 < len[2]; i2++) { |
68 | for (Index i3 = 0; i3 < len[3]; i3++) { |
69 | for (Index i4 = 0; i4 < len[4]; i4++) { |
70 | for (Index i5 = 0; i5 < len[5]; i5++) { |
71 | for (Index i6 = 0; i6 < len[6]; i6++) { |
72 | for (Index i7 = 0; i7 < len[7]; i7++) { |
73 | d(d_start[0] + i0, d_start[1] + i1, d_start[2] + i2, |
74 | d_start[3] + i3, d_start[4] + i4, d_start[5] + i5, |
75 | d_start[6] + i6, d_start[7] + i7) = |
76 | *s(s_start[0] + i0, s_start[1] + i1, s_start[2] + i2, |
77 | s_start[3] + i3, s_start[4] + i4, s_start[5] + i5, |
78 | s_start[6] + i6, s_start[7] + i7); |
79 | } |
80 | } |
81 | } |
82 | } |
83 | } |
84 | } |
85 | } |
86 | } |
87 | } |
88 | }; |
89 | |
90 | // Checkpointing of half is done by storing the raw 16 bits as a signed 32bit |
91 | // integer. To restore the checkpoint we need to do the reverse operation by |
92 | // reinterpreting the integer as a 16 bit float. This prevents us from using |
93 | // the default cast operation. |
94 | template <> |
95 | struct CopyThatWorksWithStringPointer<Eigen::half> { |
96 | template <typename SrcTensor, typename DstTensor, typename Shape> |
97 | static void Copy(const SrcTensor& s, Shape s_start, Shape len, DstTensor& d, |
98 | Shape d_start) { |
99 | typedef typename SrcTensor::Index Index; |
100 | static_assert(kTensorSliceMaxRank == 8, |
101 | "If kTensorSliceMaxRank changes, modify the loop below." ); |
102 | for (Index i0 = 0; i0 < len[0]; i0++) { |
103 | for (Index i1 = 0; i1 < len[1]; i1++) { |
104 | for (Index i2 = 0; i2 < len[2]; i2++) { |
105 | for (Index i3 = 0; i3 < len[3]; i3++) { |
106 | for (Index i4 = 0; i4 < len[4]; i4++) { |
107 | for (Index i5 = 0; i5 < len[5]; i5++) { |
108 | for (Index i6 = 0; i6 < len[6]; i6++) { |
109 | for (Index i7 = 0; i7 < len[7]; i7++) { |
110 | d(d_start[0] + i0, d_start[1] + i1, d_start[2] + i2, |
111 | d_start[3] + i3, d_start[4] + i4, d_start[5] + i5, |
112 | d_start[6] + i6, d_start[7] + i7) = |
113 | Eigen::numext::bit_cast<Eigen::half, uint16_t>( |
114 | s(s_start[0] + i0, s_start[1] + i1, s_start[2] + i2, |
115 | s_start[3] + i3, s_start[4] + i4, s_start[5] + i5, |
116 | s_start[6] + i6, s_start[7] + i7)); |
117 | } |
118 | } |
119 | } |
120 | } |
121 | } |
122 | } |
123 | } |
124 | } |
125 | } |
126 | }; |
127 | |
128 | // Given a tensor described by "shape", two slices "slice_s" and "slice_d", |
129 | // and two pointers "ptr_s" and "ptr_d", where "ptr_s" points to a chunk of |
130 | // memory that stores the data for "slice_s" and "ptr_d" points to a chunk of |
131 | // memory that stores the data for "slice_d". This function copies the data |
132 | // that belongs to the intersection of the two slices from slice_s to |
133 | // slice_d. Uses Tensor cast<DstT>() to convert from SrcT to DstT. Returns true |
134 | // iff the two slices share any intersection (and thus some data is copied). |
135 | // TODO(yangke): figure out if we can make it private. |
136 | template <typename SrcT, typename DstT> |
137 | static bool CopyDataFromTensorSliceToTensorSlice(const TensorShape& shape, |
138 | const TensorSlice& slice_s, |
139 | const TensorSlice& slice_d, |
140 | const SrcT* ptr_s, |
141 | DstT* ptr_d) { |
142 | CHECK_LE(shape.dims(), kTensorSliceMaxRank) |
143 | << "Only tensors of size up to " << kTensorSliceMaxRank |
144 | << " are supported" ; |
145 | // We need to compute the intersection of the two slices. |
146 | TensorSlice inter; |
147 | if (!slice_s.Intersect(slice_d, &inter)) { |
148 | // There is no intersection: returns false. |
149 | return false; |
150 | } else { |
151 | // We need to compute the applied shapes after applying slice_s and |
152 | // slice_d. |
153 | TensorShape shp_s, shp_d; |
154 | Status s; |
155 | s = slice_s.SliceTensorShape(shape, &shp_s); |
156 | if (!s.ok()) { |
157 | LOG(WARNING) << s; |
158 | return false; |
159 | } |
160 | s = slice_d.SliceTensorShape(shape, &shp_d); |
161 | if (!s.ok()) { |
162 | LOG(WARNING) << s; |
163 | return false; |
164 | } |
165 | |
166 | // We need to compute the relative slice of "inter" w.r.t. both slice_s and |
167 | // slice_d. |
168 | TensorSlice rel_s, rel_d; |
169 | slice_s.ComputeRelative(inter, &rel_s); |
170 | slice_d.ComputeRelative(inter, &rel_d); |
171 | |
172 | // Get the eigen tensor maps to the data. |
173 | auto t_s = GetEigenTensorMapFromTensorShape(shp_s, ptr_s); |
174 | auto t_d = GetEigenTensorMapFromTensorShape(shp_d, ptr_d); |
175 | |
176 | Eigen::DSizes<Eigen::DenseIndex, kTensorSliceMaxRank> s_start, s_len, |
177 | d_start, d_len; |
178 | |
179 | rel_s.FillIndicesAndSizes<kTensorSliceMaxRank>(shp_s, &s_start, &s_len); |
180 | rel_d.FillIndicesAndSizes<kTensorSliceMaxRank>(shp_d, &d_start, &d_len); |
181 | CopyThatWorksWithStringPointer<DstT>::Copy(t_s, s_start, s_len, t_d, |
182 | d_start); |
183 | return true; |
184 | } |
185 | } |
186 | |
187 | } // namespace |
188 | |
189 | } // namespace tensorflow |
190 | |
191 | #endif // TENSORFLOW_CORE_UTIL_TENSOR_SLICE_UTIL_H_ |
192 | |