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_FRAMEWORK_OPS_UTIL_H_ |
17 | #define TENSORFLOW_CORE_FRAMEWORK_OPS_UTIL_H_ |
18 | |
19 | // This file contains utilities for various operations. |
20 | |
21 | #include <array> |
22 | |
23 | #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" |
24 | #include "tensorflow/core/framework/tensor_shape.h" |
25 | #include "tensorflow/core/framework/tensor_types.h" |
26 | #include "tensorflow/core/lib/core/status.h" |
27 | #include "tensorflow/core/util/padding.h" |
28 | |
29 | namespace tensorflow { |
30 | |
31 | // Calculates broadcast starting index and size. For SAME padding, addition |
32 | // padding could be applied to right, left, top and bottom. Depending on the |
33 | // current index, input size, kernel size, stride, padding size, the starting |
34 | // index and size for broadcast for that dimension are different from the |
35 | // current index and kernel size. |
36 | // This is mainly used by gradient algorithms for pooling operations. |
37 | Status GetBroadcastSize(const int index, const int in_size, const int ksize, |
38 | const int stride, const int pad_size, int* bindex, |
39 | int* bsize); |
40 | |
41 | // Converts Brain's Padding to Eigen's PaddingType. |
42 | Eigen::PaddingType BrainPadding2EigenPadding(Padding padding); |
43 | |
44 | // Given a shape 's' of a tensor of type T. Returns true iff the |
45 | // number of bytes occupied by each dim 0 (i.e., &tensor(i + 1, ...) - |
46 | // &tensor(i, ...)) is multiple of EIGEN_MAX_ALIGN_BYTES. |
47 | template <typename T> |
48 | bool IsInnerDimsSizeAligned(const TensorShape& s) { |
49 | if (s.dims() == 0) return false; |
50 | const int64_t dim0_size = s.dim_size(0); |
51 | if (dim0_size == 0) return false; |
52 | #if EIGEN_MAX_ALIGN_BYTES == 0 |
53 | return true; |
54 | #else |
55 | const int64_t bytes_per_dim0 = (s.num_elements() / dim0_size) * sizeof(T); |
56 | return bytes_per_dim0 % EIGEN_MAX_ALIGN_BYTES == 0; |
57 | #endif |
58 | } |
59 | |
60 | // Given a shape 's' of a tensor of type T and the `start` and `end` index of a |
61 | // dim 0 slice, returns true iff slice is aligned with respect to original |
62 | // tensor. Here aligned implies the address is a multiple of |
63 | // EIGEN_MAX_ALIGN_BYTES. |
64 | template <typename T> |
65 | bool IsDim0SliceAligned(const TensorShape& s, int64_t start, |
66 | int64_t end_or_size) { |
67 | if (s.dims() == 1) { |
68 | #if EIGEN_MAX_ALIGN_BYTES == 0 |
69 | return true; |
70 | #else |
71 | bool start_aligned = (start * sizeof(T)) % EIGEN_MAX_ALIGN_BYTES == 0; |
72 | // End is aligned if either the explicit end index is passed and is a |
73 | // a multiple of EIGEN_MAX_ALIGN_BYTES, or the start index is aligned and |
74 | // the size is aligned. So for convenience we can either pass start and |
75 | // index, or start and size. |
76 | bool end_aligned = (end_or_size * sizeof(T)) % EIGEN_MAX_ALIGN_BYTES == 0; |
77 | return start_aligned && end_aligned; |
78 | #endif |
79 | } else { |
80 | return IsInnerDimsSizeAligned<T>(s); |
81 | } |
82 | } |
83 | |
84 | // Returns <suffix> sanitized to have only [a-zA-Z0-9-_]. |
85 | std::string SanitizeThreadSuffix(std::string suffix); |
86 | |
87 | // Helper to compute 'strides' given a tensor 'shape'. I.e., |
88 | // strides[i] = prod(shape.dim_size[(i+1):]) |
89 | template <typename T> |
90 | gtl::InlinedVector<T, 8> ComputeStride(const TensorShape& shape) { |
91 | const int ndims = shape.dims(); |
92 | gtl::InlinedVector<T, 8> strides(ndims); |
93 | T stride = 1; |
94 | for (int i = ndims - 1; i >= 0; --i) { |
95 | strides[i] = stride; |
96 | stride *= static_cast<T>(shape.dim_size(i)); |
97 | } |
98 | return strides; |
99 | } |
100 | |
101 | // Helper to compute 'strides' given an Eigen TensorDimensions |
102 | template <typename T, typename EigenDimensions> |
103 | gtl::InlinedVector<T, 8> ComputeEigenStrides(const EigenDimensions& shape) { |
104 | const int ndims = shape.rank(); |
105 | gtl::InlinedVector<T, 8> strides(ndims); |
106 | T stride = 1; |
107 | for (int i = ndims - 1; i >= 0; --i) { |
108 | strides[i] = stride; |
109 | stride *= static_cast<T>(shape[i]); |
110 | } |
111 | return strides; |
112 | } |
113 | |
114 | } // namespace tensorflow |
115 | |
116 | #endif // TENSORFLOW_CORE_FRAMEWORK_OPS_UTIL_H_ |
117 | |