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 <initializer_list>
19
20 #include "tensorflow/compiler/xla/shape_util.h"
21 #include "tensorflow/compiler/xla/test.h"
22 #include "tensorflow/compiler/xla/xla_data.pb.h"
23
24 namespace xla {
25 namespace {
26
SetMinorToMajorLayout(Shape * shape,std::vector<int64> dimensions)27 void SetMinorToMajorLayout(Shape* shape, std::vector<int64> dimensions) {
28 shape->mutable_layout()->clear_minor_to_major();
29 for (auto dimension : dimensions) {
30 shape->mutable_layout()->add_minor_to_major(dimension);
31 }
32 }
33
TEST(IndexUtilTest,VectorIndexing)34 TEST(IndexUtilTest, VectorIndexing) {
35 // Vectors are trivially laid out and the linear index should always be the
36 // same as the "multidimensional" index.
37 Shape vector_shape = ShapeUtil::MakeShape(F32, {100});
38 EXPECT_EQ(42,
39 IndexUtil::MultidimensionalIndexToLinearIndex(vector_shape, {42}));
40 std::vector<int64> multi_index =
41 IndexUtil::LinearIndexToMultidimensionalIndex(vector_shape, 42);
42 EXPECT_EQ(1, multi_index.size());
43 EXPECT_EQ(42, multi_index[0]);
44 }
45
TEST(IndexUtilTest,MatrixIndexingRowMajor)46 TEST(IndexUtilTest, MatrixIndexingRowMajor) {
47 // Set layout to [0, 1]. That is, row major.
48 Shape matrix_shape_01 = ShapeUtil::MakeShape(F32, {10, 20});
49 SetMinorToMajorLayout(&matrix_shape_01, {0, 1});
50
51 // If index is {a, b} then linear index should be: a + b * 10
52 EXPECT_EQ(0, IndexUtil::MultidimensionalIndexToLinearIndex(matrix_shape_01,
53 {0, 0}));
54 EXPECT_EQ(199, IndexUtil::MultidimensionalIndexToLinearIndex(matrix_shape_01,
55 {9, 19}));
56 EXPECT_EQ(53, IndexUtil::MultidimensionalIndexToLinearIndex(matrix_shape_01,
57 {3, 5}));
58 EXPECT_EQ(std::vector<int64>({3, 5}),
59 IndexUtil::LinearIndexToMultidimensionalIndex(matrix_shape_01, 53));
60 }
61
TEST(IndexUtilTest,MatrixIndexingColumnMajor)62 TEST(IndexUtilTest, MatrixIndexingColumnMajor) {
63 // Set layout to [1, 0]. That is, column major.
64 Shape matrix_shape_10 = ShapeUtil::MakeShape(F32, {10, 20});
65 SetMinorToMajorLayout(&matrix_shape_10, {1, 0});
66
67 // If index is {a, b} then linear index should be: a * 20 + b
68 EXPECT_EQ(0, IndexUtil::MultidimensionalIndexToLinearIndex(matrix_shape_10,
69 {0, 0}));
70 EXPECT_EQ(199, IndexUtil::MultidimensionalIndexToLinearIndex(matrix_shape_10,
71 {9, 19}));
72 EXPECT_EQ(65, IndexUtil::MultidimensionalIndexToLinearIndex(matrix_shape_10,
73 {3, 5}));
74 EXPECT_EQ(std::vector<int64>({3, 5}),
75 IndexUtil::LinearIndexToMultidimensionalIndex(matrix_shape_10, 65));
76 }
77
TEST(IndexUtilTest,ThreeDArrayIndexing210)78 TEST(IndexUtilTest, ThreeDArrayIndexing210) {
79 // Set layout to [2, 1, 0]. That is, column major.
80 Shape shape_210 = ShapeUtil::MakeShape(F32, {10, 20, 30});
81 SetMinorToMajorLayout(&shape_210, {2, 1, 0});
82
83 // If index is {a, b, c} then linear index should be:
84 // a * 20 * 30 + b * 30 + c
85 EXPECT_EQ(1957, IndexUtil::MultidimensionalIndexToLinearIndex(shape_210,
86 {3, 5, 7}));
87 EXPECT_EQ(5277, IndexUtil::MultidimensionalIndexToLinearIndex(shape_210,
88 {8, 15, 27}));
89 }
90
TEST(IndexUtilTest,ThreeDArrayIndexing120)91 TEST(IndexUtilTest, ThreeDArrayIndexing120) {
92 // Set layout to [1, 2, 0]
93 Shape shape_120 = ShapeUtil::MakeShape(F32, {10, 20, 30});
94 SetMinorToMajorLayout(&shape_120, {1, 2, 0});
95
96 // If index is {a, b, c} then linear index should be:
97 // a * 20 * 30 + b + c * 20
98 EXPECT_EQ(1945, IndexUtil::MultidimensionalIndexToLinearIndex(shape_120,
99 {3, 5, 7}));
100 EXPECT_EQ(5355, IndexUtil::MultidimensionalIndexToLinearIndex(shape_120,
101 {8, 15, 27}));
102 }
103
TEST(IndexUtilTest,FourDArrayIndexing3210)104 TEST(IndexUtilTest, FourDArrayIndexing3210) {
105 // Set layout to [3, 2, 1,0]. That is, column major.
106 Shape shape_3210 = ShapeUtil::MakeShape(F32, {10, 20, 30, 40});
107 SetMinorToMajorLayout(&shape_3210, {3, 2, 1, 0});
108
109 // If index is {a, b, c, d} then linear index should be:
110 // a * 20 * 30 * 40 + b * 30 * 40 + c * 40 + d
111 EXPECT_EQ(78289, IndexUtil::MultidimensionalIndexToLinearIndex(shape_3210,
112 {3, 5, 7, 9}));
113 EXPECT_EQ(211113, IndexUtil::MultidimensionalIndexToLinearIndex(
114 shape_3210, {8, 15, 27, 33}));
115 }
116
TEST(IndexUtilTest,LinearToMultiToLinear)117 TEST(IndexUtilTest, LinearToMultiToLinear) {
118 // Verify that converting a linear index to a multidimensional index and back
119 // always returns the same value for different crazy shapes. Shape has
120 // 1440000000 elements. Inputs are randomly-ish selected.
121 std::vector<int64> linear_indexes = {0, 1439999999, 1145567336,
122 43883404, 617295214, 1117613654};
123
124 std::vector<std::vector<int64>> minor_to_major_orders;
125 minor_to_major_orders.push_back({6, 5, 4, 3, 2, 1, 0});
126 minor_to_major_orders.push_back({0, 1, 2, 3, 4, 5, 6});
127 minor_to_major_orders.push_back({4, 5, 1, 2, 6, 0, 3});
128
129 for (auto minor_to_major_order : minor_to_major_orders) {
130 Shape shape = ShapeUtil::MakeShape(F32, {10, 20, 30, 40, 30, 20, 10});
131 SetMinorToMajorLayout(&shape, minor_to_major_order);
132 for (auto linear_index : linear_indexes) {
133 std::vector<int64> multi_index =
134 IndexUtil::LinearIndexToMultidimensionalIndex(shape, linear_index);
135 EXPECT_EQ(linear_index, IndexUtil::MultidimensionalIndexToLinearIndex(
136 shape, multi_index));
137 }
138 }
139 }
140
TEST(IndexUtilTest,BumpIndices2x2)141 TEST(IndexUtilTest, BumpIndices2x2) {
142 auto shape = ShapeUtil::MakeShape(S32, {2, 2});
143 std::vector<int64> indices = {0, 0};
144 EXPECT_TRUE(IndexUtil::BumpIndices(shape, absl::MakeSpan(indices)));
145 EXPECT_THAT(indices, ::testing::ElementsAre(0, 1));
146 EXPECT_TRUE(IndexUtil::BumpIndices(shape, absl::MakeSpan(indices)));
147 EXPECT_THAT(indices, ::testing::ElementsAre(1, 0));
148 EXPECT_TRUE(IndexUtil::BumpIndices(shape, absl::MakeSpan(indices)));
149 EXPECT_THAT(indices, ::testing::ElementsAre(1, 1));
150 EXPECT_FALSE(IndexUtil::BumpIndices(shape, absl::MakeSpan(indices)));
151 }
152
153 } // namespace
154 } // namespace xla
155