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_t> multi_index)28 /* static */ int64_t IndexUtil::MultidimensionalIndexToLinearIndex(
29 const Shape& shape, absl::Span<const int64_t> 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_t scale = 1;
76 int64_t 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_t linear_index)92 /* static */ std::vector<int64_t> IndexUtil::LinearIndexToMultidimensionalIndex(
93 const Shape& shape, int64_t 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_t> multi_index(shape.dimensions_size());
105
106 // Accumulated product D{L(0)} * D{L(1)} * ...
107 int64_t 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_t> indices)116 /* static */ bool IndexUtil::BumpIndices(const Shape& shape,
117 absl::Span<int64_t> indices) {
118 for (int64_t dimno = indices.size() - 1; dimno >= 0; --dimno) {
119 int64_t 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_t dimension)131 /* static */ int64_t IndexUtil::GetDimensionStride(const Shape& shape,
132 int64_t dimension) {
133 int64_t 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_t> index)143 /* static */ bool IndexUtil::IndexInBounds(const Shape& shape,
144 absl::Span<const int64_t> index) {
145 int64_t rank = shape.rank();
146 const int64_t index_size = index.size();
147 if (rank != index_size) {
148 return false;
149 }
150 for (int64_t 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_t> lhs,absl::Span<const int64_t> rhs)158 /* static */ int IndexUtil::CompareIndices(absl::Span<const int64_t> lhs,
159 absl::Span<const int64_t> rhs) {
160 int64_t rank = lhs.size();
161 const int64_t rhs_rank = rhs.size();
162 CHECK_EQ(rhs_rank, rank);
163 for (int64_t 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