1 /* Copyright 2017 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 #include "tensorflow/compiler/xla/index_util.h"
17
18 #include <algorithm>
19 #include <string>
20
21 #include "absl/strings/str_join.h"
22 #include "tensorflow/compiler/xla/shape_util.h"
23 #include "tensorflow/compiler/xla/types.h"
24 #include "tensorflow/core/platform/logging.h"
25
26 namespace xla {
27
MultidimensionalIndexToLinearIndex(const Shape & shape,absl::Span<const int64> multi_index)28 /* static */ int64 IndexUtil::MultidimensionalIndexToLinearIndex(
29 const Shape& shape, absl::Span<const int64> multi_index) {
30 DCHECK_EQ(shape.dimensions_size(), multi_index.size());
31
32 for (size_t i = 0; i < multi_index.size(); ++i) {
33 DCHECK_GE(multi_index[i], 0);
34 DCHECK_LT(multi_index[i], shape.dimensions(i))
35 << "indexing beyond extent in dimension " << i << ":"
36 << "\n\tindex: " << absl::StrJoin(multi_index, ",")
37 << "\n\tshape: " << ShapeUtil::HumanString(shape);
38 }
39
40 // Let the array be sized like so for dimensions i from 0 to n-1:
41 //
42 // [D{n-1} x D{n-2} x .. x D{0}]
43 //
44 // Let the order of the dimensions in the minor_to_major field in
45 // Layout be:
46 //
47 // L(0), L(1), ... , L(n-1)
48 //
49 // where L(0) is the most-minor dimension and L(n-1) the most-major. The
50 // multidimensional index:
51 //
52 // [I{0}, I{1}, ... , I{n-1}]
53 //
54 // then corresponds to the following linear index:
55 //
56 // linear_index =
57 // ((( ... + I{L(2)}) * D{L(1)} + I{L(1)}) * D{L(0)} + I{L(0)}
58 //
59 // or equivalently:
60 //
61 // linear_index =
62 // I{L(n-1)} * (D{L(n-2)} * D{L(n-3)} * D{L(n-4)} * .... D{L(0)}) +
63 // I{L(n-2)} * (D{L(n-3)} * D{L(n-4)} * .... D{L(0)}) +
64 // I{L(n-3)} * (D{L(n-4)} * .... D{L(0)}) +
65 // ... +
66 // I{L(2)} * (D{L(1)} * D{L(0)}) +
67 // I{L(1)} * D{L(0)} +
68 // I{L(0)}
69 //
70 // We compute the linear index value by accumulating the terms above from
71 // I{L(0)} up to I{L(n-1)}. Scale accumulates the product term D{L(0}} *
72 // D{L(1)} * ...
73
74 // Scale factor holding the growing product of D{L(i)} terms.
75 int64 scale = 1;
76 int64 linear_index = 0;
77 bool first = true;
78 for (auto dimension : LayoutUtil::MinorToMajor(shape)) {
79 if (first) {
80 // Avoid two multiplies on the first loop iteration
81 linear_index = multi_index[dimension];
82 scale = shape.dimensions(dimension);
83 first = false;
84 } else {
85 linear_index += scale * multi_index[dimension];
86 scale *= shape.dimensions(dimension);
87 }
88 }
89 return linear_index;
90 }
91
LinearIndexToMultidimensionalIndex(const Shape & shape,int64 linear_index)92 /* static */ std::vector<int64> IndexUtil::LinearIndexToMultidimensionalIndex(
93 const Shape& shape, int64 linear_index) {
94 DCHECK_GE(linear_index, 0);
95 DCHECK_LT(linear_index, ShapeUtil::ElementsIn(shape));
96
97 // The following formula computes each element of the multidimensional index
98 // (See comments in MultidimensionalIndexToLinearIndex for notation):
99 //
100 // I{L(0)} = linear_index % D{L(0)}
101 // I{L(1)} = (linear_index / D{L(0)}) % D{L(1)}
102 // I{L(2)} = (linear_index / (D{L(0)} * D{L(1)})) % D{L(2)}
103 // ...
104 std::vector<int64> multi_index(shape.dimensions_size());
105
106 // Accumulated product D{L(0)} * D{L(1)} * ...
107 int64 divisor = 1;
108 for (auto dimension : LayoutUtil::MinorToMajor(shape)) {
109 multi_index[dimension] =
110 (linear_index / divisor) % shape.dimensions(dimension);
111 divisor *= shape.dimensions(dimension);
112 }
113 return multi_index;
114 }
115
BumpIndices(const Shape & shape,absl::Span<int64> indices)116 /* static */ bool IndexUtil::BumpIndices(const Shape& shape,
117 absl::Span<int64> indices) {
118 for (int64 dimno = indices.size() - 1; dimno >= 0; --dimno) {
119 int64 limit = shape.dimensions(dimno);
120 if (indices[dimno] + 1 < limit) {
121 indices[dimno]++;
122 // Whenever an index of a dimension is increased, it means that all
123 // following dimensions have maxed out, so they must go to 0.
124 std::fill(indices.begin() + dimno + 1, indices.end(), 0);
125 return true;
126 }
127 }
128 return false;
129 }
130
GetDimensionStride(const Shape & shape,int64 dimension)131 /* static */ int64 IndexUtil::GetDimensionStride(const Shape& shape,
132 int64 dimension) {
133 int64 stride = 1;
134 for (auto dim : LayoutUtil::MinorToMajor(shape)) {
135 if (dim == dimension) {
136 break;
137 }
138 stride *= shape.dimensions()[dim];
139 }
140 return stride;
141 }
142
IndexInBounds(const Shape & shape,absl::Span<const int64> index)143 /* static */ bool IndexUtil::IndexInBounds(const Shape& shape,
144 absl::Span<const int64> index) {
145 int64 rank = shape.rank();
146 const int64 index_size = index.size();
147 if (rank != index_size) {
148 return false;
149 }
150 for (int64 d = 0; d < rank; ++d) {
151 if (index[d] >= shape.dimensions(d)) {
152 return false;
153 }
154 }
155 return true;
156 }
157
CompareIndices(absl::Span<const int64> lhs,absl::Span<const int64> rhs)158 /* static */ int IndexUtil::CompareIndices(absl::Span<const int64> lhs,
159 absl::Span<const int64> rhs) {
160 int64 rank = lhs.size();
161 const int64 rhs_rank = rhs.size();
162 CHECK_EQ(rhs_rank, rank);
163 for (int64 dim = 0; dim < rank; ++dim) {
164 if (lhs[dim] < rhs[dim]) {
165 return -1;
166 } else if (lhs[dim] > rhs[dim]) {
167 return 1;
168 }
169 }
170 return 0;
171 }
172
173 } // namespace xla
174